Momentum push conversion clean up for clarity and optimization

1 files changed, 72 insertions, 74 deletions

diff --git a/Source/Particles/RigidInjectedParticleContainer.cpp b/Source/Particles/RigidInjectedParticleContainer.cpp
index 7e4434278..09df32db4 100644
--- a/Source/Particles/RigidInjectedParticleContainer.cpp
+++ b/Source/Particles/RigidInjectedParticleContainer.cpp
@@ -214,56 +214,40 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
                                        Real dt)
 {
 
-    if (WarpX::do_boosted_frame_diagnostic && do_boosted_frame_diags)
-    {
-        copy_attribs(pti, xp.dataPtr(), yp.dataPtr(), zp.dataPtr());
-    }
-
     // This wraps the momentum and position advance so that inheritors can modify the call.
     auto& attribs = pti.GetAttribs();
     auto& uxp = attribs[PIdx::ux];
     auto& uyp = attribs[PIdx::uy];
     auto& uzp = attribs[PIdx::uz];
-    auto& Exp = attribs[PIdx::Ex];
-    auto& Eyp = attribs[PIdx::Ey];
-    auto& Ezp = attribs[PIdx::Ez];
-    auto& Bxp = attribs[PIdx::Bx];
-    auto& Byp = attribs[PIdx::By];
-    auto& Bzp = attribs[PIdx::Bz];
-    const long np  = pti.numParticles();
 
     // Save the position and momenta, making copies
     Cuda::ManagedDeviceVector<Real> xp_save, yp_save, zp_save;
     RealVector uxp_save, uyp_save, uzp_save;
 
-    Real* AMREX_RESTRICT x = xp.dataPtr();
-    Real* AMREX_RESTRICT y = yp.dataPtr();
-    Real* AMREX_RESTRICT z = zp.dataPtr();
-    Real* AMREX_RESTRICT gi = giv.dataPtr();
-    Real* AMREX_RESTRICT uxpp = uxp.dataPtr();
-    Real* AMREX_RESTRICT uypp = uyp.dataPtr();
-    Real* AMREX_RESTRICT uzpp = uzp.dataPtr();
-    Real* AMREX_RESTRICT Expp = Exp.dataPtr();
-    Real* AMREX_RESTRICT Eypp = Eyp.dataPtr();
-    Real* AMREX_RESTRICT Ezpp = Ezp.dataPtr();
-    Real* AMREX_RESTRICT Bxpp = Bxp.dataPtr();
-    Real* AMREX_RESTRICT Bypp = Byp.dataPtr();
-    Real* AMREX_RESTRICT Bzpp = Bzp.dataPtr();
+    Real* const AMREX_RESTRICT x = xp.dataPtr();
+    Real* const AMREX_RESTRICT y = yp.dataPtr();
+    Real* const AMREX_RESTRICT z = zp.dataPtr();
+    Real* const AMREX_RESTRICT gi = giv.dataPtr();
+    Real* const AMREX_RESTRICT ux = uxp.dataPtr();
+    Real* const AMREX_RESTRICT uy = uyp.dataPtr();
+    Real* const AMREX_RESTRICT uz = uzp.dataPtr();
+    Real* const AMREX_RESTRICT Exp = attribs[PIdx::Ex].dataPtr();
+    Real* const AMREX_RESTRICT Eyp = attribs[PIdx::Ey].dataPtr();
+    Real* const AMREX_RESTRICT Ezp = attribs[PIdx::Ez].dataPtr();
+    Real* const AMREX_RESTRICT Bxp = attribs[PIdx::Bx].dataPtr();
+    Real* const AMREX_RESTRICT Byp = attribs[PIdx::By].dataPtr();
+    Real* const AMREX_RESTRICT Bzp = attribs[PIdx::Bz].dataPtr();
 
     if (!done_injecting_lev) {
-        xp_save = xp;
-        yp_save = yp;
-        zp_save = zp;
-        uxp_save = uxp;
-        uyp_save = uyp;
-        uzp_save = uzp;
-
-        Real* AMREX_RESTRICT xp_savep = xp_save.dataPtr();
-        Real* AMREX_RESTRICT yp_savep = yp_save.dataPtr();
-        Real* AMREX_RESTRICT zp_savep = zp_save.dataPtr();
-        Real* AMREX_RESTRICT uxp_savep = uxp_save.dataPtr();
-        Real* AMREX_RESTRICT uyp_savep = uyp_save.dataPtr();
-        Real* AMREX_RESTRICT uzp_savep = uzp_save.dataPtr();
+        if (!(WarpX::do_boosted_frame_diagnostic && do_boosted_frame_diags)) {
+            // If the old values are not already saved, create copies here.
+            xp_save = xp;
+            yp_save = yp;
+            zp_save = zp;
+            uxp_save = uxp;
+            uyp_save = uyp;
+            uzp_save = uzp;
+        }
 
         // Scale the fields of particles about to cross the injection plane.
         // This only approximates what should be happening. The particles
@@ -273,12 +257,12 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
             [=] AMREX_GPU_DEVICE (long i) {
             const Real dtscale = dt - (zinject_plane_lev_previous - z[i])/(vzbeam_ave_boosted + WarpX::beta_boost*PhysConst::c);
             if (0. < dtscale && dtscale < dt) {
-                Expp[i] *= dtscale;
-                Eypp[i] *= dtscale;
-                Ezpp[i] *= dtscale;
-                Bxpp[i] *= dtscale;
-                Bypp[i] *= dtscale;
-                Bzpp[i] *= dtscale;
+                Exp[i] *= dtscale;
+                Eyp[i] *= dtscale;
+                Ezp[i] *= dtscale;
+                Bxp[i] *= dtscale;
+                Byp[i] *= dtscale;
+                Bzp[i] *= dtscale;
             }
         }
         );
@@ -293,36 +277,50 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
         const int tid = 0;
 #endif
 
-        Real* AMREX_RESTRICT xp_savep = xp_save.dataPtr();
-        Real* AMREX_RESTRICT yp_savep = yp_save.dataPtr();
-        Real* AMREX_RESTRICT zp_savep = zp_save.dataPtr();
-        Real* AMREX_RESTRICT uxp_savep = uxp_save.dataPtr();
-        Real* AMREX_RESTRICT uyp_savep = uyp_save.dataPtr();
-        Real* AMREX_RESTRICT uzp_savep = uzp_save.dataPtr();
+        Real* AMREX_RESTRICT x_save;
+        Real* AMREX_RESTRICT y_save;
+        Real* AMREX_RESTRICT z_save;
+        Real* AMREX_RESTRICT ux_save;
+        Real* AMREX_RESTRICT uy_save;
+        Real* AMREX_RESTRICT uz_save;
+        if (!(WarpX::do_boosted_frame_diagnostic && do_boosted_frame_diags)) {
+            x_save = xp_save.dataPtr();
+            y_save = yp_save.dataPtr();
+            z_save = zp_save.dataPtr();
+            ux_save = uxp_save.dataPtr();
+            uy_save = uyp_save.dataPtr();
+            uz_save = uzp_save.dataPtr();
+        } else {
+            x_save = pti.GetAttribs(particle_comps["xold"]).dataPtr();
+            y_save = pti.GetAttribs(particle_comps["yold"]).dataPtr();
+            z_save = pti.GetAttribs(particle_comps["zold"]).dataPtr();
+            ux_save = pti.GetAttribs(particle_comps["uxold"]).dataPtr();
+            uy_save = pti.GetAttribs(particle_comps["uyold"]).dataPtr();
+            uz_save = pti.GetAttribs(particle_comps["uzold"]).dataPtr();
+        }
 
         // Undo the push for particles not injected yet.
         // The zp are advanced a fixed amount.
         amrex::ParallelFor( pti.numParticles(),
             [=] AMREX_GPU_DEVICE (long i) {
             if (z[i] <= zinject_plane_lev) {
-                uxpp[i] = uxp_savep[i];
-                uypp[i] = uyp_savep[i];
-                uzpp[i] = uzp_savep[i];
-                gi[i] = 1./std::sqrt(1. + (uxpp[i]*uxpp[i] + uypp[i]*uypp[i] + uzpp[i]*uzpp[i])/(PhysConst::c*PhysConst::c));
-                x[i] = xp_savep[i];
-                y[i] = yp_savep[i];
+                ux[i] = ux_save[i];
+                uy[i] = uy_save[i];
+                uz[i] = uz_save[i];
+                gi[i] = 1./std::sqrt(1. + (ux[i]*ux[i] + uy[i]*uy[i] + uz[i]*uz[i])/(PhysConst::c*PhysConst::c));
+                x[i] = x_save[i];
+                y[i] = y_save[i];
                 if (rigid_advance) {
-                    z[i] = zp_savep[i] + dt*vzbeam_ave_boosted;
+                    z[i] = z_save[i] + dt*vzbeam_ave_boosted;
                 }
                 else {
-                    z[i] = zp_savep[i] + dt*uzpp[i]*gi[i];
+                    z[i] = z_save[i] + dt*uz[i]*gi[i];
                 }
                 done_injecting_temp[tid] = 0;
             }
         }
         );
     }
-
 }
 
 void
@@ -459,24 +457,24 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
 
             // This wraps the momentum advance so that inheritors can modify the call.
             // Extract pointers to the different particle quantities
-            Real* AMREX_RESTRICT zp = m_zp[thread_num].dataPtr();
-            Real* AMREX_RESTRICT gi = m_giv[thread_num].dataPtr();
-            Real* AMREX_RESTRICT uxpp = uxp.dataPtr();
-            Real* AMREX_RESTRICT uypp = uyp.dataPtr();
-            Real* AMREX_RESTRICT uzpp = uzp.dataPtr();
-            Real* AMREX_RESTRICT uxp_savep = uxp_save.dataPtr();
-            Real* AMREX_RESTRICT uyp_savep = uyp_save.dataPtr();
-            Real* AMREX_RESTRICT uzp_savep = uzp_save.dataPtr();
-            Real* AMREX_RESTRICT Expp = Exp.dataPtr();
-            Real* AMREX_RESTRICT Eypp = Eyp.dataPtr();
-            Real* AMREX_RESTRICT Ezpp = Ezp.dataPtr();
-            Real* AMREX_RESTRICT Bxpp = Bxp.dataPtr();
-            Real* AMREX_RESTRICT Bypp = Byp.dataPtr();
-            Real* AMREX_RESTRICT Bzpp = Bzp.dataPtr();
+            const Real* AMREX_RESTRICT zp = m_zp[thread_num].dataPtr();
+            Real* const AMREX_RESTRICT gi = m_giv[thread_num].dataPtr();
+            Real* const AMREX_RESTRICT uxpp = uxp.dataPtr();
+            Real* const AMREX_RESTRICT uypp = uyp.dataPtr();
+            Real* const AMREX_RESTRICT uzpp = uzp.dataPtr();
+            const Real* AMREX_RESTRICT uxp_savep = uxp_save.dataPtr();
+            const Real* AMREX_RESTRICT uyp_savep = uyp_save.dataPtr();
+            const Real* AMREX_RESTRICT uzp_savep = uzp_save.dataPtr();
+            const Real* AMREX_RESTRICT Expp = Exp.dataPtr();
+            const Real* AMREX_RESTRICT Eypp = Eyp.dataPtr();
+            const Real* AMREX_RESTRICT Ezpp = Ezp.dataPtr();
+            const Real* AMREX_RESTRICT Bxpp = Bxp.dataPtr();
+            const Real* AMREX_RESTRICT Bypp = Byp.dataPtr();
+            const Real* AMREX_RESTRICT Bzpp = Bzp.dataPtr();
 
             // Loop over the particles and update their momentum
             const Real q = this->charge;
-            const Real m = this-> mass;
+            const Real m = this->mass;
             if (WarpX::particle_pusher_algo == ParticlePusherAlgo::Boris){
                 amrex::ParallelFor( pti.numParticles(),
                     [=] AMREX_GPU_DEVICE (long i) {