lin_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 <nda/h5.hpp>
#include <nda/nda.hpp>
 
#include <functional>
#include <iterator>
#include <stdexcept>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
 
namespace triqs::stat {

  template <AccCompatible T> auto compress_bins(std::vector<T> const &bins, int fac) {
    // early return if the compression factor < 2
    if (fac < 2) return bins;
 
    // prepare the compressed bins
    int const nbins = bins.size() / fac;
    std::vector<T> res(nbins);
 
    // compress bins
    for (int i = 0; i < nbins; ++i) {
      res[i] = bins[fac * i];
      for (int j = 1; j < fac; ++j) { res[i] += bins[fac * i + j]; }
      res[i] /= fac;
    }
    return res;
  }

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

    using real_t = get_real_t<T>;

    using callback_t = std::function<void(lin_binning const &)>;

    lin_binning() = default;

    lin_binning(T const &sample, long max_n_bins, long bin_capacity, callback_t callback = {})
       : max_n_bins_(max_n_bins), callback_(std::move(callback)) {
      if (max_n_bins != -1 && max_n_bins < 2) throw std::runtime_error("lin_binning: max_n_bins must be -1 (unbounded) or >= 2.");
      if (bin_capacity < 1) throw std::runtime_error("lin_binning: bin_capacity must be >= 1.");
      bin_capacity_ = bin_capacity;
 
      if (max_n_bins > 0) mean_bins_.reserve(max_n_bins);
      mean_bins_.emplace_back(zeroed_sample(sample));
      mean_ = zeroed_sample(sample);
      var_  = make_real(zeroed_sample(sample));
    }

    [[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 mean_bins_.size(); }

    [[nodiscard]] long n_full_bins() const { return (last_bin_count_ == bin_capacity_ ? n_bins() : n_bins() - 1); }

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

    [[nodiscard]] std::vector<T> full_bins() const {
      return (last_bin_count_ == bin_capacity_ ? mean_bins_ : std::vector<T>(mean_bins_.begin(), std::prev(mean_bins_.end())));
    }

    [[nodiscard]] long bin_capacity() const { return bin_capacity_; }

    [[nodiscard]] long last_bin_count() const { return last_bin_count_; }

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

    [[nodiscard]] auto const &mean() const { return mean_; }

    [[nodiscard]] auto const &var_data() const { return var_; }

    [[nodiscard]] auto const &callback() const { return callback_; }

    template <typename U> auto &operator<<(U &&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));
        mean_ = zeroed_sample(s);
        var_  = make_real(zeroed_sample(s));
      }
 
      if (!is_unbounded() && n_full_bins() == max_n_bins_) {
        if (callback_) callback_(*this);
        compress(2);
      }
 
      var_ += abs_square(x - mean_) * (static_cast<double>(count_ - 1) / count_);
      mean_ += (x - mean_) / count_;
 
      if (last_bin_count_ == bin_capacity_) {
        mean_bins_.emplace_back(std::forward<U>(x));
        last_bin_count_ = 1;
      } else {
        ++last_bin_count_;
        mean_bins_.back() += (x - mean_bins_.back()) / last_bin_count_;
      }
 
      return *this;
    }

    void compress(int fac) {
      // early return if there is only 1 bin or if the compression factor < 2
      if (n_bins() == 1 || fac < 2) return;
 
      // compress full bins into new full bins
      int nbins = n_full_bins() / fac;
      for (int i = 0; i < nbins; ++i) {
        mean_bins_[i] = mean_bins_[fac * i];
        for (int j = 1; j < fac; ++j) { mean_bins_[i] += mean_bins_[fac * i + j]; }
        mean_bins_[i] /= fac;
      }
 
      // handle any left over bins
      int const left_over = n_bins() - fac * nbins;
      if (left_over != 0) {
        // currently active bin might not be full
        auto const bc   = last_bin_count_ + (left_over - 1) * bin_capacity_;
        value_t tmp_bin = mean_bins_.back() * (static_cast<double>(last_bin_count_) / bc);
        for (int j = 0; j < left_over - 1; ++j) { tmp_bin += mean_bins_[fac * nbins + j] * (static_cast<double>(bin_capacity_) / bc); }
        mean_bins_[nbins] = tmp_bin;
        last_bin_count_   = bc;
        ++nbins;
      } else {
        // no left overs, currently active bin is full w.r.t. the new bin capacity
        last_bin_count_ = bin_capacity_ * fac;
      }
 
      // resize bin vector and update bin capacity
      mean_bins_.resize(nbins);
      bin_capacity_ *= fac;
    }

    [[nodiscard]] auto mean_error_and_tau() const {
      auto [err, tau] = calculate_error_and_tau(mean_bins_, var_, count_);
      return std::make_tuple(mean_, err, tau);
    }

    [[nodiscard]] auto mean_error_and_tau(mpi::communicator c) const {
      auto [mean_red, var_red, count_red] = mpi_all_reduce(c);
      auto bins_gathered                  = mpi_all_gather(c, true);
      auto [err, tau]                     = calculate_error_and_tau(bins_gathered, var_red, count_red);
      return std::make_tuple(mean_red, err, tau);
    }

    [[nodiscard]] auto mpi_all_reduce(mpi::communicator c) const {
      auto count_red = mpi::all_reduce(count_, c);
 
      // skip the per-rank reweighting (and its 0/0) if no rank has samples
      if (count_red == 0) return std::make_tuple(zeroed_sample(mean_), make_real(zeroed_sample(mean_)), count_red);
 
      value_t mean_red = mean_ * (static_cast<double>(count_) / static_cast<double>(count_red));
      mpi::all_reduce_in_place(mean_red, c);
 
      real_t var_red = var_ + count_ * abs_square(mean_ - mean_red);
      mpi::all_reduce_in_place(var_red, c);
 
      return std::make_tuple(mean_red, var_red, count_red);
    }

    [[nodiscard]] auto mpi_all_gather(mpi::communicator c, bool same_capacity = true) const {
      auto fbins = full_bins();
 
      if (same_capacity) {
        auto bc_max = mpi::all_reduce(bin_capacity_, c, MPI_MAX);
        // ranks whose capacity does not divide the max cannot align their bins and contribute nothing
        if (bc_max % bin_capacity_ == 0)
          fbins = compress_bins(fbins, bc_max / bin_capacity_);
        else
          fbins.clear();
      }
 
      auto nbins         = mpi::all_gather(fbins.size(), c);
      auto bins_gathered = std::vector<value_t>(std::accumulate(nbins.begin(), nbins.end(), 0L));
      long start         = 0;
      for (int i = 0; i < c.size(); ++i) {
        auto const end = start + nbins[i];
        if (c.rank() == i) std::copy(fbins.begin(), fbins.end(), bins_gathered.begin() + start);
        mpi::broadcast_range(std::span{bins_gathered.begin() + start, bins_gathered.begin() + end}, c, i);
        start = end;
      }
 
      return bins_gathered;
    }

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

    friend void h5_write(h5::group g, std::string const &name, lin_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, "bin_capacity", acc.bin_capacity_);
      h5::write(gr, "last_bin_count", acc.last_bin_count_);
      h5::write(gr, "count", acc.count_);
      h5::write(gr, "mean_bins", acc.mean_bins_);
      h5::write(gr, "mean", acc.mean_);
      h5::write(gr, "var", acc.var_);
    }

    friend void h5_read(h5::group g, std::string const &name, lin_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, "bin_capacity", acc.bin_capacity_);
      h5::read(gr, "last_bin_count", acc.last_bin_count_);
      h5::read(gr, "count", acc.count_);
      h5::read(gr, "mean_bins", acc.mean_bins_);
      h5::read(gr, "mean", acc.mean_);
      h5::read(gr, "var", acc.var_);
 
      if (acc.max_n_bins_ != -1 && (acc.max_n_bins_ < 2 || acc.bin_capacity_ < 1))
        throw std::runtime_error("h5_read: invalid lin_binning state (max_n_bins must be -1 or >= 2, bin_capacity must be >= 1).");
    }
 
    private:
    // Get the mean, its standard error and an estimate for the integrated autocorrelation time.
    [[nodiscard]] auto calculate_error_and_tau(std::vector<value_t> const &bins, real_t const &q0, long ct) const {
      // need at least 2 samples and 2 bins to estimate the error and tau
      if (ct < 2 || bins.size() < 2) {
        auto nan = nan_sample(q0);
        return std::make_pair(nan, nan);
      }
 
      // unbinned estimate of the error of the mean
      real_t var0 = q0 / (static_cast<double>(ct) * static_cast<double>(ct - 1));
 
      // binned estimate of the error of the mean
      auto err = mean_and_err(bins).second;
 
      // estimate of the integrated autocorrelation time (NaN where var0 == 0)
      real_t tau = tau_estimate_from_vars(abs_square(err), var0);
 
      return std::make_pair(err, tau);
    }
 
    long max_n_bins_{-1};
    long bin_capacity_{1};
    long last_bin_count_{0};
    long count_{0};
    std::vector<value_t> mean_bins_{};
    value_t mean_{};
    real_t var_{};
    callback_t callback_{};
  };
 
} // namespace triqs::stat