aboutsummaryrefslogtreecommitdiff
path: root/Source/Diagnostics/ReducedDiags/BeamRelevant.cpp
blob: a99d476fa01f181aee0b20a68f3d36c1ee1de01b (plain) (blame)
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
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
/* Copyright 2019-2020 Yinjian Zhao
 *
 * This file is part of WarpX.
 *
 * License: BSD-3-Clause-LBNL
 */
#include "BeamRelevant.H"

#include "Diagnostics/ReducedDiags/ReducedDiags.H"
#include "Particles/MultiParticleContainer.H"
#include "Particles/WarpXParticleContainer.H"
#include "Utils/WarpXConst.H"
#include "WarpX.H"

#include <AMReX_GpuQualifiers.H>
#include <AMReX_PODVector.H>
#include <AMReX_ParallelDescriptor.H>
#include <AMReX_ParmParse.H>
#include <AMReX_ParticleReduce.H>
#include <AMReX_Particles.H>
#include <AMReX_REAL.H>
#include <AMReX_Tuple.H>

#include <algorithm>
#include <cmath>
#include <fstream>
#include <limits>
#include <map>
#include <vector>

using namespace amrex;

// constructor
BeamRelevant::BeamRelevant (std::string rd_name)
: ReducedDiags{rd_name}
{
    // read beam name
    const ParmParse pp_rd_name(rd_name);
    pp_rd_name.get("species",m_beam_name);

    // resize data array
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
    //  0, 1, 2: mean x,y,z
    //  3, 4, 5: mean px,py,pz
    //        6: gamma
    //  7, 8, 9: rms x,y,z
    // 10,11,12: rms px,py,pz
    //       13: rms gamma
    // 14,15,16: emittance x,y,z
    //    17,18: Twiss-alpha x,y
    //    19,20: beta-function x,y
    //       21: charge
    m_data.resize(22, 0.0_rt);
#elif (defined WARPX_DIM_XZ)
    //     0, 1: mean x,z
    //  2, 3, 4: mean px,py,pz
    //        5: gamma
    //     6, 7: rms x,z
    //  8, 9,10: rms px,py,pz
    //       11: rms gamma
    //    12,13: emittance x,z
    //       14: Twiss-alpha x
    //       15: beta-function x
    //       16: charge
    m_data.resize(17, 0.0_rt);
#elif (defined WARPX_DIM_1D_Z)
    //       0 : mean z
    //   1,2,3 : mean px,py,pz
    //       4 : gamma
    //       5 : rms z
    //   6,7,8 : rms px,py,pz
    //       9 : rms gamma
    //      10 : emittance z
    //      11 : charge
    m_data.resize(12, 0.0_rt);
#endif

    if (ParallelDescriptor::IOProcessor())
    {
        if ( m_IsNotRestart )
        {
            // open file
            std::ofstream ofs{m_path + m_rd_name + "." + m_extension, std::ofstream::out};
            // write header row
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
            int c = 0;
            ofs << "#";
            ofs << "[" << c++ << "]step()";           ofs << m_sep;
            ofs << "[" << c++ << "]time(s)";          ofs << m_sep;
            ofs << "[" << c++ << "]x_mean(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]y_mean(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]z_mean(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]px_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]py_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]pz_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]gamma_mean()";     ofs << m_sep;
            ofs << "[" << c++ << "]x_rms(m)";         ofs << m_sep;
            ofs << "[" << c++ << "]y_rms(m)";         ofs << m_sep;
            ofs << "[" << c++ << "]z_rms(m)";         ofs << m_sep;
            ofs << "[" << c++ << "]px_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]py_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]pz_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]gamma_rms()";      ofs << m_sep;
            ofs << "[" << c++ << "]emittance_x(m)";   ofs << m_sep;
            ofs << "[" << c++ << "]emittance_y(m)";   ofs << m_sep;
            ofs << "[" << c++ << "]emittance_z(m)";   ofs << m_sep;
            ofs << "[" << c++ << "]alpha_x()";        ofs << m_sep;
            ofs << "[" << c++ << "]alpha_y()";        ofs << m_sep;
            ofs << "[" << c++ << "]beta_x(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]beta_y(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]charge(C)";        ofs << std::endl;
#elif (defined WARPX_DIM_XZ)
            int c = 0;
            ofs << "#";
            ofs << "[" << c++ << "]step()";           ofs << m_sep;
            ofs << "[" << c++ << "]time(s)";          ofs << m_sep;
            ofs << "[" << c++ << "]x_mean(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]z_mean(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]px_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]py_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]pz_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]gamma_mean()";     ofs << m_sep;
            ofs << "[" << c++ << "]x_rms(m)";         ofs << m_sep;
            ofs << "[" << c++ << "]z_rms(m)";         ofs << m_sep;
            ofs << "[" << c++ << "]px_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]py_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]pz_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]gamma_rms()";      ofs << m_sep;
            ofs << "[" << c++ << "]emittance_x(m)";   ofs << m_sep;
            ofs << "[" << c++ << "]emittance_z(m)";   ofs << m_sep;
            ofs << "[" << c++ << "]alpha_x()";        ofs << m_sep;
            ofs << "[" << c++ << "]beta_x(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]charge(C)";        ofs << std::endl;
#elif (defined WARPX_DIM_1D_Z)
            int c = 0;
            ofs << "#";
            ofs << "[" << c++ << "]step()";           ofs << m_sep;
            ofs << "[" << c++ << "]time(s)";          ofs << m_sep;
            ofs << "[" << c++ << "]z_mean(m)";        ofs << m_sep;
            ofs << "[" << c++ << "]px_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]py_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]pz_mean(kg*m/s)";  ofs << m_sep;
            ofs << "[" << c++ << "]gamma_mean()";     ofs << m_sep;
            ofs << "[" << c++ << "]z_rms(m)";         ofs << m_sep;
            ofs << "[" << c++ << "]px_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]py_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]pz_rms(kg*m/s)";   ofs << m_sep;
            ofs << "[" << c++ << "]gamma_rms()";      ofs << m_sep;
            ofs << "[" << c++ << "]emittance_z(m)";   ofs << m_sep;
            ofs << "[" << c++ << "]charge(C)";        ofs << std::endl;
#endif
            // close file
            ofs.close();
        }
    }
}
// end constructor

// function that compute beam relevant quantities
void BeamRelevant::ComputeDiags (int step)
{
    // Judge if the diags should be done
    if (!m_intervals.contains(step+1)) { return; }

    // get MultiParticleContainer class object
    const auto & mypc = WarpX::GetInstance().GetPartContainer();

    // get number of species
    int const nSpecies = mypc.nSpecies();

    // get species names (std::vector<std::string>)
    auto const species_names = mypc.GetSpeciesNames();

    // inverse of speed of light squared
    Real constexpr inv_c2 = 1.0_rt / (PhysConst::c * PhysConst::c);

    // If 2D-XZ, p.pos(1) is z, rather than p.pos(2).
#if (defined WARPX_DIM_3D)
    int const index_z = 2;
#elif (defined WARPX_DIM_XZ || defined WARPX_DIM_RZ)
    int const index_z = 1;
#elif (defined WARPX_DIM_1D_Z)
    int const index_z = 0;
#endif

    // loop over species
    for (int i_s = 0; i_s < nSpecies; ++i_s)
    {
        // only beam species does
        if (species_names[i_s] != m_beam_name) { continue; }

        // get WarpXParticleContainer class object
        auto const & myspc = mypc.GetParticleContainer(i_s);

        // get mass and charge
        ParticleReal const m = myspc.getMass();
        ParticleReal const q = myspc.getCharge();

        using PType = typename WarpXParticleContainer::SuperParticleType;

        amrex::ReduceOps<ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum,
        ReduceOpSum,ReduceOpSum,ReduceOpSum> reduce_ops;
        auto r = amrex::ParticleReduce<amrex::ReduceData<ParticleReal,ParticleReal,
        ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal>>(
            myspc,
            [=] AMREX_GPU_DEVICE(const PType& p) noexcept -> amrex::GpuTuple
            <ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal,
            ParticleReal,ParticleReal,ParticleReal>
            {
                const ParticleReal p_ux = p.rdata(PIdx::ux);
                const ParticleReal p_uy = p.rdata(PIdx::uy);
                const ParticleReal p_uz = p.rdata(PIdx::uz);
                const ParticleReal p_us = p_ux*p_ux + p_uy*p_uy + p_uz*p_uz;
                const ParticleReal p_pos0 = p.pos(0);
                const ParticleReal p_w = p.rdata(PIdx::w);

#if (defined WARPX_DIM_RZ)
                const ParticleReal p_theta = p.rdata(PIdx::theta);
                const ParticleReal p_x_mean = p_pos0*std::cos(p_theta)*p_w;
                const ParticleReal p_y_mean = p_pos0*std::sin(p_theta)*p_w;
#else
                const ParticleReal p_pos1 = p.pos(1);
                const ParticleReal p_x_mean = p_pos0*p_w;
                const ParticleReal p_y_mean = p_pos1*p_w;
#endif
                const ParticleReal p_z_mean = p.pos(index_z)*p_w;

                const ParticleReal p_ux_mean = p_ux*p_w;
                const ParticleReal p_uy_mean = p_uy*p_w;
                const ParticleReal p_uz_mean = p_uz*p_w;
                const ParticleReal p_gm_mean = std::sqrt(1.0_rt+p_us*inv_c2)*p_w;

                return {p_w,
                        p_x_mean, p_y_mean, p_z_mean,
                        p_ux_mean, p_uy_mean, p_uz_mean,
                        p_gm_mean};
            },
            reduce_ops);

        std::vector<ParticleReal> values_per_rank_1st = {
            amrex::get<0>(r), // w
            amrex::get<1>(r), // x_mean
            amrex::get<2>(r), // y_mean
            amrex::get<3>(r), // z_mean
            amrex::get<4>(r), // ux_mean
            amrex::get<5>(r), // uy_mean
            amrex::get<6>(r), // uz_mean
            amrex::get<7>(r), // gm_mean
        };

        // reduced sum over mpi ranks (allreduce)
        amrex::ParallelAllReduce::Sum
        ( values_per_rank_1st.data(), values_per_rank_1st.size(), ParallelDescriptor::Communicator());

        const ParticleReal w_sum   = values_per_rank_1st.at(0);
        const ParticleReal x_mean  = values_per_rank_1st.at(1) /= w_sum;
        const ParticleReal y_mean  = values_per_rank_1st.at(2) /= w_sum;
        const ParticleReal z_mean  = values_per_rank_1st.at(3) /= w_sum;
        const ParticleReal ux_mean = values_per_rank_1st.at(4) /= w_sum;
        const ParticleReal uy_mean = values_per_rank_1st.at(5) /= w_sum;
        const ParticleReal uz_mean = values_per_rank_1st.at(6) /= w_sum;
        const ParticleReal gm_mean = values_per_rank_1st.at(7) /= w_sum;

        if (w_sum < std::numeric_limits<Real>::min() )
        {
            for (auto& item: m_data) item = 0.0_rt;

            return;
        }

        amrex::ReduceOps<ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum,
        ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum,ReduceOpSum> reduce_ops2;

        auto r2 = amrex::ParticleReduce<amrex::ReduceData<ParticleReal,ParticleReal,
        ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal,
        ParticleReal,ParticleReal,ParticleReal>>(
            myspc,
            [=] AMREX_GPU_DEVICE(const PType& p) noexcept -> amrex::GpuTuple
            <ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal,
            ParticleReal,ParticleReal,ParticleReal,ParticleReal,ParticleReal>
            {
                const ParticleReal p_ux = p.rdata(PIdx::ux);
                const ParticleReal p_uy = p.rdata(PIdx::uy);
                const ParticleReal p_uz = p.rdata(PIdx::uz);
                const ParticleReal p_us = p_ux*p_ux + p_uy*p_uy + p_uz*p_uz;
                const ParticleReal p_gm = std::sqrt(1.0_rt+p_us*inv_c2);
                const ParticleReal p_w = p.rdata(PIdx::w);

#if (defined WARPX_DIM_1D_Z)
                const ParticleReal p_x = 0.0;
                const ParticleReal p_y = 0.0;
#elif (defined WARPX_DIM_RZ)
                const ParticleReal p_pos0 = p.pos(0);
                const ParticleReal p_theta = p.rdata(PIdx::theta);
                const ParticleReal p_x = p_pos0*std::cos(p_theta);
                const ParticleReal p_y = p_pos0*std::sin(p_theta);
#elif (defined WARPX_DIM_XZ)
                const ParticleReal p_pos0 = p.pos(0);
                const ParticleReal p_x = p_pos0;
                const ParticleReal p_y = 0.0;
#else
                const ParticleReal p_pos0 = p.pos(0);
                const ParticleReal p_pos1 = p.pos(1);
                const ParticleReal p_x = p_pos0;
                const ParticleReal p_y = p_pos1;
#endif
                const ParticleReal p_z = p.pos(index_z);

                const ParticleReal p_x_ms = (p_x-x_mean)*(p_x-x_mean)*p_w;
                const ParticleReal p_y_ms = (p_y-y_mean)*(p_y-y_mean)*p_w;
                const ParticleReal p_z_ms = (p_z-z_mean)*(p_z-z_mean)*p_w;

                const ParticleReal p_ux_ms = (p_ux-ux_mean)*(p_ux-ux_mean)*p_w;
                const ParticleReal p_uy_ms = (p_uy-uy_mean)*(p_uy-uy_mean)*p_w;
                const ParticleReal p_uz_ms = (p_uz-uz_mean)*(p_uz-uz_mean)*p_w;
                const ParticleReal p_gm_ms = (p_gm-gm_mean)*(p_gm-gm_mean)*p_w;

                const ParticleReal p_xux = (p_x-x_mean)*(p_ux-ux_mean)*p_w;
                const ParticleReal p_yuy = (p_y-y_mean)*(p_uy-uy_mean)*p_w;
                const ParticleReal p_zuz = (p_z-z_mean)*(p_uz-uz_mean)*p_w;

                const ParticleReal p_charge = q*p_w;

                return {p_x_ms, p_y_ms, p_z_ms,
                        p_ux_ms, p_uy_ms, p_uz_ms,
                        p_gm_ms,
                        p_xux, p_yuy, p_zuz,
                        p_charge};
            },
            reduce_ops2);

        std::vector<ParticleReal> values_per_rank_2nd = {
            amrex::get<0>(r2), // x_ms
            amrex::get<1>(r2), // y_ms
            amrex::get<2>(r2), // z_ms
            amrex::get<3>(r2), // ux_ms
            amrex::get<4>(r2), // uy_ms
            amrex::get<5>(r2), // uz_ms
            amrex::get<6>(r2), // gm_ms
            amrex::get<7>(r2), // xux
            amrex::get<8>(r2), // yuy
            amrex::get<9>(r2), // zuz
            amrex::get<10>(r2) // charge
        };

        // reduced sum over mpi ranks (reduce to IO rank)
        ParallelDescriptor::ReduceRealSum
        ( values_per_rank_2nd.data(), values_per_rank_2nd.size(), ParallelDescriptor::IOProcessorNumber());

        const ParticleReal x_ms   = values_per_rank_2nd.at(0) /= w_sum;
        const ParticleReal y_ms   = values_per_rank_2nd.at(1) /= w_sum;
        const ParticleReal z_ms   = values_per_rank_2nd.at(2) /= w_sum;
        const ParticleReal ux_ms  = values_per_rank_2nd.at(3) /= w_sum;
        const ParticleReal uy_ms  = values_per_rank_2nd.at(4) /= w_sum;
        const ParticleReal uz_ms  = values_per_rank_2nd.at(5) /= w_sum;
        const ParticleReal gm_ms  = values_per_rank_2nd.at(6) /= w_sum;
        const ParticleReal xux    = values_per_rank_2nd.at(7) /= w_sum;
        const ParticleReal yuy    = values_per_rank_2nd.at(8) /= w_sum;
        const ParticleReal zuz    = values_per_rank_2nd.at(9) /= w_sum;
        const ParticleReal charge = values_per_rank_2nd.at(10);

        // save data
#if (defined WARPX_DIM_3D || defined WARPX_DIM_RZ)
        m_data[0]  = x_mean;
        m_data[1]  = y_mean;
        m_data[2]  = z_mean;
        m_data[3]  = ux_mean * m;
        m_data[4]  = uy_mean * m;
        m_data[5]  = uz_mean * m;
        m_data[6]  = gm_mean;
        m_data[7]  = std::sqrt(x_ms);
        m_data[8]  = std::sqrt(y_ms);
        m_data[9]  = std::sqrt(z_ms);
        m_data[10] = std::sqrt(ux_ms) * m;
        m_data[11] = std::sqrt(uy_ms) * m;
        m_data[12] = std::sqrt(uz_ms) * m;
        m_data[13] = std::sqrt(gm_ms);
        m_data[14] = std::sqrt(x_ms*ux_ms-xux*xux) / PhysConst::c;
        m_data[15] = std::sqrt(y_ms*uy_ms-yuy*yuy) / PhysConst::c;
        m_data[16] = std::sqrt(z_ms*uz_ms-zuz*zuz) / PhysConst::c;
        m_data[17] = - (PhysConst::c * xux) / std::sqrt(x_ms*ux_ms-xux*xux);
        m_data[18] = - (PhysConst::c * yuy) / std::sqrt(y_ms*uy_ms-yuy*yuy);
        m_data[19] = (PhysConst::c * x_ms) / std::sqrt(x_ms*ux_ms-xux*xux);
        m_data[20] = (PhysConst::c * y_ms) / std::sqrt(y_ms*uy_ms-yuy*yuy);
        m_data[21] = charge;
#elif (defined WARPX_DIM_XZ)
        m_data[0]  = x_mean;
        m_data[1]  = z_mean;
        m_data[2]  = ux_mean * m;
        m_data[3]  = uy_mean * m;
        m_data[4]  = uz_mean * m;
        m_data[5]  = gm_mean;
        m_data[6]  = std::sqrt(x_ms);
        m_data[7]  = std::sqrt(z_ms);
        m_data[8]  = std::sqrt(ux_ms) * m;
        m_data[9]  = std::sqrt(uy_ms) * m;
        m_data[10] = std::sqrt(uz_ms) * m;
        m_data[11] = std::sqrt(gm_ms);
        m_data[12] = std::sqrt(x_ms*ux_ms-xux*xux) / PhysConst::c;
        m_data[13] = std::sqrt(z_ms*uz_ms-zuz*zuz) / PhysConst::c;
        m_data[14] = - (PhysConst::c * xux) / std::sqrt(x_ms*ux_ms-xux*xux);
        m_data[15] = (PhysConst::c * x_ms) / std::sqrt(x_ms*ux_ms-xux*xux);
        m_data[16] = charge;
        amrex::ignore_unused(y_mean, y_ms, yuy);
#elif (defined WARPX_DIM_1D_Z)
        m_data[0]  = z_mean;
        m_data[1]  = ux_mean * m;
        m_data[2]  = uy_mean * m;
        m_data[3]  = uz_mean * m;
        m_data[4]  = gm_mean;
        m_data[5]  = std::sqrt(z_ms);
        m_data[6]  = std::sqrt(ux_ms) * m;
        m_data[7]  = std::sqrt(uy_ms) * m;
        m_data[8]  = std::sqrt(uz_ms) * m;
        m_data[9]  = std::sqrt(gm_ms);
        m_data[10] = std::sqrt(z_ms*uz_ms-zuz*zuz) / PhysConst::c;
        m_data[11] = charge;
        amrex::ignore_unused(x_mean, x_ms, xux, y_mean, y_ms, yuy);
#endif
    }
    // end loop over species
}
// end void BeamRelevant::ComputeDiags