Merge branch 'dev' into elementary_process

1 files changed, 21 insertions, 88 deletions

diff --git a/Source/Particles/WarpXParticleContainer.cpp b/Source/Particles/WarpXParticleContainer.cpp
index 85cac0dbe..cd5bcaaf9 100644
--- a/Source/Particles/WarpXParticleContainer.cpp
+++ b/Source/Particles/WarpXParticleContainer.cpp
@@ -30,56 +30,6 @@ WarpXParIter::WarpXParIter (ContainerType& pc, int level)
 {
 }
 
-#if (AMREX_SPACEDIM == 2)
-void
-WarpXParIter::GetPosition (Gpu::ManagedDeviceVector<ParticleReal>& xp,
-                           Gpu::ManagedDeviceVector<ParticleReal>& yp,
-                           Gpu::ManagedDeviceVector<ParticleReal>& zp) const
-{
-    amrex::ParIter<0,0,PIdx::nattribs>::GetPosition(xp, zp);
-#ifdef WARPX_DIM_RZ
-    const auto& attribs = GetAttribs();
-    const auto& thetap = attribs[PIdx::theta];
-    yp.resize(xp.size());
-    ParticleReal* const AMREX_RESTRICT x = xp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT y = yp.dataPtr();
-    const ParticleReal* const AMREX_RESTRICT theta = thetap.dataPtr();
-    amrex::ParallelFor( xp.size(),
-        [=] AMREX_GPU_DEVICE (long i) {
-        // The x stored in the particles is actually the radius
-        y[i] = x[i]*std::sin(theta[i]);
-        x[i] = x[i]*std::cos(theta[i]);
-    });
-#else
-    yp.resize(xp.size(), std::numeric_limits<ParticleReal>::quiet_NaN());
-#endif
-}
-
-void
-WarpXParIter::SetPosition (const Gpu::ManagedDeviceVector<ParticleReal>& xp,
-                           const Gpu::ManagedDeviceVector<ParticleReal>& yp,
-                           const Gpu::ManagedDeviceVector<ParticleReal>& zp)
-{
-#ifdef WARPX_DIM_RZ
-    auto& attribs = GetAttribs();
-    auto& thetap = attribs[PIdx::theta];
-    Gpu::ManagedDeviceVector<ParticleReal> rp(xp.size());
-    const ParticleReal* const AMREX_RESTRICT x = xp.dataPtr();
-    const ParticleReal* const AMREX_RESTRICT y = yp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT r = rp.dataPtr();
-    ParticleReal* const AMREX_RESTRICT theta = thetap.dataPtr();
-    amrex::ParallelFor( xp.size(),
-        [=] AMREX_GPU_DEVICE (long i) {
-        theta[i] = std::atan2(y[i], x[i]);
-        r[i] = std::sqrt(x[i]*x[i] + y[i]*y[i]);
-    });
-    amrex::ParIter<0,0,PIdx::nattribs>::SetPosition(rp, zp);
-#else
-    amrex::ParIter<0,0,PIdx::nattribs>::SetPosition(xp, zp);
-#endif
-}
-#endif
-
 WarpXParticleContainer::WarpXParticleContainer (AmrCore* amr_core, int ispecies)
     : ParticleContainer<0,0,PIdx::nattribs>(amr_core->GetParGDB())
     , species_id(ispecies)
@@ -116,10 +66,6 @@ WarpXParticleContainer::WarpXParticleContainer (AmrCore* amr_core, int ispecies)
     local_jx.resize(num_threads);
     local_jy.resize(num_threads);
     local_jz.resize(num_threads);
-    m_xp.resize(num_threads);
-    m_yp.resize(num_threads);
-    m_zp.resize(num_threads);
-
 }
 
 void
@@ -349,9 +295,7 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti,
     // CPU, tiling: deposit into local_jx
     // (same for jx and jz)
 
-    ParticleReal* AMREX_RESTRICT xp = m_xp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT zp = m_zp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT yp = m_yp[thread_num].dataPtr() + offset;
+    const auto GetPosition = GetParticlePosition(pti, offset);
 
     // Lower corner of tile box physical domain
     // Note that this includes guard cells since it is after tilebox.ngrow
@@ -362,19 +306,19 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti,
     if (WarpX::current_deposition_algo == CurrentDepositionAlgo::Esirkepov) {
         if        (WarpX::nox == 1){
             doEsirkepovDepositionShapeN<1>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                GetPosition, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_arr, jy_arr, jz_arr, np_to_depose, dt, dx, xyzmin, lo, q,
                 WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 2){
             doEsirkepovDepositionShapeN<2>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                GetPosition, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_arr, jy_arr, jz_arr, np_to_depose, dt, dx, xyzmin, lo, q,
                 WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 3){
             doEsirkepovDepositionShapeN<3>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                GetPosition, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_arr, jy_arr, jz_arr, np_to_depose, dt, dx, xyzmin, lo, q,
                 WarpX::n_rz_azimuthal_modes);
@@ -382,19 +326,19 @@ WarpXParticleContainer::DepositCurrent(WarpXParIter& pti,
     } else {
         if        (WarpX::nox == 1){
             doDepositionShapeN<1>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                GetPosition, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_fab, jy_fab, jz_fab, np_to_depose, dt, dx,
                 xyzmin, lo, q);
         } else if (WarpX::nox == 2){
             doDepositionShapeN<2>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                GetPosition, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_fab, jy_fab, jz_fab, np_to_depose, dt, dx,
                 xyzmin, lo, q);
         } else if (WarpX::nox == 3){
             doDepositionShapeN<3>(
-                xp, yp, zp, wp.dataPtr() + offset, uxp.dataPtr() + offset,
+                GetPosition, wp.dataPtr() + offset, uxp.dataPtr() + offset,
                 uyp.dataPtr() + offset, uzp.dataPtr() + offset, ion_lev,
                 jx_fab, jy_fab, jz_fab, np_to_depose, dt, dx,
                 xyzmin, lo, q);
@@ -488,9 +432,7 @@ WarpXParticleContainer::DepositCharge (WarpXParIter& pti, RealVector& wp,
     // GPU, no tiling: deposit directly in rho
     // CPU, tiling: deposit into local_rho
 
-    ParticleReal* AMREX_RESTRICT xp = m_xp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT zp = m_zp[thread_num].dataPtr() + offset;
-    ParticleReal* AMREX_RESTRICT yp = m_yp[thread_num].dataPtr() + offset;
+    const auto GetPosition = GetParticlePosition(pti, offset);
 
     // Lower corner of tile box physical domain
     // Note that this includes guard cells since it is after tilebox.ngrow
@@ -500,13 +442,13 @@ WarpXParticleContainer::DepositCharge (WarpXParIter& pti, RealVector& wp,
 
     BL_PROFILE_VAR_START(blp_ppc_chd);
     if        (WarpX::nox == 1){
-        doChargeDepositionShapeN<1>(xp, yp, zp, wp.dataPtr()+offset, ion_lev,
+        doChargeDepositionShapeN<1>(GetPosition, wp.dataPtr()+offset, ion_lev,
                                     rho_arr, np_to_depose, dx, xyzmin, lo, q);
     } else if (WarpX::nox == 2){
-        doChargeDepositionShapeN<2>(xp, yp, zp, wp.dataPtr()+offset, ion_lev,
+        doChargeDepositionShapeN<2>(GetPosition, wp.dataPtr()+offset, ion_lev,
                                     rho_arr, np_to_depose, dx, xyzmin, lo, q);
     } else if (WarpX::nox == 3){
-        doChargeDepositionShapeN<3>(xp, yp, zp, wp.dataPtr()+offset, ion_lev,
+        doChargeDepositionShapeN<3>(GetPosition, wp.dataPtr()+offset, ion_lev,
                                     rho_arr, np_to_depose, dx, xyzmin, lo, q);
     }
     BL_PROFILE_VAR_STOP(blp_ppc_chd);
@@ -545,8 +487,6 @@ WarpXParticleContainer::DepositCharge (amrex::Vector<std::unique_ptr<amrex::Mult
             const long np = pti.numParticles();
             auto& wp = pti.GetAttribs(PIdx::w);
 
-            pti.GetPosition(m_xp[thread_num], m_yp[thread_num], m_zp[thread_num]);
-
             int* AMREX_RESTRICT ion_lev;
             if (do_field_ionization){
                 ion_lev = pti.GetiAttribs(particle_icomps["ionization_level"]).dataPtr();
@@ -615,8 +555,6 @@ WarpXParticleContainer::GetChargeDensity (int lev, bool local)
             const long np = pti.numParticles();
             auto& wp = pti.GetAttribs(PIdx::w);
 
-            pti.GetPosition(m_xp[thread_num], m_yp[thread_num], m_zp[thread_num]);
-
             int* AMREX_RESTRICT ion_lev;
             if (do_field_ionization){
                 ion_lev = pti.GetiAttribs(particle_icomps["ionization_level"]).dataPtr();
@@ -772,32 +710,27 @@ WarpXParticleContainer::PushX (int lev, amrex::Real dt)
             //
             // Particle Push
             //
-            // Extract pointers to particle position and momenta, for this particle tile
-            // - positions are stored as an array of struct, in `ParticleType`
-            ParticleType * AMREX_RESTRICT pstructs = &(pti.GetArrayOfStructs()[0]);
+
+            const auto GetPosition = GetParticlePosition(pti);
+                  auto SetPosition = SetParticlePosition(pti);
+
             // - momenta are stored as a struct of array, in `attribs`
             auto& attribs = pti.GetAttribs();
             ParticleReal* AMREX_RESTRICT ux = attribs[PIdx::ux].dataPtr();
             ParticleReal* AMREX_RESTRICT uy = attribs[PIdx::uy].dataPtr();
             ParticleReal* AMREX_RESTRICT uz = attribs[PIdx::uz].dataPtr();
 #ifdef WARPX_DIM_RZ
+            auto& aos = pti.GetArrayOfStructs();
+            ParticleType* AMREX_RESTRICT const pstruct = aos().dataPtr();
             ParticleReal* AMREX_RESTRICT theta = attribs[PIdx::theta].dataPtr();
 #endif
             // Loop over the particles and update their position
             amrex::ParallelFor( pti.numParticles(),
                 [=] AMREX_GPU_DEVICE (long i) {
-                    ParticleType& p = pstructs[i]; // Particle object that gets updated
-                    ParticleReal x, y, z; // Temporary variables
-#ifndef WARPX_DIM_RZ
-                    GetPosition( x, y, z, p ); // Initialize x, y, z
-                    UpdatePosition( x, y, z, ux[i], uy[i], uz[i], dt);
-                    SetPosition( p, x, y, z ); // Update the object p
-#else
-                    // For WARPX_DIM_RZ, the particles are still pushed in 3D Cartesian
-                    GetCartesianPositionFromCylindrical( x, y, z, p, theta[i] );
-                    UpdatePosition( x, y, z, ux[i], uy[i], uz[i], dt);
-                    SetCylindricalPositionFromCartesian( p, theta[i], x, y, z );
-#endif
+                                    ParticleReal x, y, z;
+                                    GetPosition(i, x, y, z);
+                                    UpdatePosition(x, y, z, ux[i], uy[i], uz[i], dt);
+                                    SetPosition(i, x, y, z);
                 }
             );