mean_error.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2021 Simons Foundation
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You may obtain a copy of the License at
//     https://www.gnu.org/licenses/gpl-3.0.txt
//
// Authors: Philipp Dumitrescu, Olivier Parcollet, Nils Wentzell

 
#pragma once
 
#include "./concepts.hpp"
#include "./utils.hpp"
 
#include <itertools/itertools.hpp>
#include <mpi/mpi.hpp>
#include <nda/mpi.hpp>
#include <nda/nda.hpp>
 
#include <numeric>
#include <optional>
#include <ranges>
#include <utility>
 
namespace triqs::stat {


  enum class mean_tag { sum, mean };

  enum class error_tag { sum, var_data, var_mean, err_data, err_mean, jk_err };

  template <mean_tag mtag, AccCompatible T> void apply_mean_tag(T &m, [[maybe_unused]] long nsamples) {
    if constexpr (mtag == mean_tag::sum) m *= nsamples;
  }

  template <error_tag etag, AccCompatible T> void apply_error_tag(T &sum_sq_devs, [[maybe_unused]] long nsamples) {
    if constexpr (etag == error_tag::sum) return;
    auto const nd = static_cast<double>(nsamples);
    if constexpr (etag == error_tag::err_data || etag == error_tag::var_data)
      sum_sq_devs /= (nd - 1);
    else if constexpr (etag == error_tag::err_mean || etag == error_tag::var_mean)
      sum_sq_devs /= (nd * (nd - 1));
    else if constexpr (etag == error_tag::jk_err)
      sum_sq_devs *= (nd - 1) / nd;
    if constexpr (etag == error_tag::err_data || etag == error_tag::err_mean || etag == error_tag::jk_err) sum_sq_devs = nda::sqrt(sum_sq_devs);
  }

  template <mean_tag mtag = mean_tag::mean, StatCompatibleRange R> auto mean(R &&rg) { // NOLINT (ranges should not be forwarded)
    if constexpr (mtag == mean_tag::mean) {
      // calculate the arithmetic mean
      auto res = zeroed_sample(*std::ranges::begin(rg));
      for (auto const &[n, x] : itertools::enumerate(rg)) res += (x - res) / (n + 1);
      return res;
    } else {
      // calculate the simple sum
      return std::accumulate(std::ranges::begin(rg), std::ranges::end(rg), zeroed_sample(*std::ranges::begin(rg)));
    }
  }

  template <mean_tag mtag = mean_tag::mean, StatCompatibleRange R>
  auto mean_mpi(std::optional<mpi::communicator> c, R &&rg) { // NOLINT (ranges should not be forwarded)
    using value_t = get_regular_t<std::ranges::range_value_t<R>>;
 
    // early return if no communicator is provided
    if (!c) return mean<mtag>(rg);
 
    value_t res = mean<mtag>(rg);
    if constexpr (mtag == mean_tag::mean) {
      // for mtag == mean_tag::mean, we need to take care of different sample sizes
      auto const n_i = std::ranges::size(rg);
      auto const n   = mpi::all_reduce(n_i, *c);
      res *= static_cast<double>(n_i) / static_cast<double>(n);
    }
    mpi::all_reduce_in_place(res, *c);
    return res;
  }

  template <error_tag etag = error_tag::err_mean, mean_tag mtag = mean_tag::mean, StatCompatibleRange R>
  auto mean_and_err(R &&rg) { // NOLINT (ranges should not be forwarded)
    // calculate the arithmetic mean and the sum of squared deviations from the mean
    auto res_m = zeroed_sample(*std::ranges::begin(rg));
    auto res_s = make_real(res_m);
    for (auto const &[n, x] : itertools::enumerate(rg)) {
      auto const nd = static_cast<double>(n);
      res_s += abs_square(x - res_m) * nd / (nd + 1);
      res_m += (x - res_m) / (nd + 1);
    }
    // apply the mean and error tags
    apply_mean_tag<mtag>(res_m, std::ranges::size(rg));
    apply_error_tag<etag>(res_s, std::ranges::size(rg));
    return std::make_pair(res_m, res_s);
  }

  template <error_tag etag = error_tag::err_mean, mean_tag mtag = mean_tag::mean, StatCompatibleRange R>
  auto mean_and_err_mpi(std::optional<mpi::communicator> c, R &&rg) { // NOLINT (ranges should not be forwarded)
    using value_t = get_regular_t<std::ranges::range_value_t<R>>;
 
    // early return if no communicator is provided
    if (!c) return mean_and_err<etag, mtag>(rg);
 
    // local mean and sum of squared deviations from the mean
    auto [m, ssqdev] = mean_and_err<error_tag::sum>(rg);
 
    // reduce the sample size
    auto const n     = std::ranges::size(rg);
    auto const n_red = mpi::all_reduce(n, *c);
 
    // reduce the mean
    value_t res_m = m * (static_cast<double>(n) / static_cast<double>(n_red));
    mpi::all_reduce_in_place(res_m, *c);
 
    // reduce the sum of squared deviations from the mean
    ssqdev += n * abs_square(m - res_m);
    mpi::all_reduce_in_place(ssqdev, *c);
 
    // apply the mean and error tags
    apply_mean_tag<mtag>(res_m, n_red);
    apply_error_tag<etag>(ssqdev, n_red);
    return std::make_pair(res_m, ssqdev);
  }

  [[nodiscard]] auto tau_estimate_from_vars(auto const &var, auto const &var0) {
    return nda::make_regular(nda::map([](auto v, auto v0) { return (v0 == 0.0) ? nan_sample(v0) : 0.5 * (v / v0 - 1.0); })(var, var0));
  }

  template <StatCompatible T> auto tau_estimate_from_errors(T const &s_n, T const &s_0) {
    return tau_estimate_from_vars(abs_square(s_n), abs_square(s_0));
  }

 
} // namespace triqs::stat