From 2f9bd92207d60c4ebca658d8c5de15dd57a08ae8 Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Fri, 14 Aug 2026 11:00:30 -0700 Subject: [PATCH 1/8] add `t_offset` argument to `write_pmd_bunch` --- beamphysics/particles.py | 7 ++++-- beamphysics/writers.py | 17 +++++++++++++- tests/test_particlegroup.py | 46 +++++++++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/beamphysics/particles.py b/beamphysics/particles.py index 9f97a8a8..ade4b0c3 100644 --- a/beamphysics/particles.py +++ b/beamphysics/particles.py @@ -1415,7 +1415,7 @@ def write_opal(self, filePath, verbose=False, dist_type="emitted"): return write_opal(self, filePath, verbose=verbose, dist_type=dist_type) # openPMD - def write(self, h5, name=None) -> None: + def write(self, h5, name=None, t_offset=0.0) -> None: """ Write particle data to an HDF5 file or group in openPMD format. @@ -1433,6 +1433,9 @@ def write(self, h5, name=None) -> None: name : str, optional Name for the subgroup/bunch written inside the "particles" group (or provided group). If None, `write_pmd_bunch` will write directly to the "particles" group. + t_offset : float or numpy.ndarray, optional + Time offset, scalar or per-particle, written as the openPMD "timeOffset" + record. Omitted when zero. Default is 0.0. Returns ------- @@ -1465,7 +1468,7 @@ def write(self, h5, name=None) -> None: else: g = h5 - write_pmd_bunch(g, self, name=name) + write_pmd_bunch(g, self, name=name, t_offset=t_offset) # Plotting # -------- diff --git a/beamphysics/writers.py b/beamphysics/writers.py index d5fd6120..1a21b229 100644 --- a/beamphysics/writers.py +++ b/beamphysics/writers.py @@ -38,7 +38,7 @@ def pmd_field_init(h5, externalFieldPath="/ExternalFieldPath/%T/"): h5.attrs[k] = fstr(v) -def write_pmd_bunch(h5, data, name=None): +def write_pmd_bunch(h5, data, name=None, t_offset=0.0): """ Data is a dict with: np.array: 'x', 'px', 'y', 'py', 'z', 'pz', 't', 'status', 'weight' @@ -48,6 +48,9 @@ def write_pmd_bunch(h5, data, name=None): Optional data: np.array: 'id' + t_offset is a scalar or per-particle array written as the 'timeOffset' + record. It is omitted when zero. Readers add this to the 'time' record. + See inverse routine: .particles.load_bunch_data @@ -82,6 +85,18 @@ def write_pmd_bunch(h5, data, name=None): if "id" in data: g["id"] = data["id"] + # Optional time offset, with the same shape as the particle arrays. + t_offset = np.asarray(t_offset, dtype=float) + if np.any(t_offset): + n_particle = data["n_particle"] + if t_offset.ndim == 0: + t_offset = np.broadcast_to(t_offset, (n_particle,)) + elif t_offset.shape != (n_particle,): + raise ValueError( + f"t_offset shape {t_offset.shape} does not match n_particle {n_particle}" + ) + write_component_data(g, "timeOffset", t_offset, unit=pg_units("t")) + def write_pmd_field(h5, data, name=None): """ diff --git a/tests/test_particlegroup.py b/tests/test_particlegroup.py index 0721df4b..e6cdf2bf 100644 --- a/tests/test_particlegroup.py +++ b/tests/test_particlegroup.py @@ -8,6 +8,7 @@ from beamphysics import ParticleGroup from beamphysics.particles import single_particle +from beamphysics.readers import expected_record_unit_dimension, particle_array P = ParticleGroup("docs/examples/data/bmad_particles.h5") @@ -182,6 +183,51 @@ def test_write_reload_h5(tmp_path: pathlib.Path): assert P == P2 +def test_write_t_offset(tmp_path: pathlib.Path): + t_offset = 5e-9 + h5file = tmp_path / "test_offset.h5" + P.write(h5file, t_offset=t_offset) + + with h5py.File(h5file, "r") as fp: + g = fp[f"particles/{P.species}"] + + # Constant component: a group with value and shape + offset = g["timeOffset"] + assert isinstance(offset, h5py.Group) + assert offset.attrs["value"] == t_offset + assert tuple(offset.attrs["shape"]) == (len(P),) + assert offset.attrs["unitSI"] == 1.0 + assert tuple(offset.attrs["unitDimension"]) == tuple( + expected_record_unit_dimension["timeOffset"] + ) + + # The time record itself is not shifted + assert np.allclose(particle_array(g, "t", include_offset=False), P.t) + + # Readers add the offset to t + P2 = ParticleGroup(h5file) + assert np.allclose(P2.t, P.t + t_offset) + assert np.allclose(P2.x, P.x) + assert np.allclose(P2.pz, P.pz) + + +def test_write_t_offset_array(tmp_path: pathlib.Path): + t_offset = np.linspace(0, 1e-9, len(P)) + h5file = tmp_path / "test_offset_array.h5" + P.write(h5file, t_offset=t_offset) + + with h5py.File(h5file, "r") as fp: + assert isinstance(fp[f"particles/{P.species}/timeOffset"], h5py.Dataset) + + assert np.allclose(ParticleGroup(h5file).t, P.t + t_offset) + + +def test_write_t_offset_bad_shape(tmp_path: pathlib.Path): + h5file = tmp_path / "test_offset_bad.h5" + with pytest.raises(ValueError): + P.write(h5file, t_offset=np.zeros(len(P) + 1) + 1e-9) + + def test_fractional_split(): head, tail = P.fractional_split(0.5, "t") head, core, tail = P.fractional_split((0.1, 0.9), "t") From 3d8f9ef5a4f2cafa85cbf2926e6e287027762c69 Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Fri, 14 Aug 2026 11:02:02 -0700 Subject: [PATCH 2/8] convert to numpy docstring --- beamphysics/writers.py | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/beamphysics/writers.py b/beamphysics/writers.py index 1a21b229..68899280 100644 --- a/beamphysics/writers.py +++ b/beamphysics/writers.py @@ -40,20 +40,25 @@ def pmd_field_init(h5, externalFieldPath="/ExternalFieldPath/%T/"): def write_pmd_bunch(h5, data, name=None, t_offset=0.0): """ - Data is a dict with: - np.array: 'x', 'px', 'y', 'py', 'z', 'pz', 't', 'status', 'weight' - str: 'species' - int: n_particle - - Optional data: - np.array: 'id' - - t_offset is a scalar or per-particle array written as the 'timeOffset' - record. It is omitted when zero. Readers add this to the 'time' record. - - See inverse routine: - .particles.load_bunch_data - + Write bunch data in openPMD-beamphysics format. + + Parameters + ---------- + h5 : h5py.File or h5py.Group + Handle to write into. + data : dict or ParticleGroup + Requires keys 'x', 'px', 'y', 'py', 'z', 'pz', 't', 'status', 'weight' + (arrays), 'species' (str), 'n_particle' (int), 'charge' (float). + Optional key: 'id' (array). + name : str, optional + Subgroup to create for the bunch. If None, writes directly into `h5`. + t_offset : float or numpy.ndarray, optional + Time offset, scalar or per-particle, written as the 'timeOffset' record. + Omitted when zero. Readers add this to the 'time' record. Default is 0.0. + + See Also + -------- + beamphysics.particles.load_bunch_data : inverse routine """ if name: g = h5.create_group(name) From 7048233053cb167fdff694a5c29ad8317dee2cb8 Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Fri, 14 Aug 2026 11:52:17 -0700 Subject: [PATCH 3/8] break out include_offset --- beamphysics/particles.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/beamphysics/particles.py b/beamphysics/particles.py index ade4b0c3..f68a50ce 100644 --- a/beamphysics/particles.py +++ b/beamphysics/particles.py @@ -2268,7 +2268,7 @@ def _scalar_maybe_from_array(value): return value[0] -def load_bunch_data(h5): +def load_bunch_data(h5, include_offset=True): """ Load particles into structured numpy array. """ @@ -2296,7 +2296,7 @@ def load_bunch_data(h5): data["total_charge"] = attrs["totalCharge"] * attrs["chargeUnitSI"] for key in ["x", "px", "y", "py", "z", "pz", "t"]: - data[key] = particle_array(h5, key) + data[key] = particle_array(h5, key, include_offset=include_offset) if "particleStatus" in h5: data["status"] = particle_array(h5, "particleStatus") From b3f356671ccb1c404b5f683b04fe5c7d9b05e2a8 Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Fri, 14 Aug 2026 11:55:44 -0700 Subject: [PATCH 4/8] add example of using t_offset --- docs/examples/write_examples.ipynb | 61 ++++++++++++++++++++++++++++-- 1 file changed, 57 insertions(+), 4 deletions(-) diff --git a/docs/examples/write_examples.ipynb b/docs/examples/write_examples.ipynb index 39bff345..98556520 100644 --- a/docs/examples/write_examples.ipynb +++ b/docs/examples/write_examples.ipynb @@ -42,7 +42,11 @@ "source": [ "from beamphysics import ParticleGroup, pmd_init\n", "from h5py import File\n", - "import os" + "import os\n", + "import tempfile\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from beamphysics.particles import load_bunch_data, particle_paths" ] }, { @@ -173,6 +177,55 @@ "P2 == P" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Some codes read in the particle's absolute time from the particle group. This can cause numerical issues when the time offset is much larger than the relative offset of each particle in the bunch. OpenPMD-beamphysics allows writing the additional attribute `timeOffset` to the HDF5 file to avoid this. We demonstrate its use with the (extreme) example of a femtosecond long beam offset by one second in a machine, showing the difference between adding the time offset to the `ParticleGroup` attribute versus writing in the file." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Control the reference time\n", + "ref_time = 1\n", + "bunch_len = 1e-15\n", + "\n", + "# Generate a short bunch to demonstrate writing\n", + "P_short = P.copy()\n", + "P_short.t = np.random.normal(scale=bunch_len, size=len(P))\n", + "\n", + "# Create a beam with naive offset\n", + "P_offset = P_short.copy()\n", + "P_offset.t = P_offset.t + ref_time\n", + "\n", + "# Write the offset beam, read back without offsets (ie what a simulation code separating timeOffset will see)\n", + "with tempfile.NamedTemporaryFile() as f:\n", + " P_offset.write(f.name)\n", + " with File(f.name) as hf:\n", + " P_offset_loaded = ParticleGroup(\n", + " data=load_bunch_data(hf[particle_paths(hf)[0]], include_offset=False)\n", + " )\n", + "\n", + "# Write a beam using the OpenPMD `timeOffset` parameter (note `t_offset` argument in `.write`)\n", + "with tempfile.NamedTemporaryFile() as f:\n", + " P_short.write(f.name, t_offset=ref_time)\n", + " with File(f.name) as hf:\n", + " P_short_loaded = ParticleGroup(\n", + " data=load_bunch_data(hf[particle_paths(hf)[0]], include_offset=False)\n", + " )\n", + "\n", + "# Plot the longitudinal phase spaces against each other\n", + "# Note the numerical roundoff error in the t coordinate\n", + "plt.scatter(1e15 * P_short_loaded[\"delta_t\"], P_short[\"pz\"], s=1)\n", + "plt.scatter(1e15 * P_offset_loaded[\"delta_t\"], P_offset_loaded[\"pz\"], s=1)\n", + "plt.xlabel(r\"$\\Delta t$ (fs)\")\n", + "plt.ylabel(r\"$p_z$ (eV/c)\")" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -957,9 +1010,9 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python (beamphysics-dev)", "language": "python", - "name": "python3" + "name": "beamphysics-dev" }, "language_info": { "codemirror_mode": { @@ -971,7 +1024,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.7" + "version": "3.14.6" } }, "nbformat": 4, From 7bf4511f4337d3d8309ee53a19a853ef6b9adce0 Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Fri, 14 Aug 2026 12:03:41 -0700 Subject: [PATCH 5/8] add legend --- docs/examples/write_examples.ipynb | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/examples/write_examples.ipynb b/docs/examples/write_examples.ipynb index 98556520..8f51db9f 100644 --- a/docs/examples/write_examples.ipynb +++ b/docs/examples/write_examples.ipynb @@ -220,10 +220,13 @@ "\n", "# Plot the longitudinal phase spaces against each other\n", "# Note the numerical roundoff error in the t coordinate\n", - "plt.scatter(1e15 * P_short_loaded[\"delta_t\"], P_short[\"pz\"], s=1)\n", - "plt.scatter(1e15 * P_offset_loaded[\"delta_t\"], P_offset_loaded[\"pz\"], s=1)\n", + "plt.scatter(1e15 * P_short_loaded[\"delta_t\"], P_short[\"pz\"], s=1, label=\"With t_offset\")\n", + "plt.scatter(\n", + " 1e15 * P_offset_loaded[\"delta_t\"], P_offset_loaded[\"pz\"], s=1, label=\"Naive Write\"\n", + ")\n", "plt.xlabel(r\"$\\Delta t$ (fs)\")\n", - "plt.ylabel(r\"$p_z$ (eV/c)\")" + "plt.ylabel(r\"$p_z$ (eV/c)\")\n", + "plt.legend()" ] }, { From 3059f708d68ab5bc7e2d86296504454a2f9b729b Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Mon, 17 Aug 2026 09:56:24 -0700 Subject: [PATCH 6/8] resolve workflow issue --- docs/examples/write_examples.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/examples/write_examples.ipynb b/docs/examples/write_examples.ipynb index 8f51db9f..7f8e8ad8 100644 --- a/docs/examples/write_examples.ipynb +++ b/docs/examples/write_examples.ipynb @@ -1013,9 +1013,9 @@ ], "metadata": { "kernelspec": { - "display_name": "Python (beamphysics-dev)", + "display_name": "Python 3 (ipykernel)", "language": "python", - "name": "beamphysics-dev" + "name": "python3" }, "language_info": { "codemirror_mode": { From fadebd0274ffbc8639c820be09f20883454bebfb Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Wed, 19 Aug 2026 16:08:12 -0700 Subject: [PATCH 7/8] clean up notebook --- docs/examples/write_examples.ipynb | 48 +++++++++++++++++++----------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/docs/examples/write_examples.ipynb b/docs/examples/write_examples.ipynb index 7f8e8ad8..34086cfb 100644 --- a/docs/examples/write_examples.ipynb +++ b/docs/examples/write_examples.ipynb @@ -192,8 +192,15 @@ "source": [ "# Control the reference time\n", "ref_time = 1\n", - "bunch_len = 1e-15\n", - "\n", + "bunch_len = 1e-15" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ "# Generate a short bunch to demonstrate writing\n", "P_short = P.copy()\n", "P_short.t = np.random.normal(scale=bunch_len, size=len(P))\n", @@ -202,22 +209,27 @@ "P_offset = P_short.copy()\n", "P_offset.t = P_offset.t + ref_time\n", "\n", - "# Write the offset beam, read back without offsets (ie what a simulation code separating timeOffset will see)\n", - "with tempfile.NamedTemporaryFile() as f:\n", - " P_offset.write(f.name)\n", - " with File(f.name) as hf:\n", - " P_offset_loaded = ParticleGroup(\n", - " data=load_bunch_data(hf[particle_paths(hf)[0]], include_offset=False)\n", - " )\n", "\n", - "# Write a beam using the OpenPMD `timeOffset` parameter (note `t_offset` argument in `.write`)\n", - "with tempfile.NamedTemporaryFile() as f:\n", - " P_short.write(f.name, t_offset=ref_time)\n", - " with File(f.name) as hf:\n", - " P_short_loaded = ParticleGroup(\n", - " data=load_bunch_data(hf[particle_paths(hf)[0]], include_offset=False)\n", - " )\n", + "def write_read(pg, t_offset=0):\n", + " with tempfile.NamedTemporaryFile() as f:\n", + " pg.write(f.name, t_offset=t_offset)\n", + " with File(f.name) as hf:\n", + " return ParticleGroup(\n", + " data=load_bunch_data(hf[particle_paths(hf)[0]], include_offset=False)\n", + " )\n", "\n", + "\n", + "# Write the beams, read back without offsets (ie what a simulation code separating timeOffset will see)\n", + "P_offset_loaded = write_read(P_offset)\n", + "P_short_loaded = write_read(P_short, t_offset=ref_time)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ "# Plot the longitudinal phase spaces against each other\n", "# Note the numerical roundoff error in the t coordinate\n", "plt.scatter(1e15 * P_short_loaded[\"delta_t\"], P_short[\"pz\"], s=1, label=\"With t_offset\")\n", @@ -1013,9 +1025,9 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python (beamphysics-dev)", "language": "python", - "name": "python3" + "name": "beamphysics-dev" }, "language_info": { "codemirror_mode": { From fc2b82cc39d8162bcd1f83799879eb8db9a35ccb Mon Sep 17 00:00:00 2001 From: "Christopher M. Pierce" Date: Thu, 20 Aug 2026 11:46:16 -0700 Subject: [PATCH 8/8] fix notebook --- docs/examples/write_examples.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/examples/write_examples.ipynb b/docs/examples/write_examples.ipynb index 34086cfb..cd820824 100644 --- a/docs/examples/write_examples.ipynb +++ b/docs/examples/write_examples.ipynb @@ -1025,9 +1025,9 @@ ], "metadata": { "kernelspec": { - "display_name": "Python (beamphysics-dev)", + "display_name": "Python 3 (ipykernel)", "language": "python", - "name": "beamphysics-dev" + "name": "python3" }, "language_info": { "codemirror_mode": {