linear_interpolator_2d.hh

Go to the documentation of this file.

#pragma once
 
// DiFfRG
#include <DiFfRG/common/kokkos.hh>
#include <DiFfRG/common/math.hh>
#include <DiFfRG/physics/interpolation/interpolation_stencil.hh>
 
namespace DiFfRG
{
  template <typename NT, typename Coordinates> class LinearInterpolator2D
  {
    static_assert(Coordinates::dim == 2, "LinearInterpolator2D requires 2D coordinates");
 
    static constexpr bool periodic_x = is_periodic_axis_v<Coordinates, 0>;
    static constexpr bool periodic_y = is_periodic_axis_v<Coordinates, 1>;
 
    using ViewType = Kokkos::View<NT **, GPU_memory, Kokkos::MemoryTraits<Kokkos::RandomAccess>>;
    using HostViewType = typename ViewType::host_mirror_type;
 
    static constexpr bool has_separate_device =
        !std::is_same_v<typename ViewType::memory_space, typename HostViewType::memory_space>;
 
  public:
    using ctype = typename Coordinates::ctype;
    using value_type = NT;
    static constexpr size_t dim = 2;
 
    LinearInterpolator2D(const Coordinates &coordinates)
        : coordinates(coordinates), sizes(coordinates.sizes()),
          device_data("LinearInterpolator2D_data", coordinates.sizes()[0], coordinates.sizes()[1]),
          host_data(Kokkos::create_mirror_view(device_data))
    {
    }
 
    KOKKOS_DEFAULTED_FUNCTION LinearInterpolator2D(const LinearInterpolator2D &) = default;
 
    template <typename NT2> void update(const NT2 *in_data)
    {
      for (size_t i = 0; i < sizes[0]; ++i)
        for (size_t j = 0; j < sizes[1]; ++j)
          host_data(i, j) = in_data[i * sizes[1] + j];
 
      if constexpr (has_separate_device) {
        typename ViewType::execution_space exec;
        Kokkos::deep_copy(exec, device_data, host_data);
        exec.fence();
      }
    }
 
    NT operator[](size_t i) const { return host_data(i / sizes[1], i % sizes[1]); }
 
    device::array<typename Coordinates::ctype, 2> KOKKOS_FUNCTION
    index(const typename Coordinates::ctype x, const typename Coordinates::ctype y) const
    {
      return coordinates.backward(x, y);
    }
 
    NT KOKKOS_FUNCTION at(const device::array<typename Coordinates::ctype, 2> &idx) const
    {
      const auto idx_x = idx[0];
      const auto idx_y = idx[1];
      // Clamped [i, i+1] stencil on bounded axes, wrapping [i, (i+1) % n] on periodic ones
      const auto sx = make_interpolation_stencil<periodic_x>(idx_x, sizes[0]);
      const auto sy = make_interpolation_stencil<periodic_y>(idx_y, sizes[1]);
 
      const size_t x0 = sx.lower, x1 = sx.upper;
      const size_t y0 = sy.lower, y1 = sy.upper;
 
      const auto corner00 = value(x0, y0);
      const auto corner01 = value(x0, y1);
      const auto corner10 = value(x1, y0);
      const auto corner11 = value(x1, y1);
 
      const auto tx = sx.t;
      const auto ty = sy.t;
 
      if constexpr (std::is_arithmetic_v<NT>)
        return Kokkos::fma(ty, Kokkos::fma(tx, corner11, Kokkos::fma(-tx, corner01, corner01)),
                           (1 - ty) * Kokkos::fma(tx, corner10, Kokkos::fma(-tx, corner00, corner00)));
      else
        return corner00 * (1 - tx) * (1 - ty) + corner01 * (1 - tx) * ty + corner10 * tx * (1 - ty) +
               corner11 * tx * ty;
    }
 
    NT KOKKOS_FUNCTION operator()(const typename Coordinates::ctype x,
                                  const typename Coordinates::ctype y) const
    {
      return at(index(x, y));
    }
 
    const Coordinates &get_coordinates() const { return coordinates; }
 
    const NT *data() const { return host_data.data(); }
 
  private:
    KOKKOS_FORCEINLINE_FUNCTION NT value(const size_t i, const size_t j) const
    {
      KOKKOS_IF_ON_DEVICE((return device_data(i, j);))
      KOKKOS_IF_ON_HOST((return host_data(i, j);))
    }
 
    const Coordinates coordinates;
    const device::array<size_t, 2> sizes;
 
    ViewType device_data;
    HostViewType host_data;
  };
} // namespace DiFfRG