Move handling of do_not_deposit to physics callers (#3225)

* Rearrange in preparation for lifting do_not_deposit higher in the call stack

* Make do_not_deposit public, so callers can test on it

* Make all DepositCharge call sites respect WarpXParticleContainer::do_not_deposit

* Drop check for do_not_deposit in DepositCharge

* Fix comment to accurately describe np_to_depose

* Fix logic to account for other logic in AddSpaceChargeField

* Match zero charge density array to filled array in RZ PSATD case

* Structure control flow per @dpgrote's preference

* Reword comment for clarity

5 files changed, 75 insertions, 73 deletions

diff --git a/Source/Particles/LaserParticleContainer.cpp b/Source/Particles/LaserParticleContainer.cpp
index 155bc49bc..565ac3ee1 100644
--- a/Source/Particles/LaserParticleContainer.cpp
+++ b/Source/Particles/LaserParticleContainer.cpp
@@ -585,7 +585,7 @@ LaserParticleContainer::Evolve (int lev,
             plane_Yp.resize(np);
             amplitude_E.resize(np);
 
-            if (rho && ! skip_deposition) {
+            if (rho && ! skip_deposition && ! do_not_deposit) {
                 int* AMREX_RESTRICT ion_lev = nullptr;
                 DepositCharge(pti, wp, ion_lev, rho, 0, 0,
                               np_current, thread_num, lev, lev);
@@ -639,7 +639,7 @@ LaserParticleContainer::Evolve (int lev,
             }
 
 
-            if (rho && ! skip_deposition) {
+            if (rho && ! skip_deposition && ! do_not_deposit) {
                 int* AMREX_RESTRICT ion_lev = nullptr;
                 DepositCharge(pti, wp, ion_lev, rho, 1, 0,
                               np_current, thread_num, lev, lev);
diff --git a/Source/Particles/MultiParticleContainer.cpp b/Source/Particles/MultiParticleContainer.cpp
index 34dd488c1..a0b80877d 100644
--- a/Source/Particles/MultiParticleContainer.cpp
+++ b/Source/Particles/MultiParticleContainer.cpp
@@ -486,12 +486,19 @@ MultiParticleContainer::GetZeroChargeDensity (const int lev)
 {
     WarpX& warpx = WarpX::GetInstance();
 
-    BoxArray ba = warpx.boxArray(lev);
+    BoxArray nba = warpx.boxArray(lev);
     DistributionMapping dmap = warpx.DistributionMap(lev);
     const int ng_rho = warpx.get_ng_depos_rho().max();
 
-    auto zero_rho = std::make_unique<MultiFab>(amrex::convert(ba,IntVect::TheNodeVector()),
-                                               dmap,WarpX::ncomps,ng_rho);
+    bool is_PSATD_RZ = false;
+#ifdef WARPX_DIM_RZ
+    if (WarpX::maxwell_solver_id == MaxwellSolverAlgo::PSATD)
+        is_PSATD_RZ = true;
+#endif
+    if( !is_PSATD_RZ )
+        nba.surroundingNodes();
+
+    auto zero_rho = std::make_unique<MultiFab>(nba, dmap, WarpX::ncomps, ng_rho);
     zero_rho->setVal(amrex::Real(0.0));
     return zero_rho;
 }
@@ -540,6 +547,8 @@ MultiParticleContainer::DepositCharge (
     // Call the deposition kernel for each species
     for (auto& pc : allcontainers)
     {
+        if (pc->do_not_deposit) continue;
+
         bool const local = true;
         bool const reset = false;
         bool const do_rz_volume_scaling = false;
@@ -562,24 +571,19 @@ MultiParticleContainer::DepositCharge (
 std::unique_ptr<MultiFab>
 MultiParticleContainer::GetChargeDensity (int lev, bool local)
 {
-    if (allcontainers.empty())
-    {
-        std::unique_ptr<MultiFab> rho = GetZeroChargeDensity(lev);
-        return rho;
+    std::unique_ptr<MultiFab> rho = GetZeroChargeDensity(lev);
+
+    for (unsigned i = 0, n = allcontainers.size(); i < n; ++i) {
+        if (allcontainers[i]->do_not_deposit) continue;
+        std::unique_ptr<MultiFab> rhoi = allcontainers[i]->GetChargeDensity(lev, true);
+        MultiFab::Add(*rho, *rhoi, 0, 0, rho->nComp(), rho->nGrowVect());
     }
-    else
-    {
-        std::unique_ptr<MultiFab> rho = allcontainers[0]->GetChargeDensity(lev, true);
-        for (unsigned i = 1, n = allcontainers.size(); i < n; ++i) {
-            std::unique_ptr<MultiFab> rhoi = allcontainers[i]->GetChargeDensity(lev, true);
-            MultiFab::Add(*rho, *rhoi, 0, 0, rho->nComp(), rho->nGrowVect());
-        }
-        if (!local) {
-            const Geometry& gm = allcontainers[0]->Geom(lev);
-            ablastr::utils::communication::SumBoundary(*rho, WarpX::do_single_precision_comms, gm.periodicity());
-        }
-        return rho;
+    if (!local) {
+        const Geometry& gm = allcontainers[0]->Geom(lev);
+        ablastr::utils::communication::SumBoundary(*rho, WarpX::do_single_precision_comms, gm.periodicity());
     }
+
+    return rho;
 }
 
 void
diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp
index 72543982b..5b0ea8ae5 100644
--- a/Source/Particles/PhysicalParticleContainer.cpp
+++ b/Source/Particles/PhysicalParticleContainer.cpp
@@ -1910,7 +1910,7 @@ PhysicalParticleContainer::Evolve (int lev,
 
             const long np_current = (cjx) ? nfine_current : np;
 
-            if (rho && ! skip_deposition) {
+            if (rho && ! skip_deposition && ! do_not_deposit) {
                 // Deposit charge before particle push, in component 0 of MultiFab rho.
                 int* AMREX_RESTRICT ion_lev;
                 if (do_field_ionization){
@@ -2006,7 +2006,7 @@ PhysicalParticleContainer::Evolve (int lev,
                 } // end of "if do_electrostatic == ElectrostaticSolverAlgo::None"
             } // end of "if do_not_push"
 
-            if (rho && ! skip_deposition) {
+            if (rho && ! skip_deposition && ! do_not_deposit) {
                 // Deposit charge after particle push, in component 1 of MultiFab rho.
                 // (Skipped for electrostatic solver, as this may lead to out-of-bounds)
                 if (WarpX::do_electrostatic == ElectrostaticSolverAlgo::None) {
diff --git a/Source/Particles/WarpXParticleContainer.H b/Source/Particles/WarpXParticleContainer.H
index d6154a2d5..7f9e4e866 100644
--- a/Source/Particles/WarpXParticleContainer.H
+++ b/Source/Particles/WarpXParticleContainer.H
@@ -307,6 +307,7 @@ public:
     void ApplyBoundaryConditions ();
 
     bool do_splitting = false;
+    int do_not_deposit = 0;
     bool initialize_self_fields = false;
     amrex::Real self_fields_required_precision = amrex::Real(1.e-11);
     amrex::Real self_fields_absolute_tolerance = amrex::Real(0.0);
@@ -388,7 +389,6 @@ protected:
     bool m_gather_from_main_grid = false;
 
     int do_not_push = 0;
-    int do_not_deposit = 0;
     int do_not_gather = 0;
 
     // Whether to allow particles outside of the simulation domain to be
diff --git a/Source/Particles/WarpXParticleContainer.cpp b/Source/Particles/WarpXParticleContainer.cpp
index 3421fb9f5..3f0e0bcba 100644
--- a/Source/Particles/WarpXParticleContainer.cpp
+++ b/Source/Particles/WarpXParticleContainer.cpp
@@ -575,7 +575,7 @@ WarpXParticleContainer::DepositCurrent (
                         1: new value (after particle push).
  * \param offset      : Index of first particle for which charge is deposited
  * \param np_to_depose: Number of particles for which charge is deposited.
-                        Particles [offset,offset+np_tp_depose] deposit charge
+                        Particles [offset,offset+np_tp_depose) deposit charge
  * \param thread_num  : Thread number (if tiling)
  * \param lev         : Level of box that contains particles
  * \param depos_lev   : Level on which particles deposit (if buffers are used)
@@ -587,57 +587,55 @@ WarpXParticleContainer::DepositCharge (WarpXParIter& pti, RealVector const& wp,
                                        const long offset, const long np_to_depose,
                                        int thread_num, int lev, int depos_lev)
 {
-    if (!do_not_deposit) {
-        WarpX& warpx = WarpX::GetInstance();
+    WarpX& warpx = WarpX::GetInstance();
 
-        // deposition guards
-        //   note: this is smaller than rho->nGrowVect() for PSATD
-        const amrex::IntVect& ng_rho = warpx.get_ng_depos_rho();
+    // deposition guards
+    //   note: this is smaller than rho->nGrowVect() for PSATD
+    const amrex::IntVect& ng_rho = warpx.get_ng_depos_rho();
 
-        const std::array<amrex::Real,3>& dx = WarpX::CellSize(std::max(depos_lev,0));
-        amrex::IntVect ref_ratio;
-        if (lev == depos_lev) {
-            ref_ratio = IntVect(AMREX_D_DECL(1, 1, 1 ));
-        } else {
-            ref_ratio = WarpX::RefRatio(depos_lev);
-        }
-        const int nc = WarpX::ncomps;
-
-        // Get tile box where charge is deposited.
-        // The tile box is different when depositing in the buffers (depos_lev<lev)
-        // or when depositing inside the level (depos_lev=lev)
-        amrex::Box tilebox;
-        if (lev == depos_lev) {
-            tilebox = pti.tilebox();
-        } else {
-            tilebox = amrex::coarsen(pti.tilebox(), ref_ratio);
-        }
-        tilebox.grow(ng_rho);
-
-        // Lower corner of tile box physical domain
-        // Note that this includes guard cells since it is after tilebox.ngrow
-        // Take into account Galilean shift
-        const amrex::Real dt = warpx.getdt(lev);
-        const amrex::Real time_shift_delta = (icomp == 0 ? 0.0_rt : dt);
-        const std::array<amrex::Real,3>& xyzmin = WarpX::LowerCorner(tilebox, depos_lev, time_shift_delta);
-
-        // pointer to costs data
-        amrex::LayoutData<amrex::Real>* costs = WarpX::getCosts(lev);
-        amrex::Real* cost = costs ? &((*costs)[pti.index()]) : nullptr;
-
-        AMREX_ALWAYS_ASSERT(WarpX::nox == WarpX::noy);
-        AMREX_ALWAYS_ASSERT(WarpX::nox == WarpX::noz);
-
-        ablastr::particles::deposit_charge<WarpXParticleContainer>(
-            pti, wp, this->charge, ion_lev,
-            rho, local_rho[thread_num],
-            WarpX::noz, dx, xyzmin, WarpX::n_rz_azimuthal_modes,
-            ng_rho, depos_lev, ref_ratio,
-            offset, np_to_depose,
-            icomp, nc,
-            cost, WarpX::load_balance_costs_update_algo, WarpX::do_device_synchronize
-        );
+    const std::array<amrex::Real,3>& dx = WarpX::CellSize(std::max(depos_lev,0));
+    amrex::IntVect ref_ratio;
+    if (lev == depos_lev) {
+        ref_ratio = IntVect(AMREX_D_DECL(1, 1, 1 ));
+    } else {
+        ref_ratio = WarpX::RefRatio(depos_lev);
     }
+    const int nc = WarpX::ncomps;
+
+    // Get tile box where charge is deposited.
+    // The tile box is different when depositing in the buffers (depos_lev<lev)
+    // or when depositing inside the level (depos_lev=lev)
+    amrex::Box tilebox;
+    if (lev == depos_lev) {
+        tilebox = pti.tilebox();
+    } else {
+        tilebox = amrex::coarsen(pti.tilebox(), ref_ratio);
+    }
+    tilebox.grow(ng_rho);
+
+    // Lower corner of tile box physical domain
+    // Note that this includes guard cells since it is after tilebox.ngrow
+    // Take into account Galilean shift
+    const amrex::Real dt = warpx.getdt(lev);
+    const amrex::Real time_shift_delta = (icomp == 0 ? 0.0_rt : dt);
+    const std::array<amrex::Real,3>& xyzmin = WarpX::LowerCorner(tilebox, depos_lev, time_shift_delta);
+
+    // pointer to costs data
+    amrex::LayoutData<amrex::Real>* costs = WarpX::getCosts(lev);
+    amrex::Real* cost = costs ? &((*costs)[pti.index()]) : nullptr;
+
+    AMREX_ALWAYS_ASSERT(WarpX::nox == WarpX::noy);
+    AMREX_ALWAYS_ASSERT(WarpX::nox == WarpX::noz);
+
+    ablastr::particles::deposit_charge<WarpXParticleContainer>(
+        pti, wp, this->charge, ion_lev,
+        rho, local_rho[thread_num],
+        WarpX::noz, dx, xyzmin, WarpX::n_rz_azimuthal_modes,
+        ng_rho, depos_lev, ref_ratio,
+        offset, np_to_depose,
+        icomp, nc,
+        cost, WarpX::load_balance_costs_update_algo, WarpX::do_device_synchronize
+    );
 }
 
 void