Particles: Restructure

The AMReX particle logic is unfortunately heavily templated. For
simpler maintainance, we move all declarations into header file
and define needed types recursively from the particle container.

docs/source/usage/api.rst

@@ -230,7 +230,7 @@ Additional runtime attributes (Real or Int) are always in SoA memory layout.
 
 ``amrex::ParticleContainer_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>`` is implemented for both legacy (AoS+SoA) and pure SoA particle types, many number of Real and Int arguments, and allocators, e.g.,
 
-.. autoclass:: amrex.space3d.ParticleContainer_1_1_2_1_default
+.. autoclass:: amrex.space3d.ParticleContainer_2_1_3_1_default
    :members:
    :undoc-members:
 
@@ -246,7 +246,7 @@ Likewise for other classes accessible and usable on particle containers:
 
 .. autoclass:: amrex.space3d.ParConstIter_pureSoA_8_0_default
 
-.. autoclass:: amrex.space3d.ParticleInitType_1_1_2_1
+.. autoclass:: amrex.space3d.ParticleInitType_2_1_3_1
    :members:
    :undoc-members:
 

src/Particle/ArrayOfStructs.cpp → src/Particle/ArrayOfStructs.H

@@ -1,8 +1,10 @@
 /* Copyright 2022 The AMReX Community
  *
- * Authors: Ryan Sandberg
+ * Authors: Ryan Sandberg, Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
+#pragma once
+
 #include "pyAMReX.H"
 
 #include <AMReX_ArrayOfStructs.H>
@@ -15,23 +17,6 @@ namespace
 {
     using namespace amrex;
 
-    // Note - this function MUST be consistent with AMReX_Particle.H
-    Long unpack_id (uint64_t idcpu) {
-        Long r = 0;
-
-        uint64_t sign = idcpu >> 63;  // extract leftmost sign bit
-        uint64_t val  = ((idcpu >> 24) & 0x7FFFFFFFFF);  // extract next 39 id bits
-
-        Long lval = static_cast<Long>(val);  // bc we take -
-        r = (sign) ? lval : -lval;
-        return r;
-    }
-
-    // Note - this function MUST be consistent with AMReX_Particle.H
-    int unpack_cpu (uint64_t idcpu) {
-        return static_cast<int>(idcpu & 0x00FFFFFF);
-    }
-
     /** CPU: __array_interface__ v3
      *
      * https://numpy.org/doc/stable/reference/arrays.interface.html
@@ -182,12 +167,3 @@ void make_ArrayOfStructs(py::module &m)
     make_ArrayOfStructs<ParticleType, amrex::AsyncArenaAllocator> (m, "async");
 #endif
 }
-
-void init_ArrayOfStructs(py::module& m) {
-    make_ArrayOfStructs<0, 0> (m);  // WarpX 22.07, ImpactX 22.07, HiPACE++ 22.07
-    make_ArrayOfStructs<1, 1> (m);  // test in ParticleContainer
-    make_ArrayOfStructs<2, 1> (m);  // test
-
-    m.def("unpack_ids", py::vectorize(unpack_id));
-    m.def("unpack_cpus", py::vectorize(unpack_cpu));
-}

src/Particle/CMakeLists.txt

@@ -1,10 +1,6 @@
 foreach(D IN LISTS AMReX_SPACEDIM)
     target_sources(pyAMReX_${D}d
       PRIVATE
-        Particle.cpp
-        StructOfArrays.cpp
-        ArrayOfStructs.cpp
-        ParticleTile.cpp
         ParticleContainer.cpp
         ParticleContainer_HiPACE.cpp
         ParticleContainer_ImpactX.cpp

src/Particle/Particle.cpp → src/Particle/Particle.H

@@ -3,6 +3,8 @@
  * Authors: Ryan Sandberg
  * License: BSD-3-Clause-LBNL
  */
+#pragma once
+
 #include "pyAMReX.H"
 
 #include <AMReX_Config.H>
@@ -16,26 +18,12 @@
 #include <string>
 #include <sstream>
 #include <stdexcept>
+#include <unordered_set>
 #include <utility>
 #include <cmath>
 #include <regex>
 
 
-struct PIdx
-{
-    enum RealValues { // Particle Attributes for sample particle struct-of-arrays
-          w = 0,
-          vx, vy, vz,
-          Ex, Ey, Ez,
-          nRealAttribs
-    };
-
-    enum IntValues {
-        nIntAttribs
-    };
-
-};
-
 namespace
 {
     /** Build a std::array from a fixed-size C array at compile-time */
@@ -50,6 +38,17 @@ namespace
 template <int T_NReal, int T_NInt=0>
 void make_Particle(py::module &m)
 {
+    // For legacy AoS + SoA particles, we register a particle and superparticle.
+    // The latter adds the SoA attributes. Avoid to double-register those types.
+    /*
+    static std::unordered_set<std::array<int, 2>> registered;
+    std::array<int, 2> const this_p = {T_NReal, T_NInt}
+    if (auto search = registered.find(this_p); search != registered.end()) {
+        return;
+    }
+    registered.insert(this_p);
+    */
+
     using namespace amrex;
 
     using ParticleType = Particle<T_NReal, T_NInt>;
@@ -281,34 +280,3 @@ void make_Particle(py::module &m)
 #endif
     ;
 }
-
-
-
-void init_Particle(py::module& m) {
-
-    // 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
-    py::class_<PIdx> pidx(m, "PIdx");
-    py::enum_<PIdx::RealValues>(pidx, "RealValues")
-        .value("w", PIdx::RealValues::w)
-        .value("vx", PIdx::RealValues::vx)
-        .value("vy", PIdx::RealValues::vy)
-        .value("vz", PIdx::RealValues::vz)
-        .value("Ex", PIdx::RealValues::Ex)
-        .value("Ey", PIdx::RealValues::Ey)
-        .value("Ez", PIdx::RealValues::Ez)
-    ;
-    py::enum_<PIdx::IntValues>(pidx, "IntValues")
-    ;
-    make_Particle< PIdx::nRealAttribs, PIdx::nIntAttribs > (m);
-    make_Particle< 0, 0 > (m);
-    make_Particle< 1, 1 > (m);
-    make_Particle< 2, 1 > (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.03+
-    //make_Particle< 7, 0 > (m);  // WarpX 24.03+
-    make_Particle< 8, 0 > (m);  // ImpactX 24.03+
-    make_Particle< 37, 1> (m);  // HiPACE++ 22.07
-}

src/Particle/ParticleContainer.H

@@ -7,6 +7,11 @@
 
 #include "pyAMReX.H"
 
+#include "Particle.H"
+#include "ArrayOfStructs.H"
+#include "StructOfArrays.H"
+#include "ParticleTile.H"
+
 #include <AMReX_BoxArray.H>
 #include <AMReX_GpuAllocators.H>
 #include <AMReX_IntVect.H>
@@ -416,6 +421,26 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
 template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
 void make_ParticleContainer_and_Iterators (py::module &m)
 {
+    if constexpr (T_ParticleType::is_soa_particle) {
+        make_Particle<T_NArrayReal, T_NArrayInt>(m);
+        make_StructOfArrays<T_NArrayReal, T_NArrayInt, true> (m);
+    } else {
+
+        make_Particle<  // particle
+            T_ParticleType::NReal,
+            T_ParticleType::NInt
+        >(m);
+        make_Particle<  // superparticle
+            T_ParticleType::NReal + T_NArrayReal,
+            T_ParticleType::NInt  + T_NArrayInt
+        >(m);
+
+        make_ArrayOfStructs<T_ParticleType::NReal, T_ParticleType::NInt> (m);
+        make_StructOfArrays<T_NArrayReal, T_NArrayInt> (m);
+    }
+
+    make_ParticleTile<T_ParticleType, T_NArrayReal, T_NArrayInt> (m);
+
     make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);
 
     // first, because used as copy target in methods in containers with other allocators

src/Particle/ParticleContainer.cpp

@@ -1,13 +1,38 @@
 /* Copyright 2022 The AMReX Community
  *
- * Authors: Ryan Sandberg
+ * Authors: Ryan Sandberg, Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
 #include "ParticleContainer.H"
 
 #include <AMReX_Particle.H>
 
+#include <cstdint>
 
+
+namespace
+{
+    using namespace amrex;
+
+    // Note - this function MUST be consistent with AMReX_Particle.H
+    Long unpack_id (uint64_t idcpu) {
+        Long r = 0;
+
+        uint64_t sign = idcpu >> 63;  // extract leftmost sign bit
+        uint64_t val  = ((idcpu >> 24) & 0x7FFFFFFFFF);  // extract next 39 id bits
+
+        Long lval = static_cast<Long>(val);  // bc we take -
+        r = (sign) ? lval : -lval;
+        return r;
+    }
+
+    // Note - this function MUST be consistent with AMReX_Particle.H
+    int unpack_cpu (uint64_t idcpu) {
+        return static_cast<int>(idcpu & 0x00FFFFFF);
+    }
+}
+
+// forward declarations
 void init_ParticleContainer_HiPACE(py::module& m);
 void init_ParticleContainer_ImpactX(py::module& m);
 void init_ParticleContainer_WarpX(py::module& m);
@@ -17,9 +42,24 @@ void init_ParticleContainer(py::module& m) {
 
     // 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<Particle<1, 1>, 2, 1>(m);  // tests
+
+    // most common case: ND particle + runtime attributes
+#if AMREX_SPACEDIM == 1
+    make_ParticleContainer_and_Iterators<SoAParticle<1, 0>, 1, 0>(m);
+#elif AMREX_SPACEDIM == 2
+    make_ParticleContainer_and_Iterators<SoAParticle<2, 0>, 2, 0>(m);
+#elif AMREX_SPACEDIM == 3
+    make_ParticleContainer_and_Iterators<SoAParticle<3, 0>, 3, 0>(m);
+#endif
+
+    // used in tests
+    make_ParticleContainer_and_Iterators<Particle<2, 1>, 3, 1>(m);
 
     init_ParticleContainer_HiPACE(m);
     init_ParticleContainer_ImpactX(m);
     init_ParticleContainer_WarpX(m);
+
+    // for particle idcpu arrays
+    m.def("unpack_ids", py::vectorize(unpack_id));
+    m.def("unpack_cpus", py::vectorize(unpack_cpu));
 }

src/Particle/ParticleTile.cpp → src/Particle/ParticleTile.H

@@ -1,8 +1,10 @@
 /* Copyright 2022 The AMReX Community
  *
- * Authors: Ryan Sandberg
+ * Authors: Ryan Sandberg, Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
+#pragma once
+
 #include "pyAMReX.H"
 
 #include <AMReX_BoxArray.H>
@@ -14,10 +16,6 @@
 #include <sstream>
 
 
-//Forward declaration
-template <int T_NReal, int T_NInt=0>
-void make_Particle(py::module &m);
-
 template <typename T_ParticleType, int NArrayReal, int NArrayInt>
 void make_ParticleTileData(py::module &m)
 {
@@ -193,36 +191,3 @@ void make_ParticleTile(py::module &m)
                       amrex::AsyncArenaAllocator>(m, "async");
 #endif
 }
-
-void init_ParticleTile(py::module& m) {
-    using namespace amrex;
-
-    // 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>;
-    using SoAParticleType_7_0 = SoAParticle<7, 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
-#if AMREX_SPACEDIM == 1
-    make_ParticleTile<SoAParticleType_5_0, 5, 0> (m);   // WarpX 24.03+ 1D
-#elif AMREX_SPACEDIM == 2
-    make_ParticleTile<SoAParticleType_6_0, 6, 0> (m);   // WarpX 24.03+ 2D
-    make_ParticleTile<SoAParticleType_7_0, 7, 0> (m);   // WarpX 24.03+ RZ
-#elif AMREX_SPACEDIM == 3
-    make_ParticleTile<SoAParticleType_7_0, 7, 0> (m);   // WarpX 24.03+ 3D
-#endif
-    make_ParticleTile<SoAParticleType_8_0, 8, 0> (m);   // ImpactX 24.03+
-    make_ParticleTile<ParticleType_0_0, 37, 1> (m);  // HiPACE++ 22.07
-}

src/Particle/StructOfArrays.cpp → src/Particle/StructOfArrays.H

@@ -1,8 +1,10 @@
 /* Copyright 2022 The AMReX Community
  *
- * Authors: Ryan Sandberg
+ * Authors: Ryan Sandberg, Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
+#pragma once
+
 #include "pyAMReX.H"
 
 #include <AMReX_GpuAllocators.H>
@@ -99,19 +101,3 @@ void make_StructOfArrays(py::module &m)
     make_StructOfArrays<NReal, NInt, amrex::AsyncArenaAllocator, use64BitIdCpu>(m, "async");
 #endif
 }
-
-void init_StructOfArrays(py::module& m) {
-    make_StructOfArrays< 2, 1>(m);
-    make_StructOfArrays< 4, 0>(m);  // HiPACE++ 22.08 - 24.02
-    make_StructOfArrays< 5, 0>(m);  // ImpactX 22.07 - 24.02
-#if AMREX_SPACEDIM == 1
-    make_StructOfArrays< 5, 0, true>(m);  // WarpX 24.03+ 1D
-#elif AMREX_SPACEDIM == 2
-    make_StructOfArrays< 6, 0, true>(m);  // WarpX 24.03+ 2D
-    make_StructOfArrays< 7, 0, true>(m);  // WarpX 24.03+ RZ
-#elif AMREX_SPACEDIM == 3
-    make_StructOfArrays< 7, 0, true>(m);  // WarpX 24.03+ 3D
-#endif
-    make_StructOfArrays< 8, 0, true>(m);  // ImpactX 24.03+
-    make_StructOfArrays<37, 1>(m);  // HiPACE++ 22.09 - 24.02
-}

src/pyAMReX.cpp

@@ -31,10 +31,6 @@ void init_AmrMesh(py::module &);
 void init_MultiFab(py::module &);
 void init_ParallelDescriptor(py::module &);
 void init_ParmParse(py::module &);
-void init_Particle(py::module &);
-void init_StructOfArrays(py::module &);
-void init_ArrayOfStructs(py::module &);
-void init_ParticleTile(py::module &);
 void init_ParticleContainer(py::module &);
 void init_Periodicity(py::module &);
 void init_PlotFileUtil(py::module &);
@@ -107,11 +103,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) {
     init_FArrayBox(m);
     init_MultiFab(m);
     init_ParallelDescriptor(m);
-    init_Particle(m);
     init_PODVector(m);
-    init_StructOfArrays(m);
-    init_ArrayOfStructs(m);
-    init_ParticleTile(m);
 
     init_ParticleContainer(m);
     init_AmrMesh(m);