Converted field gather for RZ multimode to C++

4 files changed, 76 insertions, 27 deletions

diff --git a/Source/Particles/Gather/FieldGather.H b/Source/Particles/Gather/FieldGather.H
index 8f5e8d4cf..219a9d055 100644
--- a/Source/Particles/Gather/FieldGather.H
+++ b/Source/Particles/Gather/FieldGather.H
@@ -2,9 +2,10 @@
 #define FIELDGATHER_H_
 
 #include "ShapeFactors.H"
+#include <WarpX_Complex.H>
 
 /* \brief Field gather for particles handled by thread thread_num
- * /param xp, yp, zp   : Pointer to arrays of particle positions.
+ * \param xp, yp, zp   : Pointer to arrays of particle positions.
  * \param Exp, Eyp, Ezp: Pointer to array of electric field on particles.
  * \param Bxp, Byp, Bzp: Pointer to array of magnetic field on particles.
  * \param ex_arr ey_arr: Array4 of current density, either full array or tile.
@@ -15,6 +16,7 @@
  * \param xyzmin       : Physical lower bounds of domain.
  * \param lo           : Index lower bounds of domain.
  * \param stagger_shift: 0 if nodal, 0.5 if staggered.
+ * \param n_rz_azimuthal_modes: Number of azimuthal modes when using RZ geometry
  */
 template <int depos_order, int lower_in_v>
 void doGatherShapeN(const amrex::Real * const xp,
@@ -33,7 +35,8 @@ void doGatherShapeN(const amrex::Real * const xp,
                     const std::array<amrex::Real, 3>& dx,
                     const std::array<amrex::Real, 3> xyzmin,
                     const amrex::Dim3 lo,
-                    const amrex::Real stagger_shift)
+                    const amrex::Real stagger_shift,
+                    const long n_rz_azimuthal_modes)
 {
     const amrex::Real dxi = 1.0/dx[0];
     const amrex::Real dzi = 1.0/dx[2];
@@ -56,8 +59,8 @@ void doGatherShapeN(const amrex::Real * const xp,
             // x direction
             // Get particle position
 #ifdef WARPX_DIM_RZ
-            const amrex::Real r = std::sqrt(xp[ip]*xp[ip] + yp[ip]*yp[ip]);
-            const amrex::Real x = (r - xmin)*dxi;
+            const amrex::Real rp = std::sqrt(xp[ip]*xp[ip] + yp[ip]*yp[ip]);
+            const amrex::Real x = (rp - xmin)*dxi;
 #else
             const amrex::Real x = (xp[ip]-xmin)*dxi;
 #endif
@@ -103,7 +106,7 @@ void doGatherShapeN(const amrex::Real * const xp,
             for (int iz=0; iz<=depos_order; iz++){
                 for (int ix=0; ix<=depos_order; ix++){
                     Eyp[ip] += sx[ix]*sz[iz]*
-                        ey_arr(lo.x+j+ix, lo.y+l+iz, 0);
+                        ey_arr(lo.x+j+ix, lo.y+l+iz, 0, 0);
                 }
             }
             // Gather field on particle Exp[i] from field on grid ex_arr
@@ -111,9 +114,9 @@ void doGatherShapeN(const amrex::Real * const xp,
             for (int iz=0; iz<=depos_order; iz++){
                 for (int ix=0; ix<=depos_order-lower_in_v; ix++){
                     Exp[ip] += sx0[ix]*sz[iz]*
-                        ex_arr(lo.x+j0+ix, lo.y+l +iz, 0);
+                        ex_arr(lo.x+j0+ix, lo.y+l +iz, 0, 0);
                     Bzp[ip] += sx0[ix]*sz[iz]*
-                        bz_arr(lo.x+j0+ix, lo.y+l +iz, 0);
+                        bz_arr(lo.x+j0+ix, lo.y+l +iz, 0, 0);
                 }
             }
             // Gather field on particle Ezp[i] from field on grid ez_arr
@@ -121,30 +124,79 @@ void doGatherShapeN(const amrex::Real * const xp,
             for (int iz=0; iz<=depos_order-lower_in_v; iz++){
                 for (int ix=0; ix<=depos_order; ix++){
                     Ezp[ip] += sx[ix]*sz0[iz]*
-                        ez_arr(lo.x+j+ix, lo.y+l0 +iz, 0);
+                        ez_arr(lo.x+j+ix, lo.y+l0 +iz, 0, 0);
                     Bxp[ip] += sx[ix]*sz0[iz]*
-                        bx_arr(lo.x+j+ix, lo.y+l0 +iz, 0);
+                        bx_arr(lo.x+j+ix, lo.y+l0 +iz, 0, 0);
                 }
             }
             // Gather field on particle Byp[i] from field on grid by_arr
             for (int iz=0; iz<=depos_order-lower_in_v; iz++){
                 for (int ix=0; ix<=depos_order-lower_in_v; ix++){
                     Byp[ip] += sx0[ix]*sz0[iz]*
-                        by_arr(lo.x+j0+ix, lo.y+l0+iz, 0);
+                        by_arr(lo.x+j0+ix, lo.y+l0+iz, 0, 0);
                 }
             }
 
 #ifdef WARPX_DIM_RZ
-            // Convert Exp and Eyp (which are actually Er and Etheta) to Ex and Ey
+
             amrex::Real costheta;
             amrex::Real sintheta;
-            if (r > 0.) {
-                costheta = xp[ip]/r;
-                sintheta = yp[ip]/r;
+            if (rp > 0.) {
+                costheta = xp[ip]/rp;
+                sintheta = yp[ip]/rp;
             } else {
                 costheta = 1.;
                 sintheta = 0.;
             }
+            const Complex xy0 = Complex{costheta, -sintheta};
+            Complex xy = xy0;
+
+            for (int imode=1 ; imode < n_rz_azimuthal_modes ; imode++) {
+
+                // Gather field on particle Eyp[i] from field on grid ey_arr
+                for (int iz=0; iz<=depos_order; iz++){
+                    for (int ix=0; ix<=depos_order; ix++){
+                        const amrex::Real dEy = (+ ey_arr(lo.x+j+ix, lo.y+l+iz, 0, 2*imode)*xy.real()
+                                                 - ey_arr(lo.x+j+ix, lo.y+l+iz, 0, 2*imode+1)*xy.imag());
+                        Eyp[ip] += sx[ix]*sz[iz]*dEy;
+                    }
+                }
+                // Gather field on particle Exp[i] from field on grid ex_arr
+                // Gather field on particle Bzp[i] from field on grid bz_arr
+                for (int iz=0; iz<=depos_order; iz++){
+                    for (int ix=0; ix<=depos_order-lower_in_v; ix++){
+                        const amrex::Real dEx = (+ ex_arr(lo.x+j0+ix, lo.y+l +iz, 0, 2*imode)*xy.real()
+                                                 - ex_arr(lo.x+j0+ix, lo.y+l +iz, 0, 2*imode+1)*xy.imag());
+                        Exp[ip] += sx0[ix]*sz[iz]*dEx;
+                        const amrex::Real dBz = (+ bz_arr(lo.x+j0+ix, lo.y+l +iz, 0, 2*imode)*xy.real()
+                                                 - bz_arr(lo.x+j0+ix, lo.y+l +iz, 0, 2*imode+1)*xy.imag());
+                        Bzp[ip] += sx0[ix]*sz[iz]*dBz;
+                    }
+                }
+                // Gather field on particle Ezp[i] from field on grid ez_arr
+                // Gather field on particle Bxp[i] from field on grid bx_arr
+                for (int iz=0; iz<=depos_order-lower_in_v; iz++){
+                    for (int ix=0; ix<=depos_order; ix++){
+                        const amrex::Real dEz = (+ ez_arr(lo.x+j+ix, lo.y+l0 +iz, 0, 2*imode)*xy.real()
+                                                 - ez_arr(lo.x+j+ix, lo.y+l0 +iz, 0, 2*imode+1)*xy.imag());
+                        Ezp[ip] += sx[ix]*sz0[iz]*dEz;
+                        const amrex::Real dBx = (+ bx_arr(lo.x+j+ix, lo.y+l0 +iz, 0, 2*imode)*xy.real()
+                                                 - bx_arr(lo.x+j+ix, lo.y+l0 +iz, 0, 2*imode+1)*xy.imag());
+                        Bxp[ip] += sx[ix]*sz0[iz]*dBx;
+                    }
+                }
+                // Gather field on particle Byp[i] from field on grid by_arr
+                for (int iz=0; iz<=depos_order-lower_in_v; iz++){
+                    for (int ix=0; ix<=depos_order-lower_in_v; ix++){
+                        const amrex::Real dBy = (+ by_arr(lo.x+j0+ix, lo.y+l0+iz, 0, 2*imode)*xy.real()
+                                                 - by_arr(lo.x+j0+ix, lo.y+l0+iz, 0, 2*imode+1)*xy.imag());
+                        Byp[ip] += sx0[ix]*sz0[iz]*dBy;
+                    }
+                }
+                xy = xy*xy0;
+            }
+
+            // Convert Exp and Eyp (which are actually Er and Etheta) to Ex and Ey
             const amrex::Real Exp_save = Exp[ip];
             Exp[ip] = costheta*Exp[ip] - sintheta*Eyp[ip];
             Eyp[ip] = costheta*Eyp[ip] + sintheta*Exp_save;
diff --git a/Source/Particles/PhysicalParticleContainer.H b/Source/Particles/PhysicalParticleContainer.H
index ebd4eac2b..b80619733 100644
--- a/Source/Particles/PhysicalParticleContainer.H
+++ b/Source/Particles/PhysicalParticleContainer.H
@@ -53,7 +53,6 @@ public:
                       amrex::FArrayBox const * byfab,
                       amrex::FArrayBox const * bzfab,
                       const int ngE, const int e_is_nodal,
-                      const long n_rz_azimuthal_modes,
                       const long offset,
                       const long np_to_gather,
                       int thread_num,
diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp
index c44a9a8b7..ec360f2c4 100644
--- a/Source/Particles/PhysicalParticleContainer.cpp
+++ b/Source/Particles/PhysicalParticleContainer.cpp
@@ -899,7 +899,7 @@ PhysicalParticleContainer::FieldGather (int lev,
             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, 
-                        Ex.nGrow(), e_is_nodal, WarpX::n_rz_azimuthal_modes,
+                        Ex.nGrow(), e_is_nodal,
                         0, np, thread_num, lev, lev);
 
             if (cost) {
@@ -1176,7 +1176,7 @@ PhysicalParticleContainer::Evolve (int lev,
                 BL_PROFILE_VAR_START(blp_pxr_fg);
                 FieldGather(pti, Exp, Eyp, Ezp, Bxp, Byp, Bzp,
                             exfab, eyfab, ezfab, bxfab, byfab, bzfab, 
-                            Ex.nGrow(), e_is_nodal, WarpX::n_rz_azimuthal_modes,
+                            Ex.nGrow(), e_is_nodal,
                             0, np_gather, thread_num, lev, lev);
 
                 if (np_gather < np)
@@ -1250,7 +1250,6 @@ PhysicalParticleContainer::Evolve (int lev,
                                 cexfab, ceyfab, cezfab,
                                 cbxfab, cbyfab, cbzfab,
                                 cEx->nGrow(), e_is_nodal, 
-                                WarpX::n_rz_azimuthal_modes,
                                 nfine_gather, np-nfine_gather, 
                                 thread_num, lev, lev-1);
                 }
@@ -1594,7 +1593,7 @@ PhysicalParticleContainer::PushP (int lev, Real dt,
             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, 
-                        Ex.nGrow(), e_is_nodal, WarpX::n_rz_azimuthal_modes,
+                        Ex.nGrow(), e_is_nodal,
                         0, np, thread_num, lev, lev);
 
             // This wraps the momentum advance so that inheritors can modify the call.
@@ -1835,7 +1834,6 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                         const int ngE, const int e_is_nodal,
                                         const long offset,
                                         const long np_to_gather,
-                                        const long n_rz_azimuthal_modes,
                                         int thread_num,
                                         int lev,
                                         int gather_lev)
@@ -1890,7 +1888,7 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                 Byp.dataPtr() + offset, Bzp.dataPtr() + offset,
                                 ex_arr, ey_arr, ez_arr, bx_arr, by_arr, bz_arr,
                                 np_to_gather, dx,
-                                xyzmin, lo, stagger_shift);
+                                xyzmin, lo, stagger_shift, WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 2){
             doGatherShapeN<2,1>(xp, yp, zp,
                                 Exp.dataPtr() + offset, Eyp.dataPtr() + offset,
@@ -1898,7 +1896,7 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                 Byp.dataPtr() + offset, Bzp.dataPtr() + offset,
                                 ex_arr, ey_arr, ez_arr, bx_arr, by_arr, bz_arr,
                                 np_to_gather, dx,
-                                xyzmin, lo, stagger_shift);
+                                xyzmin, lo, stagger_shift, WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 3){
             doGatherShapeN<3,1>(xp, yp, zp,
                                 Exp.dataPtr() + offset, Eyp.dataPtr() + offset,
@@ -1906,7 +1904,7 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                 Byp.dataPtr() + offset, Bzp.dataPtr() + offset,
                                 ex_arr, ey_arr, ez_arr, bx_arr, by_arr, bz_arr,
                                 np_to_gather, dx,
-                                xyzmin, lo, stagger_shift);
+                                xyzmin, lo, stagger_shift, WarpX::n_rz_azimuthal_modes);
         }
     } else {
         if        (WarpX::nox == 1){
@@ -1916,7 +1914,7 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                 Byp.dataPtr() + offset, Bzp.dataPtr() + offset,
                                 ex_arr, ey_arr, ez_arr, bx_arr, by_arr, bz_arr,
                                 np_to_gather, dx,
-                                xyzmin, lo, stagger_shift);
+                                xyzmin, lo, stagger_shift, WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 2){
             doGatherShapeN<2,0>(xp, yp, zp,
                                 Exp.dataPtr() + offset, Eyp.dataPtr() + offset,
@@ -1924,7 +1922,7 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                 Byp.dataPtr() + offset, Bzp.dataPtr() + offset,
                                 ex_arr, ey_arr, ez_arr, bx_arr, by_arr, bz_arr,
                                 np_to_gather, dx,
-                                xyzmin, lo, stagger_shift);
+                                xyzmin, lo, stagger_shift, WarpX::n_rz_azimuthal_modes);
         } else if (WarpX::nox == 3){
             doGatherShapeN<3,0>(xp, yp, zp,
                                 Exp.dataPtr() + offset, Eyp.dataPtr() + offset,
@@ -1932,7 +1930,7 @@ PhysicalParticleContainer::FieldGather (WarpXParIter& pti,
                                 Byp.dataPtr() + offset, Bzp.dataPtr() + offset,
                                 ex_arr, ey_arr, ez_arr, bx_arr, by_arr, bz_arr,
                                 np_to_gather, dx,
-                                xyzmin, lo, stagger_shift);
+                                xyzmin, lo, stagger_shift, WarpX::n_rz_azimuthal_modes);
         }
     }
 }
diff --git a/Source/Particles/RigidInjectedParticleContainer.cpp b/Source/Particles/RigidInjectedParticleContainer.cpp
index 695955e15..a87168a75 100644
--- a/Source/Particles/RigidInjectedParticleContainer.cpp
+++ b/Source/Particles/RigidInjectedParticleContainer.cpp
@@ -426,7 +426,7 @@ RigidInjectedParticleContainer::PushP (int lev, Real dt,
             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,
-                        Ex.nGrow(), e_is_nodal, WarpX::n_rz_azimuthal_modes,
+                        Ex.nGrow(), e_is_nodal,
                         0, np, thread_num, lev, lev);
 
             // Save the position and momenta, making copies