feat: function to collect bunch across MPI ranks into a single Python dictionary #72

woodtp · 2025-08-15T21:18:52Z

This PR adds collect_bunch to orbit.bunch_utils, which accepts a Bunch and collects attributes across MPI ranks into a Python dictionary. The 6 coordinates (x, px, y, py, z, dE) are stored in the dictionary at key "coords" in an np.ndarray with shape (N, 6). In the case of either attributes related to the synchronous particle or those contained in the bunch structure retrieved via Bunch.bunchAttrDouble or Bunch.bunchAttrInt, the bunch dictionary values are stored in their own dictionaries nested at keys "sync_part" and "attributes", respectively, as either np.float64 or np.int32 where appropriate.

Since the output of collect_bunch is a dictionary of numpy objects, the user has the flexibility to pass directly to visualization libs and/or store the output in whichever format they wish. E.g.,

from orbit.bunch_utils import collect_bunch

...

bunch_dict = collect_bunch(bunch)

if mpi_rank == 0:
    # plot x-x'
    sns.histplot(x=bunch_dict["coords"][:, 0], y=bunch_dict["coords"][:, 1])
    plt.savefig("nice_plot.png")
        
    # save flattened dict with np.savez or np.savez_compressed
    np.savez("filename.npz", coords=bunch_dict["coords"], **bunch_dict["sync_part"], **bunch_dict["attributes"])
    
    # Although if you save the coords separately as a `.npy` with `np.save`, you have the ability to 
    # load it later as a mmap, which is probably better for very large bunches.
    np.save("coords.npy", bunch_dict["coords"])
    
    # sometime later...
    coords = np.load("coords.npy", mmap_mode="r")

austin-hoover · 2025-08-22T16:54:22Z

I just tested and it's working for me! One suggestion is to group the entire NX6 coordinate array into one entry called coords or array or something; is that possible?. Also would it make sense to have just one sync_part key and the value be another dictionary? @woodtp @azukov

woodtp · 2025-08-22T17:51:03Z

I just tested and it's working for me! One suggestion is to group the entire NX6 coordinate array into one entry called coords or array or something; is that possible?. Also would it make sense to have just one sync_part key and the value be another dictionary? @woodtp @azukov

I don't see why not, if that's preferable. That layout seems more consistent with how a Bunch stores this information internally.

woodtp · 2025-08-25T21:18:04Z

py/orbit/bunch_utils/collect_bunch.py

+        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)>,
+                    ...
+                }
+            }


@austin-hoover, I incorporated your request. See these lines for the updated list of keys/values:

woodtp · 2025-08-25T22:13:37Z

@azukov correctly pointed out to me that there was an issue with the implementation in the case where the global bunch size wasn't divisible by the total number of MPI nodes.

I have rewritten the utility entirely to resolve that issue while simultaneously avoiding loading the entire bunch into RAM, unless explicitly requested. It now works by creating a np.memmap backed by a file under /tmp, and each node writes to a specific slice of the memory-mapped array, in parallel.

I timed collecting a bunch with 100M particles using the new method, split across 7 nodes (to force an uneven distribution), and here are the results:

# checking that the bunch gets distributed correctly and all `local_rows` add up to 100_000_000
[DEBUG] mpi_rank=0, fname='/tmp/collect_bunch_tmpfile.dat', start_row=0, local_rows=14_285_715
[DEBUG] mpi_rank=2, fname='/tmp/collect_bunch_tmpfile.dat', start_row=28_571_430, local_rows=14_285_714
[DEBUG] mpi_rank=3, fname='/tmp/collect_bunch_tmpfile.dat', start_row=42_857_144, local_rows=14_285_714
[DEBUG] mpi_rank=1, fname='/tmp/collect_bunch_tmpfile.dat', start_row=14_285_715, local_rows=14_285_715
[DEBUG] mpi_rank=4, fname='/tmp/collect_bunch_tmpfile.dat', start_row=57_142_858, local_rows=14_285_714
[DEBUG] mpi_rank=5, fname='/tmp/collect_bunch_tmpfile.dat', start_row=71_428_572, local_rows=14_285_714
[DEBUG] mpi_rank=6, fname='/tmp/collect_bunch_tmpfile.dat', start_row=85_714_286, local_rows=14_285_714

[DEBUG] (mpi_rank=1) COLLECTING BUNCH DONE in 34.941 sec
[DEBUG] (mpi_rank=3) COLLECTING BUNCH DONE in 34.941 sec
[DEBUG] (mpi_rank=6) COLLECTING BUNCH DONE in 34.942 sec
[DEBUG] (mpi_rank=2) COLLECTING BUNCH DONE in 34.944 sec
[DEBUG] (mpi_rank=5) COLLECTING BUNCH DONE in 34.944 sec
[DEBUG] (mpi_rank=4) COLLECTING BUNCH DONE in 34.944 sec
[DEBUG] (mpi_rank=0) COLLECTING BUNCH DONE in 34.990 sec

[DEBUG] coords size in bytes:  4800.0 MB 
[DEBUG] coords shape:  (100000000, 6)

woodtp · 2025-08-26T00:49:19Z

py/orbit/bunch_utils/collect_bunch.py

+    #     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"


Note: Initially, I tried to use the built-in tempfile module to initialize the temporary file for the memmap. When attempting to broadcast the name of the file to the other nodes, it seems like I'm running into an instance of #65... So, I'm hard-coding the name of the file for now.

woodtp · 2025-08-27T14:53:09Z

Ok, it still wasn't working properly after some particles in the bunch were lost and I was getting some fun artifacts when I tried plotting emittance spectra. I was able to fix that by writing bunch coords to separate mem-mapped arrays on each MPI node, then concatenating them on the primary rank at the end of the routine.

I also squashed the previous commits for readability.

azukov · 2025-09-11T15:54:00Z

Please add benchmark for reading the bunch from binary file.

woodtp · 2025-09-11T19:37:45Z

Please add benchmark for reading the bunch from binary file.

@azukov here are some benchmarking results. I realized, after the fact, that I hadn't rebased on main after merging #75, so I included benchmarks with and without that PR as well.

FWIW, the final file sizes for this size bunch are 7.8G and 4.5G for the ASCII and .npy binary format, respectively.

…nary 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. Also introduces a FileHandler protocol, which can define the schema for handling different filetypes, e.g., numpy binaries, HDF5, etc. The desired FileHandler can be passed as an argument to the functions in `collect_bunch.py`

azukov · 2025-09-19T17:58:26Z

@woodtp I think it's all good now.
Maybe we want to merge collect_bunch.py and file_handler.py into serialize.py?

woodtp · 2025-09-22T14:30:53Z

@woodtp I think it's all good now. Maybe we want to merge collect_bunch.py and file_handler.py into serialize.py?

No problem, it's done: 077fbed.

woodtp mentioned this pull request Aug 15, 2025

Binary dump bunch function #37

Open

woodtp commented Aug 25, 2025

View reviewed changes

woodtp commented Aug 26, 2025

View reviewed changes

woodtp force-pushed the feature/py-bunch-collect branch from 46848ca to 79ee8cb Compare August 27, 2025 14:46

woodtp marked this pull request as ready for review August 27, 2025 20:20

woodtp force-pushed the feature/py-bunch-collect branch from 19150c1 to 2de7945 Compare September 2, 2025 15:51

woodtp force-pushed the feature/py-bunch-collect branch from 2de7945 to dc24558 Compare September 11, 2025 20:38

austin-hoover assigned woodtp Sep 21, 2025

austin-hoover added the enhancement New feature or request label Sep 21, 2025

Merge file_handler.py and collect_bunch.py into serialize.py

077fbed

azukov mentioned this pull request Sep 22, 2025

Clean up and sync dependencies for pip/conda, optional requirement of numpy/scipy. #84

Open

Do not fail 'import orbit.bunch_utils' if numpy is not installed

a433036

azukov merged commit b9c1bcb into PyORBIT-Collaboration:main Sep 22, 2025
5 checks passed

woodtp deleted the feature/py-bunch-collect branch September 22, 2025 17:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: function to collect bunch across MPI ranks into a single Python dictionary #72

feat: function to collect bunch across MPI ranks into a single Python dictionary #72

Uh oh!

woodtp commented Aug 15, 2025 •

edited

Loading

Uh oh!

austin-hoover commented Aug 22, 2025

Uh oh!

woodtp commented Aug 22, 2025

Uh oh!

woodtp Aug 25, 2025

Uh oh!

woodtp commented Aug 25, 2025

Uh oh!

woodtp Aug 26, 2025

Uh oh!

woodtp commented Aug 27, 2025

Uh oh!

azukov commented Sep 11, 2025

Uh oh!

woodtp commented Sep 11, 2025

Uh oh!

azukov commented Sep 19, 2025

Uh oh!

woodtp commented Sep 22, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

feat: function to collect bunch across MPI ranks into a single Python dictionary #72

feat: function to collect bunch across MPI ranks into a single Python dictionary #72

Uh oh!

Conversation

woodtp commented Aug 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

austin-hoover commented Aug 22, 2025

Uh oh!

woodtp commented Aug 22, 2025

Uh oh!

woodtp Aug 25, 2025

Choose a reason for hiding this comment

Uh oh!

woodtp commented Aug 25, 2025

Uh oh!

woodtp Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

woodtp commented Aug 27, 2025

Uh oh!

azukov commented Sep 11, 2025

Uh oh!

woodtp commented Sep 11, 2025

Uh oh!

azukov commented Sep 19, 2025

Uh oh!

woodtp commented Sep 22, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

woodtp commented Aug 15, 2025 •

edited

Loading