collisions_dimensions.hpp

// SPDX-License-Identifier: MIT
#pragma once
#include "ddc_alias_inline_functions.hpp"
#include "ddc_aliases.hpp"
 
namespace collisions_dimensions {
struct InternalSpoofGrid
{
};
 
struct InternalSpoofGridR : InternalSpoofGrid
{
};
 
struct InternalSpoofGridTheta : InternalSpoofGrid
{
};
 
template <class Grid1D>
inline constexpr bool is_spoofed_dim_v = std::is_base_of_v<InternalSpoofGrid, Grid1D>;
 
template <class Field>
struct ExtractRDim
{
    static_assert(!std::is_same_v<Field, Field>, "Unrecognised radial profile type");
};
 
template <class Field, class GridR>
struct ExtractThetaDim
{
    static_assert(!std::is_same_v<Field, Field>, "Unrecognised poloidal profile type");
};
 
template <class Grid1D, class IdxRangeFDistrib>
inline IdxRange<Grid1D> get_1d_idx_range(IdxRangeFDistrib idx_range)
{
    if constexpr (is_spoofed_dim_v<Grid1D>) {
        return IdxRange<Grid1D>(Idx<Grid1D> {0}, IdxStep<Grid1D> {1});
    } else {
        return ddc::select<Grid1D>(idx_range);
    }
}
 
template <class... Grid1D, class IdxRangeFDistrib>
inline IdxRange<Grid1D...> get_expanded_idx_range(IdxRangeFDistrib idx_range)
{
    return IdxRange<Grid1D...>(get_1d_idx_range<Grid1D>(idx_range)...);
}
 
template <>
struct ExtractRDim<double>
{
    using type = InternalSpoofGridR;
};
 
template <class GridR, class Layout>
struct ExtractRDim<
        ConstField<double, IdxRange<GridR>, Kokkos::DefaultExecutionSpace::memory_space, Layout>>
{
    using type = GridR;
};
 
template <class ElementType, class IdxRange, class MemSpace, class Layout>
struct ExtractRDim<Field<ElementType, IdxRange, MemSpace, Layout>>
{
    static_assert(
            std::is_same_v<ElementType, const double>,
            "The radial profile should be a double and should not be modifiable. Please use a "
            "constant field.");
    static_assert(
            std::is_same_v<MemSpace, Kokkos::DefaultExecutionSpace::memory_space>,
            "The radial profile should be provided on the GPU");
    static_assert(
            (ddc::type_seq_size_v<ddc::to_type_seq_t<IdxRange>>) > 1,
            "The radial profile should not be defined on more than 1 dimensions.");
};
 
template <>
struct ExtractThetaDim<double, InternalSpoofGridR>
{
    using type = InternalSpoofGridTheta;
};
 
template <class GridR, class Layout>
struct ExtractThetaDim<
        ConstField<double, IdxRange<GridR>, Kokkos::DefaultExecutionSpace::memory_space, Layout>,
        GridR>
{
    using type = InternalSpoofGridTheta;
};
 
template <class GridR, class GridTheta, class Layout>
struct ExtractThetaDim<
        ConstField<
                double,
                IdxRange<GridTheta>,
                Kokkos::DefaultExecutionSpace::memory_space,
                Layout>,
        GridR>
{
    using type = GridTheta;
};
 
template <class GridR, class GridTheta, class Layout>
struct ExtractThetaDim<
        ConstField<
                double,
                IdxRange<GridTheta, GridR>,
                Kokkos::DefaultExecutionSpace::memory_space,
                Layout>,
        GridR>
{
    using type = GridTheta;
};
 
template <class GridR, class ElementType, class IdxRange, class MemSpace, class Layout>
struct ExtractThetaDim<Field<ElementType, IdxRange, MemSpace, Layout>, GridR>
{
    static_assert(
            std::is_same_v<ElementType, const double>,
            "The poloidal profile should be a double and should not be modifiable. Please use a "
            "constant field.");
    static_assert(
            std::is_same_v<MemSpace, Kokkos::DefaultExecutionSpace::memory_space>,
            "The poloidal profile should be provided on the GPU");
    static_assert(
            (ddc::type_seq_size_v<ddc::to_type_seq_t<IdxRange>>) > 2,
            "The poloidal profile should not be defined on more than 2 dimensions.");
};
 
template <class Field>
struct ExtractRThetaTags
{
    static_assert(!std::is_same_v<Field, Field>, "Unrecognised poloidal profile type");
};
 
template <>
struct ExtractRThetaTags<double>
{
    using type = ddc::detail::TypeSeq<>;
};
 
template <class IdxRange>
struct ExtractRThetaTags<DConstField<IdxRange>>
{
    using type = ddc::to_type_seq_t<IdxRange>;
};
 
template <class Grid1D, class OrderedGrids>
constexpr bool check_dimension_location(int& idx)
{
    if constexpr (!is_spoofed_dim_v<Grid1D>) {
        bool success = ddc::type_seq_rank_v<Grid1D, OrderedGrids> == idx;
        idx += 1;
        return success;
    } else {
        return true;
    }
}
 
template <class OrderedGrids, class... ExpectedGrids>
constexpr bool order_of_last_grids()
{
    int n_real_dims = ((is_spoofed_dim_v<ExpectedGrids> ? 0 : 1) + ...);
    int idx(ddc::type_seq_size_v<OrderedGrids> - n_real_dims);
    return ((check_dimension_location<ExpectedGrids, OrderedGrids>(idx)) && ...);
}
 
template <class OrderedGrids, class... ExpectedGrids>
constexpr bool order_of_first_grids()
{
    int idx(0);
    return ((check_dimension_location<ExpectedGrids, OrderedGrids>(idx)) && ...);
}
 
} // namespace collisions_dimensions