Add t_offset Argument to ParticleGroup.write

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
     # --------
@@ -2265,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.
     """
@@ -2293,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")

beamphysics/writers.py

@@ -38,19 +38,27 @@ 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'
-        str: 'species'
-        int: n_particle
-
-    Optional data:
-        np.array: 'id'
-
-    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)
@@ -82,6 +90,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):
     """

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,70 @@
     "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"
+   ]
+  },
+  {
+   "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",
+    "\n",
+    "# Create a beam with naive offset\n",
+    "P_offset = P_short.copy()\n",
+    "P_offset.t = P_offset.t + ref_time\n",
+    "\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",
+    "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)\")\n",
+    "plt.legend()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -971,7 +1039,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.14.6"
   }
  },
  "nbformat": 4,

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")