[WIP] Add an ImpactX interface

beamphysics/particles.py

@@ -3,7 +3,7 @@
 import os
 import pathlib
 from copy import deepcopy
-from typing import Union, Optional, Sequence
+from typing import TYPE_CHECKING, Union, Optional, Sequence
 
 import numpy as np
 from h5py import File, Group
@@ -45,6 +45,9 @@
 from .wakefields import WakefieldBase
 from .writers import pmd_init, write_pmd_bunch
 
+if TYPE_CHECKING:
+    from .interfaces.impactx import ImpactXRefPart
+
 __all__ = [
     "load_bunch_data",  # Re-exported for backwards compatibility
 ]
@@ -1073,6 +1076,29 @@ def to_bmad(self, p0c=None, tref=None):
         """
         return bmad.particlegroup_to_bmad(self, p0c=p0c, tref=tref)
 
+    def to_impactx(self, ref: ImpactXRefPart) -> dict:
+        """
+        Convert to ImpactX fixed-s beam arrays.
+
+        See `beamphysics.interfaces.impactx` for the coordinate conventions;
+        note in particular that `x` and `y` are taken to be relative to the
+        reference particle already.
+
+        Parameters
+        ----------
+        ref : beamphysics.interfaces.impactx.ImpactXRefPart
+            The reference particle the ImpactX coordinates are relative to.
+
+        Returns
+        -------
+        dict
+            Arrays keyed as `ImpactXParticleContainer.to_df()` names them.
+        """
+        # imported here: interfaces.impactx imports ParticleGroup at module level
+        from .interfaces.impactx import particlegroup_to_impactx
+
+        return particlegroup_to_impactx(self, ref)
+
     @classmethod
     def from_bmad(cls, bmad_dict):
         """
@@ -1124,6 +1150,65 @@ def from_hdf5(
             data = load_species_data(group, include_time_offset=include_time_offset)
         return cls(data=data)
 
+    @classmethod
+    def from_impactx(
+        cls,
+        path: str | pathlib.Path,
+        iteration: int | None = None,
+        species_name: str = "beam",
+        species: str | None = None,
+        strict: bool = True,
+        ref: ImpactXRefPart | None = None,
+    ) -> ParticleGroup:
+        """
+        Load an ImpactX `BeamMonitor` openPMD file.
+
+        ImpactX is an s-based code, so the result is in z-coordinates: every `z`
+        is zero, the bunch length is a spread in `t`, and the transverse
+        coordinates are relative to the reference particle. See
+        `beamphysics.interfaces.impactx` for the frame conventions and for the
+        reference particle, which this does not return.
+
+        Parameters
+        ----------
+        path : str or pathlib.Path
+            Path to the file ImpactX wrote, e.g. `diags/openPMD/monitor.h5`.
+        iteration : int, optional
+            Which iteration to read. The last one when omitted.
+        species_name : str, optional
+            openPMD species to read. ImpactX always writes `"beam"`.
+        species : str, optional
+            openPMD-beamphysics species name; inferred from the reference
+            particle when omitted.
+        strict : bool, optional
+            Refuse to read a bunch carrying per-particle data ParticleGroup
+            cannot hold (non-zero spin, or a runtime component such as the
+            `s_lost` of ImpactX's `particles_lost` output). Default is True.
+        ref : beamphysics.interfaces.impactx.ImpactXRefPart, optional
+            Reference particle to interpret the coordinates against. Taken from
+            the file when omitted. ImpactX's `particles_lost` output carries a
+            zeroed reference particle, so reading it needs one; see
+            `beamphysics.interfaces.impactx.read_beam_monitor`, which this
+            delegates to, for the caveats.
+
+        Returns
+        -------
+        ParticleGroup
+        """
+        # imported here: interfaces.impactx imports ParticleGroup at module level
+        from .interfaces.impactx import read_beam_monitor_data
+
+        return cls(
+            data=read_beam_monitor_data(
+                path,
+                iteration=iteration,
+                species_name=species_name,
+                species=species,
+                strict=strict,
+                ref=ref,
+            )
+        )
+
     @classmethod
     def from_genesis4(
         cls,

docs/examples/data/impactx/generate.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""Regenerate the ImpactX BeamMonitor test data in this directory.
+
+Run with an environment that has ImpactX installed, e.g.::
+
+    conda install -c conda-forge impactx
+    python generate.py
+
+It writes ``monitor.h5`` (a two-iteration beam monitor) and
+``particles_lost.h5`` (the lost-particle output, which carries the runtime
+``s_lost`` component). Both are read by ``tests/test_impactx.py``.
+"""
+
+import os
+import pathlib
+import shutil
+import tempfile
+
+from impactx import ImpactX, distribution, elements
+
+HERE = pathlib.Path(__file__).resolve().parent
+
+# ImpactX writes its diags/ tree into the working directory, so run somewhere
+# disposable rather than leaving it next to the data files.
+_cwd = pathlib.Path.cwd()
+_tmp = tempfile.TemporaryDirectory()
+os.chdir(_tmp.name)
+
+sim = ImpactX()
+sim.space_charge = False
+sim.slice_step_diagnostics = False
+sim.particle_lost_diagnostics_backend = "h5"
+sim.init_grids()
+
+# 2 GeV electron beam, as in the ImpactX FODO example
+sim.beam.ref.set_species("electron").set_kin_energy_MeV(2.0e3)
+distr = distribution.Waterbag(
+    lambdaX=3.9984884770e-5,
+    lambdaY=3.9984884770e-5,
+    lambdaT=1.0e-3,
+    lambdaPx=2.6623538760e-5,
+    lambdaPy=2.6623538760e-5,
+    lambdaPt=2.0e-3,
+    muxpx=-0.846574929020762,
+    muypy=0.846574929020762,
+)
+sim.add_particles(1.0e-9, distr, 200)
+
+monitor = elements.BeamMonitor("monitor", backend="h5")
+sim.lattice.extend(
+    [
+        monitor,
+        elements.Drift(name="d1", ds=0.25, nslice=1),
+        # a collimator tight enough to scrape the tails, so that
+        # particles_lost.h5 has a few particles in it
+        elements.Aperture(
+            name="ap",
+            aperture_x=1.55e-4,
+            aperture_y=1.55e-4,
+            shape="rectangular",
+            action="transmit",
+        ),
+        elements.Quad(name="q1", ds=1.0, k=1.0, nslice=1),
+        monitor,
+    ]
+)
+
+sim.track_particles()
+sim.finalize()
+
+for name in ("monitor.h5", "particles_lost.h5"):
+    shutil.copy(pathlib.Path("diags") / "openPMD" / name, HERE / name)
+    print(f"wrote {HERE / name}")
+
+os.chdir(_cwd)
+_tmp.cleanup()

environment.yml

@@ -16,6 +16,7 @@ dependencies:
   - pytest-cov
   - pyyaml
   - pint  # test-only: independent oracle for beamphysics.units
+  - openpmd-api >=0.17  # test-only: reads ImpactX BeamMonitor output
   - jupyterlab>=3
   - nbconvert
   # Documentation

pyproject.toml

@@ -25,7 +25,14 @@ readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 
 [project.optional-dependencies]
-dev = ["pytest", "pytest-cov", "pytest-benchmark", "pyyaml", "pint"]
+dev = [
+  "pytest",
+  "pytest-cov",
+  "pytest-benchmark",
+  "pyyaml",
+  "pint",           # test-only: independent oracle for beamphysics.units
+  "openpmd-api>=0.17",  # test-only: reads ImpactX BeamMonitor output
+]
 doc = [
   "mkdocs",
   "mkdocs-jupyter",