discrete_field.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_DISCRETE_FIELD_HPP
#define PFC_DISCRETE_FIELD_HPP
 
#include "array.hpp"
#include "constants.hpp"
#include "openpfc/core/world.hpp"
#include "utils/show.hpp"
#include <array>
#include <cmath>
#include <cstddef>
#include <functional>
#include <ostream>
 
namespace pfc {
 
template <typename T, size_t D> class DiscreteField {
private:
  Array<T, D> m_array; 
  const std::array<double, D> m_origin; 
  const std::array<double, D>
      m_discretization;                      
  const std::array<double, D> m_coords_low;  
  const std::array<double, D> m_coords_high; 
  std::array<double, D> calculate_coords_low(const std::array<int, D> &offset) {
    std::array<double, D> coords_low;
    for (size_t i = 0; i < D; i++)
      coords_low[i] = m_origin[i] + offset[i] * m_discretization[i];
    return coords_low;
  }
 
  std::array<double, D> calculate_coords_high(const std::array<int, D> &offset,
                                              const std::array<int, D> &size) {
    std::array<double, D> coords_high;
    for (size_t i = 0; i < D; i++)
      coords_high[i] = m_origin[i] + (offset[i] + size[i]) * m_discretization[i];
    return coords_high;
  }
 
public:
  DiscreteField(const std::array<int, D> &dimensions,
                const std::array<int, D> &offsets,
                const std::array<double, D> &origin,
                const std::array<double, D> &discretization)
      : m_array(dimensions, offsets), m_origin(origin),
        m_discretization(discretization),
        m_coords_low(calculate_coords_low(offsets)),
        m_coords_high(calculate_coords_high(offsets, dimensions)) {}
 
  /* Note: Free function alternative for construction from Decomposition
   *
   * This constructor was considered but deferred in favor of explicit construction
   * to maintain clarity about what data is being used. Users can construct
   * DiscreteField directly with the required parameters.
   *
   * If needed in the future, a free function could be added:
   *
   * namespace pfc {
   *   template<int D>
   *   DiscreteField<D> make_discrete_field(const Decomposition& decomp) {
   *     return DiscreteField<D>(
   *       get_inbox_size(decomp),
   *       get_inbox_offset(decomp),
   *       get_origin(decomp.get_world()),
   *       get_spacing(decomp.get_world())
   *     );
   *   }
   * }
   */
 
  const std::array<double, D> &get_origin() const { return m_origin; }
  const std::array<double, D> &get_discretization() const {
    return m_discretization;
  }
  const std::array<double, D> &get_coords_low() const { return m_coords_low; }
  const std::array<double, D> &get_coords_high() const { return m_coords_high; }
 
  Array<T, D> &get_array() { return m_array; }
  const Array<T, D> &get_array() const { return m_array; }
  const MultiIndex<D> &get_index() { return get_array().get_index(); }
  const MultiIndex<D> &get_index() const { return get_array().get_index(); }
  std::vector<T> &get_data() { return get_array().get_data(); }
 
  T &operator[](const std::array<int, D> &indices) { return get_array()[indices]; }
 
  T &operator[](size_t idx) { return get_array()[idx]; }
 
  std::array<double, D>
  map_indices_to_coordinates(const std::array<int, D> &indices) const {
    std::array<double, D> coordinates;
    for (size_t i = 0; i < D; ++i) {
      coordinates[i] = m_origin[i] + indices[i] * m_discretization[i];
    }
    return coordinates;
  }
 
  std::array<int, D>
  map_coordinates_to_indices(const std::array<double, D> &coordinates) const {
    std::array<int, D> indices;
    for (size_t i = 0; i < D; ++i) {
      indices[i] = static_cast<int>(
          std::round((coordinates[i] - m_origin[i]) / m_discretization[i]));
    }
    return indices;
  }
 
  bool inbounds(const std::array<double, D> &coords) {
    for (size_t i = 0; i < D; i++) {
      if (m_coords_low[i] > coords[i] || coords[i] >= m_coords_high[i]) return false;
    }
    return true;
  }
 
  [[deprecated("Use pfc::interpolate(field, coords) free function instead")]] T &
  interpolate(const std::array<double, D> &coordinates) {
    // Keep original implementation (can't call free function yet - not declared)
    return get_array()[(map_coordinates_to_indices(coordinates))];
  }
 
  using FunctionND = std::function<T(std::array<double, D>)>;
  using Function1D = std::function<T(double)>;
  using Function2D = std::function<T(double, double)>;
  using Function3D = std::function<T(double, double, double)>;
 
  void apply(FunctionND &&func) {
    static_assert(
        std::is_convertible_v<
            std::invoke_result_t<FunctionND, std::array<double, D>>, T>,
        "Func must be invocable with std::array<double, D> and return a type "
        "convertible to T");
    auto index_iterator = get_index().begin();
    for (T &element : get_data()) {
      element = std::invoke(std::forward<FunctionND>(func),
                            map_indices_to_coordinates(*index_iterator));
      ++index_iterator;
    }
  }
 
  void apply(Function1D &&func) {
    auto index_iterator = get_index().begin();
    for (T &element : get_data()) {
      const auto coords = map_indices_to_coordinates(*index_iterator);
      element = std::invoke(std::forward<Function1D>(func), coords[0]);
      ++index_iterator;
    }
  }
 
  void apply(Function2D &&func) {
    auto index_iterator = get_index().begin();
    for (T &element : get_data()) {
      const auto [x, y] = map_indices_to_coordinates(*index_iterator);
      element = std::invoke(std::forward<Function2D>(func), x, y);
      ++index_iterator;
    }
  }
 
  void apply(Function3D &&func) {
    auto index_iterator = get_index().begin();
    for (T &element : get_data()) {
      const auto [x, y, z] = map_indices_to_coordinates(*index_iterator);
      element = std::invoke(std::forward<Function3D>(func), x, y, z);
      ++index_iterator;
    }
  }
 
  operator std::vector<T> &() { return get_data(); }
  operator std::vector<T> &() const { return get_data(); }
 
  const std::array<int, D> &get_size() const { return get_index().get_size(); }
 
  const std::array<int, D> &get_offset() const { return get_index().get_begin(); }
 
  void set_data(const std::vector<T> &data) { get_array().set_data(data); }
 
  friend std::ostream &operator<<(std::ostream &os,
                                  const DiscreteField<T, D> &field) {
    const Array<T, D> &array = field.get_array();
    const MultiIndex<D> &index = array.get_index();
    os << "DiscreteField<" << 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()
       << ", origin = " << utils::array_to_string(field.get_origin())
       << ", discretization = " << utils::array_to_string(field.get_discretization())
       << ", coords_low = " << utils::array_to_string(field.get_coords_low())
       << ", coords_high = " << utils::array_to_string(field.get_coords_high());
    return os;
  }
};
 
// ============================================================================
// Free Functions for DiscreteField
// ============================================================================
 
template <typename T, size_t D>
inline T &interpolate(DiscreteField<T, D> &field,
                      const std::array<double, D> &coordinates) {
  return field.get_array()[field.map_coordinates_to_indices(coordinates)];
}
 
template <typename T, size_t D>
inline const T &interpolate(const DiscreteField<T, D> &field,
                            const std::array<double, D> &coordinates) {
  // Note: const_cast needed because Array doesn't have const operator[]
  // This is safe because we return const T&
  return const_cast<DiscreteField<T, D> &>(field)
      .get_array()[field.map_coordinates_to_indices(coordinates)];
}
 
template <typename T, size_t D, typename Function>
void apply(DiscreteField<T, D> &field, Function &&func) {
  field.apply(std::forward<Function>(func));
}
 
template <typename T, size_t D> void show(DiscreteField<T, D> &field) {
  utils::show(field.get_array().get_data(), field.get_index().get_size(),
              field.get_index().get_begin());
}
 
} // namespace pfc
 
#endif