Implement: Nodal Gather and Push Momentum

src/ImpactX.H

@@ -124,7 +124,7 @@ namespace impactx
         std::unordered_map<int, amrex::MultiFab> m_rho;
         /** scalar potential per level */
         std::unordered_map<int, amrex::MultiFab> m_phi;
-        /** space charge force (vector) per level */
+        /** space charge field (rename me!!) (vector) per level */
         std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > m_space_charge_force;
 
         /** these are elements defining the accelerator lattice */

src/ImpactX.cpp

@@ -12,8 +12,9 @@
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
 #include "particles/diagnostics/DiagnosticOutput.H"
-#include "particles/spacecharge/PoissonSolve.H"
 #include "particles/spacecharge/ForceFromSelfFields.H"
+#include "particles/spacecharge/GatherAndPush.H"
+#include "particles/spacecharge/PoissonSolve.H"
 #include "particles/transformation/CoordinateTransformation.H"
 
 #include <AMReX.H>
@@ -115,7 +116,11 @@ namespace impactx
         {
             // number of slices used for the application of space charge
             int nslice = 1;
-            std::visit([&nslice](auto&& element){ nslice = element.nslice(); }, element_variant);
+            amrex::ParticleReal slice_ds; // in meters
+            std::visit([&nslice, &slice_ds](auto&& element){
+                nslice = element.nslice();
+                slice_ds = element.ds() / nslice;
+            }, element_variant);
 
             // sub-steps for space charge within the element
             for (int slice_step = 0; slice_step < nslice; ++slice_step)
@@ -156,7 +161,8 @@ namespace impactx
 
                     // gather and space-charge push in x,y,z , assuming the space-charge
                     // field is the same before/after transformation
-                    //   TODO
+                    // TODO: This is currently using linear order.
+                    spacecharge::GatherAndPush(*m_particle_container, m_space_charge_force, this->geom, slice_ds);
 
                     // transform from x,y,z to x',y',t
                     transformation::CoordinateTransformation(*m_particle_container,

src/particles/ImpactXParticleContainer.H

@@ -52,11 +52,11 @@ namespace impactx
     {
         enum
         {
-            ux,  ///< momentum in x, normalized to proper velocity
-            uy,  ///< momentum in y, normalized to proper velocity
-            pt,  ///< kinetic energy deviation, (E0-E)/mc^2
-            m_qm, ///< charge to mass ratio, in q_e/m_e (q_e/eV)
-            w,   ///< particle weight, unitless
+            ux,  ///< momentum in x, normalized to reference longitudinal momentum (unitless)
+            uy,  ///< momentum in y, normalized to reference longitudinal momentum (unitless)
+            pt,  ///< kinetic energy deviation, normalized to (E0-E)/mc^2 (unitless)
+            m_qm, ///< charge to mass ratio, normalized to q_e/m_e (q_e/eV)
+            w,   ///< particle weight, (unitless)
             nattribs ///< the number of attributes above (always last)
         };
     };

src/particles/spacecharge/CMakeLists.txt

@@ -1,5 +1,6 @@
 target_sources(ImpactX
   PRIVATE
-    PoissonSolve.cpp
     ForceFromSelfFields.cpp
+    GatherAndPush.cpp
+    PoissonSolve.cpp
 )

src/particles/spacecharge/GatherAndPush.H

@@ -0,0 +1,43 @@
+/* Copyright 2022 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: Axel Huebl, Remi Lehe
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_GATHER_AND_PUSH_H
+#define IMPACTX_GATHER_AND_PUSH_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX_Geometry.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_Vector.H>
+
+#include <unordered_map>
+#include <string>
+
+
+namespace impactx::spacecharge
+{
+    /** Gather force fields and push particles
+     *
+     * details ... ...
+     *
+     * @param[inout] pc container of the particles that deposited rho
+     * @param[in] space_charge_force space charge force component in x,y,z per level
+     * @param[in] geom geometry object
+     * @param[in] slice_ds segment length in meters
+     */
+    void GatherAndPush (
+        ImpactXParticleContainer & pc,
+        std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > const & space_charge_force,
+        const amrex::Vector<amrex::Geometry>& geom,
+        amrex::ParticleReal const slice_ds
+    );
+
+} // namespace impactx
+
+#endif // IMPACTX_GATHER_AND_PUSH_H

src/particles/spacecharge/GatherAndPush.cpp

@@ -0,0 +1,107 @@
+/* Copyright 2022 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 <ablastr/particles/NodalFieldGather.H>
+
+#include <AMReX_BLProfiler.H>
+#include <AMReX_REAL.H>       // for Real
+#include <AMReX_SPACE.H>      // for AMREX_D_DECL
+
+
+namespace impactx::spacecharge
+{
+    void GatherAndPush (
+            ImpactXParticleContainer & pc,
+            std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > const & space_charge_force,
+            const amrex::Vector<amrex::Geometry>& geom,
+            amrex::ParticleReal const slice_ds
+    )
+    {
+        BL_PROFILE("impactx::spacecharge::GatherAndPush");
+
+        using namespace amrex::literals;
+
+        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;
+            for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
+                const int np = pti.numParticles();
+                //const auto t_lev = pti.GetLevel();
+                //const auto index = pti.GetPairIndex();
+                // ...
+
+                //
+                auto scf_arr_x = space_charge_force.at(lev).at("x")[pti].array();
+                auto scf_arr_y = space_charge_force.at(lev).at("y")[pti].array();
+                auto scf_arr_z = space_charge_force.at(lev).at("z")[pti].array();
+
+                // ...
+                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 inv_gamma2 = 1.0_prt / (gamma * gamma);
+
+                amrex::ParticleReal const dt = slice_ds / pc.GetRefParticle().beta() / c0_SI;
+
+                // preparing access to particle data: AoS
+                using PType = ImpactXParticleContainer::ParticleType;
+                auto& aos = pti.GetArrayOfStructs();
+                PType* AMREX_RESTRICT aos_ptr = aos().dataPtr();
+
+                // preparing access to particle data: SoA of Reals
+                auto& soa_real = pti.GetStructOfArrays().GetRealData();
+                amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr();
+                amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr();
+                amrex::ParticleReal* const AMREX_RESTRICT part_pz = soa_real[RealSoA::pt].dataPtr(); // note: currently in z
+
+                // ...
+                amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
+                    // access AoS data such as positions and cpu/id
+                    PType const & AMREX_RESTRICT p = aos_ptr[i];
+
+                    // access SoA Real data
+                    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 (
+                            p.pos(0), p.pos(1), p.pos(2),
+                            scf_arr_x, scf_arr_y, scf_arr_z,
+                            invdr,
+                            prob_lo);
+
+                    // push momentum
+                    px += dt * charge * field_interp[0] * inv_gamma2 / pz_ref_SI;
+                    py += dt * charge * field_interp[1] * inv_gamma2 / pz_ref_SI;
+                    pz += dt * charge * field_interp[2] * inv_gamma2 / pz_ref_SI;
+
+                    // push position is done in the lattice elements
+                });
+
+
+            } // end loop over all particle boxes
+        } // env mesh-refinement level loop
+    }
+} // namespace impactx::spacecharge