log_binning.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2022 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 "./mean_error.hpp"
#include "./utils.hpp"
 
#include <h5/h5.hpp>
#include <mpi/mpi.hpp>
#include <nda/h5.hpp>
#include <nda/mpi.hpp>
#include <nda/nda.hpp>
 
#include <functional>
#include <iterator>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>
 
namespace triqs::stat {

  template <AccCompatible T> class log_binning {
    public:
    using value_t = T;

    using real_t = get_real_t<T>;

    log_binning() = default;

    log_binning(T const &sample, int max_n_bins) : max_n_bins_(max_n_bins) {
      if (max_n_bins != -1 && max_n_bins < 1) throw std::runtime_error("log_binning: max_n_bins must be -1 (unbounded) or >= 1.");
 
      if (max_n_bins > 0) {
        mean_bins_.reserve(max_n_bins);
        var_bins_.reserve(max_n_bins);
        // bare bins start at block size 2; the size-1 bin is folded into mean_bins_/var_bins_ directly
        bare_bins_.reserve(max_n_bins - 1);
        bare_counts_.reserve(max_n_bins - 1);
      }
 
      mean_bins_.emplace_back(zeroed_sample(sample));
      var_bins_.emplace_back(make_real(zeroed_sample(sample)));
      if (max_n_bins != 1) {
        bare_bins_.emplace_back(zeroed_sample(sample));
        bare_counts_.push_back(0);
      }
    }

    [[nodiscard]] long max_n_bins() const { return max_n_bins_; }

    [[nodiscard]] bool is_unbounded() const { return max_n_bins_ == -1; }

    [[nodiscard]] long n_bins() const { return var_bins_.size(); }

    [[nodiscard]] long count() const { return count_; }

    [[nodiscard]] auto effective_counts() const {
      std::vector<long> res(n_bins());
      for (int i = 0; i < n_bins(); ++i) res[i] = count_ >> i;
      return res;
    }

    [[nodiscard]] auto const &mean_bins() const { return mean_bins_; }

    [[nodiscard]] auto const &var_bins() const { return var_bins_; }

    [[nodiscard]] auto const &bare_bins() const { return bare_bins_; }

    [[nodiscard]] auto const &bare_counts() const { return bare_counts_; }

    template <typename U> log_binning<T> &operator<<(U const &x) {
      ++count_;
 
      // seed shape from the first sample for default-constructed accumulators
      if (mean_bins_.empty()) {
        T s = x;
        mean_bins_.emplace_back(zeroed_sample(s));
        var_bins_.emplace_back(make_real(zeroed_sample(s)));
        if (max_n_bins_ != 1) {
          bare_bins_.emplace_back(zeroed_sample(s));
          bare_counts_.push_back(0);
        }
      }
 
      T x_m = (x - mean_bins_[0]);
      var_bins_[0] += abs_square(x_m) * (static_cast<double>(count_ - 1) / count_);
      mean_bins_[0] += x_m / count_;
 
      if (max_n_bins_ == 1) return *this;
 
      if (count_ == (1L << bare_bins_.size()) && (is_unbounded() || n_bins() < max_n_bins_)) {
        mean_bins_.emplace_back(zeroed_sample(mean_bins_[0]));
        var_bins_.emplace_back(zeroed_sample(var_bins_[0]));
 
        // n_bins() has just been incremented; re-check whether another mean bin could still follow
        if (is_unbounded() || n_bins() < max_n_bins_) {
          bare_bins_.emplace_back(zeroed_sample(bare_bins_[0]));
          bare_counts_.push_back(0);
        }
      }
 
      bare_bins_[0] += x;
      ++bare_counts_[0];
 
      for (int i = 0; i < bare_bins_.size() && bare_counts_[i] == 2; ++i) {
        if (i + 1 < bare_bins_.size()) {
          bare_bins_[i + 1] += bare_bins_[i];
          ++bare_counts_[i + 1];
        }
 
        auto const bc = static_cast<double>(1ul << (i + 1));
        auto const k  = count_ / bc;
        x_m           = (bare_bins_[i] / bc - mean_bins_[i + 1]);
        var_bins_[i + 1] += abs_square(x_m) * ((k - 1) / k);
        mean_bins_[i + 1] += x_m / k;
 
        bare_counts_[i] = 0;
        bare_bins_[i]   = 0;
      }
 
      return *this;
    }

    [[nodiscard]] auto mean_errors_and_taus(int min_samples = 2) const {
      return calculate_mean_errors_and_taus(mean_bins_, var_bins_, effective_counts(), min_samples);
    }

    [[nodiscard]] auto mean_errors_and_taus(mpi::communicator c, int min_samples = 2) const {
      auto [mean_red, var_red, nsamples_red] = mpi_all_reduce(c);
      return calculate_mean_errors_and_taus(mean_red, var_red, nsamples_red, min_samples);
    }

    [[nodiscard]] auto mpi_all_reduce(mpi::communicator c) const {
      auto mean_red     = mean_bins_;
      auto var_red      = var_bins_;
      auto nsamples_red = effective_counts();
 
      // make sure all ranks have the same number of bins
      auto const nbins = mpi::all_reduce(n_bins(), c, MPI_MAX);
      mean_red.resize(nbins, zeroed_sample(mean_bins_[0]));
      var_red.resize(nbins, zeroed_sample(var_bins_[0]));
      nsamples_red.resize(nbins, 0);
 
      for (int i = 0; i < nbins; ++i) {
        auto const ns = nsamples_red[i];
        mpi::all_reduce_in_place(nsamples_red[i], c);
 
        // skip the per-rank reweighting (and its 0/0) if no rank has samples at this bin level
        if (nsamples_red[i] > 0) mean_red[i] *= static_cast<double>(ns) / static_cast<double>(nsamples_red[i]);
        mpi::all_reduce_in_place(mean_red[i], c);
 
        if (i < mean_bins_.size()) var_red[i] += ns * abs_square(mean_bins_[i] - mean_red[i]);
        mpi::all_reduce_in_place(var_red[i], c);
      }
 
      return std::make_tuple(mean_red, var_red, nsamples_red);
    }

    [[nodiscard]] static std::string hdf5_format() { return "log_binning"; }

    friend void h5_write(h5::group g, std::string const &name, log_binning const &acc) {
      auto gr = g.create_group(name);
      h5::write_hdf5_format(gr, acc);
      h5::write(gr, "max_n_bins", acc.max_n_bins_);
      h5::write(gr, "count", acc.count_);
      h5::write(gr, "mean_bins", acc.mean_bins_);
      h5::write(gr, "var_bins", acc.var_bins_);
      h5::write(gr, "bare_bins", acc.bare_bins_);
      h5::write(gr, "bare_counts", acc.bare_counts_);
    }

    friend void h5_read(h5::group g, std::string const &name, log_binning &acc) {
      auto gr = g.open_group(name);
      h5::assert_hdf5_format(gr, acc);
      h5::read(gr, "max_n_bins", acc.max_n_bins_);
      h5::read(gr, "count", acc.count_);
      h5::read(gr, "mean_bins", acc.mean_bins_);
      h5::read(gr, "var_bins", acc.var_bins_);
      h5::read(gr, "bare_bins", acc.bare_bins_);
      h5::read(gr, "bare_counts", acc.bare_counts_);
 
      if (acc.max_n_bins_ != -1 && acc.max_n_bins_ < 1)
        throw std::runtime_error("h5_read: invalid log_binning state (max_n_bins must be -1 or >= 1).");
    }
 
    private:
    // Get the overall mean as well as the standard error, integrated autocorrelation time and effective number of
    // samples for each bin with at least a given minimum number of samples.
    [[nodiscard]] auto calculate_mean_errors_and_taus(std::vector<value_t> const &mk, std::vector<real_t> const &qk,
                                                      std::vector<long> const &nsamples, int min_samples) const {
      // early return for empty vectors
      if (mk.empty()) return std::make_tuple(value_t{}, std::vector<real_t>{}, std::vector<real_t>{}, std::vector<long>{});
 
      // only consider bins with at least min_samples effective samples
      auto const size = std::distance(nsamples.begin(), std::ranges::upper_bound(nsamples, min_samples, std::greater{}));
      auto effs       = std::vector(nsamples.begin(), nsamples.begin() + size);
      auto errs       = std::vector(qk.begin(), qk.begin() + size);
      auto taus       = std::vector<real_t>(size);
 
      // calculate errors and taus
      if (size > 0) {
        real_t var0 = qk[0] / (static_cast<double>(nsamples[0]) * static_cast<double>(nsamples[0] - 1));
        for (int i = 0; i < size; ++i) {
          real_t var = errs[i] / (static_cast<double>(effs[i]) * static_cast<double>(effs[i] - 1));
          errs[i]    = nda::sqrt(var);
          taus[i]    = tau_estimate_from_vars(var, var0);
        }
      }
      return std::make_tuple(mk[0], errs, taus, effs);
    }
 
    long max_n_bins_{-1};
    long count_{0};
    std::vector<value_t> mean_bins_{};
    std::vector<real_t> var_bins_{};
    std::vector<value_t> bare_bins_{};
    std::vector<int> bare_counts_{};
  };
 
} // namespace triqs::stat