Generalize buffers for `SyncRho`, `SyncCurrent`, and related functions (#3995)

8 files changed, 69 insertions, 58 deletions

diff --git a/Source/Evolve/WarpXEvolve.cpp b/Source/Evolve/WarpXEvolve.cpp
index 632b97b7f..b2201eb2a 100644
--- a/Source/Evolve/WarpXEvolve.cpp
+++ b/Source/Evolve/WarpXEvolve.cpp
@@ -173,8 +173,8 @@ WarpX::Evolve (int numsteps)
             auto& current_fp_temp = m_hybrid_pic_model->current_fp_temp;
             mypc->DepositCharge(rho_fp_temp, 0._rt);
             mypc->DepositCurrent(current_fp_temp, dt[0], 0._rt);
-            SyncRho(rho_fp_temp, rho_cp);
-            SyncCurrent(current_fp_temp, current_cp);
+            SyncRho(rho_fp_temp, rho_cp, charge_buf);
+            SyncCurrent(current_fp_temp, current_cp, current_buf);
             for (int lev=0; lev <= finest_level; ++lev) {
                 // SyncCurrent does not include a call to FillBoundary, but it is needed
                 // for the hybrid-PIC solver since current values are interpolated to
@@ -556,13 +556,13 @@ void WarpX::SyncCurrentAndRho ()
             if (current_deposition_algo == CurrentDepositionAlgo::Vay)
             {
                 // TODO Replace current_cp with current_cp_vay once Vay deposition is implemented with MR
-                SyncCurrent(current_fp_vay, current_cp);
-                SyncRho(rho_fp, rho_cp);
+                SyncCurrent(current_fp_vay, current_cp, current_buf);
+                SyncRho(rho_fp, rho_cp, charge_buf);
             }
             else
             {
-                SyncCurrent(current_fp, current_cp);
-                SyncRho(rho_fp, rho_cp);
+                SyncCurrent(current_fp, current_cp, current_buf);
+                SyncRho(rho_fp, rho_cp, charge_buf);
             }
         }
         else // no periodic single box
@@ -573,8 +573,8 @@ void WarpX::SyncCurrentAndRho ()
             if (current_correction == false &&
                 current_deposition_algo != CurrentDepositionAlgo::Vay)
             {
-                SyncCurrent(current_fp, current_cp);
-                SyncRho(rho_fp, rho_cp);
+                SyncCurrent(current_fp, current_cp, current_buf);
+                SyncRho(rho_fp, rho_cp, charge_buf);
             }
 
             if (current_deposition_algo == CurrentDepositionAlgo::Vay)
@@ -587,8 +587,8 @@ void WarpX::SyncCurrentAndRho ()
     }
     else // FDTD
     {
-        SyncCurrent(current_fp, current_cp);
-        SyncRho(rho_fp, rho_cp);
+        SyncCurrent(current_fp, current_cp, current_buf);
+        SyncRho(rho_fp, rho_cp, charge_buf);
     }
 
     // Reflect charge and current density over PEC boundaries, if needed.
@@ -649,7 +649,7 @@ WarpX::OneStep_multiJ (const amrex::Real cur_time)
         // (dt[0] denotes the time step on mesh refinement level 0)
         mypc->DepositCharge(rho_fp, -dt[0]);
         // Filter, exchange boundary, and interpolate across levels
-        SyncRho(rho_fp, rho_cp);
+        SyncRho(rho_fp, rho_cp, charge_buf);
         // Forward FFT of rho
         PSATDForwardTransformRho(rho_fp, rho_cp, 0, rho_new);
     }
@@ -665,7 +665,7 @@ WarpX::OneStep_multiJ (const amrex::Real cur_time)
         // namely 'current_fp_nodal': SyncCurrent stores the result of its centering
         // into 'current_fp' and then performs both filtering, if used, and exchange
         // of guard cells.
-        SyncCurrent(current_fp, current_cp);
+        SyncCurrent(current_fp, current_cp, current_buf);
         // Forward FFT of J
         PSATDForwardTransformJ(current_fp, current_cp);
     }
@@ -699,7 +699,7 @@ WarpX::OneStep_multiJ (const amrex::Real cur_time)
         // namely 'current_fp_nodal': SyncCurrent stores the result of its centering
         // into 'current_fp' and then performs both filtering, if used, and exchange
         // of guard cells.
-        SyncCurrent(current_fp, current_cp);
+        SyncCurrent(current_fp, current_cp, current_buf);
         // Forward FFT of J
         PSATDForwardTransformJ(current_fp, current_cp);
 
@@ -713,7 +713,7 @@ WarpX::OneStep_multiJ (const amrex::Real cur_time)
             // Deposit rho at relative time t_depose_charge
             mypc->DepositCharge(rho_fp, t_depose_charge);
             // Filter, exchange boundary, and interpolate across levels
-            SyncRho(rho_fp, rho_cp);
+            SyncRho(rho_fp, rho_cp, charge_buf);
             // Forward FFT of rho
             const int rho_idx = (rho_in_time == RhoInTime::Linear) ? rho_new : rho_mid;
             PSATDForwardTransformRho(rho_fp, rho_cp, 0, rho_idx);
@@ -848,8 +848,8 @@ WarpX::OneStep_sub1 (Real curtime)
     // by only half a coarse step (first half)
     PushParticlesandDepose(coarse_lev, curtime, DtType::Full);
     StoreCurrent(coarse_lev);
-    AddCurrentFromFineLevelandSumBoundary(current_fp, current_cp, coarse_lev);
-    AddRhoFromFineLevelandSumBoundary(rho_fp, rho_cp, coarse_lev, 0, ncomps);
+    AddCurrentFromFineLevelandSumBoundary(current_fp, current_cp, current_buf, coarse_lev);
+    AddRhoFromFineLevelandSumBoundary(rho_fp, rho_cp, charge_buf, coarse_lev, 0, ncomps);
 
     EvolveB(fine_lev, PatchType::coarse, dt[fine_lev], DtType::FirstHalf);
     EvolveF(fine_lev, PatchType::coarse, dt[fine_lev], DtType::FirstHalf);
@@ -907,8 +907,8 @@ WarpX::OneStep_sub1 (Real curtime)
     // v) Push the fields on the coarse patch and mother grid
     // by only half a coarse step (second half)
     RestoreCurrent(coarse_lev);
-    AddCurrentFromFineLevelandSumBoundary(current_fp, current_cp, coarse_lev);
-    AddRhoFromFineLevelandSumBoundary(rho_fp, rho_cp, coarse_lev, ncomps, ncomps);
+    AddCurrentFromFineLevelandSumBoundary(current_fp, current_cp, current_buf, coarse_lev);
+    AddRhoFromFineLevelandSumBoundary(rho_fp, rho_cp, charge_buf, coarse_lev, ncomps, ncomps);
 
     EvolveE(fine_lev, PatchType::coarse, dt[fine_lev]);
     FillBoundaryE(fine_lev, PatchType::coarse, guard_cells.ng_FieldSolver,
diff --git a/Source/FieldSolver/ElectrostaticSolver.cpp b/Source/FieldSolver/ElectrostaticSolver.cpp
index 9fea4974a..6089cdbf2 100644
--- a/Source/FieldSolver/ElectrostaticSolver.cpp
+++ b/Source/FieldSolver/ElectrostaticSolver.cpp
@@ -207,7 +207,7 @@ WarpX::AddSpaceChargeFieldLabFrame ()
     // Deposit particle charge density (source of Poisson solver)
     mypc->DepositCharge(rho_fp, 0.0_rt);
 
-    SyncRho(rho_fp, rho_cp); // Apply filter, perform MPI exchange, interpolate across levels
+    SyncRho(rho_fp, rho_cp, charge_buf); // Apply filter, perform MPI exchange, interpolate across levels
 #ifndef WARPX_DIM_RZ
     for (int lev = 0; lev <= finestLevel(); lev++) {
         // Reflect density over PEC boundaries, if needed.
diff --git a/Source/FieldSolver/MagnetostaticSolver/MagnetostaticSolver.cpp b/Source/FieldSolver/MagnetostaticSolver/MagnetostaticSolver.cpp
index 9f91b2558..7bc4d4377 100644
--- a/Source/FieldSolver/MagnetostaticSolver/MagnetostaticSolver.cpp
+++ b/Source/FieldSolver/MagnetostaticSolver/MagnetostaticSolver.cpp
@@ -106,7 +106,7 @@ WarpX::AddMagnetostaticFieldLabFrame()
     }
 #endif
 
-    SyncCurrent(current_fp, current_cp); // Apply filter, perform MPI exchange, interpolate across levels
+    SyncCurrent(current_fp, current_cp, current_buf); // Apply filter, perform MPI exchange, interpolate across levels
 
     // set the boundary and current density potentials
     setVectorPotentialBC(vector_potential_fp_nodal);
diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index 6e39c850b..d515462d7 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -724,8 +724,8 @@ WarpX::PushPSATD ()
             PSATDBackwardTransformJ(current_fp, current_cp);
 
             // Synchronize J and rho
-            SyncCurrent(current_fp, current_cp);
-            SyncRho(rho_fp, rho_cp);
+            SyncCurrent(current_fp, current_cp, current_buf);
+            SyncRho(rho_fp, rho_cp, charge_buf);
         }
         else if (current_deposition_algo == CurrentDepositionAlgo::Vay)
         {
@@ -746,7 +746,7 @@ WarpX::PushPSATD ()
             // TODO This works only without mesh refinement
             const int lev = 0;
             SumBoundaryJ(current_fp, lev, Geom(lev).periodicity());
-            SyncRho(rho_fp, rho_cp);
+            SyncRho(rho_fp, rho_cp, charge_buf);
         }
 
         // FFT of J and rho (if used)
diff --git a/Source/Parallelization/WarpXComm.cpp b/Source/Parallelization/WarpXComm.cpp
index d0567fd39..9b7f16478 100644
--- a/Source/Parallelization/WarpXComm.cpp
+++ b/Source/Parallelization/WarpXComm.cpp
@@ -856,7 +856,8 @@ WarpX::FillBoundaryAux (int lev, IntVect ng)
 void
 WarpX::SyncCurrent (
     const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
-    const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp)
+    const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+    const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_buffer)
 {
     WARPX_PROFILE("WarpX::SyncCurrent()");
 
@@ -989,14 +990,14 @@ WarpX::SyncCurrent (
                 ablastr::coarsen::average::Coarsen(*J_cp[lev][idim],
                                                    *J_fp[lev][idim],
                                                    refRatio(lev-1));
-                if (current_buf[lev][idim])
+                if (J_buffer[lev][idim])
                 {
                     IntVect const& ng = J_cp[lev][idim]->nGrowVect();
-                    AMREX_ASSERT(ng.allLE(current_buf[lev][idim]->nGrowVect()));
-                    MultiFab::Add(*current_buf[lev][idim], *J_cp[lev][idim],
+                    AMREX_ASSERT(ng.allLE(J_buffer[lev][idim]->nGrowVect()));
+                    MultiFab::Add(*J_buffer[lev][idim], *J_cp[lev][idim],
                                   0, 0, ncomp, ng);
                     mf_comm = std::make_unique<MultiFab>
-                        (*current_buf[lev][idim], amrex::make_alias, 0, ncomp);
+                        (*J_buffer[lev][idim], amrex::make_alias, 0, ncomp);
                 }
                 else
                 {
@@ -1016,13 +1017,14 @@ WarpX::SyncCurrent (
 
 void
 WarpX::SyncRho () {
-    SyncRho(rho_fp, rho_cp);
+    SyncRho(rho_fp, rho_cp, charge_buf);
 }
 
 void
 WarpX::SyncRho (
     const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
-    const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp)
+    const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp,
+    const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_buffer)
 {
     WARPX_PROFILE("WarpX::SyncRho()");
 
@@ -1074,13 +1076,13 @@ WarpX::SyncRho (
             ablastr::coarsen::average::Coarsen(*charge_cp[lev],
                                                *charge_fp[lev],
                                                refRatio(lev-1));
-            if (charge_buf[lev])
+            if (charge_buffer[lev])
             {
                 IntVect const& ng = charge_cp[lev]->nGrowVect();
-                AMREX_ASSERT(ng.allLE(charge_buf[lev]->nGrowVect()));
-                MultiFab::Add(*charge_buf[lev], *charge_cp[lev], 0, 0, ncomp, ng);
+                AMREX_ASSERT(ng.allLE(charge_buffer[lev]->nGrowVect()));
+                MultiFab::Add(*charge_buffer[lev], *charge_cp[lev], 0, 0, ncomp, ng);
                 mf_comm = std::make_unique<MultiFab>
-                    (*charge_buf[lev], amrex::make_alias, 0, ncomp);
+                    (*charge_buffer[lev], amrex::make_alias, 0, ncomp);
             }
             else
             {
@@ -1210,6 +1212,7 @@ void WarpX::SumBoundaryJ (
 void WarpX::AddCurrentFromFineLevelandSumBoundary (
     const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
     const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+    const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_buffer,
     const int lev)
 {
     const amrex::Periodicity& period = Geom(lev).periodicity();
@@ -1233,18 +1236,18 @@ void WarpX::AddCurrentFromFineLevelandSumBoundary (
 
             const IntVect ng = J_cp[lev+1][idim]->nGrowVect();
 
-            if (use_filter && current_buf[lev+1][idim])
+            if (use_filter && J_buffer[lev+1][idim])
             {
                 ApplyFilterJ(J_cp, lev+1, idim);
-                ApplyFilterJ(current_buf, lev+1, idim);
+                ApplyFilterJ(J_buffer, lev+1, idim);
 
                 MultiFab::Add(
-                    *current_buf[lev+1][idim], *J_cp[lev+1][idim],
-                    0, 0, current_buf[lev+1][idim]->nComp(), ng);
+                    *J_buffer[lev+1][idim], *J_cp[lev+1][idim],
+                    0, 0, J_buffer[lev+1][idim]->nComp(), ng);
 
                 ablastr::utils::communication::ParallelAdd(
-                    mf, *current_buf[lev+1][idim], 0, 0,
-                    current_buf[lev+1][idim]->nComp(),
+                    mf, *J_buffer[lev+1][idim], 0, 0,
+                    J_buffer[lev+1][idim]->nComp(),
                     ng, amrex::IntVect(0),
                     do_single_precision_comms, period);
             }
@@ -1258,15 +1261,15 @@ void WarpX::AddCurrentFromFineLevelandSumBoundary (
                     ng, amrex::IntVect(0),
                     do_single_precision_comms, period);
             }
-            else if (current_buf[lev+1][idim]) // but no filter
+            else if (J_buffer[lev+1][idim]) // but no filter
             {
                 MultiFab::Add(
-                    *current_buf[lev+1][idim], *J_cp[lev+1][idim],
-                    0, 0, current_buf[lev+1][idim]->nComp(), ng);
+                    *J_buffer[lev+1][idim], *J_cp[lev+1][idim],
+                    0, 0, J_buffer[lev+1][idim]->nComp(), ng);
 
                 ablastr::utils::communication::ParallelAdd(
-                    mf, *current_buf[lev+1][idim], 0, 0,
-                    current_buf[lev+1][idim]->nComp(),
+                    mf, *J_buffer[lev+1][idim], 0, 0,
+                    J_buffer[lev+1][idim]->nComp(),
                     ng, amrex::IntVect(0),
                     do_single_precision_comms, period);
             }
@@ -1345,6 +1348,7 @@ void WarpX::ApplyFilterandSumBoundaryRho (int /*lev*/, int glev, amrex::MultiFab
 void WarpX::AddRhoFromFineLevelandSumBoundary (
     const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
     const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp,
+    const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_buffer,
     const int lev,
     const int icomp,
     const int ncomp)
@@ -1362,7 +1366,7 @@ void WarpX::AddRhoFromFineLevelandSumBoundary (
         mf.setVal(0.0);
         IntVect ng = charge_cp[lev+1]->nGrowVect();
         IntVect ng_depos_rho = get_ng_depos_rho();
-        if (use_filter && charge_buf[lev+1])
+        if (use_filter && charge_buffer[lev+1])
         {
             // coarse patch of fine level
             ng += bilinear_filter.stencil_length_each_dir-1;
@@ -1373,9 +1377,9 @@ void WarpX::AddRhoFromFineLevelandSumBoundary (
             bilinear_filter.ApplyStencil(rhofc, *charge_cp[lev+1], lev+1, icomp, 0, ncomp);
 
             // buffer patch of fine level
-            MultiFab rhofb(charge_buf[lev+1]->boxArray(),
-                           charge_buf[lev+1]->DistributionMap(), ncomp, ng);
-            bilinear_filter.ApplyStencil(rhofb, *charge_buf[lev+1], lev+1, icomp, 0, ncomp);
+            MultiFab rhofb(charge_buffer[lev+1]->boxArray(),
+                           charge_buffer[lev+1]->DistributionMap(), ncomp, ng);
+            bilinear_filter.ApplyStencil(rhofb, *charge_buffer[lev+1], lev+1, icomp, 0, ncomp);
 
             MultiFab::Add(rhofb, rhofc, 0, 0, ncomp, ng);
 
@@ -1395,16 +1399,16 @@ void WarpX::AddRhoFromFineLevelandSumBoundary (
                                                        WarpX::do_single_precision_comms, period);
             WarpXSumGuardCells( *charge_cp[lev+1], rf, period, ng_depos_rho, icomp, ncomp );
         }
-        else if (charge_buf[lev+1]) // but no filter
+        else if (charge_buffer[lev+1]) // but no filter
         {
             ng_depos_rho.min(ng);
-            MultiFab::Add(*charge_buf[lev+1],
+            MultiFab::Add(*charge_buffer[lev+1],
                           *charge_cp[lev+1], icomp, icomp, ncomp,
                            charge_cp[lev+1]->nGrowVect());
 
-            ablastr::utils::communication::ParallelAdd(mf, *charge_buf[lev + 1], icomp, 0,
+            ablastr::utils::communication::ParallelAdd(mf, *charge_buffer[lev + 1], icomp, 0,
                                                        ncomp,
-                                                       charge_buf[lev + 1]->nGrowVect(),
+                                                       charge_buffer[lev + 1]->nGrowVect(),
                                                        IntVect::TheZeroVector(), WarpX::do_single_precision_comms,
                                                        period);
             WarpXSumGuardCells(*(charge_cp[lev+1]), period, ng_depos_rho, icomp, ncomp);
diff --git a/Source/Python/WarpXWrappers.H b/Source/Python/WarpXWrappers.H
index f7d805387..65e2b796a 100644
--- a/Source/Python/WarpXWrappers.H
+++ b/Source/Python/WarpXWrappers.H
@@ -138,7 +138,8 @@ extern "C" {
   void warpx_SyncRho ();
   void warpx_SyncCurrent (
       const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
-      const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp);
+      const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+      const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_buffer);
   void warpx_UpdateAuxilaryData ();
   void warpx_PushParticlesandDepose (amrex::Real cur_time);
 
diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp
index 4f32f4ed8..afc4caa5b 100644
--- a/Source/Python/WarpXWrappers.cpp
+++ b/Source/Python/WarpXWrappers.cpp
@@ -710,9 +710,10 @@ namespace
     }
     void warpx_SyncCurrent (
         const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
-        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp) {
+        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_buffer) {
         WarpX& warpx = WarpX::GetInstance();
-        warpx.SyncCurrent(J_fp, J_cp);
+        warpx.SyncCurrent(J_fp, J_cp, J_buffer);
     }
     void warpx_UpdateAuxilaryData () {
         WarpX& warpx = WarpX::GetInstance();
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 213a5b5b2..3a6dd70fd 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -825,16 +825,19 @@ public:
      *
      * \param[in,out] J_fp reference to fine-patch current \c MultiFab (all MR levels)
      * \param[in,out] J_cp reference to coarse-patch current \c MultiFab (all MR levels)
+     * \param[in,out] J_buffer reference to buffer current \c MultiFab (all MR levels)
      */
     void SyncCurrent (
         const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
-        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp);
+        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_buffer);
 
     void SyncRho ();
 
     void SyncRho (
         const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
-        const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp);
+        const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp,
+        const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_buffer);
 
     amrex::Vector<int> getnsubsteps () const {return nsubsteps;}
     int getnsubsteps (int lev) const {return nsubsteps[lev];}
@@ -1157,6 +1160,7 @@ private:
     void AddCurrentFromFineLevelandSumBoundary (
         const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
         const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+        const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_buffer,
         const int lev);
     void StoreCurrent (const int lev);
     void RestoreCurrent (const int lev);
@@ -1196,6 +1200,7 @@ private:
     void AddRhoFromFineLevelandSumBoundary (
         const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
         const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp,
+        const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_buffer,
         const int lev,
         const int icomp,
         const int ncomp);