Particle Iterator Shortcuts (#497)

* Allow to iterate on particles of level N without a level loop.
* Allow to access named pure SoA attributes directly on the particle
iterator (`pti`).

Author: ax3l

docs/source/usage/compute.rst

@@ -92,22 +92,31 @@ Here is the general structure for computing on particles:
 
       .. tab-set::
 
-         .. tab-item:: Simple: Pandas (read-only)
+         .. tab-item:: Simple
 
             .. literalinclude:: ../../../tests/test_particleContainer.py
                :language: python3
                :dedent: 4
-               :start-after: # Manual: Pure SoA Compute PC Pandas START
-               :end-before: # Manual: Pure SoA Compute PC Pandas END
+               :start-after: # Manual: Pure SoA Compute PC Simple pti START
+               :end-before: # Manual: Pure SoA Compute PC Simple pti END
 
-         .. tab-item:: Detailed (read and write)
+         .. tab-item:: Detailed
 
             .. literalinclude:: ../../../tests/test_particleContainer.py
                :language: python3
                :dedent: 4
                :start-after: # Manual: Pure SoA Compute PC Detailed START
                :end-before: # Manual: Pure SoA Compute PC Detailed END
 
+         .. tab-item:: Pandas (read-only)
+
+            .. literalinclude:: ../../../tests/test_particleContainer.py
+               :language: python3
+               :dedent: 4
+               :start-after: # Manual: Pure SoA Compute PC Pandas START
+               :end-before: # Manual: Pure SoA Compute PC Pandas END
+
+
    .. tab-item:: Legacy (AoS + SoA) Layout
 
       .. literalinclude:: ../../../tests/test_particleContainer.py

src/Particle/ParticleContainer.H

@@ -477,9 +477,9 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
 
     // simpler particle iterator loops: return types of this particle box
     py_pc
-        .def_property_readonly_static("iterator", [](py::object /* pc */){ return py::type::of<iterator>(); },
+        .def_property_readonly_static("Iterator", [](py::object /* pc */){ return py::type::of<iterator>(); },
             "amrex iterator for particle boxes")
-        .def_property_readonly_static("const_iterator", [](py::object /* pc */){ return py::type::of<const_iterator>(); },
+        .def_property_readonly_static("ConstIterator", [](py::object /* pc */){ return py::type::of<const_iterator>(); },
             "amrex constant iterator for particle boxes (read-only)")
     ;
 }

src/Particle/StructOfArrays.H

@@ -55,6 +55,14 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
              py::return_value_policy::reference_internal,
              py::arg("index"),
              "Get access to a particle Real component Array (compile-time and runtime component)")
+        .def("get_real_data", py::overload_cast<std::string const &>(&SOAType::GetRealData),
+             py::return_value_policy::reference_internal,
+             py::arg("name"),
+             "Get access to a particle Real component Array (compile-time and runtime component)")
+        .def("get_int_data", py::overload_cast<std::string const &>(&SOAType::GetIntData),
+             py::return_value_policy::reference_internal,
+             py::arg("name"),
+             "Get access to a particle Real component Array (compile-time and runtime component)")
 
         // names
         .def_property_readonly("real_names",

src/amrex/extensions/Iterator.py

@@ -38,3 +38,18 @@ def next(self):
         raise StopIteration
 
     return self
+
+
+def getitem(self, name):
+    """Access (read/write) particle vectors."""
+    if not self.is_soa_particle:
+        raise ValueError("Only pure SoA particle containers support pti.__get__")
+
+    if name == "idcpu":
+        return self.soa().get_idcpu_data().to_xp(copy=False)
+    elif name in self.soa().real_names:
+        return self.soa().get_real_data(name).to_xp(copy=False)
+    elif name in self.soa().int_names:
+        return self.soa().get_int_data(name).to_xp(copy=False)
+    else:
+        raise KeyError(f"Unknown particle attribute name: {name}")

src/amrex/extensions/ParticleContainer.py

@@ -6,7 +6,65 @@
 License: BSD-3-Clause-LBNL
 """
 
-from .Iterator import next
+import warnings
+
+from .Iterator import getitem, next
+
+
+def iterator(self, *args, level=None):
+    """Create an iterator over all particle tiles
+
+    Parameters
+    ----------
+    self : amrex.ParticleContainer_*
+        A ParticleContainer class in pyAMReX
+    args : deprecated positional argument
+    level : int | str, optional
+        The MR level. Allowed values are [0:self.finest_level+1) and "all".
+        If there is more than one MR level, the argument is required.
+
+    Returns
+    -------
+    Iterator over all particle tiles at the specified level.
+
+    Examples
+    --------
+    >>> pc.iterator(level="all")
+    >>> pc.iterator(level=0)  # only particles on the the coarsest MR level
+    """
+    # Warn if a second positional argument is provided (ignored argument)
+    if len(args) > 0:
+        if len(args) == 1 and isinstance(args[0], int) and level is None:
+            level = args[0]
+        else:
+            warnings.warn(
+                "The second positional argument to iterator() is deprecated and ignored. "
+                "Please update your code to use iterator(self, level=...) instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+    has_mr = self.finest_level > 0
+
+    if level is None:
+        if has_mr:
+            raise ValueError(
+                "level must be specified for multi-level ParticleContainers"
+            )
+        else:
+            level = 0
+
+    if level == "all":
+        raise ValueError("level='all' is not yet supported for ParticleContainers")
+        # TODO: This does not work
+        # for lvl in range(self.finest_level + 1):
+        #     yield self.Iterator(self, level=lvl)
+    elif isinstance(level, int) and level >= 0:
+        return self.Iterator(self, level=level)
+    else:
+        raise ValueError(
+            f"level must be an integer in [0:{self.finest_level + 1}) or 'all', but got: {level}"
+        )
 
 
 def pc_to_df(self, local=True, comm=None, root_rank=0):
@@ -45,7 +103,7 @@ def pc_to_df(self, local=True, comm=None, root_rank=0):
     # local DataFrame(s)
     dfs_local = []
     for lvl in range(self.finest_level + 1):
-        for pti in self.const_iterator(self, level=lvl):
+        for pti in self.const_iterator(level=lvl):
             if pti.size == 0:
                 continue
 
@@ -130,6 +188,7 @@ def register_ParticleContainer_extension(amr):
     ):
         ParIter_type.__next__ = next
         ParIter_type.__iter__ = lambda self: self
+        ParIter_type.__getitem__ = getitem
 
     # register member functions for every ParticleContainer_* type
     for _, ParticleContainer_type in inspect.getmembers(
@@ -138,4 +197,8 @@ def register_ParticleContainer_extension(amr):
         and member.__module__ == amr.__name__
         and member.__name__.startswith("ParticleContainer_"),
     ):
+        ParticleContainer_type.iterator = iterator
+        ParticleContainer_type.const_iterator = (
+            iterator  # TODO: simplified, code duplication
+        )
         ParticleContainer_type.to_df = pc_to_df

tests/test_particleContainer.py

@@ -63,7 +63,7 @@ def particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
 
     # assign some values to runtime components
     for lvl in range(pc.finest_level + 1):
-        for pti in pc.iterator(pc, level=lvl):
+        for pti in pc.iterator(level=lvl):
             soa = pti.soa()
             soa.get_real_data(2).assign(1.2345)
             soa.get_int_data(1).assign(42)
@@ -97,7 +97,7 @@ def soa_particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
 
     # assign some values to runtime components
     for lvl in range(pc.finest_level + 1):
-        for pti in pc.iterator(pc, level=lvl):
+        for pti in pc.iterator(level=lvl):
             soa = pti.soa()
             soa.get_real_data(8).assign(1.2345)
             soa.get_int_data(0).assign(42)
@@ -212,7 +212,7 @@ def test_pc_init():
     # lvl = 0
     for lvl in range(pc.finest_level + 1):
         print(f"at level {lvl}:")
-        for pti in pc.iterator(pc, level=lvl):
+        for pti in pc.iterator(level=lvl):
             print("...")
             assert pti.num_particles == 1
             assert pti.num_real_particles == 1
@@ -243,7 +243,7 @@ def test_pc_init():
 
     # read-only
     for lvl in range(pc.finest_level + 1):
-        for pti in pc.const_iterator(pc, level=lvl):
+        for pti in pc.const_iterator(level=lvl):
             assert pti.num_particles == 1
             assert pti.num_real_particles == 1
             assert pti.num_neighbor_particles == 0
@@ -383,17 +383,17 @@ class Config:
     # iterate over mesh-refinement levels
     for lvl in range(pc.finest_level + 1):
         # loop local tiles of particles
-        for pti in pc.iterator(pc, level=lvl):
+        for pti in pc.iterator(level=lvl):
             # compile-time and runtime attributes
             soa = pti.soa().to_xp()
 
-            # print all particle ids in the chunk
+            # print all particle ids in the tile
             print("idcpu =", soa.idcpu)
 
             x = soa.real["x"]
             y = soa.real["y"]
 
-            # write to all particles in the chunk
+            # write to all particles in the tile
             # note: careful, if you change particle positions, you might need to
             #       redistribute particles before continuing the simulation step
             soa.real["x"][:] = 0.30
@@ -410,6 +410,36 @@ class Config:
                 soa_int[:] = 12
     # Manual: Pure SoA Compute PC Detailed END
 
+    # Manual: Pure SoA Compute PC Simple pti START
+    # code-specific getter function, e.g.:
+    # pc = sim.get_particles()
+    # Config = sim.extension.Config
+
+    # iterate over particles on level 0
+    for pti in pc.iterator(level=0):
+        # print all particle ids in the tile
+        print("idcpu =", pti["idcpu"])
+
+        x = pti["x"]  # this is automatically a cupy or numpy
+        y = pti["y"]  #   array, depending on Config.have_gpu
+
+        # write to all particles in the chunk
+        # note: careful, if you change particle positions, you might need to
+        #       redistribute particles before continuing the simulation step
+        pti["x"][:] = 0.30
+        pti["y"][:] = 0.35
+        pti["z"][:] = 0.40
+
+        pti["a"][:] = x[:] ** 2
+        pti["b"][:] = x[:] + y[:]
+        pti["c"][:] = 0.50
+        # ...
+
+        # int attributes
+        pti["i1"][:] = 12
+        pti["i2"][:] = 13
+    # Manual: Pure SoA Compute PC Simple pti END
+
 
 def test_pc_numpy(particle_container, Npart):
     """Used in docs/source/usage/compute.rst"""
@@ -427,7 +457,7 @@ class Config:
     # iterate over mesh-refinement levels
     for lvl in range(pc.finest_level + 1):
         # loop local tiles of particles
-        for pti in pc.iterator(pc, level=lvl):
+        for pti in pc.iterator(level=lvl):
             # default layout: AoS with positions and idcpu
             # note: not part of the new PureSoA particle container layout
             aos = (

tests/test_plotfileparticledata.py

@@ -65,7 +65,7 @@ def particle_container(Rpart, std_geometry, distmap, boxarr, std_real_box):
     particles_tile_ct = 0
     # assign some values to runtime components
     for lvl in range(pc.finest_level + 1):
-        for pti in pc.iterator(pc, level=lvl):
+        for pti in pc.iterator(level=lvl):
             aos = pti.aos()
             aos_numpy = aos.to_numpy(copy=False)
             for i, p in enumerate(aos_numpy):
@@ -81,7 +81,7 @@ def check_particles_container(pc, reference_particles):
     Checks the contents of `pc` against `reference_particles`
     """
     for lvl in range(pc.finest_level + 1):
-        for i, pti in enumerate(pc.iterator(pc, level=lvl)):
+        for i, pti in enumerate(pc.iterator(level=lvl)):
             aos = pti.aos()
             for p in aos.to_numpy(copy=True):
                 ref = reference_particles[p["idata_0"]]