From cd83f001960dd039fbb1a9516513eeda9b23077d Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 3 Nov 2025 15:17:15 -0800 Subject: [PATCH 1/5] Particle Iterator Shortcuts * 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`). --- docs/source/usage/compute.rst | 17 ++++-- src/Particle/ParticleContainer.H | 4 +- src/Particle/StructOfArrays.H | 8 +++ src/amrex/extensions/Iterator.py | 15 ++++++ src/amrex/extensions/ParticleContainer.py | 65 ++++++++++++++++++++++- tests/test_particleContainer.py | 34 +++++++++++- 6 files changed, 134 insertions(+), 9 deletions(-) diff --git a/docs/source/usage/compute.rst b/docs/source/usage/compute.rst index 5e7fced1..8e5fcc62 100644 --- a/docs/source/usage/compute.rst +++ b/docs/source/usage/compute.rst @@ -92,15 +92,15 @@ 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 @@ -108,6 +108,15 @@ Here is the general structure for computing on particles: :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 diff --git a/src/Particle/ParticleContainer.H b/src/Particle/ParticleContainer.H index 867eb73a..7daa27a0 100644 --- a/src/Particle/ParticleContainer.H +++ b/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(); }, + .def_property_readonly_static("Iterator", [](py::object /* pc */){ return py::type::of(); }, "amrex iterator for particle boxes") - .def_property_readonly_static("const_iterator", [](py::object /* pc */){ return py::type::of(); }, + .def_property_readonly_static("ConstIterator", [](py::object /* pc */){ return py::type::of(); }, "amrex constant iterator for particle boxes (read-only)") ; } diff --git a/src/Particle/StructOfArrays.H b/src/Particle/StructOfArrays.H index 1c0c3712..52245094 100644 --- a/src/Particle/StructOfArrays.H +++ b/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(&SOAType::GetRealData), + py::return_value_policy::reference_internal, + py::arg("index"), + "Get access to a particle Real component Array (compile-time and runtime component)") + .def("get_int_data", py::overload_cast(&SOAType::GetIntData), + py::return_value_policy::reference_internal, + py::arg("index"), + "Get access to a particle Real component Array (compile-time and runtime component)") // names .def_property_readonly("real_names", diff --git a/src/amrex/extensions/Iterator.py b/src/amrex/extensions/Iterator.py index 1cf8109a..c57d9fcb 100644 --- a/src/amrex/extensions/Iterator.py +++ b/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}") diff --git a/src/amrex/extensions/ParticleContainer.py b/src/amrex/extensions/ParticleContainer.py index 6e37d00e..0317573d 100644 --- a/src/amrex/extensions/ParticleContainer.py +++ b/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: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:finest_level+1) or 'all', but got: {level}" + ) def pc_to_df(self, local=True, comm=None, root_rank=0): @@ -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 diff --git a/tests/test_particleContainer.py b/tests/test_particleContainer.py index 432e57e2..ab9edff8 100644 --- a/tests/test_particleContainer.py +++ b/tests/test_particleContainer.py @@ -387,13 +387,13 @@ class Config: # 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"] + y = pti["y"] + + # 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""" From 7e0583a7bd61607c27ccb13555e76a9ed5260d7d Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 3 Nov 2025 15:35:42 -0800 Subject: [PATCH 2/5] Tests & Internal: Modernize API Avoid using the deprecated APIs. --- src/amrex/extensions/ParticleContainer.py | 2 +- tests/test_particleContainer.py | 16 ++++++++-------- tests/test_plotfileparticledata.py | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/amrex/extensions/ParticleContainer.py b/src/amrex/extensions/ParticleContainer.py index 0317573d..601e9b27 100644 --- a/src/amrex/extensions/ParticleContainer.py +++ b/src/amrex/extensions/ParticleContainer.py @@ -103,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 diff --git a/tests/test_particleContainer.py b/tests/test_particleContainer.py index ab9edff8..81ebda04 100644 --- a/tests/test_particleContainer.py +++ b/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,7 +383,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): # compile-time and runtime attributes soa = pti.soa().to_xp() @@ -420,8 +420,8 @@ class Config: # print all particle ids in the tile print("idcpu =", pti["idcpu"]) - x = pti["x"] - y = pti["y"] + 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 @@ -457,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 = ( diff --git a/tests/test_plotfileparticledata.py b/tests/test_plotfileparticledata.py index 2b4bce0b..7270eec4 100644 --- a/tests/test_plotfileparticledata.py +++ b/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"]] From a1b22ff4bdd9a63ba18fe1926bed27de8d7751d5 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 4 Nov 2025 09:27:34 -0800 Subject: [PATCH 3/5] Argument: Name Signed-off-by: Axel Huebl --- src/Particle/StructOfArrays.H | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Particle/StructOfArrays.H b/src/Particle/StructOfArrays.H index 52245094..36d05877 100644 --- a/src/Particle/StructOfArrays.H +++ b/src/Particle/StructOfArrays.H @@ -57,11 +57,11 @@ void make_StructOfArrays(py::module &m, std::string allocstr) "Get access to a particle Real component Array (compile-time and runtime component)") .def("get_real_data", py::overload_cast(&SOAType::GetRealData), py::return_value_policy::reference_internal, - py::arg("index"), + py::arg("name"), "Get access to a particle Real component Array (compile-time and runtime component)") .def("get_int_data", py::overload_cast(&SOAType::GetIntData), py::return_value_policy::reference_internal, - py::arg("index"), + py::arg("name"), "Get access to a particle Real component Array (compile-time and runtime component)") // names From d81e0a83f76e000c3410ac6585b2d0c67dfad46e Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 4 Nov 2025 09:28:59 -0800 Subject: [PATCH 4/5] Concrete Number Range Co-authored-by: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com> Signed-off-by: Axel Huebl --- src/amrex/extensions/ParticleContainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/amrex/extensions/ParticleContainer.py b/src/amrex/extensions/ParticleContainer.py index 601e9b27..5ca55fb4 100644 --- a/src/amrex/extensions/ParticleContainer.py +++ b/src/amrex/extensions/ParticleContainer.py @@ -20,7 +20,7 @@ def iterator(self, *args, level=None): A ParticleContainer class in pyAMReX args : deprecated positional argument level : int | str, optional - The MR level. Allowed values are [0:finest_level+1) and "all". + 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 @@ -63,7 +63,7 @@ def iterator(self, *args, level=None): return self.Iterator(self, level=level) else: raise ValueError( - f"level must be an integer in [0:finest_level+1) or 'all', but got: {level}" + f"level must be an integer in [0:{self.finest_level+1}) or 'all', but got: {level}" ) From 75eb275bfeee8b5a0e9dc647c3bdcf481934958f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 4 Nov 2025 17:30:50 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/amrex/extensions/ParticleContainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/amrex/extensions/ParticleContainer.py b/src/amrex/extensions/ParticleContainer.py index 5ca55fb4..044a8cd2 100644 --- a/src/amrex/extensions/ParticleContainer.py +++ b/src/amrex/extensions/ParticleContainer.py @@ -63,7 +63,7 @@ def iterator(self, *args, level=None): 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}" + f"level must be an integer in [0:{self.finest_level + 1}) or 'all', but got: {level}" )