Implementation of a pure-python function to collect bunch into dictionary of numpy objects

Each MPI node writes the bunch coordinates to a memory-mapped numpy
array in /tmp. The primary rank concatenates them into a single
memory-mapped array, and the extras are removed from disk.

Author: Wood, Tony

py/orbit/bunch_utils/__init__.py

@@ -6,6 +6,7 @@
 #
 
 from orbit.bunch_utils.particleidnumber import ParticleIdNumber
+from orbit.bunch_utils.collect_bunch import collect_bunch
 
 __all__ = []
 __all__.append("addParticleIdNumbers")

py/orbit/bunch_utils/collect_bunch.py

@@ -0,0 +1,160 @@
+import os
+# import tempfile
+from enum import IntEnum
+from typing import Optional, TypedDict
+
+from orbit.core.bunch import Bunch
+from orbit.core import orbit_mpi
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class SyncPartDict(TypedDict):
+    coords: NDArray[np.float64]
+    kin_energy: np.float64
+    momentum: np.float64
+    beta: np.float64
+    gamma: np.float64
+    time: np.float64
+
+
+class BunchDict(TypedDict):
+    coords_array: NDArray[np.float64]
+    sync_part: SyncPartDict
+    attr: dict[str, np.float64 | np.int32]
+
+
+class BunchCoord(IntEnum):
+    X  = 0
+    XP = 1
+    Y  = 2
+    YP = 3
+    Z  = 4
+    DE = 5
+
+
+def collect_bunch(
+    bunch: Bunch,
+    return_memmap: bool = True,
+    output_fname: Optional[str] = None,
+) -> BunchDict | None:
+    """Collects attributes from a PyOrbit Bunch across all MPI ranks and returns it as a dictionary.
+    Parameters
+    ----------
+    bunch : Bunch
+        The PyOrbit::Bunch object from which to collect attributes.
+    return_memmap : bool, optional
+        Return the bunch coordinates as a memory-mapped NumPy array, otherwise the
+        entire array is copied into RAM and returned as normal NDArray. Default is True.
+    Returns
+    -------
+    BunchDict | None
+        A dictionary containing the collected bunch attributes. Returns None if not on the root MPI rank or if the global bunch size is 0.
+        BunchDict structure:
+            {
+                "coords": NDArray[np.float64] of shape (N, 6) where N is the total number of macroparticles,
+                    and the 6 columns correspond to [x, xp, y, yp, z, dE] in units of [m, rad, m, rad, m, eV], respectively.
+                "sync_part": {
+                    "coords": NDArray[np.float64] of shape (3,),
+                    "kin_energy": np.float64,
+                    "momentum": np.float64,
+                    "beta": np.float64,
+                    "gamma": np.float64,
+                    "time": np.float64
+                },
+                "attributes": {
+                    <bunch attribute name>: <attribute value (np.float64 or np.int32)>,
+                    ...
+                }
+            }
+    """
+
+    global_size = bunch.getSizeGlobal()
+
+    if global_size == 0:
+        return None
+
+    mpi_comm = bunch.getMPIComm()
+    mpi_rank = orbit_mpi.MPI_Comm_rank(mpi_comm)
+
+    coords_shape = (bunch.getSizeGlobal(), 6)
+
+    local_rows = bunch.getSize()
+
+    # print(f"[DEBUG] Rank {mpi_rank}: start_row={start_row}, stop_row={stop_row}, local_rows={local_rows} bunch.getSize()={bunch.getSize()}")
+
+    # if mpi_rank == 0:
+    #     file_desc, fname = tempfile.mkstemp(suffix=".dat", prefix="collect_bunch_", dir="/tmp")
+    #     os.close(file_desc)
+    #
+    # TODO: this doesn't seem to work. "SystemError: PY_SSIZE_T_CLEAN macro must be defined for '#' formats"
+    # fname = orbit_mpi.MPI_Bcast(fname, orbit_mpi.mpi_datatype.MPI_CHAR, 0, mpi_comm)
+
+    # Using a fixed filename in the temp directory for now. Maybe that's sufficient.
+    fname = f"/tmp/collect_bunch_tmpfile_{mpi_rank}.dat"
+
+    local_shape = (local_rows, coords_shape[1])
+    dtype = np.float64
+    coords_memmap = np.memmap(fname, dtype=dtype, mode="w+", shape=local_shape)
+
+    for i in range(local_rows):
+        coords_memmap[i, BunchCoord.X] = bunch.x(i)
+        coords_memmap[i, BunchCoord.XP] = bunch.xp(i)
+        coords_memmap[i, BunchCoord.Y] = bunch.y(i)
+        coords_memmap[i, BunchCoord.YP] = bunch.yp(i)
+        coords_memmap[i, BunchCoord.Z] = bunch.z(i)
+        coords_memmap[i, BunchCoord.DE] = bunch.dE(i)
+
+    coords_memmap.flush()
+
+    bunch_dict = {"coords": None, "sync_part": {}, "attributes": {}}
+
+    if mpi_rank == 0:
+        sync_part = bunch.getSyncParticle()
+
+        bunch_dict["sync_part"] |= {
+            "coords": np.array(sync_part.pVector()),
+            "kin_energy": np.float64(sync_part.kinEnergy()),
+            "momentum": np.float64(sync_part.momentum()),
+            "beta": np.float64(sync_part.beta()),
+            "gamma": np.float64(sync_part.gamma()),
+            "time": np.float64(sync_part.time()),
+        }
+
+        for attr in bunch.bunchAttrDoubleNames():
+            bunch_dict["attributes"][attr] = np.float64(bunch.bunchAttrDouble(attr))
+
+        for attr in bunch.bunchAttrIntNames():
+            bunch_dict["attributes"][attr] = np.int32(bunch.bunchAttrInt(attr))
+
+    orbit_mpi.MPI_Barrier(mpi_comm)
+
+    if mpi_rank == 0:
+        coords_memmap = np.memmap(fname, dtype=dtype, mode="r+", shape=coords_shape)
+
+        start_row = local_rows
+
+        for r in range(1, orbit_mpi.MPI_Comm_size(mpi_comm)):
+            src_fname = f"/tmp/collect_bunch_tmpfile_{r}.dat"
+
+            if not os.path.exists(src_fname):
+                raise FileNotFoundError(f"Expected temporary file '{src_fname}' not found. Something went wrong.")
+
+            src_memmap = np.memmap(src_fname, dtype=dtype, mode="r")
+            src_memmap = src_memmap.reshape((-1, coords_shape[1]))
+
+            stop_row = start_row + src_memmap.shape[0]
+
+            coords_memmap[start_row:stop_row, :] = src_memmap[:, :]
+            coords_memmap.flush()
+
+            del src_memmap
+            os.remove(src_fname)
+            start_row = stop_row
+
+        bunch_dict["coords"] = (
+            coords_memmap if return_memmap else np.array(coords_memmap)
+        )
+
+        return bunch_dict

py/orbit/bunch_utils/meson.build

@@ -3,11 +3,12 @@
 
 py_sources = files([
     '__init__.py',
-	'particleidnumber.py'
+    'particleidnumber.py',
+    'collect_bunch.py'
 ])
 
 python.install_sources(
     py_sources,
     subdir: 'orbit/bunch_utils',
     # pure: true,
-)        
+)

tests/py/orbit/bunch_utils/test_collect.py

@@ -0,0 +1,78 @@
+from orbit.core.bunch import Bunch
+from orbit.bunch_generators import GaussDist3D
+from orbit.bunch_utils import collect_bunch
+
+from pytest import fixture
+
+
+@fixture
+def bunch():
+    bunch = Bunch()
+    bunch.mass(0.939294)
+    bunch.charge(-1.0)
+    bunch.getSyncParticle().kinEnergy(0.0025)
+    gauss_dist = GaussDist3D()
+    for i in range(10):
+        x, xp, y, yp, z, dE = gauss_dist.getCoordinates()
+        bunch.addParticle(x, xp, y, yp, z, dE)
+    bunch.macroSize(10)
+    return bunch
+
+
+def test_collect_bunch(bunch):
+    d = collect_bunch(bunch, return_memmap=False)
+
+    n_particles = bunch.getSize()
+
+    toplevel_keys = {"coords", "sync_part", "attributes"}
+
+    attribute_keys = {"charge", "classical_radius", "mass", "macro_size"}
+    sync_part_keys = {"coords", "kin_energy", "momentum", "beta", "gamma", "time"}
+
+    x, xp, y, yp, z, dE = [], [], [], [], [], []
+    for i in range(n_particles):
+        x.append(bunch.x(i))
+        xp.append(bunch.px(i))
+        y.append(bunch.y(i))
+        yp.append(bunch.py(i))
+        z.append(bunch.z(i))
+        dE.append(bunch.dE(i))
+
+    assert set(d.keys()) == toplevel_keys
+    assert set(d["sync_part"].keys()) == sync_part_keys
+    assert set(d["attributes"].keys()) == attribute_keys
+
+    assert d["coords"].shape == (n_particles, 6)
+
+    assert d["attributes"]["charge"] == bunch.bunchAttrDouble("charge")
+    assert d["attributes"]["classical_radius"] == bunch.bunchAttrDouble(
+        "classical_radius"
+    )
+    assert d["attributes"]["mass"] == bunch.bunchAttrDouble("mass")
+    assert d["attributes"]["macro_size"] == bunch.bunchAttrDouble("macro_size")
+
+    sync_part = bunch.getSyncParticle()
+    assert (d["sync_part"]["coords"] == sync_part.pVector()).all()
+    assert d["sync_part"]["kin_energy"] == sync_part.kinEnergy()
+    assert d["sync_part"]["momentum"] == sync_part.momentum()
+    assert d["sync_part"]["beta"] == sync_part.beta()
+    assert d["sync_part"]["gamma"] == sync_part.gamma()
+    assert d["sync_part"]["time"] == sync_part.time()
+
+
+def test_collect_empty_bunch():
+    bunch = Bunch()
+    d = collect_bunch(bunch)
+    assert d is None
+
+
+def test_collect_arbitrary_bunch_attr(bunch):
+    bunch.bunchAttrDouble("arbitrary_dbl_attr", 42.0)
+    bunch.bunchAttrInt("arbitrary_int_attr", 42)
+
+    d = collect_bunch(bunch)
+
+    assert "arbitrary_dbl_attr" in d["attributes"]
+    assert d["attributes"]["arbitrary_dbl_attr"] == 42.0
+    assert "arbitrary_int_attr" in d["attributes"]
+    assert d["attributes"]["arbitrary_int_attr"] == 42