diff --git a/beamphysics/interfaces/impactx.py b/beamphysics/interfaces/impactx.py new file mode 100644 index 0000000..8857c15 --- /dev/null +++ b/beamphysics/interfaces/impactx.py @@ -0,0 +1,921 @@ +"""ImpactX beam data <-> ParticleGroup. + +[ImpactX](https://impactx.readthedocs.io) enables high-performance modeling of beam +dynamics in particle accelerators with collective effects. This is the next +generation of the IMPACT-Z code. ImpactX runs on modern GPUs or CPUs alike, provides +user-friendly interfaces suitable for AI/ML workflows, has many benchmarks to ensure +its correctness, and an extensive documentation. + +ImpactX models particle beams with respect to a common ``s`` (in this repo called ``z``) +variable for the reference trajectory, with a spread in arrival time: all ``z`` equal, +``t`` varying, so the conversion is a direct algebraic map, like the Bmad interface and +unlike time-integrating codes (e.g., WarpX, ASTRA, etc.). + +Coordinates and frames +---------------------- +ImpactX describes each particle at fixed ``s`` by ``(x, y, t, px, py, pt)``: + +- ``x``, ``y`` [m] are the transverse displacement from the reference particle, in the + local (curvilinear) frame that follows the reference orbit. +- ``t`` [m] is ``c`` times the difference between the particle's and the reference + particle's arrival time, i.e. a length, not a time. +- ``px``, ``py``, ``pt`` are dimensionless, normalized by the magnitude of the + reference momentum: ``px = Delta(beta_x gamma) / (beta_0 gamma_0)`` and + ``pt = -Delta(gamma) / (beta_0 gamma_0)``. + +See: https://impactx.readthedocs.io/en/latest/theory/coordinates_units.html + +`ParticleGroup` is a lab-frame container, so the mapping has to choose a frame: + +- The transverse coordinates stay in the **local frame**: ``x`` and ``y`` are the + displacement from the reference particle and ``z`` is zero, the reference plane. + Adding ``x_ref``/``z_ref`` would be wrong wherever the reference orbit is bent, + because local ``x`` is then not lab ``x``. Use `ImpactXRefPart` (``x``, ``y``, ``z``, + ``s``, ``px``, ``py``, ``pz``) if you need to place the bunch in the lab. +- The time is **absolute**: ``t = t_ref + position_t / c``. That one is unambiguous, it + is what openPMD's ``position/t + positionOffset/t`` means in ImpactX output, and it + keeps quantities like `ParticleGroup.average_current` meaningful. + +See also: + +- ImpactX source: https://github.com/BLAST-ImpactX/impactx +- ImpactX manual: https://impactx.readthedocs.io +""" + +from __future__ import annotations + +import pathlib +import warnings +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from ..particles import ParticleGroup +from ..species import charge_of, e_charge, mass_of +from ..status import ParticleStatus +from ..units import c_light + + +__all__ = [ + "ImpactXRefPart", + "PARTICLE_STATUS_LOST", + "UnrepresentableParticleData", + "beam_monitor_iterations", + "impactx_to_particlegroup_data", + "particle_id_from_idcpu", + "particlegroup_to_impactx", + "pmd_species_of", + "read_beam_monitor", + "read_beam_monitor_data", + "refpart_from_openpmd", +] + + +#: ImpactX's built-in species names mapped to openPMD-beamphysics names. +#: ImpactX only knows these four; anything else needs explicit mass and charge. +IMPACTX_TO_PMD_SPECIES = { + "electron": "electron", + "positron": "positron", + "proton": "proton", + "Hminus": "H-", +} +PMD_TO_IMPACTX_SPECIES = {v: k for k, v in IMPACTX_TO_PMD_SPECIES.items()} + + +# -------------------------------------------------------------------------------------- +# Reference particle +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ImpactXRefPart: + """An ImpactX reference particle, detached from any live ``ImpactX`` session. + + Holding this as a plain dataclass rather than wrapping ``impactx.RefPart`` is what + lets the converters run with no ImpactX object in the process -- which the openPMD + reader and the whole test suite rely on. + + Attributes + ---------- + x, y, z : float + Lab-frame position of the reference particle, in metres. + t : float + ``c * t`` of the reference particle, in **metres** (ImpactX convention). + px, py, pz : float + Lab-frame momenta normalized by ``m * c``, i.e. ``beta_i * gamma``. + Dimensionless. + pt : float + ``-gamma`` of the reference particle. Dimensionless. + mass_MeV : float + Rest mass in MeV. + charge_qe : float + Charge in units of the elementary charge, e.g. -1 for an electron. + s : float + Integrated path length along the reference orbit, in metres. + gyromagnetic_anomaly : float + Anomalous magnetic moment, dimensionless. Carried for round-tripping only; + `ParticleGroup` has no spin. + """ + + x: float + y: float + z: float + t: float + px: float + py: float + pz: float + pt: float + mass_MeV: float + charge_qe: float + s: float = 0.0 + gyromagnetic_anomaly: float = 0.0 + + @property + def mass_eV(self) -> float: + """Rest mass in eV.""" + return self.mass_MeV * 1.0e6 + + @property + def gamma(self) -> float: + """Relativistic gamma of the reference particle.""" + return -self.pt + + @property + def beta_gamma(self) -> float: + """Magnitude of the normalized reference momentum.""" + return float(np.sqrt(self.px**2 + self.py**2 + self.pz**2)) + + @property + def qm_eV(self) -> float: + """Charge over mass in 1/eV, the form ``add_n_particles`` expects.""" + return self.charge_qe / self.mass_eV + + @property + def qm_SI(self) -> float: + """Charge over mass in C/kg, the form ``qm`` is written with in openPMD output.""" + return self.qm_eV * c_light**2 + + +def _reference_is_physical(ref: ImpactXRefPart) -> bool: + """True when ``ref`` has a positive mass and is actually moving.""" + return bool(ref.mass_MeV > 0.0 and ref.gamma >= 1.0 and ref.beta_gamma > 0.0) + + +def _check_reference_particle(ref: ImpactXRefPart) -> None: + """Refuse a reference particle the conversion cannot use. + + A default-constructed ``RefPart`` has zero mass and zero energy, which makes + ``beta_gamma`` zero and every converted momentum infinite. ImpactX writes exactly + that into its ``particles_lost`` output, so this is not a hypothetical. + + Parameters + ---------- + ref : ImpactXRefPart + The reference particle to check. + + Raises + ------ + ValueError + If the mass is not positive or the particle is not moving. + """ + if _reference_is_physical(ref): + return + raise ValueError( + f"The reference particle is not physical: mass_MeV={ref.mass_MeV}, " + f"gamma={ref.gamma}, beta_gamma={ref.beta_gamma}. A mass and an energy are " + "both required to convert ImpactX' normalized coordinates." + ) + + +def pmd_species_of(ref: ImpactXRefPart, rtol: float = 1e-6) -> str: + """Infer the openPMD-beamphysics species name from a reference particle. + + Parameters + ---------- + ref : ImpactXRefPart + The reference particle. + rtol : float + Relative tolerance for the mass match. ImpactX and openPMD-beamphysics carry + electron masses that differ in the 9th digit, so this cannot be exact. + + Returns + ------- + str + A species name such as ``"electron"``. + + Raises + ------ + ValueError + If no known species matches; pass ``species=`` explicitly in that case. + """ + for pmd_name in IMPACTX_TO_PMD_SPECIES.values(): + charge_matches = np.isclose( + ref.charge_qe, charge_of(pmd_name) / e_charge, rtol=rtol + ) + mass_matches = np.isclose(ref.mass_eV, mass_of(pmd_name), rtol=rtol) + if charge_matches and mass_matches: + return pmd_name + raise ValueError( + f"Cannot infer a species from charge_qe={ref.charge_qe} and " + f"mass_MeV={ref.mass_MeV}. Pass species= explicitly." + ) + + +def _check_species_matches_reference( + species: str, ref: ImpactXRefPart, rtol: float = 1e-6 +) -> None: + """Refuse a species that is not the one the reference particle describes. + + The momenta are un-normalized with the *reference* mass, and ``ParticleGroup`` then + reads them back with the *species* mass. If the two disagree nothing raises on its + own: the bunch is silently relabelled and comes back at the wrong energy -- a 2 GeV + electron beam read as protons reports gamma = 2.35. + + Parameters + ---------- + species : str + openPMD-beamphysics species name. + ref : ImpactXRefPart + The reference particle the coordinates are relative to. + rtol : float + Relative tolerance. ImpactX and openPMD-beamphysics carry electron masses that + differ in the 9th digit, so this cannot be exact. + + Raises + ------ + ValueError + If the species' rest mass or charge does not match the reference particle's. + """ + # named for the reference particle, not just "matches": lume-impactx, which shares + # this module verbatim, has its own _check_species_matches(pg, ref) for injection + mass_matches = np.isclose(mass_of(species), ref.mass_eV, rtol=rtol) + charge_matches = np.isclose(charge_of(species) / e_charge, ref.charge_qe, rtol=rtol) + if mass_matches and charge_matches: + return + raise ValueError( + f"species={species!r} (mass {mass_of(species):.6e} eV, charge " + f"{charge_of(species) / e_charge:+.3f} e) is not the species the reference " + f"particle describes (mass {ref.mass_eV:.6e} eV, charge {ref.charge_qe:+.3f} " + "e). The momenta are normalized by the reference mass, so relabelling alone " + "would return the bunch at the wrong energy. Pass species= only to name a " + "species that pmd_species_of() cannot infer, not to convert between species." + ) + + +# -------------------------------------------------------------------------------------- +# Coordinate conversion +# -------------------------------------------------------------------------------------- + + +def particlegroup_to_impactx(pg: ParticleGroup, ref: ImpactXRefPart) -> dict: + """Convert a ``ParticleGroup`` to ImpactX fixed-s beam arrays. + + Parameters + ---------- + pg : ParticleGroup + The bunch to convert. Unless it already sits exactly on one plane it is copied + and drifted to its own mean ``z``, so the input is never mutated. + + ``pg.x`` and ``pg.y`` are taken to be *already relative to the reference + particle*, which is the frame the reader returns and the frame ImpactX works + in; ``ref.x`` and ``ref.y`` are not subtracted. A lab-frame bunch around a + reference orbit that is off-axis must be shifted first. ``pg.z`` does not enter + the result at all: in the local frame the plane *is* the reference particle's + location. + + ``pg.status`` and ``pg.id`` are not carried -- ImpactX has no equivalent of the + first, and ``add_n_particles`` assigns its own ids -- so filter dead particles + beforehand if that matters. + ref : ImpactXRefPart + The reference particle the ImpactX coordinates are relative to. + + Returns + ------- + dict + Keys ``position_x``, ``position_y``, ``position_t``, ``momentum_x``, + ``momentum_y``, ``momentum_t`` (arrays), ``weighting`` (array, real particles + per macroparticle), ``qm`` (scalar, 1/eV) and ``species`` (str) -- the names + ``ImpactXParticleContainer.to_df()`` uses, which is also what + :func:`impactx_to_particlegroup_data` reads back. + + ``ImpactXParticleContainer.add_n_particles`` takes the same quantities under + shorter names and has no species argument, so feed it as + ``add_n_particles(x=data["position_x"], ..., px=data["momentum_x"], ..., + qm=data["qm"], w=data["weighting"])``. + """ + # The bunch must occupy a single plane. A t-coordinate bunch (spread in z) is + # drifted to its own mean z on a copy, so the input is never mutated. + _check_reference_particle(ref) + + if not pg.in_z_coordinates: + pg = pg.copy() + pg.drift_to_z() + + mass_eV = ref.mass_eV + beta_gamma = ref.beta_gamma + + position_x = pg.x + position_y = pg.y + position_t = c_light * pg.t - ref.t + + momentum_x = pg.px / mass_eV / beta_gamma + momentum_y = pg.py / mass_eV / beta_gamma + + # gamma - gamma_ref through the algebraic identity + # (gamma^2 - gamma_ref^2) / (gamma + gamma_ref), which avoids subtracting two + # numbers that are both ~4000 for a 2 GeV beam and differ by ~1e-3. The float64 + # representation of pg.p already sets the accuracy floor, so the gain over the + # plain difference is a small constant factor rather than orders of magnitude -- + # but it is free. + p_mc2 = (pg.p / mass_eV) ** 2 + gamma = np.sqrt(1.0 + p_mc2) + dgamma = (p_mc2 - beta_gamma**2) / (gamma + ref.gamma) + momentum_t = -dgamma / beta_gamma + + return { + "position_x": position_x, + "position_y": position_y, + "position_t": position_t, + "momentum_x": momentum_x, + "momentum_y": momentum_y, + "momentum_t": momentum_t, + "weighting": pg.weight / abs(charge_of(pg.species)), + "qm": ref.qm_eV, + "species": pg.species, + } + + +def impactx_to_particlegroup_data( + data: dict, + ref: ImpactXRefPart, + species: str | None = None, +) -> dict: + """Convert ImpactX fixed-s beam arrays to ``ParticleGroup`` data. + + The inverse of :func:`particlegroup_to_impactx`. + + Parameters + ---------- + data : dict + Arrays keyed as ``ImpactXParticleContainer.to_df()`` names them: + ``position_x/y/t``, ``momentum_x/y/t``, ``weighting``. Optional ``id`` and + ``status`` keys are carried through. + ref : ImpactXRefPart + The reference particle the ImpactX coordinates are relative to. + species : str, optional + openPMD-beamphysics species name. Inferred from ``ref`` when omitted. + + Returns + ------- + dict + Suitable for ``ParticleGroup(data=...)``: ``x``, ``y``, ``z`` in metres, + ``px``, ``py``, ``pz`` in eV/c, ``t`` in seconds, ``weight`` in Coulomb, + ``status`` and ``species``. The result is in z-coordinates: every ``z`` is + zero, the reference plane, and the bunch length shows up as a spread in ``t``. + + Raises + ------ + ValueError + If a particle's transverse momentum exceeds its total momentum, which makes + ``pz`` imaginary. That means ``data`` and ``ref`` do not belong together. + """ + _check_reference_particle(ref) + + if species is None: + species = pmd_species_of(ref) + else: + _check_species_matches_reference(species, ref) + + mass_eV = ref.mass_eV + beta_gamma = ref.beta_gamma + n = len(np.asarray(data["position_x"])) + + # In the local frame the reference particle is the origin: its transverse momentum + # is zero by construction and its longitudinal momentum is |p_ref|. ref.px / ref.pz + # are *lab* components and must not be mixed in here. + px_mc = beta_gamma * np.asarray(data["momentum_x"], dtype=float) + py_mc = beta_gamma * np.asarray(data["momentum_y"], dtype=float) + # ref.gamma is -ref.pt, so this is gamma_ref + (gamma - gamma_ref) = gamma. + gamma = ref.gamma - beta_gamma * np.asarray(data["momentum_t"], dtype=float) + + pz_mc2 = gamma**2 - 1.0 - px_mc**2 - py_mc**2 + n_bad = int(np.count_nonzero(pz_mc2 < 0.0)) + if n_bad: + raise ValueError( + f"{n_bad} of {n} particles have a transverse momentum larger than their " + "total momentum, so pz would be imaginary. The reference particle most " + "likely does not belong to this bunch: check mass_MeV, pt and the " + "momentum normalization." + ) + pz_mc = np.sqrt(pz_mc2) + + if "status" in data: + status = np.asarray(data["status"]) + else: + status = np.full(n, int(ParticleStatus.ALIVE)) + + pg_data = { + "x": np.asarray(data["position_x"], dtype=float), + "y": np.asarray(data["position_y"], dtype=float), + "z": np.zeros(n), # Zero by definition in z-coordinates + "px": px_mc * mass_eV, + "py": py_mc * mass_eV, + "pz": pz_mc * mass_eV, + "t": (ref.t + np.asarray(data["position_t"], dtype=float)) / c_light, + "weight": np.asarray(data["weighting"], dtype=float) * abs(charge_of(species)), + "status": status, + "species": species, + } + if "id" in data: + pg_data["id"] = np.asarray(data["id"]) + return pg_data + + +# -------------------------------------------------------------------------------------- +# Per-particle data ParticleGroup cannot hold +# -------------------------------------------------------------------------------------- + +#: The spin components, which ImpactX always allocates and always writes. They stay at +#: exactly zero unless the beam was seeded with a spin distribution -- ``sim.spin = +#: True`` alone is not enough, the gate is the ``spin_distr`` argument to +#: ``add_particles``. So testing for "any non-zero" is exact: zero means there is +#: genuinely nothing to lose. +SPIN_COLUMNS = ("spin_x", "spin_y", "spin_z") + + +class UnrepresentableParticleData(NotImplementedError): + """Raised when a bunch carries per-particle data ``ParticleGroup`` cannot hold. + + Converting anyway would return a bunch that looks right and has silently lost + physics, so the conversion refuses instead. Pass ``strict=False`` to + :func:`read_beam_monitor` to drop the extra data deliberately. + """ + + +def _unrepresentable_in(columns: dict) -> list[str]: + """Name the per-particle data in ``columns`` that ``ParticleGroup`` cannot hold. + + Parameters + ---------- + columns : dict + Extra per-particle arrays keyed by ImpactX SoA name, i.e. everything beyond + the coordinates, weighting and id that this module maps. + + Returns + ------- + list of str + Human-readable descriptions, empty when the bunch converts losslessly. + """ + carried = [] + if any( + name in columns and np.any(np.asarray(columns[name]) != 0.0) + for name in SPIN_COLUMNS + ): + carried.append("spin (spin_x/y/z)") + runtime = sorted(name for name in columns if name not in SPIN_COLUMNS) + if runtime: + carried.append(f"runtime components {runtime}") + return carried + + +def _check_representable(columns: dict) -> None: + """Refuse to convert a bunch whose extra per-particle data would be dropped. + + Parameters + ---------- + columns : dict + Extra per-particle arrays keyed by ImpactX SoA name. + + Raises + ------ + UnrepresentableParticleData + If any spin component is non-zero, or any runtime component is present. + """ + carried = _unrepresentable_in(columns) + if not carried: + return + raise UnrepresentableParticleData( + f"This bunch carries {' and '.join(carried)}, which ParticleGroup cannot " + "represent. Converting would silently drop it. Pass strict=False to drop it " + "deliberately, or work with the ImpactX particle container directly. Note " + "that ImpactX's particles_lost output always carries a runtime 's_lost' " + "component." + ) + + +# -------------------------------------------------------------------------------------- +# openPMD BeamMonitor reader +# +# ImpactX's elements.BeamMonitor writes standard openPMD with the species always named +# "beam" -- also in the particles_lost output, where only the *file* is named +# differently. It records position/{x,y,t}, momentum/{x,y,t}, positionOffset/{x,y,t}, +# weighting, qm, spin/{x,y,z} and id, plus the reference particle and the reduced beam +# characteristics as per-iteration species attributes. Verified against ImpactX 26.08. +# +# This path needs no live ImpactX object, which is what makes ImpactXRefPart a plain +# dataclass rather than a wrapper around impactx.RefPart. +# -------------------------------------------------------------------------------------- + +#: openPMD records this reader consumes or deliberately ignores. Everything else in a +#: species is per-particle data ParticleGroup has no place for -- ImpactX's ``spin``, +#: or a runtime component such as the ``s/lost`` of the particles_lost output. +#: +#: ``positionOffset`` holds ``(x_ref, y_ref, t_ref)``, not zeros: the longitudinal part +#: is applied via the ``t_ref`` attribute, the transverse part is deliberately not, see +#: the module docstring. ``qm`` is redundant with the reference particle's mass and +#: charge, which are what this reader uses. +_CONSUMED_RECORDS = frozenset( + { + "position", + "positionOffset", + "momentum", + "weighting", + "qm", + "id", + } +) + +#: ImpactX writes this zero-extent placeholder instead of particle records when a +#: BeamMonitor is configured with ``particles=False`` (moments-only output). +_EMPTY_PLACEHOLDER_RECORD = "empty" + +#: The runtime component ImpactX attaches to every particle in its ``particles_lost`` +#: output: the path length ``s`` at which the particle was lost. +_S_LOST_COLUMN = "s_lost" + +#: ``status`` for a particle ImpactX has lost. +#: +#: `ParticleStatus` defines only ``CATHODE = 0`` and ``ALIVE = 1``; `ParticleGroup` +#: counts ``status == 1`` as alive and everything else as dead, and each interface +#: passes its own source code's value straight through -- `beamphysics.interfaces.bmad` +#: hands Bmad's ``state`` over as ``status`` verbatim, loss codes and all. There is no +#: universal "lost" value to reach for, so this is a choice. +#: +#: The one value it must not be is ``0``: that is a positive claim that the particle is +#: sitting at the source, and `beamphysics.interfaces.astra` writes ``status == 0`` back +#: out as Astra's ``-1``, "at the cathode". ``2`` is outside Bmad's loss-direction range +#: and carries no such meaning -- it reads as simply "not alive". +PARTICLE_STATUS_LOST = 2 + +# AMReX packs a particle's identity into one uint64 ``idcpu``: bit 63 marks the +# particle valid, bits 24-62 hold the id, bits 0-23 the originating MPI rank. The id +# counter is per-rank, so the id on its own repeats across the ranks of a parallel run +# -- only the packed value is globally unique, and that is what becomes the +# ``ParticleGroup`` id. +_AMREX_VALID_BIT = np.uint64(1) << np.uint64(63) + + +def particle_id_from_idcpu(idcpu) -> tuple[np.ndarray, np.ndarray]: + """Split AMReX's packed ``idcpu`` into a ``ParticleGroup`` id and a validity flag. + + ImpactX writes the raw ``idcpu`` as the openPMD ``id`` record. Only the validity + bit is stripped off here, for two reasons: with it set the value exceeds the range + of a signed 64-bit integer, which is what ``ParticleGroup`` stores ids in, and + aliveness belongs in ``status`` rather than in the id. Everything identifying the + particle -- AMReX' per-rank id *and* the rank it came from -- is kept, so the + result is unique across a parallel run where the id alone would not be. + + The original value is recovered as ``np.uint64(id) | (np.uint64(1) << 63)`` for a + live particle; AMReX' own id and rank are ``id >> 24`` and ``id & 0xFFFFFF``. + + Parameters + ---------- + idcpu : array_like of uint64 + The ``id`` record as stored. + + Returns + ------- + ids : np.ndarray of int64 + ``idcpu`` with the validity bit cleared. + valid : np.ndarray of bool + True where the particle is marked valid. + """ + idcpu = np.asarray(idcpu, dtype=np.uint64) + valid = (idcpu & _AMREX_VALID_BIT).astype(bool) + return (idcpu & ~_AMREX_VALID_BIT).astype(np.int64), valid + + +def refpart_from_openpmd(species: Any) -> ImpactXRefPart: + """Rebuild a reference particle from a BeamMonitor species' attributes. + + Parameters + ---------- + species : openpmd_api.ParticleSpecies + A species from a BeamMonitor iteration. + + Returns + ------- + ImpactXRefPart + ``mass_ref`` is stored in kg and ``charge_ref`` in Coulomb, so both are + converted here to the MeV / elementary-charge units the converters use. + + Raises + ------ + KeyError + If the species carries no reference particle attributes, i.e. it was not + written by an ImpactX BeamMonitor. + """ + attributes = set(species.attributes) + required = { + "x_ref", "y_ref", "z_ref", "t_ref", + "px_ref", "py_ref", "pz_ref", "pt_ref", + "mass_ref", "charge_ref", "s_ref", + } # fmt: skip + missing = sorted(required - attributes) + if missing: + raise KeyError( + f"Species is missing the reference particle attributes {missing}. " + "This does not look like ImpactX BeamMonitor output." + ) + + get = species.get_attribute + return ImpactXRefPart( + x=get("x_ref"), + y=get("y_ref"), + z=get("z_ref"), + t=get("t_ref"), + px=get("px_ref"), + py=get("py_ref"), + pz=get("pz_ref"), + pt=get("pt_ref"), + mass_MeV=get("mass_ref") * c_light**2 / e_charge / 1.0e6, + charge_qe=get("charge_ref") / e_charge, + s=get("s_ref"), + gyromagnetic_anomaly=( + get("gyromagnetic_anomaly_ref") + if "gyromagnetic_anomaly_ref" in attributes + else 0.0 + ), + ) + + +def _import_openpmd_api(): + """Import ``openpmd_api``, with an actionable message when it is missing.""" + try: + import openpmd_api + except ImportError as exc: # pragma: no cover - depends on the install + raise ImportError( + "Reading ImpactX BeamMonitor output needs the openpmd-api Python package: " + "conda install -c conda-forge openpmd-api, or pip install openpmd-api." + ) from exc + return openpmd_api + + +def read_beam_monitor_data( + path: str | pathlib.Path, + iteration: int | None = None, + species_name: str = "beam", + species: str | None = None, + strict: bool = True, + ref: ImpactXRefPart | None = None, +) -> dict: + """Read an ImpactX ``BeamMonitor`` openPMD file into ``ParticleGroup`` data. + + See :func:`read_beam_monitor` for the parameters; this is the same reader without + the final ``ParticleGroup`` construction. + + Returns + ------- + dict + Suitable for ``ParticleGroup(data=...)``. + """ + io = _import_openpmd_api() + + series = io.Series(str(path), io.Access.read_only) + try: + iterations = list(series.iterations) + if not iterations: + raise KeyError(f"No iterations in {str(path)!r}.") + if iteration is None: + iteration = iterations[-1] + elif iteration not in iterations: + raise KeyError( + f"Iteration {iteration} not in {str(path)!r}; have {iterations}." + ) + + particles = series.iterations[iteration].particles + if species_name not in particles: + raise KeyError( + f"Species {species_name!r} not in {str(path)!r}; have " + f"{list(particles)}. ImpactX always names it 'beam', including in " + "particles_lost output." + ) + beam = particles[species_name] + + records = list(beam) + if _EMPTY_PLACEHOLDER_RECORD in records: + raise KeyError( + f"{str(path)!r} iteration {iteration} holds no particles: this " + "BeamMonitor was configured with particles=False and wrote only the " + "reduced beam characteristics, which are on the species' attributes." + ) + + def _load(record_name: str, component: str): + record_component = beam[record_name][component] + return record_component, record_component.load_chunk() + + components = { + "position_x": _load("position", "x"), + "position_y": _load("position", "y"), + "position_t": _load("position", "t"), + "momentum_x": _load("momentum", "x"), + "momentum_y": _load("momentum", "y"), + "momentum_t": _load("momentum", "t"), + "weighting": _load("weighting", io.Record_Component.SCALAR), + } + idcpu = beam["id"][io.Record_Component.SCALAR].load_chunk() + + # ParticleGroup cannot hold spin or runtime components, so collect them -- + # either to refuse loudly, or to report what strict=False threw away. + extras = {} + for name in records: + if name in _CONSUMED_RECORDS: + continue + for component, record_component in beam[name].items(): + key = ( + name + if component == io.Record_Component.SCALAR + else f"{name}_{component}" + ) + extras[key] = record_component.load_chunk() + + file_ref = refpart_from_openpmd(beam) + series.flush() + + # openPMD stores values that must be multiplied by unitSI to reach the unit + # the record declares, and ImpactX writes 1.0 throughout. The positions are + # genuinely lengths, so the factor is applied. The momenta and the weighting + # are ImpactX' own dimensionless quantities -- normalized by the reference + # momentum, and a count of real particles -- so a factor other than 1 would + # mean the file no longer follows the convention decoded below. Refuse rather + # than scale them into silent nonsense. + data = {} + for key, (record_component, chunk) in components.items(): + unit_si = record_component.unit_SI + if key.startswith("position_"): + data[key] = np.asarray(chunk) * unit_si + elif unit_si != 1.0: + raise ValueError( + f"{key} in {str(path)!r} has unitSI={unit_si}, but this reader " + "decodes ImpactX' dimensionless convention, in which the momenta " + "are normalized by the reference momentum and the weighting " + "counts real particles." + ) + else: + data[key] = np.asarray(chunk) + extras = {key: np.asarray(chunk) for key, chunk in extras.items()} + ids, valid = particle_id_from_idcpu(idcpu) + finally: + series.close() + + if ref is None: + ref = file_ref + if not _reference_is_physical(ref): + raise ValueError( + f"{str(path)!r} iteration {iteration} carries a zeroed reference " + f"particle (mass_ref={ref.mass_MeV} MeV, gamma_ref={ref.gamma}), so " + "its normalized coordinates cannot be converted. ImpactX wrote this " + "into its particles_lost output before BLAST-ImpactX/impactx#1647; " + "newer files carry a usable one. Pass ref= with the ImpactXRefPart of " + "a BeamMonitor iteration, which refpart_from_openpmd() reads from the " + "monitor file." + ) + + if strict: + _check_representable(extras) + else: + dropped = _unrepresentable_in(extras) + if dropped: + warnings.warn( + f"Dropping {' and '.join(dropped)} from {str(path)!r} iteration " + f"{iteration}: ParticleGroup cannot represent it.", + stacklevel=2, + ) + + data["id"] = ids + # AMReX' validity bit is not aliveness: ImpactX copies lost particles into a + # separate container and marks them valid *there*, so the bit is True for every + # particle in a particles_lost file and taking it at face value would report a + # bunch that is entirely alive. + # + # The s_lost runtime component identifies such a file: CollectLost always adds it. + # The zeroed reference particle is a second signature, but only because ImpactX + # fails to set one -- keyed on that alone, this would silently go back to reporting + # lost particles as alive the day ImpactX fixes it. + from_lost_file = _S_LOST_COLUMN in extras or not _reference_is_physical(file_ref) + alive = valid & (not from_lost_file) + data["status"] = np.where(alive, int(ParticleStatus.ALIVE), PARTICLE_STATUS_LOST) + + return impactx_to_particlegroup_data(data, ref, species=species) + + +def read_beam_monitor( + path: str | pathlib.Path, + iteration: int | None = None, + species_name: str = "beam", + species: str | None = None, + strict: bool = True, + ref: ImpactXRefPart | None = None, +) -> ParticleGroup: + """Read an ImpactX ``BeamMonitor`` openPMD file into a ``ParticleGroup``. + + Parameters + ---------- + path : str or pathlib.Path + Path to the file ImpactX wrote, e.g. ``diags/openPMD/monitor.h5``. Any backend + openpmd-api supports works, including ``.bp`` and a ``%T``-templated file-based + series. + iteration : int, optional + Which iteration to read. The last one when omitted. + species_name : str + openPMD species to read. ImpactX always writes ``"beam"`` -- its lost-particle + output differs in the *file* name (``particles_lost.*``), not the species name. + species : str, optional + openPMD-beamphysics species name; inferred from the reference particle when + omitted. + strict : bool + When True (the default), refuse to read a bunch that carries per-particle data + `ParticleGroup` cannot hold -- non-zero spin, or a runtime component such as + the ``s_lost`` that ImpactX's ``particles_lost`` output always carries. Set it + to False to drop that data and read the bunch anyway. + ref : ImpactXRefPart, optional + Reference particle to interpret the coordinates against. Taken from the file + when omitted, which is what you want for a BeamMonitor. ImpactX wrote a + *zeroed* reference particle into its ``particles_lost`` output before + BLAST-ImpactX/impactx#1647, so reading such a file requires passing one, e.g. + ``refpart_from_openpmd(series.iterations[n].particles["beam"])`` from the + monitor file. + + Either way, a lost particle is only converted exactly if the reference energy + did not change between where it was lost -- the file's own ``s_lost`` record -- + and the reference particle in hand. The momenta are normalized by ``beta_gamma`` + at the reference particle's own ``s``, so through an RF cavity or any other + accelerating element the two differ, and so does the reference time. Which + reference particle ImpactX stores alongside lost particles is still settling + upstream; this reader applies no correction of its own -- it uses whatever it + is given -- so pass ``ref=`` explicitly when it matters which one that is. + + Returns + ------- + ParticleGroup + 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 the + module docstring for the frame conventions. + + ``status`` follows openPMD-beamphysics, where ``1`` is alive and anything else + is not -- ``pg.n_alive`` and ``pg.n_dead`` split on exactly that. Every particle + in a ``particles_lost`` file is marked :data:`PARTICLE_STATUS_LOST`, because + they are all lost by construction. + + Raises + ------ + KeyError + If the requested iteration or species is not in the file, if the species is + not ImpactX BeamMonitor output, or if the monitor recorded no particles. + UnrepresentableParticleData + If ``strict`` and the monitor recorded spin or runtime components. + ValueError + If neither ``ref`` nor the file provides a usable reference particle, or if a + record carries a ``unitSI`` this reader cannot honour. + ImportError + If the ``openpmd-api`` package is not installed. + + Examples + -------- + >>> from beamphysics.interfaces.impactx import read_beam_monitor + >>> P = read_beam_monitor("diags/openPMD/monitor.h5") # doctest: +SKIP + >>> P.norm_emit_x # doctest: +SKIP + """ + return ParticleGroup( + data=read_beam_monitor_data( + path, + iteration=iteration, + species_name=species_name, + species=species, + strict=strict, + ref=ref, + ) + ) + + +def beam_monitor_iterations(path: str | pathlib.Path) -> list[int]: + """List the iterations available in an ImpactX ``BeamMonitor`` file. + + Parameters + ---------- + path : str or pathlib.Path + Path to the file ImpactX wrote. + + Returns + ------- + list of int + The iteration numbers, in file order. + """ + io = _import_openpmd_api() + + series = io.Series(str(path), io.Access.read_only) + try: + return list(series.iterations) + finally: + series.close() diff --git a/beamphysics/particles.py b/beamphysics/particles.py index 64aba26..a6118b7 100644 --- a/beamphysics/particles.py +++ b/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, diff --git a/docs/examples/data/impactx/generate.py b/docs/examples/data/impactx/generate.py new file mode 100755 index 0000000..05664b1 --- /dev/null +++ b/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() diff --git a/docs/examples/data/impactx/monitor.h5 b/docs/examples/data/impactx/monitor.h5 new file mode 100644 index 0000000..4229a09 Binary files /dev/null and b/docs/examples/data/impactx/monitor.h5 differ diff --git a/docs/examples/data/impactx/particles_lost.h5 b/docs/examples/data/impactx/particles_lost.h5 new file mode 100644 index 0000000..ba0357a Binary files /dev/null and b/docs/examples/data/impactx/particles_lost.h5 differ diff --git a/environment.yml b/environment.yml index be31c89..548f168 100644 --- a/environment.yml +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index dd49cb2..5fb0b59 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/test_impactx.py b/tests/test_impactx.py new file mode 100644 index 0000000..9057381 --- /dev/null +++ b/tests/test_impactx.py @@ -0,0 +1,603 @@ +"""Tests for the ImpactX interface. + +The pure conversion tests need nothing but numpy. The reader tests run against +``docs/examples/data/impactx/*.h5``, written by ImpactX 26.08 with the +``generate.py`` script next to them, and need ``openpmd-api`` but not ImpactX. +""" + +from __future__ import annotations + +import pathlib + +import numpy as np +import pytest + +from beamphysics import ParticleGroup +from beamphysics.interfaces.impactx import ( + IMPACTX_TO_PMD_SPECIES, + PARTICLE_STATUS_LOST, + ImpactXRefPart, + UnrepresentableParticleData, + _check_representable, + particle_id_from_idcpu, + beam_monitor_iterations, + impactx_to_particlegroup_data, + particlegroup_to_impactx, + pmd_species_of, + read_beam_monitor, + refpart_from_openpmd, +) +from beamphysics.species import charge_of, e_charge, mass_of +from beamphysics.status import ParticleStatus +from beamphysics.units import c_light + +DATA_DIR = ( + pathlib.Path(__file__).resolve().parent.parent + / "docs" + / "examples" + / "data" + / "impactx" +) +MONITOR = DATA_DIR / "monitor.h5" +PARTICLES_LOST = DATA_DIR / "particles_lost.h5" + +try: + import openpmd_api # noqa: F401 + + HAVE_OPENPMD_API = True +except ImportError: + HAVE_OPENPMD_API = False + +requires_data = pytest.mark.skipif( + not (HAVE_OPENPMD_API and MONITOR.exists()), + reason="needs openpmd-api and the ImpactX test data", +) + +# The reference energy of both the fixtures and the committed test data. +KIN_ENERGY_MeV = 2.0e3 + + +def make_ref(species: str = "electron", kin_energy_MeV: float = KIN_ENERGY_MeV): + """An on-axis reference particle at s = 1.5 m, using beamphysics' own constants.""" + mass_MeV = mass_of(species) / 1.0e6 + gamma = 1.0 + kin_energy_MeV / mass_MeV + return ImpactXRefPart( + x=0.0, + y=0.0, + z=1.5, + t=1.5, + px=0.0, + py=0.0, + pz=np.sqrt(gamma**2 - 1.0), + pt=-gamma, + mass_MeV=mass_MeV, + charge_qe=charge_of(species) / e_charge, + s=1.5, + ) + + +@pytest.fixture +def electron_ref() -> ImpactXRefPart: + return make_ref() + + +@pytest.fixture +def bunch(electron_ref: ImpactXRefPart) -> ParticleGroup: + """A 2 GeV electron bunch in z-coordinates, matched to `electron_ref`.""" + rng = np.random.default_rng(42) + n = 500 + mass_eV = electron_ref.mass_eV + p0 = electron_ref.beta_gamma * mass_eV # eV/c + + px = p0 * rng.normal(0.0, 1.0e-5, n) + py = p0 * rng.normal(0.0, 1.0e-5, n) + pz = p0 * (1.0 + rng.normal(0.0, 2.0e-3, n)) + + return ParticleGroup( + data={ + "x": rng.normal(0.0, 1.0e-4, n), + "y": rng.normal(0.0, 1.0e-4, n), + "z": np.full(n, electron_ref.z), + "px": px, + "py": py, + "pz": pz, + "t": electron_ref.t / c_light + rng.normal(0.0, 3.0e-12, n), + "weight": np.full(n, 1.0e-9 / n), + "status": np.full(n, int(ParticleStatus.ALIVE)), + "species": "electron", + } + ) + + +# --------------------------------------------------------------------------- +# Reference particle +# --------------------------------------------------------------------------- + + +def test_qm_units_relate_by_c_squared(electron_ref): + assert electron_ref.qm_SI == pytest.approx(electron_ref.qm_eV * c_light**2) + # -e / m_e in C/kg + assert electron_ref.qm_SI == pytest.approx(-1.75882e11, rel=1e-5) + + +def test_gamma_and_beta_gamma_agree(electron_ref): + assert electron_ref.gamma == pytest.approx(-electron_ref.pt) + assert electron_ref.beta_gamma == pytest.approx( + np.sqrt(electron_ref.gamma**2 - 1.0) + ) + + +@pytest.mark.parametrize("pmd_name", sorted(IMPACTX_TO_PMD_SPECIES.values())) +def test_species_inference(pmd_name): + assert pmd_species_of(make_ref(pmd_name)) == pmd_name + + +def test_species_inference_refuses_the_unknown(electron_ref): + from dataclasses import replace + + with pytest.raises(ValueError, match="Pass species= explicitly"): + pmd_species_of(replace(electron_ref, mass_MeV=123.456)) + + +# --------------------------------------------------------------------------- +# Coordinate conversion +# --------------------------------------------------------------------------- + + +def test_roundtrip_is_exact(bunch, electron_ref): + data = particlegroup_to_impactx(bunch, electron_ref) + back = ParticleGroup(data=impactx_to_particlegroup_data(data, electron_ref)) + + for key in ("x", "y", "px", "py", "pz", "t", "weight"): + np.testing.assert_allclose( + back[key], bunch[key], rtol=1e-11, atol=0, err_msg=key + ) + assert back.species == bunch.species + + +def test_roundtrip_preserves_beam_statistics(bunch, electron_ref): + data = particlegroup_to_impactx(bunch, electron_ref) + back = ParticleGroup(data=impactx_to_particlegroup_data(data, electron_ref)) + + assert back.charge == pytest.approx(bunch.charge, rel=1e-12) + assert back.norm_emit_x == pytest.approx(bunch.norm_emit_x, rel=1e-9) + assert back.norm_emit_y == pytest.approx(bunch.norm_emit_y, rel=1e-9) + assert back.std("t") == pytest.approx(bunch.std("t"), rel=1e-9) + assert back.avg("energy") == pytest.approx(bunch.avg("energy"), rel=1e-12) + + +def test_result_is_in_z_coordinates(bunch, electron_ref): + """The bunch comes back on the reference plane, with the length in t.""" + data = particlegroup_to_impactx(bunch, electron_ref) + back = ParticleGroup(data=impactx_to_particlegroup_data(data, electron_ref)) + + assert back.in_z_coordinates + np.testing.assert_array_equal(np.unique(back.z), [0.0]) + assert back.std("t") > 0.0 + + +def test_weighting_counts_real_particles(bunch, electron_ref): + data = particlegroup_to_impactx(bunch, electron_ref) + assert data["weighting"].sum() * e_charge == pytest.approx(bunch.charge, rel=1e-12) + + +def test_momentum_t_is_accurate(bunch, electron_ref): + """momentum_t must resolve gamma - gamma_ref, which is 1e-3 of gamma at 2 GeV. + + Errors are measured against the spread of momentum_t itself: a particle sitting + exactly on the reference momentum has momentum_t ~ 0, where a per-particle + relative error is meaningless. + """ + data = particlegroup_to_impactx(bunch, electron_ref) + + mass_eV = np.longdouble(electron_ref.mass_eV) + gamma_ref = np.longdouble(electron_ref.gamma) + beta_gamma = np.longdouble(electron_ref.beta_gamma) + gamma = np.sqrt(1.0 + (np.longdouble(bunch.p) / mass_eV) ** 2) + exact = -(gamma - gamma_ref) / beta_gamma + + naive_gamma = np.sqrt(1.0 + (bunch.p / electron_ref.mass_eV) ** 2) + naive = -(naive_gamma - electron_ref.gamma) / electron_ref.beta_gamma + + scale = float(np.std(exact)) + ours_error = float(np.max(np.abs(data["momentum_t"] - exact))) / scale + naive_error = float(np.max(np.abs(naive - exact))) / scale + assert ours_error < 1e-13 + # the algebraic form must never be worse than the plain difference + assert ours_error <= naive_error + + +def test_input_particlegroup_is_not_mutated(electron_ref): + """A t-coordinate bunch is drifted on a copy, never in place.""" + n = 32 + rng = np.random.default_rng(0) + p0 = electron_ref.beta_gamma * electron_ref.mass_eV + pg = ParticleGroup( + data={ + "x": np.zeros(n), + "y": np.zeros(n), + "z": rng.normal(0.0, 1.0e-3, n), + "px": np.zeros(n), + "py": np.zeros(n), + "pz": np.full(n, p0), + "t": np.zeros(n), + "weight": np.full(n, 1.0e-12), + "status": np.ones(n, dtype=int), + "species": "electron", + } + ) + assert pg.in_t_coordinates + z_before, t_before = pg.z.copy(), pg.t.copy() + + particlegroup_to_impactx(pg, electron_ref) + + np.testing.assert_array_equal(pg.z, z_before) + np.testing.assert_array_equal(pg.t, t_before) + + +def test_mismatched_reference_is_refused(bunch, electron_ref): + """A reference particle from another beam makes pz imaginary; say so.""" + data = particlegroup_to_impactx(bunch, electron_ref) + data["momentum_x"] = np.full(bunch.n_particle, 10.0) + + with pytest.raises(ValueError, match="transverse momentum"): + impactx_to_particlegroup_data(data, electron_ref) + + +def test_status_and_id_are_carried_through(bunch, electron_ref): + data = particlegroup_to_impactx(bunch, electron_ref) + data["id"] = np.arange(bunch.n_particle) + data["status"] = np.zeros(bunch.n_particle, dtype=int) + + back = ParticleGroup(data=impactx_to_particlegroup_data(data, electron_ref)) + np.testing.assert_array_equal(back.id, np.arange(bunch.n_particle)) + assert back.n_alive == 0 + + +# --------------------------------------------------------------------------- +# AMReX id packing +# --------------------------------------------------------------------------- + + +VALID_BIT = np.uint64(1) << np.uint64(63) + + +def pack_idcpu(ids, cpus, valid=True): + """Pack ids and ranks the way AMReX does, for the tests below.""" + ids = np.asarray(ids, dtype=np.uint64) + cpus = np.asarray(cpus, dtype=np.uint64) + packed = (ids << np.uint64(24)) | cpus + return (packed | VALID_BIT) if valid else packed + + +def test_particle_id_keeps_the_whole_idcpu(): + """Only the validity bit comes off; id and rank both stay in the value.""" + packed = pack_idcpu([1, 2, 12345], [0, 3, 7]) + + got, valid = particle_id_from_idcpu(packed) + np.testing.assert_array_equal(valid, [True, True, True]) + np.testing.assert_array_equal(got >> 24, [1, 2, 12345]) # AMReX' own id + np.testing.assert_array_equal(got & 0xFFFFFF, [0, 3, 7]) # originating rank + # int64 is what ParticleGroup stores ids in, and the value must fit + assert got.dtype == np.int64 + assert got.min() >= 0 + # and the original is recoverable + np.testing.assert_array_equal(got.astype(np.uint64) | VALID_BIT, packed) + + +def test_particle_id_survives_a_particlegroup(bunch): + """The raw idcpu overflows the int64 ParticleGroup stores ids in; the id must not. + + Stripping the validity bit is what makes it fit, so this checks the whole path, + not just the dtype: ParticleGroup coerces ids with `_round_to_int_array`, which + wraps silently rather than raising. + """ + n = bunch.n_particle + packed = pack_idcpu(np.arange(1, n + 1), np.zeros(n, dtype=np.uint64)) + assert packed.max() > np.iinfo(np.int64).max + + ids, _ = particle_id_from_idcpu(packed) + data = particlegroup_to_impactx(bunch, make_ref()) + data["id"] = ids + stored = ParticleGroup(data=impactx_to_particlegroup_data(data, make_ref())).id + + np.testing.assert_array_equal(stored, ids) + assert stored.min() > 0 + + +def test_particle_id_is_unique_across_ranks(): + """AMReX counts ids per rank, so the id alone collides on a parallel run.""" + packed = pack_idcpu([1, 1, 2, 2], [0, 1, 0, 1]) + + got, _ = particle_id_from_idcpu(packed) + assert len(np.unique(got)) == 4 + assert len(np.unique(got >> 24)) == 2 # what the AMReX id alone would give + + +def test_particle_id_reports_invalid_particles(): + packed = pack_idcpu([1, 2], [0, 0], valid=False) + + got, valid = particle_id_from_idcpu(packed) + np.testing.assert_array_equal(valid, [False, False]) + np.testing.assert_array_equal(got >> 24, [1, 2]) + + +# --------------------------------------------------------------------------- +# Unrepresentable per-particle data +# --------------------------------------------------------------------------- + + +def test_zero_spin_is_representable(): + _check_representable({name: np.zeros(4) for name in ("spin_x", "spin_y", "spin_z")}) + + +def test_nonzero_spin_refuses_loudly(): + columns = {name: np.zeros(4) for name in ("spin_x", "spin_y", "spin_z")} + columns["spin_z"][2] = 1.0 + with pytest.raises(UnrepresentableParticleData, match="spin"): + _check_representable(columns) + + +def test_runtime_component_refuses_loudly(): + with pytest.raises(UnrepresentableParticleData, match="s_lost"): + _check_representable({"s_lost": np.full(4, 0.25)}) + + +# --------------------------------------------------------------------------- +# Reference particle from openPMD attributes +# --------------------------------------------------------------------------- + + +class _FakeSpecies: + def __init__(self, attributes: dict): + self._attributes = attributes + + @property + def attributes(self): + return list(self._attributes) + + def get_attribute(self, name): + return self._attributes[name] + + +def test_refpart_from_openpmd_rejects_a_foreign_species(): + with pytest.raises(KeyError, match="ImpactX BeamMonitor"): + refpart_from_openpmd(_FakeSpecies({"x_ref": 0.0})) + + +# --------------------------------------------------------------------------- +# BeamMonitor reader +# --------------------------------------------------------------------------- + + +@requires_data +def test_beam_monitor_iterations(): + assert beam_monitor_iterations(MONITOR) == [1, 5] + + +@requires_data +def test_read_beam_monitor_matches_impactx_moments(): + """ImpactX's own reduced beam characteristics must come back out.""" + import openpmd_api as io + + P = read_beam_monitor(MONITOR) + + series = io.Series(str(MONITOR), io.Access.read_only) + beam = series.iterations[5].particles["beam"] + attrs = {name: beam.get_attribute(name) for name in beam.attributes} + series.close() + + assert P.species == "electron" + assert P.n_particle == 184 + assert P.n_alive == P.n_particle + assert P.charge == pytest.approx(abs(attrs["charge_C"]), rel=1e-12) + + assert P.std("x") == pytest.approx(attrs["sigma_x"], rel=1e-12) + assert P.std("y") == pytest.approx(attrs["sigma_y"], rel=1e-12) + assert P.std("t") * c_light == pytest.approx(attrs["sigma_t"], rel=1e-12) + assert P.avg("x") == pytest.approx(attrs["mean_x"], rel=1e-9) + assert P.avg("y") == pytest.approx(attrs["mean_y"], rel=1e-9) + assert P.avg("t") * c_light - attrs["t_ref"] == pytest.approx( + attrs["mean_t"], rel=1e-9 + ) + # beamphysics' covariance is bias-corrected and ImpactX' is not, and with equal + # weights that is exactly a factor n/(n-1) on the emittance + n = P.n_particle + bias = n / (n - 1) + assert P.norm_emit_x == pytest.approx(attrs["emittance_xn"] * bias, rel=1e-9) + assert P.norm_emit_y == pytest.approx(attrs["emittance_yn"] * bias, rel=1e-9) + + +@requires_data +def test_conversion_to_impactx_matches_the_file_it_came_from(): + """Check the write direction against real ImpactX output, not synthetic data. + + Every other conversion test starts from a bunch this module itself built, so a + paired sign or normalization error present in *both* directions would round-trip + cleanly and pass. This one converts a bunch read from the file back to ImpactX + arrays and compares against what ImpactX actually wrote. + """ + import openpmd_api as io + + series = io.Series(str(MONITOR), io.Access.read_only) + beam = series.iterations[5].particles["beam"] + ref = refpart_from_openpmd(beam) + raw = { + "position_x": beam["position"]["x"].load_chunk(), + "position_y": beam["position"]["y"].load_chunk(), + "position_t": beam["position"]["t"].load_chunk(), + "momentum_x": beam["momentum"]["x"].load_chunk(), + "momentum_y": beam["momentum"]["y"].load_chunk(), + "momentum_t": beam["momentum"]["t"].load_chunk(), + "weighting": beam["weighting"][io.Record_Component.SCALAR].load_chunk(), + } + series.flush() + series.close() + + data = particlegroup_to_impactx(read_beam_monitor(MONITOR), ref) + + for key, expected in raw.items(): + expected = np.asarray(expected) + # measure against the spread, or the level for a uniform column like weighting + scale = np.std(expected) or np.abs(np.mean(expected)) or 1.0 + error = np.max(np.abs(data[key] - expected)) / scale + assert error < 1e-12, f"{key}: {error:e}" + + +@requires_data +def test_read_beam_monitor_reference_particle(): + import openpmd_api as io + + series = io.Series(str(MONITOR), io.Access.read_only) + ref = refpart_from_openpmd(series.iterations[5].particles["beam"]) + series.close() + + assert pmd_species_of(ref) == "electron" + assert ref.charge_qe == pytest.approx(-1.0) + assert ref.mass_MeV == pytest.approx(mass_of("electron") / 1e6, rel=1e-8) + assert ref.gamma == pytest.approx(1.0 + KIN_ENERGY_MeV / ref.mass_MeV, rel=1e-8) + assert ref.s == pytest.approx(1.25) + assert ref.t == pytest.approx(ref.s, rel=1e-6) # ultrarelativistic + + +@requires_data +def test_read_beam_monitor_carries_ids(): + P = read_beam_monitor(MONITOR) + assert P.id.dtype.kind == "i" + assert len(np.unique(P.id)) == P.n_particle + # the test data is a serial run, so every rank field is 0 and the AMReX ids are + # 1..200 for the 200 macroparticles the run started with + np.testing.assert_array_equal(P.id & 0xFFFFFF, 0) + assert (P.id >> 24).min() >= 1 + assert (P.id >> 24).max() <= 200 + + +@requires_data +def test_read_beam_monitor_iteration_selection(): + first = read_beam_monitor(MONITOR, iteration=1) + last = read_beam_monitor(MONITOR, iteration=5) + + assert first.n_particle == 200 # before the collimator + assert last.n_particle == 184 + assert first.std("x") != last.std("x") + # the default is the last iteration + assert read_beam_monitor(MONITOR).n_particle == last.n_particle + + +@requires_data +def test_read_beam_monitor_rejects_bad_arguments(): + with pytest.raises(KeyError, match="Iteration 99"): + read_beam_monitor(MONITOR, iteration=99) + with pytest.raises(KeyError, match="always names it 'beam'"): + read_beam_monitor(MONITOR, species_name="particles_lost") + + +@requires_data +def test_particlegroup_from_impactx_classmethod(): + from beamphysics.testing import assert_pg_close + + assert_pg_close(ParticleGroup.from_impactx(MONITOR), read_beam_monitor(MONITOR)) + + +@requires_data +@pytest.mark.parametrize("species", ["proton", "positron"]) +def test_species_cannot_relabel_across_species(species): + """The momenta are normalized by the *reference* mass, so a relabel is not free. + + Read as protons, this 2 GeV electron beam would come back at gamma = 2.35 and no + exception. `positron` is the sharp case: same mass as the reference, opposite + charge, so only checking the mass would let it through. + """ + with pytest.raises(ValueError, match="not the species the reference particle"): + read_beam_monitor(MONITOR, species=species) + + +def test_species_can_name_what_inference_cannot(bunch): + """`species=` is for species ImpactX has no name for, not for converting.""" + from dataclasses import replace + + muon_ref = replace( + make_ref(), + mass_MeV=mass_of("muon") / 1.0e6, + charge_qe=charge_of("muon") / e_charge, + ) + with pytest.raises(ValueError, match="Pass species= explicitly"): + pmd_species_of(muon_ref) + + data = particlegroup_to_impactx(bunch, make_ref()) + back = ParticleGroup( + data=impactx_to_particlegroup_data(data, muon_ref, species="muon") + ) + assert back.species == "muon" + + +requires_lost_data = pytest.mark.skipif( + not (HAVE_OPENPMD_API and PARTICLES_LOST.exists()), + reason="needs openpmd-api and the ImpactX particles_lost test data", +) + + +@requires_lost_data +def test_particles_lost_reference_particle_either_way(): + """ImpactX wrote a default-constructed RefPart into particles_lost output before + BLAST-ImpactX/impactx#1647; newer versions store a usable one. + + The committed fixture predates that fix, so it takes the first branch. Assert + whichever the file actually holds rather than the version it was made with, so + regenerating the data against a newer ImpactX does not look like a regression. + """ + import openpmd_api as io + + series = io.Series(str(PARTICLES_LOST), io.Access.read_only) + ref = refpart_from_openpmd(series.iterations[0].particles["beam"]) + series.close() + + if ref.mass_MeV == 0.0: + assert ref.gamma == 0.0 + with pytest.raises(ValueError, match="zeroed reference particle"): + read_beam_monitor(PARTICLES_LOST) + else: + # a usable reference particle in the file means no ref= is needed + with pytest.warns(UserWarning, match="s_lost"): + lost = read_beam_monitor(PARTICLES_LOST, strict=False) + assert lost.n_alive == 0 + assert pmd_species_of(ref) == "electron" + + +@requires_lost_data +def test_particles_lost_carries_an_unrepresentable_component(): + """ImpactX' lost-particle output always has a runtime `s_lost` column.""" + import openpmd_api as io + + series = io.Series(str(MONITOR), io.Access.read_only) + ref = refpart_from_openpmd(series.iterations[5].particles["beam"]) + series.close() + + with pytest.raises(UnrepresentableParticleData, match=r"s_lost"): + read_beam_monitor(PARTICLES_LOST, ref=ref) + + with pytest.warns(UserWarning, match=r"s_lost"): + lost = read_beam_monitor(PARTICLES_LOST, ref=ref, strict=False) + kept = read_beam_monitor(MONITOR, iteration=5) + + assert lost.n_particle == 16 + assert lost.n_particle + kept.n_particle == 200 + # AMReX marks these valid -- they are live entries of the loss container -- so the + # validity bit alone would report them alive, which is exactly backwards + assert lost.n_alive == 0 + assert lost.n_dead == 16 + np.testing.assert_array_equal(lost.status, PARTICLE_STATUS_LOST) + assert kept.n_alive == kept.n_particle + # the same species is used for lost particles; ids partition the original beam + assert set(lost.id).isdisjoint(set(kept.id)) + assert len(set(lost.id) | set(kept.id)) == 200 + + +def test_unphysical_reference_particle_is_refused(bunch): + zero_ref = ImpactXRefPart( + x=0.0, y=0.0, z=0.0, t=0.0, + px=0.0, py=0.0, pz=0.0, pt=0.0, + mass_MeV=0.0, charge_qe=0.0, + ) # fmt: skip + with pytest.raises(ValueError, match="not physical"): + particlegroup_to_impactx(bunch, zero_ref)