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 {
 
/**
 * @brief A functor for describing a spline boundary value by a constant extrapolation for 2D evaluator.
 *
 * To define the value of a function on B-splines out of the domain, we here use a constant
 * extrapolation on the edge.
 */
template <class DimI>
struct ConstantExtrapolationRule
{
private:
    ddc::Coordinate<DimI> m_eval_pos;
 
public:
    /**
     * @brief Instantiate a ConstantExtrapolationRule.
     *
     * The boundary value will be the same as at the coordinate given in a dimension given.
     * The dimension of the input defines the dimension of the boundary condition.
     * No second and third parameters are needed in case of periodic splines on the
     * dimension off-interest (the complementary dimension of the boundary condition).
     *
     * @param[in] eval_pos Coordinate in the dimension given inside the domain where we will evaluate each points outside the domain.
     */
    explicit ConstantExtrapolationRule(ddc::Coordinate<DimI> eval_pos) : m_eval_pos(eval_pos) {}
 
    /**
     * @brief Get the value of the function on B-splines at a coordinate outside the domain.
     *
     * In the dimension defined in the constructor Dim1 (or Dim2), it sets the coordinate pos_1 (or pos_2)
     * given at the m_eval_pos coordinate if it is outside the domain.
     * If the coordinate on the complementary dimension of the boundary condition dimension ddc::Coordinate<DimNI>(coord_extrap) is
     * outside the domain, then it also sets the coordinate at eval_pos_not_interest_min
     * (if ddc::Coordinate<DimNI>(coord_extrap) @f$ < @f$ eval_pos_not_interest_min) or
     * at eval_pos_not_interest_max (if ddc::Coordinate<DimNI>(coord_extrap) @f$ > @f$ eval_pos_not_interest_max).
     *
     * @param[in] coord_extrap The coordinates where we want to evaluate the function on B-splines
     * @param[in] spline_coef The coefficients of the function on B-splines.
     *
     *@return A Real with the value of the function on B-splines evaluated at the coordinate.
     */
    template <class CoordType, class... BSplines, class Layout, class MemorySpace>
    KOKKOS_FUNCTION Real operator()(
            CoordType coord_extrap,
            ddc::ChunkSpan<Real const, ddc::DiscreteDomain<BSplines...>, Layout, MemorySpace> const
                    spline_coef) const
    {
        static_assert(in_tags_v<DimI, to_type_seq_t<CoordType>>);
        using TypeSeqBSplines = ddc::TypeSeq<BSplines...>;
 
        ddc::Coordinate<DimI, typename BSplines::continuous_dimension_type...> const coord_eval(
                get_eval_pos<BSplines>(ddc::select<typename BSplines::continuous_dimension_type>(
                        coord_extrap))...);
 
        auto vals_ptr = cexa::make_tuple(std::array<Real, BSplines::degree() + 1> {}...);
        auto const vals = cexa::make_tuple(
                Kokkos::mdspan<Real, Kokkos::extents<std::size_t, BSplines::degree() + 1>>(
                        cexa::get<ddc::type_seq_rank_v<BSplines, TypeSeqBSplines>>(vals_ptr)
                                .data())...);
 
        auto const jmin = cexa::make_tuple(
                ddc::discrete_space<BSplines>().eval_basis(
                        cexa::get<ddc::type_seq_rank_v<BSplines, TypeSeqBSplines>>(vals),
                        ddc::Coordinate<typename BSplines::continuous_dimension_type>(
                                coord_eval))...);
 
        static constexpr std::size_t dimension = sizeof...(BSplines);
 
        Real y = 0.0;
        for_each(
                std::array<std::size_t, dimension> {(BSplines::degree() + 1)...},
                [&](std::array<std::size_t, dimension> idx) {
                    y += spline_coef(
                                 ddc::DiscreteElement<BSplines...>((
                                         cexa::get<ddc::type_seq_rank_v<BSplines, TypeSeqBSplines>>(
                                                 jmin)
                                         + idx[ddc::type_seq_rank_v<
                                                 BSplines,
                                                 TypeSeqBSplines>])...))
                         * (cexa::get<ddc::type_seq_rank_v<BSplines, TypeSeqBSplines>>(
                                    vals)[idx[ddc::type_seq_rank_v<BSplines, TypeSeqBSplines>]]
                            * ...);
                });
 
        return y;
    }
 
private:
    template <class BSplinesNI>
    KOKKOS_FUNCTION ddc::Coordinate<typename BSplinesNI::continuous_dimension_type> get_eval_pos(
            ddc::Coordinate<typename BSplinesNI::continuous_dimension_type> coord_extrap) const
    {
        if constexpr (std::is_same_v<typename BSplinesNI::continuous_dimension_type, DimI>) {
            return m_eval_pos;
        } else {
            if constexpr (BSplinesNI::is_periodic()) {
                return ddc::Coordinate<typename BSplinesNI::continuous_dimension_type>(
                        coord_extrap);
            } else {
                return Kokkos::
                        clamp(ddc::Coordinate<typename BSplinesNI::continuous_dimension_type>(
                                      coord_extrap),
                              ddc::discrete_space<BSplinesNI>().rmin(),
                              ddc::discrete_space<BSplinesNI>().rmax());
            }
        }
    }
 
    template <std::size_t N, class Functor, class... Is>
    KOKKOS_FUNCTION static void for_each(
            std::array<std::size_t, N> const& bounds,
            Functor const& f,
            Is... is)
    {
        static constexpr std::size_t I = sizeof...(Is);
        if constexpr (I == N) {
            f(std::array<std::size_t, N> {is...});
        } else {
            for (std::size_t i = 0; i < bounds[I]; ++i) {
                for_each(bounds, f, is..., i);
            }
        }
    }
};
 
} // namespace ddc