diff --git a/beamphysics/__init__.py b/beamphysics/__init__.py index c569eeec..18825456 100644 --- a/beamphysics/__init__.py +++ b/beamphysics/__init__.py @@ -7,7 +7,7 @@ from .particles import ParticleGroup, single_particle from .readers import particle_paths from .status import ParticleStatus - from .wavefront import Wavefront, WavefrontK + from .wavefront import Wavefront, WavefrontAttrs, WavefrontK from .writers import pmd_init try: @@ -23,6 +23,7 @@ "particle_paths": ".readers", "ParticleStatus": ".status", "Wavefront": ".wavefront", + "WavefrontAttrs": ".wavefront", "WavefrontK": ".wavefront", "pmd_init": ".writers", } @@ -36,6 +37,7 @@ "pmd_init", "single_particle", "Wavefront", + "WavefrontAttrs", "WavefrontK", ] diff --git a/beamphysics/interfaces/genesis.py b/beamphysics/interfaces/genesis.py index 02e9d7e4..e30914d1 100644 --- a/beamphysics/interfaces/genesis.py +++ b/beamphysics/interfaces/genesis.py @@ -1033,7 +1033,6 @@ def wavefront_write_genesis4( w, h5: File, polarization: str = None, - refposition: float = 0, ) -> None: """ Write the wavefront field data to a Genesis4-style HDF5 file. @@ -1068,15 +1067,13 @@ def wavefront_write_genesis4( - If only `Ey` exists, it will be written. - If both components exist, a `ValueError` is raised. - refposition : float, optional - The reference position in meters, stored as metadata in the output file. Default is `0`. - Raises ------ ValueError - If both `Ex` and `Ey` exist but no polarization is explicitly specified. - If `nx != ny`, as Genesis4 requires a square grid. - If `dx != dy`, as Genesis4 requires equal grid spacing in both transverse directions. + - If any grid offset is nonzero, as Genesis4 has no way to store a grid origin. - If `polarization` is specified but not `"x"` or `"y"`. Notes @@ -1084,6 +1081,9 @@ def wavefront_write_genesis4( - The function ensures that the grid size and spacing meet Genesis4's requirements. - The data is stored in slices, following the indexing convention of Genesis4: The x-coordinates are stored as the inner loop, requiring a transpose before flattening. + - `refposition` is written from `w.s_position`, the position of this dump along + the undulator line. There is no write-time override, so the file cannot + disagree with the wavefront it came from. """ nx, ny, nz = w.shape @@ -1113,9 +1113,17 @@ def wavefront_write_genesis4( if dx != dy: raise ValueError(f"Genesis4 requires dx = dy. This data has {dx=}, {dy=}") + if (w.xmid, w.ymid, w.zmid) != (0.0, 0.0, 0.0): + raise ValueError( + "Genesis4 stores only a grid point count and spacing, so it has nowhere " + "to record a grid origin. Writing this wavefront would silently move it " + f"to a centered grid. This data has xmid={w.xmid}, ymid={w.ymid}, " + f"zmid={w.zmid}." + ) + h5["gridpoints"] = np.asarray([nx]) h5["gridsize"] = np.asarray([dx]) - h5["refposition"] = np.asarray([refposition]) + h5["refposition"] = np.asarray([w.s_position]) h5["wavelength"] = np.asarray([wavelength]) h5["slicecount"] = np.asarray([nz]) h5["slicespacing"] = np.asarray([dz]) diff --git a/beamphysics/wavefront/__init__.py b/beamphysics/wavefront/__init__.py index 97cea780..3e852b84 100644 --- a/beamphysics/wavefront/__init__.py +++ b/beamphysics/wavefront/__init__.py @@ -1,3 +1,4 @@ +from .openpmd import WavefrontAttrs from .wavefront import Wavefront, WavefrontK -__all__ = ["Wavefront", "WavefrontK"] +__all__ = ["Wavefront", "WavefrontAttrs", "WavefrontK"] diff --git a/beamphysics/wavefront/openpmd.py b/beamphysics/wavefront/openpmd.py new file mode 100644 index 00000000..a1ec6597 --- /dev/null +++ b/beamphysics/wavefront/openpmd.py @@ -0,0 +1,761 @@ +""" +openPMD EXT_Wavefront reading and writing for the `Wavefront` class. + +The layout below follows the ``Wavefront`` extension of the openPMD standard +(branch ``upcoming-2.0.0``), which is authoritative:: + + / openPMD "2.0.0", openPMDextension "Wavefront", + basePath "/data/%T/", meshesPath "meshes/", + iterationEncoding "groupBased" + /data// one iteration per file. Slices of a single pulse are + simultaneous, so the slice axis is a *mesh* axis and + never the openPMD iteration. + .../meshes/electricField the mesh record. All required attributes live on the + record: the base standard's geometry, axisLabels, + gridSpacing, gridGlobalOffset, gridUnitSI, + gridUnitDimension, unitDimension and timeOffset, plus + the extension's photonEnergy [J], temporalDomain, + spatialDomain and zCoordinate. + .../electricField/x, y complex compound ``{r, i}`` datasets in V/m, which + h5py maps to complex dtypes natively. The ``z`` + component is never written: a paraxial field has none. + +Datasets are stored in ``(z, y, x)`` order -- declared by ``axisLabels`` -- so that each +transverse slice is one contiguous block in the file. The class's ``(nx, ny, nz)`` +convention is recovered by a transpose that is applied one slice at a time, keeping peak +memory at one field plus one transverse slice rather than two full copies. + +Only the real-space, time-domain case is implemented; ``spatialDomain='k'`` and +``temporalDomain='frequency'`` are refused by name. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields + +import numpy as np +from scipy.constants import h as h_planck + +from ..readers import constant_component_value, is_constant_component +from ..tools import decode_attr, encode_attr +from ..units import c_light, dimension + +__all__ = [ + "WavefrontAttrs", + "load_wavefront_openpmd", + "write_wavefront_openpmd", +] + + +# openPMD defaults, used when a file omits these root attributes. +_DEFAULT_BASE_PATH = "/data/%T/" +_DEFAULT_MESHES_PATH = "meshes/" + +_RECORD_NAME = "electricField" + +# Stored (file) axis order. See the module docstring for why. +_STORED_AXIS_LABELS = ("z", "y", "x") + +# Axis order of the `Wavefront` class arrays. +_CLASS_AXIS_LABELS = ("x", "y", "z") + +# 7-tuples of base-SI exponents, from the package's own table. +_UNIT_DIMENSION_E_FIELD = dimension("electric_field") +_UNIT_DIMENSION_LENGTH = dimension("length") + +# Attributes on the mesh record that this module writes from the wavefront itself. +# They may not be supplied by the caller, and are not carried in `WavefrontAttrs.other` +# on read, because the value in the file is redundant with the class's own state. +_COMPUTED_RECORD_ATTRS = ( + # Extension. + "photonEnergy", # derived from `wavelength` + "temporalDomain", # always 'time' for this class + "spatialDomain", # always 'r' for this class + "zCoordinate", # `Wavefront.s_position`, which `drift` advances + # Base standard, written from the grid and the field dtype. + "geometry", + "geometryParameters", + "axisLabels", + "gridSpacing", + "gridGlobalOffset", + "gridUnitSI", + "gridUnitDimension", + "unitDimension", + "timeOffset", + "dataOrder", # openPMD 1.x only; 2.0 uses axisLabels alone +) + + +def _pmd(key): + """ + Dataclass field metadata tagging a field with its openPMD attribute name. + + Parameters + ---------- + key : str + The attribute name as it appears in the file. + + Returns + ------- + dict + """ + return {"pmd_key": key} + + +@dataclass +class WavefrontAttrs: + """ + openPMD EXT_Wavefront attributes carried alongside a `Wavefront`. + + Field names are Python-style; the corresponding openPMD attribute names, used in + the file, are camelCase and recorded in each field's metadata. Attributes the + class derives from its own state -- `photonEnergy`, `temporalDomain`, + `spatialDomain`, `zCoordinate` and the grid attributes -- are deliberately + absent. `zCoordinate` in particular is `Wavefront.s_position`, which is a + coordinate the propagators advance rather than provenance carried through I/O. + + Parameters + ---------- + beamline : str, optional + Name of the beamline this wavefront belongs to. + radius_of_curvature_x : float, optional + Radius of curvature in x, in m. + radius_of_curvature_y : float, optional + Radius of curvature in y, in m. + delta_radius_of_curvature_x : float, optional + Uncertainty in `radius_of_curvature_x`, in m. + delta_radius_of_curvature_y : float, optional + Uncertainty in `radius_of_curvature_y`, in m. + other : dict, optional + Record attributes that are not part of the extension as this module knows + it, keyed by their openPMD name. Populated on read and written back + verbatim, so that a file using a newer revision of the extension survives a + round trip. Names this module computes are refused. Named after the `other` + dict that `readers.load_field_attrs` uses for the same purpose. + + Raises + ------ + ValueError + If `other` holds a name that has a field of its own, or one that the writer + computes. + + Examples + -------- + >>> attrs = WavefrontAttrs(beamline="SXR", radius_of_curvature_x=12.5) + >>> attrs.to_pmd()["radiusOfCurvatureX"] + 12.5 + + A misspelled attribute is a `TypeError` at construction rather than a silent + omission at write time: + + >>> WavefrontAttrs(radius_of_curvature_z=1.0) + Traceback (most recent call last): + ... + TypeError: ... + """ + + beamline: str | None = field(default=None, metadata=_pmd("beamline")) + radius_of_curvature_x: float | None = field( + default=None, metadata=_pmd("radiusOfCurvatureX") + ) + radius_of_curvature_y: float | None = field( + default=None, metadata=_pmd("radiusOfCurvatureY") + ) + delta_radius_of_curvature_x: float | None = field( + default=None, metadata=_pmd("deltaRadiusOfCurvatureX") + ) + delta_radius_of_curvature_y: float | None = field( + default=None, metadata=_pmd("deltaRadiusOfCurvatureY") + ) + other: dict = field(default_factory=dict) + + def __post_init__(self): + known = self.pmd_keys() + for name in self.other: + if name in known.values(): + raise ValueError( + f"{name!r} has a field of its own; set it directly rather than " + "through `other`" + ) + if name in _COMPUTED_RECORD_ATTRS: + raise ValueError( + f"{name!r} is written from the wavefront itself and cannot be " + "set through `other`" + ) + + @classmethod + def pmd_keys(cls): + """ + Map field names to their openPMD attribute names. + + Returns + ------- + dict + ``{field_name: pmd_key}``, excluding `other`. + """ + return { + fld.name: fld.metadata["pmd_key"] + for fld in fields(cls) + if "pmd_key" in fld.metadata + } + + @classmethod + def from_pmd(cls, attrs): + """ + Build from a mapping keyed by openPMD attribute names. + + Parameters + ---------- + attrs : mapping + Attribute names to values. Field names are also accepted, so that a + dict written in Python style round-trips. A nested `other` mapping is + merged into `other` rather than nested inside it. Anything unrecognized + is kept in `other`. + + Returns + ------- + WavefrontAttrs + + Raises + ------ + ValueError + If an attribute is given under both its openPMD name and its field + name. Silently preferring one would write the other to the file as a + nonstandard attribute holding a conflicting value. + """ + if isinstance(attrs, cls): + return attrs.copy() + + remaining = dict(attrs) + kwargs = {} + for name, key in cls.pmd_keys().items(): + has_key = key in remaining + # Some fields spell the two the same way, which is not a collision. + has_name = name != key and name in remaining + if has_key and has_name: + raise ValueError( + f"{key!r} and {name!r} are two spellings of the same attribute " + f"and both were given, with values {remaining[key]!r} and " + f"{remaining[name]!r}" + ) + if has_key: + kwargs[name] = remaining.pop(key) + elif has_name: + kwargs[name] = remaining.pop(name) + + # `other` is a field of this class, not a record attribute, so a mapping + # carrying one means Python-style input rather than something read from a + # file. Merge it instead of nesting it, which would produce an `other` + # entry whose value is a dict and fail at write time. + other = dict(remaining.pop("other", {})) + other.update(remaining) + + return cls(**kwargs, other=other) + + def to_pmd(self): + """ + Render as a mapping keyed by openPMD attribute names. + + Returns + ------- + dict + Set fields plus `other`. Fields left as None are omitted, since the + extension treats them as absent rather than zero. + """ + out = {} + for name, key in self.pmd_keys().items(): + value = getattr(self, name) + if value is not None: + out[key] = value + out.update(self.other) + return out + + def copy(self): + """ + Return an independent copy, including `other`. + + Returns + ------- + WavefrontAttrs + """ + return type(self)( + **{name: getattr(self, name) for name in self.pmd_keys()}, + other=dict(self.other), + ) + + +def photon_energy_joules(wavelength): + """ + Central photon energy in joules for a given wavelength. + + Parameters + ---------- + wavelength : float + Central wavelength in m. + + Returns + ------- + float + Photon energy ``h c / wavelength`` in J. + + Notes + ----- + The extension gives `photonEnergy` a `unitDimension` of energy but, as an + attribute, it carries no `unitSI`. SI (joules) is written here. + """ + return h_planck * c_light / wavelength + + +def wavelength_from_photon_energy(photon_energy): + """ + Central wavelength for a photon energy in joules. + + Parameters + ---------- + photon_energy : float + Photon energy in J. + + Returns + ------- + float + Wavelength ``h c / photon_energy`` in m. + """ + return h_planck * c_light / photon_energy + + +def _axis_permutation(source, target): + """ + Permutation taking an array indexed by `source` axes to `target` axis order. + + Parameters + ---------- + source : sequence of str + Axis labels of the array being permuted, in its own axis order. + target : sequence of str + Desired axis label order. + + Returns + ------- + tuple of int + Argument for `numpy.ndarray.transpose`. + """ + return tuple(source.index(label) for label in target) + + +def _write_attr(group, name, value): + """ + Write one attribute, encoded the way the rest of the package encodes them. + + Parameters + ---------- + group : h5py.Group + Target group. + name : str + Attribute name. + value : object + Attribute value. + """ + group.attrs[name] = encode_attr(value) + + +def write_wavefront_openpmd( + wavefront, + h5, + iteration=1, + **extension_attrs, +): + """ + Write a `Wavefront` into an open HDF5 group as an openPMD EXT_Wavefront series. + + Parameters + ---------- + wavefront : Wavefront + Real-space, time-domain wavefront. `Ex` and `Ey` are in V/m with shape + ``(nx, ny, nz)``; an absent polarization is not written. Its `s_position` + is written as the extension's required `zCoordinate`. + h5 : h5py.Group + Group to use as the openPMD series root. + iteration : int, default=1 + openPMD iteration index to write under `basePath`. + **extension_attrs + `WavefrontAttrs` field names (`beamline`, `radius_of_curvature_x`, ...), + taking precedence over the same fields in ``wavefront.attrs``. + + Raises + ------ + TypeError + If an attribute name is not a `WavefrontAttrs` field. + + Notes + ----- + Datasets are stored in ``(z, y, x)`` order and written one transverse slice at a + time, so no full transposed copy of the field is ever materialized. + """ + attrs = WavefrontAttrs.from_pmd(wavefront.attrs) + for name, value in extension_attrs.items(): + # Assigning an unknown name would silently stick a new instance attribute on + # the dataclass, so check first and let the constructor raise TypeError. + if name not in attrs.pmd_keys(): + WavefrontAttrs(**{name: value}) + setattr(attrs, name, value) + + carried = attrs.to_pmd() + + # Series (root) attributes. + _write_attr(h5, "openPMD", "2.0.0") + _write_attr(h5, "openPMDextension", "Wavefront") + _write_attr(h5, "basePath", _DEFAULT_BASE_PATH) + _write_attr(h5, "meshesPath", _DEFAULT_MESHES_PATH) + _write_attr(h5, "iterationEncoding", "groupBased") + _write_attr(h5, "iterationFormat", _DEFAULT_BASE_PATH) + + base = _DEFAULT_BASE_PATH.replace("%T", str(int(iteration))).strip("/") + iteration_group = h5.require_group(base) + iteration_group.attrs["time"] = 0.0 + iteration_group.attrs["dt"] = 0.0 + iteration_group.attrs["timeUnitSI"] = 1.0 + + mesh = h5.require_group(f"{base}/{_DEFAULT_MESHES_PATH}{_RECORD_NAME}") + + # Grid quantities, ordered like the stored axes. `gridGlobalOffset` is the + # position of the beginning of the first cell; the components declare + # `position = 0`, so that is exactly the first sample of each axis. + spacing = {"x": wavefront.dx, "y": wavefront.dy, "z": wavefront.dz} + offset = {"x": wavefront.xmin, "y": wavefront.ymin, "z": wavefront.zmin} + + # Base standard, required on the mesh record. + _write_attr(mesh, "geometry", "cartesian") + mesh.attrs["axisLabels"] = encode_attr(_STORED_AXIS_LABELS) + mesh.attrs["gridSpacing"] = np.array( + [spacing[label] for label in _STORED_AXIS_LABELS], dtype=float + ) + mesh.attrs["gridGlobalOffset"] = np.array( + [offset[label] for label in _STORED_AXIS_LABELS], dtype=float + ) + # openPMD 2.0 makes gridUnitSI one value per axis. + mesh.attrs["gridUnitSI"] = np.ones(len(_STORED_AXIS_LABELS), dtype=float) + mesh.attrs["gridUnitDimension"] = np.array( + _UNIT_DIMENSION_LENGTH * len(_STORED_AXIS_LABELS), dtype=float + ) + mesh.attrs["unitDimension"] = np.array(_UNIT_DIMENSION_E_FIELD, dtype=float) + mesh.attrs["timeOffset"] = 0.0 + + # Extension, required on the mesh record. + mesh.attrs["photonEnergy"] = photon_energy_joules(wavefront.wavelength) + _write_attr(mesh, "temporalDomain", "time") + _write_attr(mesh, "spatialDomain", "r") + mesh.attrs["zCoordinate"] = float(wavefront.s_position) + + # Extension, optional. + for name, value in carried.items(): + _write_attr(mesh, name, value) + + # Components. `z` is never written. + to_stored = _axis_permutation(_CLASS_AXIS_LABELS, _STORED_AXIS_LABELS) + for name, field_array in (("x", wavefront.Ex), ("y", wavefront.Ey)): + if field_array is None: + continue + + # A view, not a copy: the transpose is realized one slice at a time below. + stored_view = np.asarray(field_array).transpose(to_stored) + + dataset = mesh.create_dataset( + name, shape=stored_view.shape, dtype=field_array.dtype + ) + for islice in range(stored_view.shape[0]): + dataset[islice] = stored_view[islice] + + dataset.attrs["unitSI"] = 1.0 + dataset.attrs["position"] = np.zeros(len(_STORED_AXIS_LABELS), dtype=float) + + +def _iteration_group(h5, iteration=None): + """ + Return the openPMD iteration group, honoring `basePath`. + + Parameters + ---------- + h5 : h5py.Group + Series root. + iteration : int, optional + Iteration to select. If None, the sole iteration is used. + + Returns + ------- + h5py.Group + + Raises + ------ + ValueError + If the base path is missing, if there are no iterations, or if `iteration` + is None and the file holds more than one. + """ + base_path = decode_attr(h5.attrs.get("basePath", _DEFAULT_BASE_PATH)) + + if "%T" not in base_path: + raise ValueError(f"basePath {base_path!r} has no %T iteration placeholder") + + parent_path = base_path.split("%T")[0].strip("/") + if parent_path not in h5: + raise ValueError( + f"basePath {base_path!r} points at {parent_path!r}, which is not in the file" + ) + parent = h5[parent_path] + + if iteration is not None: + key = str(int(iteration)) + if key not in parent: + raise ValueError( + f"iteration {key} not in the file. Available: {sorted(parent)}" + ) + return parent[key] + + available = sorted(parent) + if not available: + raise ValueError(f"no iterations under {parent_path!r}") + if len(available) > 1: + raise ValueError( + f"file holds {len(available)} iterations {available}; " + "pass iteration= to select one" + ) + return parent[available[0]] + + +def _mesh_record(h5, iteration=None): + """ + Return the `electricField` mesh record group. + + Parameters + ---------- + h5 : h5py.Group + Series root. + iteration : int, optional + Iteration to select. + + Returns + ------- + h5py.Group + + Raises + ------ + ValueError + If the record is absent. + """ + iteration_group = _iteration_group(h5, iteration) + meshes_path = decode_attr(h5.attrs.get("meshesPath", _DEFAULT_MESHES_PATH)) + + record_path = f"{meshes_path}{_RECORD_NAME}" + if record_path not in iteration_group: + raise ValueError( + f"no {_RECORD_NAME!r} mesh record at " + f"{iteration_group.name}/{record_path}: not an EXT_Wavefront file" + ) + return iteration_group[record_path] + + +def _required_attr(mesh, name, fallback=None): + """ + Read a required attribute from the mesh record, with an optional fallback group. + + Parameters + ---------- + mesh : h5py.Group + The mesh record. + name : str + Attribute name. + fallback : h5py.Group, optional + Group to consult if `mesh` lacks the attribute. Used for `photonEnergy`, + which the extension places ambiguously. + + Returns + ------- + object + The decoded attribute value. + + Raises + ------ + ValueError + If the attribute is present in neither location. + """ + if name in mesh.attrs: + return decode_attr(mesh.attrs[name]) + if fallback is not None and name in fallback.attrs: + return decode_attr(fallback.attrs[name]) + raise ValueError(f"required EXT_Wavefront attribute {name!r} is missing") + + +def load_wavefront_openpmd(h5, iteration=None): + """ + Read an openPMD EXT_Wavefront series into `Wavefront` constructor arguments. + + Parameters + ---------- + h5 : h5py.Group + Series root. + iteration : int, optional + Iteration to read. If None, the file must hold exactly one. + + Returns + ------- + dict + Keyword arguments for `Wavefront`: `Ex`, `Ey`, `dx`, `dy`, `dz`, + `wavelength`, `attrs` (a `WavefrontAttrs`) and, when the file declares a + `gridGlobalOffset`, `xmid`, `ymid` and `zmid`. + + Raises + ------ + ValueError + If a required attribute is missing, or if the file holds something this + class cannot represent: a frequency-domain field, a k-space field, or + `axisLabels` that are not a permutation of x, y and z. + + Notes + ----- + Any `axisLabels` permutation of ``(x, y, z)`` is honored, with `gridSpacing` + permuted alongside it. Scalar attributes stored as length-1 arrays and strings + stored as bytes are both tolerated. `gridUnitSI`, `gridGlobalOffset` and the + per-component `unitSI` are applied, so a file written in non-SI units reads + correctly; `gridUnitSI` is accepted either as the openPMD 2.0 per-axis array or + as the 1.x scalar. + """ + mesh = _mesh_record(h5, iteration) + + temporal_domain = _required_attr(mesh, "temporalDomain") + if temporal_domain != "time": + raise ValueError( + f"temporalDomain {temporal_domain!r}: only 'time' (a field in V/m) " + "is implemented" + ) + + spatial_domain = _required_attr(mesh, "spatialDomain") + if spatial_domain != "r": + raise ValueError( + f"spatialDomain {spatial_domain!r}: only 'r' (cartesian space) " + "is implemented" + ) + + photon_energy = float(_required_attr(mesh, "photonEnergy", fallback=h5)) + + raw_labels = _required_attr(mesh, "axisLabels") + labels = tuple(decode_attr(label) for label in np.atleast_1d(raw_labels)) + if sorted(labels) != sorted(_CLASS_AXIS_LABELS): + raise ValueError( + f"axisLabels {labels}: only a permutation of " + f"{_CLASS_AXIS_LABELS} is implemented" + ) + + raw_spacing = np.atleast_1d(_required_attr(mesh, "gridSpacing")) + if len(raw_spacing) != len(labels): + raise ValueError( + f"gridSpacing has {len(raw_spacing)} values but axisLabels has " + f"{len(labels)}" + ) + + # gridSpacing is in the file's own units; gridUnitSI converts it to meters. + # openPMD 2.0 makes gridUnitSI one value per axis, but implementations written + # against 1.x emit a single scalar, so accept either. + raw_grid_unit = np.atleast_1d(decode_attr(mesh.attrs.get("gridUnitSI", 1.0))) + if raw_grid_unit.size == 1: + grid_unit = np.full(len(labels), float(raw_grid_unit[0])) + elif raw_grid_unit.size == len(labels): + grid_unit = raw_grid_unit.astype(float) + else: + raise ValueError( + f"gridUnitSI has {raw_grid_unit.size} values but axisLabels has " + f"{len(labels)}" + ) + + spacing = { + label: float(value) * float(unit) + for label, value, unit in zip(labels, raw_spacing, grid_unit) + } + + # gridGlobalOffset is the position of the first cell, in the same units as + # gridSpacing. A file that omits it is read as a centered grid, which is this + # class's own default rather than an assertion about the file. + global_offset = None + if "gridGlobalOffset" in mesh.attrs: + raw_offset = np.atleast_1d(decode_attr(mesh.attrs["gridGlobalOffset"])) + if len(raw_offset) != len(labels): + raise ValueError( + f"gridGlobalOffset has {len(raw_offset)} values but axisLabels has " + f"{len(labels)}" + ) + global_offset = { + label: float(value) * float(unit) + for label, value, unit in zip(labels, raw_offset, grid_unit) + } + + # Everything on the record that this module does not write from the wavefront + # itself, including names it does not know: those land in `other` so that a file + # written against a newer revision of the extension survives a round trip. + pmd_attrs = { + name: decode_attr(value) + for name, value in mesh.attrs.items() + if name not in _COMPUTED_RECORD_ATTRS + } + if "zCoordinate" not in mesh.attrs: + raise ValueError("required EXT_Wavefront attribute 'zCoordinate' is missing") + attrs = WavefrontAttrs.from_pmd(pmd_attrs) + + kwargs = { + "dx": spacing["x"], + "dy": spacing["y"], + "dz": spacing["z"], + "wavelength": wavelength_from_photon_energy(photon_energy), + "s_position": float(decode_attr(mesh.attrs["zCoordinate"])), + "Ex": None, + "Ey": None, + "attrs": attrs, + } + + # Permutation of the class's (x, y, z) axes into the file's stored order, so the + # transposed view of the output array lines up with the dataset slice by slice. + class_to_stored = _axis_permutation(_CLASS_AXIS_LABELS, labels) + for name, key in (("x", "Ex"), ("y", "Ey")): + if name not in mesh: + continue + dataset = mesh[name] + + # openPMD allows a uniform component to be stored as a group carrying only + # `value` and `shape` instead of a dataset. It has no axes to permute. + if is_constant_component(dataset): + stored_shape = tuple(decode_attr(dataset.attrs["shape"])) + kwargs[key] = np.full( + tuple( + stored_shape[labels.index(label)] for label in _CLASS_AXIS_LABELS + ), + constant_component_value(dataset), + ) + continue + + shape = tuple( + dataset.shape[labels.index(label)] for label in _CLASS_AXIS_LABELS + ) + out = np.empty(shape, dtype=dataset.dtype) + + # A view of `out` in the file's axis order. Filling it slice by slice avoids + # allocating a second full-size array for the transpose. + stored_view = out.transpose(class_to_stored) + for islice in range(dataset.shape[0]): + stored_view[islice] = dataset[islice] + + # Component values are in the file's own units; unitSI converts to V/m. + # Multiply in place so that a complex64 field is not silently widened. + unit_si = float(np.atleast_1d(decode_attr(dataset.attrs.get("unitSI", 1.0)))[0]) + if unit_si != 1.0: + out *= unit_si + + kwargs[key] = out + + # The class's grid fields are the axis *midpoint*, while gridGlobalOffset is + # the first sample, so converting between them needs the sample count. + sample = kwargs["Ex"] if kwargs["Ex"] is not None else kwargs["Ey"] + if global_offset is not None and sample is not None: + counts = dict(zip(_CLASS_AXIS_LABELS, sample.shape)) + for label in _CLASS_AXIS_LABELS: + kwargs[f"{label}mid"] = ( + global_offset[label] + (counts[label] - 1) * spacing[label] / 2 + ) + + return kwargs diff --git a/beamphysics/wavefront/propagators.py b/beamphysics/wavefront/propagators.py index 0e892dd5..63066381 100644 --- a/beamphysics/wavefront/propagators.py +++ b/beamphysics/wavefront/propagators.py @@ -56,7 +56,13 @@ def fftfreq(*args, **kwargs): continue new_fields.append(kernel * field) Ex_drifted, Ey_drifted = new_fields - return replace(w, Ex=Ex_drifted, Ey=Ey_drifted) + return replace( + w, + Ex=Ex_drifted, + Ey=Ey_drifted, + attrs=w.attrs.copy(), + s_position=w.s_position + z, + ) # r-space calculation @@ -78,9 +84,11 @@ def fftfreq(*args, **kwargs): # Allocate an output array of the same shape/dtype # Handle dtype (TODO: more general) - if device == "mps" or field.dtype == backend.complex64: - dtype = backend.complex64 - elif field.dtype == backend.float32: + if ( + device == "mps" + or field.dtype == backend.complex64 + or field.dtype == backend.float32 + ): dtype = backend.complex64 elif field.dtype == backend.float64: dtype = backend.complex128 @@ -98,7 +106,16 @@ def fftfreq(*args, **kwargs): Ex_drifted, Ey_drifted = new_fields - return replace(w, Ex=Ex_drifted, Ey=Ey_drifted) + # The mesh is co-moving, so `s_position` is the only place a record of the + # propagation can go. `attrs` is copied because it is mutable and shared + # otherwise: `replace` rebinds the same object onto the new wavefront. + return replace( + w, + Ex=Ex_drifted, + Ey=Ey_drifted, + attrs=w.attrs.copy(), + s_position=w.s_position + z, + ) def drift_wavefront_advanced(w, z, backend=np, device="cpu", curvature=1.0): @@ -109,11 +126,13 @@ def drift_wavefront_advanced(w, z, backend=np, device="cpu", curvature=1.0): if z == 0: return w.copy() + s_final = w.s_position + z + x_mesh, y_mesh, _ = backend.meshgrid(w.xvec, w.yvec, w.zvec, indexing="ij") curv = np.exp(-1j * backend.pi * (x_mesh**2 + y_mesh**2) / w.wavelength * curvature) - w = replace(w) + w = replace(w, attrs=w.attrs.copy()) w.Ex = w.Ex * curv if w.Ex is not None else None w.Ey = w.Ey * curv if w.Ey is not None else None @@ -137,4 +156,16 @@ def drift_wavefront_advanced(w, z, backend=np, device="cpu", curvature=1.0): w.dx = w.dx / M w.dy = w.dy / M + # Both quadratic phases above are referenced to x = y = 0, so this is a + # magnification about the optical axis: the map is x -> x/M. The grid midpoint + # has to scale with the spacing, or an off-axis feature would keep its + # coordinate while the grid around it expanded. + w.xmid = w.xmid / M + w.ymid = w.ymid / M + + # The beamline advances by the requested distance. `drift_wavefront_basic` above + # was handed `z_eff`, which is a scaled-frame device rather than a distance + # anything travels, so its increment has to be overwritten rather than kept. + w.s_position = s_final + return w diff --git a/beamphysics/wavefront/wavefront.py b/beamphysics/wavefront/wavefront.py index 6466314d..51bdd620 100644 --- a/beamphysics/wavefront/wavefront.py +++ b/beamphysics/wavefront/wavefront.py @@ -3,8 +3,9 @@ import pathlib from abc import ABC, abstractmethod from copy import deepcopy -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from enum import Enum +from html import escape from math import pi from typing import ClassVar, Union @@ -22,6 +23,11 @@ from ..plot import plot_1d_density, plot_2d_density_with_marginals from ..statistics import mean_calc, mean_variance_calc from ..units import Z0, c_light +from ..wavefront.openpmd import ( + WavefrontAttrs, + load_wavefront_openpmd, + write_wavefront_openpmd, +) from ..wavefront.propagators import drift_wavefront @@ -75,6 +81,39 @@ class WavefrontBase(ABC): wavelength: float = 1.0 # m axis_labels: ClassVar[tuple[str, str, str]] = ("", "", "") + # Declared as a real dataclass field on each concrete subclass. + attrs: WavefrontAttrs + + # Midpoint of the grid in m, defaulting to a grid centered on the origin. + # `xmin` and `xmax` follow from these. They are real-space quantities even on + # `WavefrontK`, which carries them through the transform without interpreting + # them: a real-space shift is a linear phase in k-space, not a shift of the + # k grid. + xmid: float = 0.0 + ymid: float = 0.0 + zmid: float = 0.0 + + # Position of this plane along the beamline in m, in whatever frame the caller + # picked. Advanced by `drift`, since the mesh is co-moving and has nowhere else + # to record that the wavefront has propagated. Written as the extension's + # `zCoordinate` and as Genesis4's `refposition`. + s_position: float = 0.0 + + def __setattr__(self, name, value): + """ + Keep `attrs` a `WavefrontAttrs` however it is set. + + A mapping keyed by either openPMD or field names is accepted for + convenience, but the attribute itself is always the typed object. Coercing + only in `__post_init__` would leave a plain assignment such as + ``w.attrs = {...}`` holding a raw dict, so that `w.attrs.beamline` + raises `AttributeError` and the mapping's contents go unvalidated until + write time. + """ + if name == "attrs" and not isinstance(value, WavefrontAttrs): + value = WavefrontAttrs.from_pmd(value) + object.__setattr__(self, name, value) + def __post_init__(self): """ Validate inputs after dataclass initialization @@ -192,18 +231,18 @@ def xmin(self): """ Minimum x coordinate in m. - xmin = -(nx-1) * dx / 2 + xmin = xmid - (nx-1) * dx / 2 """ - return -((self.nx - 1) * self.dx) / 2 + return self.xmid - ((self.nx - 1) * self.dx) / 2 @property def xmax(self): """ Maximum x coordinate in m. - xmax = (nx-1) * dx / 2 + xmax = xmid + (nx-1) * dx / 2 """ - return ((self.nx - 1) * self.dx) / 2 + return self.xmid + ((self.nx - 1) * self.dx) / 2 @property def xvec(self): @@ -218,18 +257,18 @@ def ymin(self): """ Minimum y coordinate in m. - ymin = -(ny-1) * dy / 2 + ymin = ymid - (ny-1) * dy / 2 """ - return -((self.ny - 1) * self.dy) / 2 + return self.ymid - ((self.ny - 1) * self.dy) / 2 @property def ymax(self): """ Maximum y coordinate in m. - ymax = (ny-1) * dy / 2 + ymax = ymid + (ny-1) * dy / 2 """ - return ((self.ny - 1) * self.dy) / 2 + return self.ymid + ((self.ny - 1) * self.dy) / 2 @property def yvec(self): @@ -244,18 +283,18 @@ def zmin(self): """ Minimum z coordinate in m. - zmin = -(nz-1) * dz / 2 + zmin = zmid - (nz-1) * dz / 2 """ - return -((self.nz - 1) * self.dz) / 2 + return self.zmid - ((self.nz - 1) * self.dz) / 2 @property def zmax(self): """ Maximum z coordinate in m. - zmax = (nz-1) * dz / 2 + zmax = zmid + (nz-1) * dz / 2 """ - return ((self.nz - 1) * self.dz) / 2 + return self.zmid + ((self.nz - 1) * self.dz) / 2 @property def zvec(self): @@ -440,6 +479,25 @@ def _std(self, key): def drift(self, z, curvature=0): return drift_wavefront(self, z, curvature=curvature) + def _resized_mids(self, nx, ny, nz): + """ + Grid midpoint fields after each axis is resized. + + Parameters + ---------- + nx, ny, nz : tuple of (int, int) + Elements removed from the (start, end) of each axis. Negative values + mean elements were added. + + Returns + ------- + dict + Fields to pass to `dataclasses.replace`. Empty here, because the axes + of a k-space wavefront are not real-space axes: discarding k samples + does not move the real-space grid. + """ + return {} + def pad(self, nx=(0, 0), ny=(0, 0), nz=(0, 0)): """ zero-pad the field arrays. @@ -455,7 +513,9 @@ def pad(self, nx=(0, 0), ny=(0, 0), nz=(0, 0)): Ex = np.pad(self.Ex, (nx, ny, nz)) if self.Ex is not None else None Ey = np.pad(self.Ey, (nx, ny, nz)) if self.Ey is not None else None - return replace(self, Ex=Ex, Ey=Ey) + # Padding adds elements, which is a removal of a negative count. + mids = self._resized_mids((-nx[0], -nx[1]), (-ny[0], -ny[1]), (-nz[0], -nz[1])) + return replace(self, Ex=Ex, Ey=Ey, attrs=self.attrs.copy(), **mids) def crop(self, nx=(0, 0), ny=(0, 0), nz=(0, 0)): """ @@ -482,6 +542,11 @@ def crop(self, nx=(0, 0), ny=(0, 0), nz=(0, 0)): ------ ValueError If crop amounts exceed array dimensions. + + Notes + ----- + The surviving samples keep their physical coordinates: an asymmetric crop + moves the grid origin rather than re-centering the grid. """ nx = (nx, nx) if np.isscalar(nx) else nx ny = (ny, ny) if np.isscalar(ny) else ny @@ -520,7 +585,13 @@ def crop(self, nx=(0, 0), ny=(0, 0), nz=(0, 0)): else None ) - return replace(self, Ex=Ex, Ey=Ey) + return replace( + self, + Ex=Ex, + Ey=Ey, + attrs=self.attrs.copy(), + **self._resized_mids(nx, ny, nz), + ) def auto_crop(self, threshold: float = 1e-6, apply: bool = True): """ @@ -548,16 +619,28 @@ def auto_crop(self, threshold: float = 1e-6, apply: bool = True): Examples -------- - >>> # Get crop amounts without applying - >>> crop_info = w.auto_crop(threshold=1e-4, apply=False) - >>> print(crop_info) # {'nx': 10, 'ny': 8, 'nz': 5} + >>> from beamphysics.wavefront import Wavefront + >>> w = Wavefront.from_gaussian( + ... shape=(64, 64, 16), dx=1e-6, dy=1e-6, dz=1e-6, + ... wavelength=1e-9, sigma0=3e-6, sigma_z=2e-6, + ... ) - >>> # Apply auto-crop directly - >>> w_cropped = w.auto_crop(threshold=1e-4) + Get the crop amounts without applying them: - >>> # Manual two-step process >>> crop_info = w.auto_crop(threshold=1e-4, apply=False) - >>> w_cropped = w.crop(**crop_info) + >>> crop_info + {'nx': 19, 'ny': 19, 'nz': 0} + + Apply the crop directly, removing that many elements from each end: + + >>> w.auto_crop(threshold=1e-4).shape + (26, 26, 16) + + The two steps are equivalent, so the amounts can be inspected or adjusted + first: + + >>> w.crop(**crop_info).shape + (26, 26, 16) """ def find_symmetric_crop(profile: np.ndarray, threshold: float) -> int: @@ -585,7 +668,8 @@ def find_symmetric_crop(profile: np.ndarray, threshold: float) -> int: crop_start = first_idx crop_end = len(profile) - 1 - last_idx - return min(crop_start, crop_end) + # A plain int, so that the returned dict reprs and serializes cleanly. + return int(min(crop_start, crop_end)) # Get 1D intensity projections along each axis intensity = self.intensity @@ -609,86 +693,108 @@ def copy(self): """Returns a deep copy""" return deepcopy(self) + def _summary_rows(self): + """ + Label/value pairs describing this wavefront, shared by both reprs. + + Rows that carry no information are omitted, so that the common case stays + short: a grid centered on the origin, a wavefront at the origin of the + beamline frame, and unset metadata each contribute nothing. Metadata is + listed field by field rather than as the `WavefrontAttrs` repr, which + spells out every unset field. + + Returns + ------- + list of (str, str) + """ + present = [name for name in ("Ex", "Ey") if getattr(self, name) is not None] + + # %g rather than a fixed exponent, so that a spacing of 1 reads as "1" + # instead of "1.000e+00". + def g(value): + return f"{value:g}" + + # Report each class's own axes: dx/dy/dz here, dkx/dky/dkz in k-space. + spacing = ", ".join( + f"d{label}={g(getattr(self, 'd' + label))}" for label in self.axis_labels + ) + + rows = [ + ("wavelength", f"{g(self.wavelength)} m"), + ("photon energy", f"{g(self.photon_energy)} eV"), + ("grid shape", f"{self.shape}"), + ("spacing", f"{spacing} {'m' if self.in_rspace else 'rad/m'}"), + ("fields", ", ".join(present) if present else "None"), + ] + + if (self.xmid, self.ymid, self.zmid) != (0.0, 0.0, 0.0): + # Always a real-space quantity, so say so when the axes are not. + label = "grid midpoint" if self.in_rspace else "grid midpoint (r-space)" + rows.append( + (label, f"x={g(self.xmid)}, y={g(self.ymid)}, z={g(self.zmid)} m") + ) + if self.s_position != 0.0: + rows.append(("s_position", f"{g(self.s_position)} m")) + + for name in self.attrs.pmd_keys(): + value = getattr(self.attrs, name) + if value is not None: + rows.append((name, str(value))) + # Keyed by openPMD name, since that is how they arrived and how they leave. + for name, value in self.attrs.other.items(): + rows.append((name, str(value))) + + return rows + + def __repr__(self): + # The subclasses set `repr=False` so that this is used instead of the + # dataclass repr, which prints the entire field array. Deliberately short, + # since this appears inside containers and tracebacks; the full summary is + # in the pretty and HTML reprs. + return ( + f"{self.__class__.__name__}(shape={self.shape}, " + f"wavelength={self.wavelength:g} m, s_position={self.s_position:g} m)" + ) + def _repr_pretty_(self, p, cycle): """IPython/Jupyter pretty-print representation""" if cycle: p.text(f"{self.__class__.__name__}(...)") return - def summarize_field(field, name): - if field is None: - return f"{name}: None" - return f"{name}: {field.shape}" - - lines = [ - f"{self.__class__.__name__}(", - f" wavelength: {self.wavelength:.6e} m", - f" shape: {self.shape}", - f" spacing: dx={self.dx:.3e}, dy={self.dy:.3e}, dz={self.dz:.3e} m", - f" {summarize_field(self.Ex, 'Ex')}", - f" {summarize_field(self.Ey, 'Ey')}", - ")", - ] + lines = [f"{self.__class__.__name__}("] + lines += [f" {label}: {value}" for label, value in self._summary_rows()] + lines.append(")") p.text("\n".join(lines)) def _repr_html_(self): """Rich HTML representation for Jupyter notebooks""" - # Determine which fields exist - fields = [] - if self.Ex is not None: - fields.append("Ex") - if self.Ey is not None: - fields.append("Ey") - field_str = ", ".join(fields) if fields else "None" - - # Add photon energy for easier reference - photon_energy_str = "" - if hasattr(self, "photon_energy"): - photon_energy_str = f""" - - photon energy - {self.photon_energy:.6e} eV - - """ - - fmt = "" - - html = f""" + # Labels and values can come from a file, since `attrs.other` carries + # whatever the file held. This string is injected into the notebook DOM, + # so escape both rather than trusting them. + cells = [] + for i, (label, value) in enumerate(self._summary_rows()): + shade = ( + ' style="background-color: rgba(128,128,128,0.1);"' + if i % 2 == 0 + else "" + ) + cells.append( + f"{escape(label)}" + f"{escape(value)}" + ) + + return f"""

{self.__class__.__name__}

- - - - - {photon_energy_str} - - - - - - - - - - - - - - - - - - - - + {"".join(cells)}
wavelength{self.wavelength:{fmt}} m
grid shape{self.shape}
dx{self.dx:{fmt}} m
dy{self.dy:{fmt}} m
dz{self.dz:{fmt}} m
fields{field_str}
""" - return html -@dataclass +@dataclass(repr=False) class WavefrontK(WavefrontBase): """ K-space (Fourier) representation of electromagnetic wavefront fields. @@ -714,6 +820,17 @@ class WavefrontK(WavefrontBase): Grid spacing in kz direction (rad/m) wavelength : float, default=1.0 Central wavelength (m) + attrs : WavefrontAttrs or mapping, optional + openPMD EXT_Wavefront attributes carried through file I/O, such as + `beamline`. Not interpreted by this class. A mapping keyed by either + openPMD or field names is coerced to `WavefrontAttrs`. + xmid, ymid, zmid : float, default=0.0 + Midpoint of the *real-space* grid in m, carried through from the real-space + wavefront. Not interpreted here: a real-space shift is a linear phase in + k-space, not a shift of the k grid. + s_position : float, default=0.0 + Beamline position of this plane in m, carried through from the real-space + wavefront. Attributes ---------- @@ -771,6 +888,15 @@ class WavefrontK(WavefrontBase): wavelength: float = 1.0 # m + # Provenance carried through from file I/O. Not interpreted by this class. + attrs: WavefrontAttrs = field(default_factory=WavefrontAttrs) + + # Real-space grid midpoint, carried through the transform. See the class docstring. + xmid: float = 0.0 # m + ymid: float = 0.0 # m + zmid: float = 0.0 # m + s_position: float = 0.0 # m + axis_labels: ClassVar[tuple[str, str, str]] = ("kx", "ky", "kz") @property @@ -838,6 +964,11 @@ def to_rspace(self, *, inplace: bool = False) -> Wavefront: dy=self.dy, dz=self.dz, wavelength=self.wavelength, + attrs=self.attrs.copy(), + xmid=self.xmid, + ymid=self.ymid, + zmid=self.zmid, + s_position=self.s_position, ) @property @@ -1074,7 +1205,7 @@ def sigma_thetay(self): return self._std("thetay") -@dataclass +@dataclass(repr=False) class Wavefront(WavefrontBase): """ Real-space representation of electromagnetic wavefront fields. @@ -1101,6 +1232,18 @@ class Wavefront(WavefrontBase): Grid spacing in z direction (m) wavelength : float, default=1.0 Central wavelength (m) + attrs : WavefrontAttrs or mapping, optional + openPMD EXT_Wavefront attributes carried through file I/O, such as + `beamline`. Not interpreted by this class. A mapping keyed by either + openPMD or field names is coerced to `WavefrontAttrs`. + xmid, ymid, zmid : float, default=0.0 + Midpoint of the grid in m. The default of 0.0 gives a grid centered on the + origin, so that ``xvec`` runs from ``-(nx-1)*dx/2`` to ``+(nx-1)*dx/2``. + Note that `z` here is the intra-pulse coordinate, not the beamline + position: see the note on frames below. + s_position : float, default=0.0 + Position of this plane along the beamline in m, in whatever frame the + caller picked. Advanced by `drift`. See the note on frames below. Attributes ---------- @@ -1141,6 +1284,10 @@ class Wavefront(WavefrontBase): Create Wavefront from Genesis4 HDF5 field file write_genesis4(file) Write Wavefront to Genesis4 HDF5 format + from_openpmd(file, iteration=None) + Create Wavefront from an openPMD EXT_Wavefront file + write_openpmd(file, iteration=1, ...) + Write Wavefront to an openPMD EXT_Wavefront file Notes ----- @@ -1151,6 +1298,24 @@ class Wavefront(WavefrontBase): ∫∫∫ |E(x,y,z)|² dx dy dz = ∫∫∫ |Ẽ(kx,ky,kz)|² dkx dky dkz - Statistical moments (mean, sigma) are intensity-weighted + **Frames.** There are two longitudinal quantities and they belong to different + frames. The mesh's own `z` axis, positioned by `zmid`, is the intra-pulse + coordinate: it is co-moving, and `drift` leaves it alone because the paraxial + kernel is transverse, applied slice by slice with no piston term. `s_position` + is the position of this plane along the beamline, and `drift` does advance it, + by the propagation distance. + + That split follows `ParticleGroup`, where `drift` moves the coordinates that + exist rather than a separate register: there, `z` records that the particles + propagated. Here the mesh is co-moving and cannot record it, so `s_position` is + the only place the information can go. Operations that do not move the + wavefront along the beamline -- `crop`, `pad`, `to_kspace`, an applied lens + phase -- leave it alone, just as they leave `ParticleGroup.z` alone. + + `s_position` is measured in whatever frame the caller picked, so the default of + 0.0 is the origin of that frame rather than an unknown. It is written as the + EXT_Wavefront `zCoordinate` and as Genesis4's `refposition`. + """ Ex: np.ndarray | None = None # V/m @@ -1161,8 +1326,44 @@ class Wavefront(WavefrontBase): dz: float = 1.0 # m # type: ignore[override] wavelength: float = 1.0 # m + # Provenance carried through from file I/O. Not interpreted by this class. + attrs: WavefrontAttrs = field(default_factory=WavefrontAttrs) + + # Grid midpoint. The default of 0 is a grid centered on the origin. + xmid: float = 0.0 # m + ymid: float = 0.0 # m + zmid: float = 0.0 # m + + # Beamline position of this plane. Advanced by `drift`. + s_position: float = 0.0 # m + axis_labels: ClassVar[tuple[str, str, str]] = ("x", "y", "z") + def _resized_mids(self, nx, ny, nz): + """ + Grid midpoint fields that keep the surviving samples where they were. + + Removing `a` elements from the start of an axis and `b` from the end moves + the midpoint of the remaining grid by ``(a - b) * d / 2``, so the stored + midpoint has to move with it. Without this, an asymmetric crop would + silently translate the field. + + Parameters + ---------- + nx, ny, nz : tuple of (int, int) + Elements removed from the (start, end) of each axis. Negative values + mean elements were added. + + Returns + ------- + dict + """ + return { + "xmid": self.xmid + (nx[0] - nx[1]) * self.dx / 2, + "ymid": self.ymid + (ny[0] - ny[1]) * self.dy / 2, + "zmid": self.zmid + (nz[0] - nz[1]) * self.dz / 2, + } + @property def spatial_domain(self) -> SpatialDomain: return SpatialDomain.R @@ -1323,6 +1524,11 @@ def to_kspace(self, backend=np, *, inplace: bool = False) -> WavefrontK: dky=self.dky, dkz=self.dkz, wavelength=self.wavelength, + attrs=self.attrs.copy(), + xmid=self.xmid, + ymid=self.ymid, + zmid=self.zmid, + s_position=self.s_position, ) def plot_power( @@ -1544,6 +1750,8 @@ def from_genesis4( ----- - The field data is extracted and converted into an electric field representation. - The grid spacing (`dx` and `dz`) and wavelength are obtained from the file metadata. + - Genesis4's `refposition` is the position of this dump along the undulator + line, so it is read into `s_position`. """ if isinstance(file, (str, pathlib.Path)): @@ -1567,6 +1775,7 @@ def from_genesis4( dy=float(dx), dz=float(dz), wavelength=float(wavelength), + s_position=float(param["refposition"]), ) def write_genesis4( @@ -1593,6 +1802,8 @@ def write_genesis4( ----- - If `file` is a path or string, a new HDF5 file is created and written to. - If `file` is an `h5py.Group`, the data is written directly into the provided group. + - `s_position` is written as Genesis4's `refposition`; `attrs` has nowhere + to go in this format and is dropped. """ if isinstance(file, (str, pathlib.Path)): with h5py.File(file, "w") as h5: @@ -1606,6 +1817,114 @@ def write_genesis4( f"file must be a str, pathlib.Path, or h5py.Group, got {type(file)}" ) + @classmethod + def from_openpmd( + cls, + file: Union[pathlib.Path, str, h5py.Group], + iteration: int | None = None, + ): + """ + Create a Wavefront from an openPMD file using the EXT_Wavefront extension. + + Parameters + ---------- + file : Union[pathlib.Path, str, h5py.Group] + Path to the openPMD HDF5 file, or an open `h5py.Group` to use as the + series root. + iteration : int, optional + openPMD iteration to read. If None, the file must hold exactly one. + + Returns + ------- + Wavefront + A new instance holding the field data, with the extension's + `zCoordinate` in `s_position` and any other record attributes carried + in `.attrs`. + + Raises + ------ + ValueError + If the file is not a valid path, string, or `h5py.Group`; if a required + attribute is missing; or if the file holds a frequency-domain or k-space + field, neither of which this class can represent. + + Notes + ----- + Any `axisLabels` permutation of (x, y, z) is honored on read, though the + writer always stores (z, y, x). + """ + if isinstance(file, (str, pathlib.Path)): + with h5py.File(file, "r") as h5: + kwargs = load_wavefront_openpmd(h5, iteration=iteration) + elif isinstance(file, h5py.Group): + kwargs = load_wavefront_openpmd(file, iteration=iteration) + else: + raise ValueError( + f"file must be a str, pathlib.Path, or h5py.Group, got {type(file)}" + ) + + return cls(**kwargs) + + def write_openpmd( + self, + file: pathlib.Path | str | h5py.Group, + iteration: int = 1, + **extension_attrs, + ): + """ + Write the Wavefront as an openPMD file using the EXT_Wavefront extension. + + Parameters + ---------- + file : Union[pathlib.Path, str, h5py.Group] + Path for a new HDF5 file, or an open `h5py.Group` to use as the series + root. + iteration : int, default=1 + openPMD iteration index to write. + **extension_attrs + `WavefrontAttrs` field names (`beamline`, `radius_of_curvature_x`, + `radius_of_curvature_y`, `delta_radius_of_curvature_x`, + `delta_radius_of_curvature_y`), taking precedence over `.attrs`. + + Raises + ------ + ValueError + If the file is not a valid path, string, or `h5py.Group`. + TypeError + If an attribute name is not a `WavefrontAttrs` field. + + Notes + ----- + Both transverse polarizations are written into one file as components `x` + and `y` of the `electricField` record; the `z` component is never written. + + The required `zCoordinate` is written from `s_position`. There is no + write-time override: it is a coordinate the propagators maintain, so a + keyword here could put a value in the file that disagrees with the + wavefront it came from. Set `s_position` on the wavefront instead. + """ + if isinstance(file, (str, pathlib.Path)): + with h5py.File(file, "w") as h5: + write_wavefront_openpmd( + self, + h5, + iteration=iteration, + **extension_attrs, + ) + return + elif isinstance(file, h5py.Group): + write_wavefront_openpmd( + self, + file, + iteration=iteration, + **extension_attrs, + ) + return + else: + raise ValueError( + f"file must be a str, pathlib.Path, or h5py.Group, got {type(file)}" + ) + def estimate_curvature( self, axis: str = "x", diff --git a/docs/api/wavefront.md b/docs/api/wavefront.md new file mode 100644 index 00000000..df8834f6 --- /dev/null +++ b/docs/api/wavefront.md @@ -0,0 +1,5 @@ +::: beamphysics.Wavefront + +::: beamphysics.WavefrontK + +::: beamphysics.wavefront.openpmd diff --git a/docs/examples/wavefront/wavefront.ipynb b/docs/examples/wavefront/wavefront.ipynb index 9ec31d8e..436d17ac 100644 --- a/docs/examples/wavefront/wavefront.ipynb +++ b/docs/examples/wavefront/wavefront.ipynb @@ -52,7 +52,6 @@ "source": [ "W = Wavefront(Ex=np.zeros((11, 11, 2)))\n", "W.Ex[6:8, 7:9, :] = 1\n", - "\n", "W" ] }, @@ -765,11 +764,95 @@ "# Cleanup\n", "os.remove(\"genesis4_field.h5\")" ] + }, + { + "cell_type": "markdown", + "id": "69ef86ab", + "metadata": {}, + "source": [ + "## openPMD\n", + "\n", + "The `.write_openpmd` method writes an openPMD 2.0 file using the [EXT_Wavefront](https://github.com/openPMD/openPMD-standard) extension, and `.from_openpmd` reads one back.\n", + "\n", + "The extension requires a `zCoordinate`: the position of this plane along the beamline, which is distinct from the mesh's own `z` axis (the intra-pulse coordinate). Because those two are easy to confuse, the Python side names it `s_position`, after the accelerator-physics `s`.\n", + "\n", + "`s_position` is a coordinate rather than metadata, so it lives on the wavefront itself and `.drift` advances it — the mesh is co-moving and has nowhere else to record that the wavefront propagated. This follows `ParticleGroup`, where `.drift` likewise moves the coordinates that exist. Operations that change the representation rather than the position, such as `.crop` or `.to_kspace`, leave it alone.\n", + "\n", + "`.attrs` carries the rest: it is a `WavefrontAttrs` dataclass, not a free-form dict. Its fields use snake_case Python names (`beamline`, `radius_of_curvature_x`, ...) that map to the extension's camelCase names on disk, so a misspelled attribute raises `TypeError` instead of being silently dropped. Attribute names the extension does not define yet are carried verbatim in `.attrs.other`, the same way `FieldMesh` handles nonstandard attributes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ec1825c", + "metadata": {}, + "outputs": [], + "source": [ + "W.s_position = 12.5\n", + "W.write_openpmd(\"wavefront_openpmd.h5\", beamline=\"example\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6568e8c", + "metadata": {}, + "outputs": [], + "source": [ + "import h5py\n", + "\n", + "# The field is stored slowest-varying first, so the dataset shape is (nz, ny, nx).\n", + "with h5py.File(\"wavefront_openpmd.h5\") as h5:\n", + " mesh = h5[\"data/1/meshes/electricField\"]\n", + " for key, value in mesh.attrs.items():\n", + " print(f\"{key:>20}: {value}\")\n", + " print(f\"{'Ex dataset shape':>20}: {mesh['x'].shape}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "049c4b86", + "metadata": {}, + "outputs": [], + "source": [ + "W3 = Wavefront.from_openpmd(\"wavefront_openpmd.h5\")\n", + "W3" + ] + }, + { + "cell_type": "markdown", + "id": "26ae4494", + "metadata": {}, + "source": [ + "The round trip is exact, and the extension attributes come back in `.attrs`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d519157b", + "metadata": {}, + "outputs": [], + "source": [ + "np.array_equal(W.Ex, W3.Ex), W3.attrs, W3.s_position" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "304f8fa9", + "metadata": {}, + "outputs": [], + "source": [ + "# Cleanup\n", + "os.remove(\"wavefront_openpmd.h5\")" + ] } ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "beamphysics-dev", "language": "python", "name": "python3" }, @@ -783,7 +866,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.9" + "version": "3.14.3" } }, "nbformat": 4, diff --git a/mkdocs.yml b/mkdocs.yml index b8f45f5f..61bbfcb1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -35,6 +35,7 @@ nav: - Particles: api/particles.md - Fields: api/fields.md - Wakefields: api/wakefields.md + - Wavefront: api/wavefront.md - Standards: - Statistics Standard: api/statistics_standard.md diff --git a/tests/test_wavefront.py b/tests/test_wavefront.py index f673b251..4823c9ce 100644 --- a/tests/test_wavefront.py +++ b/tests/test_wavefront.py @@ -1,5 +1,12 @@ +from dataclasses import replace + +import h5py import matplotlib.pyplot as plt import pytest +from scipy.constants import h as h_planck + +from beamphysics.units import c_light, dimension, e_charge +from beamphysics.wavefront.openpmd import WavefrontAttrs from beamphysics.wavefront.wavefront import Wavefront from beamphysics.wavefront.propagators import ( @@ -269,3 +276,887 @@ def test_from_gaussian_polarization(): Wy = Wavefront.from_gaussian(shape=(10, 10, 10), sigma0=1e-6, polarization="y") assert Wy.Ex is None assert Wy.Ey is not None + + +# ----------------------------------------------------------------------------- +# openPMD EXT_Wavefront I/O +# ----------------------------------------------------------------------------- + + +def make_small(shape=(9, 11, 7), dtype=None, polarization="x"): + """ + Small non-cubic Wavefront for I/O tests. + + Parameters + ---------- + shape : tuple of int, default=(9, 11, 7) + Grid shape (nx, ny, nz). Deliberately non-cubic and with all three axes + distinct, so that an axis-order mistake cannot pass silently. + dtype : numpy dtype, optional + If given, the field arrays are cast to this dtype. + polarization : {'x', 'y', 'xy'}, default='x' + Which components to populate. + + Returns + ------- + Wavefront + """ + kwargs = { + "shape": shape, + "dx": 1e-6, + "dy": 2e-6, + "dz": 3e-6, + "wavelength": 1.5e-9, + "sigma0": 5e-6, + "energy": 1.0, + } + if polarization == "xy": + Wx = Wavefront.from_gaussian(polarization="x", **kwargs) + Wy = Wavefront.from_gaussian(polarization="y", **kwargs) + W = replace(Wx, Ey=2.0 * Wy.Ey) + else: + W = Wavefront.from_gaussian(polarization=polarization, **kwargs) + + if dtype is not None: + W = replace( + W, + Ex=None if W.Ex is None else W.Ex.astype(dtype), + Ey=None if W.Ey is None else W.Ey.astype(dtype), + ) + return W + + +@pytest.mark.parametrize("polarization", ["x", "y", "xy"]) +def test_openpmd_round_trip(tmp_path, polarization): + """Arrays and grid come back exactly, for each polarization case.""" + W = make_small(polarization=polarization) + path = tmp_path / "wavefront.h5" + W.write_openpmd(path) + W2 = Wavefront.from_openpmd(path) + + for original, restored in ((W.Ex, W2.Ex), (W.Ey, W2.Ey)): + if original is None: + assert restored is None + else: + assert np.array_equal(original, restored) + + assert W2.shape == W.shape + assert (W2.dx, W2.dy, W2.dz) == (W.dx, W.dy, W.dz) + # Wavelength goes out as a photon energy and comes back through the same + # SI-exact constants, so it is bit-identical. + assert W2.wavelength == W.wavelength + + +def test_openpmd_round_trip_preserves_dtype(tmp_path): + """complex64 stays complex64; it is not silently widened.""" + W = make_small(dtype=np.complex64) + assert W.Ex.dtype == np.complex64 + + path = tmp_path / "wavefront.h5" + W.write_openpmd(path) + + with h5py.File(path) as h5: + dataset = h5["data/1/meshes/electricField/x"] + # h5py maps the {r, i} compound back to a numpy complex dtype on read, so + # inspect the HDF5 type itself to see what is actually on disk. + hdf5_type = dataset.id.get_type() + assert hdf5_type.get_nmembers() == 2 + assert [hdf5_type.get_member_name(i).decode() for i in range(2)] == [ + "r", + "i", + ], "real part first, then imaginary, per FORMAT_HDF5" + assert ( + hdf5_type.get_size() == 8 + ), "complex64 must stay 32+32 bits, not be widened to 64+64" + + W2 = Wavefront.from_openpmd(path) + assert W2.Ex.dtype == np.complex64 + assert np.array_equal(W.Ex, W2.Ex) + + +def test_openpmd_layout(tmp_path): + """ + The written layout is a contract with other codes; assert it directly. + + This test is deliberately about bytes on disk, not about round tripping. + """ + W = replace(make_small(polarization="xy"), s_position=12.5) + path = tmp_path / "wavefront.h5" + W.write_openpmd(path) + + nx, ny, nz = W.shape + + with h5py.File(path) as h5: + assert h5.attrs["openPMD"].decode() == "2.0.0" + assert h5.attrs["openPMDextension"].decode() == "Wavefront" + assert h5.attrs["basePath"].decode() == "/data/%T/" + assert h5.attrs["meshesPath"].decode() == "meshes/" + assert h5.attrs["iterationEncoding"].decode() == "groupBased" + + iteration = h5["data/1"] + assert iteration.attrs["timeUnitSI"] == 1.0 + + mesh = iteration["meshes/electricField"] + + # Stored slowest-varying first: (z, y, x). + assert [label.decode() for label in mesh.attrs["axisLabels"]] == ["z", "y", "x"] + assert mesh["x"].shape == (nz, ny, nx) + assert mesh["y"].shape == (nz, ny, nx) + + # gridSpacing follows axisLabels, so it is (dz, dy, dx). + assert np.array_equal(mesh.attrs["gridSpacing"], [W.dz, W.dy, W.dx]) + + # gridGlobalOffset is the first cell of each axis, in the same order. + # Components declare position = 0, so that is the first sample. + assert np.allclose( + mesh.attrs["gridGlobalOffset"], [W.zmin, W.ymin, W.xmin], rtol=1e-15 + ) + + # openPMD 2.0 makes gridUnitSI one value per axis. + assert mesh.attrs["gridUnitSI"].shape == (3,) + assert np.array_equal(mesh.attrs["gridUnitSI"], np.ones(3)) + assert mesh.attrs["gridUnitDimension"].shape == (21,) + + # electricField is V/m, the package's own "electric_field" dimension. + assert np.array_equal(mesh.attrs["unitDimension"], dimension("electric_field")) + assert np.array_equal(mesh.attrs["gridUnitDimension"], dimension("length") * 3) + + assert mesh.attrs["geometry"].decode() == "cartesian" + assert mesh.attrs["temporalDomain"].decode() == "time" + assert mesh.attrs["spatialDomain"].decode() == "r" + assert mesh.attrs["zCoordinate"] == 12.5 + + # photonEnergy is in joules, not eV. + expected = h_planck * c_light / W.wavelength + assert mesh.attrs["photonEnergy"] == expected + assert not np.isclose(mesh.attrs["photonEnergy"], expected / e_charge) + + for component in ("x", "y"): + assert mesh[component].attrs["unitSI"] == 1.0 + assert np.array_equal(mesh[component].attrs["position"], np.zeros(3)) + + # The class has no longitudinal field, so no z component is written. + assert "z" not in mesh + + +def test_openpmd_transposes_rather_than_reshapes(tmp_path): + """ + The stored array is a real transpose, not a reshape of the same buffer. + + A reshape would produce the right dataset shape while scrambling the data, so + check an individual element against its transposed index. + """ + W = make_small() + path = tmp_path / "wavefront.h5" + W.write_openpmd(path) + + with h5py.File(path) as h5: + stored = h5["data/1/meshes/electricField/x"][()] + + assert np.array_equal(stored, W.Ex.transpose(2, 1, 0)) + assert stored[3, 2, 1] == W.Ex[1, 2, 3] + + +def _write_foreign_file(path, W, labels=("z", "y", "x")): + """ + Write a file by hand the way a foreign code plausibly would. + + Scalars are stored as shape-(1,) arrays and strings as bytes, which is what a + Fortran HDF5 writer tends to emit. The axis order is configurable so that the + reader's handling of an arbitrary `axisLabels` permutation can be exercised + without the writer being involved. + + Parameters + ---------- + path : pathlib.Path + File to create. + W : Wavefront + Source of the field data and grid. + labels : tuple of str, default=('z', 'y', 'x') + Axis order to store. + """ + spacing = {"x": W.dx, "y": W.dy, "z": W.dz} + first = {"x": W.xmin, "y": W.ymin, "z": W.zmin} + permutation = tuple(("x", "y", "z").index(label) for label in labels) + + with h5py.File(path, "w") as h5: + h5.attrs["openPMD"] = np.bytes_("2.0.0") + h5.attrs["openPMDextension"] = np.bytes_("Wavefront") + h5.attrs["basePath"] = np.bytes_("/data/%T/") + h5.attrs["meshesPath"] = np.bytes_("meshes/") + + mesh = h5.create_group("data/7/meshes/electricField") + mesh.attrs["geometry"] = np.bytes_("cartesian") + mesh.attrs["axisLabels"] = np.array([np.bytes_(label) for label in labels]) + mesh.attrs["gridSpacing"] = np.array([spacing[label] for label in labels]) + mesh.attrs["gridGlobalOffset"] = np.array([first[label] for label in labels]) + mesh.attrs["gridUnitSI"] = np.ones(3) + + # Scalars as length-1 arrays, and photonEnergy up at the series root. + h5.attrs["photonEnergy"] = np.array([h_planck * c_light / W.wavelength]) + mesh.attrs["temporalDomain"] = np.bytes_("time") + mesh.attrs["spatialDomain"] = np.bytes_("r") + mesh.attrs["zCoordinate"] = np.array([3.25]) + mesh.attrs["beamline"] = np.bytes_("undulator") + + mesh["x"] = W.Ex.transpose(permutation) + mesh["x"].attrs["unitSI"] = np.array([1.0]) + + +def test_openpmd_read_foreign_file(tmp_path): + """Bytes strings, length-1 array scalars and a non-default iteration all read.""" + W = make_small() + path = tmp_path / "foreign.h5" + _write_foreign_file(path, W) + + W2 = Wavefront.from_openpmd(path) + + assert np.array_equal(W2.Ex, W.Ex) + assert W2.Ey is None + assert (W2.dx, W2.dy, W2.dz) == (W.dx, W.dy, W.dz) + assert np.isclose(W2.wavelength, W.wavelength) + + # Decoded to str and float, not left as bytes and ndarray. + assert W2.attrs.beamline == "undulator" + assert W2.s_position == 3.25 + + +@pytest.mark.parametrize( + "labels", + [("z", "y", "x"), ("x", "y", "z"), ("y", "x", "z"), ("z", "x", "y")], +) +def test_openpmd_read_axis_permutations(tmp_path, labels): + """Any permutation of (x, y, z) in axisLabels is honored on read.""" + W = make_small() + path = tmp_path / "permuted.h5" + _write_foreign_file(path, W, labels=labels) + + W2 = Wavefront.from_openpmd(path) + + assert W2.shape == W.shape + assert (W2.dx, W2.dy, W2.dz) == (W.dx, W.dy, W.dz) + assert np.array_equal(W2.Ex, W.Ex) + + # gridGlobalOffset is permuted alongside gridSpacing, so a reader that + # permuted one but not the other would put the grid in the wrong place. + assert np.allclose([W2.xmin, W2.ymin, W2.zmin], [W.xmin, W.ymin, W.zmin]) + + +@pytest.mark.parametrize("grid_unit_si", [np.full(3, 1e-3), np.array([1e-3])]) +def test_openpmd_read_non_si_units(tmp_path, grid_unit_si): + """ + gridUnitSI and unitSI are applied, so a file in non-SI units reads correctly. + + A reader that ignored them would return a grid that is wrong by a factor of a + thousand while looking entirely plausible. gridUnitSI is written both as the + openPMD 2.0 per-axis array and as the 1.x scalar that older EXT_Wavefront + implementations still emit. + """ + W = make_small() + path = tmp_path / "millimeters.h5" + W.write_openpmd(path) + + # Restate the same grid in mm and the same field in kV/m. + with h5py.File(path, "r+") as h5: + mesh = h5["data/1/meshes/electricField"] + mesh.attrs["gridSpacing"] = mesh.attrs["gridSpacing"] * 1e3 + # gridUnitSI scales gridGlobalOffset too, not just gridSpacing. + mesh.attrs["gridGlobalOffset"] = mesh.attrs["gridGlobalOffset"] * 1e3 + mesh.attrs["gridUnitSI"] = grid_unit_si + mesh["x"][...] = mesh["x"][()] / 1e3 + mesh["x"].attrs["unitSI"] = 1e3 + + W2 = Wavefront.from_openpmd(path) + + assert np.allclose(W2.Ex, W.Ex, rtol=1e-15, atol=0.0) + assert np.allclose([W2.dx, W2.dy, W2.dz], [W.dx, W.dy, W.dz], rtol=1e-15) + assert np.allclose([W2.xmin, W2.ymin, W2.zmin], [W.xmin, W.ymin, W.zmin]) + + +def _spike(W, index=(6, 3, 2)): + """A wavefront whose only nonzero sample sits at `index`.""" + Ex = np.zeros(W.shape, dtype=complex) + Ex[index] = 1.0 + return replace(W, Ex=Ex) + + +def _feature_x(W): + """x coordinate of the largest |Ex| sample.""" + index = np.unravel_index(np.argmax(np.abs(W.Ex)), W.Ex.shape) + return W.xvec[index[0]] + + +def test_offset_shifts_the_grid(): + """The grid midpoint moves the grid without touching its spacing or extent.""" + W = make_small() + S = replace(W, xmid=1e-5) + + assert S.dx == W.dx + assert S.xmin == pytest.approx(W.xmin + 1e-5) + assert S.xmax == pytest.approx(W.xmax + 1e-5) + assert np.allclose(S.xvec, W.xvec + 1e-5) + + # Translating the grid translates the mean and leaves the width alone. + assert S.mean_x == pytest.approx(W.mean_x + 1e-5) + assert S.sigma_x == pytest.approx(W.sigma_x) + + +def test_crop_and_pad_preserve_physical_coordinates(): + """ + An asymmetric crop or pad must not translate the physics. + + Before the grid carried an origin, cropping two samples off the front of the + x axis silently moved every remaining sample by one dx, because the grid was + re-centered on whatever was left. + """ + W = _spike(make_small()) + x0 = _feature_x(W) + + assert _feature_x(W.crop(nx=(2, 0))) == pytest.approx(x0) + assert _feature_x(W.crop(nx=(0, 2))) == pytest.approx(x0) + assert _feature_x(W.pad(nx=(3, 0))) == pytest.approx(x0) + assert _feature_x(W.pad(nx=(0, 3))) == pytest.approx(x0) + + # A crop and an equal pad put the grid back exactly where it started. + restored = W.crop(nx=(1, 3)).pad(nx=(1, 3)) + assert restored.xmid == pytest.approx(W.xmid) + assert np.allclose(restored.xvec, W.xvec) + + +def test_kspace_carries_grid_midpoint_through_unchanged(): + """ + `WavefrontK` holds the real-space grid midpoint as inert state. + + A real-space shift is a linear phase in k-space, not a shift of the k grid, so + the transform must neither apply it nor lose it. Cropping in k-space discards + k samples, which does not move the real-space grid either. + """ + W = replace(make_small(), xmid=1e-5, ymid=-2e-5, zmid=3e-5) + mids = (W.xmid, W.ymid, W.zmid) + + K = W.to_kspace() + assert (K.xmid, K.ymid, K.zmid) == mids + assert (K.crop(nx=(2, 0)).xmid, K.pad(ny=(1, 0)).ymid) == mids[:2] + + W2 = K.to_rspace() + assert (W2.xmid, W2.ymid, W2.zmid) == mids + assert np.allclose(W2.Ex, W.Ex) + + +def test_drift_preserves_grid_midpoint(): + """Drift is a shift-invariant convolution, so the grid rides along.""" + W = replace(make_small(), xmid=1e-5) + assert drift_wavefront(W, 1.0).xmid == W.xmid + + +def test_drift_advances_s_position(): + """ + Propagating records itself, as `ParticleGroup.drift` does by moving `z`. + + The mesh here is co-moving, so `zmid` cannot record the propagation and + `s_position` is the only place the information can go. + """ + W = replace(make_small(), s_position=3.0) + + W2 = drift_wavefront(W, 2.0) + assert W2.s_position == pytest.approx(5.0) + assert W2.zmid == W.zmid + + # It accumulates, so two steps land where one step of the total would. + assert drift_wavefront(W2, 4.0).s_position == pytest.approx(9.0) + + # A drift of zero is a no-op, and a backwards drift walks it back. + assert drift_wavefront(W, 0.0).s_position == pytest.approx(3.0) + assert drift_wavefront(W2, -2.0).s_position == pytest.approx(3.0) + + +def test_advanced_drift_advances_s_position_by_the_physical_distance(): + """ + The curved propagator advances `s_position` by `z`, not by its internal `z_eff`. + + It delegates to the basic propagator with a scaled distance `z / (1 + curv*z)`, + which is a change of variables rather than a distance anything travels. Letting + that increment stand would silently under-report how far the wavefront went. + """ + z, curvature = 1.0, 0.5 + W = replace(make_small(), s_position=3.0) + + W2 = drift_wavefront(W, z, curvature=curvature) + + assert W2.s_position == pytest.approx(3.0 + z) + # The scaled distance is a different number, so this is not a vacuous check. + assert z / (1 + curvature * z) != pytest.approx(z) + + +def test_transforms_and_resizes_leave_s_position_alone(): + """ + Only propagation moves the wavefront along the beamline. + + `crop`, `pad` and the domain transforms change the representation, not the + position, exactly as they would leave `ParticleGroup.z` alone. + """ + W = replace(make_small(), s_position=3.0) + + assert W.to_kspace().s_position == 3.0 + assert W.to_kspace().to_rspace().s_position == 3.0 + assert W.crop(nx=(1, 1)).s_position == 3.0 + assert W.pad(nz=(2, 2)).s_position == 3.0 + + +def test_derived_wavefronts_do_not_share_attrs(): + """ + Every derivation must deep-enough copy `attrs`. + + `attrs` is a mutable dataclass, and `dataclasses.replace` rebinds the very same + object onto the new instance. Without an explicit copy, editing metadata on a + propagated or cropped wavefront silently reaches back into the original. + """ + W = replace(make_small(), attrs={"beamline": "SXR"}) + + derived = [ + W.drift(1.0), + drift_wavefront(W.to_kspace(), 1.0), + drift_wavefront(W, 1.0, curvature=0.5), + W.crop(nx=(1, 1)), + W.pad(nx=(1, 1)), + W.to_kspace(), + W.to_kspace().to_rspace(), + ] + + for w in derived: + assert w.attrs is not W.attrs + assert w.attrs.beamline == "SXR" + w.attrs.beamline = "HXR" + + assert W.attrs.beamline == "SXR" + + +def test_advanced_drift_magnifies_the_grid_midpoint(): + """ + The curved propagator rescales the grid about the optical axis. + + It scales dx and dy by 1/M; both of its quadratic phases are referenced to + x = y = 0, so the map is x -> x/M and the midpoint has to scale by the same + factor. Leaving it fixed would expand the grid around an off-axis feature while + the feature kept its old coordinate. + """ + W = replace(make_small(), xmid=1e-5, ymid=-2e-5) + z, curvature = 1.0, 0.5 + + W2 = drift_wavefront_advanced(W, z, curvature=curvature) + + M = (z / (1 + curvature * z)) / z + assert W2.dx == pytest.approx(W.dx / M) + assert W2.xmid == pytest.approx(W.xmid / M) + assert W2.ymid == pytest.approx(W.ymid / M) + + # The whole grid magnifies as one, so the extent scales by the same factor. + assert W2.xmin == pytest.approx(W.xmin / M) + assert W2.xmax == pytest.approx(W.xmax / M) + + # The intra-pulse axis is untouched by a transverse rescaling. + assert W2.zmid == W.zmid + + +def test_openpmd_grid_midpoint_round_trip(tmp_path): + """A grid midpoint survives a write and read, which it previously did not.""" + W = replace(make_small(), xmid=1e-5, ymid=-2e-5, zmid=3e-5) + path = tmp_path / "offset.h5" + W.write_openpmd(path) + + W2 = Wavefront.from_openpmd(path) + + assert np.allclose([W2.xmid, W2.ymid, W2.zmid], [W.xmid, W.ymid, W.zmid]) + assert np.allclose(W2.xvec, W.xvec) + assert np.allclose(W2.yvec, W.yvec) + assert np.allclose(W2.zvec, W.zvec) + assert np.array_equal(W2.Ex, W.Ex) + + +def test_openpmd_read_without_grid_global_offset(tmp_path): + """ + A file with no gridGlobalOffset reads as a centered grid. + + Defaulting the attribute to zero and treating that as the first sample would + instead shift the whole grid so that it started at the origin. + """ + W = replace(make_small(), xmid=1e-5) + path = tmp_path / "no_offset.h5" + W.write_openpmd(path) + + with h5py.File(path, "r+") as h5: + del h5["data/1/meshes/electricField"].attrs["gridGlobalOffset"] + + W2 = Wavefront.from_openpmd(path) + + assert (W2.xmid, W2.ymid, W2.zmid) == (0.0, 0.0, 0.0) + assert W2.xmin == pytest.approx(-(W.nx - 1) * W.dx / 2) + + +def test_genesis4_write_refuses_offset_grid(tmp_path): + """ + Genesis4 stores a point count and a spacing, with nowhere to put an origin. + + Writing anyway would silently re-center the wavefront. + """ + W = replace(make_small(shape=(8, 8, 4)), dy=1e-6, xmid=1e-5) + + with pytest.raises(ValueError, match="grid origin"): + W.write_genesis4(tmp_path / "genesis.h5") + + +def test_genesis4_round_trips_s_position_as_refposition(tmp_path): + """ + Genesis4's `refposition` is the same quantity as `s_position`. + + It is the position of the dump along the undulator line, so it is written from + `s_position` and read back into it rather than being dropped at both ends. + """ + W = replace(make_small(shape=(8, 8, 4)), dy=1e-6, s_position=7.25) + path = tmp_path / "genesis.h5" + + W.write_genesis4(path) + + with h5py.File(path) as h5: + assert h5["refposition"][0] == pytest.approx(7.25) + + assert Wavefront.from_genesis4(path).s_position == pytest.approx(7.25) + + +def test_openpmd_read_constant_component(tmp_path): + """ + A uniform component stored as a `value`/`shape` group is read, not crashed on. + + openPMD allows this compression, and the package's other readers support it via + `readers.is_constant_component`. + """ + W = make_small() + path = tmp_path / "constant.h5" + W.write_openpmd(path) + + nx, ny, nz = W.shape + with h5py.File(path, "r+") as h5: + mesh = h5["data/1/meshes/electricField"] + del mesh["x"] + constant = mesh.create_group("x") + constant.attrs["value"] = 2.0 + 3.0j + constant.attrs["shape"] = np.array([nz, ny, nx]) + constant.attrs["unitSI"] = 1.0 + + W2 = Wavefront.from_openpmd(path) + + assert W2.shape == (nx, ny, nz) + assert np.all(W2.Ex == 2.0 + 3.0j) + + +def test_openpmd_attrs_round_trip(tmp_path): + """Extension attributes survive a round trip and follow domain transforms.""" + W = replace(make_small(), s_position=4.0) + path = tmp_path / "attrs.h5" + W.write_openpmd( + path, + beamline="SXR", + radius_of_curvature_x=2.5, + radius_of_curvature_y=3.5, + ) + W2 = Wavefront.from_openpmd(path) + + assert W2.s_position == 4.0 + assert W2.attrs.beamline == "SXR" + assert W2.attrs.radius_of_curvature_x == 2.5 + assert W2.attrs.radius_of_curvature_y == 3.5 + # Unset optional attributes stay absent rather than defaulting to zero. + assert W2.attrs.delta_radius_of_curvature_x is None + + # attrs are carried across the k-space transform, and are copies, not aliases. + Wk = W2.to_kspace() + assert Wk.attrs == W2.attrs + assert Wk.attrs is not W2.attrs + assert Wk.to_rspace().attrs == W2.attrs + + # `s_position` is a coordinate on the wavefront, so it survives a rewrite. + path2 = tmp_path / "attrs2.h5" + W2.write_openpmd(path2) + assert Wavefront.from_openpmd(path2).s_position == 4.0 + + +def test_openpmd_default_attrs(): + """A Wavefront built in memory has default attrs, not None.""" + assert make_small().attrs == WavefrontAttrs() + assert make_small().to_kspace().attrs == WavefrontAttrs() + assert make_small().attrs.beamline is None + + +def test_wavefront_attrs_accepts_mapping(): + """A mapping keyed by either openPMD or field names is coerced.""" + Ex = make_small().Ex + assert ( + Wavefront(Ex=Ex, attrs={"radiusOfCurvatureX": 2.0}).attrs.radius_of_curvature_x + == 2.0 + ) + assert ( + Wavefront( + Ex=Ex, attrs={"radius_of_curvature_x": 2.0} + ).attrs.radius_of_curvature_x + == 2.0 + ) + + +def test_wavefront_attrs_assignment_is_coerced(): + """ + Assigning a mapping coerces it, rather than leaving a raw dict on the instance. + + Coercing only at construction meant `w.attrs = {...}` left a dict behind, so + `w.attrs.beamline` raised `AttributeError` and the mapping went unvalidated + until write time. + """ + W = make_small() + W.attrs = {"radiusOfCurvatureX": 2.0} + + assert isinstance(W.attrs, WavefrontAttrs) + assert W.attrs.radius_of_curvature_x == 2.0 + + with pytest.raises(ValueError, match="radiusOfCurvatureX"): + W.attrs = {"radiusOfCurvatureX": 1.0, "radius_of_curvature_x": 2.0} + + +def test_wavefront_attrs_unknown_field_is_type_error(): + """A misspelled attribute fails at construction, not silently at write time.""" + with pytest.raises(TypeError, match="radius_of_curvature_z"): + WavefrontAttrs(radius_of_curvature_z=1.0) + + +def test_wavefront_attrs_other_rejects_computed(): + """`other` cannot override values the writer takes from the wavefront.""" + with pytest.raises(ValueError, match="photonEnergy"): + WavefrontAttrs(other={"photonEnergy": 1.0}) + with pytest.raises(ValueError, match="gridSpacing"): + WavefrontAttrs(other={"gridSpacing": [1.0, 2.0, 3.0]}) + + # `zCoordinate` is `Wavefront.s_position`, so it is not settable through attrs + # either: a value here would conflict with the one the propagators maintain. + with pytest.raises(ValueError, match="zCoordinate"): + WavefrontAttrs(other={"zCoordinate": 1.0}) + with pytest.raises(ValueError, match="zCoordinate"): + Wavefront(Ex=make_small().Ex, attrs={"zCoordinate": 1.0}) + + +def test_wavefront_attrs_rejects_conflicting_spellings(): + """ + Giving an attribute under both spellings is an error, not a silent preference. + + Preferring the openPMD name would drop the field-name value into `other`, where + the writer would emit it as a nonstandard attribute holding a conflicting value. + """ + with pytest.raises(ValueError, match="radiusOfCurvatureX"): + WavefrontAttrs.from_pmd( + {"radiusOfCurvatureX": 1.0, "radius_of_curvature_x": 2.0} + ) + + # A field whose two spellings are identical is not a conflict. + assert WavefrontAttrs.from_pmd({"beamline": "SXR"}).beamline == "SXR" + + +def test_wavefront_attrs_unwraps_nested_other(): + """ + An `other` key in a mapping is merged, not nested. + + Nesting it produced `other={'other': {...}}`, which the writer would then try to + store as a dict-valued HDF5 attribute. + """ + attrs = WavefrontAttrs.from_pmd( + {"beamline": "SXR", "other": {"someFutureAttr": 3.0}} + ) + + assert attrs.beamline == "SXR" + assert attrs.other == {"someFutureAttr": 3.0} + + # The nested form is still validated against the computed names. + with pytest.raises(ValueError, match="photonEnergy"): + WavefrontAttrs.from_pmd({"other": {"photonEnergy": 1.0}}) + + +def test_wavefront_attrs_is_importable_from_package(): + """`WavefrontAttrs` is part of the public API, so it must be reachable there.""" + import beamphysics + + assert beamphysics.WavefrontAttrs is WavefrontAttrs + assert "WavefrontAttrs" in beamphysics.__all__ + + +def test_wavefront_attrs_other_rejects_known_field(): + """`other` cannot shadow an attribute that has a field of its own.""" + with pytest.raises(ValueError, match="beamline"): + WavefrontAttrs(other={"beamline": "SXR"}) + + +def test_openpmd_unknown_attr_round_trips_via_other(tmp_path): + """ + A record attribute this module does not know is preserved across a round trip. + + EXT_Wavefront is still in flux, so a file written against a newer revision must + not silently lose attributes by passing through this class. + """ + W = make_small() + path = tmp_path / "future.h5" + W.write_openpmd(path) + + with h5py.File(path, "r+") as h5: + h5["data/1/meshes/electricField"].attrs["someFutureAttr"] = 7.0 + + W2 = Wavefront.from_openpmd(path) + assert W2.attrs.other == {"someFutureAttr": 7.0} + + path2 = tmp_path / "future2.h5" + W2.write_openpmd(path2) + with h5py.File(path2, "r") as h5: + assert h5["data/1/meshes/electricField"].attrs["someFutureAttr"] == 7.0 + + +def test_openpmd_iteration_selection(tmp_path): + """Multiple iterations must be disambiguated explicitly.""" + W1 = make_small() + W2 = replace(W1, Ex=2.0 * W1.Ex) + + path = tmp_path / "multi.h5" + with h5py.File(path, "w") as h5: + W1.write_openpmd(h5, iteration=1) + W2.write_openpmd(h5, iteration=2) + + with pytest.raises(ValueError, match="2 iterations"): + Wavefront.from_openpmd(path) + + assert np.array_equal(Wavefront.from_openpmd(path, iteration=2).Ex, W2.Ex) + assert np.array_equal(Wavefront.from_openpmd(path, iteration=1).Ex, W1.Ex) + + with pytest.raises(ValueError, match="iteration 3 not in the file"): + Wavefront.from_openpmd(path, iteration=3) + + +def test_openpmd_group_and_path_agree(tmp_path): + """Passing an open group is equivalent to passing a path.""" + W = make_small() + by_path = tmp_path / "by_path.h5" + by_group = tmp_path / "by_group.h5" + + W.write_openpmd(by_path) + with h5py.File(by_group, "w") as h5: + W.write_openpmd(h5) + + with h5py.File(by_group) as h5: + from_group = Wavefront.from_openpmd(h5) + + assert np.array_equal(from_group.Ex, Wavefront.from_openpmd(by_path).Ex) + + +def test_openpmd_bad_file_argument(): + """Only str, Path and h5py.Group are accepted.""" + W = make_small() + with pytest.raises(ValueError, match="h5py.Group"): + W.write_openpmd(42) + with pytest.raises(ValueError, match="h5py.Group"): + Wavefront.from_openpmd(42) + + +def test_openpmd_rejects_unknown_extension_attr(tmp_path): + """An unrecognized attribute is refused rather than silently written.""" + W = make_small() + with pytest.raises(TypeError, match="radius_of_curvature_z"): + W.write_openpmd(tmp_path / "bad.h5", radius_of_curvature_z=1.0) + + +def _corrupt(path, mutate): + """ + Apply `mutate` to the mesh record of an existing file. + + Parameters + ---------- + path : pathlib.Path + File to modify in place. + mutate : callable + Called with the `h5py.Group` of the mesh record. + """ + with h5py.File(path, "r+") as h5: + mutate(h5["data/1/meshes/electricField"]) + + +@pytest.mark.parametrize( + "mutate, match", + [ + ( + lambda mesh: mesh.attrs.__setitem__( + "temporalDomain", np.bytes_("frequency") + ), + "temporalDomain", + ), + ( + lambda mesh: mesh.attrs.__setitem__("spatialDomain", np.bytes_("k")), + "spatialDomain", + ), + ( + lambda mesh: mesh.attrs.__setitem__( + "axisLabels", np.array([np.bytes_(s) for s in ("z", "y", "r")]) + ), + "axisLabels", + ), + (lambda mesh: mesh.attrs.__delitem__("zCoordinate"), "zCoordinate"), + (lambda mesh: mesh.attrs.__delitem__("photonEnergy"), "photonEnergy"), + (lambda mesh: mesh.attrs.__delitem__("gridSpacing"), "gridSpacing"), + ], +) +def test_openpmd_refusals(tmp_path, mutate, match): + """ + Unsupported or incomplete files are refused with a message naming the problem. + + A silent wrong answer here would be worse than a crash: the numbers would look + plausible. + """ + path = tmp_path / "wavefront.h5" + make_small().write_openpmd(path) + _corrupt(path, mutate) + + with pytest.raises(ValueError, match=match): + Wavefront.from_openpmd(path) + + +def test_openpmd_refuses_non_wavefront_file(tmp_path): + """A file without the mesh record is named as such, not left to KeyError.""" + path = tmp_path / "empty.h5" + with h5py.File(path, "w") as h5: + h5.create_group("data/1/meshes") + + with pytest.raises(ValueError, match="not an EXT_Wavefront file"): + Wavefront.from_openpmd(path) + + +def test_openpmd_energy_agrees_with_genesis4(tmp_path): + """ + The two formats must not disagree about how much energy is in the pulse. + + Measured on this grid: the openPMD round trip is exact, and Genesis4 differs + from it by 4.4e-16 relative, i.e. two ulp of float64, from its own unit + conversions. The assertions below are set just above those measured levels + rather than at a round number. + """ + # Genesis4 requires a square transverse grid with equal spacing. + W = Wavefront.from_gaussian( + shape=(33, 33, 17), + dx=10e-6, + dy=10e-6, + dz=3e-6, + wavelength=1.5e-9, + sigma0=50e-6, + energy=1.0, + ) + + openpmd_path = tmp_path / "wavefront.h5" + genesis_path = tmp_path / "wavefront.genesis.h5" + W.write_openpmd(openpmd_path) + W.write_genesis4(genesis_path) + + from_openpmd = Wavefront.from_openpmd(openpmd_path) + from_genesis = Wavefront.from_genesis4(genesis_path) + + assert from_openpmd.energy == W.energy + assert abs(from_genesis.energy - W.energy) / W.energy < 1e-15 + assert abs(from_openpmd.energy - from_genesis.energy) / W.energy < 1e-15