Update PODVector & SoA (#157)

Add length and additional component-wise access methods.
AMReX-Codes · Jul 21, 2023 · b284984 · b284984
1 parent 1992c18
commit b284984
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/src/Base/PODVector.cpp b/src/Base/PODVector.cpp
@@ -58,12 +58,13 @@ void make_PODVector(py::module &m, std::string typestr, std::string allocstr)
              }
         )
         .def(py::init<>())
-        .def(py::init<std::size_t>())
-        .def(py::init<PODVector_type&>())
+        .def(py::init<std::size_t>(), py::arg("size"))
+        .def(py::init<PODVector_type&>(), py::arg("other"))
         .def("push_back", py::overload_cast<const T&>(&PODVector_type::push_back))
         .def("pop_back", &PODVector_type::pop_back)
         .def("clear", &PODVector_type::clear)
         .def("size", &PODVector_type::size)
+        .def("__len__", &PODVector_type::size)
         // .def("max_size", &PODVector_type::max_size)
         .def("capacity", &PODVector_type::capacity)
         .def("empty", &PODVector_type::empty)

diff --git a/src/Particle/StructOfArrays.cpp b/src/Particle/StructOfArrays.cpp
@@ -24,13 +24,32 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
     py::class_<SOAType>(m, soa_name.c_str())
         .def(py::init())
         .def("define", &SOAType::define)
-        .def("NumRealComps", &SOAType::NumRealComps)
-        .def("NumIntComps", &SOAType::NumIntComps)
+        .def("NumRealComps", &SOAType::NumRealComps,
+             "Get the number of compile-time + runtime Real components")
+        .def("NumIntComps", &SOAType::NumIntComps,
+             "Get the number of compile-time + runtime Int components")
+
+        // compile-time components
         .def("GetRealData", py::overload_cast<>(&SOAType::GetRealData),
-            py::return_value_policy::reference_internal)
+            py::return_value_policy::reference_internal,
+            "Get access to the particle Real Arrays (only compile-time components)")
         .def("GetIntData", py::overload_cast<>(&SOAType::GetIntData),
-            py::return_value_policy::reference_internal)
-        .def("size", &SOAType::size)
+            py::return_value_policy::reference_internal,
+            "Get access to the particle Int Arrays (only compile-time components)")
+        // compile-time and runtime components
+        .def("GetRealData", py::overload_cast<const int>(&SOAType::GetRealData),
+             py::return_value_policy::reference_internal,
+             py::arg("index"),
+             "Get access to a particle Real component Array (compile-time and runtime component)")
+        .def("GetIntData", py::overload_cast<const int>(&SOAType::GetIntData),
+             py::return_value_policy::reference_internal,
+             py::arg("index"),
+             "Get access to a particle Real component Array (compile-time and runtime component)")
+
+        .def("size", &SOAType::size,
+             "Get the number of particles")
+        .def("__len__", &SOAType::size,
+             "Get the number of particles")
         .def("numParticles", &SOAType::numParticles)
         .def("numRealParticles", &SOAType::numRealParticles)
         .def("numTotalParticles", &SOAType::numTotalParticles)