added binding for WarpXParIter - needed to port `libwarpx.depositCh…

…argeDensity()` which is an ongoing effort

Python/pywarpx/_libwarpx.py

@@ -232,7 +232,7 @@ def getistep(self, level=0):
             The refinement level to reference
         '''
 
-        warpx = libwarpx.libwarpx_so.get_instance()
+        warpx = self.libwarpx_so.get_instance()
         return warpx.getistep(level)
 
     def gett_new(self, level=0):
@@ -247,7 +247,7 @@ def gett_new(self, level=0):
             The refinement level to reference
         '''
 
-        warpx = libwarpx.libwarpx_so.get_instance()
+        warpx = self.libwarpx_so.get_instance()
         return warpx.gett_new(level)
 
     def evolve(self, num_steps=-1):
@@ -965,15 +965,36 @@ def depositChargeDensity(self, species_name, level, clear_rho=True, sync_rho=Tru
         sync_rho       : bool
             If True, perform MPI exchange and properly set boundary cells for rho_fp.
         '''
+        warpx = self.libwarpx_so.get_instance()
+        mypc = warpx.multi_particle_container()
+        myspc = mypc.get_particle_container_from_name(species_name)
+
+        rho_fp = warpx.multifab(f'rho_fp[level={level}]')
+
+        if rho_fp is None:
+            raise RuntimeWarning("Multifab `rho_fp` is not allocated.")
+            # ablastr::warn_manager::WMRecordWarning(
+            #     "WarpXWrappers", "rho_fp is not allocated",
+            #     ablastr::warn_manager::WarnPriority::low
+            # );
+            return
 
         if clear_rho:
             from . import fields
             fields.RhoFPWrapper(level, True)[...] = 0.0
-        self.libwarpx_so.warpx_depositChargeDensity(
-            ctypes.c_char_p(species_name.encode('utf-8')), level
-        )
-        if sync_rho:
-            self.libwarpx_so.warpx_SyncRho()
+
+        for pti in self.libwarpx_so.WarpXParIter(myspc, level):
+            print(pti.num_particles)
+            # TODO: Implement the bindings for the remaining logic here
+            # wp = pti.GetAttribs(PIdx::w);
+            # # Do this unconditionally, ignoring myspc.do_not_deposit, to support diagnostic uses
+            # myspc.deposit_charge(pti, wp, nullptr, rho_fp, 0, 0, np, 0, lev, lev);
+
+        # if self.geometry_dim == 'rz':
+        #     warpx.apply_inverse_volume_scaling_to_charge_density(rho_fp, lev);
+
+        # if sync_rho:
+        #     warpx.sync_rho()
 
     def set_potential_EB(self, potential):
         """

Source/Python/WarpXParticleContainer.cpp

@@ -9,6 +9,15 @@
 
 namespace py = pybind11;
 
+void init_WarpXParIter (py::module& m)
+{
+    py::class_<WarpXParIter, amrex::ParIter<0,0,PIdx::nattribs>
+    > wpi (m, "WarpXParIter");
+    wpi
+        .def(py::init<WarpXParticleContainer&, int>())
+        ;
+}
+
 void init_WarpXParticleContainer (py::module& m)
 {
     py::class_<
@@ -29,5 +38,18 @@ void init_WarpXParticleContainer (py::module& m)
             py::arg("nattr_int"), py::arg("attr_int"),
             py::arg("uniqueparticles"), py::arg("id")
         )
+        .def("deposit_charge",
+            static_cast<void (WarpXParticleContainer::*)(
+                WarpXParIter&,
+                amrex::Gpu::DeviceVector<amrex::Real> const &,
+                const int * const, amrex::MultiFab*,
+                const int, const long, const long,
+                const int, const int, const int
+            )>(&WarpXParticleContainer::DepositCharge),
+            py::arg("pti"), py::arg("wp"), py::arg("ion_lev"),
+            py::arg("rho"), py::arg("icomp"),
+            py::arg("offset"), py::arg("np_to_depose"),
+            py::arg("thread_num"), py::arg("lev"), py::arg("depos_lev")
+        )
         ;
 }

Source/Python/WarpXWrappers.H

@@ -101,24 +101,13 @@ extern "C" {
 
     void warpx_clearParticleBoundaryBuffer ();
 
-    /**
-     * \brief 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()
-     *
-     * @param[in] species_name specifying the name of the species to deposit
-     * @param[in] lev mesh refinement level
-     */
-    void warpx_depositChargeDensity (const char* species_name, int lev);
-
   void warpx_ComputeDt ();
   void warpx_MoveWindow (int step, bool move_j);
 
   void warpx_EvolveE (amrex::Real dt);
   void warpx_EvolveB (amrex::Real dt, DtType a_dt_type);
   void warpx_FillBoundaryE ();
   void warpx_FillBoundaryB ();
-  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,

Source/Python/WarpXWrappers.cpp

@@ -356,36 +356,6 @@
         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) {
-            ablastr::warn_manager::WMRecordWarning(
-                "WarpXWrappers", "rho_fp is not allocated",
-                ablastr::warn_manager::WarnPriority::low
-            );
-            return;
-        }
-
-        for (WarpXParIter pti(myspc, lev); pti.isValid(); ++pti)
-        {
-            const long np = pti.numParticles();
-            auto& wp = pti.GetAttribs(PIdx::w);
-            // Do this unconditionally, ignoring myspc.do_not_deposit, to support diagnostic uses
-            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();
@@ -411,10 +381,6 @@
         WarpX& warpx = WarpX::GetInstance();
         warpx.FillBoundaryB(warpx.getngEB());
     }
-    void warpx_SyncRho () {
-        WarpX& warpx = WarpX::GetInstance();
-        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,

Source/Python/pyWarpX.cpp

@@ -30,6 +30,7 @@ namespace py = pybind11;
 
 
 // forward declarations of exposed classes
+void init_WarpXParIter (py::module&);
 void init_WarpXParticleContainer (py::module&);
 void init_MultiParticleContainer (py::module&);
 void init_WarpX(py::module&);
@@ -55,6 +56,7 @@ PYBIND11_MODULE(PYWARPX_MODULE_NAME, m) {
     )pbdoc";
 
     // note: order from parent to child classes
+    init_WarpXParIter(m);
     init_WarpXParticleContainer(m);
     init_MultiParticleContainer(m);
     init_WarpX(m);