bsl_advection_rp.hpp

// SPDX-License-Identifier: MIT
 
#pragma once
#include <sll/mapping/metric_tensor.hpp>
 
#include "advection_domain.hpp"
#include "ddc_alias_inline_functions.hpp"
#include "ddc_aliases.hpp"
#include "directional_tag.hpp"
#include "geometry.hpp"
#include "i_interpolator_2d_rp.hpp"
#include "iadvectionrp.hpp"
#include "spline_foot_finder.hpp"
#include "spline_interpolator_2d_rp.hpp"
#include "vector_field.hpp"
#include "vector_field_mem.hpp"
 
 
 
template <class FootFinder, class Mapping>
class BslAdvectionRTheta : public IAdvectionRTheta
{
private:
    PreallocatableSplineInterpolatorRTheta<ddc::NullExtrapolationRule> const& m_interpolator;
 
    FootFinder const& m_find_feet;
 
    Mapping const& m_mapping;
 
 
public:
    BslAdvectionRTheta(
            PreallocatableSplineInterpolatorRTheta<ddc::NullExtrapolationRule> const&
                    function_interpolator,
            FootFinder const& foot_finder,
            Mapping const& mapping)
        : m_interpolator(function_interpolator)
        , m_find_feet(foot_finder)
        , m_mapping(mapping)
    {
    }
 
    ~BslAdvectionRTheta() override = default;
 
 
    host_t<DFieldRTheta> operator()(
            host_t<DFieldRTheta> allfdistribu,
            host_t<DConstVectorFieldRTheta<X, Y>> advection_field_xy,
            double dt) const override
    {
        // Pre-allocate some memory to prevent allocation later in loop
        std::unique_ptr<IInterpolatorRTheta> const interpolator_ptr = m_interpolator.preallocate();
 
        // Initialise the feet
        host_t<FieldMemRTheta<CoordRTheta>> feet_rp(get_idx_range(advection_field_xy));
        ddc::for_each(get_idx_range(advection_field_xy), [&](IdxRTheta const irp) {
            feet_rp(irp) = ddc::coordinate(irp);
        });
 
        // Compute the characteristic feet at tn ----------------------------------------------------
        m_find_feet(get_field(feet_rp), advection_field_xy, dt);
 
        // Interpolate the function on the characteristic feet. -------------------------------------
        (*interpolator_ptr)(allfdistribu, get_const_field(feet_rp));
 
        return allfdistribu;
    }
 
 
    host_t<DFieldRTheta> operator()(
            host_t<DFieldRTheta> allfdistribu,
            host_t<DConstVectorFieldRTheta<R, Theta>> advection_field_rp,
            CoordXY const& advection_field_xy_center,
            double dt) const override
    {
        Kokkos::Profiling::pushRegion("PolarAdvection");
        IdxRangeRTheta grid(get_idx_range<GridR, GridTheta>(allfdistribu));
 
        const int npoints_p = IdxRangeTheta(grid).size();
        IdxRangeRTheta const grid_without_Opoint(grid.remove_first(IdxStepRTheta(1, 0)));
        IdxRangeRTheta const Opoint_grid(grid.take_first(IdxStepRTheta(1, npoints_p)));
 
 
        // Convert advection field on RTheta to advection field on XY
        host_t<DVectorFieldMemRTheta<X, Y>> advection_field_xy(grid);
 
        MetricTensor<Mapping, CoordRTheta> metric_tensor(m_mapping);
 
        ddc::for_each(grid_without_Opoint, [&](IdxRTheta const irp) {
            CoordRTheta const coord_rp(ddc::coordinate(irp));
 
            std::array<std::array<double, 2>, 2> J; // Jacobian matrix
            m_mapping.jacobian_matrix(coord_rp, J);
            std::array<std::array<double, 2>, 2> G; // Metric tensor
            metric_tensor(G, coord_rp);
 
            ddcHelper::get<X>(advection_field_xy)(irp)
                    = ddcHelper::get<R>(advection_field_rp)(irp) * J[1][1] / std::sqrt(G[1][1])
                      + ddcHelper::get<Theta>(advection_field_rp)(irp) * -J[1][0]
                                / std::sqrt(G[0][0]);
            ddcHelper::get<Y>(advection_field_xy)(irp)
                    = ddcHelper::get<R>(advection_field_rp)(irp) * -J[0][1] / std::sqrt(G[1][1])
                      + ddcHelper::get<Theta>(advection_field_rp)(irp) * J[0][0]
                                / std::sqrt(G[0][0]);
        });
 
        ddc::for_each(Opoint_grid, [&](IdxRTheta const irp) {
            ddcHelper::get<X>(advection_field_xy)(irp) = CoordX(advection_field_xy_center);
            ddcHelper::get<Y>(advection_field_xy)(irp) = CoordY(advection_field_xy_center);
        });
 
        // Pre-allocate some memory to prevent allocation later in loop
        std::unique_ptr<IInterpolatorRTheta> const interpolator_ptr = m_interpolator.preallocate();
 
        // Initialise the feet
        host_t<FieldMemRTheta<CoordRTheta>> feet_rp(grid);
        ddc::for_each(grid, [&](IdxRTheta const irp) { feet_rp(irp) = ddc::coordinate(irp); });
 
        // Compute the characteristic feet at tn ----------------------------------------------------
        m_find_feet(get_field(feet_rp), get_const_field(advection_field_xy), dt);
 
        // Interpolate the function on the characteristic feet. -------------------------------------
        (*interpolator_ptr)(allfdistribu, get_const_field(feet_rp));
        Kokkos::Profiling::popRegion();
 
        return allfdistribu;
    }
};