Skip to content

Commit

Permalink
Add option for Cartesian velocity grid with `VelocityCoincidenceThinn…
Browse files Browse the repository at this point in the history 
…ing` (#4886)

* initial implementation of the grid based merging routine

* add compiler directives for different dimensions

* avoid implicit capture of `this` in lambda

* code clean-up

* use `PIdx::x` for r-coordinate while 4667 is still being reviewed

* fix clang-tidy errors

* another clang-tidy fix

* improve doc-strings

* use iterative heap sort since "SYCL kernel cannot call a recursive function"

* fix clang-tidy error

* fix sorting bug; add documentation and CI test

* add dummy CI benchmark file to get proper values

* update benchmark values based on Azure results

* rename `GridBasedMerging` -> `VelocityCoincidenceThinning`

* use `Algorithms::KineticEnergy`

* reorganize merging loop

* update CI benchmark values

* Relativistic correction in product particles' velocity calculation

* update benchmark values after changing energy calculation

* handle edge case with zero cluster momentum

* call redistribute after particle resampling to remove invalid particles

* use unsigned ints for indexing

* Revert "use unsigned ints for indexing"

This reverts commit abe027f.

* call `Redistribute` before merging

* code clean-up

* also check for `std::isnan` in edge case handling

* add reference for grid based merging

* check that cluster has total weight > 0 before merging

* add defense against numerical error leading to nan

* make resampling message more verbose

* use `deleteInvalidParticles` instead of `Redistribute`

* remove default values for merging parameters

* remove doc-string default specifications in picmi.py

* apply suggestions from code review

* update benchmark values; avoid possible nans

* add assert to prevent merging of photons

* add `BackwardCompatibility` check to `LevelingThinning`

* implement option for Cartesian velocity mesh

* use `enum` for velocity grid types

* fix Windows `uint` issue and clang-tidy error

* use `Reduce` functions for GPU compatibility

* merge kernels finding min and max velocities

* take array of du values in picmi

* avoid checking velocity grid type for every particle

* fix issue in picmi

* use array input for velocity bin sizes
  • Loading branch information
@roelof-groenewald
roelof-groenewald authored May 2, 2024
1 parent c5a5732 commit 8a9bedc
Show file tree
Hide file tree
Showing 3 changed files with 187 additions and 45 deletions.
24 changes: 20 additions & 4 deletions Python/pywarpx/picmi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,17 +147,27 @@ class Species(picmistandard.PICMI_Species):
warpx_resampling_algorithm: str, default="leveling_thinning"
Resampling algorithm to use.
warpx_resampling_algorithm_velocity_grid_type: str, default="spherical"
Type of grid to use when clustering particles in velocity space. Only
applicable with the `velocity_coincidence_thinning` algorithm.
warpx_resampling_algorithm_delta_ur: float
Size of velocity window used for clustering particles during grid-based
merging.
merging, with `velocity_grid_type == "spherical"`.
warpx_resampling_algorithm_n_theta: int
Number of bins to use in theta when clustering particle velocities
during grid-based merging.
during grid-based merging, with `velocity_grid_type == "spherical"`.
warpx_resampling_algorithm_n_phi: int
Number of bins to use in phi when clustering particle velocities
during grid-based merging.
during grid-based merging, with `velocity_grid_type == "spherical"`.
warpx_resampling_algorithm_delta_u: array of floats or float
Size of velocity window used in ux, uy and uz for clustering particles
during grid-based merging, with `velocity_grid_type == "cartesian"`. If
a single number is given the same du value will be used in all three
directions.
"""
def init(self, kw):

Expand Down Expand Up @@ -239,9 +249,13 @@ def init(self, kw):
self.resampling_min_ppc = kw.pop('warpx_resampling_min_ppc', None)
self.resampling_trigger_intervals = kw.pop('warpx_resampling_trigger_intervals', None)
self.resampling_triggering_max_avg_ppc = kw.pop('warpx_resampling_trigger_max_avg_ppc', None)
self.resampling_algorithm_velocity_grid_type = kw.pop('warpx_resampling_algorithm_velocity_grid_type', None)
self.resampling_algorithm_delta_ur = kw.pop('warpx_resampling_algorithm_delta_ur', None)
self.resampling_algorithm_n_theta = kw.pop('warpx_resampling_algorithm_n_theta', None)
self.resampling_algorithm_n_phi = kw.pop('warpx_resampling_algorithm_n_phi', None)
self.resampling_algorithm_delta_u = kw.pop('warpx_resampling_algorithm_delta_u', None)
if np.size(self.resampling_algorithm_delta_u) == 1:
self.resampling_algorithm_delta_u = [self.resampling_algorithm_delta_u]*3

def species_initialize_inputs(self, layout,
initialize_self_fields = False,
Expand Down Expand Up @@ -284,9 +298,11 @@ def species_initialize_inputs(self, layout,
resampling_min_ppc=self.resampling_min_ppc,
resampling_trigger_intervals=self.resampling_trigger_intervals,
resampling_trigger_max_avg_ppc=self.resampling_triggering_max_avg_ppc,
resampling_algorithm_velocity_grid_type=self.resampling_algorithm_velocity_grid_type,
resampling_algorithm_delta_ur=self.resampling_algorithm_delta_ur,
resampling_algorithm_n_theta=self.resampling_algorithm_n_theta,
resampling_algorithm_n_phi=self.resampling_algorithm_n_phi)
resampling_algorithm_n_phi=self.resampling_algorithm_n_phi,
resampling_algorithm_delta_u=self.resampling_algorithm_delta_u)

# add reflection models
self.species.add_new_attr("reflection_model_xlo(E)", self.reflection_model_xlo)
Expand Down
95 changes: 92 additions & 3 deletions Source/Particles/Resampling/VelocityCoincidenceThinning.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@
#ifndef WARPX_VELOCITY_COINCIDENCE_THINNING_H_
#define WARPX_VELOCITY_COINCIDENCE_THINNING_H_

#include "Particles/Algorithms/KineticEnergy.H"
#include "Resampling.H"

#include "Utils/Parser/ParserUtils.H"
#include "Utils/ParticleUtils.H"

/**
* \brief This class implements a particle merging scheme wherein particles
Expand All @@ -33,6 +35,11 @@ public:
*/
VelocityCoincidenceThinning (const std::string& species_name);

enum struct VelocityGridType {
Spherical = 0,
Cartesian = 1
};

/**
* \brief A method that performs merging for the considered species.
*
Expand Down Expand Up @@ -106,10 +113,92 @@ public:
}
};

/**
* \brief Struct used to assign velocity space bin numbers to a given set
* of particles.
*/
struct VelocityBinCalculator {

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void labelOnSphericalVelocityGrid (const amrex::ParticleReal ux[],
const amrex::ParticleReal uy[],
const amrex::ParticleReal uz[],
const unsigned int indices[],
int bin_array[], int index_array[],
const int cell_start, const int cell_stop ) const
{
for (int i = cell_start; i < cell_stop; ++i)
{
// get polar components of the velocity vector
auto u_mag = std::sqrt(
ux[indices[i]]*ux[indices[i]] +
uy[indices[i]]*uy[indices[i]] +
uz[indices[i]]*uz[indices[i]]
);
auto u_theta = std::atan2(uy[indices[i]], ux[indices[i]]) + MathConst::pi;
auto u_phi = std::acos(uz[indices[i]]/u_mag);

const int ii = static_cast<int>(u_theta / dutheta);
const int jj = static_cast<int>(u_phi / duphi);
const int kk = static_cast<int>(u_mag / dur);

bin_array[i] = ii + jj * n1 + kk * n1 * n2;
index_array[i] = i;
}
}

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void labelOnCartesianVelocityGrid (const amrex::ParticleReal ux[],
const amrex::ParticleReal uy[],
const amrex::ParticleReal uz[],
const unsigned int indices[],
int bin_array[], int index_array[],
const int cell_start, const int cell_stop ) const
{
for (int i = cell_start; i < cell_stop; ++i)
{
const int ii = static_cast<int>((ux[indices[i]] - ux_min) / dux);
const int jj = static_cast<int>((uy[indices[i]] - uy_min) / duy);
const int kk = static_cast<int>((uz[indices[i]] - uz_min) / duz);

bin_array[i] = ii + jj * n1 + kk * n1 * n2;
index_array[i] = i;
}
}

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (const amrex::ParticleReal ux[], const amrex::ParticleReal uy[],
const amrex::ParticleReal uz[], const unsigned int indices[],
int bin_array[], int index_array[],
const int cell_start, const int cell_stop) const
{
if (velocity_grid_type == VelocityGridType::Spherical) {
labelOnSphericalVelocityGrid(
ux, uy, uz, indices, bin_array, index_array, cell_start,
cell_stop
);
}
else if (velocity_grid_type == VelocityGridType::Cartesian) {
labelOnCartesianVelocityGrid(
ux, uy, uz, indices, bin_array, index_array, cell_start,
cell_stop
);
}
}

VelocityGridType velocity_grid_type;
int n1, n2;
amrex::ParticleReal dur, dutheta, duphi;
amrex::ParticleReal dux, duy, duz;
amrex::ParticleReal ux_min, uy_min, uz_min, ux_max, uy_max;
};

private:
VelocityGridType m_velocity_grid_type;

int m_min_ppc = 1;
int m_ntheta;
int m_nphi;
int m_ntheta, m_nphi;
amrex::ParticleReal m_delta_ur;
amrex::Vector<amrex::ParticleReal> m_delta_u;
};
#endif // WARPX_VELOCITY_COINCIDENCE_THINNING_H_
113 changes: 75 additions & 38 deletions Source/Particles/Resampling/VelocityCoincidenceThinning.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,6 @@

#include "VelocityCoincidenceThinning.H"

#include "Particles/Algorithms/KineticEnergy.H"
#include "Utils/Parser/ParserUtils.H"
#include "Utils/ParticleUtils.H"


VelocityCoincidenceThinning::VelocityCoincidenceThinning (const std::string& species_name)
{
Expand All @@ -28,15 +24,33 @@ VelocityCoincidenceThinning::VelocityCoincidenceThinning (const std::string& spe
"Resampling min_ppc should be greater than or equal to 1"
);

utils::parser::getWithParser(
pp_species_name, "resampling_algorithm_delta_ur", m_delta_ur
);
utils::parser::getWithParser(
pp_species_name, "resampling_algorithm_n_theta", m_ntheta
);
utils::parser::getWithParser(
pp_species_name, "resampling_algorithm_n_phi", m_nphi
std::string velocity_grid_type_str = "spherical";
pp_species_name.query(
"resampling_algorithm_velocity_grid_type", velocity_grid_type_str
);
if (velocity_grid_type_str == "spherical") {
m_velocity_grid_type = VelocityGridType::Spherical;
utils::parser::getWithParser(
pp_species_name, "resampling_algorithm_delta_ur", m_delta_ur
);
utils::parser::getWithParser(
pp_species_name, "resampling_algorithm_n_theta", m_ntheta
);
utils::parser::getWithParser(
pp_species_name, "resampling_algorithm_n_phi", m_nphi
);
}
else if (velocity_grid_type_str == "cartesian") {
m_velocity_grid_type = VelocityGridType::Cartesian;
utils::parser::getArrWithParser(
pp_species_name, "resampling_algorithm_delta_u", m_delta_u
);
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(m_delta_u.size() == 3,
"resampling_algorithm_delta_u must have three components.");
}
else {
WARPX_ABORT_WITH_MESSAGE("Unkown velocity grid type.");
}
}

void VelocityCoincidenceThinning::operator() (WarpXParIter& pti, const int lev,
Expand All @@ -45,6 +59,7 @@ void VelocityCoincidenceThinning::operator() (WarpXParIter& pti, const int lev,
using namespace amrex::literals;

auto& ptile = pc->ParticlesAt(lev, pti);
const auto n_parts_in_tile = pti.numParticles();
auto& soa = ptile.GetStructOfArrays();
#if defined(WARPX_DIM_XZ) || defined(WARPX_DIM_3D)
auto * const AMREX_RESTRICT x = soa.GetRealData(PIdx::x).data();
Expand Down Expand Up @@ -79,21 +94,57 @@ void VelocityCoincidenceThinning::operator() (WarpXParIter& pti, const int lev,
);

// create a GPU vector to hold the momentum cluster index for each particle
amrex::Gpu::DeviceVector<int> momentum_bin_number(bins.numItems());
amrex::Gpu::DeviceVector<int> momentum_bin_number(n_parts_in_tile);
auto* momentum_bin_number_data = momentum_bin_number.data();

// create a GPU vector to hold the index sorting for the momentum bins
amrex::Gpu::DeviceVector<int> sorted_indices(bins.numItems());
amrex::Gpu::DeviceVector<int> sorted_indices(n_parts_in_tile);
auto* sorted_indices_data = sorted_indices.data();

const auto Ntheta = m_ntheta;
const auto Nphi = m_nphi;

const auto dr = m_delta_ur;
const auto dtheta = 2.0_prt * MathConst::pi / Ntheta;
const auto dphi = MathConst::pi / Nphi;
constexpr auto c2 = PhysConst::c * PhysConst::c;

auto velocityBinCalculator = VelocityBinCalculator();
velocityBinCalculator.velocity_grid_type = m_velocity_grid_type;
if (m_velocity_grid_type == VelocityGridType::Spherical) {
velocityBinCalculator.dur = m_delta_ur;
velocityBinCalculator.n1 = m_ntheta;
velocityBinCalculator.n2 = m_nphi;
velocityBinCalculator.dutheta = 2.0_prt * MathConst::pi / m_ntheta;
velocityBinCalculator.duphi = MathConst::pi / m_nphi;
}
else if (m_velocity_grid_type == VelocityGridType::Cartesian) {
velocityBinCalculator.dux = m_delta_u[0];
velocityBinCalculator.duy = m_delta_u[1];
velocityBinCalculator.duz = m_delta_u[2];

// get the minimum and maximum velocities to determine the velocity space
// grid boundaries
{
using ReduceOpsT = amrex::TypeMultiplier<amrex::ReduceOps,
amrex::ReduceOpMin[3],
amrex::ReduceOpMax[2]>;
using ReduceDataT = amrex::TypeMultiplier<amrex::ReduceData, amrex::ParticleReal[5]>;
ReduceOpsT reduce_op;
ReduceDataT reduce_data(reduce_op);
using ReduceTuple = typename ReduceDataT::Type;
reduce_op.eval(n_parts_in_tile, reduce_data, [=] AMREX_GPU_DEVICE(int i) -> ReduceTuple {
return {ux[i], uy[i], uz[i], ux[i], uy[i]};
});
auto hv = reduce_data.value(reduce_op);
velocityBinCalculator.ux_min = amrex::get<0>(hv);
velocityBinCalculator.uy_min = amrex::get<1>(hv);
velocityBinCalculator.uz_min = amrex::get<2>(hv);
velocityBinCalculator.ux_max = amrex::get<3>(hv);
velocityBinCalculator.uy_max = amrex::get<4>(hv);
}

velocityBinCalculator.n1 = static_cast<int>(
std::ceil((velocityBinCalculator.ux_max - velocityBinCalculator.ux_min) / m_delta_u[0])
);
velocityBinCalculator.n2 = static_cast<int>(
std::ceil((velocityBinCalculator.uy_max - velocityBinCalculator.uy_min) / m_delta_u[1])
);
}
auto heapSort = HeapSort();

// Loop over cells
Expand All @@ -114,24 +165,10 @@ void VelocityCoincidenceThinning::operator() (WarpXParIter& pti, const int lev,

// Loop over particles and label them with the appropriate momentum bin
// number. Also assign initial ordering to the sorted_indices array.
for (int i = cell_start; i < cell_stop; ++i)
{
// get polar components of the velocity vector
auto u_mag = std::sqrt(
ux[indices[i]]*ux[indices[i]] +
uy[indices[i]]*uy[indices[i]] +
uz[indices[i]]*uz[indices[i]]
);
auto u_theta = std::atan2(uy[indices[i]], ux[indices[i]]) + MathConst::pi;
auto u_phi = std::acos(uz[indices[i]]/u_mag);

const auto ii = static_cast<int>(u_theta / dtheta);
const auto jj = static_cast<int>(u_phi / dphi);
const auto kk = static_cast<int>(u_mag / dr);

momentum_bin_number_data[i] = ii + jj * Ntheta + kk * Ntheta * Nphi;
sorted_indices_data[i] = i;
}
velocityBinCalculator(
ux, uy, uz, indices, momentum_bin_number_data, sorted_indices_data,
cell_start, cell_stop
);

// sort indices based on comparing values in momentum_bin_number
heapSort(sorted_indices_data, momentum_bin_number_data, cell_start, cell_numparts);
Expand Down

0 comments on commit 8a9bedc

Please sign in to comment.