impactx/src/particles/spacecharge/GatherAndPush.cpp at 98ccaee1964ea37d68ac66d2d14cb7879cb371b0 · BLAST-ImpactX/impactx · GitHub

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
/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
 *           Berkeley National Laboratory (subject to receipt of any required
 *           approvals from the U.S. Dept. of Energy). All rights reserved.
 *
 * This file is part of ImpactX.
 *
 * Authors: Marco Garten, Axel Huebl
 * License: BSD-3-Clause-LBNL
 */
#include "GatherAndPush.H"

#include "initialization/Algorithms.H"

#include <ablastr/particles/NodalFieldGather.H>

#include <AMReX_BLProfiler.H>
#include <AMReX_REAL.H>       // for Real
#include <AMReX_SPACE.H>      // for AMREX_D_DECL


namespace impactx::particles::spacecharge
{
    void GatherAndPush (
        ImpactXParticleContainer & pc,
        std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > const & space_charge_field,
        const amrex::Vector<amrex::Geometry>& geom,
        amrex::ParticleReal const slice_ds
    )
    {
        BL_PROFILE("impactx::spacecharge::GatherAndPush");

        using namespace amrex::literals;

        auto space_charge = get_space_charge_algo();

        amrex::ParticleReal const charge = pc.GetRefParticle().charge;

        // loop over refinement levels
        int const nLevel = pc.finestLevel();
        for (int lev = 0; lev <= nLevel; ++lev)
        {
            // get simulation geometry information
            auto const &gm = geom[lev];
            auto const dr = gm.CellSizeArray();
            amrex::GpuArray<amrex::Real, 3> const invdr{AMREX_D_DECL(1_rt/dr[0], 1_rt/dr[1], 1_rt/dr[2])};
            const auto prob_lo = gm.ProbLoArray();

            // loop over all particle boxes
            using ParIt = ImpactXParticleContainer::iterator;
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
            for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
                const int np = pti.numParticles();

                // get the device pointer-wrapper Array4 for 3D field access
                auto const scf_arr_x = space_charge_field.at(lev).at("x")[pti].array();
                auto const scf_arr_y = space_charge_field.at(lev).at("y")[pti].array();
                auto const scf_arr_z = space_charge_field.at(lev).at("z")[pti].array();

                // physical constants and reference quantities
                amrex::ParticleReal const c0_SI = 2.99792458e8;  // TODO move out
                amrex::ParticleReal const mc_SI = pc.GetRefParticle().mass * c0_SI;
                amrex::ParticleReal const pz_ref_SI = pc.GetRefParticle().beta_gamma() * mc_SI;
                amrex::ParticleReal const gamma = pc.GetRefParticle().gamma();
                amrex::ParticleReal const beta_gamma = pc.GetRefParticle().beta_gamma();
                amrex::ParticleReal const beta = beta_gamma / gamma;
                amrex::ParticleReal const inv_gamma2 = 1.0_prt / (gamma * gamma);

                amrex::ParticleReal const dt = slice_ds / pc.GetRefParticle().beta() / c0_SI;

                // preparing access to particle data: SoA of Reals
                auto& soa_real = pti.GetStructOfArrays().GetRealData();
                amrex::ParticleReal* const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr();
                amrex::ParticleReal* const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr();
                amrex::ParticleReal* const AMREX_RESTRICT part_z = soa_real[RealSoA::z].dataPtr(); // note: currently for a fixed t
                amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::px].dataPtr();
                amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr();
                amrex::ParticleReal* const AMREX_RESTRICT part_pz = soa_real[RealSoA::pz].dataPtr(); // note: currently for a fixed t

                // group together constants for the momentum push
                amrex::ParticleReal const push_consts = dt * charge * inv_gamma2 / pz_ref_SI;

                // gather to each particle and push momentum
                if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
                    // flatten 3rd dimension
                    auto prob_lo_2D = gm.ProbLoArray();
                    prob_lo_2D[2] = 0.0_rt;

                    // TODO: add in z-dependent scaling by current?

                    amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
                        // access SoA Real data
                        amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
                        amrex::ParticleReal & AMREX_RESTRICT y = part_y[i];
                        amrex::ParticleReal z = 0.0_prt;  // flatten 3rd dimension
                        amrex::ParticleReal & AMREX_RESTRICT px = part_px[i];
                        amrex::ParticleReal & AMREX_RESTRICT py = part_py[i];
                        amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i];

                        // force gather
                        amrex::GpuArray<amrex::Real, 3> const field_interp =
                            ablastr::particles::doGatherVectorFieldNodal<2>(
                                x, y, z,
                                scf_arr_x, scf_arr_y, scf_arr_z,
                                invdr,
                                prob_lo_2D
                            );

                        // push momentum
                        px += field_interp[0] * push_consts * dr[2] / (beta * c0_SI);
                        py += field_interp[1] * push_consts * dr[2] / (beta * c0_SI);
                        pz += 0.0_rt;
                        //pz += field_interp[2] * push_consts;  // TODO: non-zero in 2.5D, but we will add a toggle to turn it off there, too

                        // push position is done in the lattice elements
                    });
                }
                if (space_charge == SpaceChargeAlgo::True_3D) {
                    amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
                        // access SoA Real data
                        amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
                        amrex::ParticleReal & AMREX_RESTRICT y = part_y[i];
                        amrex::ParticleReal & AMREX_RESTRICT z = part_z[i];
                        amrex::ParticleReal & AMREX_RESTRICT px = part_px[i];
                        amrex::ParticleReal & AMREX_RESTRICT py = part_py[i];
                        amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i];

                        // force gather
                        amrex::GpuArray<amrex::Real, 3> const field_interp =
                            ablastr::particles::doGatherVectorFieldNodal(
                                x, y, z,
                                scf_arr_x, scf_arr_y, scf_arr_z,
                                invdr,
                                prob_lo
                            );

                        // push momentum
                        px += field_interp[0] * push_consts;
                        py += field_interp[1] * push_consts;
                        pz += field_interp[2] * push_consts;

                        // push position is done in the lattice elements
                    });
                }
            } // end loop over all particle boxes
        } // env mesh-refinement level loop
    }
}  // namespace impactx::particles::spacecharge