From 76781d5e6d3bb6168fa55834d9d311828adb1626 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 22 Sep 2025 15:52:49 -0400 Subject: [PATCH 1/2] Add BunchHistogram diagnostic --- .gitignore | 1 + py/orbit/diagnostics/__init__.py | 5 + py/orbit/diagnostics/histogram.py | 205 ++++++++++++++++++++++++++++++ py/orbit/diagnostics/meson.build | 3 +- 4 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 py/orbit/diagnostics/histogram.py diff --git a/.gitignore b/.gitignore index 47650c0e..9291a157 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ PyORBIT.egg-info *.so .*.swp .eggs +.ipynb_checkpoints diff --git a/py/orbit/diagnostics/__init__.py b/py/orbit/diagnostics/__init__.py index c1d197a3..4f12a0ac 100644 --- a/py/orbit/diagnostics/__init__.py +++ b/py/orbit/diagnostics/__init__.py @@ -13,6 +13,10 @@ from .TeapotDiagnosticsNode import TeapotStatLatsNode, TeapotStatLatsNodeSetMember from .TeapotDiagnosticsNode import TeapotMomentsNode, TeapotMomentsNodeSetMember from .TeapotDiagnosticsNode import TeapotTuneAnalysisNode +from .histogram import BunchHistogram +from .histogram import BunchHistogram1D +from .histogram import BunchHistogram2D +from .histogram import BunchHistogram3D __all__ = [] @@ -30,3 +34,4 @@ __all__.append("addTeapotMomentsNodeSet") __all__.append("TeapotTuneAnalysisNode") __all__.append("profiles") +__all__.append("histogram") diff --git a/py/orbit/diagnostics/histogram.py b/py/orbit/diagnostics/histogram.py new file mode 100644 index 00000000..e8ee3979 --- /dev/null +++ b/py/orbit/diagnostics/histogram.py @@ -0,0 +1,205 @@ +import os +import sys +import time +from typing import Any +from typing import Callable +from typing import Union + +import numpy as np +import xarray as xr + +from orbit.core import orbit_mpi +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import Grid1D +from orbit.core.spacecharge import Grid2D +from orbit.core.spacecharge import Grid3D +from orbit.lattice import AccLattice +from orbit.lattice import AccNode + + +Grid = Union[Grid1D, Grid2D, Grid3D] + + +def get_grid_points(grid_coords: list[np.ndarray]) -> np.ndarray: + if len(grid_coords) == 1: + return grid_coords[0] + return np.vstack([c.ravel() for c in np.meshgrid(*grid_coords, indexing="ij")]).T + + +def grid_edges_to_coords(grid_edges: np.ndarray) -> np.ndarray: + return 0.5 * (grid_edges[:-1] + grid_edges[1:]) + + +def make_grid(shape: tuple[int, ...], limits: list[tuple[float, float]]) -> Grid: + + ndim = len(shape) + + grid = None + if ndim == 1: + grid = Grid1D(shape[0] + 1, limits[0][0], limits[0][1]) + elif ndim == 2: + grid = Grid2D( + shape[0] + 1, + shape[1] + 1, + limits[0][0], + limits[0][1], + limits[1][0], + limits[1][1], + ) + elif ndim == 3: + grid = Grid3D( + shape[0] + 1, + shape[1] + 1, + shape[2] + 1, + ) + grid.setGridX(limits[0][0], limits[0][1]) + grid.setGridY(limits[1][0], limits[1][1]) + grid.setGridZ(limits[2][0], limits[2][1]) + else: + raise ValueError + + return grid + + +class BunchHistogram: + """MPI-compatible bunch histogram.""" + def __init__( + self, + axis: tuple[int, ...], + shape: tuple[int, ...], + limits: list[tuple[float, float]], + method: str = None, + transform: Callable = None, + normalize: bool = True, + output_dir: str = None, + verbose: int = 2, + **kwargs + ) -> None: + """Constructor. + + Args: + axis: Axis on which to compute the histogram. + shape: Number of bins along each axis. + limits: Min/max coordinates along each axis. + method: Smoothing method {"bilinear", "nine-point", None}. + transform: Transforms bunch before histogram is calculated. + Call signature is `bunch_new = transform(bunch)`. + normalize: Whehter to normalize values to PDF. + output_dir: Output directory for saved files. + verbose: Whether to print update messages. + """ + self.mpi_comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD + self.mpi_rank = orbit_mpi.MPI_Comm_rank(self.mpi_comm) + self.output_dir = output_dir + self.verbose = verbose + + self.axis = axis + self.ndim = len(axis) + self.method = method + self.transform = transform + self.normalize = normalize + + self.index = 0 # number of calls to `track` method + self.node = None + + if self.ndim > 2: + raise NotImplementedError( + "BunchHistogram does not yet support 3D grids. See " + "https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/46" + " and " + "https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/47" + ) + + # Dimension names + self.dims = ["x", "xp", "y", "yp", "z", "dE"] + self.dims = [self.dims[i] for i in self.axis] + + # Create grid + self.grid_shape = shape + self.grid_limits = limits + self.grid_edges = [ + np.linspace(self.grid_limits[i][0], self.grid_limits[i][1], self.grid_shape[i] + 1) + for i in range(self.ndim) + ] + self.grid_coords = [grid_edges_to_coords(e) for e in self.grid_edges] + self.grid_values = np.zeros(shape) + self.grid_points = get_grid_points(self.grid_coords) + self.grid = make_grid(self.grid_shape, self.grid_limits) + + # Store cell volume for normalization + self.cell_volume = np.prod([e[1] - e[0] for e in self.grid_edges]) + + def sync_mpi(self) -> None: + self.grid.synchronizeMPI(self.mpi_comm) + + def bin_bunch(self, bunch: Bunch) -> None: + macrosize = bunch.macroSize() + if macrosize == 0: + bunch.macroSize(1.0) + + if self.method == "bilinear": + self.grid.binBunchBilinear(bunch, *self.axis) + else: + self.grid.binBunch(bunch, *self.axis) + + bunch.macroSize(macrosize) + + def compute_histogram(self, bunch: Bunch) -> np.ndarray: + self.bin_bunch(bunch) + self.sync_mpi() + + values = np.zeros(self.grid_points.shape[0]) + if self.method == "bilinear": + for i, point in enumerate(self.grid_points): + values[i] = self.grid.getValueBilinear(*point) + elif self.method == "nine-point": + for i, point in enumerate(self.grid_points): + values[i] = self.grid.getValue(*point) + else: + for i, indices in enumerate(np.ndindex(*self.grid_shape)): + values[i] = self.grid.getValueOnGrid(*indices) + values = np.reshape(values, self.grid_shape) + + if self.normalize: + values_sum = np.sum(values) + if values_sum > 0.0: + values /= values_sum + values /= self.cell_volume + return values + + def track(self, bunch: Bunch) -> None: + bunch_copy = Bunch() + bunch.copyBunchTo(bunch_copy) + if self.transform is not None: + bunch_copy = self.transform(bunch_copy) + + self.grid.setZero() + self.grid_values = self.compute_histogram(bunch_copy) + + if self.output_dir is not None: + array = xr.DataArray(self.grid_values, coords=self.grid_coords, dims=self.dims) + array.to_netcdf(path=self.get_filename()) + + self.index += 1 + + def get_filename(self) -> str: + filename = "hist_" + "-".join([str(i) for i in self.axis]) + filename = "{}_{:04.0f}".format(filename, self.index) + filename = "{}.nc".format(filename) + filename = os.path.join(self.output_dir, filename) + return filename + + +class BunchHistogram1D(BunchHistogram): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + + +class BunchHistogram2D(BunchHistogram): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + + +class BunchHistogram3D(BunchHistogram): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) diff --git a/py/orbit/diagnostics/meson.build b/py/orbit/diagnostics/meson.build index 26e40c8e..f355d8dc 100644 --- a/py/orbit/diagnostics/meson.build +++ b/py/orbit/diagnostics/meson.build @@ -6,7 +6,8 @@ py_sources = files([ 'TeapotDiagnosticsNode.py', 'diagnosticsLatticeModifications.py', '__init__.py', - 'profiles.py' + 'profiles.py', + 'histogram.py', ]) python.install_sources( From 5b5615af6b3ec9dc80b230858e00b446d3e41b23 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 22 Sep 2025 15:53:28 -0400 Subject: [PATCH 2/2] Add histogram diagnostic examples --- examples/Diagnostics/Histogram/style.mplstyle | 7 ++ .../Diagnostics/Histogram/test_hist_1d.py | 75 ++++++++++++++++ .../Diagnostics/Histogram/test_hist_2d.py | 85 +++++++++++++++++++ 3 files changed, 167 insertions(+) create mode 100644 examples/Diagnostics/Histogram/style.mplstyle create mode 100644 examples/Diagnostics/Histogram/test_hist_1d.py create mode 100644 examples/Diagnostics/Histogram/test_hist_2d.py diff --git a/examples/Diagnostics/Histogram/style.mplstyle b/examples/Diagnostics/Histogram/style.mplstyle new file mode 100644 index 00000000..ecd17ad7 --- /dev/null +++ b/examples/Diagnostics/Histogram/style.mplstyle @@ -0,0 +1,7 @@ +axes.linewidth: 1.25 +figure.constrained_layout.use: True +xtick.minor.visible: True +ytick.minor.visible: True + +savefig.format: "png" +savefig.dpi: 300 \ No newline at end of file diff --git a/examples/Diagnostics/Histogram/test_hist_1d.py b/examples/Diagnostics/Histogram/test_hist_1d.py new file mode 100644 index 00000000..3cdc51c2 --- /dev/null +++ b/examples/Diagnostics/Histogram/test_hist_1d.py @@ -0,0 +1,75 @@ +import argparse +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core import orbit_mpi +from orbit.core.bunch import Bunch +from orbit.diagnostics import BunchHistogram +from orbit.diagnostics import BunchHistogram1D +from orbit.diagnostics import BunchHistogram2D +from orbit.diagnostics import BunchHistogram3D +from orbit.bunch_utils import collect_bunch + +plt.style.use("style.mplstyle") + + +# Parse args +parser = argparse.ArgumentParser() +parser.add_argument("--nsamp", type=int, default=10_000) +parser.add_argument("--nbins", type=int, default=64) +args = parser.parse_args() + +# Make output directory +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + +# Setup MPI +_mpi_comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD +_mpi_rank = orbit_mpi.MPI_Comm_rank(_mpi_comm) + +# Generate particle distribution +rng = np.random.default_rng(123) +X = rng.normal(size=(args.nsamp, 2)) +X = X / np.linalg.norm(X, axis=1)[:, None] +X[:, 0] -= 0.75 * X[:, 1] ** 2 +X = X + rng.normal(size=X.shape, scale=0.25) +X = X / np.std(X, axis=0) +X = np.hstack([X, np.zeros((X.shape[0], 4))]) + +# Create bunch +bunch = Bunch() +for i in range(X.shape[0]): + bunch.addParticle(*X[i, :]) + +# Create histogram diagnostic +grid_limits = [(-4.0, 4.0)] +grid_shape = (args.nbins,) + +axis = (0,) +diag = BunchHistogram( + axis=axis, + shape=grid_shape, + limits=grid_limits, + method=None, + normalize=True, + output_dir=output_dir, +) + +# Compute histogram +diag.track(bunch) + +# Plot histogram +if _mpi_rank == 0: + fig, ax = plt.subplots(figsize=(5, 2)) + ax.plot(diag.grid_coords[0], diag.grid_values, lw=1.5, color=None, label="PyORBIT") + grid_values_np, _ = np.histogram(X[:, axis], bins=diag.grid_edges[0], density=True) + ax.plot(diag.grid_coords[0], grid_values_np, lw=1.5, color=None, label="NumPy") + ax.legend() + + filename = os.path.join(output_dir, "fig_hist") + plt.savefig(filename) + plt.show() diff --git a/examples/Diagnostics/Histogram/test_hist_2d.py b/examples/Diagnostics/Histogram/test_hist_2d.py new file mode 100644 index 00000000..aa2f54de --- /dev/null +++ b/examples/Diagnostics/Histogram/test_hist_2d.py @@ -0,0 +1,85 @@ +import argparse +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core import orbit_mpi +from orbit.core.bunch import Bunch +from orbit.diagnostics import BunchHistogram +from orbit.diagnostics import BunchHistogram1D +from orbit.diagnostics import BunchHistogram2D +from orbit.diagnostics import BunchHistogram3D +from orbit.bunch_utils import collect_bunch + +plt.style.use("style.mplstyle") + + +# Parse args +parser = argparse.ArgumentParser() +parser.add_argument("--nsamp", type=int, default=10_000) +parser.add_argument("--nbins", type=int, default=64) +args = parser.parse_args() + +# Make output directory +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + +# Setup MPI +_mpi_comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD +_mpi_rank = orbit_mpi.MPI_Comm_rank(_mpi_comm) + +# Generate particle distribution +rng = np.random.default_rng(123) +X = rng.normal(size=(args.nsamp, 2)) +X = X / np.linalg.norm(X, axis=1)[:, None] +X[:, 0] -= 0.75 * X[:, 1] ** 2 +X = X + rng.normal(size=X.shape, scale=0.25) +X = X / np.std(X, axis=0) +X = np.hstack([X, np.zeros((X.shape[0], 4))]) + +# Create bunch +bunch = Bunch() +for i in range(X.shape[0]): + bunch.addParticle(*X[i, :]) + +# Create histogram diagnostic +grid_limits = 2 * [(-4.0, 4.0)] +grid_shape = (args.nbins, args.nbins) + +axis = (0, 1) +diag = BunchHistogram( + axis=axis, + shape=grid_shape, + limits=grid_limits, + method=None, + normalize=True, + output_dir=output_dir, +) + +# Compute histogram +diag.track(bunch) + +# Plot histogram +if _mpi_rank == 0: + fig, axs = plt.subplots(ncols=2, figsize=(6.0, 3.0), sharex=True, sharey=True) + axs[0].pcolormesh( + diag.grid_coords[0], + diag.grid_coords[1], + diag.grid_values.T, + ) + + grid_values_np, _ = np.histogramdd(X[:, axis], bins=diag.grid_edges, density=True) + axs[1].pcolormesh( + diag.grid_coords[0], + diag.grid_coords[1], + grid_values_np.T, + ) + axs[0].set_title("BunchHistogram", fontsize="medium") + axs[1].set_title("NumPy", fontsize="medium") + + filename = os.path.join(output_dir, "fig_hist") + plt.savefig(filename) + plt.show()