Skip to content

Commit b452f85

Browse files
Yin-YinjianZhaoYin-YinjianZhao
authored
Make BeamRelevant reduced diagnostic doable for RZ (#1183)
* Add RZ in BeamRelevant.cpp * Fix z * change to defined WARPX_DIM_3D
1 parent a22aa21 commit b452f85

File tree

1 file changed

+66
-20
lines changed

1 file changed

+66
-20
lines changed

‎Source/Diagnostics/ReducedDiags/BeamRelevant.cpp‎

Lines changed: 66 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -23,18 +23,12 @@ BeamRelevant::BeamRelevant (std::string rd_name)
2323
: ReducedDiags{rd_name}
2424
{
2525

26-
#if (defined WARPX_DIM_RZ)
27-
// RZ coordinate is not working
28-
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(false,
29-
"BeamRelevant reduced diagnostics does not work for RZ coordinate.");
30-
#endif
31-
3226
// read beam name
3327
ParmParse pp(rd_name);
3428
pp.get("species",m_beam_name);
3529

3630
// resize data array
37-
#if (AMREX_SPACEDIM == 3)
31+
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
3832
// 0, 1, 2: mean x,y,z
3933
// 3, 4, 5: mean px,py,pz
4034
// 6: gamma
@@ -44,7 +38,7 @@ BeamRelevant::BeamRelevant (std::string rd_name)
4438
// 14,15,16: emittance x,y,z
4539
// 17: charge
4640
m_data.resize(18,0.0);
47-
#elif (AMREX_SPACEDIM == 2)
41+
#elif (defined WARPX_DIM_XZ)
4842
// 0, 1: mean x,z
4943
// 2, 3, 4: mean px,py,pz
5044
// 5: gamma
@@ -65,7 +59,7 @@ BeamRelevant::BeamRelevant (std::string rd_name)
6559
ofs.open(m_path + m_rd_name + "." + m_extension,
6660
std::ofstream::out | std::ofstream::app);
6761
// write header row
68-
#if (AMREX_SPACEDIM == 3)
62+
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
6963
ofs << "#";
7064
ofs << "[1]step()"; ofs << m_sep;
7165
ofs << "[2]time(s)"; ofs << m_sep;
@@ -87,7 +81,7 @@ BeamRelevant::BeamRelevant (std::string rd_name)
8781
ofs << "[18]emittance_y(m)"; ofs << m_sep;
8882
ofs << "[19]emittance_z(m)"; ofs << m_sep;
8983
ofs << "[20]charge(C)"; ofs << std::endl;
90-
#elif (AMREX_SPACEDIM == 2)
84+
#elif (defined WARPX_DIM_XZ)
9185
ofs << "#";
9286
ofs << "[1]step()"; ofs << m_sep;
9387
ofs << "[2]time(s)"; ofs << m_sep;
@@ -135,9 +129,9 @@ void BeamRelevant::ComputeDiags (int step)
135129
Real constexpr inv_c2 = 1.0 / (PhysConst::c * PhysConst::c);
136130

137131
// If 2D-XZ, p.pos(1) is z, rather than p.pos(2).
138-
#if (AMREX_SPACEDIM == 3)
132+
#if (defined WARPX_DIM_3D)
139133
int const index_z = 2;
140-
#elif (AMREX_SPACEDIM == 2)
134+
#elif (defined WARPX_DIM_XZ || defined WARPX_DIM_RZ)
141135
int const index_z = 1;
142136
#endif
143137

@@ -173,16 +167,28 @@ void BeamRelevant::ComputeDiags (int step)
173167
return;
174168
}
175169

170+
#if (defined WARPX_DIM_RZ)
171+
// x mean
172+
Real x_mean = ReduceSum( myspc,
173+
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
174+
{ return p.pos(0)*std::cos(p.rdata(PIdx::theta)) * p.rdata(PIdx::w); });
175+
#else
176176
// x mean
177177
Real x_mean = ReduceSum( myspc,
178178
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
179179
{ return p.pos(0) * p.rdata(PIdx::w); });
180+
#endif
180181

181-
#if (AMREX_SPACEDIM == 3)
182+
#if (defined WARPX_DIM_3D)
182183
// y mean
183184
Real y_mean = ReduceSum( myspc,
184185
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
185186
{ return p.pos(1) * p.rdata(PIdx::w); });
187+
#elif (defined WARPX_DIM_RZ)
188+
// y mean
189+
Real y_mean = ReduceSum( myspc,
190+
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
191+
{ return p.pos(0)*std::sin(p.rdata(PIdx::theta)) * p.rdata(PIdx::w); });
186192
#endif
187193

188194
// z mean
@@ -218,7 +224,7 @@ void BeamRelevant::ComputeDiags (int step)
218224

219225
// reduced sum over mpi ranks
220226
ParallelDescriptor::ReduceRealSum(x_mean); x_mean /= w_sum;
221-
#if (AMREX_SPACEDIM == 3)
227+
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
222228
ParallelDescriptor::ReduceRealSum(y_mean); y_mean /= w_sum;
223229
#endif
224230
ParallelDescriptor::ReduceRealSum(z_mean); z_mean /= w_sum;
@@ -227,22 +233,42 @@ void BeamRelevant::ComputeDiags (int step)
227233
ParallelDescriptor::ReduceRealSum(uz_mean); uz_mean /= w_sum;
228234
ParallelDescriptor::ReduceRealSum(gm_mean); gm_mean /= w_sum;
229235

236+
#if (defined WARPX_DIM_RZ)
237+
// x mean square
238+
Real x_ms = ReduceSum( myspc,
239+
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
240+
{
241+
Real const x = p.pos(0)*std::cos(p.rdata(PIdx::theta));
242+
Real const a = (x-x_mean) * (x-x_mean);
243+
return a * p.rdata(PIdx::w);
244+
});
245+
#else
230246
// x mean square
231247
Real x_ms = ReduceSum( myspc,
232248
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
233249
{
234250
Real const a = (p.pos(0)-x_mean) * (p.pos(0)-x_mean);
235251
return a * p.rdata(PIdx::w);
236252
});
253+
#endif
237254

238-
#if (AMREX_SPACEDIM == 3)
255+
#if (defined WARPX_DIM_3D)
239256
// y mean square
240257
Real y_ms = ReduceSum( myspc,
241258
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
242259
{
243260
Real const a = (p.pos(1)-y_mean) * (p.pos(1)-y_mean);
244261
return a * p.rdata(PIdx::w);
245262
});
263+
#elif (defined WARPX_DIM_RZ)
264+
// y mean square
265+
Real y_ms = ReduceSum( myspc,
266+
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
267+
{
268+
Real const y = p.pos(0)*std::sin(p.rdata(PIdx::theta));
269+
Real const a = (y-y_mean) * (y-y_mean);
270+
return a * p.rdata(PIdx::w);
271+
});
246272
#endif
247273

248274
// z mean square
@@ -293,22 +319,42 @@ void BeamRelevant::ComputeDiags (int step)
293319
return a * p.rdata(PIdx::w);
294320
});
295321

322+
#if (defined WARPX_DIM_RZ)
323+
// x times ux
324+
Real xux = ReduceSum( myspc,
325+
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
326+
{
327+
Real const x = p.pos(0)*std::cos(p.rdata(PIdx::theta));
328+
Real const a = (x-x_mean) * (p.rdata(PIdx::ux)-ux_mean);
329+
return a * p.rdata(PIdx::w);
330+
});
331+
#else
296332
// x times ux
297333
Real xux = ReduceSum( myspc,
298334
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
299335
{
300336
Real const a = (p.pos(0)-x_mean) * (p.rdata(PIdx::ux)-ux_mean);
301337
return a * p.rdata(PIdx::w);
302338
});
339+
#endif
303340

304-
#if (AMREX_SPACEDIM == 3)
341+
#if (defined WARPX_DIM_3D)
305342
// y times uy
306343
Real yuy = ReduceSum( myspc,
307344
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
308345
{
309346
Real const a = (p.pos(1)-y_mean) * (p.rdata(PIdx::uy)-uy_mean);
310347
return a * p.rdata(PIdx::w);
311348
});
349+
#elif (defined WARPX_DIM_RZ)
350+
// y times uy
351+
Real yuy = ReduceSum( myspc,
352+
[=] AMREX_GPU_HOST_DEVICE (const PType& p) -> Real
353+
{
354+
Real const y = p.pos(0)*std::sin(p.rdata(PIdx::theta));
355+
Real const a = (y-y_mean) * (p.rdata(PIdx::uy)-uy_mean);
356+
return a * p.rdata(PIdx::w);
357+
});
312358
#endif
313359

314360
// z times uz
@@ -331,7 +377,7 @@ void BeamRelevant::ComputeDiags (int step)
331377
ParallelDescriptor::ReduceRealSum
332378
( x_ms, ParallelDescriptor::IOProcessorNumber());
333379
x_ms /= w_sum;
334-
#if (AMREX_SPACEDIM == 3)
380+
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
335381
ParallelDescriptor::ReduceRealSum
336382
( y_ms, ParallelDescriptor::IOProcessorNumber());
337383
y_ms /= w_sum;
@@ -354,7 +400,7 @@ void BeamRelevant::ComputeDiags (int step)
354400
ParallelDescriptor::ReduceRealSum
355401
( xux, ParallelDescriptor::IOProcessorNumber());
356402
xux /= w_sum;
357-
#if (AMREX_SPACEDIM == 3)
403+
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
358404
ParallelDescriptor::ReduceRealSum
359405
( yuy, ParallelDescriptor::IOProcessorNumber());
360406
yuy /= w_sum;
@@ -366,7 +412,7 @@ void BeamRelevant::ComputeDiags (int step)
366412
( charge, ParallelDescriptor::IOProcessorNumber());
367413

368414
// save data
369-
#if (AMREX_SPACEDIM == 3)
415+
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
370416
m_data[0] = x_mean;
371417
m_data[1] = y_mean;
372418
m_data[2] = z_mean;
@@ -385,7 +431,7 @@ void BeamRelevant::ComputeDiags (int step)
385431
m_data[15] = std::sqrt(y_ms*uy_ms-yuy*yuy) / PhysConst::c;
386432
m_data[16] = std::sqrt(z_ms*uz_ms-zuz*zuz) / PhysConst::c;
387433
m_data[17] = charge;
388-
#elif (AMREX_SPACEDIM == 2)
434+
#elif (defined WARPX_DIM_XZ)
389435
m_data[0] = x_mean;
390436
m_data[1] = z_mean;
391437
m_data[2] = ux_mean * m;

0 commit comments

Comments
 (0)