WarpX
SplitAndScatterFunc.H
Go to the documentation of this file.
1 /* Copyright 2023-2024 The WarpX Community
2  *
3  * This file is part of WarpX.
4  *
5  * Authors: Roelof Groenewald (TAE Technologies)
6  *
7  * License: BSD-3-Clause-LBNL
8  */
9 #ifndef WARPX_SPLIT_AND_SCATTER_FUNC_H_
10 #define WARPX_SPLIT_AND_SCATTER_FUNC_H_
11 
16 #include "Utils/ParticleUtils.H"
17 
23 {
24  // Define shortcuts for frequently-used type names
31 
32 
33 public:
37  SplitAndScatterFunc () = default;
38 
45  SplitAndScatterFunc (const std::string& collision_name, MultiParticleContainer const * mypc);
46 
55  const index_type& n_total_pairs,
56  // Tile& ptile1, Tile& ptile2,
57  const SoaData_type& /*soa_1*/, const SoaData_type& /*soa_2*/,
58  const amrex::Vector<WarpXParticleContainer*>& pc_products,
59  ParticleTileType** AMREX_RESTRICT tile_products,
60  const amrex::ParticleReal m1, const amrex::ParticleReal m2,
61  const amrex::Vector<amrex::ParticleReal>& /*products_mass*/,
62  const index_type* AMREX_RESTRICT mask,
63  const amrex::Vector<index_type>& products_np,
64  const SmartCopy* AMREX_RESTRICT copy_species1,
65  const SmartCopy* AMREX_RESTRICT copy_species2,
66  const index_type* AMREX_RESTRICT p_pair_indices_1,
67  const index_type* AMREX_RESTRICT p_pair_indices_2,
68  const amrex::ParticleReal* AMREX_RESTRICT p_pair_reaction_weight ) const
69  {
70  using namespace amrex::literals;
71 
72  if (n_total_pairs == 0) { return amrex::Vector<int>(m_num_product_species, 0); }
73 
74  amrex::Gpu::DeviceVector<index_type> offsets(n_total_pairs);
75  index_type* AMREX_RESTRICT offsets_data = offsets.data();
76  const index_type* AMREX_RESTRICT p_offsets = offsets.dataPtr();
77 
78  // The following is used to calculate the appropriate offsets. Note that
79  // a standard cummulative sum is not appropriate since the mask is also
80  // used to specify the type of collision and can therefore have values >1
81  auto const total = amrex::Scan::PrefixSum<index_type>(n_total_pairs,
82  [=] AMREX_GPU_DEVICE (index_type i) -> index_type { return mask[i] ? 1 : 0; },
83  [=] AMREX_GPU_DEVICE (index_type i, index_type s) { offsets_data[i] = s; },
85  );
86 
88  for (int i = 0; i < m_num_product_species; i++)
89  {
90  // How many particles of product species i are created.
91  const index_type num_added = total * m_num_products_host[i];
92  num_added_vec[i] = static_cast<int>(num_added);
93  tile_products[i]->resize(products_np[i] + num_added);
94  }
95 
96  // this works for DSMC since the colliding particles are also products
97  const auto soa_1 = tile_products[0]->getParticleTileData();
98  const auto soa_2 = tile_products[1]->getParticleTileData();
99 
100  amrex::ParticleReal* AMREX_RESTRICT w1 = soa_1.m_rdata[PIdx::w];
101  amrex::ParticleReal* AMREX_RESTRICT w2 = soa_2.m_rdata[PIdx::w];
102  uint64_t* AMREX_RESTRICT idcpu1 = soa_1.m_idcpu;
103  uint64_t* AMREX_RESTRICT idcpu2 = soa_2.m_idcpu;
104 
105  // Create necessary GPU vectors, that will be used in the kernel below
106  amrex::Vector<SoaData_type> soa_products;
107  for (int i = 0; i < m_num_product_species; i++)
108  {
109  soa_products.push_back(tile_products[i]->getParticleTileData());
110  }
111 #ifdef AMREX_USE_GPU
114 
115  amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, soa_products.begin(),
116  soa_products.end(),
117  device_soa_products.begin());
118  amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, products_np.begin(),
119  products_np.end(),
120  device_products_np.begin());
121 
123  SoaData_type* AMREX_RESTRICT soa_products_data = device_soa_products.data();
124  const index_type* AMREX_RESTRICT products_np_data = device_products_np.data();
125 #else
126  SoaData_type* AMREX_RESTRICT soa_products_data = soa_products.data();
127  const index_type* AMREX_RESTRICT products_np_data = products_np.data();
128 #endif
129 
130  const int* AMREX_RESTRICT p_num_products_device = m_num_products_device.data();
131 
132  amrex::ParallelForRNG(n_total_pairs,
133  [=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept
134  {
135  if (mask[i])
136  {
137  // for now we ignore the possibility of having actual reaction
138  // products - only duplicating (splitting) of the colliding
139  // particles is supported.
140 
141  const auto product1_index = products_np_data[0] +
142  (p_offsets[i]*p_num_products_device[0] + 0);
143  // Make a copy of the particle from species 1
144  copy_species1[0](soa_products_data[0], soa_1, static_cast<int>(p_pair_indices_1[i]),
145  static_cast<int>(product1_index), engine);
146  // Set the weight of the new particles to p_pair_reaction_weight[i]
147  soa_products_data[0].m_rdata[PIdx::w][product1_index] = p_pair_reaction_weight[i];
148 
149  const auto product2_index = products_np_data[1] +
150  (p_offsets[i]*p_num_products_device[1] + 0);
151  // Make a copy of the particle from species 2
152  copy_species2[1](soa_products_data[1], soa_2, static_cast<int>(p_pair_indices_2[i]),
153  static_cast<int>(product2_index), engine);
154  // Set the weight of the new particles to p_pair_reaction_weight[i]
155  soa_products_data[1].m_rdata[PIdx::w][product2_index] = p_pair_reaction_weight[i];
156 
157  // Remove p_pair_reaction_weight[i] from the colliding particles' weights
158  amrex::Gpu::Atomic::AddNoRet(&w1[p_pair_indices_1[i]],
159  -p_pair_reaction_weight[i]);
160  amrex::Gpu::Atomic::AddNoRet(&w2[p_pair_indices_2[i]],
161  -p_pair_reaction_weight[i]);
162 
163  // Note: Particle::atomicSetID should also be provided as a standalone helper function in AMReX
164  // to replace the following lambda.
165  auto const atomicSetIdInvalid = [] AMREX_GPU_DEVICE (uint64_t & idcpu)
166  {
167 #if defined(AMREX_USE_OMP)
168 #pragma omp atomic write
170 #else
172  (unsigned long long *)&idcpu,
173  (unsigned long long)amrex::ParticleIdCpus::Invalid
174  );
175 #endif
176  };
177 
178  // If the colliding particle weight decreases to zero, remove particle by
179  // setting its id to invalid
180  if (w1[p_pair_indices_1[i]] <= std::numeric_limits<amrex::ParticleReal>::min())
181  {
182  atomicSetIdInvalid(idcpu1[p_pair_indices_1[i]]);
183  }
184  if (w2[p_pair_indices_2[i]] <= std::numeric_limits<amrex::ParticleReal>::min())
185  {
186  atomicSetIdInvalid(idcpu2[p_pair_indices_2[i]]);
187  }
188 
189  // Set the child particle properties appropriately
190  auto& ux1 = soa_products_data[0].m_rdata[PIdx::ux][product1_index];
191  auto& uy1 = soa_products_data[0].m_rdata[PIdx::uy][product1_index];
192  auto& uz1 = soa_products_data[0].m_rdata[PIdx::uz][product1_index];
193  auto& ux2 = soa_products_data[1].m_rdata[PIdx::ux][product2_index];
194  auto& uy2 = soa_products_data[1].m_rdata[PIdx::uy][product2_index];
195  auto& uz2 = soa_products_data[1].m_rdata[PIdx::uz][product2_index];
196 
197  // for simplicity (for now) we assume non-relativistic particles
198  // and simply calculate the center-of-momentum velocity from the
199  // rest masses
200  auto const uCOM_x = (m1 * ux1 + m2 * ux2) / (m1 + m2);
201  auto const uCOM_y = (m1 * uy1 + m2 * uy2) / (m1 + m2);
202  auto const uCOM_z = (m1 * uz1 + m2 * uz2) / (m1 + m2);
203 
204  // transform to COM frame
205  ux1 -= uCOM_x;
206  uy1 -= uCOM_y;
207  uz1 -= uCOM_z;
208  ux2 -= uCOM_x;
209  uy2 -= uCOM_y;
210  uz2 -= uCOM_z;
211 
212  if (mask[i] == int(ScatteringProcessType::ELASTIC)) {
213  // randomly rotate the velocity vector for the first particle
215  ux1, uy1, uz1, std::sqrt(ux1*ux1 + uy1*uy1 + uz1*uz1), engine
216  );
217  // set the second particles velocity so that the total momentum
218  // is zero
219  ux2 = -ux1 * m1 / m2;
220  uy2 = -uy1 * m1 / m2;
221  uz2 = -uz1 * m1 / m2;
222  } else if (mask[i] == int(ScatteringProcessType::BACK)) {
223  // reverse the velocity vectors of both particles
224  ux1 *= -1.0_prt;
225  uy1 *= -1.0_prt;
226  uz1 *= -1.0_prt;
227  ux2 *= -1.0_prt;
228  uy2 *= -1.0_prt;
229  uz2 *= -1.0_prt;
230  } else if (mask[i] == int(ScatteringProcessType::CHARGE_EXCHANGE)) {
231  if (std::abs(m1 - m2) < 1e-28) {
232  auto const temp_ux = ux1;
233  auto const temp_uy = uy1;
234  auto const temp_uz = uz1;
235  ux1 = ux2;
236  uy1 = uy2;
237  uz1 = uz2;
238  ux2 = temp_ux;
239  uy2 = temp_uy;
240  uz2 = temp_uz;
241  }
242  else {
243  amrex::Abort("Uneven mass charge-exchange not implemented yet.");
244  }
245  }
246  else {
247  amrex::Abort("Unknown scattering process.");
248  }
249  // transform back to labframe
250  ux1 += uCOM_x;
251  uy1 += uCOM_y;
252  uz1 += uCOM_z;
253  ux2 += uCOM_x;
254  uy2 += uCOM_y;
255  uz2 += uCOM_z;
256  }
257  });
258 
259  // Initialize the user runtime components
260  for (int i = 0; i < m_num_product_species; i++)
261  {
262  const int start_index = int(products_np[i]);
263  const int stop_index = int(products_np[i] + num_added_vec[i]);
265  0, 0,
266  pc_products[i]->getUserRealAttribs(), pc_products[i]->getUserIntAttribs(),
267  pc_products[i]->getParticleComps(), pc_products[i]->getParticleiComps(),
268  pc_products[i]->getUserRealAttribParser(),
269  pc_products[i]->getUserIntAttribParser(),
270 #ifdef WARPX_QED
271  false, // do not initialize QED quantities, since they were initialized
272  // when calling the SmartCopy functors
273  pc_products[i]->get_breit_wheeler_engine_ptr(),
274  pc_products[i]->get_quantum_sync_engine_ptr(),
275 #endif
276  pc_products[i]->getIonizationInitialLevel(),
277  start_index, stop_index);
278  }
279 
281  return num_added_vec;
282  }
283 
284 private:
285  // How many different type of species the collision produces
287  // Vectors of size m_num_product_species storing how many particles of a given species are
288  // produced by a collision event. These vectors are duplicated (one version for host and one
289  // for device) which is necessary with GPUs but redundant on CPU.
293 };
294 #endif // WARPX_SPLIT_AND_SCATTER_FUNC_H_
#define AMREX_RESTRICT
#define AMREX_INLINE
#define AMREX_GPU_DEVICE
Array4< int const > mask
CollisionType
Definition: BinaryCollisionUtils.H:17
Definition: MultiParticleContainer.H:66
This class defines an operator to create product particles from DSMC collisions and sets the particle...
Definition: SplitAndScatterFunc.H:23
AMREX_INLINE amrex::Vector< int > operator()(const index_type &n_total_pairs, const SoaData_type &, const SoaData_type &, const amrex::Vector< WarpXParticleContainer * > &pc_products, ParticleTileType **AMREX_RESTRICT tile_products, const amrex::ParticleReal m1, const amrex::ParticleReal m2, const amrex::Vector< amrex::ParticleReal > &, const index_type *AMREX_RESTRICT mask, const amrex::Vector< index_type > &products_np, const SmartCopy *AMREX_RESTRICT copy_species1, const SmartCopy *AMREX_RESTRICT copy_species2, const index_type *AMREX_RESTRICT p_pair_indices_1, const index_type *AMREX_RESTRICT p_pair_indices_2, const amrex::ParticleReal *AMREX_RESTRICT p_pair_reaction_weight) const
Function that performs the particle scattering and injection due to binary collisions.
Definition: SplitAndScatterFunc.H:54
typename ParticleBins::index_type index_type
Definition: SplitAndScatterFunc.H:29
typename ParticleTileType::ParticleTileDataType ParticleTileDataType
Definition: SplitAndScatterFunc.H:27
int m_num_product_species
Definition: SplitAndScatterFunc.H:286
SplitAndScatterFunc()=default
Default constructor of the SplitAndScatterFunc class.
typename WarpXParticleContainer::ParticleType ParticleType
Definition: SplitAndScatterFunc.H:25
CollisionType m_collision_type
Definition: SplitAndScatterFunc.H:292
typename WarpXParticleContainer::ParticleTileType::ParticleTileDataType SoaData_type
Definition: SplitAndScatterFunc.H:30
typename WarpXParticleContainer::ParticleTileType ParticleTileType
Definition: SplitAndScatterFunc.H:26
amrex::Gpu::HostVector< int > m_num_products_host
Definition: SplitAndScatterFunc.H:291
amrex::Gpu::DeviceVector< int > m_num_products_device
Definition: SplitAndScatterFunc.H:290
unsigned int index_type
T * data() noexcept
iterator begin() noexcept
T * dataPtr() noexcept
ParticleTile< ParticleType, NArrayReal, NArrayInt, Allocator > ParticleTileType
void DefaultInitializeRuntimeAttributes(PTile &ptile, const int n_external_attr_real, const int n_external_attr_int, const std::vector< std::string > &user_real_attribs, const std::vector< std::string > &user_int_attribs, const std::map< std::string, int > &particle_comps, const std::map< std::string, int > &particle_icomps, const std::vector< amrex::Parser * > &user_real_attrib_parser, const std::vector< amrex::Parser * > &user_int_attrib_parser, const bool do_qed_comps, BreitWheelerEngine *p_bw_engine, QuantumSynchrotronEngine *p_qs_engine, const int ionization_initial_level, int start, int stop)
Default initialize runtime attributes in a tile. This routine does not initialize the first n_externa...
Definition: DefaultInitialization.H:118
AMREX_GPU_HOST_DEVICE AMREX_INLINE void RandomizeVelocity(amrex::ParticleReal &ux, amrex::ParticleReal &uy, amrex::ParticleReal &uz, const amrex::ParticleReal vp, amrex::RandomEngine const &engine)
Function to perform scattering of a particle that results in a random velocity vector with given magn...
Definition: ParticleUtils.H:162
typename ParticleTileType::ParticleTileDataType ParticleTileDataType
Definition: ParticleUtils.cpp:33
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE T Exch(T *address, T val) noexcept
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void AddNoRet(T *sum, T value) noexcept
void synchronize() noexcept
void copyAsync(HostToDevice, InIter begin, InIter end, OutIter result) noexcept
static constexpr HostToDevice hostToDevice
void streamSynchronize() noexcept
constexpr std::uint64_t Invalid
static constexpr struct amrex::Scan::Type::Exclusive exclusive
static constexpr RetSum retSum
void Abort(const std::string &msg)
const int[]
AMREX_ATTRIBUTE_FLATTEN_FOR void ParallelForRNG(T n, L const &f) noexcept
i
Definition: check_interp_points_and_weights.py:174
@ uz
Definition: NamedComponentParticleContainer.H:34
@ w
weight
Definition: NamedComponentParticleContainer.H:33
@ uy
Definition: NamedComponentParticleContainer.H:34
@ ux
Definition: NamedComponentParticleContainer.H:34
This is a functor for performing a "smart copy" that works in both host and device code.
Definition: SmartCopy.H:34
ParticleTileData< StorageParticleType, NArrayReal, NArrayInt > ParticleTileDataType