1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
|
#ifndef ELEMENTARYPROCESS_H_
#define ELEMENTARYPROCESS_H_
#include "WarpXParticleContainer.H"
#include "CopyParticle.H"
#include "TransformParticle.H"
/**
* \brief Base class for particle creation processes
*
* Particles in a source particle container are copied to a product particle
* container if they are flagged. Both source and product particles can be
* modified.
*
* initialize_copy initializes a general copy functor
* createParticles formats the data to prepare for the copy, then
* calls copyAndTransformParticles
* copyAndTransformParticles loops over particles, performs the copy and
* transform source and product particles if needed.
*
* The class is templated on the process type. This gives flexibility
* on source and product particle transformation.
*/
template<elementaryProcessType ProcessT>
class elementaryProcess
{
public:
/**
* \brief initialize and return functor for particle copy
*
* \param cpuid id of MPI rank
* \param do_boosted_product whether to copy old attribs
* \param runtime_uold_source momentum of source particles
* \param attribs_source attribs of source particles
* \param attribs_product attribs of product particles
* \param runtime_attribs_product runtime attribs for product, to store old attribs
*/
copyParticle initialize_copy(
const int cpuid, const int do_boosted_product,
const amrex::GpuArray<amrex::ParticleReal*,3> runtime_uold_source,
const amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_source,
const amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_product,
const amrex::GpuArray<amrex::ParticleReal*,6> runtime_attribs_product) const noexcept
{
return copyParticle (
cpuid, do_boosted_product, runtime_uold_source,
attribs_source, attribs_product, runtime_attribs_product);
};
/**
* \brief create particles in product particle containers
*
* For particle i in mfi, if is_flagged[i]=1, copy particle
* particle i from container pc_source into pc_product
*
* \param lev MR level
* \param mfi MFIter where source particles are located
* \param pc_source source particle container
* \param v_pc_product vector of product particle containers
* \param is_flagged particles for which is_flagged is 1 are copied
* \param v_do_boosted_product vector: whether to copy old attribs for
* each product container
*/
void createParticles (
int lev, const amrex::MFIter& mfi,
std::unique_ptr< WarpXParticleContainer>& pc_source,
amrex::Vector<WarpXParticleContainer*> v_pc_product,
amrex::Gpu::ManagedDeviceVector<int>& is_flagged,
const amrex::Vector<int> v_do_boosted_product)
{
BL_PROFILE("createIonizedParticles");
const int grid_id = mfi.index();
const int tile_id = mfi.LocalTileIndex();
// Get source particle data
auto& ptile_source = pc_source->GetParticles(lev)[std::make_pair(grid_id,tile_id)];
const int np_source = ptile_source.GetArrayOfStructs().size();
if (np_source == 0) return;
// --- source AoS particle data
WarpXParticleContainer::ParticleType* particles_source = ptile_source.GetArrayOfStructs()().data();
// --- source SoA particle data
auto& soa_source = ptile_source.GetStructOfArrays();
amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_source;
for (int ia = 0; ia < PIdx::nattribs; ++ia) {
attribs_source[ia] = soa_source.GetRealData(ia).data();
}
// --- source runtime attribs
amrex::GpuArray<amrex::ParticleReal*,3> runtime_uold_source;
// Prepare arrays for boosted frame diagnostics.
runtime_uold_source[0] = soa_source.GetRealData(PIdx::ux).data();
runtime_uold_source[1] = soa_source.GetRealData(PIdx::uy).data();
runtime_uold_source[2] = soa_source.GetRealData(PIdx::uz).data();
// --- source runtime integer attribs
amrex::GpuArray<int*,1> runtime_iattribs_source;
std::map<std::string, int> icomps_source = pc_source->getParticleiComps();
runtime_iattribs_source[0] = soa_source.GetIntData(icomps_source["ionization_level"]).data();
int nproducts = v_pc_product.size();
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
v_do_boosted_product.size() == nproducts,
"v_pc_product and v_do_boosted_product must have the same size.");
// Indices of product particle for each ionized source particle.
// i_product[i]-1 is the location in product tile of product particle
// from source particle i.
amrex::Gpu::ManagedDeviceVector<int> i_product;
i_product.resize(np_source);
// 0<i<np_source
// 0<i_product<np_ionized
// Strictly speaking, i_product should be an exclusive_scan of
// is_flagged. However, for indices where is_flagged is 1, the
// inclusive scan gives the same result with an offset of 1.
// The advantage of inclusive_scan is that the sum of is_flagged
// is in the last element, so no other reduction is required to get
// number of particles.
// Gpu::inclusive_scan runs on the current GPU stream, and synchronizes
// with the CPU, so that the next line (executed by the CPU) has the
// updated values of i_product
amrex::Gpu::inclusive_scan(is_flagged.begin(), is_flagged.end(), i_product.begin());
int np_ionized = i_product[np_source-1];
if (np_ionized == 0) return;
// amrex::Vector<copyParticle> v_copy_functor(nproducts);
amrex::Vector<copyParticle> v_copy_functor;
amrex::Vector<int> v_pid_product(nproducts);
amrex::Gpu::ManagedDeviceVector<WarpXParticleContainer::ParticleType*> v_particles_product(nproducts);
for (int iproduct=0; iproduct<nproducts; iproduct++){
WarpXParticleContainer*& pc_product = v_pc_product[iproduct];
// Get product particle data
auto& ptile_product = pc_product->GetParticles(lev)[std::make_pair(grid_id,tile_id)];
// old and new (i.e., including ionized particles) number of particles
// for product species
const int np_product_old = ptile_product.GetArrayOfStructs().size();
const int np_product_new = np_product_old + np_ionized;
// Allocate extra space in product species for ionized particles.
ptile_product.resize(np_product_new);
// --- product AoS particle data
// First element is the first newly-created product particle
v_particles_product[iproduct] = ptile_product.GetArrayOfStructs()().data() + np_product_old;
// --- product SoA particle data
auto& soa_product = ptile_product.GetStructOfArrays();
amrex::GpuArray<amrex::ParticleReal*,PIdx::nattribs> attribs_product;
for (int ia = 0; ia < PIdx::nattribs; ++ia) {
// First element is the first newly-created product particle
attribs_product[ia] = soa_product.GetRealData(ia).data() + np_product_old;
}
// --- product runtime attribs
amrex::GpuArray<amrex::ParticleReal*,6> runtime_attribs_product;
const int do_boost = v_do_boosted_product[iproduct];
if (do_boost) {
std::map<std::string, int> comps_product = pc_product->getParticleComps();
runtime_attribs_product[0] = soa_product.GetRealData(comps_product[ "xold"]).data() + np_product_old;
runtime_attribs_product[1] = soa_product.GetRealData(comps_product[ "yold"]).data() + np_product_old;
runtime_attribs_product[2] = soa_product.GetRealData(comps_product[ "zold"]).data() + np_product_old;
runtime_attribs_product[3] = soa_product.GetRealData(comps_product["uxold"]).data() + np_product_old;
runtime_attribs_product[4] = soa_product.GetRealData(comps_product["uyold"]).data() + np_product_old;
runtime_attribs_product[5] = soa_product.GetRealData(comps_product["uzold"]).data() + np_product_old;
}
// --- product runtime integer attribs
int pid_product;
#pragma omp critical (doFieldIonization_nextid)
{
// ID of first newly-created product particle
pid_product = pc_product->NextID();
// Update NextID to include particles created in this function
pc_product->setNextID(pid_product+np_ionized);
}
const int cpuid = amrex::ParallelDescriptor::MyProc();
// Create instance of copy functor
v_copy_functor.push_back (initialize_copy(
cpuid, v_do_boosted_product[iproduct],
runtime_uold_source,
attribs_source,
attribs_product,
runtime_attribs_product) );
v_pid_product[iproduct] = pid_product;
}
// Loop over source particles and create product particles
copyAndTransformParticles(is_flagged, i_product, np_source, v_pid_product,
v_particles_product, particles_source, v_copy_functor,
runtime_iattribs_source);
}
/**
* \brief Loop over source particles and create product particles
*
* \param is_flagged particles for which is_flagged is 1 are copied
* \param i_product relative indices of product particle. This is the same
* for all product particle containers.
* \param np_source number of particles in source container
* \param v_pid_product for each product particle container: ID of the
* first product particle. Others IDs are incremented from this one
* \param v_particles_product for each product particle container: pointer
* to the AoS particle data
* \param particles_source pointter to the AoS source particle data
* \param v_copy_functor for each product particle container: copy functor
* \param runtime_iattribs_source pointer to source runtime integer attribs
*/
void copyAndTransformParticles(
amrex::Gpu::ManagedDeviceVector<int>& is_flagged,
amrex::Gpu::ManagedDeviceVector<int>& i_product,
int np_source,
const amrex::Vector<int> v_pid_product,
const amrex::Gpu::ManagedDeviceVector<WarpXParticleContainer::ParticleType*> v_particles_product,
WarpXParticleContainer::ParticleType* particles_source,
const amrex::Vector<copyParticle>& v_copy_functor,
amrex::GpuArray<int*,1> runtime_iattribs_source)
{
int const * const AMREX_RESTRICT p_is_flagged = is_flagged.dataPtr();
int const * const AMREX_RESTRICT p_i_product = i_product.dataPtr();
WarpXParticleContainer::ParticleType* const * const AMREX_RESTRICT p_particles_product = v_particles_product.dataPtr();
// Loop over all source particles. If is_flagged, copy particle data
// to corresponding product particle.
amrex::For(
np_source, [=] AMREX_GPU_DEVICE (int is) noexcept
{
if(p_is_flagged[is]){
int nproducts = v_particles_product.size();
// offset of 1 due to inclusive scan
int ip = p_i_product[is]-1;
WarpXParticleContainer::ParticleType& p_source = particles_source[is];
for (int iproduct=0; iproduct<nproducts; iproduct++){
// is: index of ionized particle in source species
// ip: index of corresponding new particle in product species
WarpXParticleContainer::ParticleType& p_product = v_particles_product[iproduct][ip];
// copyParticle copy_functor (v_copy_functor[iproduct]);
copyParticle copy_functor = v_copy_functor[iproduct];
int pid_product = v_pid_product[iproduct];
copy_functor(is, ip, pid_product, p_source, p_product);
}
// Transform source particle
transformSourceParticle<ProcessT>(is, p_source, runtime_iattribs_source);
// Transform product particles. The loop over product
// species is done inside the function to allow for
// more flexibility.
transformProductParticle<ProcessT>(ip, p_particles_product);
}
}
);
}
};
// Derived class for ionization process
class IonizationProcess: public elementaryProcess<elementaryProcessType::Ionization>
{};
#endif // ELEMENTARYPROCESS_H_
|
