Particles: Restructure

The AMReX particle logic is 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

@@ -218,19 +218,19 @@ Additional runtime attributes (Real or Int) are always in SoA memory layout.
 
 ``amrex::ParticleTile<T_ParticleType, NArrayReal, 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.ParticleTile_pureSoA_0_0_8_0_default
+.. autoclass:: amrex.space3d.ParticleTile_pureSoA_8_0_default
    :members:
    :undoc-members:
 
 ``amrex::ParticleTileData<T_ParticleType, NArrayReal>`` is implemented for both legacy (AoS+SoA) and pure SoA particle types, many number of Real and Int arguments, e.g.,
 
-.. autoclass:: amrex.space3d.ParticleTileData_pureSoA_0_0_8_0
+.. autoclass:: amrex.space3d.ParticleTileData_pureSoA_8_0
    :members:
    :undoc-members:
 
 ``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:
 
@@ -262,12 +262,12 @@ This is for the legacy, AoS + SoA particle containers only:
 
 ``amrex::ArrayOfStructs<T_ParticleType, Allocator>`` is implemented for many numbers of extra Real and Int arguments, and allocators, e.g.,
 
-.. autoclass:: amrex.space3d.ArrayOfStructs_0_0_default
+.. autoclass:: amrex.space3d.ArrayOfStructs_2_1_default
    :members:
    :undoc-members:
 
 ``amrex::Particle<T_NReal, T_NInt>`` is implemented for many numbers of extra Real and Int arguments, e.g.,
 
-.. autoclass:: amrex.space3d.Particle_0_0
+.. autoclass:: amrex.space3d.Particle_2_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,25 @@ 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
+    //   pure SoA
+    make_ParticleContainer_and_Iterators<
+        SoAParticle<AMREX_SPACEDIM, 0>,
+                    AMREX_SPACEDIM, 0
+    >(m);
+    //   legacy AoS + SoA
+    //make_ParticleContainer_and_Iterators<Particle<0, 0>, 0, 0>(m);
+
+    // used in tests
+    make_ParticleContainer_and_Iterators<Particle<2, 1>, 3, 1>(m);
+
+    // application codes
     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));
 }