Beam Caching

.gitattributes

@@ -0,0 +1 @@
+*.h5 filter=lfs diff=lfs merge=lfs -text

.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

.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

.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

.gitignore

@@ -213,6 +213,7 @@ __marimo__/
 .vscode/
 
 *.h5
+!virtual_accelerator/beams/**/*.h5
 *.png
 *.svg
 

.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

README.md

@@ -48,7 +48,6 @@ conda install -c conda-forge impact-t=*=mpi_mpich*
 And the examples require installing ipykernel and register as a Jupyter kernel.
 
 
-<<<<<<< HEAD
 Optional Dependency Keys by Model:
 | Model / Factory Function | Optional dependency key(s) | Notes |
 | --- | --- | --- |
@@ -60,7 +59,6 @@ Optional Dependency Keys by Model:
 | `get_cu_hxr_zfel_model` | `zfel` | CU HXR taper model using the 1D ZFEL backend. |
 | `virtual_accelerator.models.runners` CLI | `pva` (+ model backend key) | Runner requires `pva`; selected model backend must also be installed. |
 | `get_cu_inj_impact_model` | `Impact` | Requires impact pip install AND conda install, both detailed above |
-=======
 ## Loading a model
 
 Use `get_model()` to build a single model or a staged chain. See
@@ -112,7 +110,6 @@ Standard staged chains (build with `get_model([upstream, downstream], ...)`):
 | `fast_facet2_s2e` | `surrogate_f2e_inj` | `bmad_f2_elec` | PR10241 |
 
 The `Runner` CLI additionally needs the `pva` extra.
->>>>>>> upstream/main
 
 The package now lazily imports backend-specific dependencies. If you call a model
 whose optional dependency is not installed, you will get an actionable error with
@@ -123,6 +120,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 `<name>.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,

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:]))

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."
+}

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."
+}

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."
+}