Complete class WarpXParticleContainer

ax3l · Jul 12, 2023 · c198175 · c198175
1 parent ef154af
commit c198175
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 5 deletions.
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
@@ -784,10 +784,8 @@ def add_real_comp(self, species_name, pid_name, comm=True):
         '''
         warpx = self.libwarpx_so.get_instance()
         mpc = warpx.multi_particle_container()
-        self.libwarpx_so.warpx_addRealComp(
-            ctypes.c_char_p(species_name.encode('utf-8')),
-            ctypes.c_char_p(pid_name.encode('utf-8')), comm
-        )
+        pc = mpc.get_particle_container_from_name(species_name)
+        pc.add_real_comp(pid_name, comm)
 
     def get_species_charge_sum(self, species_name, local=False):
         '''

diff --git a/Source/Python/CMakeLists.txt b/Source/Python/CMakeLists.txt
@@ -18,6 +18,7 @@ foreach(D IN LISTS WarpX_DIMS)
             # pybind11
             WarpX.cpp
             MultiParticleContainer.cpp
+            WarpXParticleContainer.cpp
         )
     endif()
 endforeach()
diff --git a/Source/Python/MultiParticleContainer.cpp b/Source/Python/MultiParticleContainer.cpp
@@ -6,10 +6,18 @@
 #include "pyWarpX.H"
 
 #include <Particles/MultiParticleContainer.H>
+#include <Particles/WarpXParticleContainer.H>
 
 namespace py = pybind11;
 
 void init_MultiParticleContainer (py::module& m)
 {
     py::class_<MultiParticleContainer> mpc(m, "MultiParticleContainer");
+    mpc
+        .def("get_particle_container_from_name",
+            &MultiParticleContainer::GetParticleContainerFromName,
+            py::arg("name"),
+            py::return_value_policy::reference_internal
+        )
+    ;
 }
diff --git a/Source/Python/WarpXParticleContainer.cpp b/Source/Python/WarpXParticleContainer.cpp
@@ -0,0 +1,24 @@
+/* Copyright 2021-2022 The WarpX Community
+ *
+ * Authors: Axel Huebl, Remi Lehe
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyWarpX.H"
+
+#include <Particles/WarpXParticleContainer.H>
+
+namespace py = pybind11;
+
+void init_WarpXParticleContainer (py::module& m)
+{
+    py::class_<
+        WarpXParticleContainer,
+        amrex::ParticleContainer<0, 0, PIdx::nattribs, 0>
+    > wpc (m, "WarpXParticleContainer");
+    wpc
+        .def("add_real_comp",
+            [](WarpXParticleContainer& pc, const std::string& name, bool const comm) { pc.AddRealComp(name, comm); },
+            py::arg("name"), py::arg("comm")
+        )
+    ;
+}
diff --git a/Source/Python/pyWarpX.cpp b/Source/Python/pyWarpX.cpp
@@ -30,7 +30,8 @@ namespace py = pybind11;
 
 
 // forward declarations of exposed classes
-void init_MultiParticleContainer (py::module& m);
+void init_WarpXParticleContainer (py::module&);
+void init_MultiParticleContainer (py::module&);
 void init_WarpX(py::module&);
 
 PYBIND11_MODULE(PYWARPX_MODULE_NAME, m) {
@@ -54,6 +55,7 @@ PYBIND11_MODULE(PYWARPX_MODULE_NAME, m) {
     )pbdoc";
 
     // note: order from parent to child classes
+    init_WarpXParticleContainer(m);
     init_MultiParticleContainer(m);
     init_WarpX(m);