Particle Container: Support Pure SoA

Transition particle containers, iterators, etc. to support also the
new pure SoA layout.

src/Particle/ParticleContainer.H

@@ -11,6 +11,7 @@
 #include <AMReX_GpuAllocators.H>
 #include <AMReX_IntVect.H>
 #include <AMReX_ParIter.H>
+#include <AMReX_Particle.H>
 #include <AMReX_Particles.H>
 #include <AMReX_ParticleContainerBase.H>
 #include <AMReX_ParticleContainer.H>
@@ -24,38 +25,49 @@
 #include <sstream>
 
 namespace py = pybind11;
-using namespace amrex;
 
+template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
+std::string particle_type_suffix ()
+{
+    std::string suffix;
+    if constexpr (T_ParticleType::is_soa_particle)
+        suffix = "pureSoA_" +
+                 std::to_string(T_NArrayReal) + "_" +
+                 std::to_string(T_NArrayInt);
+    else
+        suffix = std::to_string(T_ParticleType::NReal) + "_" +
+                 std::to_string(T_ParticleType::NInt) + "_" +
+                 std::to_string(T_NArrayReal) + "_" +
+                 std::to_string(T_NArrayInt);
+
+    return suffix;
+}
 
 template <bool is_const, typename T_ParIterBase>
 void make_Base_Iterators (py::module &m, std::string allocstr)
 {
     using iterator_base = T_ParIterBase;
     using container = typename iterator_base::ContainerType;
-    constexpr int NStructReal = container::NStructReal;
-    constexpr int NStructInt = container::NStructInt;
+    using ParticleType = typename container::ParticleType;
     constexpr int NArrayReal = container::NArrayReal;
     constexpr int NArrayInt = container::NArrayInt;
 
+    std::string const suffix = particle_type_suffix<ParticleType, NArrayReal, NArrayInt>();
     std::string particle_it_base_name = std::string("Par");
     if (is_const) particle_it_base_name += "Const";
-    particle_it_base_name += "IterBase_" +
-        std::to_string(NStructReal) + "_" + std::to_string(NStructInt) + "_" +
-        std::to_string(NArrayReal) + "_" + std::to_string(NArrayInt) + "_" +
-        allocstr;
-    py::class_<iterator_base, MFIter>(m, particle_it_base_name.c_str())
+    particle_it_base_name += "IterBase_" + suffix + "_" + allocstr;
+    auto py_it_base = py::class_<iterator_base, MFIter>(m, particle_it_base_name.c_str())
         .def(py::init<container&, int>(),
             py::arg("particle_container"), py::arg("level"))
         .def(py::init<container&, int, MFItInfo&>(),
              py::arg("particle_container"), py::arg("level"), py::arg("info"))
 
         .def("particle_tile", &iterator_base::GetParticleTile,
                                py::return_value_policy::reference_internal)
-        .def("aos", &iterator_base::GetArrayOfStructs,
-                               py::return_value_policy::reference_internal)
         .def("soa", &iterator_base::GetStructOfArrays,
                                py::return_value_policy::reference_internal)
 
+        .def_property_readonly_static("is_soa_particle", [](const py::object&){ return ParticleType::is_soa_particle;})
         .def_property_readonly("num_particles", &iterator_base::numParticles)
         .def_property_readonly("num_real_particles", &iterator_base::numRealParticles)
         .def_property_readonly("num_neighbor_particles", &iterator_base::numNeighborParticles)
@@ -75,26 +87,30 @@ void make_Base_Iterators (py::module &m, std::string allocstr)
              py::return_value_policy::reference_internal
         )
     ;
+
+    // only legacy particle has an AoS data structure for positions and id+cpu
+    if constexpr (!ParticleType::is_soa_particle)
+        py_it_base.def("aos",
+                       &iterator_base::GetArrayOfStructs,
+                       py::return_value_policy::reference_internal);
 }
 
 template <bool is_const, typename T_ParIter, template<class> class Allocator=DefaultAllocator>
 void make_Iterators (py::module &m, std::string allocstr)
 {
     using iterator = T_ParIter;
     using container = typename iterator::ContainerType;
-    constexpr int NStructReal = container::NStructReal;
-    constexpr int NStructInt = container::NStructInt;
+    using ParticleType = typename container::ParticleType;
     constexpr int NArrayReal = container::NArrayReal;
     constexpr int NArrayInt = container::NArrayInt;
 
-    using iterator_base = amrex::ParIterBase<is_const, NStructReal, NStructInt, NArrayReal, NArrayInt, Allocator>;
+    using iterator_base = amrex::ParIterBase_impl<is_const, ParticleType, NArrayReal, NArrayInt, Allocator>;
     make_Base_Iterators< is_const, iterator_base >(m, allocstr);
 
+    std::string const suffix = particle_type_suffix<ParticleType, NArrayReal, NArrayInt>();
     auto particle_it_name = std::string("Par");
     if (is_const) particle_it_name += "Const";
-    particle_it_name += std::string("Iter_") + std::to_string(NStructReal) + "_" +
-                        std::to_string(NStructInt) + "_" + std::to_string(NArrayReal) + "_" +
-                        std::to_string(NArrayInt) + "_" + allocstr;
+    particle_it_name += std::string("Iter_") + suffix + "_" + allocstr;
     py::class_<iterator, iterator_base>(m, particle_it_name.c_str())
         .def("__repr__",
              [particle_it_name](iterator const & pti) {
@@ -108,53 +124,55 @@ void make_Iterators (py::module &m, std::string allocstr)
              py::arg("particle_container"), py::arg("level"))
         .def(py::init<container&, int, MFItInfo&>(),
              py::arg("particle_container"), py::arg("level"), py::arg("info"))
+        .def_property_readonly_static("is_soa_particle", [](const py::object&){ return ParticleType::is_soa_particle;})
     ;
 }
 
-template <int T_NStructReal, int T_NStructInt=0, int T_NArrayReal=0, int T_NArrayInt=0>
+template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
 void make_ParticleInitData (py::module &m) {
-    using ParticleInitData = ParticleInitType<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt>;
+    using ParticleType = T_ParticleType;
+    using ParticleInitData = ParticleInitType<ParticleType::NReal, ParticleType::NInt, T_NArrayReal, T_NArrayInt>;
 
+    std::string const suffix = particle_type_suffix<T_ParticleType, T_NArrayReal, T_NArrayInt>();
     auto const particle_init_data_type =
-        std::string("ParticleInitType_") + std::to_string(T_NStructReal) + "_" +
-        std::to_string(T_NStructInt) + "_" + std::to_string(T_NArrayReal) + "_" +
-        std::to_string(T_NArrayInt);
+        std::string("ParticleInitType_") + suffix;
     py::class_<ParticleInitData>(m, particle_init_data_type.c_str())
         .def(py::init<>())
+        .def_property_readonly_static("is_soa_particle", [](const py::object&){ return ParticleType::is_soa_particle;})
         .def_readwrite("real_struct_data", &ParticleInitData::real_struct_data)
         .def_readwrite("int_struct_data", &ParticleInitData::int_struct_data)
         .def_readwrite("real_array_data", &ParticleInitData::real_array_data)
         .def_readwrite("int_array_data", &ParticleInitData::int_array_data)
     ;
 }
 
-template <int T_NStructReal, int T_NStructInt=0, int T_NArrayReal=0, int T_NArrayInt=0,
+template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0,
           template<class> class Allocator=DefaultAllocator>
 void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
 {
-    using ParticleContainerType = ParticleContainer<
-        T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    using ParticleType = T_ParticleType;
+    using ParticleContainerType = ParticleContainer_impl<
+        ParticleType, T_NArrayReal, T_NArrayInt,
         Allocator
     >;
-    using ParticleInitData = typename ParticleContainerType::ParticleInitData;
+    //using ParticleInitData = typename ParticleContainerType::ParticleInitData; // TODO pure SoA
     using ParticleTileType = typename ParticleContainerType::ParticleTileType;
-    //using AoS = typename ParticleContainerType::AoS;
 
-    auto const particle_container_type = std::string("ParticleContainer_") + std::to_string(T_NStructReal) + "_" +
-                                         std::to_string(T_NStructInt) + "_" + std::to_string(T_NArrayReal) + "_" +
-                                         std::to_string(T_NArrayInt) + "_" + allocstr;
+    std::string const suffix = particle_type_suffix<T_ParticleType, T_NArrayReal, T_NArrayInt>();
+    auto const particle_container_type = std::string("ParticleContainer_") + suffix + "_" + allocstr;
     py::class_<ParticleContainerType>(m, particle_container_type.c_str())
         .def(py::init())
         .def(py::init<const Geometry&, const DistributionMapping&, const BoxArray&>())
         .def(py::init<const Vector<Geometry>&,
-                        const Vector<DistributionMapping>&,
-                        const Vector<BoxArray>&,
-                        const Vector<int>&>())
+                      const Vector<DistributionMapping>&,
+                      const Vector<BoxArray>&,
+                      const Vector<int>&>())
         .def(py::init<const Vector<Geometry>&,
-                        const Vector<DistributionMapping>&,
-                        const Vector<BoxArray>&,
-                        const Vector<IntVect>&>())
+                      const Vector<DistributionMapping>&,
+                      const Vector<BoxArray>&,
+                      const Vector<IntVect>&>())
 
+        .def_property_readonly_static("is_soa_particle", [](const py::object&){return ParticleType::is_soa_particle;})
         .def_property_readonly_static("NStructReal", [](const py::object&){return ParticleContainerType::NStructReal; })
         .def_property_readonly_static("NStructInt", [](const py::object&){return ParticleContainerType::NStructInt; })
         .def_property_readonly_static("NArrayReal", [](const py::object&){return ParticleContainerType::NArrayReal; })
@@ -203,13 +221,13 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
         //                  const ParticleInitData& mass,
         //                  bool serialize = false, RealBox bx = RealBox());
 
-        .def("InitRandom", py::overload_cast<Long, ULong, const ParticleInitData&, bool, RealBox>(&ParticleContainerType::InitRandom))
+        //.def("InitRandom", py::overload_cast<Long, ULong, const ParticleInitData&, bool, RealBox>(&ParticleContainerType::InitRandom)) // TODO pure SoA
 
-        .def("InitRandomPerBox", py::overload_cast<Long, ULong, const ParticleInitData&>(&ParticleContainerType::InitRandomPerBox))
-        .def("InitOnePerCell", &ParticleContainerType::InitOnePerCell)
+        //.def("InitRandomPerBox", py::overload_cast<Long, ULong, const ParticleInitData&>(&ParticleContainerType::InitRandomPerBox)) // TODO pure SoA
+        //.def("InitOnePerCell", &ParticleContainerType::InitOnePerCell) // TODO pure SoA
 
-        .def("Increment", &ParticleContainerType::Increment)
-        .def("IncrementWithTotal", &ParticleContainerType::IncrementWithTotal, py::arg("mf"), py::arg("level"), py::arg("local")=false)
+        //.def("Increment", &ParticleContainerType::Increment) // TODO pure SoA
+        //.def("IncrementWithTotal", &ParticleContainerType::IncrementWithTotal, py::arg("mf"), py::arg("level"), py::arg("local")=false) // TODO pure SoA
         .def("Redistribute", &ParticleContainerType::Redistribute, py::arg("lev_min")=0, py::arg("lev_max")=-1,
                                             py::arg("nGrow")=0, py::arg("local")=0, py::arg("remove_negative")=true)
         .def("SortParticlesByCell", &ParticleContainerType::SortParticlesByCell)
@@ -235,12 +253,12 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
         // void CreateVirtualParticles (int level, AoS& virts) const;
         //.def("CreateVirtualParticles", py::overload_cast<int, AoS&>(&ParticleContainerType::CreateVirtualParticles, py::const_),
         //    py::arg("level"), py::arg("virts"))
-        .def("CreateVirtualParticles", py::overload_cast<int, ParticleTileType&>(&ParticleContainerType::CreateVirtualParticles, py::const_),
-             py::arg("level"), py::arg("virts"))
+        //.def("CreateVirtualParticles", py::overload_cast<int, ParticleTileType&>(&ParticleContainerType::CreateVirtualParticles, py::const_),
+        //     py::arg("level"), py::arg("virts")) // TODO pure SoA
         //.def("CreateGhostParticles", py::overload_cast<int, int, AoS&>(&ParticleContainerType::CreateGhostParticles, py::const_),
         //     py::arg("level"), py::arg("ngrow"), py::arg("ghosts"))
-        .def("CreateGhostParticles", py::overload_cast<int, int, ParticleTileType&>(&ParticleContainerType::CreateGhostParticles, py::const_),
-             py::arg("level"), py::arg("ngrow"), py::arg("ghosts"))
+        //.def("CreateGhostParticles", py::overload_cast<int, int, ParticleTileType&>(&ParticleContainerType::CreateGhostParticles, py::const_),
+        //     py::arg("level"), py::arg("ngrow"), py::arg("ghosts")) // TODO pure SoA
         //.def("AddParticlesAtLevel", py::overload_cast<AoS&, int, int>(&ParticleContainerType::AddParticlesAtLevel),
         //    py::arg("particles"), py::arg("level"), py::arg("ngrow")=0)
         .def("AddParticlesAtLevel", py::overload_cast<ParticleTileType&, int, int>(&ParticleContainerType::AddParticlesAtLevel),
@@ -342,46 +360,48 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
         // }
     ;
 
-    using iterator = amrex::ParIter<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt, Allocator>;
+    using iterator = amrex::ParIter_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>;
     make_Iterators< false, iterator, Allocator >(m, allocstr);
-    using const_iterator = amrex::ParConstIter<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt, Allocator>;
+    using const_iterator = amrex::ParConstIter_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>;
     make_Iterators< true, const_iterator, Allocator >(m, allocstr);
 }
 
-template <int T_NStructReal, int T_NStructInt=0, int T_NArrayReal=0, int T_NArrayInt=0>
+/** Create ParticleContainers and Iterators
+ */
+template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
 void make_ParticleContainer_and_Iterators (py::module &m)
 {
-    make_ParticleInitData<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt>(m);
+    make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);
 
     // see Src/Base/AMReX_GpuContainers.H
     //   !AMREX_USE_GPU: DefaultAllocator = std::allocator
     //    AMREX_USE_GPU: DefaultAllocator = amrex::ArenaAllocator
 
     //   work-around for https://github.com/pybind/pybind11/pull/4581
-    //make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    //make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
     //                                     std::allocator>(m, "std");            // CPU DefaultAllocator
-    //make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    //make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
     //                                     amrex::ArenaAllocator>(m, "arena");   // GPU DefaultAllocator
 #ifdef AMREX_USE_GPU
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
                                          std::allocator>(m, "std");
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
                                          amrex::DefaultAllocator>(m, "default");   // amrex::ArenaAllocator
 #else
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
-                                         amrex::DefaultAllocator>(m, "default");   // std::allocator
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
-                                         amrex::ArenaAllocator>(m, "arena");
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+            amrex::DefaultAllocator>(m, "default");   // std::allocator
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+            amrex::ArenaAllocator>(m, "arena");
 #endif
     //   end work-around
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
-                                         amrex::PinnedArenaAllocator>(m, "pinned");
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+            amrex::PinnedArenaAllocator>(m, "pinned");
 #ifdef AMREX_USE_GPU
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
                                          amrex::DeviceArenaAllocator>(m, "device");
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
                                          amrex::ManagedArenaAllocator>(m, "managed");
-    make_ParticleContainer_and_Iterators<T_NStructReal, T_NStructInt, T_NArrayReal, T_NArrayInt,
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
                                          amrex::AsyncArenaAllocator>(m, "async");
 #endif
-}
+}

src/Particle/ParticleContainer.cpp

@@ -5,14 +5,19 @@
  */
 #include "ParticleContainer.H"
 
+#include <AMReX_Particle.H>
+
+
 void init_ParticleContainer_HiPACE(py::module& m);
 void init_ParticleContainer_ImpactX(py::module& m);
 void init_ParticleContainer_WarpX(py::module& m);
 
 void init_ParticleContainer(py::module& m) {
+    using namespace amrex;
+
     // 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_ParticleContainer_and_Iterators< 1, 1, 2, 1> (m);  // tests
+    make_ParticleContainer_and_Iterators<Particle<1, 1>, 2, 1>(m);  // tests
 
     init_ParticleContainer_HiPACE(m);
     init_ParticleContainer_ImpactX(m);

src/Particle/ParticleContainer_HiPACE.cpp

@@ -5,10 +5,14 @@
  */
 #include "ParticleContainer.H"
 
+#include <AMReX_Particle.H>
+
 
 void init_ParticleContainer_HiPACE(py::module& /* m */) {
+    using namespace amrex;
+
     // 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_ParticleContainer_and_Iterators< 0, 0, 4, 0> (m);   // HiPACE++ 22.07
-    //make_ParticleContainer_and_Iterators< 0, 0, 37, 1> (m);  // HiPACE++ 22.07
+    //make_ParticleContainer_and_Iterators<Particle<0, 0>, 4, 0> (m);   // HiPACE++ 22.07
+    //make_ParticleContainer_and_Iterators<Particle<0, 0>, 37, 1> (m);  // HiPACE++ 22.07
 }

src/Particle/ParticleContainer_ImpactX.cpp

@@ -5,9 +5,15 @@
  */
 #include "ParticleContainer.H"
 
+#include <AMReX_Particle.H>
+#include <AMReX_ParticleTile.H>
+
 
 void init_ParticleContainer_ImpactX(py::module& m) {
+    using namespace amrex;
+
     // 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_ParticleContainer_and_Iterators< 0, 0, 5, 0> (m);   // ImpactX 22.07
+    make_ParticleContainer_and_Iterators<Particle<0, 0>, 5, 0>(m);     // ImpactX 22.07 - 23.04
+    make_ParticleContainer_and_Iterators<SoAParticle<8, 2>, 8, 2>(m);  // ImpactX 23.05+
 }

src/Particle/ParticleContainer_WarpX.cpp

@@ -5,10 +5,14 @@
  */
 #include "ParticleContainer.H"
 
+#include <AMReX_Particle.H>
+
 
 void init_ParticleContainer_WarpX(py::module& m) {
+    using namespace amrex;
+
     // 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_ParticleContainer_and_Iterators< 0, 0, 4, 0> (m);   // WarpX 22.07 1D-3D
-    //make_ParticleContainer_and_Iterators< 0, 0, 5, 0> (m);   // WarpX 22.07 RZ
+    make_ParticleContainer_and_Iterators<Particle<0, 0>, 4, 0>(m);   // WarpX 22.07 - 23.04 1D-3D
+    //make_ParticleContainer_and_Iterators<Particle<0, 0>, 5, 0> (m);   // WarpX 22.07 - 23.04 RZ
 }