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 <molpro/mpi.h>
#include <vector>
 
#include <molpro/linalg/array/Span.h>
 
namespace molpro::linalg::array {
namespace util {
template <typename Ind>
class Distribution;
}
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;
 
  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 value_type at(index_type ind) const = 0;
  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);
  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;
 
  value_type operator[](size_t index) { return (*this->local_buffer())[index]; };
 
protected:
  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);
} // namespace util
} // namespace molpro::linalg::array
 
#endif // LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_DISTRARRAY_H