AMReX SoA Named Components (#382)

Replace our own named SoA particle components implementation with the
now upstream support in AMReX. Early draft, unfinished.

- [x] add new APIs from AMReX-Codes/amrex#4300
- [x] rebase after #404

---------

Signed-off-by: Axel Huebl <axel.huebl@plasma.ninja>

src/Particle/ParticleContainer.H

@@ -204,6 +204,9 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
         .def_property_readonly("num_position_components", [](const py::object&){ return AMREX_SPACEDIM; })
         .def_property_readonly("byte_spread", &ParticleContainerType::ByteSpread)
 
+        // compile-time components
+        .def("set_soa_compile_time_names", &ParticleContainerType::SetSoACompileTimeNames)
+
         // runtime components
         .def("add_real_comp", py::overload_cast<int>(&ParticleContainerType::AddRealComp),
              py::arg("communicate")=1,
@@ -224,6 +227,33 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
              "add a new runtime component with type Int"
         )
 
+        .def_property_readonly("real_soa_names",
+            &ParticleContainerType::GetRealSoANames,
+            "Get the names for the Real SoA components"
+        )
+        .def_property_readonly("int_soa_names",
+            &ParticleContainerType::GetIntSoANames,
+            "Get the names for the int SoA components"
+        )
+
+        .def("has_real_comp",
+            &ParticleContainerType::HasRealComp,
+            "Check if a container has an ParticleReal component"
+        )
+        .def("has_int_comp",
+            &ParticleContainerType::HasIntComp,
+            "Check if a container has an Integer component"
+        )
+
+        .def("get_real_comp_index",
+            &ParticleContainerType::GetRealCompIndex,
+            "Get the ParticleReal SoA index of a component"
+        )
+        .def("get_int_comp_index",
+            &ParticleContainerType::GetIntCompIndex,
+            "Get the Integer SoA index of a component"
+        )
+
         .def_property_readonly("finest_level", &ParticleContainerBase::finestLevel)
 
         // ParticleContainer ( const ParticleContainer &) = delete;

src/Particle/StructOfArrays.H

@@ -56,6 +56,16 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
              py::arg("index"),
              "Get access to a particle Real component Array (compile-time and runtime component)")
 
+        // names
+        .def_property_readonly("real_names",
+            &SOAType::GetRealNames,
+            "Names for the Real SoA components"
+        )
+        .def_property_readonly("int_names",
+            &SOAType::GetIntNames,
+            "Names for the int SoA components"
+        )
+
         .def("__len__", &SOAType::size,
              "Get the number of particles")
         .def_property_readonly("size", &SOAType::size,

src/amrex/extensions/StructOfArrays.py

@@ -9,71 +9,6 @@
 from collections import namedtuple
 
 
-def soa_real_comps(self, num_comps, spacedim=3, rotate=True):
-    """
-    Name the ParticleReal components in SoA.
-
-    Parameters
-    ----------
-    self : SoA Type
-      maybe unused, depending on implementation
-    num_comps : int
-      number of components to generate names for.
-    spacedim : int
-      AMReX dimensionality
-    rotate : bool = True
-      start with "x", "y", "z", "a", "b", ...
-
-    Returns
-    -------
-    A list of length num_comps with values
-    rotate=True (for pure SoA layout):
-    - 3D: "x", "y", "z", "a", "b", ... "w", "r0", "r1", ...
-    - 2D: "x", "y", "a", "b", ... "w", "r0", "r1", ...
-    - 1D: "x", "a", "b", ... "w", "r0", "r1", ...
-    rotate=False (for legacy layout):
-    - 1D-3D: "a", "b", ... "w", "r0", "r1", ...
-    """
-    import string
-
-    # x, y, z, a, b, ...
-    comp_names = list(string.ascii_lowercase)
-    if rotate:
-        # rotate x, y, z to be beginning (positions)
-        comp_names = comp_names[-3:] + comp_names[:-3]
-    else:
-        # cut off x, y, z to avoid confusion
-        comp_names = comp_names[:-3]
-
-    num_named = len(comp_names)
-    if num_comps < num_named:
-        comp_names = list(comp_names)[0:num_comps]
-    elif num_comps > num_named:
-        comp_names.extend(["r" + str(i) for i in range(num_comps - num_named)])
-
-    return comp_names
-
-
-def soa_int_comps(self, num_comps):
-    """
-    Name the int components in SoA.
-
-    Parameters
-    ----------
-    self : SoA Type
-      maybe unused, depending on implementation
-    num_comps : int
-      number of components to generate names for.
-
-    Returns
-    -------
-    A list of length num_comps with values "i1", "i2", "i3", ...
-    """
-    comp_names = ["i" + str(i) for i in range(num_comps)]
-
-    return comp_names
-
-
 def soa_to_numpy(self, copy=False):
     """
     Provide NumPy views into a StructOfArrays.
@@ -107,18 +42,17 @@ def soa_to_numpy(self, copy=False):
     else:
         soa_view = SoA_np({}, {}, None)
 
-    # for the legacy data layout, do not start with x, y, z but with a, b, c, ...
-    if self.has_idcpu:
-        real_comp_names = self.soa_real_comps(self.num_real_comps)
-    else:
-        real_comp_names = self.soa_real_comps(self.num_real_comps, rotate=False)
-
+    real_comp_names = self.real_names
+    if len(real_comp_names) != self.num_real_comps:
+        raise ValueError("Missing names for SoA Real components.")
     for idx_real in range(self.num_real_comps):
         soa_view.real[real_comp_names[idx_real]] = self.get_real_data(
             idx_real
         ).to_numpy(copy=copy)
 
-    int_comp_names = self.soa_int_comps(self.num_int_comps)
+    int_comp_names = self.int_names
+    if len(int_comp_names) != self.num_int_comps:
+        raise ValueError("Missing names for SoA int components.")
     for idx_int in range(self.num_int_comps):
         soa_view.int[int_comp_names[idx_int]] = self.get_int_data(idx_int).to_numpy(
             copy=copy
@@ -165,18 +99,17 @@ def soa_to_cupy(self, copy=False):
     else:
         soa_view = SoA_cp({}, {}, None)
 
-    # for the legacy data layout, do not start with x, y, z but with a, b, c, ...
-    if self.has_idcpu:
-        real_comp_names = self.soa_real_comps(self.num_real_comps)
-    else:
-        real_comp_names = self.soa_real_comps(self.num_real_comps, rotate=False)
-
+    real_comp_names = self.real_names
+    if len(real_comp_names) != self.num_real_comps:
+        raise ValueError("Missing names for SoA Real components.")
     for idx_real in range(self.num_real_comps):
         soa_view.real[real_comp_names[idx_real]] = self.get_real_data(idx_real).to_cupy(
             copy=copy
         )
 
-    int_comp_names = self.soa_int_comps(self.num_int_comps)
+    int_comp_names = self.int_names
+    if len(int_comp_names) != self.num_int_comps:
+        raise ValueError("Missing names for SoA int components.")
     for idx_int in range(self.num_int_comps):
         soa_view.int[int_comp_names[idx_int]] = self.get_int_data(idx_int).to_cupy(
             copy=copy
@@ -226,10 +159,6 @@ def register_SoA_extension(amr):
         and member.__module__ == amr.__name__
         and member.__name__.startswith("StructOfArrays_"),
     ):
-        # name providers
-        SoA_type.soa_real_comps = soa_real_comps
-        SoA_type.soa_int_comps = soa_int_comps
-
         # converters
         SoA_type.to_numpy = soa_to_numpy
         SoA_type.to_cupy = soa_to_cupy

tests/test_particleContainer.py

@@ -57,9 +57,9 @@ def particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
     pc.init_random(Npart, iseed, myt, False, std_real_box)
 
     # add runtime components: 1 real 2 int
-    pc.add_real_comp(True)
-    pc.add_int_comp(True)
-    pc.add_int_comp(True)
+    pc.add_real_comp("b", True)
+    pc.add_int_comp("i1", True)
+    pc.add_int_comp("i2", True)
 
     # assign some values to runtime components
     for lvl in range(pc.finest_level + 1):
@@ -79,13 +79,21 @@ def soa_particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
     myt.real_array_data = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
     myt.int_array_data = []
 
+    with pytest.raises(Exception):
+        pc.set_soa_compile_time_names(
+            ["x", "y", "z", "z", "b", "c", "d", "e"], []
+        )  # error: z added twice
+    pc.set_soa_compile_time_names(["x", "y", "z", "a", "b", "c", "d", "e"], [])
+
     iseed = 1
     pc.init_random(Npart, iseed, myt, False, std_real_box)
 
     # add runtime components: 1 real 2 int
-    pc.add_real_comp(True)
-    pc.add_int_comp(True)
-    pc.add_int_comp(True)
+    with pytest.raises(Exception):
+        pc.add_real_comp("a", True)  # already used as a compile-time component
+    pc.add_real_comp("f", True)
+    pc.add_int_comp("i1", True)
+    pc.add_int_comp("i2", True)
 
     # assign some values to runtime components
     for lvl in range(pc.finest_level + 1):