Skip to content

Commit

Permalink
Avoid N² Operation in Binary-Pairing Coulomb Collsions (ECP-WarpX#5066)
Browse files Browse the repository at this point in the history 
* moved density and temperature calc from each pair to once for each cell.
* pulled out n12 calc to cell level.
* fixed dV issue for cyl geom.
  • Loading branch information
@JustinRayAngus
JustinRayAngus authored Jul 20, 2024
1 parent 8bf0fb0 commit c6bb95a
Show file tree
Hide file tree
Showing 5 changed files with 207 additions and 60 deletions.
172 changes: 172 additions & 0 deletions Source/Particles/Collision/BinaryCollision/BinaryCollision.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#define WARPX_PARTICLES_COLLISION_BINARYCOLLISION_H_

#include "Particles/Collision/BinaryCollision/Coulomb/PairWiseCoulombCollisionFunc.H"
#include "Particles/Collision/BinaryCollision/Coulomb/ComputeTemperature.H"
#include "Particles/Collision/BinaryCollision/DSMC/DSMCFunc.H"
#include "Particles/Collision/BinaryCollision/NuclearFusion/NuclearFusionFunc.H"
#include "Particles/Collision/BinaryCollision/ParticleCreationFunc.H"
Expand Down Expand Up @@ -360,6 +361,21 @@ public:
End of calculations only required when creating product particles
*/

// create vectors to store density and temperature on cell level
amrex::Gpu::DeviceVector<amrex::ParticleReal> n1_vec;
amrex::Gpu::DeviceVector<amrex::ParticleReal> n12_vec;
amrex::Gpu::DeviceVector<amrex::ParticleReal> T1_vec;
if (binary_collision_functor.m_computeSpeciesDensities) {
n1_vec.resize(n_cells);
n12_vec.resize(n_cells);
}
if (binary_collision_functor.m_computeSpeciesTemperatures) {
T1_vec.resize(n_cells);
}
amrex::ParticleReal* AMREX_RESTRICT n1_in_each_cell = n1_vec.dataPtr();
amrex::ParticleReal* AMREX_RESTRICT n12_in_each_cell = n12_vec.dataPtr();
amrex::ParticleReal* AMREX_RESTRICT T1_in_each_cell = T1_vec.dataPtr();

// Loop over cells
amrex::ParallelForRNG( n_cells,
[=] AMREX_GPU_DEVICE (int i_cell, amrex::RandomEngine const& engine) noexcept
Expand All @@ -368,12 +384,60 @@ public:
// given by the `indices_1[cell_start_1:cell_stop_1]`
index_type const cell_start_1 = cell_offsets_1[i_cell];
index_type const cell_stop_1 = cell_offsets_1[i_cell+1];
index_type const cell_half_1 = (cell_start_1+cell_stop_1)/2;

// Do not collide if there is only one particle in the cell
if ( cell_stop_1 - cell_start_1 <= 1 ) { return; }

// compute local density [1/m^3]
if (binary_collision_functor.m_computeSpeciesDensities) {
amrex::ParticleReal wtot1 = 0.0;
amrex::ParticleReal * const AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
for (index_type i1=cell_start_1; i1<cell_stop_1; ++i1) {
wtot1 += w1[ indices_1[i1] ];
}
#if defined WARPX_DIM_RZ
const int ri = (i_cell - i_cell%nz) / nz;
auto dV = MathConst::pi*(2.0_prt*ri+1.0_prt)*dr*dr*dz;
#endif
n1_in_each_cell[i_cell] = wtot1/dV;
}

// compute local temperature [Joules]
if (binary_collision_functor.m_computeSpeciesTemperatures) {
amrex::ParticleReal * const AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
amrex::ParticleReal * const AMREX_RESTRICT u1x = soa_1.m_rdata[PIdx::ux];
amrex::ParticleReal * const AMREX_RESTRICT u1y = soa_1.m_rdata[PIdx::uy];
amrex::ParticleReal * const AMREX_RESTRICT u1z = soa_1.m_rdata[PIdx::uz];
T1_in_each_cell[i_cell] = ComputeTemperature( cell_start_1, cell_stop_1, indices_1,
w1, u1x, u1y, u1z, m1 );
}

// shuffle
ShuffleFisherYates(indices_1, cell_start_1, cell_stop_1, engine);

// compute n12 for intra-species
if (binary_collision_functor.m_computeSpeciesDensities) {
amrex::ParticleReal n12 = 0.0;
index_type const NI1 = cell_half_1 - cell_start_1;
index_type const NI2 = cell_stop_1 - cell_half_1;
index_type const max_N = amrex::max(NI1,NI2);
index_type i1 = cell_start_1;
index_type i2 = cell_half_1;
amrex::ParticleReal * const AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
for (index_type k=0; k<max_N; ++k) {
n12 += amrex::min( w1[ indices_1[i1] ], w1[ indices_1[i2] ] );
++i1; if ( i1 == cell_half_1 ) { i1 = cell_start_1; }
++i2; if ( i2 == cell_stop_1 ) { i2 = cell_half_1; }
}
#if defined WARPX_DIM_RZ
const int ri = (i_cell - i_cell%nz) / nz;
auto dV = MathConst::pi*(2.0_prt*ri+1.0_prt)*dr*dr*dz;
#endif
n12 = 2.0_prt * n12 / dV;
n12_in_each_cell[i_cell] = n12;
}

}
);

Expand Down Expand Up @@ -410,6 +474,19 @@ public:
const int ri = (i_cell - i_cell%nz) / nz;
auto dV = MathConst::pi*(2.0_prt*ri+1.0_prt)*dr*dr*dz;
#endif

// Get the local density and temperature for this cell
amrex::ParticleReal n1 = 0.0;
amrex::ParticleReal n12 = 0.0;
amrex::ParticleReal T1 = 0.0;
if (binary_collision_functor.m_computeSpeciesDensities) {
n1 = n1_in_each_cell[i_cell];
n12 = n12_in_each_cell[i_cell];
}
if (binary_collision_functor.m_computeSpeciesTemperatures) {
T1 = T1_in_each_cell[i_cell];
}

// Call the function in order to perform collisions
// If there are product species, mask, p_pair_indices_1/2, and
// p_pair_reaction_weight are filled here
Expand All @@ -418,6 +495,7 @@ public:
cell_half_1, cell_stop_1,
indices_1, indices_1,
soa_1, soa_1, get_position_1, get_position_1,
n1, n1, n12, T1, T1,
q1, q1, m1, m1, dt, dV, coll_idx,
cell_start_pair, p_mask, p_pair_indices_1, p_pair_indices_2,
p_pair_reaction_weight, engine);
Expand Down Expand Up @@ -562,6 +640,24 @@ public:
End of calculations only required when creating product particles
*/

// create vectors to store density and temperature on cell level
amrex::Gpu::DeviceVector<amrex::ParticleReal> n1_vec, n2_vec;
amrex::Gpu::DeviceVector<amrex::ParticleReal> n12_vec;
amrex::Gpu::DeviceVector<amrex::ParticleReal> T1_vec, T2_vec;
if (binary_collision_functor.m_computeSpeciesDensities) {
n1_vec.resize(n_cells);
n2_vec.resize(n_cells);
n12_vec.resize(n_cells);
}
if (binary_collision_functor.m_computeSpeciesTemperatures) {
T1_vec.resize(n_cells);
T2_vec.resize(n_cells);
}
amrex::ParticleReal* AMREX_RESTRICT n1_in_each_cell = n1_vec.dataPtr();
amrex::ParticleReal* AMREX_RESTRICT n2_in_each_cell = n2_vec.dataPtr();
amrex::ParticleReal* AMREX_RESTRICT n12_in_each_cell = n12_vec.dataPtr();
amrex::ParticleReal* AMREX_RESTRICT T1_in_each_cell = T1_vec.dataPtr();
amrex::ParticleReal* AMREX_RESTRICT T2_in_each_cell = T2_vec.dataPtr();

// Loop over cells
amrex::ParallelForRNG( n_cells,
Expand All @@ -579,13 +675,73 @@ public:
// ux_1[ indices_1[i] ], where i is between
// cell_start_1 (inclusive) and cell_start_2 (exclusive)

// compute local densities [1/m^3]
if (binary_collision_functor.m_computeSpeciesDensities) {
amrex::ParticleReal w1tot = 0.0;
amrex::ParticleReal w2tot = 0.0;
amrex::ParticleReal * const AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
amrex::ParticleReal * const AMREX_RESTRICT w2 = soa_2.m_rdata[PIdx::w];
for (index_type i1=cell_start_1; i1<cell_stop_1; ++i1) {
w1tot += w1[ indices_1[i1] ];
}
for (index_type i2=cell_start_2; i2<cell_stop_2; ++i2) {
w2tot += w2[ indices_2[i2] ];
}
#if defined WARPX_DIM_RZ
const int ri = (i_cell - i_cell%nz) / nz;
auto dV = MathConst::pi*(2.0_prt*ri+1.0_prt)*dr*dr*dz;
#endif
n1_in_each_cell[i_cell] = w1tot/dV;
n2_in_each_cell[i_cell] = w2tot/dV;
}

// compute local temperatures [Joules]
if (binary_collision_functor.m_computeSpeciesTemperatures) {
amrex::ParticleReal * const AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
amrex::ParticleReal * const AMREX_RESTRICT u1x = soa_1.m_rdata[PIdx::ux];
amrex::ParticleReal * const AMREX_RESTRICT u1y = soa_1.m_rdata[PIdx::uy];
amrex::ParticleReal * const AMREX_RESTRICT u1z = soa_1.m_rdata[PIdx::uz];
T1_in_each_cell[i_cell] = ComputeTemperature( cell_start_1, cell_stop_1, indices_1,
w1, u1x, u1y, u1z, m1 );

amrex::ParticleReal * const AMREX_RESTRICT w2 = soa_2.m_rdata[PIdx::w];
amrex::ParticleReal * const AMREX_RESTRICT u2x = soa_2.m_rdata[PIdx::ux];
amrex::ParticleReal * const AMREX_RESTRICT u2y = soa_2.m_rdata[PIdx::uy];
amrex::ParticleReal * const AMREX_RESTRICT u2z = soa_2.m_rdata[PIdx::uz];
T2_in_each_cell[i_cell] = ComputeTemperature( cell_start_2, cell_stop_2, indices_2,
w2, u2x, u2y, u2z, m2 );
}

// Do not collide if one species is missing in the cell
if ( cell_stop_1 - cell_start_1 < 1 ||
cell_stop_2 - cell_start_2 < 1 ) { return; }

// shuffle
ShuffleFisherYates(indices_1, cell_start_1, cell_stop_1, engine);
ShuffleFisherYates(indices_2, cell_start_2, cell_stop_2, engine);

// compute n12 for inter-species
if (binary_collision_functor.m_computeSpeciesDensities) {
amrex::ParticleReal n12 = 0.0;
index_type const NI1 = cell_stop_1 - cell_start_1;
index_type const NI2 = cell_stop_2 - cell_start_2;
index_type const max_N = amrex::max(NI1,NI2);
index_type i1 = cell_start_1;
index_type i2 = cell_start_2;
amrex::ParticleReal * const AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
amrex::ParticleReal * const AMREX_RESTRICT w2 = soa_2.m_rdata[PIdx::w];
for (index_type k=0; k<max_N; ++k) {
n12 += amrex::min( w1[ indices_1[i1] ], w2[ indices_1[i2] ] );
++i1; if ( i1 == cell_stop_1 ) { i1 = cell_start_1; }
++i2; if ( i2 == cell_stop_2 ) { i2 = cell_start_2; }
}
#if defined WARPX_DIM_RZ
const int ri = (i_cell - i_cell%nz) / nz;
auto dV = MathConst::pi*(2.0_prt*ri+1.0_prt)*dr*dr*dz;
#endif
n12 = n12 / dV;
n12_in_each_cell[i_cell] = n12;
}
}
);

Expand Down Expand Up @@ -628,13 +784,29 @@ public:
const int ri = (i_cell - i_cell%nz) / nz;
auto dV = MathConst::pi*(2.0_prt*ri+1.0_prt)*dr*dr*dz;
#endif

// Get the local densities and temperatures for this cell
amrex::ParticleReal n1 = 0.0, n2 = 0.0;
amrex::ParticleReal n12 = 0.0;
amrex::ParticleReal T1 = 0.0, T2 = 0.0;
if (binary_collision_functor.m_computeSpeciesDensities) {
n1 = n1_in_each_cell[i_cell];
n2 = n2_in_each_cell[i_cell];
n12 = n12_in_each_cell[i_cell];
}
if (binary_collision_functor.m_computeSpeciesTemperatures) {
T1 = T1_in_each_cell[i_cell];
T2 = T2_in_each_cell[i_cell];
}

// Call the function in order to perform collisions
// If there are product species, p_mask, p_pair_indices_1/2, and
// p_pair_reaction_weight are filled here
binary_collision_functor(
cell_start_1, cell_stop_1, cell_start_2, cell_stop_2,
indices_1, indices_2,
soa_1, soa_2, get_position_1, get_position_2,
n1, n2, n12, T1, T2,
q1, q2, m1, m2, dt, dV, coll_idx,
cell_start_pair, p_mask, p_pair_indices_1, p_pair_indices_2,
p_pair_reaction_weight, engine);
Expand Down
69 changes: 13 additions & 56 deletions Source/Particles/Collision/BinaryCollision/Coulomb/ElasticCollisionPerez.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
#ifndef WARPX_PARTICLES_COLLISION_ELASTIC_COLLISION_PEREZ_H_
#define WARPX_PARTICLES_COLLISION_ELASTIC_COLLISION_PEREZ_H_

#include "ComputeTemperature.H"
#include "UpdateMomentumPerezElastic.H"
#include "Particles/WarpXParticleContainer.H"
#include "Utils/WarpXConst.H"
Expand All @@ -25,18 +24,16 @@
* @param[in] I1e,I2e is the stop index for I1,I2 (exclusive).
* @param[in] I1,I2 the index arrays. They determine all elements that will be used.
* @param[in,out] soa_1,soa_2 the struct of array for species 1/2
* @param[in] n1,n2 density of species 1/2
* @param[in] n12 density parameter for s12
* @param[in] T1,T2 temperature [Joules] of species 1/2
* only used if L <= 0
* @param[in] q1,q2 charge of species 1/2
* @param[in] m1,m2 mass of species 1/2
* @param[in] T1 temperature (Joule) of species 1
* and will be used if greater than zero,
* otherwise will be computed.
* @param[in] T2 temperature (Joule) of species 2, @see T1
* @param[in] dt is the time step length between two collision calls.
* @param[in] L is the Coulomb log and will be used if greater than zero,
* otherwise will be computed.
* @param[in] dV is the volume of the corresponding cell.
* @param[in] engine the random number generator state & factory
* @param[in] isSameSpecies whether this is an intra-species collision process
* @param[in] coll_idx is the collision index offset.
*/

Expand All @@ -48,12 +45,14 @@ void ElasticCollisionPerez (
T_index const* AMREX_RESTRICT I1,
T_index const* AMREX_RESTRICT I2,
SoaData_type soa_1, SoaData_type soa_2,
T_PR const n1, T_PR const n2,
T_PR const n12,
T_PR const T1, T_PR const T2,
T_PR const q1, T_PR const q2,
T_PR const m1, T_PR const m2,
T_PR const T1, T_PR const T2,
T_R const dt, T_PR const L, T_R const dV,
T_R const dt, T_PR const L,
amrex::RandomEngine const& engine,
bool const isSameSpecies, T_index coll_idx)
T_index coll_idx)
{
const T_index NI1 = I1e - I1s;
const T_index NI2 = I2e - I2s;
Expand All @@ -70,53 +69,11 @@ void ElasticCollisionPerez (
T_PR * const AMREX_RESTRICT u2y = soa_2.m_rdata[PIdx::uy];
T_PR * const AMREX_RESTRICT u2z = soa_2.m_rdata[PIdx::uz];

// get local T1t and T2t
T_PR T1t; T_PR T2t;
if ( T1 <= T_PR(0.0) && L <= T_PR(0.0) )
{
T1t = ComputeTemperature(I1s,I1e,I1,w1,u1x,u1y,u1z,m1);
}
else { T1t = T1; }
if ( T2 <= T_PR(0.0) && L <= T_PR(0.0) )
{
T2t = ComputeTemperature(I2s,I2e,I2,w2,u2x,u2y,u2z,m2);
}
else { T2t = T2; }

// local density
T_PR n1 = T_PR(0.0);
T_PR n2 = T_PR(0.0);
T_PR n12 = T_PR(0.0);
for (T_index i1=I1s; i1<I1e; ++i1) { n1 += w1[ I1[i1] ]; }
for (T_index i2=I2s; i2<I2e; ++i2) { n2 += w2[ I2[i2] ]; }
// Intra-species: the density is in fact the sum of the density of
// each sub-group of particles
if (isSameSpecies) {
n1 = n1 + n2;
n2 = n1;
}
n1 = n1 / dV; n2 = n2 / dV;
{
T_index i1 = I1s; T_index i2 = I2s;
for (T_index k = 0; k < max_N; ++k)
{
n12 += amrex::min( w1[ I1[i1] ], w2[ I2[i2] ] );
++i1; if ( i1 == I1e ) { i1 = I1s; }
++i2; if ( i2 == I2e ) { i2 = I2s; }
}
n12 = n12 / dV;
}
// Intra-species: Apply correction in collision rate
if (isSameSpecies) { n12 *= T_PR(2.0); }

// compute Debye length lmdD
T_PR lmdD;
if ( T1t < T_PR(0.0) || T2t < T_PR(0.0) ) {
lmdD = T_PR(0.0);
}
else {
lmdD = T_PR(1.0)/std::sqrt( n1*q1*q1/(T1t*PhysConst::ep0) +
n2*q2*q2/(T2t*PhysConst::ep0) );
T_PR lmdD = T_PR(-1.0);
if ( L <= T_PR(0.0) ) {
lmdD = T_PR(1.0)/std::sqrt( n1*q1*q1/(T1*PhysConst::ep0) +
n2*q2*q2/(T2*PhysConst::ep0) );
}
T_PR rmin = std::pow( T_PR(4.0) * MathConst::pi / T_PR(3.0) *
amrex::max(n1,n2), T_PR(-1.0/3.0) );
Expand Down
Loading

0 comments on commit c6bb95a

Please sign in to comment.