remove the copies between soa and aos for the particle positions

1 files changed, 80 insertions, 78 deletions

diff --git a/Source/Particles/RigidInjectedParticleContainer.cpp b/Source/Particles/RigidInjectedParticleContainer.cpp
index f3b502a0a..ad2c89428 100644
--- a/Source/Particles/RigidInjectedParticleContainer.cpp
+++ b/Source/Particles/RigidInjectedParticleContainer.cpp
@@ -78,8 +78,6 @@ RigidInjectedParticleContainer::RemapParticles()
                 // Note that the particles are already in the boosted frame.
                 // This value is saved to advance the particles not injected yet
 
-                Gpu::ManagedDeviceVector<ParticleReal> xp, yp, zp;
-
                 for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
                 {
                     auto& attribs = pti.GetAttribs();
@@ -87,9 +85,9 @@ RigidInjectedParticleContainer::RemapParticles()
                     auto& uyp = attribs[PIdx::uy];
                     auto& uzp = attribs[PIdx::uz];
 
-                    // Copy data from particle container to temp arrays
-                    pti.GetPosition(xp, yp, zp);
-
+                    auto& aos = pti.GetArrayOfStructs();
+                    ParticleType* AMREX_RESTRICT const pstruct = aos().dataPtr();
+                    
                     // Loop over particles
                     const long np = pti.numParticles();
                     for (int i=0 ; i < np ; i++) {
@@ -98,20 +96,16 @@ RigidInjectedParticleContainer::RemapParticles()
                         const Real vzpr = uzp[i]/gammapr;
 
                         // Back out the value of z_lab
-                        const Real z_lab = (zp[i] + uz_boost*t_lab + WarpX::gamma_boost*t_lab*vzpr)/(WarpX::gamma_boost + uz_boost*vzpr/csq);
+                        const Real z_lab = (pstruct[i].pos(2) + uz_boost*t_lab + WarpX::gamma_boost*t_lab*vzpr)/(WarpX::gamma_boost + uz_boost*vzpr/csq);
 
                         // Time of the particle in the boosted frame given its position in the lab frame at t=0.
                         const Real tpr = WarpX::gamma_boost*t_lab - uz_boost*z_lab/csq;
 
                         // Adjust the position, taking away its motion from its own velocity and adding
                         // the motion from the average velocity
-                        zp[i] = zp[i] + tpr*vzpr - tpr*vzbeam_ave_boosted;
+                        pstruct[i].pos(2) += tpr*vzpr - tpr*vzbeam_ave_boosted;
 
                     }
-
-                    // Copy the data back to the particle container
-                    pti.SetPosition(xp, yp, zp);
-
                 }
             }
         }
@@ -138,8 +132,6 @@ RigidInjectedParticleContainer::BoostandRemapParticles()
 #pragma omp parallel
 #endif
     {
-        Gpu::ManagedDeviceVector<ParticleReal> xp, yp, zp;
-
         for (WarpXParIter pti(*this, 0); pti.isValid(); ++pti)
         {
 
@@ -148,8 +140,8 @@ RigidInjectedParticleContainer::BoostandRemapParticles()
             auto& uyp = attribs[PIdx::uy];
             auto& uzp = attribs[PIdx::uz];
 
-            // Copy data from particle container to temp arrays
-            pti.GetPosition(xp, yp, zp);
+            auto& aos = pti.GetArrayOfStructs();
+            ParticleType* AMREX_RESTRICT const pstruct = aos().dataPtr();
 
             // Loop over particles
             const long np = pti.numParticles();
@@ -164,24 +156,24 @@ RigidInjectedParticleContainer::BoostandRemapParticles()
                 // t0_lab is the time in the lab frame that the particles reaches z=0
                 // The location and time (z=0, t=0) is a synchronization point between the
                 // lab and boosted frames.
-                const Real t0_lab = -zp[i]/vz_lab;
+                const Real t0_lab = -pstruct[i].pos(2)/vz_lab;
 
                 if (!projected) {
-                    xp[i] += t0_lab*vx_lab;
-                    yp[i] += t0_lab*vy_lab;
+                    pstruct[i].pos(0) += t0_lab*vx_lab;
+                    pstruct[i].pos(1) += t0_lab*vy_lab;
                 }
                 if (focused) {
                     // Correct for focusing effect from shift from z=0 to zinject
                     const Real tfocus = -zinject_plane*WarpX::gamma_boost/vz_lab;
-                    xp[i] -= tfocus*vx_lab;
-                    yp[i] -= tfocus*vy_lab;
+                    pstruct[i].pos(0) -= tfocus*vx_lab;
+                    pstruct[i].pos(1) -= tfocus*vy_lab;
                 }
 
                 // Time of the particle in the boosted frame given its position in the lab frame at t=0.
-                const Real tpr = -WarpX::gamma_boost*WarpX::beta_boost*zp[i]/PhysConst::c;
+                const Real tpr = -WarpX::gamma_boost*WarpX::beta_boost*pstruct[i].pos(2)/PhysConst::c;
 
                 // Position of the particle in the boosted frame given its position in the lab frame at t=0.
-                const Real zpr = WarpX::gamma_boost*zp[i];
+                const Real zpr = WarpX::gamma_boost*pstruct[i].pos(2);
 
                 // Momentum of the particle in the boosted frame (assuming that it is fixed).
                 uzp[i] = WarpX::gamma_boost*(uzp[i] - WarpX::beta_boost*PhysConst::c*gamma_lab);
@@ -189,30 +181,22 @@ RigidInjectedParticleContainer::BoostandRemapParticles()
                 // Put the particle at the location in the boosted frame at boost frame t=0,
                 if (rigid_advance) {
                     // with the particle moving at the average velocity
-                    zp[i] = zpr - vzbeam_ave_boosted*tpr;
+                    pstruct[i].pos(2) = zpr - vzbeam_ave_boosted*tpr;
                 }
                 else {
                     // with the particle moving with its own velocity
                     const Real gammapr = std::sqrt(1. + (uxp[i]*uxp[i] + uyp[i]*uyp[i] + uzp[i]*uzp[i])/(PhysConst::c*PhysConst::c));
                     const Real vzpr = uzp[i]/gammapr;
-                    zp[i] = zpr - vzpr*tpr;
+                    pstruct[i].pos(2) = zpr - vzpr*tpr;
                 }
 
             }
-
-            // Copy the data back to the particle container
-            pti.SetPosition(xp, yp, zp);
-
         }
     }
 }
 
 void
-RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
-                                       Gpu::ManagedDeviceVector<ParticleReal>& xp,
-                                       Gpu::ManagedDeviceVector<ParticleReal>& yp,
-                                       Gpu::ManagedDeviceVector<ParticleReal>& zp,
-                                       Real dt, DtType a_dt_type)
+RigidInjectedParticleContainer::PushPX (WarpXParIter& pti, Real dt, DtType a_dt_type)
 {
 
     // This wraps the momentum and position advance so that inheritors can modify the call.
@@ -225,9 +209,9 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
     Gpu::ManagedDeviceVector<ParticleReal> xp_save, yp_save, zp_save;
     RealVector uxp_save, uyp_save, uzp_save;
 
-    ParticleReal* const AMREX_RESTRICT x = xp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT y = yp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT z = zp.dataPtr();
+    auto& aos = pti.GetArrayOfStructs();
+    ParticleType* AMREX_RESTRICT const pstruct = aos().dataPtr();
+        
     ParticleReal* const AMREX_RESTRICT ux = uxp.dataPtr();
     ParticleReal* const AMREX_RESTRICT uy = uyp.dataPtr();
     ParticleReal* const AMREX_RESTRICT uz = uzp.dataPtr();
@@ -237,15 +221,37 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
     ParticleReal* const AMREX_RESTRICT Bxp = attribs[PIdx::Bx].dataPtr();
     ParticleReal* const AMREX_RESTRICT Byp = attribs[PIdx::By].dataPtr();
     ParticleReal* const AMREX_RESTRICT Bzp = attribs[PIdx::Bz].dataPtr();
-
-    if (!done_injecting_lev) {
+    
+    if (!done_injecting_lev)
+    {
         // 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;
+        const auto np = pti.numParticles();
+
+        xp_save.resize(np);
+        yp_save.resize(np);
+        zp_save.resize(np);
+
+        uxp_save.resize(np);
+        uyp_save.resize(np);
+        uzp_save.resize(np);
+
+        amrex::Real* const AMREX_RESTRICT xp_save_ptr = xp_save.dataPtr();
+        amrex::Real* const AMREX_RESTRICT yp_save_ptr = yp_save.dataPtr();
+        amrex::Real* const AMREX_RESTRICT zp_save_ptr = zp_save.dataPtr();
+
+        amrex::Real* const AMREX_RESTRICT uxp_save_ptr = uxp_save.dataPtr();
+        amrex::Real* const AMREX_RESTRICT uyp_save_ptr = uyp_save.dataPtr();
+        amrex::Real* const AMREX_RESTRICT uzp_save_ptr = uzp_save.dataPtr();
+        
+        amrex::ParallelFor( np,
+                            [=] AMREX_GPU_DEVICE (long i) {
+                                xp_save_ptr[i] = pstruct[i].pos(0);
+                                yp_save_ptr[i] = pstruct[i].pos(1);
+                                zp_save_ptr[i] = pstruct[i].pos(2);
+                                uxp_save_ptr[i] = ux[i];
+                                uyp_save_ptr[i] = uy[i];
+                                uzp_save_ptr[i] = uz[i];
+                            });        
 
         // Scale the fields of particles about to cross the injection plane.
         // This only approximates what should be happening. The particles
@@ -256,7 +262,7 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
         const Real vz_ave_boosted = vzbeam_ave_boosted;
         amrex::ParallelFor( pti.numParticles(),
             [=] AMREX_GPU_DEVICE (long i) {
-            const Real dtscale = dt - (z_plane_previous - z[i])/(vz_ave_boosted + v_boost);
+            const Real dtscale = dt - (z_plane_previous - pstruct[i].pos(2))/(vz_ave_boosted + v_boost);
             if (0. < dtscale && dtscale < dt) {
                 Exp[i] *= dtscale;
                 Eyp[i] *= dtscale;
@@ -269,7 +275,7 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
         );
     }
 
-    PhysicalParticleContainer::PushPX(pti, xp, yp, zp, dt, a_dt_type);
+    PhysicalParticleContainer::PushPX(pti, dt, a_dt_type);
 
     if (!done_injecting_lev) {
 
@@ -287,23 +293,22 @@ RigidInjectedParticleContainer::PushPX(WarpXParIter& pti,
         const bool rigid = rigid_advance;
         const Real inv_csq = 1./(PhysConst::c*PhysConst::c);
         amrex::ParallelFor( pti.numParticles(),
-            [=] AMREX_GPU_DEVICE (long i) {
-            if (z[i] <= z_plane_lev) {
-                ux[i] = ux_save[i];
-                uy[i] = uy_save[i];
-                uz[i] = uz_save[i];
-                x[i] = x_save[i];
-                y[i] = y_save[i];
-                if (rigid) {
-                    z[i] = z_save[i] + dt*vz_ave_boosted;
-                }
-                else {
-                    const Real gi = 1./std::sqrt(1. + (ux[i]*ux[i] + uy[i]*uy[i] + uz[i]*uz[i])*inv_csq);
-                    z[i] = z_save[i] + dt*uz[i]*gi;
-                }
-            }
-        }
-        );
+                            [=] AMREX_GPU_DEVICE (long i) {
+                                if (pstruct[i].pos(2) <= z_plane_lev) {
+                                    ux[i] = ux_save[i];
+                                    uy[i] = uy_save[i];
+                                    uz[i] = uz_save[i];
+                                    pstruct[i].pos(0) = x_save[i];
+                                    pstruct[i].pos(1) = y_save[i];
+                                    if (rigid) {
+                                        pstruct[i].pos(2) = z_save[i] + dt*vz_ave_boosted;
+                                    }
+                                    else {
+                                        const Real gi = 1./std::sqrt(1. + (ux[i]*ux[i] + uy[i]*uy[i] + uz[i]*uz[i])*inv_csq);
+                                        pstruct[i].pos(2) = z_save[i] + dt*uz[i]*gi;
+                                    }
+                                }
+                            });
     }
 }
 
@@ -392,11 +397,6 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
             const FArrayBox& byfab = By[pti];
             const FArrayBox& bzfab = Bz[pti];
 
-            //
-            // copy data from particle container to temp arrays
-            //
-            pti.GetPosition(m_xp[thread_num], m_yp[thread_num], m_zp[thread_num]);
-
             int e_is_nodal = Ex.is_nodal() and Ey.is_nodal() and Ez.is_nodal();
             FieldGather(pti, Exp, Eyp, Ezp, Bxp, Byp, Bzp,
                         &exfab, &eyfab, &ezfab, &bxfab, &byfab, &bzfab,
@@ -410,7 +410,9 @@ 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
-            const ParticleReal* const AMREX_RESTRICT zp = m_zp[thread_num].dataPtr();
+            const auto& aos = pti.GetArrayOfStructs();
+            const ParticleType* AMREX_RESTRICT const pstruct = aos().dataPtr();
+
             ParticleReal* const AMREX_RESTRICT uxpp = uxp.dataPtr();
             ParticleReal* const AMREX_RESTRICT uypp = uyp.dataPtr();
             ParticleReal* const AMREX_RESTRICT uzpp = uzp.dataPtr();
@@ -476,15 +478,15 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
             const ParticleReal* const AMREX_RESTRICT uz_save = uzp_save.dataPtr();
             const ParticleReal zz = zinject_plane_levels[lev];
             amrex::ParallelFor( pti.numParticles(),
-                [=] AMREX_GPU_DEVICE (long i) {
-                if (zp[i] <= zz) {
-                    uxpp[i] = ux_save[i];
-                    uypp[i] = uy_save[i];
-                    uzpp[i] = uz_save[i];
-                }
-            }
-            );
-
+                                [=] AMREX_GPU_DEVICE (long i) {
+                                    if (pstruct[i].pos(2) <= zz) {
+                                        uxpp[i] = ux_save[i];
+                                        uypp[i] = uy_save[i];
+                                        uzpp[i] = uz_save[i];
+                                    }
+                                }
+                                );
+            
         }
     }
 }