Skip to content

Commit

Permalink
Merge branch 'ebourne_collision_readability' into 'main'
Browse files Browse the repository at this point in the history 
Extract collisions_dimensions from CollisionSpVparMu

See merge request gysela-developpers/gyselalibxx!596
  • Loading branch information
@EmilyBourne
EmilyBourne committed Jul 25, 2024
1 parent 6ad7e85 commit 5cb4661
Show file tree
Hide file tree
Showing 3 changed files with 276 additions and 218 deletions.
8 changes: 4 additions & 4 deletions simulations/geometry5D/testcollisions/testcollisions.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,11 +230,11 @@ int main(int argc, char** argv)
dom_sp_tor3D_v2D,
coeff_intdmu.span_cview(),
coeff_intdvpar.span_cview(),
nustar0_r.span_view(),
nustar0_r.span_cview(),
collisions_interspecies,
field_grid_tor1.span_view(),
safety_factor.span_view(),
B_norm.span_view());
field_grid_tor1.span_cview(),
safety_factor.span_cview(),
B_norm.span_cview());

steady_clock::time_point const start = steady_clock::now();

Expand Down
247 changes: 33 additions & 214 deletions src/collisions/CollisionSpVparMu.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,93 +6,9 @@
#include <ddc_helper.hpp>
#include <species_info.hpp>

#include "collisions_dimensions.hpp"
#include "koliop_interface.hpp"

/**
* @brief A namespace to collect classes which are necessary to create Chunks with the
* correct number of dimensions to be compatible with Koliop.
*/
namespace collisions_dimensions {
/**
* Create dimensions to act as the radial/poloidal/toroidal tag for Koliop even if these dims don't
* exist in the simulation
*/

/// Class from which fake dimensions inherit. These fake dimensions are inserted in order to match Koliop interface
struct InternalSpoofGrid
{
};

/// Fake radial dimension to be used if there is no radial dimension in the simulation
struct InternalSpoofGridR : InternalSpoofGrid
{
};

/// Fake poloidal dimension to be used if there is no poloidal dimension in the simulation
struct InternalSpoofGridTheta : InternalSpoofGrid
{
};

/// Check if a dimension is spoofed but is not present in the actual simulation
template <class Grid>
inline constexpr bool is_spoofed_dim_v = std::is_base_of_v<InternalSpoofGrid, Grid>;

/**
* Class to get the type of the radial dimension from a field containing a radial profile.
* @tparam Field The type of the field containing the radial profile.
*/
template <class Field>
struct ExtractRDim
{
static_assert(!std::is_same_v<Field, Field>, "Unrecognised radial profile type");
};

/**
* Class to get the type of the poloidal dimension from a field containing a profile on the poloidal plane.
* @tparam Field The type of the field containing the profile on the poloidal plane.
* @tparam GridR The tag for the discrete radial dimension.
*/
template <class Field, class GridR>
struct ExtractThetaDim
{
static_assert(!std::is_same_v<Field, Field>, "Unrecognised poloidal profile type");
};

/**
* @brief Get the index range for a specific grid from a multi-D index range.
* If the dimension is spoofed and does not appear in the multi-D index range then an index
* range which only iterates over the index(0) is returned.
*
* @tparam Grid The tag for the specific grid.
* @param idx_range The multi-D index range.
* @returns The index range for the specific grid.
*/
template <class Grid, class FDistribDomain>
inline ddc::DiscreteDomain<Grid> get_1d_idx_range(FDistribDomain idx_range)
{
if constexpr (is_spoofed_dim_v<Grid>) {
return ddc::DiscreteDomain<
Grid>(ddc::DiscreteElement<Grid> {0}, ddc::DiscreteVector<Grid> {1});
} else {
return ddc::select<Grid>(idx_range);
}
}

/**
* @brief Get the index range for specific grid dimensions from a multi-D index range.
*
* @tparam Grid The tags for the specific grid dimensions.
* @param idx_range The multi-D index range.
* @returns The index range for the specific grid.
*/
template <class... Grid, class FDistribDomain>
inline ddc::DiscreteDomain<Grid...> get_idx_range(FDistribDomain idx_range)
{
return ddc::DiscreteDomain<Grid...>(get_1d_idx_range<Grid>(idx_range)...);
}

} // namespace collisions_dimensions

/**
* @brief A class which computes the collision operator in (vpar,mu)
*/
Expand All @@ -101,19 +17,23 @@ template <
class GridVpar,
class GridMu,
class InputDFieldR,
class InputDFieldRTheta>
class InputDFieldThetaR>
class CollisionSpVparMu /* : public IRightHandSide */
{
private:
using Species = IDimSp;
using fdistrib_domain_tags = ddc::to_type_seq_t<FDistribDomain>;
using GridR = typename collisions_dimensions::ExtractRDim<InputDFieldR>::type;
using GridTheta =
typename collisions_dimensions::ExtractThetaDim<InputDFieldThetaR, GridR>::type;

// Validate template types
static_assert(ddc::is_discrete_domain_v<FDistribDomain>);
static_assert(FDistribDomain::rank() >= 3 && FDistribDomain::rank() <= 6);
static_assert((std::is_same_v<InputDFieldR, double>) || ddc::is_borrowed_chunk_v<InputDFieldR>);
static_assert(
(std::is_same_v<InputDFieldRTheta, double>)
|| (ddc::is_borrowed_chunk_v<InputDFieldRTheta>));
(std::is_same_v<InputDFieldThetaR, double>)
|| (ddc::is_borrowed_chunk_v<InputDFieldThetaR>));
// Ensure expected types appear in distribution domain
static_assert(
ddc::in_tags_v<Species, fdistrib_domain_tags>,
Expand All @@ -124,37 +44,31 @@ class CollisionSpVparMu /* : public IRightHandSide */
static_assert(
ddc::in_tags_v<GridMu, fdistrib_domain_tags>,
"Mu is missing from distribution function domain");
// Ensure that uniform
// [TODO] Restore it as soon as the geometry5D is deleted
//static_assert(ddc::is_uniform_point_sampling_v<GridVpar>);

// Ensure expected types appear in distribution domain at expected position
static_assert(
ddc::type_seq_rank_v<Species, fdistrib_domain_tags> == 0,
"Species should appear first in the distribution function domain");
static_assert(
ddc::type_seq_rank_v<GridMu, fdistrib_domain_tags> == (FDistribDomain::rank() - 1),
"Mu should appear last in the distribution function domain");
static_assert(
ddc::type_seq_rank_v<GridVpar, fdistrib_domain_tags> == (FDistribDomain::rank() - 2),
"Vpar should appear second to last in the distribution function domain");
// Ensure that uniform
// [TODO] Restore it as soon as the geometry5D is deleted
//static_assert(ddc::is_uniform_point_sampling_v<GridVpar>);

private:
using GridR = typename collisions_dimensions::ExtractRDim<InputDFieldR>::type;
using GridTheta =
typename collisions_dimensions::ExtractThetaDim<InputDFieldRTheta, GridR>::type;
collisions_dimensions::
order_of_last_grids<fdistrib_domain_tags, GridTheta, GridR, GridVpar, GridMu>(),
"Misordered distribution function domain detected. Koliop expects (Sp, Phi, Theta, R, "
"Vpar, Mu)");

public:
/// Type alias for the domain of the radial points
using DDomR = ddc::DiscreteDomain<GridR>;
using DDomRTheta = ddc::DiscreteDomain<GridR, GridTheta>;
using DDomSpRThetaVpar = ddc::DiscreteDomain<GridVpar, GridR, GridTheta, Species>;
/// Type alias for the domain of the poloidal plane
using DDomThetaR = ddc::DiscreteDomain<GridTheta, GridR>;
/// Type alias for the domain containing (species, theta, r, vpar)
using DDomSpThetaRVpar = ddc::DiscreteDomain<Species, GridTheta, GridR, GridVpar>;

public:
/// Type alias for the domain of the magnetic moment.
using DDomMu = ddc::DiscreteDomain<GridMu>;
/// Type alias for the domain of the velocity parallel to the magnetic field.
using DDomVpar = ddc::DiscreteDomain<GridVpar>;

public:
/// Type alias for a field on a grid of species
using DFieldSp = device_t<ddc::Chunk<double, IDomainSp>>;
/// Type alias for a field on a grid of radial values
Expand All @@ -164,13 +78,13 @@ class CollisionSpVparMu /* : public IRightHandSide */
/// Type alias for a field on a grid of parallel velocities
using DFieldVpar = device_t<ddc::Chunk<double, DDomVpar>>;
/// Type alias for a field on a grid on a poloidal plane
using DFieldRTheta = device_t<ddc::Chunk<double, DDomRTheta>>;
using DFieldThetaR = device_t<ddc::Chunk<double, DDomThetaR>>;
/// Type alias for a field on a grid of species, poloidal plane and parallel velocities
using DFieldSpRThetaVpar = device_t<ddc::Chunk<double, DDomSpRThetaVpar>>;
using DFieldSpThetaRVpar = device_t<ddc::Chunk<double, DDomSpThetaRVpar>>;
/// Type alias for a span of a field defined on a grid of radial values
using DSpanR = device_t<ddc::ChunkSpan<double, DDomR>>;
/// Type alias for a span of a field defined on a grid on a poloidal plane
using DSpanRTheta = device_t<ddc::ChunkSpan<double, DDomRTheta>>;
using DSpanThetaR = device_t<ddc::ChunkSpan<double, DDomThetaR>>;
/// Type alias for a constant reference to a Chunk on GPU defined on a grid of magnetic moments.
using DViewMu = ddc::ChunkSpan<
double const,
Expand Down Expand Up @@ -214,7 +128,7 @@ class CollisionSpVparMu /* : public IRightHandSide */
* @param[in] src The source from which the data is copied.
* @param[out] dst The destination into which the data is copied.
*/
void deepcopy_poloidal_plane(DSpanRTheta dst, InputDFieldRTheta src)
void deepcopy_poloidal_plane(DSpanThetaR dst, InputDFieldThetaR src)
{
if constexpr ((collisions_dimensions::is_spoofed_dim_v<GridR>)&&(
collisions_dimensions::is_spoofed_dim_v<GridTheta>)) {
Expand All @@ -228,7 +142,7 @@ class CollisionSpVparMu /* : public IRightHandSide */
ddc::parallel_for_each(
Kokkos::DefaultExecutionSpace(),
dst.domain(),
KOKKOS_LAMBDA(ddc::DiscreteElement<GridR, GridTheta> idx) {
KOKKOS_LAMBDA(ddc::DiscreteElement<GridTheta, GridR> idx) {
dst(idx) = src(ddc::select<NonSpoofDim>(idx));
});
}
Expand Down Expand Up @@ -265,7 +179,7 @@ class CollisionSpVparMu /* : public IRightHandSide */
std::int8_t const collisions_interspecies,
InputDFieldR rg,
InputDFieldR safety_factor,
InputDFieldRTheta B_norm)
InputDFieldThetaR B_norm)
: m_operator_handle {}
, m_comb_mat {Kokkos::view_alloc(Kokkos::WithoutInitializing, "m_comb_mat")}
, m_hat_As {"m_hat_As", collisions_dimensions::get_idx_range<Species>(fdistrib_domain)}
Expand All @@ -274,41 +188,14 @@ class CollisionSpVparMu /* : public IRightHandSide */
, m_rg {"m_rg", collisions_dimensions::get_idx_range<GridR>(fdistrib_domain)}
, m_safety_factor {"m_safety_factor", collisions_dimensions::get_idx_range<GridR>(fdistrib_domain)}
, m_mask_buffer_r {"m_mask_buffer_r", collisions_dimensions::get_idx_range<GridR>(fdistrib_domain)}
, m_mask_LIM {"m_mask_LIM", DDomRTheta {collisions_dimensions::get_idx_range<GridR, GridTheta>(fdistrib_domain)}}
, m_B_norm {"m_B_norm", DDomRTheta {collisions_dimensions::get_idx_range<GridR, GridTheta>(fdistrib_domain)}}
, m_Bstar_s {"m_Bstar_s", DDomSpRThetaVpar {collisions_dimensions::get_idx_range<Species, GridR, GridTheta, GridVpar>(fdistrib_domain)}}
, m_mask_LIM {"m_mask_LIM", DDomThetaR {collisions_dimensions::get_idx_range<GridTheta, GridR>(fdistrib_domain)}}
, m_B_norm {"m_B_norm", DDomThetaR {collisions_dimensions::get_idx_range<GridTheta, GridR>(fdistrib_domain)}}
, m_Bstar_s {"m_Bstar_s", DDomSpThetaRVpar {collisions_dimensions::get_idx_range<Species, GridTheta, GridR, GridVpar>(fdistrib_domain)}}
, m_mug {"m_mug", ddc::select<GridMu>(fdistrib_domain)}
, m_vparg {"m_vparg", ddc::select<GridVpar>(fdistrib_domain)}
{
using namespace collisions_dimensions;
// Check that the distribution function is correctly ordered
if constexpr (!is_spoofed_dim_v<GridR>) {
if constexpr (!is_spoofed_dim_v<GridTheta>) {
static_assert(
ddc::type_seq_rank_v<
GridR,
fdistrib_domain_tags> == (FDistribDomain::rank() - 3),
"R should appear third to last in the distribution function domain");
static_assert(
ddc::type_seq_rank_v<
GridTheta,
fdistrib_domain_tags> == (FDistribDomain::rank() - 4),
"Theta should appear fourth to last in the distribution function domain");
} else {
static_assert(
ddc::type_seq_rank_v<
GridR,
fdistrib_domain_tags> == (FDistribDomain::rank() - 3),
"R should appear third to last in the distribution function domain");
}
} else if constexpr (!is_spoofed_dim_v<GridTheta>) {
static_assert(
ddc::type_seq_rank_v<
GridTheta,
fdistrib_domain_tags> == (FDistribDomain::rank() - 3),
"Theta should appear third to last in the distribution function domain");
}

koliop_interface::DoCombMatComputation(m_comb_mat);

IDomainSp idxrange_sp = ddc::select<Species>(fdistrib_domain);
Expand Down Expand Up @@ -428,83 +315,15 @@ class CollisionSpVparMu /* : public IRightHandSide */
// [TODO]: This mask should maybe be deleted in C++ version
DFieldR m_mask_buffer_r;
/// Limiter mask in (r,theta)
DFieldRTheta m_mask_LIM;
DFieldThetaR m_mask_LIM;
/// B norm in (r,theta)
// [TODO] Attention this must be 3D for generalization to 3D geometry--> transfer it in a 1D array ?
DFieldRTheta m_B_norm;
DFieldThetaR m_B_norm;
/// Bstar(species,r,theta,vpar)
// [TODO] Must be 5D for full 3D geometry
DFieldSpRThetaVpar m_Bstar_s;
DFieldSpThetaRVpar m_Bstar_s;
/// grid in mu direction
DFieldMu m_mug;
/// grid in vpar direction
DFieldVpar m_vparg;
};

namespace collisions_dimensions {

/// If radial profile is stored in a double then the grid tag must be spoofed.
template <>
struct ExtractRDim<double>
{
using type = InternalSpoofGridR;
};

/// If radial profile is stored in a 1D chunk then the grid tag is extracted.
template <class GridR, class Layout>
struct ExtractRDim<ddc::ChunkSpan<
double,
ddc::DiscreteDomain<GridR>,
Layout,
Kokkos::DefaultExecutionSpace::memory_space>>
{
using type = GridR;
};

/// If the profile on the poloidal plane is stored in a double then the grid tag must be spoofed.
template <>
struct ExtractThetaDim<double, InternalSpoofGridR>
{
using type = InternalSpoofGridTheta;
};

/// If the profile on the poloidal plane is stored in a chunk on radial values then the grid tag must be spoofed.
template <class GridR, class Layout>
struct ExtractThetaDim<
ddc::ChunkSpan<
double,
ddc::DiscreteDomain<GridR>,
Layout,
Kokkos::DefaultExecutionSpace::memory_space>,
GridR>
{
using type = InternalSpoofGridTheta;
};

/// If the profile on the poloidal plane is stored in a chunk on poloidal values then the grid tag is extracted.
template <class GridR, class GridTheta, class Layout>
struct ExtractThetaDim<
ddc::ChunkSpan<
double,
ddc::DiscreteDomain<GridTheta>,
Layout,
Kokkos::DefaultExecutionSpace::memory_space>,
GridR>
{
using type = GridTheta;
};

/// If the profile on the poloidal plane is stored in a 2D chunk then the grid tag is extracted.
template <class GridR, class GridTheta, class Layout>
struct ExtractThetaDim<
ddc::ChunkSpan<
double,
ddc::DiscreteDomain<GridTheta, GridR>,
Layout,
Kokkos::DefaultExecutionSpace::memory_space>,
GridR>
{
using type = GridTheta;
};

} // namespace collisions_dimensions
Loading

0 comments on commit 5cb4661

Please sign in to comment.