spline_interpolator_1d.hh

Go to the documentation of this file.

#pragma once
 
// DiFfRG
#include <DiFfRG/common/kokkos.hh>
#include <DiFfRG/common/math.hh>
#include <DiFfRG/discretization/coordinates/coordinates.hh>
 
// std
#include <cmath>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <vector>
 
namespace DiFfRG
{
  template <typename NT, typename Coordinates> class SplineInterpolator1D
  {
    static_assert(Coordinates::dim == 1, "SplineInterpolator1D requires 1D coordinates");
    // The spline coefficients come from a non-cyclic tridiagonal solve with boundary conditions at the grid edges,
    // which cannot close a periodic axis. Use LinearInterpolator1D there instead.
    static_assert(!is_periodic_coordinate_v<Coordinates>,
                  "SplineInterpolator1D does not support periodic coordinates; use LinearInterpolator1D.");
 
    // SoA layout: separate views for values and spline coefficients, for coalesced GPU access
    using ValueViewType = Kokkos::View<NT *, GPU_memory, Kokkos::MemoryTraits<Kokkos::RandomAccess>>;
    using CoeffViewType = ValueViewType;
    using HostValueViewType = typename ValueViewType::host_mirror_type;
    using HostCoeffViewType = typename CoeffViewType::host_mirror_type;
 
    static constexpr bool has_separate_device =
        !std::is_same_v<typename ValueViewType::memory_space, typename HostValueViewType::memory_space>;
 
  public:
    using ctype = typename Coordinates::ctype;
    using value_type = NT;
    static constexpr size_t dim = 1;
 
    SplineInterpolator1D(const Coordinates &coordinates)
        : coordinates(coordinates), size(coordinates.size()),
          device_values("SplineInterpolator1D_values", coordinates.size()),
          device_coeffs("SplineInterpolator1D_coeffs", coordinates.size()),
          host_values(Kokkos::create_mirror_view(device_values)),
          host_coeffs(Kokkos::create_mirror_view(device_coeffs))
    {
    }
 
    KOKKOS_DEFAULTED_FUNCTION SplineInterpolator1D(const SplineInterpolator1D &) = default;
 
    template <typename NT2>
    void update(const NT2 *in_data, const ctype lower_y1 = std::numeric_limits<ctype>::max(),
                const ctype upper_y1 = std::numeric_limits<ctype>::max())
    {
      // Copy values from input data
      for (size_t i = 0; i < size; ++i)
        host_values(i) = in_data[i];
 
      // Build the spline coefficients
      build_y2(lower_y1, upper_y1);
 
      if constexpr (has_separate_device) {
        typename ValueViewType::execution_space exec;
        Kokkos::deep_copy(exec, device_values, host_values);
        Kokkos::deep_copy(exec, device_coeffs, host_coeffs);
        exec.fence();
      }
    }
 
    NT operator[](size_t i) const { return host_values(i); }
 
    ctype KOKKOS_FUNCTION index(const typename Coordinates::ctype x) const { return coordinates.backward(x); }
 
    NT KOKKOS_FUNCTION at(const ctype raw_idx) const
    {
      // Clamp the index to the range [0, size - 1]
      const ctype idx = Kokkos::max(static_cast<ctype>(0),
                                    Kokkos::min(raw_idx, static_cast<ctype>(size - 1)));
      const size_t lidx = Kokkos::min(size_t(Kokkos::floor(idx)), size - 2);
      const size_t uidx = lidx + 1;
      // t is the fractional part of the index
      const ctype t = idx - lidx;
 
      const NT lower = value(lidx);
      const NT upper = value(uidx);
      const NT cl = coeff(lidx);
      const NT cu = coeff(uidx);
 
      const ctype tm1 = t - 1;
      const NT cubic = t * tm1 * ((t + 1) * cl - (t - 2) * cu);
 
      if constexpr (std::is_arithmetic_v<NT>)
        return Kokkos::fma(t, upper, Kokkos::fma(-t, lower, lower)) + cubic; // linear + cubic
      else
        return t * upper + (1 - t) * lower + cubic; // linear + cubic
    }
 
    NT KOKKOS_FUNCTION operator()(const typename Coordinates::ctype x) const { return at(index(x)); }
 
    const Coordinates &get_coordinates() const { return coordinates; }
 
    const NT *data() const { return host_values.data(); }
 
  private:
    KOKKOS_FORCEINLINE_FUNCTION NT value(const size_t i) const
    {
      KOKKOS_IF_ON_DEVICE((return device_values(i);))
      KOKKOS_IF_ON_HOST((return host_values(i);))
    }
 
    KOKKOS_FORCEINLINE_FUNCTION NT coeff(const size_t i) const
    {
      KOKKOS_IF_ON_DEVICE((return device_coeffs(i);))
      KOKKOS_IF_ON_HOST((return host_coeffs(i);))
    }
 
    const Coordinates coordinates;
    const size_t size;
 
    ValueViewType device_values;
    CoeffViewType device_coeffs;
    HostValueViewType host_values;
    HostCoeffViewType host_coeffs;
 
    void build_y2(const ctype lower_y1, const ctype upper_y1)
    {
      NT p, qn, sig, un;
      std::vector<NT> u(size - 1);
 
      if (!std::isfinite(lower_y1) || lower_y1 >= std::numeric_limits<ctype>::max() / 2)
        host_coeffs(0) = u[0] = 0.0;
      else {
        host_coeffs(0) = -0.5;
        u[0] = 3.0 * ((host_values(1) - host_values(0)) - lower_y1);
      }
      for (size_t i = 1; i < size - 1; i++) {
        sig = 0.5;
        p = sig * host_coeffs(i - 1) + 2.0;
        host_coeffs(i) = (sig - 1.0) / p;
        u[i] = (host_values(i + 1) - host_values(i)) - (host_values(i) - host_values(i - 1));
        u[i] = (6.0 * u[i] / 2. - sig * u[i - 1]) / p;
      }
      if (!std::isfinite(upper_y1) || upper_y1 >= std::numeric_limits<ctype>::max() / 2)
        qn = un = 0.0;
      else {
        qn = 0.5;
        un = 3.0 * (upper_y1 - (host_values(size - 1) - host_values(size - 2)));
      }
      host_coeffs(size - 1) = (un - qn * u[size - 2]) / (qn * host_coeffs(size - 2) + 1);
      for (int k = size - 2; k >= 0; k--)
        host_coeffs(k) = host_coeffs(k) * host_coeffs(k + 1) + u[k];
 
      // Precompute division by 6 so operator() avoids per-call divides
      for (size_t k = 0; k < size; ++k)
        host_coeffs(k) /= (ctype)6;
    }
  };
} // namespace DiFfRG