CollisionSpVparMu.hpp

// SPDX-License-Identifier: MIT
#pragma once
#include <ddc/ddc.hpp>
 
#include "assert.hpp"
#include "collisions_dimensions.hpp"
#include "ddc_alias_inline_functions.hpp"
#include "ddc_aliases.hpp"
#include "ddc_helper.hpp"
#include "koliop_interface.hpp"
#include "species_info.hpp"
 
template <
        class CollInfo,
        class IdxRangeFDistrib,
        class GridVpar,
        class GridMu,
        class InputDFieldThetaR>
class CollisionSpVparMu /* : public IRightHandSide */
{
private:
    using InputDFieldR = typename CollInfo::radial_chunk_type;
    using fdistrib_grids = ddc::to_type_seq_t<IdxRangeFDistrib>;
    using GridR = typename collisions_dimensions::ExtractRDim<InputDFieldR>::type;
    using GridTheta =
            typename collisions_dimensions::ExtractThetaDim<InputDFieldThetaR, GridR>::type;
    using InputThetaRTags =
            typename collisions_dimensions::ExtractRThetaTags<InputDFieldThetaR>::type;
    using InputSpThetaRVparTags = ddc::type_seq_merge_t<
            ddc::detail::TypeSeq<Species>,
            ddc::type_seq_merge_t<InputThetaRTags, ddc::detail::TypeSeq<GridVpar>>>;
    using InputIdxRangeSpThetaRVpar
            = ddc::detail::convert_type_seq_to_discrete_domain_t<InputSpThetaRVparTags>;
 
    // Validate template types
    static_assert(ddc::is_discrete_domain_v<IdxRangeFDistrib>);
    static_assert(IdxRangeFDistrib::rank() >= 3 && IdxRangeFDistrib::rank() <= 6);
    static_assert((std::is_same_v<InputDFieldR, double>) || ddc::is_borrowed_chunk_v<InputDFieldR>);
    static_assert(
            (std::is_same_v<InputDFieldThetaR, double>)
            || (ddc::is_borrowed_chunk_v<InputDFieldThetaR>));
    // Ensure expected types appear in distribution index range
    static_assert(
            ddc::in_tags_v<Species, fdistrib_grids>,
            "The distribution function must be defined over the species");
    static_assert(
            ddc::in_tags_v<GridVpar, fdistrib_grids>,
            "The distribution function must be defined over the Vpar direction");
    static_assert(
            ddc::in_tags_v<GridMu, fdistrib_grids>,
            "The distribution function must be defined over the Mu direction");
    // Ensure that uniform
    static_assert(
            ddc::is_uniform_point_sampling_v<GridVpar>,
            "The grid must be uniform in the vpar direction.");
    static_assert(
            ddc::is_uniform_point_sampling_v<GridMu>,
            "The grid must be uniform in the mu direction.");
 
    static_assert(
            ddc::type_seq_rank_v<Species, fdistrib_grids> == 0,
            "Species should be the first index in the distribution function index set");
    static_assert(
            collisions_dimensions::
                    order_of_last_grids<fdistrib_grids, GridTheta, GridR, GridVpar, GridMu>(),
            "Misordered inex set for the distribution function. Koliop expects (Sp, Phi, Theta, R, "
            "Vpar, Mu)");
 
public:
    using InputDFieldSpThetaRVpar = DConstField<InputIdxRangeSpThetaRVpar>;
 
    using IdxRangeR = IdxRange<GridR>;
    using IdxRangeThetaR = IdxRange<GridTheta, GridR>;
    using IdxRangeSpThetaRVpar = IdxRange<Species, GridTheta, GridR, GridVpar>;
    using IdxRangeMu = IdxRange<GridMu>;
    using IdxRangeVpar = IdxRange<GridVpar>;
 
    using IdxThetaR = Idx<GridTheta, GridR>;
    using IdxSpThetaRVpar = Idx<Species, GridTheta, GridR, GridVpar>;
 
    using DFieldMemR = DFieldMem<IdxRangeR>;
    using DFieldMemMu = DFieldMem<IdxRangeMu>;
    using DFieldMemVpar = DFieldMem<IdxRangeVpar>;
    using DFieldMemThetaR = DFieldMem<IdxRangeThetaR>;
    using DFieldMemSpThetaRVpar = DFieldMem<IdxRangeSpThetaRVpar>;
    using DFieldR = DField<IdxRangeR>;
    using DFieldThetaR = DField<IdxRangeThetaR>;
    using DFieldSpThetaRVpar = DField<IdxRangeSpThetaRVpar>;
    using DConstFieldMu = DConstField<
            IdxRangeMu,
            Kokkos::DefaultExecutionSpace::
                    memory_space>; // Equivalent to Field<double const, IdxRangeMu>
    using DConstFieldVpar = DConstField<
            IdxRangeVpar,
            Kokkos::DefaultExecutionSpace::
                    memory_space>; // Equivalent to Field<double const, IdxRangeMu>
 
    using FDistribField = DField<IdxRangeFDistrib>;
 
private:
    void deepcopy_radial_profile(DFieldR dst, InputDFieldR src)
    {
        if constexpr (collisions_dimensions::is_spoofed_dim_v<GridR>) {
            // If InputDFieldR is a double because R is not provided
            ddc::parallel_fill(Kokkos::DefaultExecutionSpace(), dst, src);
        } else {
            // If InputDFieldR and DFieldR are the same type
            ddc::parallel_deepcopy(dst, src);
        }
    }
 
public:
    void deepcopy_poloidal_plane(DFieldThetaR dst, InputDFieldThetaR src)
    {
        if constexpr ((collisions_dimensions::is_spoofed_dim_v<GridR>)&&(
                              collisions_dimensions::is_spoofed_dim_v<GridTheta>)) {
            // If InputDFieldThetaR is a double because R and Theta are not provided
            ddc::parallel_fill(Kokkos::DefaultExecutionSpace(), dst, src);
        } else if constexpr ((!collisions_dimensions::is_spoofed_dim_v<GridR>)&&(
                                     !collisions_dimensions::is_spoofed_dim_v<GridTheta>)) {
            // If InputDFieldThetaR and DFieldThetaR are the same type
            ddc::parallel_deepcopy(dst, src);
        } else {
            using NonSpoofDim = std::
                    conditional_t<collisions_dimensions::is_spoofed_dim_v<GridR>, GridTheta, GridR>;
            ddc::parallel_for_each(
                    Kokkos::DefaultExecutionSpace(),
                    get_idx_range(dst),
                    KOKKOS_LAMBDA(Idx<GridTheta, GridR> idx) {
                        dst(idx) = src(ddc::select<NonSpoofDim>(idx));
                    });
        }
    }
 
    void deepcopy_Bstar(DFieldSpThetaRVpar dst, InputDFieldSpThetaRVpar src)
    {
        if constexpr (std::is_same_v<DFieldSpThetaRVpar, InputDFieldSpThetaRVpar>) {
            ddc::parallel_deepcopy(dst, src);
        } else {
            using InputIdxSpThetaRVpar = typename InputIdxRangeSpThetaRVpar::discrete_element_type;
            ddc::parallel_for_each(
                    Kokkos::DefaultExecutionSpace(),
                    get_idx_range(dst),
                    KOKKOS_LAMBDA(IdxSpThetaRVpar idx) {
                        InputIdxSpThetaRVpar src_idx(idx);
                        dst(idx) = src(src_idx);
                    });
        }
    }
 
public:
    CollisionSpVparMu(
            CollInfo const& collision_info,
            IdxRangeFDistrib idxrange_fdistrib,
            DConstFieldMu coeff_intdmu,
            DConstFieldVpar coeff_intdvpar,
            InputDFieldThetaR B_norm,
            InputDFieldSpThetaRVpar Bstar_s)
        : m_operator_handle {}
        , m_hat_As {"m_hat_As", collisions_dimensions::get_expanded_idx_range<Species>(idxrange_fdistrib)}
        , m_hat_Zs {"m_hat_Zs", collisions_dimensions::get_expanded_idx_range<Species>(idxrange_fdistrib)}
        , m_mask_buffer_r {"m_mask_buffer_r", collisions_dimensions::get_expanded_idx_range<GridR>(idxrange_fdistrib)}
        , m_mask_LIM {"m_mask_LIM", IdxRangeThetaR {collisions_dimensions::get_expanded_idx_range<GridTheta, GridR>(idxrange_fdistrib)}}
        , m_B_norm {"m_B_norm", IdxRangeThetaR {collisions_dimensions::get_expanded_idx_range<GridTheta, GridR>(idxrange_fdistrib)}}
        , m_Bstar_s {"m_Bstar_s", IdxRangeSpThetaRVpar {collisions_dimensions::get_expanded_idx_range<Species, GridTheta, GridR, GridVpar>(idxrange_fdistrib)}}
        , m_coeff_AD {"m_coeff_AD", collisions_dimensions::get_expanded_idx_range<GridR>(idxrange_fdistrib)}
        , m_mug {"m_mug", ddc::select<GridMu>(idxrange_fdistrib)}
        , m_vparg {"m_vparg", ddc::select<GridVpar>(idxrange_fdistrib)}
    {
        IdxRangeVpar idxrange_vpar(idxrange_fdistrib);
        if (idxrange_vpar.size() % 2 != 0) {
            throw std::runtime_error("The number of points in the vpar direction must be a "
                                     "multiple of 2. This ensures that there is no grid point at "
                                     "vpar=0 (this would cause division by 0).");
        }
        IdxRangeSp idxrange_sp(idxrange_fdistrib);
        // --> Initialize the mass species
        host_t<DConstFieldSp> hat_As_host
                = ddc::host_discrete_space<Species>().masses()[idxrange_sp];
        ddc::parallel_deepcopy(get_field(m_hat_As), hat_As_host);
        // --> Initialize the charge species
        host_t<DConstFieldSp> hat_Zs_host
                = ddc::host_discrete_space<Species>().charges()[idxrange_sp];
        ddc::parallel_deepcopy(get_field(m_hat_Zs), hat_Zs_host);
 
        // --> Initialize the other quantities needed in koliop
        deepcopy_radial_profile(get_field(m_coeff_AD), collision_info.coeff_AD());
        deepcopy_poloidal_plane(get_field(m_B_norm), B_norm);
        deepcopy_Bstar(get_field(m_Bstar_s), Bstar_s);
 
        // --> Initialization of the masks that have no sense here to 0.
        ddc::parallel_fill(get_field(m_mask_buffer_r), 0.0); // Masked if >= 0.99
        ddc::parallel_fill(get_field(m_mask_LIM), 0.0); // Masked if >= 0.99
 
        // --> Initialization of vpar and mu grids
        ddcHelper::dump_coordinates(Kokkos::DefaultExecutionSpace(), get_field(m_mug));
        ddcHelper::dump_coordinates(Kokkos::DefaultExecutionSpace(), get_field(m_vparg));
 
        // NOTE: We need to fence because DumpCoordinates is asynchronous on GPUs.
        Kokkos::fence();
 
        std::size_t const n_mu = ddc::select<GridMu>(idxrange_fdistrib).size();
        std::size_t const n_vpar = idxrange_vpar.size();
        std::size_t const n_r = get_idx_range<GridR>(m_mask_buffer_r).size();
        std::size_t const n_theta = get_idx_range<GridTheta>(m_B_norm).size();
        std::size_t const n_sp = idxrange_sp.size();
        std::size_t const n_batch
                = idxrange_fdistrib.size() / (n_mu * n_vpar * n_r * n_theta * n_sp);
 
        m_operator_handle = koliop_interface::DoOperatorInitialization(
                n_mu,
                n_vpar,
                n_r,
                n_theta,
                n_batch,
                n_sp,
                collision_info.collisions_interspecies(),
                /* the_local_index range_r_offset */ 0 + n_r - 1,
                m_hat_As.data_handle(),
                m_hat_Zs.data_handle(),
                m_mug.data_handle(),
                m_vparg.data_handle(),
                coeff_intdmu.data_handle(),
                coeff_intdvpar.data_handle(),
                m_coeff_AD.data_handle(),
                m_mask_buffer_r.data_handle(),
                m_mask_LIM.data_handle(),
                m_B_norm.data_handle(),
                m_Bstar_s.data_handle());
    }
 
    ~CollisionSpVparMu()
    {
        koliop_interface::DoOperatorDeinitialization(
                static_cast<::koliop_Operator>(m_operator_handle));
    }
 
    void operator()(FDistribField all_f_distribution, double deltat_coll) const
    {
        if (::koliop_Collision(
                    static_cast<::koliop_Operator>(m_operator_handle),
                    deltat_coll,
                    all_f_distribution.data_handle())
            != KOLIOP_STATUS_SUCCESS) {
            GSLX_ASSERT(false);
        }
 
        // NOTE: While gslx is not fully stream compatible, just fence at the end of
        // the computation.
        if (::koliop_Fence(static_cast<::koliop_Operator>(m_operator_handle))
            != KOLIOP_STATUS_SUCCESS) {
            GSLX_ASSERT(false);
        }
    }
 
protected:
    ::koliop_Operator m_operator_handle;
    // NOTE: Some of these arrays should come from a parent class or manager.
    // They resides in this class while we wait for their implementation.
    koliop_interface::MDL<double[6][6]> m_comb_mat;
    DFieldMemSp m_hat_As;
    DFieldMemSp m_hat_Zs;
    DFieldMemR m_coeff_AD;
    // [TODO]: This mask should maybe be deleted in C++ version
    DFieldMemR m_mask_buffer_r;
    DFieldMemThetaR m_mask_LIM;
    // [TODO] Attention this must be 3D for generalization to 3D geometry--> transfer it in a 1D array ?
    DFieldMemThetaR m_B_norm;
    // [TODO] Must be 5D for full 3D geometry
    DFieldMemSpThetaRVpar m_Bstar_s;
    DFieldMemMu m_mug;
    DFieldMemVpar m_vparg;
};