Add new options to genesis4_par_to_data: wrap, z0, equal_weights, cut…

[ChristopherMayes]

tldr;

Summary

This branch overhauls how Genesis4 .par (raw particle) HDF5 files are read into openPMD-beamphysics. It adds a first-class ParticleGroup.from_genesis4 constructor, significantly expands the conversion options, makes the conversion reproducible via an explicit RNG, and supports downsampling huge (especially one4one) files at read time without ever loading the full beam. The example data is no longer committed to the repository; it is generated on the fly by a script that installs Genesis4 into an isolated conda environment, used by both the docs build and the test suite in CI.

What changed

New public API

• ParticleGroup.from_genesis4(...) classmethod (beamphysics/particles.py) — load a Genesis4 .par file directly into a ParticleGroup. This is the recommended entry point; it forwards to genesis4_par_to_data.

genesis4_par_to_data rework

beamphysics/interfaces/genesis.py

New keyword arguments:

Argument	Default	Purpose
wrap	False	Wrap the in-slice z position modulo the slice length (theta taken mod 2π).
z0	0.0	Absolute z offset applied to slice index 1.
slices	None	Select specific slice indices instead of all slices.
equal_weights	False	Resample particles per slice so all output particles carry equal charge weight. Only existing particles are used (never upsampled), so the count generally shrinks. Mutually exclusive with n_particle.
cutoff	0.0	Skip quiet-loaded slices whose per-particle weight falls below this charge; intended for filtering extremely low-charge macroparticles (e.g. cutoff=1.6e-19 drops sub-electron slices). Defaults to 0.0 (keep everything). Zero-current / non-finite slices are always dropped. Ignored for one4one beams.
n_particle	None	Thin the beam on read to exactly this many particles, retaining the per-slice weights (rescaled once so total charge is conserved exactly). Only the selected rows are read from the file, so huge files are downsampled immediately. Values >= the file's count are a no-op. Mutually exclusive with equal_weights.
rng	None	Explicit numpy.random.Generator or integer seed for smearing and resampling (reproducible).

Behavioral and implementation changes:

• smear default changed True → False (BREAKING for genesis4_par_to_data). Particles are now loaded at their literal slice positions by default, which produces a comb-like longitudinal distribution when sample > 1. Smearing is opt-in (smear=True).
  • Rationale: smearing injects RNG-based randomness, so an off-by-default keeps the load deterministic/reproducible and returns exactly what Genesis4 wrote.
  • Migration: pass smear=True explicitly to restore the previous behavior.
• smear and wrap are now independent (intentional design change). The old code always applied theta/(2π) % 1 before smearing; wrapping is now a separate, explicit user directive. Callers who want the old always-wrapped smear should pass smear=True, wrap=True. Without wrap, theta values outside [0, 2π) (e.g. from slippage in post-undulator dumps) place particles outside their nominal slice, faithfully reflecting the file contents.
• Absolute z reconstruction. z is computed directly from the slice index as (ix - 1) * slicespacing + z0 rather than from a running i0 counter. This is robust to skipped/zero-current/missing slices (verified equivalent to the old counter for well-formed files).
• Metadata pre-pass. Slice filtering (empty, zero-current/nan, below cutoff) and per-slice weights are now computed up front from HDF5 metadata only (dataset shapes and the 1-element current datasets), before any bulk reads. Consequences:
  • n_particle apportionment runs over the surviving slices, so filtered slices cannot silently consume part of the requested count — the output hits the target exactly, even in combination with cutoff.
  • With n_particle, only the selected rows are read from the file (load_parfile_slice_data(group, sel=...) with sorted indices), so bulk I/O scales with n_particle, not with the file size, and the full beam is never materialized.
• n_particle and equal_weights are mutually exclusive and raise a ValueError when combined. They serve different needs: n_particle thins the beam while retaining the relative slice weights; equal_weights resamples the existing particles to a single common weight (never upsampling).
• Reproducible randomness. Smearing and resampling use np.random.default_rng(rng); subsample selections use shuffle=False and sorted indices (statistically equivalent, faster, and required for direct HDF5 row reads).
• one4one support. one4one beams (each macroparticle is a single electron) have variable per-slice particle counts and can contain empty slices; both are handled. Weights are assigned exactly qe (matching Genesis4's own &sort, which resets each slice current to npart·qe), so sorted and unsorted files reconstruct identically. The cutoff is disabled for one4one beams.
• equal_weights short-circuit. When every slice already has the same per-particle weight (e.g. an unsubsampled one4one beam), the per-slice resampling is skipped.
• sample = round(slicespacing / slicelength) (was int(...)), avoiding float truncation.
• Missing-file paths raise h5py's FileNotFoundError directly (no racy os.path.exists pre-check).
• Modern type hints (PEP 604 unions, e.g. rng: Optional[int | np.random.Generator]).

New helpers and a typed slice container

• genesis4_parfile_scalars(h5) — read scalar metadata (slicelength, slicespacing, slicecount, …).
• genesis4_parfile_slice_groups(h5) — slice group names, matched positively by name pattern (slice + digits) and structure and sorted by slice index, so new metadata entries in future Genesis4 versions are ignored automatically.
• genesis4_parfile_slice_counts(h5) — per-slice particle counts from dataset shapes only (no bulk data); the building block for both the count helper and the n_particle pre-pass.
• genesis4_parfile_n_particle(h5) — total particle count (sum of the slice counts); cheap for huge files and handy for choosing an n_particle target.
• load_parfile_slice_data(group, sel=None) → ParfileSliceData namedtuple (x, px, y, py, gamma, theta, current); sel reads only the given (sorted) rows.
• All helpers accept str, pathlib.Path, or an open h5py.File handle.

Example data is generated, not committed

• The .par.h5 / .out.h5 files are no longer stored in the repository (they are gitignored). They are generated from the committed inputs docs/examples/data/genesis4/{genesis4.in, one4one.in, genesis4.lat} by run.sh (genesis4 genesis4.in and genesis4 one4one.in, with set -e).
• scripts/generate_genesis4_example_data.bash (new) creates an isolated conda environment (genesis4-example-data) containing only the conda-forge genesis4 package — Genesis4 is not a dependency of environment.yml — and runs run.sh with it. The environment is reused if it already exists.
• CI wiring: gh-pages.yml runs the script before executing the notebooks; tests.yaml (conda) runs it before pytest. The pip test workflow has no Genesis4, so the Genesis4 particle tests skip cleanly there.
• The inputs are seeded (seed = 123456), so regeneration is reproducible for a given Genesis4 version. The quiet-loading example is a 100 pC Gaussian with sample = 10; the one4one example is a low-charge Gaussian (934 electrons) so that tail slices are empty and per-slice counts vary.

Tests and docs

• tests/test_genesis4_particles.py (new, 23 tests): scalars/slice-group helpers, default charge (100 pC within 1e-3), smear reproducibility, z0/wrap/slices behavior, cutoff filtering, one4one loading (uniform qe weights, no division by zero, cutoff not applied), n_particle subsampling (exact counts, exact charge conservation, uniform one4one weights, no-op above the total, exact count in combination with cutoff), equal_weights (uniform weights, charge preserved, never upsamples), and the n_particle/equal_weights mutual-exclusion error. A session fixture generates the data via run.sh when missing and skips if the genesis4 executable is unavailable. Test parameters carry type annotations.
• docs/examples/read_examples.ipynb: the Genesis4 section documents (in full sentences) the default comb-like load, smear with rng reproducibility, equal_weights (never upsampling), n_particle thinning in the quiet-loading mode (9216 → 1000, plotted), the one4one mode (uniform qe weights, cutoff inapplicability), and subsampling on read for huge one4one files (count helper, 934 → 200 with conserved charge, and the ≥-total no-op). The notebook does not generate the data itself; the markdown points readers at the generation script.

Verification

• Full test suite passes in beamphysics-dev (1842 passed, 1 xpassed), including the 23 Genesis4 particle tests.
• read_examples.ipynb executes end-to-end with no error cells; outputs show 99.74 pC recovered against the 100 pC input target, exact n_particle counts, and exact charge conservation on subsampling.
• The isolated-environment generation script is validated for syntax and wired into both workflows, but its first end-to-end run happens in CI (local environment creation was intentionally not exercised).

Genesis4 source verification

Cross-checked against the Genesis-1.3-Version4 source to confirm the particle
interpretation is correct:

• Slice metadata (src/IO/writeBeamHDF5.cpp, writeGlobal): the HDF5 dataset names are swapped relative to the internal C++ variable names, and the interface reads them by the HDF5 name (correctly):
  • HDF5 slicelength = beam->reflength = λ₀ (reference wavelength / bucket length) → used as ds_slice.
  • HDF5 slicespacing = beam->slicelength = sample·λ₀ (distance between written slices) → used as s_spacing.
  • Hence sample = round(slicespacing / slicelength) is exact. ✓
• px / py (include/Particle.h, src/Loading/QuietLoading.cpp L195–197, src/Loading/ImportTransformation.cpp L128) store γ·x′ ≈ γ·βₓ (the file's "rad" unit label refers to px/γ = x′). So p_x[eV/c] = px · mₑc² as the interface does. ✓
• theta is the ponderomotive phase in radians; one bucket = 2π, so the in-slice longitudinal position z = (theta / 2π)·λ₀ is correct. ✓
• Slice groups are 1-based (slice%6.6d, index i+1), so absolute z = (ix − 1)·slicespacing + z0 places slice 1 at z0. ✓
• Charge: each written slice represents sample buckets, i.e. a length slicespacing, so Q_slice = current · slicespacing / c and the per-particle weight Q_slice / n1 is correct. Confirmed empirically: the example recovers 99.74 pC against a 100 pC input target.

Conclusion: the particle model and unit conversions match the Genesis4 implementation.

Downstream usage (lume-genesis)

The only known external consumer of genesis4_par_to_data is lume-genesis, via LoadableParticleGroupH5File.load() → load_particle_group(h5, smear=True) → genesis4_par_to_data(h5, smear=smear).

• The smear default change is safe for lume-genesis. Its load_particle_group defines smear: bool = True and always passes smear=smear explicitly, so flipping the beamphysics default does not change which value it passes.
• The smear/wrap separation does affect lume-genesis for files whose theta extends outside [0, 2π): the old code always wrapped before smearing, while the new code wraps only with wrap=True. If the previous behavior is desired downstream, lume-genesis should pass wrap=True alongside smear=True.
• Follow-up (lume-genesis repo): the new options (wrap, z0, slices, equal_weights, cutoff, n_particle, rng) are not reachable through lume-genesis yet; load_particle_group only forwards smear. They would need to be plumbed through (or the call switched to ParticleGroup.from_genesis4) to be usable downstream.

Caveats and follow-ups

• Breaking default: smear=True → False changes output for existing callers of genesis4_par_to_data. Worth calling out in release notes.
• Breaking wrap behavior for smear=True callers: wrapping is no longer implied by smearing (intentional; see above). Pass wrap=True to reproduce the old always-wrapped positions.
• RNG streams changed relative to earlier iterations of this branch (shuffle=False, sorted selections): seeded results differ from earlier commits, but same-seed calls remain reproducible.
• equal_weights in the reader vs ParticleGroup.resample(equal_weights=True): the reader resamples per slice, without replacement, with a seeded RNG (preserving the slice charge profile exactly); statistics.resample_particles samples globally, with replacement, using global NumPy random state. A possible future cleanup is to lift the per-slice algorithm into statistics.resample_particles and delegate.
• First CI run of the data-generation script should be watched: the isolated-environment path (mamba/conda create + conda run) was validated for syntax but first exercised end-to-end in CI.
• Removing the committed binaries: after merging the generation machinery, drop the files with
  git rm --cached docs/examples/data/genesis4/{end.par.h5,genesis4.out.h5,one4one.par.h5,one4one.out.h5}.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 9 changed files in this pull request and generated 11 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 13 changed files in this pull request and generated 8 comments.

[ken-lauer]

So long as your .h5 files remain in git history, they still take up space in the repository.

In any case:

I'd rather we version control the data elsewhere (or make it available as a tarball to download) than generate a conda env + test data on the fly during the test suite. That feels rather fragile.

Or just go back to version-controlling it here - it's not that big of a deal overall.

Regardless of what you do here, rebase to clean up the history and remove any unnecessary h5 objects there.

[ChristopherMayes]

I made the files smaller with sample = 10 -> sample = 100, committed.

[ken-lauer]

You missed this important bit from my comment:

First add from __future__ import annotations to the top of this module, and then

That will fix the ParticleGroup name error while importing.