Distribution.h

#ifndef LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_UTIL_DISTRIBUTION_H
#define LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_UTIL_DISTRIBUTION_H
#include <algorithm>
#include <cassert>
#include <numeric>
#include <utility>
#include <vector>
 
namespace molpro::linalg::array::util {
 
template <typename Ind = size_t>
class Distribution {
public:
  using index_type = Ind;
 
  Distribution() = default;
  Distribution(const Distribution<Ind> &) = default;
  Distribution(Distribution<Ind> &&) noexcept = default;
  Distribution<Ind> &operator=(const Distribution<Ind> &) = default;
  Distribution<Ind> &operator=(Distribution<Ind> &&) noexcept = default;
  ~Distribution() = default;
 
  friend void swap(Distribution<index_type> &l, Distribution<index_type> &r) {
    using std::swap;
    swap(l.m_chunk_borders, r.m_chunk_borders);
  }
 
  Distribution(const std::vector<index_type> &indices) : m_chunk_borders{indices} {
    assert("indices must be sorted" && std::is_sorted(begin(m_chunk_borders), end(m_chunk_borders)));
  };
 
  std::pair<int, int> cover(index_type lo, index_type hi) const {
    assert(lo < hi);
    return {cover(lo), cover(hi - 1)};
  };
 
  int cover(index_type ind) const {
    if (ind < border().first || ind > border().second)
      return size();
    auto iter = std::upper_bound(begin(m_chunk_borders), end(m_chunk_borders), ind);
    int chunk = iter == begin(m_chunk_borders) ? size() : distance(begin(m_chunk_borders), prev(iter));
    return chunk;
  };
 
  std::pair<index_type, index_type> range(int chunk) const {
    return {m_chunk_borders[chunk], m_chunk_borders[chunk + 1]};
  };
 
  std::pair<index_type, index_type> border() const { return {m_chunk_borders.front(), m_chunk_borders.back()}; }
 
  size_t size() const { return m_chunk_borders.size() ? m_chunk_borders.size() - 1 : 0; };
 
  const std::vector<index_type> &chunk_borders() const { return m_chunk_borders; }
 
  bool compatible(const Distribution<Ind> &other) const {
    return size() == other.size() &&
           std::inner_product(begin(m_chunk_borders), end(m_chunk_borders), begin(other.m_chunk_borders), true,
                              std::logical_and<>(), std::equal_to<>());
  }
 
protected:
  std::vector<index_type> m_chunk_borders;
};
 
template <typename Ind>
Distribution<Ind> make_distribution_spread_remainder(size_t dimension, int n_chunks) {
  assert(n_chunks > 0);
  auto chunk_borders = std::vector<Ind>{0};
  chunk_borders.reserve(n_chunks + 1);
  auto block = dimension / n_chunks;
  auto n_extra = dimension % n_chunks;
  std::fill_n(std::back_inserter(chunk_borders), n_extra, block + 1);
  std::fill_n(std::back_inserter(chunk_borders), n_chunks - n_extra, block);
  std::partial_sum(begin(chunk_borders), end(chunk_borders), begin(chunk_borders));
  return {chunk_borders};
}
 
} // namespace molpro::linalg::array::util
 
#endif // LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_UTIL_DISTRIBUTION_H