array.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2025 VTT Technical Research Centre of Finland Ltd
// SPDX-License-Identifier: AGPL-3.0-or-later
 
#ifndef PFC_ARRAY_HPP
#define PFC_ARRAY_HPP
 
#include "core/decomposition.hpp"
#include "fft.hpp"
#include "multi_index.hpp"
#include "utils/array_to_string.hpp"
#include "utils/show.hpp"
#include "utils/typename.hpp"
#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <functional>
#include <type_traits>
#include <typeinfo>
#include <vector>
 
namespace pfc {
 
template <typename T, size_t D> class Array {
 
private:
  const MultiIndex<D> index;
  std::vector<T> data;
 
  Array(const FFT &fft, std::true_type)
      : index(get_outbox(fft).size, get_outbox(fft).low) {}
 
  Array(const FFT &fft, std::false_type)
      : index(get_inbox(fft).size, get_inbox(fft).low) {}
 
  // Custom type trait to check if a type is complex
  template <typename U> struct is_complex : std::false_type {};
  template <typename U> struct is_complex<std::complex<U>> : std::true_type {};
 
public:
  Array(const std::array<int, D> &dimensions,
        const std::array<int, D> &offsets = {0})
      : index(dimensions, offsets) {
    data.resize(index.get_linear_size());
  }
 
  Array(const Decomposition &decomp) : Array(decomp, is_complex<T>()) {
    data.resize(index.get_linear_size());
  }
 
  typename std::vector<T>::iterator begin() { return data.begin(); }
  typename std::vector<T>::iterator end() { return data.end(); }
  typename std::vector<T>::const_iterator begin() const { return data.begin(); }
  typename std::vector<T>::const_iterator end() const { return data.end(); }
 
  const MultiIndex<D> &get_index() const { return index; }
 
  std::vector<T> &get_data() { return data; }
 
  T &operator[](const std::array<int, D> &indices) {
    return operator[](index.to_linear(indices));
  }
 
  T &operator[](int idx) { return data.operator[](idx); }
 
  T &operator()(const std::array<int, D> &indices) { return operator[](indices); }
 
  std::array<int, D> get_size() const { return index.get_size(); }
 
  std::array<int, D> get_offset() const { return index.get_offset(); }
 
  bool inbounds(const std::array<int, D> &indices) {
    return index.inbounds(indices);
  }
 
  template <typename Func> void apply(Func &&func) {
    static_assert(
        std::is_convertible_v<std::invoke_result_t<Func, std::array<int, D>>, T>,
        "Func must be invocable with std::array<int, D> and return a type "
        "convertible to T");
    auto it = index.begin();
    for (T &element : get_data())
      element = std::invoke(std::forward<Func>(func), *(it++));
  }
 
  operator std::vector<T> &() { return data; }
 
  void set_data(const std::vector<T> &new_data) {
    if (new_data.size() != index.get_linear_size()) {
      throw std::runtime_error("Dimension mismatch, set_data failed");
    }
    this->data = std::move(new_data);
  }
 
  friend std::ostream &operator<<(std::ostream &os, const Array<T, D> &array) {
    const MultiIndex<D> &index = array.get_index();
    os << "Array<" << TypeName<T>::get() << "," << D
       << ">(begin = " << utils::array_to_string(index.get_begin())
       << ", end = " << utils::array_to_string(index.get_end())
       << ", size = " << utils::array_to_string(index.get_size())
       << ", linear_size = " << index.get_linear_size() << ")";
    return os;
  }
};
 
template <typename T, size_t D> void show(Array<T, D> &array) {
  utils::show(array.get_data(), array.get_index().get_size(),
              array.get_index().get_begin());
}
 
} // namespace pfc
 
#endif