Wavefront openPMD I/O, grid origins, and beamline position

[ChristopherMayes]

Summary

Adds openPMD EXT_Wavefront read/write support for Wavefront, gives the class a grid origin and a beamline position, and fixes several correctness bugs found along the way. beamphysics/wavefront/openpmd.py is new; wavefront.py, propagators.py, and interfaces/genesis.py are modified.

openPMD support

Wavefront.from_openpmd and Wavefront.write_openpmd read and write openPMD 2.0 files using the EXT_Wavefront extension. Both transverse polarizations are written into one electricField record as components x and y; the z component is never written. Datasets are stored in (z, y, x) order so that each transverse slice is one contiguous block, and are written a slice at a time so no transposed copy of the field is materialized. Any axisLabels permutation of (x, y, z) is honored on read.

Extension attributes are carried by Wavefront.attrs, a typed WavefrontAttrs dataclass rather than a free-form dict. 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. Attributes the extension does not define yet round-trip verbatim through attrs.other, following what readers.load_field_attrs and writers.write_pmd_field already do for FieldMesh. This matters more here than there: EXT_Wavefront is still a draft, so a file written against a newer revision would otherwise lose data, silently, just by passing through this class.

Frequency-domain and k-space files are rejected on read rather than silently misinterpreted, as Wavefront cannot represent them.

Grid origin

Wavefront and WavefrontK gain xmid/ymid/zmid, the position of the grid midpoint, defaulting to 0.0 for a grid centered on the origin. xmin, xmax, and xvec derive from them.

This fixes a real bug: crop and pad previously discarded the coordinates of the surviving samples, so an asymmetric operation silently translated the physics. Both now shift the midpoint to compensate, through a _resized_mids hook that is a no-op on WavefrontK because discarding k samples does not move the real-space grid.

The name completes the xmin/xmid/xmax family and matches the underscore-free convention already used by dx, nx, and xvec. It is deliberately not x0, which already means beam center on from_gaussian, and not x_offset, which reads as openPMD's gridGlobalOffset — a different quantity, being the first sample rather than the midpoint, half a grid cell away.

Beamline position

s_position is a first-class field on Wavefront and WavefrontK: the position of this plane along the beamline, in whatever frame the caller picked. On disk it is the extension's required zCoordinate and Genesis4's refposition, both of which now round-trip. Genesis4 previously wrote a hardcoded 0.

drift advances it by the propagation distance. The mesh's own z axis is the intra-pulse coordinate and is co-moving — the paraxial kernel is purely transverse, applied slice by slice with no piston term — so s_position is the only place a record of the propagation can go. This follows ParticleGroup, where drift advances the real t/z coordinates rather than a separate register. Operations that change the representation rather than the position (crop, pad, to_kspace, an applied lens phase) leave it alone, exactly as they leave ParticleGroup.z alone.

The name avoids zref and similar, which would read as a member of the mesh family (zmin/zmid/zmax/zvec) while belonging to a different axis and a different frame. Bare s is also avoided because the two communities invert the convention: in Genesis4 beam files s is the in-bunch coordinate and z/refposition is the undulator position, and this repo already encodes the FEL side by mapping Genesis4's zpos to a column named s.

Bug fixes

• crop and pad translated the field on asymmetric operations (above).
• drift_wavefront_advanced scaled dx/dy by the magnification but not xmid/ymid, so an off-axis wavefront was magnified about the wrong point. Both of its quadratic phases are referenced to x = y = 0.
• crop, pad, and both drift_wavefront_basic returns shared the mutable attrs object with the original, because dataclasses.replace rebinds rather than copies. Editing metadata on a propagated wavefront reached back into the source. to_kspace and to_rspace already copied, so the behavior was inconsistent.
• auto_crop(apply=False) returned np.int64 despite an int annotation, so the returned dict did not serialize cleanly.
• Five auto_crop doctests referenced an undefined name and could never have run. Replaced with a working example.
• WavefrontAttrs.from_pmd wrote both spellings when handed an alias collision such as {"zCoordinate": 1, "s_position": 2}, and nested an other key into other={'other': {...}} instead of merging it. Both now behave correctly, the former by raising.
• w.attrs = {...} after construction left a raw dict in place, so w.attrs.beamline raised AttributeError. Coercion moved from __post_init__ to WavefrontBase.__setattr__.
• The Jupyter _repr_html_ reported none of the new state. It now shows the grid midpoint, s_position, and each set attribute, all omitted when they carry no information, and shares one row builder with _repr_pretty_ so the two cannot drift. Values are HTML-escaped, since attrs.other can carry arbitrary strings from a file straight into the notebook DOM. The k-space table now reports dkx/dky/dkz in rad/m rather than the derived real-space spacing.
• _repr_pretty_ never ran. IPython walks the MRO and takes the first class whose __dict__ holds _repr_pretty_ or __repr__, and the dataclass-generated __repr__ on Wavefront shadowed the base class method, so repr() of a wavefront printed the entire field array. The two dataclasses now set repr=False and the base supplies a short __repr__ giving shape, wavelength and s_position.

Changes visible to other codes

These change bytes on disk and are stated explicitly so downstream readers can move with them. Record and attribute names are unchanged.

• gridUnitSI is now a 3-element array rather than a scalar, and gridUnitDimension is written.
• gridGlobalOffset is written from xmin/ymin/zmin instead of zeros, and is applied on read. A file missing it reads as a centered grid, not as xmin = 0.
• The reader now applies gridUnitSI and unitSI rather than assuming SI.
• Genesis4 refposition carries s_position instead of always 0.
• zCoordinate accumulates across drifts rather than being whatever the caller passed at write time.

Breaking API changes

• Wavefront.write_openpmd(..., s_position=...), write_wavefront_openpmd(..., s_position=...), and wavefront_write_genesis4(..., refposition=...) no longer accept those keywords. Both writers take the value from the wavefront. A write-time override let the file disagree with the object it came from, which defeats the point of making the position a coordinate. Set w.s_position or replace(w, s_position=...) before writing. Of these, only refposition predates this branch; nothing in the repo passed it.
• No deprecation shim is provided for the attrs renames, as that layer is new in this branch.

Testing

1898 passing, up from a much smaller wavefront suite. tests/test_wavefront.py gains coverage for the openPMD layout as a contract (asserting bytes on disk, not just round trips), foreign files written in either axis order, corrupted and incomplete files, unknown attributes surviving a round trip, iteration disambiguation, the grid origin under asymmetric crop and pad, s_position under propagation versus representation changes, Genesis4 refposition in both directions, and metadata independence across every derivation path.

Known limitations, not addressed here

• WavefrontAttrs.copy() is one level deep, so a nested mutable inside other is still shared between a wavefront and its derivatives. Deepcopying would cost on every drift.
• For even n, xvec's origin at index (n-1)/2 disagrees with the ifftshift origin at index n//2 by dx/2. Odd n is exact. Changing it would move existing plots.
• drift_wavefront_basic loops over field[:, :, iz] out of a C-order (nx, ny, nz) array, which is the maximally strided slice. A (nz, ny, nx) internal order would make slices contiguous but breaks the public shape contract and every plotting and statistics path.
• Genesis4 has no place to store a grid origin, so write_genesis4 raises on a nonzero xmid. An asymmetric crop can produce one, making crop-then-write-Genesis4 an error by design.

Acknowledgment

Developed with Claude (Anthropic) via GitHub Copilot in VS Code. All design decisions, physics judgments, and the final review were made by the author; the model drafted code, tests, and documentation under direction.