diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..1bccc1f --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.h5 filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/test-bmad.yml b/.github/workflows/test-bmad.yml index 000a554..6f8654f 100644 --- a/.github/workflows/test-bmad.yml +++ b/.github/workflows/test-bmad.yml @@ -26,6 +26,7 @@ jobs: with: submodules: recursive fetch-depth: 1 + lfs: true - name: Checkout lattice repos uses: ./.github/actions/checkout-lattices diff --git a/.github/workflows/test-core.yml b/.github/workflows/test-core.yml index e1d134e..a32b028 100644 --- a/.github/workflows/test-core.yml +++ b/.github/workflows/test-core.yml @@ -19,6 +19,7 @@ jobs: with: submodules: recursive fetch-depth: 1 + lfs: true - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2e0e10c..217f5eb 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -27,6 +27,7 @@ jobs: - uses: actions/checkout@4f1f4aec02e41874fa0262ea8ff5172d7978ad1e with: fetch-depth: 1 + lfs: true - name: Checkout lattice repos uses: ./.github/actions/checkout-lattices - name: Set up conda environment diff --git a/.gitignore b/.gitignore index 0ee86f8..1834b7e 100644 --- a/.gitignore +++ b/.gitignore @@ -213,6 +213,7 @@ __marimo__/ .vscode/ *.h5 +!virtual_accelerator/beams/**/*.h5 *.png *.svg diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3bef9dc..5a3392d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,3 +24,12 @@ repos: types_or: [python, pyi, jupyter] - id: ruff-format types_or: [python, pyi, jupyter] + + - repo: local + hooks: + - id: beam-h5-has-sidecar + name: Every beams/*.h5 has a matching .h5.meta.json sidecar + entry: python scripts/check_beam_sidecars.py + language: system + files: ^virtual_accelerator/beams/.*\.h5$ + pass_filenames: true diff --git a/README.md b/README.md index 78dbf53..34adbaa 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,7 @@ conda install -c conda-forge impact-t=*=mpi_mpich* And the examples require installing ipykernel and register as a Jupyter kernel. + Optional Dependency Keys by Model: | Model / Factory Function | Optional dependency key(s) | Notes | | --- | --- | --- | @@ -120,6 +121,143 @@ Creating model instances also requires the `$LCLS_LATTICE` environment variable contents of the lcls-lattice repo https://github.com/slaclab/lcls-lattice or the facet2-lattice repo https://github.com/slaclab/facet2-lattice. +## Cached beam distributions + +Reference beam distributions at handoff planes (e.g. FACET PR10241, L0AFEND) are stored +under `virtual_accelerator/beams/` via **Git LFS**. Before cloning or pulling this repo, +install Git LFS once per machine: + +``` +brew install git-lfs # or: conda install -c conda-forge git-lfs +git lfs install +``` + +If you have already cloned without LFS, run `git lfs pull` to fetch the beam blobs. + +### Layout + +Beams are grouped by scenario (one subdirectory per date-tagged run) and named by +handoff element and particle count. Each `.h5` ships with a `.meta.json` sidecar +of the same base name (e.g. `PR10241_100000.h5` ↔ `PR10241_100000.meta.json`): + +``` +virtual_accelerator/beams/ + 2024-10-22_facet2_oneBunch/ + L0AFEND_100000.h5 # + L0AFEND_100000.meta.json + PR10241_100000.h5 # + PR10241_100000.meta.json +``` + +### Loading beams from Python + +Use the top-level `virtual_accelerator.beams` API instead of hard-coding paths. +The registry scans `virtual_accelerator/beams/` once on first use and reads +every `.meta.json` sidecar to build an in-memory index. + +#### `get_beam(beamline, element, mode) -> list[ParticleGroup]` + +Load every cached beam matching a `(beamline, element, mode)` triple. + +| Parameter | Type | Example | +| ---------- | ----- | ------------------------------------------ | +| `beamline` | `str` | `"facet2"`, `"cu_inj"`, `"cu_hxr"` | +| `element` | `str` | `"PR10241"`, `"L0AFEND"`, `"YAG03"` | +| `mode` | `str` | `"nominal"`, `"nominal_one_bunch"`, `"two_bunch"` | + +Returns `list[pmd_beamphysics.ParticleGroup]` (always a list — index `[0]` when +you know there's a single match). Raises `KeyError` with the list of available +`(element, mode)` pairs when nothing matches for the beamline, or the list of +known beamlines when the beamline itself is unknown. + +```python +from virtual_accelerator.beams import get_beam + +[pg] = get_beam(beamline="facet2", element="L0AFEND", mode="nominal_one_bunch") +pg_small = pg.resample(1000) # for a smaller particle count, resample on the fly +``` + +#### `list_beams(beamline=None, element=None, mode=None) -> list[BeamEntry]` + +Enumerate metadata entries without opening the HDF5 blobs. Every parameter is +optional; passing `None` skips that filter, and filters are AND-combined. + +```python +from virtual_accelerator.beams import list_beams + +# List every beam the registry knows about. +for entry in list_beams(): + print(entry.beamline, entry.element, entry.mode, entry.n_particles) + +# Filter by mode (e.g. all two-bunch beams across every beamline). +for entry in list_beams(mode="two_bunch"): + print(entry.path) + +# Each entry exposes .path (the .h5 file) plus every sidecar field, and +# .load() returns the ParticleGroup lazily. +entry = list_beams(beamline="facet2", element="PR10241")[0] +pg = entry.load() +``` + +`BeamEntry` fields mirror the sidecar schema: `path`, `beamline`, `element`, +`mode`, `s_m`, `ref_energy_eV`, `generator`, `n_particles`, `charge_C`, +`species`, `date_generated`, `source`, `notes`, plus a catch-all `extra` dict +for future sidecar fields. + +### Adding a new beam + +1. Place the `.h5` in an appropriate subdirectory of `virtual_accelerator/beams/` + (create a new date-tagged subdirectory if the scenario is new). The `*.h5` + LFS filter in `.gitattributes` handles the tracking automatically. +2. Write a `.meta.json` sidecar next to it. Required fields: + + ```json + { + "beamline": "facet2", + "element": "PR10241", + "s_m": 0.942, + "ref_energy_eV": 6.099e6, + "generator": "impact", + "n_particles": 100000, + "charge_C": 1.6e-9, + "species": "electron", + "date_generated": "YYYY-MM-DD", + "mode": "nominal_one_bunch", + "source": "where this beam came from (upstream repo path, notebook, run command, ...)", + "notes": "" + } + ``` + + - `beamline` names the machine (`facet2`, `cu_inj`, `cu_hxr`) — this is the + key `get_beam` searches on. + - `element` names the handoff plane (`PR10241`, `L0AFEND`, `YAG03`, ...). + - `mode` names the operating scenario (`nominal`, `nominal_one_bunch`, + `two_bunch`, ...). Multiple beams may share `(beamline, element, mode)` — + e.g. different particle counts of the same physical distribution — and + `get_beam` returns all of them. + - `s_m`, `ref_energy_eV`, and `n_particles` can be filled from + `pmd_beamphysics.ParticleGroup(path).avg("z")`, `.avg("energy")`, + `.n_particle`. + +3. The pre-commit hook (`scripts/check_beam_sidecars.py`) refuses commits that + add an `.h5` without a matching sidecar, and `tests/test_beam_sidecars.py` + enforces the schema in CI. +4. The new beam is now discoverable — no Python changes needed: + + ```python + from virtual_accelerator.beams import get_beam + [pg] = get_beam(beamline="facet2", element="PR10241", mode="nominal_one_bunch") + ``` + +### Downsampling + +We commit one canonical distribution per `(beamline, element, mode)` (with +occasional smaller-count copies of the same distribution for fast tests). To +run a model with fewer particles at runtime, use `ParticleGroup.resample`: + +```python +[pg] = get_beam(beamline="facet2", element="L0AFEND", mode="nominal_one_bunch") +small = pg.resample(5000) +``` + ## Running the models You can use the runner script to start the model. The script allows you to specify the model backend, diff --git a/scripts/check_beam_sidecars.py b/scripts/check_beam_sidecars.py new file mode 100644 index 0000000..3d1ba6f --- /dev/null +++ b/scripts/check_beam_sidecars.py @@ -0,0 +1,34 @@ +"""Pre-commit hook: every ``beams/*.h5`` file must have a matching ``.meta.json`` sidecar. + +Invoked with a list of paths (staged ``.h5`` files under ``virtual_accelerator/beams/``) +and exits non-zero if any of them lack a sidecar in the Git index (tracked or staged). +""" + +import subprocess +import sys +from pathlib import Path + + +def index_has(path: Path) -> bool: + result = subprocess.run( + ["git", "ls-files", "--cached", "--error-unmatch", "--", str(path)], + capture_output=True, + ) + return result.returncode == 0 + + +def main(argv: list[str]) -> int: + missing = [] + for arg in argv: + h5 = Path(arg) + sidecar = h5.with_suffix(".meta.json") + if not index_has(sidecar): + missing.append(f"{h5}: missing sidecar {sidecar.name} in the commit") + if missing: + print("\n".join(missing), file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/L0AFEND_100000.h5 b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/L0AFEND_100000.h5 new file mode 100644 index 0000000..57aee02 --- /dev/null +++ b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/L0AFEND_100000.h5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:997f2a8cff9797e4f738f15136dd028b23430c722e192daf5b919bbccfc604b2 +size 4809504 diff --git a/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/L0AFEND_100000.meta.json b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/L0AFEND_100000.meta.json new file mode 100644 index 0000000..e7bd8ba --- /dev/null +++ b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/L0AFEND_100000.meta.json @@ -0,0 +1,14 @@ +{ + "beamline": "facet2", + "element": "L0AFEND", + "s_m": 4.128, + "ref_energy_eV": 6.862e7, + "generator": "impact", + "n_particles": 100000, + "charge_C": 1.6e-9, + "species": "electron", + "date_generated": "2024-10-22", + "mode": "nominal_one_bunch", + "source": "FACET2-S2E 2024-10-22 one-bunch scenario", + "notes": "Default beam for the FACET-II Bmad model (resolved by virtual_accelerator.beams.get_beam). Mean z and energy computed with pmd_beamphysics.ParticleGroup: z=4.128 m, energy=68.62 MeV. The FACET2-S2E 2024-10-22_Impact_OneBunch run stops at PR10241 upstream of L0AF." +} diff --git a/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/PR10241_100000.h5 b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/PR10241_100000.h5 new file mode 100644 index 0000000..179bfc7 --- /dev/null +++ b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/PR10241_100000.h5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4169195ad94d364669ca2802ed7fae2bc43ddcaf75ca15d2b65dfa45356f3c8 +size 4811512 diff --git a/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/PR10241_100000.meta.json b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/PR10241_100000.meta.json new file mode 100644 index 0000000..56127f6 --- /dev/null +++ b/virtual_accelerator/beams/2024-10-22_facet2_oneBunch/PR10241_100000.meta.json @@ -0,0 +1,14 @@ +{ + "beamline": "facet2", + "element": "PR10241", + "s_m": 0.942, + "ref_energy_eV": 6.099e6, + "generator": "impact", + "n_particles": 100000, + "charge_C": 1.6e-9, + "species": "electron", + "date_generated": "2024-10-22", + "mode": "nominal_one_bunch", + "source": "IMPACT-T output from the FACET2-S2E 2024-10-22 one-bunch scenario", + "notes": "Mean z and energy computed with pmd_beamphysics.ParticleGroup: z=0.942 m, energy=6.099 MeV. Consistent with PR10241 (s=0.942 m, ~6 MeV) upstream of L0AF in the gun/solenoid region." +} diff --git a/virtual_accelerator/beams/2024-10-22_oneBunch.h5 b/virtual_accelerator/beams/2024-10-22_oneBunch.h5 deleted file mode 100644 index 851d6ed..0000000 Binary files a/virtual_accelerator/beams/2024-10-22_oneBunch.h5 and /dev/null differ diff --git a/virtual_accelerator/beams/2026-09-17_cu_inj_impact/YAG03_50000.h5 b/virtual_accelerator/beams/2026-09-17_cu_inj_impact/YAG03_50000.h5 new file mode 100644 index 0000000..d474b37 --- /dev/null +++ b/virtual_accelerator/beams/2026-09-17_cu_inj_impact/YAG03_50000.h5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:26a652889bba64bda2e4f23ffb634f273a2646766727228bff4c2a66743357d8 +size 2411512 diff --git a/virtual_accelerator/beams/2026-09-17_cu_inj_impact/YAG03_50000.meta.json b/virtual_accelerator/beams/2026-09-17_cu_inj_impact/YAG03_50000.meta.json new file mode 100644 index 0000000..c03ebab --- /dev/null +++ b/virtual_accelerator/beams/2026-09-17_cu_inj_impact/YAG03_50000.meta.json @@ -0,0 +1,14 @@ +{ + "beamline": "cu_inj", + "element": "YAG03", + "s_m": 4.614, + "ref_energy_eV": 6.386e7, + "generator": "impact", + "n_particles": 50000, + "charge_C": 2.5e-10, + "species": "electron", + "date_generated": "2026-09-17", + "mode": "nominal", + "source": "IMPACT-T output from the LCLS Cu injector run terminating at YAG03 (originally impact_inj_50k_YAG03.h5)", + "notes": "Mean z and energy computed with pmd_beamphysics.ParticleGroup: z=4.614 m, energy=63.86 MeV. Total charge 2.5e-10 C (50k macroparticles). Intended as an injector-handoff beam for the Cu HXR impact model at end_element=YAG03." +} diff --git a/virtual_accelerator/beams/PR10241_impact_10000.h5 b/virtual_accelerator/beams/PR10241_impact_10000.h5 deleted file mode 100644 index c3fb493..0000000 Binary files a/virtual_accelerator/beams/PR10241_impact_10000.h5 and /dev/null differ diff --git a/virtual_accelerator/beams/PR10241_impact_100000.h5 b/virtual_accelerator/beams/PR10241_impact_100000.h5 deleted file mode 100644 index 501d7d5..0000000 Binary files a/virtual_accelerator/beams/PR10241_impact_100000.h5 and /dev/null differ diff --git a/virtual_accelerator/beams/__init__.py b/virtual_accelerator/beams/__init__.py new file mode 100644 index 0000000..ecb4fd4 --- /dev/null +++ b/virtual_accelerator/beams/__init__.py @@ -0,0 +1,270 @@ +"""Registry for cached beam distributions under ``virtual_accelerator/beams/``. + +Each ``.h5`` beam distribution ships with a matching ``.meta.json`` sidecar +describing which beamline it belongs to, which element it was captured at, the +operating mode, and other physical parameters. This module discovers those +sidecars lazily on first use and exposes two entry points: + +* :func:`get_beam` — load beams by ``beamline``, ``element``, and ``mode``. +* :func:`list_beams` — enumerate metadata entries without opening the HDF5. + +The registry walks the beams directory the first time it is queried and caches +the result for the lifetime of the process. Call :func:`clear_cache` in tests +that mutate the beams directory. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from virtual_accelerator.utils.optional_dependencies import import_optional_symbol + + +BEAMS_DIR = Path(__file__).resolve().parent + + +@dataclass(frozen=True) +class BeamEntry: + """One cached beam distribution and its sidecar metadata. + + ``path`` points at the ``.h5`` file; the remaining fields mirror the + ``.meta.json`` schema documented in the top-level README. Use + :meth:`load` to materialize the beam as an openPMD ``ParticleGroup``. + """ + + path: Path + beamline: str + element: str + mode: str + s_m: float + ref_energy_eV: float + generator: str + n_particles: int + charge_C: float + species: str + date_generated: str + source: str + notes: str + extra: dict[str, Any] = field(default_factory=dict) + + def load(self): + """Materialize the on-disk beam as a ``pmd_beamphysics.ParticleGroup``. + + Opens ``self.path`` with h5py and hands the relevant group to + ``ParticleGroup``. Two on-disk layouts are supported transparently: + + * **root layout** — particle datasets live at the file root + (``/position/x``, ``/momentum/px``, ...). The entire file handle is + passed straight through. + * **openPMD tree** — datasets live under ``/particles//``. + The species named in this entry's sidecar is preferred; if it is + not present in the file, the first (and typically only) species + group is used instead. + + Returns + ------- + pmd_beamphysics.ParticleGroup + The loaded beam distribution. + + Raises + ------ + ImportError + If ``pmd_beamphysics`` is not installed. Install the ``bmad`` + extra to pull it in. + ValueError + If the file contains an openPMD ``/particles`` group with no + species inside. + OSError + If ``self.path`` cannot be opened by h5py. + """ + ParticleGroup = import_optional_symbol( + "pmd_beamphysics", + "ParticleGroup", + feature="loading cached beam distributions", + extra="bmad", + ) + import h5py + + with h5py.File(str(self.path), "r") as f: + if "particles" in f: + species_names = list(f["particles"].keys()) + if not species_names: + raise ValueError(f"{self.path} has an empty /particles group") + # Prefer the sidecar's species when it's present; otherwise + # take the only species we find. + species = ( + self.species if self.species in species_names else species_names[0] + ) + return ParticleGroup(h5=f[f"particles/{species}"]) + return ParticleGroup(h5=f) + + +_KNOWN_FIELDS = { + "beamline", + "element", + "mode", + "s_m", + "ref_energy_eV", + "generator", + "n_particles", + "charge_C", + "species", + "date_generated", + "source", + "notes", +} + + +_cache: list[BeamEntry] | None = None + + +def _scan() -> list[BeamEntry]: + entries: list[BeamEntry] = [] + for sidecar in sorted(BEAMS_DIR.rglob("*.meta.json")): + data = json.loads(sidecar.read_text()) + # Strip the ".meta.json" double-suffix to get the paired .h5 path. + h5 = sidecar.with_name(sidecar.name.removesuffix(".meta.json") + ".h5") + if not h5.exists(): + # Sidecar without a beam file — skip; the pre-commit hook and + # test_beam_sidecars would already flag this. + continue + known = {k: data[k] for k in _KNOWN_FIELDS if k in data} + extra = {k: v for k, v in data.items() if k not in _KNOWN_FIELDS} + entries.append(BeamEntry(path=h5, extra=extra, **known)) + return entries + + +def _entries() -> list[BeamEntry]: + global _cache + if _cache is None: + _cache = _scan() + return _cache + + +def clear_cache() -> None: + """Discard the in-memory registry so the next lookup re-scans the disk. + + The registry walks ``BEAMS_DIR`` once and memoizes the result for the + lifetime of the process, which is fine for normal use but stale in tests + that add, remove, or rewrite beam files at runtime. Call this after any + such mutation so the next :func:`list_beams` or :func:`get_beam` call + sees the new state. + + Returns + ------- + None + """ + global _cache + _cache = None + + +def list_beams( + *, + beamline: str | None = None, + element: str | None = None, + mode: str | None = None, +) -> list[BeamEntry]: + """Return cached :class:`BeamEntry` records matching the given filters. + + All filters are keyword-only and ANDed together; passing ``None`` (the + default) skips that filter and matches every value. Only the sidecar + metadata is consulted — the ``.h5`` blobs are never opened, so this call + is cheap and safe to use for discovery, listing, or completion. + + Parameters + ---------- + beamline : str, optional + Match only entries whose ``beamline`` field equals this value. + element : str, optional + Match only entries whose ``element`` field equals this value. + mode : str, optional + Match only entries whose ``mode`` field equals this value. + + Returns + ------- + list[BeamEntry] + Matching entries in the order they were discovered on disk. Empty + list when nothing matches — this function never raises for a miss; + use :func:`get_beam` if you want an error instead. + + Examples + -------- + >>> list_beams(beamline="FACET2") # doctest: +SKIP + >>> list_beams(beamline="FACET2", element="L0AFEND") # doctest: +SKIP + """ + results = _entries() + if beamline is not None: + results = [e for e in results if e.beamline == beamline] + if element is not None: + results = [e for e in results if e.element == element] + if mode is not None: + results = [e for e in results if e.mode == mode] + return results + + +def get_beam( + beamline: str, + element: str, + mode: str, +): + """Load every cached beam matching ``(beamline, element, mode)``. + + Looks up sidecars via :func:`list_beams` and materializes each match by + calling :meth:`BeamEntry.load`. All three keys are required — this is + the strict, "give me the beam or fail loudly" entry point; use + :func:`list_beams` for filtered discovery that never raises. + + Parameters + ---------- + beamline : str + Beamline identifier from the sidecar (e.g. ``"FACET2"``). + element : str + Element identifier at which the beam was captured + (e.g. ``"L0AFEND"``). + mode : str + Operating mode label from the sidecar (e.g. ``"oneBunch"``). + + Returns + ------- + list[pmd_beamphysics.ParticleGroup] + One entry per matching sidecar, in discovery order. Multiple hits + are possible when several cached beams share the same + ``(beamline, element, mode)`` triple (for example, different + particle counts or generation dates). + + Raises + ------ + KeyError + If no cached beam matches. The message enumerates the available + ``(element, mode)`` pairs for that beamline, or the known beamlines + when the beamline itself is unknown. + ImportError + Propagated from :meth:`BeamEntry.load` when ``pmd_beamphysics`` is + not installed. + """ + matches = list_beams(beamline=beamline, element=element, mode=mode) + if not matches: + available = sorted({(e.element, e.mode) for e in list_beams(beamline=beamline)}) + if not available: + raise KeyError( + f"No beams cached for beamline={beamline!r}. " + f"Known beamlines: {sorted({e.beamline for e in _entries()})}" + ) + raise KeyError( + f"No beam for beamline={beamline!r}, element={element!r}, " + f"mode={mode!r}. Available (element, mode) for {beamline!r}: " + f"{available}" + ) + return [entry.load() for entry in matches] + + +__all__ = [ + "BeamEntry", + "BEAMS_DIR", + "list_beams", + "get_beam", + "clear_cache", +] diff --git a/virtual_accelerator/models/facet2.py b/virtual_accelerator/models/facet2.py index a2f8953..2d9d416 100644 --- a/virtual_accelerator/models/facet2.py +++ b/virtual_accelerator/models/facet2.py @@ -234,6 +234,20 @@ def get_facet_bmad_model( - is_on=false for fixer elements """ from virtual_accelerator.bmad.factory import BmadModelSpec, build_bmad_model + from virtual_accelerator.beams import list_beams + + if track_beam and custom_beam_path is None and start_element == "L0AFEND": + matches = list_beams( + beamline="facet2", element="L0AFEND", mode="nominal_one_bunch" + ) + if not matches: + raise RuntimeError( + "No cached L0AFEND beam found for FACET-II. " + "Expected a facet2/L0AFEND/nominal_one_bunch entry under " + "virtual_accelerator/beams/." + ) + # Pick the largest-particle-count match as the default. + custom_beam_path = str(max(matches, key=lambda e: e.n_particles).path) spec = BmadModelSpec( feature="FACET-II Bmad model", @@ -241,7 +255,6 @@ def get_facet_bmad_model( tao_init_relpath="bmad/models/f2_elec/tao.init", mapping_beampath=None, profmon_config_filename="facet2_profmon_info.yaml", - default_beam_relpath="../beams/2024-10-22_oneBunch.h5", default_track_start="L0AFEND", ) model = build_bmad_model( diff --git a/virtual_accelerator/models/runners.py b/virtual_accelerator/models/runners.py index c7abee4..5625e93 100644 --- a/virtual_accelerator/models/runners.py +++ b/virtual_accelerator/models/runners.py @@ -59,7 +59,11 @@ def main(): if args.model == "cu_hxr_bmad": from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model - model = get_cu_hxr_bmad_model(start_element=args.start_element, end_element=args.end_element, track_beam=True) + model = get_cu_hxr_bmad_model( + start_element=args.start_element, + end_element=args.end_element, + track_beam=True, + ) elif args.model == "cu_hxr_staged": from virtual_accelerator.models.cu_hxr import get_cu_hxr_staged_model @@ -71,7 +75,11 @@ def main(): if args.start_element == "CATHODE": args.start_element = "CATHODEF" - model = get_facet_bmad_model(start_element=args.start_element, end_element=args.end_element, track_beam=True) + model = get_facet_bmad_model( + start_element=args.start_element, + end_element=args.end_element, + track_beam=True, + ) elif args.model == "facet_staged": from virtual_accelerator.models.facet2 import get_facet_staged_model diff --git a/virtual_accelerator/tests/test_beam_registry.py b/virtual_accelerator/tests/test_beam_registry.py new file mode 100644 index 0000000..1cf0c99 --- /dev/null +++ b/virtual_accelerator/tests/test_beam_registry.py @@ -0,0 +1,135 @@ +"""Tests for the ``virtual_accelerator.beams`` registry.""" + +import json + +import pytest + +from virtual_accelerator import beams as beams_module +from virtual_accelerator.beams import ( + BEAMS_DIR, + BeamEntry, + clear_cache, + list_beams, + get_beam, +) + + +@pytest.fixture(autouse=True) +def _reset_cache(): + clear_cache() + yield + clear_cache() + + +def test_scan_discovers_every_on_disk_sidecar(): + on_disk = sorted(p.name for p in BEAMS_DIR.rglob("*.meta.json")) + entries = list_beams() + discovered = sorted(e.path.name.replace(".h5", ".meta.json") for e in entries) + assert discovered == on_disk, ( + f"Registry scan missed sidecars. On disk: {on_disk}, discovered: {discovered}" + ) + + +def test_list_beams_filters_by_beamline(): + facet_beams = list_beams(beamline="facet2") + assert facet_beams, "expected at least one facet2 beam in the cache" + assert all(e.beamline == "facet2" for e in facet_beams) + + +def test_list_beams_filters_are_anded(): + entries = list_beams(beamline="facet2", element="PR10241", mode="nominal_one_bunch") + assert entries + for e in entries: + assert e.beamline == "facet2" + assert e.element == "PR10241" + assert e.mode == "nominal_one_bunch" + + +def test_list_beams_returns_empty_list_when_nothing_matches(): + assert list_beams(beamline="does_not_exist") == [] + + +def test_get_beam_raises_with_available_options_when_element_missing(): + with pytest.raises(KeyError) as excinfo: + get_beam(beamline="facet2", element="NOPE", mode="nominal_one_bunch") + msg = str(excinfo.value) + assert "facet2" in msg + assert "NOPE" in msg + # The error must list at least one real (element, mode) pair for the beamline. + assert "PR10241" in msg or "L0AFEND" in msg + + +def test_get_beam_raises_when_beamline_unknown(): + with pytest.raises(KeyError) as excinfo: + get_beam(beamline="nonexistent", element="X", mode="Y") + assert "Known beamlines" in str(excinfo.value) + + +def test_get_beam_returns_list_matching_list_beams(): + """get_beam should return one entry per list_beams match.""" + pytest.importorskip("pmd_beamphysics") + matches = list_beams(beamline="facet2", element="L0AFEND", mode="nominal_one_bunch") + if not matches: + pytest.skip("no L0AFEND beams cached") + # Skip if any matching .h5 is an unresolved LFS pointer stub. + for entry in matches: + if entry.path.stat().st_size < 1024: + head = entry.path.read_bytes()[:64] + if head.startswith(b"version https://git-lfs"): + pytest.skip(f"{entry.path.name} is an unresolved Git LFS pointer") + beams_list = get_beam( + beamline="facet2", element="L0AFEND", mode="nominal_one_bunch" + ) + assert isinstance(beams_list, list) + assert len(beams_list) == len(matches) + + +def test_cache_is_reused_across_calls(): + first = list_beams() + second = list_beams() + # Same underlying list object — no re-scan. + assert beams_module._cache is not None + assert first is second + + +def test_clear_cache_forces_rescan(): + list_beams() + assert beams_module._cache is not None + clear_cache() + assert beams_module._cache is None + list_beams() + assert beams_module._cache is not None + + +def test_beam_entry_captures_unknown_sidecar_fields(tmp_path, monkeypatch): + """Fields beyond the known schema land in `extra` instead of erroring.""" + fake_beams = tmp_path / "beams" + (fake_beams / "scenario").mkdir(parents=True) + h5 = fake_beams / "scenario" / "TEST_100.h5" + h5.write_bytes(b"not a real h5 but the registry doesn't open it") + sidecar = fake_beams / "scenario" / "TEST_100.meta.json" + sidecar.write_text( + json.dumps( + { + "beamline": "test_line", + "element": "TEST", + "mode": "nominal", + "s_m": 0.0, + "ref_energy_eV": 1e6, + "generator": "synthetic", + "n_particles": 100, + "charge_C": 1e-12, + "species": "electron", + "date_generated": "2026-01-01", + "source": "unit test", + "notes": "", + "future_field": "hello", + } + ) + ) + monkeypatch.setattr(beams_module, "BEAMS_DIR", fake_beams) + clear_cache() + entries = list_beams(beamline="test_line") + assert len(entries) == 1 + assert isinstance(entries[0], BeamEntry) + assert entries[0].extra == {"future_field": "hello"} diff --git a/virtual_accelerator/tests/test_beam_sidecars.py b/virtual_accelerator/tests/test_beam_sidecars.py new file mode 100644 index 0000000..4dfaf89 --- /dev/null +++ b/virtual_accelerator/tests/test_beam_sidecars.py @@ -0,0 +1,101 @@ +"""Structural checks on cached beam files under ``virtual_accelerator/beams/``. + +Every ``.h5`` beam distribution must ship with a matching ``.meta.json`` +sidecar of the same base name, and every sidecar must contain a minimal set of +fields so a reader can tell what the beam represents without opening the +(possibly LFS-only) HDF5 blob. +""" + +import json +from pathlib import Path + +import pytest + + +BEAMS_DIR = Path(__file__).resolve().parent.parent / "beams" + +REQUIRED_SIDECAR_FIELDS = frozenset( + { + "beamline", + "element", + "s_m", + "ref_energy_eV", + "generator", + "n_particles", + "charge_C", + "species", + "date_generated", + "mode", + "source", + "notes", + } +) + + +def _h5_files() -> list[Path]: + return sorted(BEAMS_DIR.rglob("*.h5")) + + +@pytest.mark.parametrize( + "h5_path", _h5_files(), ids=lambda p: str(p.relative_to(BEAMS_DIR)) +) +def test_h5_beam_has_json_sidecar(h5_path: Path) -> None: + sidecar = h5_path.with_suffix(".meta.json") + assert sidecar.exists(), ( + f"{h5_path.relative_to(BEAMS_DIR)} has no sidecar. " + f"Expected {sidecar.name} alongside it." + ) + + +@pytest.mark.parametrize( + "h5_path", _h5_files(), ids=lambda p: str(p.relative_to(BEAMS_DIR)) +) +def test_sidecar_has_required_fields(h5_path: Path) -> None: + sidecar = h5_path.with_suffix(".meta.json") + if not sidecar.exists(): + pytest.skip("sidecar missing; covered by test_h5_beam_has_json_sidecar") + data = json.loads(sidecar.read_text()) + missing = REQUIRED_SIDECAR_FIELDS - data.keys() + assert not missing, ( + f"{sidecar.relative_to(BEAMS_DIR)} is missing required fields: " + f"{sorted(missing)}" + ) + + +@pytest.mark.parametrize( + "h5_path", _h5_files(), ids=lambda p: str(p.relative_to(BEAMS_DIR)) +) +def test_h5_loads_as_openpmd_particle_group(h5_path: Path) -> None: + """Each cached ``.h5`` must be readable as an openPMD ``ParticleGroup``. + + Supports both on-disk layouts we ship: datasets at the root + (``/position/x`` ...) and datasets under ``/particles//``. + """ + pmd_beamphysics = pytest.importorskip("pmd_beamphysics") + h5py = pytest.importorskip("h5py") + + # Skip LFS pointer stubs so this test is a no-op when blobs weren't pulled. + if h5_path.stat().st_size < 1024: + head = h5_path.read_bytes()[:64] + if head.startswith(b"version https://git-lfs"): + pytest.skip(f"{h5_path.name} is an unresolved Git LFS pointer") + + with h5py.File(str(h5_path), "r") as f: + if "particles" in f: + species = next(iter(f["particles"].keys())) + pg = pmd_beamphysics.ParticleGroup(h5=f[f"particles/{species}"]) + else: + pg = pmd_beamphysics.ParticleGroup(h5=f) + + assert pg.n_particle > 0, ( + f"{h5_path.relative_to(BEAMS_DIR)} loaded but reports zero particles" + ) + + sidecar = h5_path.with_suffix(".meta.json") + if sidecar.exists(): + expected = json.loads(sidecar.read_text()).get("n_particles") + if expected is not None: + assert pg.n_particle == expected, ( + f"{h5_path.relative_to(BEAMS_DIR)} has {pg.n_particle} particles " + f"but sidecar declares n_particles={expected}" + )