3 files changed, 50 insertions, 8 deletions

diff --git a/Source/Diagnostics/ComputeDiagFunctors/RhoFunctor.cpp b/Source/Diagnostics/ComputeDiagFunctors/RhoFunctor.cpp
index c8ccb2ded..e48703e91 100644
--- a/Source/Diagnostics/ComputeDiagFunctors/RhoFunctor.cpp
+++ b/Source/Diagnostics/ComputeDiagFunctors/RhoFunctor.cpp
@@ -2,6 +2,10 @@
 #include "RhoFunctor.H"
 #include "Utils/CoarsenIO.H"
 
+#ifdef WARPX_DIM_RZ
+#include "FieldSolver/SpectralSolver/SpectralFieldData.H"
+#endif
+
 #include <AMReX.H>
 
 RhoFunctor::RhoFunctor (const int lev,
@@ -21,17 +25,36 @@ RhoFunctor::operator() ( amrex::MultiFab& mf_dst, const int dcomp, const int /*i
     auto& warpx = WarpX::GetInstance();
     std::unique_ptr<amrex::MultiFab> rho;
 
+    // Deposit charge density
+    // Call this with local=true since the parallel transfers will be handled
+    // by ApplyFilterandSumBoundaryRho
+
     // Dump total rho
     if (m_species_index == -1) {
         auto& mypc = warpx.GetPartContainer();
-        rho = mypc.GetChargeDensity(m_lev);
+        rho = mypc.GetChargeDensity(m_lev, true);
     }
     // Dump rho per species
     else {
         auto& mypc = warpx.GetPartContainer().GetParticleContainer(m_species_index);
-        rho = mypc.GetChargeDensity(m_lev);
+        rho = mypc.GetChargeDensity(m_lev, true);
     }
 
+    // Handle the parallel transfers of guard cells and
+    // apply the filtering if requested.
+    warpx.ApplyFilterandSumBoundaryRho(m_lev, m_lev, *rho, 0, rho->nComp());
+
+#if (defined WARPX_DIM_RZ) && (defined WARPX_USE_PSATD)
+    using Idx = SpectralAvgFieldIndex;
+    if (WarpX::use_kspace_filter) {
+        auto & solver = warpx.get_spectral_solver_fp(m_lev);
+        solver.ForwardTransform(*rho, Idx::rho_new);
+        solver.ApplyFilter(Idx::rho_new);
+        solver.BackwardTransform(*rho, Idx::rho_new);
+    }
+#endif
+
+
 #ifdef WARPX_DIM_RZ
     if (m_convertRZmodes2cartesian) {
         // In cylindrical geometry, sum real part of all modes of rho in
diff --git a/Source/Parallelization/WarpXComm.cpp b/Source/Parallelization/WarpXComm.cpp
index 667be0696..daee24054 100644
--- a/Source/Parallelization/WarpXComm.cpp
+++ b/Source/Parallelization/WarpXComm.cpp
@@ -873,17 +873,23 @@ void
 WarpX::ApplyFilterandSumBoundaryRho (int lev, PatchType patch_type, int icomp, int ncomp)
 {
     const int glev = (patch_type == PatchType::fine) ? lev : lev-1;
-    const auto& period = Geom(glev).periodicity();
     auto& r = (patch_type == PatchType::fine) ? rho_fp[lev] : rho_cp[lev];
     if (r == nullptr) return;
+    ApplyFilterandSumBoundaryRho(lev, glev, *r, icomp, ncomp);
+}
+
+void
+WarpX::ApplyFilterandSumBoundaryRho (int lev, int glev, amrex::MultiFab& rho, int icomp, int ncomp)
+{
+    const auto& period = Geom(glev).periodicity();
     if (use_filter) {
-        IntVect ng = r->nGrowVect();
+        IntVect ng = rho.nGrowVect();
         ng += bilinear_filter.stencil_length_each_dir-1;
-        MultiFab rf(r->boxArray(), r->DistributionMap(), ncomp, ng);
-        bilinear_filter.ApplyStencil(rf, *r, icomp, 0, ncomp);
-        WarpXSumGuardCells(*r, rf, period, icomp, ncomp );
+        MultiFab rf(rho.boxArray(), rho.DistributionMap(), ncomp, ng);
+        bilinear_filter.ApplyStencil(rf, rho, icomp, 0, ncomp);
+        WarpXSumGuardCells(rho, rf, period, icomp, ncomp );
     } else {
-        WarpXSumGuardCells(*r, period, icomp, ncomp);
+        WarpXSumGuardCells(rho, period, icomp, ncomp);
     }
 }
 
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 68dd45949..928e40b42 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -543,6 +543,8 @@ public:
      */
     void ComputeCostsHeuristic (amrex::Vector<std::unique_ptr<amrex::LayoutData<amrex::Real> > >& costs);
 
+    void ApplyFilterandSumBoundaryRho (int lev, int glev, amrex::MultiFab& rho, int icomp, int ncomp);
+
 protected:
 
     /**
@@ -877,7 +879,18 @@ private:
         amrex::Vector<std::unique_ptr<SpectralSolver>> spectral_solver_fp;
         amrex::Vector<std::unique_ptr<SpectralSolver>> spectral_solver_cp;
 #   endif
+
+public:
+
+#   ifdef WARPX_DIM_RZ
+        SpectralSolverRZ&
+#   else
+        SpectralSolver&
+#   endif
+            get_spectral_solver_fp (int lev) {return *spectral_solver_fp[lev];}
 #endif
+
+private:
     amrex::Vector<std::unique_ptr<FiniteDifferenceSolver>> m_fdtd_solver_fp;
     amrex::Vector<std::unique_ptr<FiniteDifferenceSolver>> m_fdtd_solver_cp;
 };