Cleaning: Docs, Formatting, Constants, Constness

ECP-WarpX · Sep 19, 2022 · 560165f · 560165f
1 parent 1cc30ea
commit 560165f
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp
@@ -124,8 +124,9 @@ namespace impactx
                 {
 
                     // transform from x',y',t to x,y,z
-                    transformation::CoordinateTransformation(*m_particle_container,
-                                                             transformation::Direction::to_fixed_t);
+                    transformation::CoordinateTransformation(
+                        *m_particle_container,
+                        transformation::Direction::to_fixed_t);
 
                     // Note: The following operation assume that
                     // the particles are in x, y, z coordinates.
@@ -150,7 +151,10 @@ namespace impactx
                     // gather and space-charge push in x,y,z , assuming the space-charge
                     // field is the same before/after transformation
                     // TODO: This is currently using linear order.
-                    spacecharge::GatherAndPush(*m_particle_container, m_space_charge_field, this->geom, slice_ds);
+                    spacecharge::GatherAndPush(*m_particle_container,
+                                               m_space_charge_field,
+                                               this->geom,
+                                               slice_ds);
 
                     // transform from x,y,z to x',y',t
                     transformation::CoordinateTransformation(*m_particle_container,

diff --git a/src/particles/spacecharge/GatherAndPush.H b/src/particles/spacecharge/GatherAndPush.H
@@ -22,9 +22,12 @@
 
 namespace impactx::spacecharge
 {
-    /** Gather force fields and push particles
+    /** Gather force fields and push particles in x,y,z
      *
-     * details ... ...
+     * This gathers the space charge field with respect to particle position
+     * and shape. The momentum of all particles is then pushed using a common
+     * time step given by the reference particle speed and ds slice. The
+     * position push is done in the lattice elements and not here.
      *
      * @param[inout] pc container of the particles that deposited rho
      * @param[in] space_charge_field space charge force component in x,y,z per level

diff --git a/src/particles/spacecharge/GatherAndPush.cpp b/src/particles/spacecharge/GatherAndPush.cpp
@@ -45,16 +45,13 @@ namespace impactx::spacecharge
             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_field.at(lev).at("x")[pti].array();
-                auto scf_arr_y = space_charge_field.at(lev).at("y")[pti].array();
-                auto scf_arr_z = space_charge_field.at(lev).at("z")[pti].array();
+                // 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;
@@ -65,16 +62,19 @@ namespace impactx::spacecharge
 
                 // preparing access to particle data: AoS
                 using PType = ImpactXParticleContainer::ParticleType;
-                auto& aos = pti.GetArrayOfStructs();
-                PType* AMREX_RESTRICT aos_ptr = aos().dataPtr();
+                auto const & aos = pti.GetArrayOfStructs();
+                PType const * const 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
 
-                // ...
+                // 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
                 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];
@@ -93,9 +93,9 @@ namespace impactx::spacecharge
                             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;
+                    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
                 });