ParticleContainer: WarpX 1D and 2D

More specializations.
AMReX-Codes · Dec 22, 2023 · fc21e7d · fc21e7d
1 parent eb24d03
commit fc21e7d
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 12 deletions.
diff --git a/src/Base/PODVector.cpp b/src/Base/PODVector.cpp
@@ -137,4 +137,5 @@ void make_PODVector(py::module &m, std::string typestr)
 void init_PODVector(py::module& m) {
     make_PODVector<ParticleReal> (m, "real");
     make_PODVector<int> (m, "int");
+    make_PODVector<uint64_t> (m, "uint64");
 }
diff --git a/src/Particle/Particle.cpp b/src/Particle/Particle.cpp
@@ -327,6 +327,7 @@ void init_Particle(py::module& m) {
     make_Particle< 3, 2 > (m);
     make_Particle< 4, 0 > (m);  // HiPACE++ 22.07
     make_Particle< 5, 0 > (m);  // ImpactX 22.07
+    make_Particle< 6, 0 > (m);  // WarpX 24.01+
     //make_Particle< 7, 0 > (m);  // WarpX 24.01+
     make_Particle< 8, 0 > (m);  // ImpactX 24.01+
     make_Particle< 37, 1> (m);  // HiPACE++ 22.07

diff --git a/src/Particle/ParticleContainer_WarpX.cpp b/src/Particle/ParticleContainer_WarpX.cpp
@@ -16,5 +16,11 @@ void init_ParticleContainer_WarpX(py::module& m) {
     make_ParticleContainer_and_Iterators<Particle<0, 0>, 4, 0>(m);   // WarpX 22.07 - 23.12 1D-3D
     //make_ParticleContainer_and_Iterators<Particle<0, 0>, 5, 0> (m);   // WarpX 22.07 - 23.12 RZ
 
-    make_ParticleContainer_and_Iterators<SoAParticle<7, 0>, 7, 0>(m);  // WarpX 24.01+ 1D-3D
+#if AMREX_SPACEDIM == 1
+    make_ParticleContainer_and_Iterators<SoAParticle<5, 0>, 5, 0>(m);  // WarpX 24.01+ 1D
+#elif AMREX_SPACEDIM == 2
+    make_ParticleContainer_and_Iterators<SoAParticle<6, 0>, 6, 0>(m);  // WarpX 24.01+ 2D
+#elif AMREX_SPACEDIM == 3
+    make_ParticleContainer_and_Iterators<SoAParticle<7, 0>, 7, 0>(m);  // WarpX 24.01+ 3D
+#endif
 }
diff --git a/src/Particle/ParticleTile.cpp b/src/Particle/ParticleTile.cpp
@@ -30,10 +30,15 @@ void make_ParticleTileData(py::module &m)
     using ParticleTileDataType = ParticleTileData<T_ParticleType, NArrayReal, NArrayInt>;
     using SuperParticleType = Particle<NStructReal + NArrayReal, NStructInt + NArrayInt>;
 
-    auto const particle_tile_data_type =
-        std::string("ParticleTileData_") + std::to_string(NStructReal) + "_" +
-        std::to_string(NStructInt) + "_" + std::to_string(NArrayReal) + "_" +
+    auto particle_tile_data_type = std::string("ParticleTileData_");
+    if (T_ParticleType::is_soa_particle)
+        particle_tile_data_type += "pureSoA_";
+    particle_tile_data_type +=
+        std::to_string(NStructReal) + "_" +
+        std::to_string(NStructInt) + "_" +
+        std::to_string(NArrayReal) + "_" +
         std::to_string(NArrayInt);
+
     py::class_<ParticleTileDataType>(m, particle_tile_data_type.c_str())
             .def(py::init())
             .def_readonly("m_size", &ParticleTileDataType::m_size)
@@ -63,9 +68,14 @@ void make_ParticleTile(py::module &m, std::string allocstr)
     using ParticleTileType = ParticleTile<T_ParticleType, NArrayReal, NArrayInt, Allocator>;
     using SuperParticleType = Particle<NStructReal + NArrayReal, NStructInt + NArrayInt>;
 
-    auto const particle_tile_type = std::string("ParticleTile_") + std::to_string(NStructReal) + "_" +
-                                    std::to_string(NStructInt) + "_" + std::to_string(NArrayReal) + "_" +
-                                    std::to_string(NArrayInt) + "_" + allocstr;
+    auto particle_tile_type = std::string("ParticleTile_");
+    if (T_ParticleType::is_soa_particle)
+        particle_tile_type += "pureSoA_";
+    particle_tile_type += std::to_string(NStructReal) + "_" +
+                          std::to_string(NStructInt) + "_" +
+                          std::to_string(NArrayReal) + "_" +
+                          std::to_string(NArrayInt) + "_" + allocstr;
+
     auto py_particle_tile = py::class_<ParticleTileType>(m, particle_tile_type.c_str())
         .def(py::init())
         .def_readonly_static("NAR", &ParticleTileType::NAR)
@@ -184,15 +194,27 @@ void init_ParticleTile(py::module& m) {
     // AMReX legacy AoS position + id/cpu particle ype
     using ParticleType_0_0 = Particle<0, 0>;
     using ParticleType_1_1 = Particle<1, 1>;
+#if AMREX_SPACEDIM == 1
+    using SoAParticleType_5_0 = SoAParticle<5, 0>;
+#elif AMREX_SPACEDIM == 2
+    using SoAParticleType_6_0 = SoAParticle<6, 0>;
+#elif AMREX_SPACEDIM == 3
     using SoAParticleType_7_0 = SoAParticle<7, 0>;
+#endif
     using SoAParticleType_8_0 = SoAParticle<8, 0>;
 
     // TODO: we might need to move all or most of the defines in here into a
     //       test/example submodule, so they do not collide with downstream projects
     make_ParticleTile<ParticleType_1_1, 2, 1> (m);
     make_ParticleTile<ParticleType_0_0, 4, 0> (m);   // HiPACE++ 22.07
     make_ParticleTile<ParticleType_0_0, 5, 0> (m);   // ImpactX 22.07
-    make_ParticleTile<SoAParticleType_7_0, 7, 0> (m);   // ImpactX 24.01+
+#if AMREX_SPACEDIM == 1
+    make_ParticleTile<SoAParticleType_5_0, 5, 0> (m);   // WarpX 24.01+ 1D
+#elif AMREX_SPACEDIM == 2
+    make_ParticleTile<SoAParticleType_6_0, 6, 0> (m);   // WarpX 24.01+ 2D
+#elif AMREX_SPACEDIM == 3
+    make_ParticleTile<SoAParticleType_7_0, 7, 0> (m);   // WarpX 24.01+ 3D
+#endif
     make_ParticleTile<SoAParticleType_8_0, 8, 0> (m);   // ImpactX 24.01+
     make_ParticleTile<ParticleType_0_0, 37, 1> (m);  // HiPACE++ 22.07
 }
diff --git a/src/Particle/StructOfArrays.cpp b/src/Particle/StructOfArrays.cpp
@@ -20,9 +20,14 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
 
     using SOAType = StructOfArrays<NReal, NInt, Allocator, use64BitIdCpu>;
 
-    auto const soa_name = std::string("StructOfArrays_") + std::to_string(NReal) + "_" +
-                          std::to_string(NInt) + "_" + allocstr;
-    py::class_<SOAType>(m, soa_name.c_str())
+    auto soa_name = std::string("StructOfArrays_") + std::to_string(NReal) + "_" +
+                    std::to_string(NInt);
+    if (use64BitIdCpu)
+        soa_name += "_idcpu";
+    soa_name += "_" + allocstr;
+
+    py::class_<SOAType> py_SoA(m, soa_name.c_str());
+    py_SoA
         .def(py::init())
         .def("define", &SOAType::define)
         .def_property_readonly("num_real_comps", &SOAType::NumRealComps,
@@ -58,6 +63,10 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
         .def("getNumNeighbors", &SOAType::getNumNeighbors)
         .def("resize", &SOAType::resize)
     ;
+    if (use64BitIdCpu)
+        py_SoA.def("GetIdCPUData", py::overload_cast<>(&SOAType::GetIdCPUData),
+                   py::return_value_policy::reference_internal,
+                   "Get access to a particle IdCPU component Array");
 }
 
 template <int NReal, int NInt, bool use64BitIdCpu=false>
@@ -92,7 +101,13 @@ void init_StructOfArrays(py::module& m) {
     make_StructOfArrays< 2, 1>(m);
     make_StructOfArrays< 4, 0>(m);  // HiPACE++ 22.08 - 23.12
     make_StructOfArrays< 5, 0>(m);  // ImpactX 22.07 - 23.12
-    make_StructOfArrays< 7, 0, true>(m);  // WarpX 24.01+
+#if AMREX_SPACEDIM == 1
+    make_StructOfArrays< 5, 0, true>(m);  // WarpX 24.01+ 1D
+#elif AMREX_SPACEDIM == 2
+    make_StructOfArrays< 6, 0, true>(m);  // WarpX 24.01+ 2D
+#elif AMREX_SPACEDIM == 3
+    make_StructOfArrays< 7, 0, true>(m);  // WarpX 24.01+ 3D
+#endif
     make_StructOfArrays< 8, 0, true>(m);  // ImpactX 24.01+
     make_StructOfArrays<37, 1>(m);  // HiPACE++ 22.09 - 23.12
 }