`SyncRho`: Pass References To Charge MultiFabs (#3814)

* `SyncRho`: Pass References To Charge MultiFabs

* Keep a function SyncRho for the Python interface

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

5 files changed, 26 insertions, 15 deletions

diff --git a/Source/Evolve/WarpXEvolve.cpp b/Source/Evolve/WarpXEvolve.cpp
index 94379ad3d..ddf052d35 100644
--- a/Source/Evolve/WarpXEvolve.cpp
+++ b/Source/Evolve/WarpXEvolve.cpp
@@ -505,12 +505,12 @@ void WarpX::SyncCurrentAndRho ()
             {
                 // TODO Replace current_cp with current_cp_vay once Vay deposition is implemented with MR
                 SyncCurrent(current_fp_vay, current_cp);
-                SyncRho();
+                SyncRho(rho_fp, rho_cp);
             }
             else
             {
                 SyncCurrent(current_fp, current_cp);
-                SyncRho();
+                SyncRho(rho_fp, rho_cp);
             }
         }
         else // no periodic single box
@@ -522,7 +522,7 @@ void WarpX::SyncCurrentAndRho ()
                 current_deposition_algo != CurrentDepositionAlgo::Vay)
             {
                 SyncCurrent(current_fp, current_cp);
-                SyncRho();
+                SyncRho(rho_fp, rho_cp);
             }
 
             if (current_deposition_algo == CurrentDepositionAlgo::Vay)
@@ -536,7 +536,7 @@ void WarpX::SyncCurrentAndRho ()
     else // FDTD
     {
         SyncCurrent(current_fp, current_cp);
-        SyncRho();
+        SyncRho(rho_fp, rho_cp);
     }
 }
 
@@ -576,7 +576,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();
+        SyncRho(rho_fp, rho_cp);
         // Forward FFT of rho
         PSATDForwardTransformRho(rho_fp, rho_cp, 0, rho_new);
     }
@@ -640,7 +640,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();
+            SyncRho(rho_fp, rho_cp);
             // 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);
diff --git a/Source/FieldSolver/ElectrostaticSolver.cpp b/Source/FieldSolver/ElectrostaticSolver.cpp
index 42a98dab1..e43c98675 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(); // Apply filter, perform MPI exchange, interpolate across levels
+    SyncRho(rho_fp, rho_cp); // Apply filter, perform MPI exchange, interpolate across levels
 
     // beta is zero in lab frame
     // Todo: use simpler finite difference form with beta=0
diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index fbc189460..f633422b2 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -724,7 +724,7 @@ WarpX::PushPSATD ()
 
             // Synchronize J and rho
             SyncCurrent(current_fp, current_cp);
-            SyncRho();
+            SyncRho(rho_fp, rho_cp);
         }
         else if (current_deposition_algo == CurrentDepositionAlgo::Vay)
         {
@@ -745,7 +745,7 @@ WarpX::PushPSATD ()
             // TODO This works only without mesh refinement
             const int lev = 0;
             SumBoundaryJ(current_fp, lev, Geom(lev).periodicity());
-            SyncRho();
+            SyncRho(rho_fp, rho_cp);
         }
 
         // FFT of J and rho (if used)
diff --git a/Source/Parallelization/WarpXComm.cpp b/Source/Parallelization/WarpXComm.cpp
index 8585da632..a2b25cbf4 100644
--- a/Source/Parallelization/WarpXComm.cpp
+++ b/Source/Parallelization/WarpXComm.cpp
@@ -902,20 +902,27 @@ WarpX::SyncCurrent (
 }
 
 void
-WarpX::SyncRho ()
+WarpX::SyncRho () {
+    SyncRho(rho_fp, rho_cp);
+}
+
+void
+WarpX::SyncRho (
+    const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
+    const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp)
 {
     WARPX_PROFILE("WarpX::SyncRho()");
 
-    if (!rho_fp[0]) return;
-    const int ncomp = rho_fp[0]->nComp();
+    if (!charge_fp[0]) return;
+    const int ncomp = charge_fp[0]->nComp();
 
     // Restrict fine patch onto the coarse patch,
     // before summing the guard cells of the fine patch
     for (int lev = 1; lev <= finest_level; ++lev)
     {
-        rho_cp[lev]->setVal(0.0);
+        charge_cp[lev]->setVal(0.0);
         const IntVect& refinement_ratio = refRatio(lev-1);
-        ablastr::coarsen::average::Coarsen(*rho_cp[lev], *rho_fp[lev], refinement_ratio );
+        ablastr::coarsen::average::Coarsen(*charge_cp[lev], *charge_fp[lev], refinement_ratio );
     }
 
     // For each level
@@ -923,7 +930,7 @@ WarpX::SyncRho ()
     // - add the coarse patch/buffer of `lev+1` into the fine patch of `lev`
     // - sum guard cells of the coarse patch of `lev+1` and fine patch of `lev`
     for (int lev=0; lev <= finest_level; ++lev) {
-        AddRhoFromFineLevelandSumBoundary(rho_fp, rho_cp, lev, 0, ncomp);
+        AddRhoFromFineLevelandSumBoundary(charge_fp, charge_cp, lev, 0, ncomp);
     }
 }
 
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 07bd268d8..a45d27583 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -778,6 +778,10 @@ public:
 
     void SyncRho ();
 
+    void SyncRho (
+        const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
+        const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp);
+
     amrex::Vector<int> getnsubsteps () const {return nsubsteps;}
     int getnsubsteps (int lev) const {return nsubsteps[lev];}
     amrex::Vector<int> getistep () const {return istep;}