DistrArray.h

#ifndef LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_DISTRARRAY_H
#define LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_DISTRARRAY_H
#include <cmath>
#include <cstdlib>
#include <list>
#include <map>
#include <memory>
#include <string>
#include <molpro/mpi.h>
#include <vector>
#include <type_traits>
#include <compare>
#include <iterator>
#include <ranges>
 
#include <molpro/linalg/array/Span.h>
#include <molpro/linalg/array/util/Distribution.h>
 
namespace molpro::linalg::array {
namespace util {
  template<typename Array>
  class ValueProxy {
  public:
    ValueProxy(const ValueProxy &) = delete;
    ValueProxy(ValueProxy &&) = delete;
 
    // Assigning a proxy writes the *value* it refers to into this proxy's cell,
    // rather than rebinding this proxy to a different array/index -- that keeps
    // *it1 = *it2 (as used by std::copy et al.) working between two DistrArrays.
    ValueProxy &operator=(const ValueProxy &other) { return *this = static_cast<typename Array::value_type>(other); }
    ValueProxy &operator=(ValueProxy &&other) { return *this = static_cast<typename Array::value_type>(other); }
    template <typename OtherArray>
    ValueProxy &operator=(const ValueProxy<OtherArray> &other) {
      return *this = static_cast<typename OtherArray::value_type>(other);
    }
 
    ValueProxy &operator=(Array::value_type val) {
      m_arr.set(m_idx, std::move(val));
 
      return *this;
    }
 
    operator typename Array::value_type() const {
      return m_arr.at(m_idx);
    }
 
  private:
    std::add_lvalue_reference_t<Array> m_arr;
    Array::index_type m_idx;
 
    friend std::remove_const_t<Array>;
 
    ValueProxy(std::add_lvalue_reference_t<Array> arr, Array::index_type idx) : m_arr(arr), m_idx(idx) {}
  };
}
class DistrArray {
public:
  using distributed_array = void; 
  using value_type = double;
  using index_type = size_t;
  using SparseArray = std::map<index_type, double>;
  using Distribution = util::Distribution<index_type>;
 
protected:
  index_type m_dimension = 0; 
  MPI_Comm m_communicator;    
  DistrArray(size_t dimension, MPI_Comm commun);
  DistrArray() = default;
 
public:
  virtual ~DistrArray() = default;
  // Declared explicitly because the user-declared destructor suppresses the implicit move operations
  DistrArray(const DistrArray&) = default;
  DistrArray(DistrArray&&) noexcept = default;
  DistrArray& operator=(const DistrArray&) = default;
  DistrArray& operator=(DistrArray&&) noexcept = default;
 
  MPI_Comm communicator() const { return m_communicator; }
  virtual void sync() const;
  size_t size() const { return m_dimension; };
  bool compatible(const DistrArray &other) const;
 
protected:
  class LocalBuffer : public span::Span<value_type> {
  public:
    virtual ~LocalBuffer() = default;
    using span::Span<value_type>::Span;
    friend void swap(LocalBuffer &, LocalBuffer &) = delete;
    bool compatible(const LocalBuffer &other) const { return start() == other.start() && size() == other.size(); };
    size_type start() const { return m_start; }
 
  protected:
    size_type m_start = 0; 
  };
 
public:
  [[nodiscard]] virtual std::unique_ptr<LocalBuffer> local_buffer() = 0;
  [[nodiscard]] virtual std::unique_ptr<const LocalBuffer> local_buffer() const = 0;
  [[nodiscard]] virtual const Distribution &distribution() const = 0;
  [[nodiscard]] virtual size_t disk_page_size() const { return 0; }
 
  [[nodiscard]] virtual value_type at(index_type ind) const = 0;
  [[nodiscard]] value_type operator[](index_type ind) const;
  [[nodiscard]] util::ValueProxy<DistrArray> operator[](index_type ind);
  virtual void set(index_type ind, value_type val) = 0;
  virtual void get(index_type lo, index_type hi, value_type *buf) const = 0;
  [[nodiscard]] virtual std::vector<value_type> get(index_type lo, index_type hi) const = 0;
  virtual void put(index_type lo, index_type hi, const value_type *data) = 0;
  virtual void acc(index_type lo, index_type hi, const value_type *data) = 0;
  [[nodiscard]] virtual std::vector<value_type> gather(const std::vector<index_type> &indices) const = 0;
  virtual void scatter(const std::vector<index_type> &indices, const std::vector<value_type> &data) = 0;
  virtual void scatter_acc(std::vector<index_type> &indices, const std::vector<value_type> &data) = 0;
  [[nodiscard]] virtual std::vector<value_type> vec() const = 0;
 
  virtual void fill(value_type val);
  virtual void copy(const DistrArray &y);
  virtual void copy_patch(const DistrArray &y, index_type start, index_type end);
  virtual void axpy(value_type a, const DistrArray &y);
  virtual void axpy(value_type a, const SparseArray &y);
  void copy(const SparseArray &y);
 
  void copy(Span<const value_type> y);
  void axpy(value_type a, Span<const value_type> y);
  void copy_to(Span<value_type> y) const;
  void axpy_to(value_type a, Span<value_type> y) const;
  [[nodiscard]] value_type dot(Span<const value_type> y) const;
  [[nodiscard]] std::map<size_t, value_type> select_max_dot(size_t n, Span<const value_type> y) const;
 
  virtual void scal(value_type a);
  virtual void add(const DistrArray &y);
  virtual void add(value_type a);
  virtual void sub(const DistrArray &y);
  virtual void sub(value_type a);
  virtual void recip();
  virtual void times(const DistrArray &y);
  virtual void times(const DistrArray &y, const DistrArray &z);
 
 
  [[nodiscard]] virtual value_type dot(const DistrArray &y) const;
  [[nodiscard]] virtual value_type dot(const SparseArray &y) const;
 
  void divide(const DistrArray &y, const DistrArray &z, value_type shift = 0, bool append = false,
              bool negative = false) {
    _divide(y, z, shift, append, negative);
  }
 
  [[nodiscard]] std::list<std::pair<index_type, value_type>> min_n(int n) const;
 
  [[nodiscard]] std::list<std::pair<index_type, value_type>> max_n(int n) const;
 
  [[nodiscard]] std::list<std::pair<index_type, value_type>> min_abs_n(int n) const;
 
  [[nodiscard]] std::list<std::pair<index_type, value_type>> max_abs_n(int n) const;
 
  [[nodiscard]] std::vector<index_type> min_loc_n(int n) const;
 
  [[nodiscard]] std::map<size_t, value_type> select_max_dot(size_t n, const DistrArray &y) const;
  [[nodiscard]] std::map<size_t, value_type> select_max_dot(size_t n, const SparseArray &y) const;
  [[nodiscard]] std::map<size_t, value_type> select(size_t n, bool max = false, bool ignore_sign = false) const;
 
  virtual void zero();
 
  virtual void error(const std::string &message) const;
 
protected:
  void check_compatible(Span<const value_type> y, const std::string &name) const;
  virtual void _divide(const DistrArray &y, const DistrArray &z, value_type shift, bool append, bool negative);
};
 
namespace util {
template <typename T, class Compare>
struct CompareAbs {
  constexpr bool operator()(const T &lhs, const T &rhs) const { return Compare()(std::abs(lhs), std::abs(rhs)); }
};
template <class Compare>
[[nodiscard]] std::list<std::pair<DistrArray::index_type, DistrArray::value_type>> extrema(const DistrArray &x, int n);
std::map<size_t, double> select_max_dot_broadcast(size_t n, std::map<size_t, double> &local_selection,
                                                  MPI_Comm communicator);
#ifdef HAVE_MPI_H
inline MPI_Datatype mpi_index_type() {
  static_assert(sizeof(DistrArray::index_type) == sizeof(unsigned long) ||
                sizeof(DistrArray::index_type) == sizeof(unsigned long long));
  return sizeof(DistrArray::index_type) == sizeof(unsigned long) ? MPI_UNSIGNED_LONG : MPI_UNSIGNED_LONG_LONG;
}
#endif
 
template<bool is_const>
class DistrArrayIteratorImpl {
  public:
    using value_type = std::conditional_t<
      is_const, DistrArray::value_type, ValueProxy<DistrArray>>;
    using reference = value_type;
    using const_reference = reference;
    using difference_type = std::ptrdiff_t;
    // *this dereferences to a proxy, not an lvalue reference, so there is no meaningful
    // pointer type. iterator_category is spelled out explicitly because std::iterator_traits'
    // legacy (pre-C++20) fallback cannot deduce past input_iterator_tag for a proxy iterator,
    // which would make std::fill/std::copy etc. reject the single-pass-only overloads even
    // though this iterator satisfies std::random_access_iterator (see the static_assert below).
    using pointer = void;
    using iterator_category = std::random_access_iterator_tag;
 
    DistrArrayIteratorImpl() = default;
    DistrArrayIteratorImpl(std::conditional_t<is_const, const DistrArray, DistrArray> &array,
        DistrArray::index_type pos) : m_arr(&array), m_pos(pos) {}
 
    DistrArrayIteratorImpl &operator+=(difference_type val) {
      m_pos += val;
      return *this;
    }
    friend DistrArrayIteratorImpl operator+(const DistrArrayIteratorImpl &it, difference_type val) {
      return DistrArrayIteratorImpl(*it.m_arr, it.m_pos + val);
    }
    friend DistrArrayIteratorImpl operator+(difference_type val, const DistrArrayIteratorImpl &it) {
      return it + val;
    }
    DistrArrayIteratorImpl &operator++() {
      ++m_pos;
      return *this;
    }
    DistrArrayIteratorImpl operator++(int) {
      DistrArrayIteratorImpl copy(*this);
      ++m_pos;
      return copy;
    }
 
    DistrArrayIteratorImpl &operator-=(difference_type val) {
      m_pos -= val;
      return *this;
    }
    friend DistrArrayIteratorImpl operator-(const DistrArrayIteratorImpl &it, difference_type val) {
      return DistrArrayIteratorImpl(*it.m_arr, it.m_pos - val);
    }
    DistrArrayIteratorImpl &operator--() {
      --m_pos;
      return *this;
    }
    DistrArrayIteratorImpl operator--(int) {
      DistrArrayIteratorImpl copy(*this);
      --m_pos;
      return copy;
    }
 
    reference operator*() const { return (*m_arr)[m_pos]; }
 
    reference operator[](difference_type offset) const { return (*m_arr)[m_pos + offset]; }
 
    difference_type operator-(const DistrArrayIteratorImpl &other) const { return m_pos - other.m_pos; }
 
    bool operator==(const DistrArrayIteratorImpl &) const = default;
    std::strong_ordering operator<=>(const DistrArrayIteratorImpl &) const = default;
 
  private:
    std::conditional_t<is_const, const DistrArray, DistrArray> *m_arr = nullptr;
    DistrArray::index_type m_pos = 0;
};
 
static_assert(std::random_access_iterator<DistrArrayIteratorImpl<true>>);
static_assert(std::random_access_iterator<DistrArrayIteratorImpl<false>>);
 
} // namespace util
 
using DistrArrayIterator = util::DistrArrayIteratorImpl<false>;
using DistrArrayConstIterator = util::DistrArrayIteratorImpl<true>;
 
inline DistrArrayIterator begin(DistrArray &array) { return DistrArrayIterator(array, 0); }
inline DistrArrayIterator end(DistrArray &array) { return DistrArrayIterator(array, array.size()); }
inline DistrArrayConstIterator begin(const DistrArray &array) { return DistrArrayConstIterator(array, 0); }
inline DistrArrayConstIterator end(const DistrArray &array) { return DistrArrayConstIterator(array, array.size()); }
inline DistrArrayConstIterator cbegin(const DistrArray &array) { return begin(array); }
inline DistrArrayConstIterator cend(const DistrArray &array) { return end(array); }
 
} // namespace molpro::linalg::array
 
#endif // LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_DISTRARRAY_H