constant_extrapolation_rule.hpp

// Copyright (C) The DDC development team, see COPYRIGHT.md file
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <array>
#include <cstddef>
#include <type_traits>
 
#include <ddc/ddc.hpp>
 
#include <Kokkos_Core.hpp>
 
namespace ddc {
 
template <class DimI, class... Dim>
struct ConstantExtrapolationRule
{
};
 
template <class DimI>
struct ConstantExtrapolationRule<DimI>
{
private:
    ddc::Coordinate<DimI> m_eval_pos;
 
public:
    explicit ConstantExtrapolationRule(ddc::Coordinate<DimI> eval_pos) : m_eval_pos(eval_pos) {}
 
    template <class CoordType, class BSplines, class Layout, class MemorySpace>
    KOKKOS_FUNCTION double operator()(
            [[maybe_unused]] CoordType pos,
            ddc::ChunkSpan<double const, ddc::DiscreteDomain<BSplines>, Layout, MemorySpace> const
                    spline_coef) const
    {
        static_assert(in_tags_v<DimI, to_type_seq_t<CoordType>>);
 
        std::array<double, BSplines::degree() + 1> vals_ptr;
        Kokkos::mdspan<double, Kokkos::extents<std::size_t, BSplines::degree() + 1>> const vals(
                vals_ptr.data());
 
        ddc::DiscreteElement<BSplines> const idx
                = ddc::discrete_space<BSplines>().eval_basis(vals, m_eval_pos);
 
        double y = 0.0;
        for (std::size_t i = 0; i < BSplines::degree() + 1; ++i) {
            y += spline_coef(idx + i) * vals[i];
        }
        return y;
    }
};
 
template <class DimI, class DimNI>
struct ConstantExtrapolationRule<DimI, DimNI>
{
private:
    ddc::Coordinate<DimI> m_eval_pos;
    ddc::Coordinate<DimNI> m_eval_pos_not_interest_min;
    ddc::Coordinate<DimNI> m_eval_pos_not_interest_max;
 
public:
    explicit ConstantExtrapolationRule(
            ddc::Coordinate<DimI> eval_pos,
            ddc::Coordinate<DimNI> eval_pos_not_interest_min,
            ddc::Coordinate<DimNI> eval_pos_not_interest_max)
        : m_eval_pos(eval_pos)
        , m_eval_pos_not_interest_min(eval_pos_not_interest_min)
        , m_eval_pos_not_interest_max(eval_pos_not_interest_max)
    {
    }
 
    template <class SFINAEDimNI = DimNI, std::enable_if_t<SFINAEDimNI::PERIODIC, int> = 0>
    explicit ConstantExtrapolationRule(ddc::Coordinate<DimI> eval_pos)
        : m_eval_pos(eval_pos)
        , m_eval_pos_not_interest_min(0.)
        , m_eval_pos_not_interest_max(0.)
    {
    }
 
    template <class CoordType, class BSplines1, class BSplines2, class Layout, class MemorySpace>
    KOKKOS_FUNCTION double operator()(
            CoordType coord_extrap,
            ddc::ChunkSpan<
                    double const,
                    ddc::DiscreteDomain<BSplines1, BSplines2>,
                    Layout,
                    MemorySpace> const spline_coef) const
    {
        static_assert(
                in_tags_v<DimI, to_type_seq_t<CoordType>>
                && in_tags_v<DimNI, to_type_seq_t<CoordType>>);
 
        ddc::Coordinate<DimI, DimNI> eval_pos;
        if constexpr (DimNI::PERIODIC) {
            eval_pos = ddc::
                    Coordinate<DimI, DimNI>(m_eval_pos, ddc::Coordinate<DimNI>(coord_extrap));
        } else {
            eval_pos = ddc::Coordinate<DimI, DimNI>(
                    m_eval_pos,
                    Kokkos::
                            clamp(ddc::Coordinate<DimNI>(coord_extrap),
                                  m_eval_pos_not_interest_min,
                                  m_eval_pos_not_interest_max));
        }
 
        std::array<double, BSplines1::degree() + 1> vals1_ptr;
        Kokkos::mdspan<double, Kokkos::extents<std::size_t, BSplines1::degree() + 1>> const vals1(
                vals1_ptr.data());
        std::array<double, BSplines2::degree() + 1> vals2_ptr;
        Kokkos::mdspan<double, Kokkos::extents<std::size_t, BSplines2::degree() + 1>> const vals2(
                vals2_ptr.data());
 
        ddc::DiscreteElement<BSplines1> const idx1 = ddc::discrete_space<BSplines1>().eval_basis(
                vals1,
                ddc::Coordinate<typename BSplines1::continuous_dimension_type>(eval_pos));
        ddc::DiscreteElement<BSplines2> const idx2 = ddc::discrete_space<BSplines2>().eval_basis(
                vals2,
                ddc::Coordinate<typename BSplines2::continuous_dimension_type>(eval_pos));
 
        double y = 0.0;
        for (std::size_t i = 0; i < BSplines1::degree() + 1; ++i) {
            for (std::size_t j = 0; j < BSplines2::degree() + 1; ++j) {
                y += spline_coef(idx1 + i, idx2 + j) * vals1[i] * vals2[j];
            }
        }
 
        return y;
    }
};
 
} // namespace ddc