Access species specific charge density from python (#2710)

* added python wrapper function to deposit a specific species density in rho_fp

* added 1D ES input file with MCC that uses the charge deposition functionality

* reset rho_fp[lev] before depositing

* updated documentation

* switch to using simulation.extension in Poisson solver

* Apply suggestion from code review

Co-authored-by: Phil Miller <phil.miller@intensecomputing.com>

* suggested changes from code review

* add comment explaining why a direct Poisson solver is used

* removed direct solver in 1D example since it is actually slower than the MLMG solver

* Apply suggestions from code review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* added docstring for warpx_depositChargeDensity

* fixed order of imports in new PICMI input file

Co-authored-by: Phil Miller <phil.miller@intensecomputing.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

1 files changed, 31 insertions, 0 deletions

diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp
index 7350bbbb6..061e80c6a 100644
--- a/Source/Python/WarpXWrappers.cpp
+++ b/Source/Python/WarpXWrappers.cpp
@@ -606,6 +606,37 @@ namespace
         particle_buffers.clearParticles();
     }
 
+    void warpx_depositChargeDensity (const char* char_species_name, int lev) {
+        // this function is used to deposit a given species' charge density
+        // in the rho_fp multifab which can then be accessed from python via
+        // pywarpx.fields.RhoFPWrapper()
+        WarpX& warpx = WarpX::GetInstance();
+        const auto & mypc = warpx.GetPartContainer();
+        const std::string species_name(char_species_name);
+        auto & myspc = mypc.GetParticleContainerFromName(species_name);
+        auto * rho_fp = warpx.get_pointer_rho_fp(lev);
+
+        if (rho_fp == nullptr) {
+            warpx.RecordWarning(
+                "WarpXWrappers", "rho_fp is not allocated", WarnPriority::low
+            );
+            return;
+        }
+
+        // reset rho before depositing
+        rho_fp->setVal(0.);
+
+        for (WarpXParIter pti(myspc, lev); pti.isValid(); ++pti)
+        {
+            const long np = pti.numParticles();
+            auto& wp = pti.GetAttribs(PIdx::w);
+            myspc.DepositCharge(pti, wp, nullptr, rho_fp, 0, 0, np, 0, lev, lev);
+        }
+#ifdef WARPX_DIM_RZ
+        warpx.ApplyInverseVolumeScalingToChargeDensity(rho_fp, lev);
+#endif
+    }
+
     void warpx_ComputeDt () {
         WarpX& warpx = WarpX::GetInstance();
         warpx.ComputeDt ();