From a8c8becb1a17398efad1e7390f14cd4d55d3bd25 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Tue, 1 Sep 2026 16:36:54 -0700 Subject: [PATCH 01/15] registry --- docs/model_registry_design.md | 375 +++++++++++++++++++++ virtual_accelerator/registry/__init__.py | 319 ++++++++++++++++++ virtual_accelerator/registry/models.py | 162 +++++++++ virtual_accelerator/tests/test_registry.py | 195 +++++++++++ 4 files changed, 1051 insertions(+) create mode 100644 docs/model_registry_design.md create mode 100644 virtual_accelerator/registry/__init__.py create mode 100644 virtual_accelerator/registry/models.py create mode 100644 virtual_accelerator/tests/test_registry.py diff --git a/docs/model_registry_design.md b/docs/model_registry_design.md new file mode 100644 index 0000000..639205b --- /dev/null +++ b/docs/model_registry_design.md @@ -0,0 +1,375 @@ +# Model registry + unified `get_model` — design proposal + +Status: **implemented for LCLS** in `virtual_accelerator/registry/`. FACET-II entries are not +registered yet; §7 (screen PV prefixes) should land before they are. + +Goal: one entry point for every model, a registry of what exists and how to configure it, and +enough per-model metadata to stitch stages together safely. + +Scope decision driving this revision: **handoff points are diagnostics only.** That single +constraint removes the need for an element catalog, because diagnostics are the one category of +element whose names already agree across engines (see §1). Each model carries a plain ordered list +of the diagnostics it exposes, plus an optional name-override dict for future divergence. + +--- + +## 1. Evidence: why a list of names is enough + +Checked against the real lattice files, not assumed. + +**Diagnostics already agree between Bmad and IMPACT, in both facilities.** The suspected +`YAG03`/`YAG3` mismatch does not exist. Confirmed screen-by-screen: + +| facility | diagnostic | Bmad | IMPACT-T | +|---|---|---|---| +| LCLS | YAG02, YAG03, OTR1, OTR2 | same | same | +| FACET | PR10241, PR10465, PR10471, PR10571 | same | same | + +Bmad aliases match `utils/cu_hxr_profmon_info.yaml` exactly (`OTR2[alias]=OTRS:IN20:571`, +`YAG03[alias]=YAGS:IN20:351`), so the profmon YAMLs already are a de-facto standard-name list — +this proposal just makes them addressable per model. + +Note that the *element names* agreeing is separate from the *PVs* agreeing — the latter is broken +for FACET, see §7. + +**The divergences that exist are all in non-diagnostic elements**, i.e. things that are never +handoff points: + +| element | Bmad | IMPACT-T | +|---|---|---| +| gun / cathode | `CATHODE` / `CATHODEF` | `GUN` / `GUNF` | +| L0A cavity | `L0A` (one lcavity) | `L0A_entrance`, `L0A_body_1`, `L0A_body_2`, `L0A_exit` | + +So **the name-override dict is empty for both facilities today.** It exists so a future rename has +somewhere to live, not because it is currently load-bearing. + +Two facts that don't need machinery but should be written down, because they are the reason not to +validate handoffs numerically: + +- IMPACT z is element *entrance*; the lattice CSV's `SumL` is element *centre*. They coincide only + for zero-length monitors — which is another reason to restrict handoffs to diagnostics. +- FACET IMPACT geometry disagrees with the Bmad geometry by up to ~3.9 mm at `PR10571`, growing + downstream. LCLS agrees exactly. Real, but a physics caveat rather than something the registry + should police. + +--- + +## 2. Registry + +One table. Models only; no separate element registry. + +```python +@dataclass(frozen=True) +class ModelEntry: + name: str # "bmad_cu_hxr" + description: str # what models_available prints + facility: str # "lcls" | "facet2" + engine: str # "bmad" | "impact" | "surrogate" | "cheetah" + builder: str # "virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model" + extras: tuple[str, ...] # pip extras needed, e.g. ("bmad",) + params: dict[str, Any] # param name -> default, for validation + discovery + broadcast_params: frozenset[str] # params safe to send to every stage, e.g. {"n_particles"} + + diagnostics: tuple[str, ...] # standard names, ORDERED by lattice position + start_param: str | None # builder kwarg for start, None if fixed + end_param: str | None # builder kwarg for end, None if fixed + default_start: str | None + default_end: str | None + + images_diagnostics: bool # publishes full screen-image PVs? False for surrogates + element_after: dict[str, str] # diagnostic -> element immediately downstream of it + + element_aliases: dict[str, str] = field(default_factory=dict) + # standard name -> this engine's local name. Empty today; escape hatch only. +``` + +`start_param` / `end_param` replace the earlier `can_start_anywhere` flag: they carry the same +information (`None` means the extent is fixed) while also recording what each builder *calls* the +parameter, which differs between engines. + +`images_diagnostics` and `element_after` exist because of §5's collision rule — see there for why. + +Three choices worth calling out: + +**`builder` is a `"module:function"` string, not a callable.** The existing builders lazily import +`pytao` / `torch` / `impact` inside the function body via `import_optional`. A real callable in the +registry would force importing every engine module just to import the registry, breaking +`models_available` for anyone without all extras installed. + +**`diagnostics` is ordered by lattice position.** Ordering is then checkable by list index instead +of by metres, which is what lets §4 drop positions without losing the ordering check. + +**No data files, no generator script.** The lists are hand-maintained Python literals in +`registry/models.py`. They are short, they change only when a lattice model changes, and a diff is +readable. Revisit only if they start drifting from the lattice. + +``` +virtual_accelerator/registry/ +├── __init__.py # get_model, models_available, list_diagnostics +└── models.py # the ModelEntry table +``` + +--- + +## 3. Model catalog + +| registry name | facility | engine | builder | key params | diagnostics (ordered) | +|---|---|---|---|---|---| +| `impact_cu_inj` | lcls | impact | `get_cu_inj_impact_model` | `n_particles=100`, `end_element="OTR2"` | YAG02, YAG03, OTR1, OTR2 | +| `bmad_cu_hxr` | lcls | bmad | `get_cu_hxr_bmad_model` | `start_element="OTR2"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | all 12 profmon screens, YAG02 → OTRDMP | +| `surrogate_cu_inj` | lcls | surrogate | `get_cu_hxr_injector_surrogate_model` | `n_particles=1000` | OTR2 (fixed end) | +| `cheetah_cu_hxr` | lcls | cheetah | `get_cu_hxr_cheetah_model` | `n_particles=1000` | — | + +Naming convention: `_`. + +**Not yet registered** — FACET-II, pending §7: + +| registry name | facility | engine | builder | key params | diagnostics (ordered) | +|---|---|---|---|---|---| +| `impact_f2e_inj` | facet2 | impact | `get_facet_impact_model` | `n_particles=100`, `end_element="PR10571"` | PR10241, PR10465, PR10471, PR10571 | +| `bmad_f2_elec` | facet2 | bmad | `get_facet_bmad_model` | `start_element="L0AFEND"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | PR10241, PR10465, PR10471, PR10571, PR10711 | +| `surrogate_f2e_inj` | facet2 | surrogate | ⚠ **does not exist yet** | `n_particles=10000`, `surrogate_inputs="machine"` | PR10241 (fixed end) | + +Notes on the lists: + +- `impact_cu_inj` omits `OTRH1`/`OTRH2` (laser heater is unmodeled in the deck: `!!! Unmodeled: + Laser Heater from 9.076892 m to 10.690580 m`) and `OTR3`/`YAG01` (lines commented out — + `YAG01` is marked `!!! Broken:`). `OTR4` is past `stop_1` at z=16.5. +- `impact_f2e_inj` includes `PR10571` because the file actually loaded is + `ImpactT_template.in` (per `ImpactT.yaml`'s `input_file:` key). The checked-in `ImpactT.in` is a + stale truncated artifact stopping at z=12.0 and omitting it — reading that file gives the wrong + answer about available stop points. Worth an upstream issue. +- `bmad_f2_elec` currently defaults `start_element="L0AFEND"`, which is a marker rather than a + diagnostic. `L0AFEND` exists in both engines (Bmad superimposed marker, IMPACT write-beam + element) so it's a legitimate handoff plane; treat markers as admissible where both engines have + them, and keep them in `diagnostics` despite the field name. (Or rename the field + `handoff_points` — probably clearer.) + +**Prerequisite:** `surrogate_f2e_inj` has no standalone builder — the `BeamOutputModel` is built +inline inside `get_facet_staged_model`. It needs extracting to +`get_facet_injector_surrogate_model()` to mirror `get_cu_hxr_injector_surrogate_model`, otherwise +it can't be a registry entry. + +`get_cu_hxr_staged_model` / `get_facet_staged_model` become thin back-compat wrappers over +`get_model([...])`. + +--- + +## 4. `get_model` + +```python +def get_model( + spec: str | Sequence[str], + *, + handoff_loc: str | Sequence[str] | None = None, + start_ele: str | None = None, + end_ele: str | None = None, + **kwargs, +) -> LUMEModel: +``` + +```python +# single +model = get_model("bmad_cu_hxr", end_ele="OTR4", track_beam=True) + +# staged, explicit handoff +model = get_model(["impact_cu_inj", "bmad_cu_hxr"], + handoff_loc="YAG03", end_ele="OTR4", n_particles=1000) + +# staged, handoff inferred from the surrogate's fixed end (OTR2) +model = get_model(["surrogate_cu_inj", "bmad_cu_hxr"], end_ele="OTR4", n_particles=10000) +``` + +Every example above is registered and working today, except that the `impact_cu_inj` stage needs +the IMPACT-T executable (`conda install -c conda-forge impact-t`) on top of the Python package. + +FACET-II is **not** registered yet, so this raises `KeyError` for now — see §3 and §7: + +```python +model = get_model(["impact_f2e_inj", "bmad_f2_elec"], handoff_loc="L0AFEND") # not yet +``` + +`start_ele` / `end_ele` on a staged call are the *overall* extent — first stage's start, last +stage's end. Interior extents come from `handoff_loc`. + +### Discovery + +```python +from virtual_accelerator.registry import models_available, list_diagnostics + +print(models_available) +# impact_cu_inj IMPACT-T LCLS injector, cathode -> OTR2 +# bmad_cu_hxr Bmad CU-HXR, gun -> undulator/dump +# ... + +list_diagnostics("bmad_cu_hxr") +# ('YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', +# 'OTR11', 'OTR12', 'OTR21', 'OTRDMP') +``` + +All 12 screens in `cu_hxr_profmon_info.yaml` were verified present in the `cu_hxr` lattice, in the +order shown. + +### Kwarg routing + +Because the registry declares each model's params, routing is a lookup, not signature +introspection: + +1. Param in `broadcast_params` (e.g. `n_particles`) → sent to every stage that declares it. +2. Declared by exactly one stage → routed there. +3. Declared by more than one stage and not broadcast → `ValueError` naming the candidates. +4. `"."` always wins. +5. Matching no declared param → rejected with a suggestion, not silently forwarded. + +`track_beam=True` is forced on **every** stage that declares it, regardless of what the user +passes. An earlier draft said "non-terminal stages only", which is wrong and was caught by testing: +a non-final stage must *produce* `final_particles`, but a non-first stage must also *accept* +`initial_particles`, and `lume_bmad.model` raises `Cannot set initial_particles when track_type is +not 'beam'` otherwise. In a two-stage surrogate → Bmad chain the Bmad stage is terminal and still +needs it. + +--- + +## 5. Compatibility checking + +`StagedModel.validate_lume_model_instances` already checks mixin presence and duplicate variable +names — but only *after* both models are constructed, and `build_impact_model` calls +`impact.run()` during construction. So a duplicate-variable error costs a full IMPACT run before +it surfaces. The registry pre-validates from metadata alone, before instantiating anything. + +All checks are name-based. No positions involved. + +**C1 — same facility.** Staging `bmad_cu_hxr` onto `impact_f2e_inj` is rejected immediately. + +**C2 — handoff is a legal point in both stages.** `handoff_loc` must be in the upstream stage's +`diagnostics` and in the downstream stage's. Plus the downstream stage must have a `start_param`, +which is what makes `get_model(["bmad_cu_hxr", "impact_cu_inj"], ...)` fail with "IMPACT models can +only start at the cathode" instead of something inscrutable from inside `set_stop_location`. + +**C3 — the downstream stage must not re-image the handoff diagnostic.** This one was found by +testing and is the most important rule in practice. + +IMPACT's `set_stop_location` prunes to `s <= stop`, so a model stopped at `YAG03` *keeps* `YAG03` +and publishes its six `YAGS:IN20:351:*` image PVs. Slicing Bmad from `YAG03` also includes the +screen and publishes the same six PVs. `StagedModel` then rejects the pair on duplicate variables — +after paying for a full IMPACT run. + +Measured on the real lattice: + +| Bmad `start_element` | n_vars | `YAGS:IN20:351:*` PVs | +|---|---|---| +| `YAG03` | 292 | 6 | +| `DL02A2` | 286 | 0 | + +So `DL02A2` in `examples/staged_example.ipynb` was **not** arbitrary and not merely a naming +inconsistency — it is a deliberate workaround for this collision. An earlier draft of this document +mischaracterised it as a rename opportunity; that was wrong. + +The rule: if the upstream and downstream stages both image the handoff diagnostic, the downstream +stage starts at `element_after[handoff]` instead. Surrogates set `images_diagnostics=False` (they +publish only `XRMS`/`YRMS` scalars), so surrogate → Bmad hands off *at* the diagnostic and needs no +skip — which is why the existing `get_cu_hxr_staged_model` works with `start_element="OTR2"`. + +`element_after` was generated from the lattice and each value verified so that its entrance face +sits exactly at the screen's `s` (Bmad's `s` is the exit face, which is why the drift *after* a +screen begins at the screen). `OTR11` and `OTR21` are omitted because both are followed by an +element named `DDG4`, which is ambiguous in the lattice and so unusable as a slice start; the error +message names the field to edit if anyone hits that case. + +**C4 — beam handoff mechanics.** One place, in the registry, resolving an existing inconsistency: +`get_facet_staged_model` writes `final_particles` to a `NamedTemporaryFile` and passes it as +`custom_beam_path`, while `get_cu_hxr_staged_model` does neither and relies solely on the +`FinalParticlesMixIn` wiring in `StagedModel._set`. One of those is redundant or one is a latent +bug; the registry should own this in exactly one place. + +Also worth noting for C4: `get_facet_staged_model` hardcodes `t0=3.15391398e-09`, `p0c=6.3e06`, +`z0=0.9420843`. That `z0` is exactly `PR10241`'s s-position — these are handoff-plane quantities, +so if a second FACET handoff plane is ever used they will need to move somewhere per-plane rather +than staying inline. + +--- + +## 6. Deliberately deferred + +Cut from the previous draft, with the trigger for reconsidering each: + +| deferred | add it when | +|---|---| +| Element catalog generated from `lcls_elements.csv` | a handoff point is needed that isn't a diagnostic, or the hand-maintained lists start drifting from the lattice | +| s-positions on handoff points | we want to *numerically* verify a handoff rather than trust name equality | +| Tolerance-based position agreement check | ditto — this is where the FACET ~3.9 mm discrepancy would resurface | +| Physical-extent overlap check between stages | `StagedModel`'s duplicate-variable check proves insufficient in practice | +| PV names accepted as `start_ele` / `end_ele` | a control-room user asks for it; cheap to add later | + +--- + +## 7. Screen PV prefixes — resolved, and a bug it exposes + +**Resolved rule (supervisor, 2026-09-01): `PROF:` is correct for FACET VAs; `YAGS:`/`OTRS:` are +correct for LCLS / LCLS-II VAs.** + +This explains an asymmetry in the current code that otherwise looks arbitrary: +`get_facet_bmad_model` carries a `custom_aliases` dict while `get_cu_hxr_bmad_model` carries none. +FACET needs it precisely *because* the lattice aliases (`YAGS:/OTRS:IN10:*`) are wrong for VA +purposes; LCLS needs none because its lattice aliases are already right. + +### How each engine currently derives a screen's PV + +- **Bmad** — `bmad/variables.py:375`: `base_pv = tao.ele(screen_name).head.alias`, i.e. the Tao + element alias. `build_bmad_model` applies `custom_aliases` (factory.py:72-78) *before* + `get_variables` (factory.py:86), so the override ordering is correct. +- **IMPACT** — `impact/variables.py`: `alias_dict[element_name]`, where `alias_dict` is + `Element -> Control System Name` from `bmad/conversion/from_oracle/lcls_elements.csv`. There is + **no** `custom_aliases` equivalent on this path. +- **Neither** reads the `name:` field of `utils/*_profmon_info.yaml`. Bmad uses that file only for + `shape` and `pixel_size` (variables.py:379-382). So the `name:` field is currently dead data — + even though it holds the correct value in every case. + +### Consequence: FACET screen PVs are wrong today + +| screen | Bmad yields | IMPACT yields | correct (per rule) | +|---|---|---|---| +| PR10241 | `PROF:IN10:241` ✓ | `YAGS:IN10:241` ✗ | `PROF:IN10:241` | +| PR10465 | `OTRS:IN10:465` ✗ | `OTRS:IN10:465` ✗ | `PROF:IN10:465` | +| PR10471 | `OTRS:IN10:471` ✗ | `OTRS:IN10:471` ✗ | `PROF:IN10:471` | +| PR10571 | `PROF:IN10:571` ✓ | `OTRS:IN10:571` ✗ | `PROF:IN10:571` | +| PR10711 | `PROF:IN10:711` ✓ | `OTRS:IN10:711` ✗ | `PROF:IN10:711` | + +FACET Bmad is wrong for 2 of 5 screens (`custom_aliases` simply omits `PR10465`/`PR10471`); FACET +IMPACT is wrong for all 5. This directly affects `impact_f2e_inj` — the model about to be staged. + +LCLS is consistent across all three sources (lattice alias, CSV `Control System Name`, and +`cu_hxr_profmon_info.yaml`), so it is unaffected either way. + +### Proposed fix — no new data structure + +Make `utils/*_profmon_info.yaml` the single source of truth for screen PVs, and have **both** +engines read `screen_config[name]["name"]` instead of the lattice alias / CSV. Then: + +- Every FACET screen is fixed in both engines at once, because the YAML already holds the right + values for all five. +- LCLS is a **no-op** — its YAML values already equal its lattice aliases and CSV entries. +- `custom_aliases` in `models/facet2.py` loses its screen entries entirely, keeping only + `TCY10490 -> KLYS:LI10:51` (a TCAV, not a screen, and a separate question). +- The dead `name:` field becomes load-bearing, so it can no longer silently drift. + +This is deliberately *not* the element catalog returning: it reuses a file that already exists and +already has the answer. Worth doing as a small standalone PR **before** the registry work, since +the registry shouldn't inherit a known-wrong PV mapping. + +One fragility worth fixing alongside: `alias_dict[element_name]` in `impact/variables.py` is an +unguarded dict lookup, so any element absent from the CSV raises a bare `KeyError`. + +--- + +## 8. Open questions + +1. **Field name** — `diagnostics` or `handoff_points`? The FACET default start `L0AFEND` is a + marker, not a diagnostic, so the latter is more honest. +2. **`TCY10490 -> KLYS:LI10:51`** — the one non-screen entry in FACET's `custom_aliases`. The + lattice says `TCY10490[alias]=TCAV:IN10:490` and the CSV agrees. Is the `KLYS:` override + deliberate in the same way the `PROF:` ones were? Same question, different device class. +3. **`models/runners.py` CLI** — its four hardcoded `--model` choices become `get_model(args.model)`, + which is a strict improvement but changes the accepted values. +4. **`BmadModelSpec.database_relpath` is dead config** — declared at `bmad/factory.py:24`, never + read in `build_bmad_model`. Unrelated to this work; noted while reading. diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py new file mode 100644 index 0000000..8c058ec --- /dev/null +++ b/virtual_accelerator/registry/__init__.py @@ -0,0 +1,319 @@ +"""Single entry point for building virtual-accelerator models by name. + +from virtual_accelerator.registry import get_model, models_available + +print(models_available) +model = get_model("bmad_cu_hxr", end_ele="OTR4", track_beam=True) +model = get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="YAG03") +""" + +import importlib +from typing import Any + +from virtual_accelerator.registry.models import MODELS, ModelEntry + +__all__ = ["get_model", "models_available", "list_models", "list_diagnostics"] + + +class _ModelCatalog(dict): + """Mapping of model name -> description that prints as an aligned table.""" + + def __repr__(self) -> str: + if not self: + return "(no models registered)" + width = max(len(name) for name in self) + return "\n".join(f"{name:<{width}} {desc}" for name, desc in self.items()) + + +models_available = _ModelCatalog( + (name, entry.description) for name, entry in MODELS.items() +) + + +def list_models(facility: str | None = None, engine: str | None = None) -> list[str]: + """Names of registered models, optionally filtered by facility or engine.""" + return [ + name + for name, entry in MODELS.items() + if (facility is None or entry.facility == facility) + and (engine is None or entry.engine == engine) + ] + + +def list_diagnostics(model_name: str) -> tuple[str, ...]: + """Diagnostics usable as start/end/handoff points, in lattice order.""" + return _entry(model_name).diagnostics + + +def _entry(name: str) -> ModelEntry: + try: + return MODELS[name] + except KeyError: + raise KeyError( + f"Unknown model {name!r}. Available: {', '.join(sorted(MODELS))}" + ) from None + + +def _load_builder(entry: ModelEntry): + module_path, _, func_name = entry.builder.partition(":") + try: + module = importlib.import_module(module_path) + except ImportError as exc: + extras = ", ".join(entry.extras) or "none" + raise ImportError( + f"Cannot import builder for {entry.name!r} ({entry.builder}). " + f"Required extras: {extras}. Install with " + f'`pip install "virtual-accelerator[{",".join(entry.extras)}]"`.' + ) from exc + return getattr(module, func_name) + + +def _normalize(name: str | None) -> str | None: + """Canonicalise a user-supplied element name to upper case. + + Lattice element names are upper case everywhere. Tao is case-insensitive so a + lower-case name would appear to work, but IMPACT's ``impact.ele[...]`` is a + plain dict lookup, and more importantly the registry's own diagnostics and + ``element_after`` lookups would silently miss -- which would defeat the + handoff collision check in ``_downstream_start``. + """ + return name if name is None else name.upper() + + +def _resolve_element(entry: ModelEntry, name: str) -> str: + """Translate a standard element name into this model's engine-local name.""" + return entry.element_aliases.get(name, name) + + +def _check_element(entry: ModelEntry, name: str, role: str) -> None: + """Validate a handoff element. Non-diagnostic names are allowed through. + + Markers and drifts (END, TD11, DL02A2) are legitimate start/end elements but + are deliberately not enumerated, so only reject a name that looks like a + diagnostic this model does not have. + """ + if name in entry.diagnostics: + return + looks_like_diagnostic = name.startswith(("OTR", "YAG", "PR")) + if looks_like_diagnostic: + raise ValueError( + f"{name!r} is not an available {role} diagnostic for {entry.name!r}. " + f"Available: {', '.join(entry.diagnostics)}" + ) + + +def _route_kwargs( + entries: list[ModelEntry], kwargs: dict[str, Any] +) -> list[dict[str, Any]]: + """Distribute flat kwargs across stages using the registry's declared params. + + Routing is a table lookup, not signature introspection, so the error + messages can name the candidate stages. + """ + routed: list[dict[str, Any]] = [{} for _ in entries] + by_name = {entry.name: i for i, entry in enumerate(entries)} + + for key, value in kwargs.items(): + stage_name, sep, param = key.partition(".") + if sep: + if stage_name not in by_name: + raise ValueError( + f"{key!r} targets stage {stage_name!r}, which is not in this model. " + f"Stages: {', '.join(by_name)}" + ) + index = by_name[stage_name] + if param not in entries[index].params: + raise ValueError( + f"{param!r} is not a parameter of {stage_name!r}. " + f"Accepted: {', '.join(sorted(entries[index].params))}" + ) + routed[index][param] = value + continue + + accepting = [i for i, entry in enumerate(entries) if key in entry.params] + if not accepting: + known = sorted({p for entry in entries for p in entry.params}) + raise ValueError( + f"{key!r} is not a parameter of any stage. Accepted: {', '.join(known)}" + ) + + broadcast = any(key in entries[i].broadcast_params for i in accepting) + if len(accepting) > 1 and not broadcast: + names = ", ".join(entries[i].name for i in accepting) + raise ValueError( + f"{key!r} is ambiguous across stages ({names}). " + f'Qualify it, e.g. "{entries[accepting[0]].name}.{key}=...".' + ) + for i in accepting: + routed[i][key] = value + + return routed + + +def _build( + entry: ModelEntry, + call_kwargs: dict[str, Any], + start_ele: str | None, + end_ele: str | None, +) -> Any: + kwargs = dict(call_kwargs) + + if start_ele is not None: + if entry.start_param is None: + raise ValueError( + f"{entry.name!r} has a fixed start and does not accept start_ele." + ) + _check_element(entry, start_ele, "start") + kwargs[entry.start_param] = _resolve_element(entry, start_ele) + + if end_ele is not None: + if entry.end_param is None: + raise ValueError( + f"{entry.name!r} has a fixed end and does not accept end_ele." + ) + _check_element(entry, end_ele, "end") + kwargs[entry.end_param] = _resolve_element(entry, end_ele) + + return _load_builder(entry)(**kwargs) + + +def _resolve_handoffs( + entries: list[ModelEntry], handoff_loc: str | list[str] | None +) -> list[str]: + """Determine the handoff element between each consecutive pair of stages.""" + n_handoffs = len(entries) - 1 + + if handoff_loc is None: + handoffs = [] + for upstream in entries[:-1]: + if upstream.default_end is None: + raise ValueError( + f"handoff_loc is required: {upstream.name!r} has no default end " + "to infer it from." + ) + handoffs.append(upstream.default_end) + return handoffs + + handoffs = [handoff_loc] if isinstance(handoff_loc, str) else list(handoff_loc) + if len(handoffs) != n_handoffs: + raise ValueError( + f"{len(entries)} stages need {n_handoffs} handoff location(s), " + f"got {len(handoffs)}." + ) + return handoffs + + +def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) -> None: + if upstream.facility != downstream.facility: + raise ValueError( + f"Cannot stage {upstream.name!r} ({upstream.facility}) onto " + f"{downstream.name!r} ({downstream.facility}): different facilities." + ) + + if downstream.start_param is None: + reason = ( + "IMPACT models can only start at the cathode" + if downstream.engine == "impact" + else f"{downstream.name!r} has a fixed start" + ) + raise ValueError(f"{downstream.name!r} cannot be a downstream stage: {reason}.") + + _check_element(upstream, handoff, "end") + _check_element(downstream, handoff, "start") + + +def _downstream_start( + upstream: ModelEntry, downstream: ModelEntry, handoff: str +) -> str: + """Where the downstream stage should actually begin tracking. + + If both stages image the handoff diagnostic they would publish identical + screen PVs and StagedModel would reject the pair, so the downstream stage + starts at the element immediately after it instead. + """ + both_image = upstream.images_diagnostics and downstream.images_diagnostics + if not (both_image and handoff in downstream.diagnostics): + return handoff + + try: + return downstream.element_after[handoff] + except KeyError: + raise ValueError( + f"Both {upstream.name!r} and {downstream.name!r} image {handoff!r}, so " + f"{downstream.name!r} must start just after it, but no downstream element " + f"is recorded for {handoff!r}. Add it to {downstream.name}.element_after " + "in virtual_accelerator/registry/models.py." + ) from None + + +def get_model( + spec: str | list[str], + *, + handoff_loc: str | list[str] | None = None, + start_ele: str | None = None, + end_ele: str | None = None, + **kwargs: Any, +): + """Build a model, or a staged chain of models, by registry name. + + Parameters + ---------- + spec + A registry name, or an ordered list of names to stage together. + handoff_loc + Element where each consecutive pair hands the beam over. Inferred from + the upstream stage's fixed end when it has one. A list is required for + more than two stages. + start_ele, end_ele + Overall extent. For a staged model these apply to the first and last + stage respectively; interior extents come from ``handoff_loc``. + **kwargs + Builder parameters. For staged models, qualify an ambiguous parameter as + ``"."``. + """ + start_ele, end_ele = _normalize(start_ele), _normalize(end_ele) + + if isinstance(spec, str): + entry = _entry(spec) + (routed,) = _route_kwargs([entry], kwargs) + return _build(entry, routed, start_ele, end_ele) + + names = list(spec) + if len(names) < 2: + raise ValueError("Staging requires at least two models.") + + entries = [_entry(name) for name in names] + handoffs = [_normalize(h) for h in _resolve_handoffs(entries, handoff_loc)] + routed = _route_kwargs(entries, kwargs) + + for upstream, downstream, handoff in zip(entries, entries[1:], handoffs): + _validate_pair(upstream, downstream, handoff) + + stages = [] + for i, entry in enumerate(entries): + stage_start = ( + start_ele + if i == 0 + else _downstream_start(entries[i - 1], entry, handoffs[i - 1]) + ) + stage_end = end_ele if i == len(entries) - 1 else handoffs[i] + + stage_kwargs = dict(routed[i]) + # Every stage needs tracking on: a non-final stage has to produce + # final_particles, and a non-first stage has to accept initial_particles + # (lume_bmad rejects those unless track_type is 'beam'). + if "track_beam" in entry.params: + stage_kwargs["track_beam"] = True + + stages.append( + _build( + entry, + stage_kwargs, + stage_start if entry.start_param else None, + stage_end if entry.end_param else None, + ) + ) + + from lume.staged_model import StagedModel + + return StagedModel(stages) diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py new file mode 100644 index 0000000..6675d2d --- /dev/null +++ b/virtual_accelerator/registry/models.py @@ -0,0 +1,162 @@ +"""Registry of available virtual-accelerator models. + +Currently LCLS only; FACET-II entries are not registered yet. +""" + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True) +class ModelEntry: + """Metadata describing one model and how to configure it. + + ``builder`` is a ``"module:function"`` string rather than a callable so that + importing this module does not import pytao / torch / impact. Discovery must + work with no optional dependencies installed. + """ + + name: str + description: str + facility: str + engine: str + builder: str + extras: tuple[str, ...] + + params: dict[str, Any] + """Configurable parameter name -> default. Also the allow-list for kwargs.""" + + diagnostics: tuple[str, ...] + """Standard-named diagnostics usable as handoff points, in lattice order.""" + + start_param: str | None = None + """Builder kwarg controlling the start element, or None if not configurable.""" + + end_param: str | None = None + """Builder kwarg controlling the end element, or None if not configurable.""" + + default_start: str | None = None + default_end: str | None = None + + broadcast_params: frozenset[str] = frozenset() + """Params safe to send to every stage of a staged model, e.g. n_particles.""" + + images_diagnostics: bool = True + """Whether this model publishes full screen-image PVs for its diagnostics. + + Two stages that both image the handoff diagnostic would publish the same PVs + and be rejected by StagedModel, so the downstream stage has to start just + past it. Surrogates publish only scalars (XRMS/YRMS) and so never collide. + """ + + element_after: dict[str, str] = field(default_factory=dict) + """Diagnostic -> the element immediately downstream of it. + + Only needed for entries that can be a downstream stage, and only for + diagnostics used as handoff points. Bmad's ``s`` is the exit face, so + slicing from ``DL02A2`` begins at exactly YAG03's position while excluding + the screen itself. + """ + + element_aliases: dict[str, str] = field(default_factory=dict) + """Standard name -> this engine's local name. + + Empty for every current entry: diagnostic names already agree between Bmad + and IMPACT in both facilities. Kept as the place a future rename would go. + """ + + @property + def configurable_extent(self) -> bool: + return self.start_param is not None or self.end_param is not None + + +_ALL_CU_HXR_SCREENS = ( + "YAG02", + "YAG03", + "OTRH1", + "OTRH2", + "OTR1", + "OTR2", + "OTR3", + "OTR4", + "OTR11", + "OTR12", + "OTR21", + "OTRDMP", +) + +MODELS: dict[str, ModelEntry] = { + "impact_cu_inj": ModelEntry( + name="impact_cu_inj", + description="IMPACT-T LCLS injector, cathode -> OTR2", + facility="lcls", + engine="impact", + builder="virtual_accelerator.models.cu_hxr:get_cu_inj_impact_model", + extras=("impact",), + params={"n_particles": 100, "end_element": "OTR2"}, + # YAG01 and OTR3 exist in the deck but their lines are commented out; + # OTR4 is past stop_1 at z=16.5. + diagnostics=("YAG02", "YAG03", "OTR1", "OTR2"), + end_param="end_element", + default_end="OTR2", + broadcast_params=frozenset({"n_particles"}), + ), + "bmad_cu_hxr": ModelEntry( + name="bmad_cu_hxr", + description="Bmad CU-HXR, gun -> undulator/dump", + facility="lcls", + engine="bmad", + builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model", + extras=("bmad",), + params={ + "start_element": "OTR2", + "end_element": "END", + "track_beam": False, + "custom_beam_path": None, + }, + diagnostics=_ALL_CU_HXR_SCREENS, + start_param="start_element", + end_param="end_element", + default_start="OTR2", + default_end="END", + # Verified against the lattice: each value's entrance face sits exactly at + # the screen's s. OTR11/OTR21 are omitted because both are followed by an + # element named DDG4, which is ambiguous and so unusable as a slice start. + element_after={ + "YAG02": "DL01G", + "YAG03": "DL02A2", + "OTRH1": "DH03A", + "OTRH2": "DH02B", + "OTR1": "DE05C", + "OTR2": "DE06D", + "OTR3": "DE07", + "OTR4": "DB00B", + }, + ), + "surrogate_cu_inj": ModelEntry( + name="surrogate_cu_inj", + description="NN LCLS injector surrogate, fixed cathode -> OTR2", + facility="lcls", + engine="surrogate", + builder=( + "virtual_accelerator.models.cu_hxr:get_cu_hxr_injector_surrogate_model" + ), + extras=("surrogate",), + params={"n_particles": 1000}, + diagnostics=("OTR2",), + default_end="OTR2", + broadcast_params=frozenset({"n_particles"}), + images_diagnostics=False, + ), + "cheetah_cu_hxr": ModelEntry( + name="cheetah_cu_hxr", + description="Cheetah nc_hxr", + facility="lcls", + engine="cheetah", + builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_cheetah_model", + extras=("cheetah",), + params={"n_particles": 1000}, + diagnostics=(), + broadcast_params=frozenset({"n_particles"}), + ), +} diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py new file mode 100644 index 0000000..92333e9 --- /dev/null +++ b/virtual_accelerator/tests/test_registry.py @@ -0,0 +1,195 @@ +"""Registry tests that need no engine dependencies or lattice checkouts.""" + +import pytest + +from virtual_accelerator.registry import ( + _downstream_start, + _normalize, + _resolve_handoffs, + _route_kwargs, + get_model, + list_diagnostics, + list_models, + models_available, +) +from virtual_accelerator.registry.models import MODELS + + +class TestDiscovery: + def test_all_entries_listed(self): + assert set(models_available) == set(MODELS) + + def test_repr_is_aligned_table(self): + text = repr(models_available) + assert "impact_cu_inj" in text + assert len(text.splitlines()) == len(MODELS) + + def test_filter_by_engine_and_facility(self): + assert list_models(engine="bmad") == ["bmad_cu_hxr"] + assert set(list_models(facility="lcls")) == set(MODELS) + assert list_models(facility="facet2") == [] + + def test_diagnostics_are_lattice_ordered(self): + diags = list_diagnostics("bmad_cu_hxr") + assert diags.index("YAG02") < diags.index("YAG03") < diags.index("OTR2") + + def test_impact_omits_unavailable_screens(self): + # YAG01/OTR3 are commented out in the deck; OTR4 is past stop_1. + diags = list_diagnostics("impact_cu_inj") + assert "OTR2" in diags + for absent in ("YAG01", "OTR3", "OTR4"): + assert absent not in diags + + +class TestEntryIntegrity: + @pytest.mark.parametrize("name", sorted(MODELS)) + def test_builder_is_importable_path(self, name): + module_path, sep, func = MODELS[name].builder.partition(":") + assert sep and module_path.startswith("virtual_accelerator.") and func + + @pytest.mark.parametrize("name", sorted(MODELS)) + def test_extent_params_are_declared(self, name): + entry = MODELS[name] + for param in (entry.start_param, entry.end_param): + if param is not None: + assert param in entry.params + + @pytest.mark.parametrize("name", sorted(MODELS)) + def test_broadcast_params_are_declared(self, name): + entry = MODELS[name] + assert entry.broadcast_params <= set(entry.params) + + @pytest.mark.parametrize("name", sorted(MODELS)) + def test_element_after_keys_are_diagnostics(self, name): + entry = MODELS[name] + assert set(entry.element_after) <= set(entry.diagnostics) + + @pytest.mark.parametrize("name", sorted(MODELS)) + def test_defaults_are_consistent(self, name): + entry = MODELS[name] + if entry.default_start and entry.start_param: + assert entry.params[entry.start_param] == entry.default_start + if entry.default_end and entry.end_param: + assert entry.params[entry.end_param] == entry.default_end + + +class TestValidation: + def test_unknown_model(self): + with pytest.raises(KeyError, match="Unknown model"): + get_model("bmad_does_not_exist") + + def test_rejects_unavailable_diagnostic(self): + with pytest.raises(ValueError, match="not an available end diagnostic"): + get_model("impact_cu_inj", end_ele="OTR4") + + def test_rejects_impact_as_downstream_stage(self): + with pytest.raises(ValueError, match="only start at the cathode"): + get_model(["bmad_cu_hxr", "impact_cu_inj"], handoff_loc="OTR2") + + def test_rejects_start_ele_on_fixed_extent_model(self): + with pytest.raises(ValueError, match="fixed start"): + get_model("surrogate_cu_inj", start_ele="OTR2") + + def test_rejects_single_model_list(self): + with pytest.raises(ValueError, match="at least two models"): + get_model(["bmad_cu_hxr"]) + + def test_rejects_wrong_handoff_count(self): + with pytest.raises(ValueError, match="handoff location"): + get_model(["surrogate_cu_inj", "bmad_cu_hxr"], handoff_loc=["OTR2", "OTR3"]) + + +class TestKwargRouting: + def test_unknown_kwarg_rejected(self): + with pytest.raises(ValueError, match="not a parameter of any stage"): + _route_kwargs([MODELS["bmad_cu_hxr"]], {"n_particle": 5}) + + def test_broadcast_reaches_every_declaring_stage(self): + entries = [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]] + routed = _route_kwargs(entries, {"n_particles": 42}) + assert routed == [{"n_particles": 42}, {"n_particles": 42}] + + def test_routes_to_single_declaring_stage(self): + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] + routed = _route_kwargs(entries, {"track_beam": True}) + assert routed == [{}, {"track_beam": True}] + + def test_dotted_form_targets_one_stage(self): + entries = [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]] + routed = _route_kwargs(entries, {"cheetah_cu_hxr.n_particles": 7}) + assert routed == [{}, {"n_particles": 7}] + + def test_dotted_form_rejects_unknown_stage(self): + with pytest.raises(ValueError, match="not in this model"): + _route_kwargs([MODELS["bmad_cu_hxr"]], {"nope.track_beam": True}) + + def test_dotted_form_rejects_unknown_param(self): + with pytest.raises(ValueError, match="not a parameter of"): + _route_kwargs([MODELS["bmad_cu_hxr"]], {"bmad_cu_hxr.bogus": 1}) + + +class TestHandoffResolution: + def test_inferred_from_upstream_fixed_end(self): + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] + assert _resolve_handoffs(entries, None) == ["OTR2"] + + def test_scalar_upstream_hands_off_at_the_diagnostic(self): + # The surrogate publishes XRMS/YRMS only, so no PV collision. + start = _downstream_start( + MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"], "OTR2" + ) + assert start == "OTR2" + + @pytest.mark.parametrize( + ("handoff", "expected"), [("YAG03", "DL02A2"), ("OTR2", "DE06D")] + ) + def test_imaging_upstream_starts_after_the_diagnostic(self, handoff, expected): + # Both stages image the screen, so the downstream stage must skip it or + # StagedModel would reject the pair on duplicate PVs. + start = _downstream_start( + MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"], handoff + ) + assert start == expected + + def test_unrecorded_handoff_gives_actionable_error(self): + entry = MODELS["bmad_cu_hxr"] + stripped = type(entry)(**{**entry.__dict__, "element_after": {}}) + with pytest.raises(ValueError, match="element_after"): + _downstream_start(MODELS["impact_cu_inj"], stripped, "OTR2") + + +class TestElementNameCase: + """Element names are normalised at the API boundary. + + Tao is case-insensitive so lower case would appear to work, but IMPACT's + impact.ele[...] is a dict lookup and the registry's own diagnostics and + element_after lookups would silently miss. + """ + + @pytest.mark.parametrize("given", ["OTR4", "otr4", "Otr4", "oTr4"]) + def test_normalize_is_idempotent_upper(self, given): + assert _normalize(given) == "OTR4" + + def test_normalize_passes_none_through(self): + assert _normalize(None) is None + + def test_lowercase_bad_diagnostic_is_still_rejected(self): + # Before normalisation this slipped past validation and failed later + # inside Tao with a far worse message. + with pytest.raises(ValueError, match="not an available end diagnostic"): + get_model("impact_cu_inj", end_ele="otr99") + + def test_lowercase_valid_diagnostic_is_accepted(self): + # Reaches the builder (and fails only because the extra is absent here), + # proving validation no longer rejects it. + with pytest.raises((ImportError, ValueError)) as excinfo: + get_model("impact_cu_inj", end_ele="yag03") + assert "not an available" not in str(excinfo.value) + + def test_lowercase_handoff_still_skips_the_screen(self): + # Mirrors what get_model does: normalise, then resolve the handoff. + handoff = _normalize("yag03") + start = _downstream_start( + MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"], handoff + ) + assert start == "DL02A2" From a5af25dc32afd025dc35a0aac37dd52379226ea6 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Wed, 2 Sep 2026 11:39:26 -0700 Subject: [PATCH 02/15] Addressing comments --- docs/model_registry_design.md | 115 ++++++------ docs/registry_decisions_and_questions.md | 194 +++++++++++++++++++++ virtual_accelerator/registry/__init__.py | 113 +++++++----- virtual_accelerator/registry/models.py | 60 +++---- virtual_accelerator/tests/test_registry.py | 134 +++++++++----- 5 files changed, 435 insertions(+), 181 deletions(-) create mode 100644 docs/registry_decisions_and_questions.md diff --git a/docs/model_registry_design.md b/docs/model_registry_design.md index 639205b..1ea0d9f 100644 --- a/docs/model_registry_design.md +++ b/docs/model_registry_design.md @@ -6,10 +6,12 @@ registered yet; §7 (screen PV prefixes) should land before they are. Goal: one entry point for every model, a registry of what exists and how to configure it, and enough per-model metadata to stitch stages together safely. -Scope decision driving this revision: **handoff points are diagnostics only.** That single -constraint removes the need for an element catalog, because diagnostics are the one category of -element whose names already agree across engines (see §1). Each model carries a plain ordered list -of the diagnostics it exposes, plus an optional name-override dict for future divergence. +Scope: **any lattice element may be a start/end point**, with `CATHODE` as the canonical name for +models starting at the front of the machine. No element catalog is needed, because the only element +names that differ between engines are the gun and the RF cavity segments (see §1), so a small +per-model alias dict covers it. Screens are enumerated per model as a discovery aid and a typo check. + +Reference plane is the element **entrance** — see §5. --- @@ -40,14 +42,14 @@ handoff points: | gun / cathode | `CATHODE` / `CATHODEF` | `GUN` / `GUNF` | | L0A cavity | `L0A` (one lcavity) | `L0A_entrance`, `L0A_body_1`, `L0A_body_2`, `L0A_exit` | -So **the name-override dict is empty for both facilities today.** It exists so a future rename has -somewhere to live, not because it is currently load-bearing. +So the alias dict carries only the gun (`CATHODE` → IMPACT `GUN`). Cavity segmentation is one-to-many +and cavities are not sensible handoff planes, so they are simply not registered as handoff points. -Two facts that don't need machinery but should be written down, because they are the reason not to -validate handoffs numerically: +Two facts that don't need machinery but should be written down: -- IMPACT z is element *entrance*; the lattice CSV's `SumL` is element *centre*. They coincide only - for zero-length monitors — which is another reason to restrict handoffs to diagnostics. +- IMPACT z and Bmad slice starts are both the element *entrance*; the lattice CSV's `SumL` is the + element *centre*. For a quad these differ by half the length (FACET `QA10361`: 4.3103893 vs + 4.4123893). This is why the entrance is the reference plane and the CSV's `SumL` is not used for it. - FACET IMPACT geometry disagrees with the Bmad geometry by up to ~3.9 mm at `PR10571`, growing downstream. LCLS agrees exactly. Real, but a physics caveat rather than something the registry should police. @@ -70,15 +72,12 @@ class ModelEntry: params: dict[str, Any] # param name -> default, for validation + discovery broadcast_params: frozenset[str] # params safe to send to every stage, e.g. {"n_particles"} - diagnostics: tuple[str, ...] # standard names, ORDERED by lattice position + handoff_points: tuple[str, ...] # suggested elements, ORDERED by lattice position start_param: str | None # builder kwarg for start, None if fixed end_param: str | None # builder kwarg for end, None if fixed default_start: str | None default_end: str | None - images_diagnostics: bool # publishes full screen-image PVs? False for surrogates - element_after: dict[str, str] # diagnostic -> element immediately downstream of it - element_aliases: dict[str, str] = field(default_factory=dict) # standard name -> this engine's local name. Empty today; escape hatch only. ``` @@ -87,7 +86,8 @@ class ModelEntry: information (`None` means the extent is fixed) while also recording what each builder *calls* the parameter, which differs between engines. -`images_diagnostics` and `element_after` exist because of §5's collision rule — see there for why. +`handoff_points` is a discovery aid and typo check, **not** a restriction — any lattice element may +be used as a start/end point. See §5 for how duplicate PVs at the handoff are resolved. Three choices worth calling out: @@ -96,8 +96,8 @@ Three choices worth calling out: registry would force importing every engine module just to import the registry, breaking `models_available` for anyone without all extras installed. -**`diagnostics` is ordered by lattice position.** Ordering is then checkable by list index instead -of by metres, which is what lets §4 drop positions without losing the ordering check. +**`handoff_points` is ordered by lattice position.** Ordering is then checkable by list index +instead of by metres, which is what lets the design drop positions entirely. **No data files, no generator script.** The lists are hand-maintained Python literals in `registry/models.py`. They are short, they change only when a lattice model changes, and a diff is @@ -105,7 +105,7 @@ readable. Revisit only if they start drifting from the lattice. ``` virtual_accelerator/registry/ -├── __init__.py # get_model, models_available, list_diagnostics +├── __init__.py # get_model, models_available, list_handoff_points └── models.py # the ModelEntry table ``` @@ -113,18 +113,18 @@ virtual_accelerator/registry/ ## 3. Model catalog -| registry name | facility | engine | builder | key params | diagnostics (ordered) | +| registry name | facility | engine | builder | key params | suggested handoff points | |---|---|---|---|---|---| -| `impact_cu_inj` | lcls | impact | `get_cu_inj_impact_model` | `n_particles=100`, `end_element="OTR2"` | YAG02, YAG03, OTR1, OTR2 | -| `bmad_cu_hxr` | lcls | bmad | `get_cu_hxr_bmad_model` | `start_element="OTR2"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | all 12 profmon screens, YAG02 → OTRDMP | -| `surrogate_cu_inj` | lcls | surrogate | `get_cu_hxr_injector_surrogate_model` | `n_particles=1000` | OTR2 (fixed end) | +| `impact_cu_inj` | lcls | impact | `get_cu_inj_impact_model` | `n_particles=100`, `end_element="OTR2"` | CATHODE, YAG02, YAG03, OTR1, OTR2 | +| `bmad_cu_hxr` | lcls | bmad | `get_cu_hxr_bmad_model` | `start_element="OTR2"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | CATHODE, all 12 profmon screens, END | +| `surrogate_cu_inj` | lcls | surrogate | `get_cu_hxr_injector_surrogate_model` | `n_particles=1000` | CATHODE, OTR2 (fixed extent) | | `cheetah_cu_hxr` | lcls | cheetah | `get_cu_hxr_cheetah_model` | `n_particles=1000` | — | Naming convention: `_`. **Not yet registered** — FACET-II, pending §7: -| registry name | facility | engine | builder | key params | diagnostics (ordered) | +| registry name | facility | engine | builder | key params | suggested handoff points | |---|---|---|---|---|---| | `impact_f2e_inj` | facet2 | impact | `get_facet_impact_model` | `n_particles=100`, `end_element="PR10571"` | PR10241, PR10465, PR10471, PR10571 | | `bmad_f2_elec` | facet2 | bmad | `get_facet_bmad_model` | `start_element="L0AFEND"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | PR10241, PR10465, PR10471, PR10571, PR10711 | @@ -139,11 +139,10 @@ Notes on the lists: `ImpactT_template.in` (per `ImpactT.yaml`'s `input_file:` key). The checked-in `ImpactT.in` is a stale truncated artifact stopping at z=12.0 and omitting it — reading that file gives the wrong answer about available stop points. Worth an upstream issue. -- `bmad_f2_elec` currently defaults `start_element="L0AFEND"`, which is a marker rather than a - diagnostic. `L0AFEND` exists in both engines (Bmad superimposed marker, IMPACT write-beam - element) so it's a legitimate handoff plane; treat markers as admissible where both engines have - them, and keep them in `diagnostics` despite the field name. (Or rename the field - `handoff_points` — probably clearer.) +- `bmad_f2_elec` defaults `start_element="L0AFEND"`, a marker rather than a screen. It exists in + both engines (Bmad superimposed marker, IMPACT write-beam element) so it is a legitimate handoff + plane. Since any element is now allowed, markers need no special treatment — they are simply listed + in `handoff_points` where useful. **Prerequisite:** `surrogate_f2e_inj` has no standalone builder — the `BeamOutputModel` is built inline inside `get_facet_staged_model`. It needs extracting to @@ -195,14 +194,14 @@ stage's end. Interior extents come from `handoff_loc`. ### Discovery ```python -from virtual_accelerator.registry import models_available, list_diagnostics +from virtual_accelerator.registry import models_available, list_handoff_points print(models_available) # impact_cu_inj IMPACT-T LCLS injector, cathode -> OTR2 # bmad_cu_hxr Bmad CU-HXR, gun -> undulator/dump # ... -list_diagnostics("bmad_cu_hxr") +list_handoff_points("bmad_cu_hxr") # ('YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', # 'OTR11', 'OTR12', 'OTR21', 'OTRDMP') ``` @@ -246,35 +245,33 @@ All checks are name-based. No positions involved. which is what makes `get_model(["bmad_cu_hxr", "impact_cu_inj"], ...)` fail with "IMPACT models can only start at the cathode" instead of something inscrutable from inside `set_stop_location`. -**C3 — the downstream stage must not re-image the handoff diagnostic.** This one was found by -testing and is the most important rule in practice. - -IMPACT's `set_stop_location` prunes to `s <= stop`, so a model stopped at `YAG03` *keeps* `YAG03` -and publishes its six `YAGS:IN20:351:*` image PVs. Slicing Bmad from `YAG03` also includes the -screen and publishes the same six PVs. `StagedModel` then rejects the pair on duplicate variables — -after paying for a full IMPACT run. - -Measured on the real lattice: - -| Bmad `start_element` | n_vars | `YAGS:IN20:351:*` PVs | -|---|---|---| -| `YAG03` | 292 | 6 | -| `DL02A2` | 286 | 0 | - -So `DL02A2` in `examples/staged_example.ipynb` was **not** arbitrary and not merely a naming -inconsistency — it is a deliberate workaround for this collision. An earlier draft of this document -mischaracterised it as a rename opportunity; that was wrong. - -The rule: if the upstream and downstream stages both image the handoff diagnostic, the downstream -stage starts at `element_after[handoff]` instead. Surrogates set `images_diagnostics=False` (they -publish only `XRMS`/`YRMS` scalars), so surrogate → Bmad hands off *at* the diagnostic and needs no -skip — which is why the existing `get_cu_hxr_staged_model` works with `start_element="OTR2"`. - -`element_after` was generated from the lattice and each value verified so that its entrance face -sits exactly at the screen's `s` (Bmad's `s` is the exit face, which is why the drift *after* a -screen begins at the screen). `OTR11` and `OTR21` are omitted because both are followed by an -element named `DDG4`, which is ambiguous in the lattice and so unusable as a slice start; the error -message names the field to edit if anyone hits that case. +**C3 — duplicate PVs at the handoff are removed from the downstream stage.** This is the most +important rule in practice, and it was revised on 2026-09-02. + +Both stages include the handoff element, so both publish its PVs. IMPACT's `set_stop_location` +prunes to `s <= stop`, so a model stopped at `YAG03` *keeps* the screen and publishes its six +`YAGS:IN20:351:*` image PVs; Bmad sliced from `YAG03` publishes the same six. `StagedModel` then +rejects the pair as duplicates — after paying for a full IMPACT run. + +Resolution: compute the overlap after construction and unregister it from the **downstream** stage. +The upstream stage owns those PVs because it is the stage that actually tracks the beam to that +plane. `lume.actions` exposes `unregister_action_variable(name)` and `supported_variables` is a +property over `_action_variable_by_name`, so this is clean public API. Measured on a real Bmad model: +318 vars → 312 after removal, model still functional. + +This lets the handoff be named for the real element (`YAG03`) rather than the drift after it. An +earlier draft instead moved the downstream start to `element_after[handoff]` (`DL02A2`); that dict is +gone. It also explains `start_element="DL02A2"` in `examples/staged_example.ipynb` — a workaround for +this collision, no longer needed. + +**Safety rule:** a *writable* overlap raises rather than being dropped. Writable overlap means both +stages are driving the same magnet — extents overlapping rather than meeting at a plane — and +dropping it downstream would leave that stage tracking with a stale value. The real YAG03 overlap was +measured to be entirely read-only, so this does not fire in normal use. + +**Caveat:** the only staged path runnable without the IMPACT-T binary does not exercise this. +Measured overlap between `surrogate_cu_inj` and `bmad_cu_hxr` is **zero** — the surrogate publishes +`OTRS:IN20:571:XRMS`/`YRMS` while Bmad publishes `Image:*`/`RESOLUTION`/`X`. **C4 — beam handoff mechanics.** One place, in the registry, resolving an existing inconsistency: `get_facet_staged_model` writes `final_particles` to a `NamedTemporaryFile` and passes it as diff --git a/docs/registry_decisions_and_questions.md b/docs/registry_decisions_and_questions.md new file mode 100644 index 0000000..e827248 --- /dev/null +++ b/docs/registry_decisions_and_questions.md @@ -0,0 +1,194 @@ +# Model registry — design decisions and open questions + +Companion to `docs/model_registry_design.md`. LCLS registry is implemented and verified against +real models; FACET-II is not registered yet. + +--- + +## Part A — Design decisions + +### A1. Any element may be a start/end point; screens are enumerated as a hint + +**Revised 2026-09-02 per supervisor.** An earlier draft restricted handoffs to diagnostics. That is +lifted: any element in the lattice may be used, with `CATHODE` as the canonical name for models +starting at the front of the machine. + +Screens are still listed per model (`handoff_points`) for two reasons: discovery, and as a typo +check — screens are enumerated exhaustively, so a screen-shaped name absent from the list is +definitely wrong and is rejected early rather than failing deep inside Tao. Anything else passes +through to the engine. There are 3323 elements in `cu_hxr`, so exhaustive validation is not on the +table. + +`CATHODE` is the first entry that makes `element_aliases` non-empty: Bmad calls it `CATHODE`, IMPACT +calls it `GUN`. + +### A2. Reference plane is the element ENTRANCE + +**A midpoint default was considered and rejected on evidence.** Bmad's `-slice_lattice` begins at an +element's *entrance* and cannot begin at a midpoint: + +| element | L | entrance | centre | Bmad slice begins at | +|---|---|---|---|---| +| `QE01` | 0.108 | 8.440049 | 8.494049 | **8.440049** | +| `L0A` | 3.095 | 1.459000 | 3.006622 | **1.459000** | + +IMPACT's `impact.ele[name]["s"]` is also the entrance, though IMPACT *could* express a midpoint +since `impact.stop` is a float. So a midpoint default would be implementable in IMPACT but not in +Bmad, making the two engines silently disagree by half an element length — the exact class of bug +this design is meant to prevent. Entrance is the only plane both engines express identically. + +The separate issue of identical elements sitting at slightly different `s` in the two lattices +(FACET differs by up to ~3.9 mm, growing downstream; LCLS agrees exactly) is unaffected by this +choice and remains unpoliced — see A3. + +### A2b. Duplicate PVs at the handoff are unregistered from the downstream stage + +**Revised 2026-09-02 per supervisor.** An earlier draft moved the downstream start to the *next* +element (an `element_after` dict) to dodge the collision. That dict is now gone. + +Both stages include the handoff element, so both publish its PVs and `StagedModel` would reject the +pair as duplicates. Instead the overlap is computed after construction and removed from the +downstream stage, which lets the handoff be expressed as the real element name (`YAG03`, not +`DL02A2`). `lume.actions` provides `unregister_action_variable(name)` and `supported_variables` is a +property over `_action_variable_by_name`, so this is clean public API. Verified on a real Bmad model: +318 vars → 312 after removing the six `YAGS:IN20:351:*` PVs, model still functional. + +The upstream stage owns the handoff PVs because it is the stage that actually tracks the beam to +that plane. + +**One safety rule added:** a *writable* overlap raises instead of being dropped. Writable overlap +means both stages are driving the same magnet — the extents overlap rather than meeting at a plane — +and silently dropping it downstream would leave that stage tracking with a stale value. Measured +that the real YAG03 overlap is entirely read-only, so the rule does not fire in normal use. + +### A3. No positions stored — compatibility is name equality plus list-index ordering + +`handoff_points` lists are ordered by lattice position, so ordering is checkable by index without +storing metres. + +**Consequence:** the FACET IMPACT geometry disagrees with the Bmad geometry by up to ~3.9 mm at +`PR10571`, growing downstream (LCLS agrees exactly, to nine decimals). This is documented as a +physics caveat and deliberately **not** policed by the registry. Trigger for revisiting is recorded +in the design doc's "deliberately deferred" table. + +### A4. Builders referenced as `"module:function"` strings, resolved lazily + +The existing builders import `pytao` / `torch` / `impact` inside their function bodies via +`import_optional`. Holding real callables in the registry would force importing every engine module +just to import the registry. As strings, `models_available` and `list_handoff_points` work with **zero +optional dependencies installed** — verified. + +### A5. Registry declares each model's params; unknown kwargs are rejected + +Kwarg routing across stages is a table lookup, not signature introspection, so error messages can +name the candidate stages. Trade-off: adding a parameter to a builder means also adding it to the +registry. That cost buys real error messages and a truthful `models_available`. + +Routing rules: broadcast params (`n_particles`) go to every stage declaring them; a param declared +by exactly one stage routes there; declared by more than one and not broadcast is an error naming +the candidates; `"."` always wins. + +### A6. Element names normalised to upper case at the API boundary + +Found by testing. Without it, `handoff_loc='yag03'` silently bypassed the A2b collision check and +started Bmad *at* the screen, resurfacing the duplicate-PV failure. Also fixes IMPACT, where +`impact.ele[...]` is a case-sensitive dict lookup. + +### A7. `track_beam=True` forced on every stage in a chain + +An earlier draft said non-terminal stages only. Wrong, and caught by testing: a non-final stage must +*produce* `final_particles`, but a non-first stage must also *accept* `initial_particles`, and +`lume_bmad` raises `Cannot set initial_particles when track_type is not 'beam'`. In a two-stage +surrogate → Bmad chain the Bmad stage is terminal and still needs it. + +### A8. Hand-maintained Python literals — no generated data files, no CI regenerator + +Per the "start simple" steer. The lists are short and change only when a lattice model changes, so a +diff is readable. Revisit if they drift from the lattice. + +### A9. Registry naming convention `_` + +`impact_cu_inj`, `bmad_cu_hxr`, `surrogate_cu_inj`, `cheetah_cu_hxr`. See Q6. + +### A10. Existing builders keep working; migration is deferred + +`get_cu_hxr_staged_model` etc. are untouched so far. Proven equivalent to the registry path: +identical variable sets, identical Bmad slice, and beam moments differing by 1.6e-05 relative — +less than the 4.6e-05 run-to-run variation of the existing builder against *itself* (distgen +sampling noise). Migration to thin wrappers is queued behind end-to-end verification of the +IMPACT → Bmad path. + +--- + +## Part B — Questions + +### Blocking now + +**Q1. `TCY10490 → KLYS:LI10:51` — deliberate, like the `PROF:` overrides?** +`models/facet2.py` overrides this alias, but the FACET lattice says `TCY10490[alias]=TCAV:IN10:490` +and the elements CSV agrees. This is the same shape of question as the screen-prefix one, but for a +TCAV rather than a screen. *Blocks the FACET PV fix, since that PR touches `custom_aliases`.* + +**Q2. ANSWERED 2026-09-02.** `PROF:` applies to both engines, and more importantly: *"the csv +file usually doesn't have correct values"*. So `lcls_elements.csv` is **not** authoritative for PVs. +That confirms the fix — make `utils/*_profmon_info.yaml` the single source of truth for screen PVs +for both Bmad and IMPACT. It already holds correct values for both facilities, in a `name:` field +that neither engine currently reads. LCLS is a no-op. Note this narrows the CSV's role to element +*names* and `SumL`; its `Control System Name` column should not be trusted. + +**Q3. ANSWERED 2026-09-02.** `element_after` is gone — replaced by unregistering overlapping +variables from the downstream stage (A2b). Simpler and more general. + +### Naming and API — cheap now, expensive later + +**Q4. Field name: `diagnostics` or `handoff_points`?** +`L0AFEND` is a marker, not a screen, so `diagnostics` is slightly dishonest (A1 caveat). + +**Q5. Which stage should own the handoff screen?** +I gave it to the **upstream** stage, on the reasoning that it physically images the screen as its +last element. The alternative is to have IMPACT stop just *before* the screen and let the +downstream Bmad model own it. That changes which model reports the handoff-plane measurement, so it +is a physics/operations call rather than a code one. + +**Q6. Are `impact_cu_inj` / `bmad_cu_hxr` / `surrogate_cu_inj` / `cheetah_cu_hxr` the names we want +users typing?** These become the public interface and are awkward to change once notebooks and +scripts use them. + +**Q7. Should `get_model` accept PVs as `start_ele` / `end_ele`** (e.g. `"OTRS:IN20:571"` as well as +`"OTR2"`)? Cheap to add; plausibly what a control-room user reaches for first. + +### Scope + +**Q8. `runners.py` CLI back-compatibility.** Routing it through the registry changes the accepted +`--model` values (`cu_hxr_bmad` → `bmad_cu_hxr`, etc.). Alias the old names, or make a clean break? + +**Q9. What else should be registered beyond LCLS cu_hxr and FACET-II?** +`examples/cheetah_diag0_model.ipynb` builds a diag0 Cheetah model inline with no factory function, +and a `models/sc_diag0.py` exists in stale build artifacts but not in current source. Is diag0 +in scope? Are LCLS-II / nc_sxr models expected? This determines whether the flat +`_` naming holds up. + +**Q10. Who owns keeping the diagnostics lists in sync with the lattice?** +A8 chose hand-maintained literals. If upstream renames or adds a screen, nothing detects it +automatically. Acceptable, or should there be a CI check that reconciles the lists against +`$LCLS_LATTICE`? + +--- + +## Verification status + +| path | status | +|---|---| +| Discovery / validation / routing / overlap logic, no extras installed | 51 unit tests, ~1 s | +| `bmad_cu_hxr` single model | built on real lattice, correct slice and PVs | +| `surrogate_cu_inj` -> `bmad_cu_hxr` staged | built, quad set, beam propagated, real OTR4 image | +| Equivalence with `get_cu_hxr_staged_model` | identical vars and slice; within sampling noise | +| Removing overlapping variables, on a real model | 318 -> 312 vars, model still functional | +| `impact_cu_inj` -> `bmad_cu_hxr` staged | **NOT verified** -- needs the IMPACT-T executable | + +**The last row matters more after this revision.** Removing overlapping variables is now the core +handoff mechanism, and the only staged path runnable here does not exercise it: measured overlap +between `surrogate_cu_inj` and `bmad_cu_hxr` is **zero**, because the surrogate publishes +`OTRS:IN20:571:XRMS`/`YRMS` while Bmad publishes `Image:*`/`RESOLUTION`/`X`. So the mechanism is +covered by unit tests and by a direct measurement on a Bmad model, but has never run inside a real +two-stage build. Closing that needs `conda install -c conda-forge impact-t`. diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index 8c058ec..1b0c625 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -8,11 +8,14 @@ """ import importlib +import logging from typing import Any from virtual_accelerator.registry.models import MODELS, ModelEntry -__all__ = ["get_model", "models_available", "list_models", "list_diagnostics"] +logger = logging.getLogger(__name__) + +__all__ = ["get_model", "models_available", "list_models", "list_handoff_points"] class _ModelCatalog(dict): @@ -40,9 +43,12 @@ def list_models(facility: str | None = None, engine: str | None = None) -> list[ ] -def list_diagnostics(model_name: str) -> tuple[str, ...]: - """Diagnostics usable as start/end/handoff points, in lattice order.""" - return _entry(model_name).diagnostics +def list_handoff_points(model_name: str) -> tuple[str, ...]: + """Suggested start/end/handoff elements, in lattice order. + + A discovery aid, not an exhaustive list -- any lattice element may be used. + """ + return _entry(model_name).handoff_points def _entry(name: str) -> ModelEntry: @@ -73,9 +79,8 @@ def _normalize(name: str | None) -> str | None: Lattice element names are upper case everywhere. Tao is case-insensitive so a lower-case name would appear to work, but IMPACT's ``impact.ele[...]`` is a - plain dict lookup, and more importantly the registry's own diagnostics and - ``element_after`` lookups would silently miss -- which would defeat the - handoff collision check in ``_downstream_start``. + plain dict lookup, and the registry's own ``handoff_points`` and + ``element_aliases`` lookups would silently miss. """ return name if name is None else name.upper() @@ -86,19 +91,20 @@ def _resolve_element(entry: ModelEntry, name: str) -> str: def _check_element(entry: ModelEntry, name: str, role: str) -> None: - """Validate a handoff element. Non-diagnostic names are allowed through. + """Validate a start/end/handoff element. - Markers and drifts (END, TD11, DL02A2) are legitimate start/end elements but - are deliberately not enumerated, so only reject a name that looks like a - diagnostic this model does not have. + Any element in the underlying lattice is allowed, so this cannot be an + exhaustive check -- quads, markers and drifts are all legitimate and there are + thousands of them. Screens *are* enumerated exhaustively though, so a + screen-shaped name missing from ``handoff_points`` is a typo worth catching + early rather than letting it fail deep inside Tao. """ - if name in entry.diagnostics: + if name in entry.handoff_points: return - looks_like_diagnostic = name.startswith(("OTR", "YAG", "PR")) - if looks_like_diagnostic: + if name.startswith(("OTR", "YAG", "PR")): raise ValueError( - f"{name!r} is not an available {role} diagnostic for {entry.name!r}. " - f"Available: {', '.join(entry.diagnostics)}" + f"{name!r} is not an available {role} screen for {entry.name!r}. " + f"Suggested points: {', '.join(entry.handoff_points)}" ) @@ -222,28 +228,54 @@ def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) - _check_element(downstream, handoff, "start") -def _downstream_start( - upstream: ModelEntry, downstream: ModelEntry, handoff: str -) -> str: - """Where the downstream stage should actually begin tracking. +def _strip_overlapping_variables(upstream, downstream, upstream_name, downstream_name): + """Remove variables the downstream stage shares with the upstream stage. - If both stages image the handoff diagnostic they would publish identical - screen PVs and StagedModel would reject the pair, so the downstream stage - starts at the element immediately after it instead. + Both stages include the handoff element, so both publish its PVs and + ``StagedModel`` would reject the pair as duplicates. The upstream stage owns + them -- it is the stage that actually tracks the beam to that plane -- so they + are unregistered from the downstream stage. + + A *writable* overlap means something different and worse: both stages would be + driving the same magnet, and dropping it downstream would silently leave that + stage tracking with a stale value. That is a slicing error, so it raises. """ - both_image = upstream.images_diagnostics and downstream.images_diagnostics - if not (both_image and handoff in downstream.diagnostics): - return handoff + from lume.actions import WritableActionMixin - try: - return downstream.element_after[handoff] - except KeyError: + downstream_vars = downstream.supported_variables + overlap = sorted(set(upstream.supported_variables) & set(downstream_vars)) + if not overlap: + return [] + + writable = [ + name + for name in overlap + if isinstance(downstream_vars[name], WritableActionMixin) + ] + if writable: raise ValueError( - f"Both {upstream.name!r} and {downstream.name!r} image {handoff!r}, so " - f"{downstream.name!r} must start just after it, but no downstream element " - f"is recorded for {handoff!r}. Add it to {downstream.name}.element_after " - "in virtual_accelerator/registry/models.py." - ) from None + f"{upstream_name!r} and {downstream_name!r} both control " + f"{len(writable)} writable variable(s), so their extents overlap rather " + f"than meeting at a plane: {', '.join(writable[:5])}" + f"{' ...' if len(writable) > 5 else ''}. Check the handoff element." + ) + + if not hasattr(downstream, "unregister_action_variable"): + raise TypeError( + f"{downstream_name!r} shares {len(overlap)} variable(s) with " + f"{upstream_name!r} but does not support unregister_action_variable, so " + "the duplicates cannot be resolved." + ) + + for name in overlap: + downstream.unregister_action_variable(name) + logger.debug( + "Removed %d variable(s) from %s already provided by %s", + len(overlap), + downstream_name, + upstream_name, + ) + return overlap def get_model( @@ -291,11 +323,7 @@ def get_model( stages = [] for i, entry in enumerate(entries): - stage_start = ( - start_ele - if i == 0 - else _downstream_start(entries[i - 1], entry, handoffs[i - 1]) - ) + stage_start = start_ele if i == 0 else handoffs[i - 1] stage_end = end_ele if i == len(entries) - 1 else handoffs[i] stage_kwargs = dict(routed[i]) @@ -314,6 +342,13 @@ def get_model( ) ) + # Both stages include the handoff element and so publish its PVs. Resolve the + # duplicates before StagedModel validation, which would otherwise reject them. + for i in range(1, len(stages)): + _strip_overlapping_variables( + stages[i - 1], stages[i], entries[i - 1].name, entries[i].name + ) + from lume.staged_model import StagedModel return StagedModel(stages) diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index 6675d2d..c9d11fd 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -26,8 +26,19 @@ class ModelEntry: params: dict[str, Any] """Configurable parameter name -> default. Also the allow-list for kwargs.""" - diagnostics: tuple[str, ...] - """Standard-named diagnostics usable as handoff points, in lattice order.""" + handoff_points: tuple[str, ...] + """Suggested start/end/handoff elements, in lattice order. + + A discovery aid, **not** a restriction: any element name in the underlying + lattice may be used. Screens are enumerated exhaustively, so a screen-shaped + name absent from this tuple is a typo and is rejected; anything else passes + through to the engine. + + Positions refer to the **entrance** face of the element. Bmad's + ``-slice_lattice`` begins at an element's entrance and cannot begin at a + midpoint, and IMPACT's ``impact.ele[name]["s"]`` is also the entrance, so the + entrance is the only reference plane both engines express identically. + """ start_param: str | None = None """Builder kwarg controlling the start element, or None if not configurable.""" @@ -41,28 +52,12 @@ class ModelEntry: broadcast_params: frozenset[str] = frozenset() """Params safe to send to every stage of a staged model, e.g. n_particles.""" - images_diagnostics: bool = True - """Whether this model publishes full screen-image PVs for its diagnostics. - - Two stages that both image the handoff diagnostic would publish the same PVs - and be rejected by StagedModel, so the downstream stage has to start just - past it. Surrogates publish only scalars (XRMS/YRMS) and so never collide. - """ - - element_after: dict[str, str] = field(default_factory=dict) - """Diagnostic -> the element immediately downstream of it. - - Only needed for entries that can be a downstream stage, and only for - diagnostics used as handoff points. Bmad's ``s`` is the exit face, so - slicing from ``DL02A2`` begins at exactly YAG03's position while excluding - the screen itself. - """ - element_aliases: dict[str, str] = field(default_factory=dict) """Standard name -> this engine's local name. - Empty for every current entry: diagnostic names already agree between Bmad - and IMPACT in both facilities. Kept as the place a future rename would go. + Diagnostics already agree between Bmad and IMPACT, so this only carries the + handful of elements that genuinely differ -- currently just the gun, which + Bmad calls CATHODE and IMPACT calls GUN. """ @property @@ -96,10 +91,11 @@ def configurable_extent(self) -> bool: params={"n_particles": 100, "end_element": "OTR2"}, # YAG01 and OTR3 exist in the deck but their lines are commented out; # OTR4 is past stop_1 at z=16.5. - diagnostics=("YAG02", "YAG03", "OTR1", "OTR2"), + handoff_points=("CATHODE", "YAG02", "YAG03", "OTR1", "OTR2"), end_param="end_element", default_end="OTR2", broadcast_params=frozenset({"n_particles"}), + element_aliases={"CATHODE": "GUN"}, ), "bmad_cu_hxr": ModelEntry( name="bmad_cu_hxr", @@ -114,24 +110,11 @@ def configurable_extent(self) -> bool: "track_beam": False, "custom_beam_path": None, }, - diagnostics=_ALL_CU_HXR_SCREENS, + handoff_points=("CATHODE", *_ALL_CU_HXR_SCREENS, "END"), start_param="start_element", end_param="end_element", default_start="OTR2", default_end="END", - # Verified against the lattice: each value's entrance face sits exactly at - # the screen's s. OTR11/OTR21 are omitted because both are followed by an - # element named DDG4, which is ambiguous and so unusable as a slice start. - element_after={ - "YAG02": "DL01G", - "YAG03": "DL02A2", - "OTRH1": "DH03A", - "OTRH2": "DH02B", - "OTR1": "DE05C", - "OTR2": "DE06D", - "OTR3": "DE07", - "OTR4": "DB00B", - }, ), "surrogate_cu_inj": ModelEntry( name="surrogate_cu_inj", @@ -143,10 +126,9 @@ def configurable_extent(self) -> bool: ), extras=("surrogate",), params={"n_particles": 1000}, - diagnostics=("OTR2",), + handoff_points=("CATHODE", "OTR2"), default_end="OTR2", broadcast_params=frozenset({"n_particles"}), - images_diagnostics=False, ), "cheetah_cu_hxr": ModelEntry( name="cheetah_cu_hxr", @@ -156,7 +138,7 @@ def configurable_extent(self) -> bool: builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_cheetah_model", extras=("cheetah",), params={"n_particles": 1000}, - diagnostics=(), + handoff_points=(), broadcast_params=frozenset({"n_particles"}), ), } diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index 92333e9..4eb6c9e 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -2,13 +2,15 @@ import pytest +from lume.actions import WritableActionMixin + from virtual_accelerator.registry import ( - _downstream_start, _normalize, _resolve_handoffs, _route_kwargs, + _strip_overlapping_variables, get_model, - list_diagnostics, + list_handoff_points, list_models, models_available, ) @@ -29,13 +31,13 @@ def test_filter_by_engine_and_facility(self): assert set(list_models(facility="lcls")) == set(MODELS) assert list_models(facility="facet2") == [] - def test_diagnostics_are_lattice_ordered(self): - diags = list_diagnostics("bmad_cu_hxr") + def test_handoff_points_are_lattice_ordered(self): + diags = list_handoff_points("bmad_cu_hxr") assert diags.index("YAG02") < diags.index("YAG03") < diags.index("OTR2") def test_impact_omits_unavailable_screens(self): # YAG01/OTR3 are commented out in the deck; OTR4 is past stop_1. - diags = list_diagnostics("impact_cu_inj") + diags = list_handoff_points("impact_cu_inj") assert "OTR2" in diags for absent in ("YAG01", "OTR3", "OTR4"): assert absent not in diags @@ -60,9 +62,9 @@ def test_broadcast_params_are_declared(self, name): assert entry.broadcast_params <= set(entry.params) @pytest.mark.parametrize("name", sorted(MODELS)) - def test_element_after_keys_are_diagnostics(self, name): + def test_element_alias_keys_are_handoff_points(self, name): entry = MODELS[name] - assert set(entry.element_after) <= set(entry.diagnostics) + assert set(entry.element_aliases) <= set(entry.handoff_points) @pytest.mark.parametrize("name", sorted(MODELS)) def test_defaults_are_consistent(self, name): @@ -78,8 +80,8 @@ def test_unknown_model(self): with pytest.raises(KeyError, match="Unknown model"): get_model("bmad_does_not_exist") - def test_rejects_unavailable_diagnostic(self): - with pytest.raises(ValueError, match="not an available end diagnostic"): + def test_rejects_unavailable_screen(self): + with pytest.raises(ValueError, match="not an available end screen"): get_model("impact_cu_inj", end_ele="OTR4") def test_rejects_impact_as_downstream_stage(self): @@ -133,37 +135,84 @@ def test_inferred_from_upstream_fixed_end(self): entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] assert _resolve_handoffs(entries, None) == ["OTR2"] - def test_scalar_upstream_hands_off_at_the_diagnostic(self): - # The surrogate publishes XRMS/YRMS only, so no PV collision. - start = _downstream_start( - MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"], "OTR2" - ) - assert start == "OTR2" - - @pytest.mark.parametrize( - ("handoff", "expected"), [("YAG03", "DL02A2"), ("OTR2", "DE06D")] - ) - def test_imaging_upstream_starts_after_the_diagnostic(self, handoff, expected): - # Both stages image the screen, so the downstream stage must skip it or - # StagedModel would reject the pair on duplicate PVs. - start = _downstream_start( - MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"], handoff - ) - assert start == expected - - def test_unrecorded_handoff_gives_actionable_error(self): - entry = MODELS["bmad_cu_hxr"] - stripped = type(entry)(**{**entry.__dict__, "element_after": {}}) - with pytest.raises(ValueError, match="element_after"): - _downstream_start(MODELS["impact_cu_inj"], stripped, "OTR2") + def test_explicit_handoff_is_used_verbatim(self): + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] + assert _resolve_handoffs(entries, "YAG03") == ["YAG03"] + + +class _FakeStage: + """Minimal stand-in for a LUMEModel with registerable action variables.""" + + def __init__(self, variables): + self._vars = dict(variables) + + @property + def supported_variables(self): + return dict(self._vars) + + def unregister_action_variable(self, name): + return self._vars.pop(name) + + +class _ReadOnlyVar: + pass + + +class _WritableVar(WritableActionMixin): + def _get(self, simulator): # pragma: no cover - never invoked + raise NotImplementedError + + def _set(self, simulator, value): # pragma: no cover - never invoked + raise NotImplementedError + + +class TestOverlapRemoval: + """The handoff element belongs to both stages, so both publish its PVs. + + The upstream stage owns them since it tracks the beam to that plane, so they + are unregistered downstream rather than moving the downstream start element. + """ + + def test_shared_read_only_variables_are_removed_downstream(self): + up = _FakeStage({"SCREEN:IMAGE": _ReadOnlyVar(), "UP:ONLY": _ReadOnlyVar()}) + down = _FakeStage({"SCREEN:IMAGE": _ReadOnlyVar(), "DOWN:ONLY": _ReadOnlyVar()}) + removed = _strip_overlapping_variables(up, down, "up", "down") + assert removed == ["SCREEN:IMAGE"] + assert set(down.supported_variables) == {"DOWN:ONLY"} + # the upstream stage keeps its copy + assert "SCREEN:IMAGE" in up.supported_variables + + def test_no_overlap_is_a_no_op(self): + up = _FakeStage({"UP:ONLY": _ReadOnlyVar()}) + down = _FakeStage({"DOWN:ONLY": _ReadOnlyVar()}) + assert _strip_overlapping_variables(up, down, "up", "down") == [] + assert set(down.supported_variables) == {"DOWN:ONLY"} + + def test_writable_overlap_raises_instead_of_silently_dropping(self): + # Both stages driving the same magnet means the extents overlap rather + # than meeting at a plane; dropping it downstream would leave that stage + # tracking with a stale value. + up = _FakeStage({"QUAD:BCTRL": _WritableVar()}) + down = _FakeStage({"QUAD:BCTRL": _WritableVar()}) + with pytest.raises(ValueError, match="writable variable"): + _strip_overlapping_variables(up, down, "up", "down") + assert "QUAD:BCTRL" in down.supported_variables + + def test_stage_without_unregister_support_raises(self): + class Fixed: + supported_variables = {"SHARED": _ReadOnlyVar()} + + up = _FakeStage({"SHARED": _ReadOnlyVar()}) + with pytest.raises(TypeError, match="unregister_action_variable"): + _strip_overlapping_variables(up, Fixed(), "up", "down") class TestElementNameCase: """Element names are normalised at the API boundary. Tao is case-insensitive so lower case would appear to work, but IMPACT's - impact.ele[...] is a dict lookup and the registry's own diagnostics and - element_after lookups would silently miss. + impact.ele[...] is a dict lookup and the registry's own handoff_points and + element_aliases lookups would silently miss. """ @pytest.mark.parametrize("given", ["OTR4", "otr4", "Otr4", "oTr4"]) @@ -173,23 +222,20 @@ def test_normalize_is_idempotent_upper(self, given): def test_normalize_passes_none_through(self): assert _normalize(None) is None - def test_lowercase_bad_diagnostic_is_still_rejected(self): + def test_lowercase_bad_screen_is_still_rejected(self): # Before normalisation this slipped past validation and failed later # inside Tao with a far worse message. - with pytest.raises(ValueError, match="not an available end diagnostic"): + with pytest.raises(ValueError, match="not an available end screen"): get_model("impact_cu_inj", end_ele="otr99") - def test_lowercase_valid_diagnostic_is_accepted(self): + def test_lowercase_valid_screen_is_accepted(self): # Reaches the builder (and fails only because the extra is absent here), # proving validation no longer rejects it. with pytest.raises((ImportError, ValueError)) as excinfo: get_model("impact_cu_inj", end_ele="yag03") assert "not an available" not in str(excinfo.value) - def test_lowercase_handoff_still_skips_the_screen(self): - # Mirrors what get_model does: normalise, then resolve the handoff. - handoff = _normalize("yag03") - start = _downstream_start( - MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"], handoff - ) - assert start == "DL02A2" + def test_lowercase_handoff_normalises_before_resolution(self): + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] + handoffs = [_normalize(h) for h in _resolve_handoffs(entries, "yag03")] + assert handoffs == ["YAG03"] From d93f1702431cd6aaa7589117c37972e7dbde2712 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 3 Sep 2026 12:06:25 -0700 Subject: [PATCH 03/15] 2 available lcls bmad models for now --- virtual_accelerator/registry/__init__.py | 16 ++++++ virtual_accelerator/registry/models.py | 58 ++++++++++++++-------- virtual_accelerator/tests/test_registry.py | 35 ++++++++----- 3 files changed, 76 insertions(+), 33 deletions(-) diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index 1b0c625..c6d492e 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -227,6 +227,22 @@ def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) - _check_element(upstream, handoff, "end") _check_element(downstream, handoff, "start") + # Each linac model is registered for one standard handoff plane, so a mismatch + # means the wrong pair was chosen -- silently re-slicing would leave a gap. + if downstream.default_start is not None and handoff != downstream.default_start: + alternatives = [ + name + for name, entry in MODELS.items() + if entry.facility == downstream.facility + and entry.engine == downstream.engine + and entry.default_start == handoff + ] + hint = f" Use {alternatives[0]!r} instead." if alternatives else "" + raise ValueError( + f"{upstream.name!r} hands off at {handoff!r} but {downstream.name!r} " + f"starts at {downstream.default_start!r}.{hint}" + ) + def _strip_overlapping_variables(upstream, downstream, upstream_name, downstream_name): """Remove variables the downstream stage shares with the upstream stage. diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index c9d11fd..0e177cd 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -80,41 +80,57 @@ def configurable_extent(self) -> bool: "OTRDMP", ) + +def _bmad_cu_hxr(name: str, start: str, note: str) -> ModelEntry: + """One Bmad CU-HXR entry per standard handoff plane. + + LCLS needs two because its two injector models end at different planes and + neither can move: the NN surrogate predicts OTRS:IN20:571 (OTR2) at 135 MeV + and cannot produce a beam at YAG03, which sits before L0B at 64 MeV. + """ + return ModelEntry( + name=name, + description=f"Bmad CU-HXR linac, {start} -> END ({note})", + facility="lcls", + engine="bmad", + builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model", + extras=("bmad",), + params={ + "start_element": start, + "end_element": "END", + "track_beam": False, + "custom_beam_path": None, + }, + handoff_points=("CATHODE", *_ALL_CU_HXR_SCREENS, "END"), + start_param="start_element", + end_param="end_element", + default_start=start, + default_end="END", + ) + + MODELS: dict[str, ModelEntry] = { "impact_cu_inj": ModelEntry( name="impact_cu_inj", - description="IMPACT-T LCLS injector, cathode -> OTR2", + description="IMPACT-T LCLS injector, cathode -> YAG03", facility="lcls", engine="impact", builder="virtual_accelerator.models.cu_hxr:get_cu_inj_impact_model", extras=("impact",), - params={"n_particles": 100, "end_element": "OTR2"}, + params={"n_particles": 100, "end_element": "YAG03"}, # YAG01 and OTR3 exist in the deck but their lines are commented out; # OTR4 is past stop_1 at z=16.5. handoff_points=("CATHODE", "YAG02", "YAG03", "OTR1", "OTR2"), end_param="end_element", - default_end="OTR2", + default_end="YAG03", broadcast_params=frozenset({"n_particles"}), element_aliases={"CATHODE": "GUN"}, ), - "bmad_cu_hxr": ModelEntry( - name="bmad_cu_hxr", - description="Bmad CU-HXR, gun -> undulator/dump", - facility="lcls", - engine="bmad", - builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model", - extras=("bmad",), - params={ - "start_element": "OTR2", - "end_element": "END", - "track_beam": False, - "custom_beam_path": None, - }, - handoff_points=("CATHODE", *_ALL_CU_HXR_SCREENS, "END"), - start_param="start_element", - end_param="end_element", - default_start="OTR2", - default_end="END", + "bmad_cu_hxr_yag03": _bmad_cu_hxr( + "bmad_cu_hxr_yag03", "YAG03", "pairs with impact_cu_inj" + ), + "bmad_cu_hxr_otr2": _bmad_cu_hxr( + "bmad_cu_hxr_otr2", "OTR2", "pairs with surrogate_cu_inj" ), "surrogate_cu_inj": ModelEntry( name="surrogate_cu_inj", diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index 4eb6c9e..4c88842 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -27,12 +27,12 @@ def test_repr_is_aligned_table(self): assert len(text.splitlines()) == len(MODELS) def test_filter_by_engine_and_facility(self): - assert list_models(engine="bmad") == ["bmad_cu_hxr"] + assert list_models(engine="bmad") == ["bmad_cu_hxr_yag03", "bmad_cu_hxr_otr2"] assert set(list_models(facility="lcls")) == set(MODELS) assert list_models(facility="facet2") == [] def test_handoff_points_are_lattice_ordered(self): - diags = list_handoff_points("bmad_cu_hxr") + diags = list_handoff_points("bmad_cu_hxr_otr2") assert diags.index("YAG02") < diags.index("YAG03") < diags.index("OTR2") def test_impact_omits_unavailable_screens(self): @@ -86,7 +86,7 @@ def test_rejects_unavailable_screen(self): def test_rejects_impact_as_downstream_stage(self): with pytest.raises(ValueError, match="only start at the cathode"): - get_model(["bmad_cu_hxr", "impact_cu_inj"], handoff_loc="OTR2") + get_model(["bmad_cu_hxr_otr2", "impact_cu_inj"], handoff_loc="OTR2") def test_rejects_start_ele_on_fixed_extent_model(self): with pytest.raises(ValueError, match="fixed start"): @@ -94,17 +94,28 @@ def test_rejects_start_ele_on_fixed_extent_model(self): def test_rejects_single_model_list(self): with pytest.raises(ValueError, match="at least two models"): - get_model(["bmad_cu_hxr"]) + get_model(["bmad_cu_hxr_otr2"]) def test_rejects_wrong_handoff_count(self): with pytest.raises(ValueError, match="handoff location"): - get_model(["surrogate_cu_inj", "bmad_cu_hxr"], handoff_loc=["OTR2", "OTR3"]) + get_model( + ["surrogate_cu_inj", "bmad_cu_hxr_otr2"], handoff_loc=["OTR2", "OTR3"] + ) + + def test_rejects_mismatched_injector_and_linac_pairing(self): + # impact ends at YAG03; bmad_cu_hxr_otr2 starts at OTR2 -> 9.6 m gap. + with pytest.raises(ValueError, match="bmad_cu_hxr_yag03"): + get_model(["impact_cu_inj", "bmad_cu_hxr_otr2"]) + + def test_rejects_surrogate_paired_with_yag03_linac(self): + with pytest.raises(ValueError, match="bmad_cu_hxr_otr2"): + get_model(["surrogate_cu_inj", "bmad_cu_hxr_yag03"]) class TestKwargRouting: def test_unknown_kwarg_rejected(self): with pytest.raises(ValueError, match="not a parameter of any stage"): - _route_kwargs([MODELS["bmad_cu_hxr"]], {"n_particle": 5}) + _route_kwargs([MODELS["bmad_cu_hxr_otr2"]], {"n_particle": 5}) def test_broadcast_reaches_every_declaring_stage(self): entries = [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]] @@ -112,7 +123,7 @@ def test_broadcast_reaches_every_declaring_stage(self): assert routed == [{"n_particles": 42}, {"n_particles": 42}] def test_routes_to_single_declaring_stage(self): - entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr_otr2"]] routed = _route_kwargs(entries, {"track_beam": True}) assert routed == [{}, {"track_beam": True}] @@ -123,20 +134,20 @@ def test_dotted_form_targets_one_stage(self): def test_dotted_form_rejects_unknown_stage(self): with pytest.raises(ValueError, match="not in this model"): - _route_kwargs([MODELS["bmad_cu_hxr"]], {"nope.track_beam": True}) + _route_kwargs([MODELS["bmad_cu_hxr_otr2"]], {"nope.track_beam": True}) def test_dotted_form_rejects_unknown_param(self): with pytest.raises(ValueError, match="not a parameter of"): - _route_kwargs([MODELS["bmad_cu_hxr"]], {"bmad_cu_hxr.bogus": 1}) + _route_kwargs([MODELS["bmad_cu_hxr_otr2"]], {"bmad_cu_hxr_otr2.bogus": 1}) class TestHandoffResolution: def test_inferred_from_upstream_fixed_end(self): - entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr_otr2"]] assert _resolve_handoffs(entries, None) == ["OTR2"] def test_explicit_handoff_is_used_verbatim(self): - entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr_yag03"]] assert _resolve_handoffs(entries, "YAG03") == ["YAG03"] @@ -236,6 +247,6 @@ def test_lowercase_valid_screen_is_accepted(self): assert "not an available" not in str(excinfo.value) def test_lowercase_handoff_normalises_before_resolution(self): - entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr_yag03"]] handoffs = [_normalize(h) for h in _resolve_handoffs(entries, "yag03")] assert handoffs == ["YAG03"] From 041ff5f356a6a6beae9efb027bed4cf9b0af4ff4 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 3 Sep 2026 14:17:43 -0700 Subject: [PATCH 04/15] usage --- docs/model_registry_usage.md | 176 +++++++++++++++++++++++++++++++++++ 1 file changed, 176 insertions(+) create mode 100644 docs/model_registry_usage.md diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md new file mode 100644 index 0000000..498770b --- /dev/null +++ b/docs/model_registry_usage.md @@ -0,0 +1,176 @@ +# Virtual Accelerator (VA) Registry — README +## Overview +The virtual_accelerator.registry module provides a unified interface for loading, configuring, and chaining accelerator simulation models for the LCLS copper linac (CU). Models can be used standalone or chained together to simulate the full beamline from cathode to end. + +### Installation +```bash +pip install git+https://github.com/lume-science/lume-base.git +``` + +### Quick Start +```python +from virtual_accelerator.registry import get_model, models_available, list_handoff_points +from virtual_accelerator.registry.models import MODELS +``` + +### Available Models +Print all registered models and their descriptions: + +```python +>>> print(models_available) +impact_cu_inj IMPACT-T LCLS injector, cathode -> YAG03 +bmad_cu_hxr_yag03 Bmad CU-HXR linac, YAG03 -> END (pairs with impact_cu_inj) +bmad_cu_hxr_otr2 Bmad CU-HXR linac, OTR2 -> END (pairs with surrogate_cu_inj) +surrogate_cu_inj NN LCLS injector surrogate, fixed cathode -> OTR2 +cheetah_cu_hxr Cheetah nc_hxr +``` + +### Handoff Points +Each model exposes a set of handoff points — named screen locations where beam tracking can start or stop, and where chained models exchange beam state. + +```python +>>> for m in ["impact_cu_inj", "bmad_cu_hxr_yag03", "bmad_cu_hxr_otr2", "surrogate_cu_inj", "cheetah_cu_hxr"]: +... print(m, list_handoff_points(m)) + +impact_cu_inj ('CATHODE', 'YAG02', 'YAG03', 'OTR1', 'OTR2') +bmad_cu_hxr_yag03 ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') +bmad_cu_hxr_otr2 ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') +surrogate_cu_inj ('CATHODE', 'OTR2') +cheetah_cu_hxr () +``` + +### Element Aliases +Some handoff point names are aliases for internal model element names: + +```python +>>> print(MODELS["impact_cu_inj"].element_aliases["CATHODE"]) +GUN +``` + +### Loading a Single Model +Use get_model() with a model ID and an optional end_ele to stop tracking at a specific screen. + +```python +>>> get_model("bmad_cu_hxr_otr2", end_ele="TD11") + + +>>> get_model("impact_cu_inj", end_ele="OTR2") + +``` +Note: Unrecognized element types (e.g. solrf, Lcavity, Unknown) are skipped during model initialisation. These warnings are informational and do not prevent the model from loading. + +### Error: Unknown Model Name +Model IDs must be exact. Partial names are not supported: + +```python +>>> get_model("bmad_cu_hxr", end_ele="TD11") +KeyError: "Unknown model 'bmad_cu_hxr'. Available: bmad_cu_hxr_otr2, bmad_cu_hxr_yag03, cheetah_cu_hxr, impact_cu_inj, surrogate_cu_inj" +``` + +### Error: Invalid End Element +end_ele must be one of the model's listed handoff points: + +```python +>>> get_model("impact_cu_inj", end_ele="otr99") +ValueError: 'OTR99' is not an available end screen for 'impact_cu_inj'. +Suggested points: CATHODE, YAG02, YAG03, OTR1, OTR2 +``` + +## Staged Models +Pass a list of two model IDs to get_model() to chain an injector model into a linac model. The upstream model hands off beam particles to the downstream model at a shared handoff point. + +surrogate_cu_inj → bmad_cu_hxr_otr2 +```python +>>> m = get_model(["surrogate_cu_inj", "bmad_cu_hxr_otr2"], end_ele="OTR4", n_particles=500) + +>>> m.set({"QUAD:IN20:525:BCTRL": -10.0}) + +>>> print(m.get("OTR4_beam")["norm_emit_y"]) +5.850087235892218e-07 + +>>> print(m.get("OTRS:IN20:711:Image:ArrayData").shape) +(1040, 1392) + +>>> print([n.split("#")[0] for n in m.lume_model_instances[1].get("name")][:3]) +['BEGINNING', 'OTR2', 'DE06D'] +``` + +impact_cu_inj → bmad_cu_hxr_yag03 +When impact_cu_inj is the upstream model, use bmad_cu_hxr_yag03 (which starts at YAG03) and specify handoff_loc: + +```python +>>> model = get_model( +... ["impact_cu_inj", "bmad_cu_hxr_yag03"], +... handoff_loc="YAG03", +... end_ele="TD11", +... n_particles=1000, +... ) + +>>> model.set({"QUAD:IN20:525:BCTRL": -7.5}) + +>>> print(model.get("OTR4_beam")["norm_emit_y"]) +2.3638227838794528e-07 +``` + +### Pairing Rules +Upstream and downstream models must share a compatible handoff location. Mismatched pairs raise a descriptive error: + +```python +>>> get_model(["surrogate_cu_inj", "bmad_cu_hxr_yag03"], end_ele="TD11", n_particles=500) +ValueError: 'surrogate_cu_inj' hands off at 'OTR2' but 'bmad_cu_hxr_yag03' starts at 'YAG03'. +Use 'bmad_cu_hxr_otr2' instead. + +``` +Valid pairs: + +| Upstream | Downstream | Handoff Location |\ +| impact_cu_inj | bmad_cu_hxr_yag03 | YAG03 |\ +| surrogate_cu_inj | bmad_cu_hxr_otr2 | OTR2| + + +### Advanced: Stripping Overlapping Variables +When building chained models manually, use _strip_overlapping_variables to remove output variables from the downstream model that are already owned by the upstream model. This avoids variable conflicts. + +```python +from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model +from virtual_accelerator.registry import _strip_overlapping_variables + +bm = get_cu_hxr_bmad_model(start_element="YAG03", end_element="OTR4", track_beam=True) + +# Identify overlapping variables +yag = sorted(v for v in bm.supported_variables if "IN20:351" in v) + +# Construct a mock upstream model exposing only those variables +class Up: + supported_variables = {n: bm.supported_variables[n] for n in yag} + +# Strip them from the downstream model +removed = _strip_overlapping_variables(Up(), bm, "upstream", "bmad_cu_hxr") + +>>> print(removed) +['YAGS:IN20:351:Image:ArrayData', 'YAGS:IN20:351:Image:ArraySize0_RBV', + 'YAGS:IN20:351:Image:ArraySize1_RBV', 'YAGS:IN20:351:RESOLUTION', + 'YAGS:IN20:351:X', 'YAGS:IN20:351:Y'] + +>>> print(len(bm.supported_variables)) # 318 → 312 +312 +``` + +Note: _strip_overlapping_variables is a private utility. Prefer using get_model() with a list of model IDs, which handles this automatically. + +## API Reference +```get_model(spec, *, end_ele=None, start_ele=None, handoff_loc=None, n_particles=None)``` + +|Parameter | Type | Description |\ +|spec | str or list[str] | Model ID or [upstream, downstream] pair|\ +|end_ele | str | Screen name to stop tracking at|\ +|start_ele | str | Screen name to start tracking from|\ +|handoff_loc | str | Explicit handoff point when chaining models|\ +|n_particles | int | Number of macro-particles for beam tracking| + +```models_available``` +Printable summary of all registered models and their descriptions. + +```list_handoff_points(model_id: str) -> tuple[str, ...]``` + +Returns the available handoff point names for a given model. From fc7baad40c2bc6b35904ffdde78b1efdf075c19d Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 3 Sep 2026 15:01:56 -0700 Subject: [PATCH 05/15] usage --- docs/model_registry_usage.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index 498770b..0c42698 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -4,7 +4,7 @@ The virtual_accelerator.registry module provides a unified interface for loading ### Installation ```bash -pip install git+https://github.com/lume-science/lume-base.git +pip install git+https://github.com/slaclab/virtual-accelerator ``` ### Quick Start @@ -96,6 +96,7 @@ surrogate_cu_inj → bmad_cu_hxr_otr2 ``` impact_cu_inj → bmad_cu_hxr_yag03 + When impact_cu_inj is the upstream model, use bmad_cu_hxr_yag03 (which starts at YAG03) and specify handoff_loc: ```python From 499cf96414d22c53f03e5d68c067f43867e77bff Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 3 Sep 2026 15:03:11 -0700 Subject: [PATCH 06/15] usage --- docs/model_registry_usage.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index 0c42698..fa6601d 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -4,7 +4,7 @@ The virtual_accelerator.registry module provides a unified interface for loading ### Installation ```bash -pip install git+https://github.com/slaclab/virtual-accelerator +pip install git+https://github.com/slaclab/virtual-accelerator.git ``` ### Quick Start From 6842903d859bf7c801931d12218b112c5e620b40 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 3 Sep 2026 16:18:13 -0700 Subject: [PATCH 07/15] fixes --- docs/model_registry_usage.md | 140 +++++++++++++-------- virtual_accelerator/registry/__init__.py | 114 ++++++++++++++--- virtual_accelerator/registry/models.py | 61 ++++----- virtual_accelerator/tests/test_registry.py | 89 +++++++++---- 4 files changed, 268 insertions(+), 136 deletions(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index fa6601d..d3d0c30 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -9,7 +9,12 @@ pip install git+https://github.com/slaclab/virtual-accelerator.git ### Quick Start ```python -from virtual_accelerator.registry import get_model, models_available, list_handoff_points +from virtual_accelerator.registry import ( + get_model, + models_available, + list_handoff_points, + common_handoff_points, +) from virtual_accelerator.registry.models import MODELS ``` @@ -18,25 +23,39 @@ Print all registered models and their descriptions: ```python >>> print(models_available) -impact_cu_inj IMPACT-T LCLS injector, cathode -> YAG03 -bmad_cu_hxr_yag03 Bmad CU-HXR linac, YAG03 -> END (pairs with impact_cu_inj) -bmad_cu_hxr_otr2 Bmad CU-HXR linac, OTR2 -> END (pairs with surrogate_cu_inj) -surrogate_cu_inj NN LCLS injector surrogate, fixed cathode -> OTR2 -cheetah_cu_hxr Cheetah nc_hxr +impact_cu_inj IMPACT-T LCLS injector, cathode -> YAG03 +bmad_cu_hxr Bmad CU-HXR linac, injector handoff -> END +surrogate_cu_inj NN LCLS injector surrogate, cathode -> OTR2 +cheetah_cu_hxr Cheetah nc_hxr, cathode -> END ``` ### Handoff Points -Each model exposes a set of handoff points — named screen locations where beam tracking can start or stop, and where chained models exchange beam state. +Each model exposes a set of suggested handoff points — named locations where beam tracking can start or stop, and where chained models exchange beam state. ```python ->>> for m in ["impact_cu_inj", "bmad_cu_hxr_yag03", "bmad_cu_hxr_otr2", "surrogate_cu_inj", "cheetah_cu_hxr"]: +>>> for m in ["impact_cu_inj", "bmad_cu_hxr", "surrogate_cu_inj", "cheetah_cu_hxr"]: ... print(m, list_handoff_points(m)) -impact_cu_inj ('CATHODE', 'YAG02', 'YAG03', 'OTR1', 'OTR2') -bmad_cu_hxr_yag03 ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') -bmad_cu_hxr_otr2 ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') -surrogate_cu_inj ('CATHODE', 'OTR2') -cheetah_cu_hxr () +impact_cu_inj ('CATHODE', 'YAG02', 'YAG03') +bmad_cu_hxr ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') +surrogate_cu_inj ('CATHODE', 'OTR2') +cheetah_cu_hxr ('CATHODE', 'END') +``` + +### Shared Handoff Points +`common_handoff_points()` returns the locations two models can actually hand over at +— the intersection of their handoff points, with `CATHODE` excluded since nothing is +upstream of it. + +```python +>>> common_handoff_points("impact_cu_inj", "bmad_cu_hxr") +('YAG02', 'YAG03') + +>>> common_handoff_points("surrogate_cu_inj", "bmad_cu_hxr") +('OTR2',) + +>>> common_handoff_points("cheetah_cu_hxr", "bmad_cu_hxr") +() ``` ### Element Aliases @@ -51,20 +70,19 @@ GUN Use get_model() with a model ID and an optional end_ele to stop tracking at a specific screen. ```python ->>> get_model("bmad_cu_hxr_otr2", end_ele="TD11") +>>> get_model("bmad_cu_hxr", end_ele="TD11") ->>> get_model("impact_cu_inj", end_ele="OTR2") +>>> get_model("impact_cu_inj", end_ele="YAG03") ``` -Note: Unrecognized element types (e.g. solrf, Lcavity, Unknown) are skipped during model initialisation. These warnings are informational and do not prevent the model from loading. ### Error: Unknown Model Name Model IDs must be exact. Partial names are not supported: ```python ->>> get_model("bmad_cu_hxr", end_ele="TD11") -KeyError: "Unknown model 'bmad_cu_hxr'. Available: bmad_cu_hxr_otr2, bmad_cu_hxr_yag03, cheetah_cu_hxr, impact_cu_inj, surrogate_cu_inj" +>>> get_model("bmad_cu_hx", end_ele="TD11") +KeyError: "Unknown model 'bmad_cu_hx'. Available: bmad_cu_hxr, cheetah_cu_hxr, impact_cu_inj, surrogate_cu_inj" ``` ### Error: Invalid End Element @@ -73,15 +91,16 @@ end_ele must be one of the model's listed handoff points: ```python >>> get_model("impact_cu_inj", end_ele="otr99") ValueError: 'OTR99' is not an available end screen for 'impact_cu_inj'. -Suggested points: CATHODE, YAG02, YAG03, OTR1, OTR2 +Suggested points: CATHODE, YAG02, YAG03 ``` ## Staged Models Pass a list of two model IDs to get_model() to chain an injector model into a linac model. The upstream model hands off beam particles to the downstream model at a shared handoff point. -surrogate_cu_inj → bmad_cu_hxr_otr2 +surrogate_cu_inj → bmad_cu_hxr — no `handoff_loc` needed, it is inferred from the +surrogate's fixed end (OTR2): ```python ->>> m = get_model(["surrogate_cu_inj", "bmad_cu_hxr_otr2"], end_ele="OTR4", n_particles=500) +>>> m = get_model(["surrogate_cu_inj", "bmad_cu_hxr"], end_ele="OTR4", n_particles=500) >>> m.set({"QUAD:IN20:525:BCTRL": -10.0}) @@ -95,13 +114,11 @@ surrogate_cu_inj → bmad_cu_hxr_otr2 ['BEGINNING', 'OTR2', 'DE06D'] ``` -impact_cu_inj → bmad_cu_hxr_yag03 - -When impact_cu_inj is the upstream model, use bmad_cu_hxr_yag03 (which starts at YAG03) and specify handoff_loc: +impact_cu_inj → bmad_cu_hxr — hand off at YAG03: ```python >>> model = get_model( -... ["impact_cu_inj", "bmad_cu_hxr_yag03"], +... ["impact_cu_inj", "bmad_cu_hxr"], ... handoff_loc="YAG03", ... end_ele="TD11", ... n_particles=1000, @@ -113,51 +130,58 @@ When impact_cu_inj is the upstream model, use bmad_cu_hxr_yag03 (which starts at 2.3638227838794528e-07 ``` -### Pairing Rules -Upstream and downstream models must share a compatible handoff location. Mismatched pairs raise a descriptive error: +### Handoff Validation +`handoff_loc` must be a point both stages share. Anything else is rejected before any +model is built, so you do not pay for an IMPACT run to find out: ```python ->>> get_model(["surrogate_cu_inj", "bmad_cu_hxr_yag03"], end_ele="TD11", n_particles=500) -ValueError: 'surrogate_cu_inj' hands off at 'OTR2' but 'bmad_cu_hxr_yag03' starts at 'YAG03'. -Use 'bmad_cu_hxr_otr2' instead. +>>> get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="OTR4") +ValueError: 'OTR4' is not a shared handoff point for 'impact_cu_inj' -> 'bmad_cu_hxr'. +Available: YAG02, YAG03 +>>> get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="CATHODE") +ValueError: 'CATHODE' cannot be a handoff location: nothing is upstream of it. ``` -Valid pairs: -| Upstream | Downstream | Handoff Location |\ -| impact_cu_inj | bmad_cu_hxr_yag03 | YAG03 |\ -| surrogate_cu_inj | bmad_cu_hxr_otr2 | OTR2| +Standard chains: +| Upstream | Downstream | Handoff | +|---|---|---| +| `impact_cu_inj` | `bmad_cu_hxr` | YAG03 | +| `surrogate_cu_inj` | `bmad_cu_hxr` | OTR2 (inferred) | -### Advanced: Stripping Overlapping Variables -When building chained models manually, use _strip_overlapping_variables to remove output variables from the downstream model that are already owned by the upstream model. This avoids variable conflicts. - -```python -from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model -from virtual_accelerator.registry import _strip_overlapping_variables +LCLS needs two handoff planes because its injector models end at different places and +neither can move. The NN surrogate predicts `OTRS:IN20:571` (OTR2) at 135 MeV and +cannot produce a beam at YAG03, which sits before L0B at 64 MeV. -bm = get_cu_hxr_bmad_model(start_element="YAG03", end_element="OTR4", track_beam=True) +### Overlapping Variables Are Handled For You +Both stages include the handoff element, so both publish its PVs — an IMPACT model +stopped at `YAG03` keeps the screen (it prunes to `s <= stop`) and so does a Bmad model +sliced from `YAG03`. `StagedModel` would reject the pair as duplicates. -# Identify overlapping variables -yag = sorted(v for v in bm.supported_variables if "IN20:351" in v) +`get_model()` resolves this automatically: the upstream stage owns those PVs, because it +is the stage that tracks the beam to that plane, so they are unregistered from the +downstream stage before the chain is assembled. Nothing is required of the caller. -# Construct a mock upstream model exposing only those variables -class Up: - supported_variables = {n: bm.supported_variables[n] for n in yag} +The removal is surgical — only genuine collisions go. At YAG03 the IMPACT stage publishes +four PVs; the Bmad stage publishes those four plus `:X` and `:Y` centroid readbacks that +IMPACT does not provide. So the four move to IMPACT and the two Bmad-only ones stay: -# Strip them from the downstream model -removed = _strip_overlapping_variables(Up(), bm, "upstream", "bmad_cu_hxr") +```python +>>> m = get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="YAG03", end_ele="TD11") +>>> imp, bmad = m.lume_model_instances ->>> print(removed) +>>> sorted(v for v in imp.supported_variables if "IN20:351" in v) ['YAGS:IN20:351:Image:ArrayData', 'YAGS:IN20:351:Image:ArraySize0_RBV', - 'YAGS:IN20:351:Image:ArraySize1_RBV', 'YAGS:IN20:351:RESOLUTION', - 'YAGS:IN20:351:X', 'YAGS:IN20:351:Y'] + 'YAGS:IN20:351:Image:ArraySize1_RBV', 'YAGS:IN20:351:RESOLUTION'] ->>> print(len(bm.supported_variables)) # 318 → 312 -312 +>>> sorted(v for v in bmad.supported_variables if "IN20:351" in v) +['YAGS:IN20:351:X', 'YAGS:IN20:351:Y'] ``` -Note: _strip_overlapping_variables is a private utility. Prefer using get_model() with a list of model IDs, which handles this automatically. +A *writable* overlap raises instead of being dropped. That means both stages drive the +same magnet — their extents overlap rather than meeting at a plane — and dropping it +downstream would leave that stage tracking a stale value. ## API Reference ```get_model(spec, *, end_ele=None, start_ele=None, handoff_loc=None, n_particles=None)``` @@ -174,4 +198,10 @@ Printable summary of all registered models and their descriptions. ```list_handoff_points(model_id: str) -> tuple[str, ...]``` -Returns the available handoff point names for a given model. +Returns the suggested handoff point names for a given model, in lattice order. A +discovery aid, not a restriction. + +```common_handoff_points(*model_ids: str) -> tuple[str, ...]``` + +Returns the handoff points shared by all named models, in lattice order, excluding +`CATHODE`. Use it to see where two models can legally hand over. diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index c6d492e..d8431e9 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -15,7 +15,13 @@ logger = logging.getLogger(__name__) -__all__ = ["get_model", "models_available", "list_models", "list_handoff_points"] +__all__ = [ + "get_model", + "models_available", + "list_models", + "list_handoff_points", + "common_handoff_points", +] class _ModelCatalog(dict): @@ -51,6 +57,50 @@ def list_handoff_points(model_name: str) -> tuple[str, ...]: return _entry(model_name).handoff_points +CATHODE = "CATHODE" +"""Start-of-machine marker. Never a valid handoff: nothing is upstream of it.""" + + +def common_handoff_points(*model_names: str) -> tuple[str, ...]: + """Elements every named model can hand off at, in lattice order. + + ``CATHODE`` is always excluded -- it marks the front of the machine, so + nothing can hand over to a stage that begins there. + + This is the set intersection, not the union. A union would admit planes only + one stage can reach: ``impact_cu_inj`` stops by z=16.5 m and so cannot reach + ``OTR4`` at 17.80 m, but ``bmad_cu_hxr`` lists it, so the union would wrongly + accept ``handoff_loc="OTR4"`` for that pair. + + Parameters + ---------- + *model_names + Two or more registry names. + + Returns + ------- + tuple[str, ...] + Shared handoff elements, ordered by the first model's lattice order. + + Examples + -------- + >>> common_handoff_points("impact_cu_inj", "bmad_cu_hxr") + ('YAG02', 'YAG03') + >>> common_handoff_points("surrogate_cu_inj", "bmad_cu_hxr") + ('OTR2',) + """ + if len(model_names) < 2: + raise ValueError("Need at least two models to find common handoff points.") + + entries = [_entry(name) for name in model_names] + shared = set(entries[0].handoff_points) + for entry in entries[1:]: + shared &= set(entry.handoff_points) + shared.discard(CATHODE) + + return tuple(name for name in entries[0].handoff_points if name in shared) + + def _entry(name: str) -> ModelEntry: try: return MODELS[name] @@ -224,23 +274,16 @@ def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) - ) raise ValueError(f"{downstream.name!r} cannot be a downstream stage: {reason}.") - _check_element(upstream, handoff, "end") - _check_element(downstream, handoff, "start") - - # Each linac model is registered for one standard handoff plane, so a mismatch - # means the wrong pair was chosen -- silently re-slicing would leave a gap. - if downstream.default_start is not None and handoff != downstream.default_start: - alternatives = [ - name - for name, entry in MODELS.items() - if entry.facility == downstream.facility - and entry.engine == downstream.engine - and entry.default_start == handoff - ] - hint = f" Use {alternatives[0]!r} instead." if alternatives else "" + if handoff == CATHODE: raise ValueError( - f"{upstream.name!r} hands off at {handoff!r} but {downstream.name!r} " - f"starts at {downstream.default_start!r}.{hint}" + f"{CATHODE!r} cannot be a handoff location: nothing is upstream of it." + ) + + shared = common_handoff_points(upstream.name, downstream.name) + if handoff not in shared: + raise ValueError( + f"{handoff!r} is not a shared handoff point for {upstream.name!r} -> " + f"{downstream.name!r}. Available: {', '.join(shared) or 'none'}" ) @@ -309,8 +352,9 @@ def get_model( spec A registry name, or an ordered list of names to stage together. handoff_loc - Element where each consecutive pair hands the beam over. Inferred from - the upstream stage's fixed end when it has one. A list is required for + Element where each consecutive pair hands the beam over. Inferred from the + upstream stage's standard end when omitted. Must be a shared handoff point + of both stages -- see :func:`common_handoff_points`. A list is required for more than two stages. start_ele, end_ele Overall extent. For a staged model these apply to the first and last @@ -318,6 +362,38 @@ def get_model( **kwargs Builder parameters. For staged models, qualify an ambiguous parameter as ``"."``. + + Returns + ------- + LUMEModel + A single model, or a ``StagedModel`` wrapping the chain. + + Notes + ----- + For staged chains this handles two things that are easy to get wrong by hand. + + **Duplicate variables at the handoff are removed automatically.** Both stages + include the handoff element, so both publish its PVs -- an IMPACT model stopped + at ``YAG03`` keeps the screen (it prunes to ``s <= stop``) and so does a Bmad + model sliced from ``YAG03``. ``StagedModel`` would reject the pair as + duplicates. The upstream stage owns them, because it is the stage that tracks + the beam to that plane, so they are unregistered from the downstream stage + before the chain is assembled. A *writable* overlap raises instead: that means + both stages drive the same magnet, so their extents overlap rather than meeting + at a plane, and dropping it downstream would leave that stage tracking a stale + value. + + **Beam tracking is forced on for every stage that supports it**, since a + non-final stage must produce ``final_particles`` and a non-first stage must + accept ``initial_particles``. + + Examples + -------- + >>> get_model("bmad_cu_hxr", end_ele="TD11") # doctest: +SKIP + >>> get_model(["surrogate_cu_inj", "bmad_cu_hxr"], # doctest: +SKIP + ... end_ele="OTR4", n_particles=500) + >>> get_model(["impact_cu_inj", "bmad_cu_hxr"], # doctest: +SKIP + ... handoff_loc="YAG03", end_ele="TD11") """ start_ele, end_ele = _normalize(start_ele), _normalize(end_ele) diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index 0e177cd..b4c6e11 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -81,34 +81,6 @@ def configurable_extent(self) -> bool: ) -def _bmad_cu_hxr(name: str, start: str, note: str) -> ModelEntry: - """One Bmad CU-HXR entry per standard handoff plane. - - LCLS needs two because its two injector models end at different planes and - neither can move: the NN surrogate predicts OTRS:IN20:571 (OTR2) at 135 MeV - and cannot produce a beam at YAG03, which sits before L0B at 64 MeV. - """ - return ModelEntry( - name=name, - description=f"Bmad CU-HXR linac, {start} -> END ({note})", - facility="lcls", - engine="bmad", - builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model", - extras=("bmad",), - params={ - "start_element": start, - "end_element": "END", - "track_beam": False, - "custom_beam_path": None, - }, - handoff_points=("CATHODE", *_ALL_CU_HXR_SCREENS, "END"), - start_param="start_element", - end_param="end_element", - default_start=start, - default_end="END", - ) - - MODELS: dict[str, ModelEntry] = { "impact_cu_inj": ModelEntry( name="impact_cu_inj", @@ -120,21 +92,36 @@ def _bmad_cu_hxr(name: str, start: str, note: str) -> ModelEntry: params={"n_particles": 100, "end_element": "YAG03"}, # YAG01 and OTR3 exist in the deck but their lines are commented out; # OTR4 is past stop_1 at z=16.5. - handoff_points=("CATHODE", "YAG02", "YAG03", "OTR1", "OTR2"), + handoff_points=("CATHODE", "YAG02", "YAG03"), end_param="end_element", default_end="YAG03", broadcast_params=frozenset({"n_particles"}), element_aliases={"CATHODE": "GUN"}, ), - "bmad_cu_hxr_yag03": _bmad_cu_hxr( - "bmad_cu_hxr_yag03", "YAG03", "pairs with impact_cu_inj" - ), - "bmad_cu_hxr_otr2": _bmad_cu_hxr( - "bmad_cu_hxr_otr2", "OTR2", "pairs with surrogate_cu_inj" + "bmad_cu_hxr": ModelEntry( + name="bmad_cu_hxr", + description="Bmad CU-HXR linac, injector handoff -> END", + facility="lcls", + engine="bmad", + builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model", + extras=("bmad",), + params={ + "start_element": "OTR2", + "end_element": "END", + "track_beam": False, + "custom_beam_path": None, + }, + # Starts wherever the upstream injector hands off: YAG03 from + # impact_cu_inj, OTR2 from surrogate_cu_inj. + handoff_points=("CATHODE", *_ALL_CU_HXR_SCREENS, "END"), + start_param="start_element", + end_param="end_element", + default_start="OTR2", + default_end="END", ), "surrogate_cu_inj": ModelEntry( name="surrogate_cu_inj", - description="NN LCLS injector surrogate, fixed cathode -> OTR2", + description="NN LCLS injector surrogate, cathode -> OTR2", facility="lcls", engine="surrogate", builder=( @@ -148,13 +135,13 @@ def _bmad_cu_hxr(name: str, start: str, note: str) -> ModelEntry: ), "cheetah_cu_hxr": ModelEntry( name="cheetah_cu_hxr", - description="Cheetah nc_hxr", + description="Cheetah nc_hxr, cathode -> END", facility="lcls", engine="cheetah", builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_cheetah_model", extras=("cheetah",), params={"n_particles": 1000}, - handoff_points=(), + handoff_points=("CATHODE", "END"), broadcast_params=frozenset({"n_particles"}), ), } diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index 4c88842..b319589 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -9,6 +9,7 @@ _resolve_handoffs, _route_kwargs, _strip_overlapping_variables, + common_handoff_points, get_model, list_handoff_points, list_models, @@ -27,19 +28,21 @@ def test_repr_is_aligned_table(self): assert len(text.splitlines()) == len(MODELS) def test_filter_by_engine_and_facility(self): - assert list_models(engine="bmad") == ["bmad_cu_hxr_yag03", "bmad_cu_hxr_otr2"] + assert list_models(engine="bmad") == ["bmad_cu_hxr"] assert set(list_models(facility="lcls")) == set(MODELS) assert list_models(facility="facet2") == [] def test_handoff_points_are_lattice_ordered(self): - diags = list_handoff_points("bmad_cu_hxr_otr2") + diags = list_handoff_points("bmad_cu_hxr") assert diags.index("YAG02") < diags.index("YAG03") < diags.index("OTR2") - def test_impact_omits_unavailable_screens(self): - # YAG01/OTR3 are commented out in the deck; OTR4 is past stop_1. + def test_impact_lists_only_its_standard_extent(self): + # Standard injector extent is cathode -> YAG03. Screens further downstream + # are excluded: YAG01/OTR3 are commented out in the deck, OTR4 is past + # stop_1 at z=16.5, and OTR1/OTR2 are past the standard handoff. diags = list_handoff_points("impact_cu_inj") - assert "OTR2" in diags - for absent in ("YAG01", "OTR3", "OTR4"): + assert diags == ("CATHODE", "YAG02", "YAG03") + for absent in ("YAG01", "OTR1", "OTR2", "OTR3", "OTR4"): assert absent not in diags @@ -86,7 +89,7 @@ def test_rejects_unavailable_screen(self): def test_rejects_impact_as_downstream_stage(self): with pytest.raises(ValueError, match="only start at the cathode"): - get_model(["bmad_cu_hxr_otr2", "impact_cu_inj"], handoff_loc="OTR2") + get_model(["bmad_cu_hxr", "impact_cu_inj"], handoff_loc="OTR2") def test_rejects_start_ele_on_fixed_extent_model(self): with pytest.raises(ValueError, match="fixed start"): @@ -94,28 +97,26 @@ def test_rejects_start_ele_on_fixed_extent_model(self): def test_rejects_single_model_list(self): with pytest.raises(ValueError, match="at least two models"): - get_model(["bmad_cu_hxr_otr2"]) + get_model(["bmad_cu_hxr"]) def test_rejects_wrong_handoff_count(self): with pytest.raises(ValueError, match="handoff location"): - get_model( - ["surrogate_cu_inj", "bmad_cu_hxr_otr2"], handoff_loc=["OTR2", "OTR3"] - ) + get_model(["surrogate_cu_inj", "bmad_cu_hxr"], handoff_loc=["OTR2", "OTR3"]) - def test_rejects_mismatched_injector_and_linac_pairing(self): - # impact ends at YAG03; bmad_cu_hxr_otr2 starts at OTR2 -> 9.6 m gap. - with pytest.raises(ValueError, match="bmad_cu_hxr_yag03"): - get_model(["impact_cu_inj", "bmad_cu_hxr_otr2"]) + def test_rejects_cathode_as_handoff(self): + with pytest.raises(ValueError, match="nothing is upstream"): + get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="CATHODE") - def test_rejects_surrogate_paired_with_yag03_linac(self): - with pytest.raises(ValueError, match="bmad_cu_hxr_otr2"): - get_model(["surrogate_cu_inj", "bmad_cu_hxr_yag03"]) + def test_rejects_handoff_not_shared_by_both_stages(self): + # OTR4 is past impact_cu_inj's stop at z=16.5, so it cannot hand off there. + with pytest.raises(ValueError, match="not a shared handoff point"): + get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="OTR4") class TestKwargRouting: def test_unknown_kwarg_rejected(self): with pytest.raises(ValueError, match="not a parameter of any stage"): - _route_kwargs([MODELS["bmad_cu_hxr_otr2"]], {"n_particle": 5}) + _route_kwargs([MODELS["bmad_cu_hxr"]], {"n_particle": 5}) def test_broadcast_reaches_every_declaring_stage(self): entries = [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]] @@ -123,7 +124,7 @@ def test_broadcast_reaches_every_declaring_stage(self): assert routed == [{"n_particles": 42}, {"n_particles": 42}] def test_routes_to_single_declaring_stage(self): - entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr_otr2"]] + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] routed = _route_kwargs(entries, {"track_beam": True}) assert routed == [{}, {"track_beam": True}] @@ -134,23 +135,61 @@ def test_dotted_form_targets_one_stage(self): def test_dotted_form_rejects_unknown_stage(self): with pytest.raises(ValueError, match="not in this model"): - _route_kwargs([MODELS["bmad_cu_hxr_otr2"]], {"nope.track_beam": True}) + _route_kwargs([MODELS["bmad_cu_hxr"]], {"nope.track_beam": True}) def test_dotted_form_rejects_unknown_param(self): with pytest.raises(ValueError, match="not a parameter of"): - _route_kwargs([MODELS["bmad_cu_hxr_otr2"]], {"bmad_cu_hxr_otr2.bogus": 1}) + _route_kwargs([MODELS["bmad_cu_hxr"]], {"bmad_cu_hxr.bogus": 1}) class TestHandoffResolution: def test_inferred_from_upstream_fixed_end(self): - entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr_otr2"]] + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] assert _resolve_handoffs(entries, None) == ["OTR2"] def test_explicit_handoff_is_used_verbatim(self): - entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr_yag03"]] + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] assert _resolve_handoffs(entries, "YAG03") == ["YAG03"] +class TestCommonHandoffPoints: + def test_intersection_of_the_two_standard_chains(self): + assert common_handoff_points("impact_cu_inj", "bmad_cu_hxr") == ( + "YAG02", + "YAG03", + ) + assert common_handoff_points("surrogate_cu_inj", "bmad_cu_hxr") == ("OTR2",) + + def test_cathode_is_always_excluded(self): + # CATHODE is in both models' handoff_points but is never a valid handoff. + assert "CATHODE" in MODELS["impact_cu_inj"].handoff_points + assert "CATHODE" in MODELS["bmad_cu_hxr"].handoff_points + assert "CATHODE" not in common_handoff_points("impact_cu_inj", "bmad_cu_hxr") + + def test_is_intersection_not_union(self): + # OTR4 is only reachable by bmad_cu_hxr; a union would wrongly include it. + shared = common_handoff_points("impact_cu_inj", "bmad_cu_hxr") + assert "OTR4" in MODELS["bmad_cu_hxr"].handoff_points + assert "OTR4" not in shared + + def test_ordered_by_lattice_position(self): + shared = common_handoff_points("impact_cu_inj", "bmad_cu_hxr") + assert list(shared) == sorted( + shared, key=MODELS["impact_cu_inj"].handoff_points.index + ) + + def test_no_shared_points_gives_empty_tuple(self): + assert common_handoff_points("cheetah_cu_hxr", "bmad_cu_hxr") == () + + def test_requires_at_least_two_models(self): + with pytest.raises(ValueError, match="at least two models"): + common_handoff_points("bmad_cu_hxr") + + def test_unknown_model_name(self): + with pytest.raises(KeyError, match="Unknown model"): + common_handoff_points("impact_cu_inj", "nope") + + class _FakeStage: """Minimal stand-in for a LUMEModel with registerable action variables.""" @@ -247,6 +286,6 @@ def test_lowercase_valid_screen_is_accepted(self): assert "not an available" not in str(excinfo.value) def test_lowercase_handoff_normalises_before_resolution(self): - entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr_yag03"]] + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] handoffs = [_normalize(h) for h in _resolve_handoffs(entries, "yag03")] assert handoffs == ["YAG03"] From 07897301dfadf437e2f82dded97b39a7e12c371d Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Fri, 4 Sep 2026 09:27:37 -0700 Subject: [PATCH 08/15] kwargs --- docs/model_registry_usage.md | 71 +++++++++++++++++++--- virtual_accelerator/registry/__init__.py | 18 +++--- virtual_accelerator/registry/models.py | 17 ++---- virtual_accelerator/tests/test_registry.py | 26 +++++--- 4 files changed, 93 insertions(+), 39 deletions(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index d3d0c30..9054529 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -183,15 +183,70 @@ A *writable* overlap raises instead of being dropped. That means both stages dri same magnet — their extents overlap rather than meeting at a plane — and dropping it downstream would leave that stage tracking a stale value. +### Targeting One Stage With kwargs +Plain kwargs apply to whichever stage declares them. `n_particles` is a *broadcast* +parameter, so it reaches every stage that accepts one; `start_ele` and `end_ele` apply to +the first and last stage respectively. + +To target a specific stage, prefix the parameter with the model ID: + +```python +>>> m = get_model( +... ["impact_cu_inj", "bmad_cu_hxr"], +... handoff_loc="YAG03", +... **{"impact_cu_inj.n_particles": 200, "bmad_cu_hxr.end_ele": "TD11"}, +... ) +``` + +The dotted form always wins and is never ambiguous. Both the `get_model` spelling +(`end_ele`, `start_ele`) and the underlying builder spelling (`end_element`, +`start_element`) are accepted: + +```python +>>> "bmad_cu_hxr.end_ele" # same as +>>> "bmad_cu_hxr.end_element" +``` + +Ambiguity is an error rather than a guess. A non-broadcast parameter declared by more than +one stage must be qualified: + +```python +>>> get_model(["impact_cu_inj", "bmad_cu_hxr"], end_element="TD11") +ValueError: 'end_element' is ambiguous across stages (impact_cu_inj, bmad_cu_hxr). +Qualify it, e.g. "impact_cu_inj.end_element=...". +``` + +Unknown parameters are rejected outright, listing what is accepted: + +```python +>>> get_model("bmad_cu_hxr", n_particle=5) +ValueError: 'n_particle' is not a parameter of any stage. +Accepted: custom_beam_path, end_element, start_element, track_beam +``` + +To see what a model accepts: + +```python +>>> MODELS["bmad_cu_hxr"].params +{'start_element': 'OTR2', 'end_element': 'END', 'track_beam': False, 'custom_beam_path': None} + +>>> MODELS["impact_cu_inj"].broadcast_params +frozenset({'n_particles'}) +``` + ## API Reference -```get_model(spec, *, end_ele=None, start_ele=None, handoff_loc=None, n_particles=None)``` - -|Parameter | Type | Description |\ -|spec | str or list[str] | Model ID or [upstream, downstream] pair|\ -|end_ele | str | Screen name to stop tracking at|\ -|start_ele | str | Screen name to start tracking from|\ -|handoff_loc | str | Explicit handoff point when chaining models|\ -|n_particles | int | Number of macro-particles for beam tracking| +```get_model(spec, *, handoff_loc=None, start_ele=None, end_ele=None, **kwargs)``` + +| Parameter | Type | Description | +|---|---|---| +| `spec` | str or list[str] | Model ID, or `[upstream, downstream]` to chain | +| `handoff_loc` | str | Where the stages exchange beam. Inferred from the upstream model's standard end when omitted. Must be in `common_handoff_points()` | +| `start_ele` | str | Element to start tracking from (first stage) | +| `end_ele` | str | Element to stop tracking at (last stage) | +| `**kwargs` | any | Builder parameters. Prefix with `"."` to target one stage | + +For staged chains `get_model()` also removes variables that both stages publish at the +handoff, and forces beam tracking on for every stage that supports it. ```models_available``` Printable summary of all registered models and their descriptions. diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index d8431e9..247b83d 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -129,17 +129,12 @@ def _normalize(name: str | None) -> str | None: Lattice element names are upper case everywhere. Tao is case-insensitive so a lower-case name would appear to work, but IMPACT's ``impact.ele[...]`` is a - plain dict lookup, and the registry's own ``handoff_points`` and - ``element_aliases`` lookups would silently miss. + plain dict lookup, and the registry's own ``handoff_points`` lookups would + silently miss. """ return name if name is None else name.upper() -def _resolve_element(entry: ModelEntry, name: str) -> str: - """Translate a standard element name into this model's engine-local name.""" - return entry.element_aliases.get(name, name) - - def _check_element(entry: ModelEntry, name: str, role: str) -> None: """Validate a start/end/handoff element. @@ -178,6 +173,11 @@ def _route_kwargs( f"Stages: {', '.join(by_name)}" ) index = by_name[stage_name] + # Accept the get_model spelling per stage, e.g. "bmad_cu_hxr.end_ele". + param = { + "start_ele": entries[index].start_param, + "end_ele": entries[index].end_param, + }.get(param, param) if param not in entries[index].params: raise ValueError( f"{param!r} is not a parameter of {stage_name!r}. " @@ -220,7 +220,7 @@ def _build( f"{entry.name!r} has a fixed start and does not accept start_ele." ) _check_element(entry, start_ele, "start") - kwargs[entry.start_param] = _resolve_element(entry, start_ele) + kwargs[entry.start_param] = start_ele if end_ele is not None: if entry.end_param is None: @@ -228,7 +228,7 @@ def _build( f"{entry.name!r} has a fixed end and does not accept end_ele." ) _check_element(entry, end_ele, "end") - kwargs[entry.end_param] = _resolve_element(entry, end_ele) + kwargs[entry.end_param] = end_ele return _load_builder(entry)(**kwargs) diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index b4c6e11..a5525b8 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -3,7 +3,7 @@ Currently LCLS only; FACET-II entries are not registered yet. """ -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any @@ -52,14 +52,6 @@ class ModelEntry: broadcast_params: frozenset[str] = frozenset() """Params safe to send to every stage of a staged model, e.g. n_particles.""" - element_aliases: dict[str, str] = field(default_factory=dict) - """Standard name -> this engine's local name. - - Diagnostics already agree between Bmad and IMPACT, so this only carries the - handful of elements that genuinely differ -- currently just the gun, which - Bmad calls CATHODE and IMPACT calls GUN. - """ - @property def configurable_extent(self) -> bool: return self.start_param is not None or self.end_param is not None @@ -92,11 +84,10 @@ def configurable_extent(self) -> bool: params={"n_particles": 100, "end_element": "YAG03"}, # YAG01 and OTR3 exist in the deck but their lines are commented out; # OTR4 is past stop_1 at z=16.5. - handoff_points=("CATHODE", "YAG02", "YAG03"), + handoff_points=("YAG02", "YAG03"), end_param="end_element", default_end="YAG03", broadcast_params=frozenset({"n_particles"}), - element_aliases={"CATHODE": "GUN"}, ), "bmad_cu_hxr": ModelEntry( name="bmad_cu_hxr", @@ -129,7 +120,7 @@ def configurable_extent(self) -> bool: ), extras=("surrogate",), params={"n_particles": 1000}, - handoff_points=("CATHODE", "OTR2"), + handoff_points=("OTR2",), default_end="OTR2", broadcast_params=frozenset({"n_particles"}), ), @@ -141,7 +132,7 @@ def configurable_extent(self) -> bool: builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_cheetah_model", extras=("cheetah",), params={"n_particles": 1000}, - handoff_points=("CATHODE", "END"), + handoff_points=(), broadcast_params=frozenset({"n_particles"}), ), } diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index b319589..c60907e 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -41,10 +41,17 @@ def test_impact_lists_only_its_standard_extent(self): # are excluded: YAG01/OTR3 are commented out in the deck, OTR4 is past # stop_1 at z=16.5, and OTR1/OTR2 are past the standard handoff. diags = list_handoff_points("impact_cu_inj") - assert diags == ("CATHODE", "YAG02", "YAG03") + assert diags == ("YAG02", "YAG03") for absent in ("YAG01", "OTR1", "OTR2", "OTR3", "OTR4"): assert absent not in diags + def test_cathode_is_only_listed_where_it_is_usable(self): + # bmad accepts start_ele="CATHODE"; the injectors have a fixed start, so + # listing it there would advertise something that cannot be passed. + assert "CATHODE" in list_handoff_points("bmad_cu_hxr") + for fixed in ("impact_cu_inj", "surrogate_cu_inj", "cheetah_cu_hxr"): + assert "CATHODE" not in list_handoff_points(fixed) + class TestEntryIntegrity: @pytest.mark.parametrize("name", sorted(MODELS)) @@ -64,11 +71,6 @@ def test_broadcast_params_are_declared(self, name): entry = MODELS[name] assert entry.broadcast_params <= set(entry.params) - @pytest.mark.parametrize("name", sorted(MODELS)) - def test_element_alias_keys_are_handoff_points(self, name): - entry = MODELS[name] - assert set(entry.element_aliases) <= set(entry.handoff_points) - @pytest.mark.parametrize("name", sorted(MODELS)) def test_defaults_are_consistent(self, name): entry = MODELS[name] @@ -137,6 +139,13 @@ def test_dotted_form_rejects_unknown_stage(self): with pytest.raises(ValueError, match="not in this model"): _route_kwargs([MODELS["bmad_cu_hxr"]], {"nope.track_beam": True}) + def test_dotted_form_accepts_end_ele_per_stage(self): + entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] + routed = _route_kwargs( + entries, {"impact_cu_inj.end_ele": "YAG02", "bmad_cu_hxr.end_ele": "TD11"} + ) + assert routed == [{"end_element": "YAG02"}, {"end_element": "TD11"}] + def test_dotted_form_rejects_unknown_param(self): with pytest.raises(ValueError, match="not a parameter of"): _route_kwargs([MODELS["bmad_cu_hxr"]], {"bmad_cu_hxr.bogus": 1}) @@ -161,10 +170,9 @@ def test_intersection_of_the_two_standard_chains(self): assert common_handoff_points("surrogate_cu_inj", "bmad_cu_hxr") == ("OTR2",) def test_cathode_is_always_excluded(self): - # CATHODE is in both models' handoff_points but is never a valid handoff. - assert "CATHODE" in MODELS["impact_cu_inj"].handoff_points + # bmad lists CATHODE, so the intersection must drop it explicitly. assert "CATHODE" in MODELS["bmad_cu_hxr"].handoff_points - assert "CATHODE" not in common_handoff_points("impact_cu_inj", "bmad_cu_hxr") + assert "CATHODE" not in common_handoff_points("bmad_cu_hxr", "bmad_cu_hxr") def test_is_intersection_not_union(self): # OTR4 is only reachable by bmad_cu_hxr; a union would wrongly include it. From d32f98a160369cd4fc48fff50c109036a3b17206 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Fri, 4 Sep 2026 09:53:45 -0700 Subject: [PATCH 09/15] kwargs corrected --- docs/model_registry_design.md | 22 ++++--- docs/model_registry_usage.md | 76 +++++++++++++++------- docs/registry_decisions_and_questions.md | 15 +++-- virtual_accelerator/registry/__init__.py | 28 +++++++- virtual_accelerator/registry/models.py | 16 +++-- virtual_accelerator/tests/test_registry.py | 29 +++++++-- 6 files changed, 132 insertions(+), 54 deletions(-) diff --git a/docs/model_registry_design.md b/docs/model_registry_design.md index 1ea0d9f..5b78659 100644 --- a/docs/model_registry_design.md +++ b/docs/model_registry_design.md @@ -42,8 +42,9 @@ handoff points: | gun / cathode | `CATHODE` / `CATHODEF` | `GUN` / `GUNF` | | L0A cavity | `L0A` (one lcavity) | `L0A_entrance`, `L0A_body_1`, `L0A_body_2`, `L0A_exit` | -So the alias dict carries only the gun (`CATHODE` → IMPACT `GUN`). Cavity segmentation is one-to-many -and cavities are not sensible handoff planes, so they are simply not registered as handoff points. +No alias dict is needed. The gun would have been the one entry (`CATHODE` → IMPACT `GUN`) but it is +unreachable: IMPACT has no start parameter to pass it to. Cavity segmentation is one-to-many and +cavities are not sensible handoff planes, so they are simply not registered as handoff points. Two facts that don't need machinery but should be written down: @@ -70,16 +71,13 @@ class ModelEntry: builder: str # "virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model" extras: tuple[str, ...] # pip extras needed, e.g. ("bmad",) params: dict[str, Any] # param name -> default, for validation + discovery - broadcast_params: frozenset[str] # params safe to send to every stage, e.g. {"n_particles"} + shared_params: frozenset[str] # must match across stages, e.g. {"n_particles"} handoff_points: tuple[str, ...] # suggested elements, ORDERED by lattice position start_param: str | None # builder kwarg for start, None if fixed end_param: str | None # builder kwarg for end, None if fixed default_start: str | None default_end: str | None - - element_aliases: dict[str, str] = field(default_factory=dict) - # standard name -> this engine's local name. Empty today; escape hatch only. ``` `start_param` / `end_param` replace the earlier `can_start_anywhere` flag: they carry the same @@ -115,9 +113,9 @@ virtual_accelerator/registry/ | registry name | facility | engine | builder | key params | suggested handoff points | |---|---|---|---|---|---| -| `impact_cu_inj` | lcls | impact | `get_cu_inj_impact_model` | `n_particles=100`, `end_element="OTR2"` | CATHODE, YAG02, YAG03, OTR1, OTR2 | +| `impact_cu_inj` | lcls | impact | `get_cu_inj_impact_model` | `n_particles=100`, `end_element="YAG03"` | YAG02, YAG03 | | `bmad_cu_hxr` | lcls | bmad | `get_cu_hxr_bmad_model` | `start_element="OTR2"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | CATHODE, all 12 profmon screens, END | -| `surrogate_cu_inj` | lcls | surrogate | `get_cu_hxr_injector_surrogate_model` | `n_particles=1000` | CATHODE, OTR2 (fixed extent) | +| `surrogate_cu_inj` | lcls | surrogate | `get_cu_hxr_injector_surrogate_model` | `n_particles=1000` | OTR2 (fixed extent) | | `cheetah_cu_hxr` | lcls | cheetah | `get_cu_hxr_cheetah_model` | `n_particles=1000` | — | Naming convention: `_`. @@ -214,9 +212,13 @@ order shown. Because the registry declares each model's params, routing is a lookup, not signature introspection: -1. Param in `broadcast_params` (e.g. `n_particles`) → sent to every stage that declares it. +1. Param in `shared_params` (e.g. `n_particles`) → sent to every stage that declares it, and + the per-stage form is **rejected**: the beam flows through the stages, so differing values + would break a physical invariant rather than configure anything. 2. Declared by exactly one stage → routed there. -3. Declared by more than one stage and not broadcast → `ValueError` naming the candidates. +3. Declared by more than one stage and not shared → `ValueError` naming the candidates. The + builder spellings `end_element`/`start_element` are rejected flat, since they do not say + which stage they mean; use `end_ele`/`start_ele`, or qualify per stage. 4. `"."` always wins. 5. Matching no declared param → rejected with a suggestion, not silently forwarded. diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index 9054529..7285637 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -36,12 +36,17 @@ Each model exposes a set of suggested handoff points — named locations where b >>> for m in ["impact_cu_inj", "bmad_cu_hxr", "surrogate_cu_inj", "cheetah_cu_hxr"]: ... print(m, list_handoff_points(m)) -impact_cu_inj ('CATHODE', 'YAG02', 'YAG03') +impact_cu_inj ('YAG02', 'YAG03') bmad_cu_hxr ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') -surrogate_cu_inj ('CATHODE', 'OTR2') -cheetah_cu_hxr ('CATHODE', 'END') +surrogate_cu_inj ('OTR2',) +cheetah_cu_hxr () ``` +`CATHODE` appears only for `bmad_cu_hxr`, the one model whose start is configurable — you +can slice it from the front of the machine with `start_ele="CATHODE"`. The injector models +always begin at the cathode and cannot be told otherwise, so listing it there would +advertise something you cannot pass. + ### Shared Handoff Points `common_handoff_points()` returns the locations two models can actually hand over at — the intersection of their handoff points, with `CATHODE` excluded since nothing is @@ -58,14 +63,6 @@ upstream of it. () ``` -### Element Aliases -Some handoff point names are aliases for internal model element names: - -```python ->>> print(MODELS["impact_cu_inj"].element_aliases["CATHODE"]) -GUN -``` - ### Loading a Single Model Use get_model() with a model ID and an optional end_ele to stop tracking at a specific screen. @@ -91,7 +88,7 @@ end_ele must be one of the model's listed handoff points: ```python >>> get_model("impact_cu_inj", end_ele="otr99") ValueError: 'OTR99' is not an available end screen for 'impact_cu_inj'. -Suggested points: CATHODE, YAG02, YAG03 +Suggested points: YAG02, YAG03 ``` ## Staged Models @@ -184,36 +181,65 @@ same magnet — their extents overlap rather than meeting at a plane — and dro downstream would leave that stage tracking a stale value. ### Targeting One Stage With kwargs -Plain kwargs apply to whichever stage declares them. `n_particles` is a *broadcast* -parameter, so it reaches every stage that accepts one; `start_ele` and `end_ele` apply to -the first and last stage respectively. +Parameters fall into two kinds, and which one it is decides how you pass it. + +**Shared parameters must hold the same value in every stage.** `n_particles` is the only +one: the beam flows through the stages, so a particle count that differs between them is +physically meaningless. Pass it flat and it reaches every stage that declares one. + +```python +>>> m = get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="YAG03", n_particles=1000) +``` + +Setting a shared parameter per stage is refused — divergence would break the invariant +rather than configure anything: + +```python +>>> get_model([...], **{"impact_cu_inj.n_particles": 200}) +ValueError: 'n_particles' must be the same in every stage, so it cannot be set per stage. +Pass n_particles=... instead of 'impact_cu_inj.n_particles'. +``` + +**Everything else means something different to each stage**, so it is either unambiguous +or you say which stage. `track_beam` and `custom_beam_path` are only declared by +`bmad_cu_hxr`, so they route there on their own: + +```python +>>> m = get_model(["surrogate_cu_inj", "bmad_cu_hxr"], custom_beam_path="beam.h5") +``` + +Start and end elements need naming, because both stages have them. Use `start_ele` / +`end_ele` for the overall extent — first and last stage respectively: + +```python +>>> m = get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="YAG03", end_ele="TD11") +``` -To target a specific stage, prefix the parameter with the model ID: +Or prefix with the model ID to set one stage: ```python >>> m = get_model( ... ["impact_cu_inj", "bmad_cu_hxr"], ... handoff_loc="YAG03", -... **{"impact_cu_inj.n_particles": 200, "bmad_cu_hxr.end_ele": "TD11"}, +... **{"bmad_cu_hxr.end_ele": "TD11"}, ... ) ``` -The dotted form always wins and is never ambiguous. Both the `get_model` spelling -(`end_ele`, `start_ele`) and the underlying builder spelling (`end_element`, -`start_element`) are accepted: +The dotted form accepts either spelling — the registry's (`end_ele`, `start_ele`) or the +underlying builder's (`end_element`, `start_element`): ```python >>> "bmad_cu_hxr.end_ele" # same as >>> "bmad_cu_hxr.end_element" ``` -Ambiguity is an error rather than a guess. A non-broadcast parameter declared by more than -one stage must be qualified: +Passing a builder spelling *flat* is refused, since it does not say which stage it means: ```python >>> get_model(["impact_cu_inj", "bmad_cu_hxr"], end_element="TD11") -ValueError: 'end_element' is ambiguous across stages (impact_cu_inj, bmad_cu_hxr). -Qualify it, e.g. "impact_cu_inj.end_element=...". +ValueError: Do not pass 'end_element' directly -- it is the builder's own name and does +not say which stage it applies to. Use end_ele=... for the overall extent, or +".end_element=..." to target one stage. ``` Unknown parameters are rejected outright, listing what is accepted: @@ -230,7 +256,7 @@ To see what a model accepts: >>> MODELS["bmad_cu_hxr"].params {'start_element': 'OTR2', 'end_element': 'END', 'track_beam': False, 'custom_beam_path': None} ->>> MODELS["impact_cu_inj"].broadcast_params +>>> MODELS["impact_cu_inj"].shared_params frozenset({'n_particles'}) ``` diff --git a/docs/registry_decisions_and_questions.md b/docs/registry_decisions_and_questions.md index e827248..acf8523 100644 --- a/docs/registry_decisions_and_questions.md +++ b/docs/registry_decisions_and_questions.md @@ -19,8 +19,11 @@ definitely wrong and is rejected early rather than failing deep inside Tao. Anyt through to the engine. There are 3323 elements in `cu_hxr`, so exhaustive validation is not on the table. -`CATHODE` is the first entry that makes `element_aliases` non-empty: Bmad calls it `CATHODE`, IMPACT -calls it `GUN`. +`CATHODE` is listed only for `bmad_cu_hxr`, the one model whose start is configurable. The +injectors always begin at the cathode and cannot be told otherwise, so an alias mapping Bmad's +`CATHODE` to IMPACT's `GUN` was added and then removed again -- it was unreachable, since IMPACT +has no start parameter to pass it to. `element_aliases` is gone with it; diagnostics already agree +between the engines, so there was nothing left for it to carry. ### A2. Reference plane is the element ENTRANCE @@ -84,9 +87,11 @@ Kwarg routing across stages is a table lookup, not signature introspection, so e name the candidate stages. Trade-off: adding a parameter to a builder means also adding it to the registry. That cost buys real error messages and a truthful `models_available`. -Routing rules: broadcast params (`n_particles`) go to every stage declaring them; a param declared -by exactly one stage routes there; declared by more than one and not broadcast is an error naming -the candidates; `"."` always wins. +Routing rules: *shared* params (`n_particles`) go to every stage declaring them and cannot be set +per stage, since the beam flows through and differing values would break a physical invariant; a +param declared by exactly one stage routes there; declared by more than one is an error naming the +candidates; `"."` targets one stage. The builders' own `end_element` / +`start_element` are rejected flat in favour of `end_ele` / `start_ele`. ### A6. Element names normalised to upper case at the API boundary diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index 247b83d..aed0612 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -164,6 +164,15 @@ def _route_kwargs( routed: list[dict[str, Any]] = [{} for _ in entries] by_name = {entry.name: i for i, entry in enumerate(entries)} + # Builder spellings of the extent params. Flat use is rejected: which stage a + # bare "end_element" means is ambiguous, and start_ele/end_ele already say it. + extent_params = { + param + for entry in entries + for param in (entry.start_param, entry.end_param) + if param is not None + } + for key, value in kwargs.items(): stage_name, sep, param = key.partition(".") if sep: @@ -178,6 +187,11 @@ def _route_kwargs( "start_ele": entries[index].start_param, "end_ele": entries[index].end_param, }.get(param, param) + if param in entries[index].shared_params: + raise ValueError( + f"{param!r} must be the same in every stage, so it cannot be set " + f"per stage. Pass {param}=... instead of {key!r}." + ) if param not in entries[index].params: raise ValueError( f"{param!r} is not a parameter of {stage_name!r}. " @@ -186,6 +200,16 @@ def _route_kwargs( routed[index][param] = value continue + if key in extent_params: + role = ( + "start_ele" if any(e.start_param == key for e in entries) else "end_ele" + ) + raise ValueError( + f"Do not pass {key!r} directly -- it is the builder's own name and does " + f"not say which stage it applies to. Use {role}=... for the overall " + f'extent, or ".{key}=..." to target one stage.' + ) + accepting = [i for i, entry in enumerate(entries) if key in entry.params] if not accepting: known = sorted({p for entry in entries for p in entry.params}) @@ -193,8 +217,8 @@ def _route_kwargs( f"{key!r} is not a parameter of any stage. Accepted: {', '.join(known)}" ) - broadcast = any(key in entries[i].broadcast_params for i in accepting) - if len(accepting) > 1 and not broadcast: + shared = any(key in entries[i].shared_params for i in accepting) + if len(accepting) > 1 and not shared: names = ", ".join(entries[i].name for i in accepting) raise ValueError( f"{key!r} is ambiguous across stages ({names}). " diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index a5525b8..e1ccc2b 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -49,8 +49,14 @@ class ModelEntry: default_start: str | None = None default_end: str | None = None - broadcast_params: frozenset[str] = frozenset() - """Params safe to send to every stage of a staged model, e.g. n_particles.""" + shared_params: frozenset[str] = frozenset() + """Params that must hold the same value in every stage of a chain. + + The beam flows through the stages, so a particle count differing between them + is physically meaningless. These are broadcast to every stage that declares + them, and the per-stage ``"."`` form is rejected for them -- + letting the values diverge would break the invariant, not configure anything. + """ @property def configurable_extent(self) -> bool: @@ -87,7 +93,7 @@ def configurable_extent(self) -> bool: handoff_points=("YAG02", "YAG03"), end_param="end_element", default_end="YAG03", - broadcast_params=frozenset({"n_particles"}), + shared_params=frozenset({"n_particles"}), ), "bmad_cu_hxr": ModelEntry( name="bmad_cu_hxr", @@ -122,7 +128,7 @@ def configurable_extent(self) -> bool: params={"n_particles": 1000}, handoff_points=("OTR2",), default_end="OTR2", - broadcast_params=frozenset({"n_particles"}), + shared_params=frozenset({"n_particles"}), ), "cheetah_cu_hxr": ModelEntry( name="cheetah_cu_hxr", @@ -133,6 +139,6 @@ def configurable_extent(self) -> bool: extras=("cheetah",), params={"n_particles": 1000}, handoff_points=(), - broadcast_params=frozenset({"n_particles"}), + shared_params=frozenset({"n_particles"}), ), } diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index c60907e..d27c123 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -67,9 +67,9 @@ def test_extent_params_are_declared(self, name): assert param in entry.params @pytest.mark.parametrize("name", sorted(MODELS)) - def test_broadcast_params_are_declared(self, name): + def test_shared_params_are_declared(self, name): entry = MODELS[name] - assert entry.broadcast_params <= set(entry.params) + assert entry.shared_params <= set(entry.params) @pytest.mark.parametrize("name", sorted(MODELS)) def test_defaults_are_consistent(self, name): @@ -120,7 +120,7 @@ def test_unknown_kwarg_rejected(self): with pytest.raises(ValueError, match="not a parameter of any stage"): _route_kwargs([MODELS["bmad_cu_hxr"]], {"n_particle": 5}) - def test_broadcast_reaches_every_declaring_stage(self): + def test_shared_param_reaches_every_declaring_stage(self): entries = [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]] routed = _route_kwargs(entries, {"n_particles": 42}) assert routed == [{"n_particles": 42}, {"n_particles": 42}] @@ -131,14 +131,29 @@ def test_routes_to_single_declaring_stage(self): assert routed == [{}, {"track_beam": True}] def test_dotted_form_targets_one_stage(self): - entries = [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]] - routed = _route_kwargs(entries, {"cheetah_cu_hxr.n_particles": 7}) - assert routed == [{}, {"n_particles": 7}] + entries = [MODELS["surrogate_cu_inj"], MODELS["bmad_cu_hxr"]] + routed = _route_kwargs(entries, {"bmad_cu_hxr.track_beam": True}) + assert routed == [{}, {"track_beam": True}] def test_dotted_form_rejects_unknown_stage(self): with pytest.raises(ValueError, match="not in this model"): _route_kwargs([MODELS["bmad_cu_hxr"]], {"nope.track_beam": True}) + def test_shared_param_cannot_be_set_per_stage(self): + # n_particles must match across stages -- the beam flows through them. + with pytest.raises(ValueError, match="same in every stage"): + _route_kwargs( + [MODELS["surrogate_cu_inj"], MODELS["cheetah_cu_hxr"]], + {"cheetah_cu_hxr.n_particles": 7}, + ) + + @pytest.mark.parametrize("key", ["end_element", "start_element"]) + def test_flat_builder_spelling_is_rejected(self, key): + with pytest.raises(ValueError, match="does not say which stage"): + _route_kwargs( + [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]], {key: "TD11"} + ) + def test_dotted_form_accepts_end_ele_per_stage(self): entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]] routed = _route_kwargs( @@ -270,7 +285,7 @@ class TestElementNameCase: Tao is case-insensitive so lower case would appear to work, but IMPACT's impact.ele[...] is a dict lookup and the registry's own handoff_points and - element_aliases lookups would silently miss. + handoff_points lookups would silently miss. """ @pytest.mark.parametrize("given", ["OTR4", "otr4", "Otr4", "oTr4"]) From cccc8481224f76f90365c243d475a42d5e04e1ba Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Fri, 4 Sep 2026 11:44:16 -0700 Subject: [PATCH 10/15] FACET models --- docs/model_registry_usage.md | 80 ++++++++++++++++++---- virtual_accelerator/models/facet2.py | 46 ++++++++++--- virtual_accelerator/registry/models.py | 56 +++++++++++++++ virtual_accelerator/tests/test_registry.py | 32 +++++++-- 4 files changed, 187 insertions(+), 27 deletions(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index 7285637..eb21ec8 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -1,6 +1,9 @@ # Virtual Accelerator (VA) Registry — README ## Overview -The virtual_accelerator.registry module provides a unified interface for loading, configuring, and chaining accelerator simulation models for the LCLS copper linac (CU). Models can be used standalone or chained together to simulate the full beamline from cathode to end. +The virtual_accelerator.registry module provides a unified interface for loading, configuring, and chaining accelerator simulation models for the LCLS copper linac (CU) and FACET-II. Models can be used standalone or chained together to simulate the full beamline from cathode to end. + +Building a model needs the matching lattice checkout on the environment: +`$LCLS_LATTICE` for the `*_cu_*` models, `$FACET2_LATTICE` for the `*_f2*` ones. ### Installation ```bash @@ -12,6 +15,7 @@ pip install git+https://github.com/slaclab/virtual-accelerator.git from virtual_accelerator.registry import ( get_model, models_available, + list_models, list_handoff_points, common_handoff_points, ) @@ -23,10 +27,23 @@ Print all registered models and their descriptions: ```python >>> print(models_available) -impact_cu_inj IMPACT-T LCLS injector, cathode -> YAG03 -bmad_cu_hxr Bmad CU-HXR linac, injector handoff -> END -surrogate_cu_inj NN LCLS injector surrogate, cathode -> OTR2 -cheetah_cu_hxr Cheetah nc_hxr, cathode -> END +impact_cu_inj IMPACT-T LCLS injector, cathode -> YAG03 +bmad_cu_hxr Bmad CU-HXR linac, injector handoff -> END +surrogate_cu_inj NN LCLS injector surrogate, cathode -> OTR2 +cheetah_cu_hxr Cheetah nc_hxr, cathode -> END +impact_f2e_inj IMPACT-T FACET-II injector, cathode -> PR10241 +surrogate_f2e_inj NN FACET-II injector surrogate, cathode -> PR10241 +bmad_f2_elec Bmad FACET-II e- linac, injector handoff -> END +``` + +Filter by facility or engine: + +```python +>>> list_models(facility="facet2") +['impact_f2e_inj', 'surrogate_f2e_inj', 'bmad_f2_elec'] + +>>> list_models(engine="bmad") +['bmad_cu_hxr', 'bmad_f2_elec'] ``` ### Handoff Points @@ -36,16 +53,24 @@ Each model exposes a set of suggested handoff points — named locations where b >>> for m in ["impact_cu_inj", "bmad_cu_hxr", "surrogate_cu_inj", "cheetah_cu_hxr"]: ... print(m, list_handoff_points(m)) -impact_cu_inj ('YAG02', 'YAG03') -bmad_cu_hxr ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') -surrogate_cu_inj ('OTR2',) -cheetah_cu_hxr () +impact_cu_inj ('YAG02', 'YAG03') +bmad_cu_hxr ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') +surrogate_cu_inj ('OTR2',) +cheetah_cu_hxr () +impact_f2e_inj ('PR10241',) +surrogate_f2e_inj ('PR10241',) +bmad_f2_elec ('CATHODEF', 'PR10241', 'L0AFEND', 'PR10465', 'PR10471', 'PR10571', 'PR10711', 'END') ``` -`CATHODE` appears only for `bmad_cu_hxr`, the one model whose start is configurable — you -can slice it from the front of the machine with `start_ele="CATHODE"`. The injector models -always begin at the cathode and cannot be told otherwise, so listing it there would -advertise something you cannot pass. +The cathode appears only on the Bmad models, whose start is configurable — you can slice +from the front of the machine with `start_ele="CATHODE"` (LCLS) or `start_ele="CATHODEF"` +(FACET; the element carries an `F` suffix in that lattice). The injector models always +begin at the cathode and cannot be told otherwise, so listing it there would advertise +something you cannot pass. + +FACET's injectors list only `PR10241`, which is what restricts every FACET chain to that +one handoff plane. `bmad_f2_elec` still lists the downstream screens so they remain usable +as `end_ele`. ### Shared Handoff Points `common_handoff_points()` returns the locations two models can actually hand over at @@ -61,6 +86,9 @@ upstream of it. >>> common_handoff_points("cheetah_cu_hxr", "bmad_cu_hxr") () + +>>> common_handoff_points("impact_f2e_inj", "bmad_f2_elec") +('PR10241',) ``` ### Loading a Single Model @@ -127,6 +155,32 @@ impact_cu_inj → bmad_cu_hxr — hand off at YAG03: 2.3638227838794528e-07 ``` +### FACET-II +Both FACET chains hand off at PR10241, so `handoff_loc` can be left out: + +```python +>>> m = get_model(["surrogate_f2e_inj", "bmad_f2_elec"], end_ele="PR10711", n_particles=2000) +>>> m = get_model(["impact_f2e_inj", "bmad_f2_elec"], end_ele="PR10711", n_particles=200) +``` + +Standalone: + +```python +>>> get_model("bmad_f2_elec", end_ele="PR10711", track_beam=True) + +``` + +Anything other than PR10241 is refused, as is mixing facilities: + +```python +>>> get_model(["impact_f2e_inj", "bmad_f2_elec"], handoff_loc="PR10571") +ValueError: 'PR10571' is not a shared handoff point for 'impact_f2e_inj' -> 'bmad_f2_elec'. +Available: PR10241 + +>>> get_model(["impact_cu_inj", "bmad_f2_elec"], handoff_loc="YAG03") +ValueError: Cannot stage 'impact_cu_inj' (lcls) onto 'bmad_f2_elec' (facet2): different facilities. +``` + ### Handoff Validation `handoff_loc` must be a point both stages share. Anything else is rejected before any model is built, so you do not pay for an IMPACT run to find out: diff --git a/virtual_accelerator/models/facet2.py b/virtual_accelerator/models/facet2.py index eef6b5b..c8bece3 100644 --- a/virtual_accelerator/models/facet2.py +++ b/virtual_accelerator/models/facet2.py @@ -131,6 +131,42 @@ def get_facet_bmad_model( return model +def get_facet_injector_surrogate_model( + n_particles: int = 10000, surrogate_inputs: str = "machine" +): + """ + Get the surrogate model for the FACET-II injector to PR10241. + + Parameters + ---------- + n_particles: int, optional + Number of particles to generate in the output beam. Default is 10000. + surrogate_inputs: str, optional + Input for the surrogate model, either "machine" or "sim". Default is "machine". + + Returns + ------- + BeamOutputModel + Injector surrogate whose output beam is defined at PR10241. + + Notes + ----- + ``t0``, ``p0c`` and ``z0`` describe the PR10241 handoff plane -- ``z0`` is that + element's s position. They would need to move per-plane if a second FACET + handoff location is ever used. + """ + from facet2_inj_ml_model import load_model + from virtual_accelerator.surrogates.beam_output import BeamOutputModel + + return BeamOutputModel( + load_model(surrogate_inputs), + n_particles=n_particles, + t0=3.15391398e-09, + p0c=6.3e06, + z0=0.9420843, + ) + + def get_facet_staged_model(n_particles=10000, surrogate_inputs="machine", **kwargs): """ Get the StagedModel for the FACET-II lattice from PR10241 to END, with an injector surrogate model. @@ -149,16 +185,10 @@ def get_facet_staged_model(n_particles=10000, surrogate_inputs="machine", **kwar StagedModel Instance of the StagedModel for the FACET-II lattice. """ - from facet2_inj_ml_model import load_model - from virtual_accelerator.surrogates.beam_output import BeamOutputModel from lume.staged_model import StagedModel - injector_surrogate = BeamOutputModel( - load_model(surrogate_inputs), - n_particles=n_particles, - t0=3.15391398e-09, - p0c=6.3e06, - z0=0.9420843, + injector_surrogate = get_facet_injector_surrogate_model( + n_particles=n_particles, surrogate_inputs=surrogate_inputs ) tmp = tempfile.NamedTemporaryFile(suffix=".h5") diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index e1ccc2b..20db6d7 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -141,4 +141,60 @@ def configurable_extent(self) -> bool: handoff_points=(), shared_params=frozenset({"n_particles"}), ), + "impact_f2e_inj": ModelEntry( + name="impact_f2e_inj", + description="IMPACT-T FACET-II injector, cathode -> PR10241", + facility="facet2", + engine="impact", + builder="virtual_accelerator.models.facet2:get_facet_impact_model", + extras=("impact",), + params={"n_particles": 100, "end_element": "PR10241"}, + handoff_points=("PR10241",), + end_param="end_element", + default_end="PR10241", + shared_params=frozenset({"n_particles"}), + ), + "surrogate_f2e_inj": ModelEntry( + name="surrogate_f2e_inj", + description="NN FACET-II injector surrogate, cathode -> PR10241", + facility="facet2", + engine="surrogate", + builder="virtual_accelerator.models.facet2:get_facet_injector_surrogate_model", + extras=("surrogate",), + params={"n_particles": 10000, "surrogate_inputs": "machine"}, + handoff_points=("PR10241",), + default_end="PR10241", + shared_params=frozenset({"n_particles"}), + ), + "bmad_f2_elec": ModelEntry( + name="bmad_f2_elec", + description="Bmad FACET-II e- linac, injector handoff -> END", + facility="facet2", + engine="bmad", + builder="virtual_accelerator.models.facet2:get_facet_bmad_model", + extras=("bmad",), + params={ + "start_element": "L0AFEND", + "end_element": "END", + "track_beam": False, + "custom_beam_path": None, + }, + # Both FACET injectors end at PR10241, so that is the only shared handoff. + # The screens downstream of it are listed so they can be used as end_ele. + # CATHODEF, not CATHODE -- FACET's cathode element carries the "F" suffix. + handoff_points=( + "CATHODEF", + "PR10241", + "L0AFEND", + "PR10465", + "PR10471", + "PR10571", + "PR10711", + "END", + ), + start_param="start_element", + end_param="end_element", + default_start="L0AFEND", + default_end="END", + ), } diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index d27c123..9f17f16 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -28,9 +28,15 @@ def test_repr_is_aligned_table(self): assert len(text.splitlines()) == len(MODELS) def test_filter_by_engine_and_facility(self): - assert list_models(engine="bmad") == ["bmad_cu_hxr"] - assert set(list_models(facility="lcls")) == set(MODELS) - assert list_models(facility="facet2") == [] + assert list_models(engine="bmad") == ["bmad_cu_hxr", "bmad_f2_elec"] + assert list_models(facility="facet2") == [ + "impact_f2e_inj", + "surrogate_f2e_inj", + "bmad_f2_elec", + ] + assert set(list_models(facility="lcls")) | set( + list_models(facility="facet2") + ) == set(MODELS) def test_handoff_points_are_lattice_ordered(self): diags = list_handoff_points("bmad_cu_hxr") @@ -46,11 +52,18 @@ def test_impact_lists_only_its_standard_extent(self): assert absent not in diags def test_cathode_is_only_listed_where_it_is_usable(self): - # bmad accepts start_ele="CATHODE"; the injectors have a fixed start, so - # listing it there would advertise something that cannot be passed. + # The bmad models accept a cathode start; the injectors have a fixed start, + # so listing it there would advertise something that cannot be passed. + # FACET's element is CATHODEF, not CATHODE. assert "CATHODE" in list_handoff_points("bmad_cu_hxr") + assert "CATHODEF" in list_handoff_points("bmad_f2_elec") for fixed in ("impact_cu_inj", "surrogate_cu_inj", "cheetah_cu_hxr"): - assert "CATHODE" not in list_handoff_points(fixed) + assert not {"CATHODE", "CATHODEF"} & set(list_handoff_points(fixed)) + + def test_facet_handoff_is_restricted_to_pr10241(self): + for inj in ("impact_f2e_inj", "surrogate_f2e_inj"): + assert list_handoff_points(inj) == ("PR10241",) + assert common_handoff_points(inj, "bmad_f2_elec") == ("PR10241",) class TestEntryIntegrity: @@ -89,6 +102,10 @@ def test_rejects_unavailable_screen(self): with pytest.raises(ValueError, match="not an available end screen"): get_model("impact_cu_inj", end_ele="OTR4") + def test_rejects_cross_facility_staging(self): + with pytest.raises(ValueError, match="different facilities"): + get_model(["impact_cu_inj", "bmad_f2_elec"], handoff_loc="YAG03") + def test_rejects_impact_as_downstream_stage(self): with pytest.raises(ValueError, match="only start at the cathode"): get_model(["bmad_cu_hxr", "impact_cu_inj"], handoff_loc="OTR2") @@ -204,6 +221,9 @@ def test_ordered_by_lattice_position(self): def test_no_shared_points_gives_empty_tuple(self): assert common_handoff_points("cheetah_cu_hxr", "bmad_cu_hxr") == () + def test_cross_facility_pairs_share_nothing(self): + assert common_handoff_points("impact_cu_inj", "bmad_f2_elec") == () + def test_requires_at_least_two_models(self): with pytest.raises(ValueError, match="at least two models"): common_handoff_points("bmad_cu_hxr") From 848517bd48eb9128d04f6c119ab6386e51ae3bf8 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Fri, 4 Sep 2026 15:42:02 -0700 Subject: [PATCH 11/15] PR ready --- docs/model_registry_design.md | 374 ----------------------- docs/registry_decisions_and_questions.md | 199 ------------ 2 files changed, 573 deletions(-) delete mode 100644 docs/model_registry_design.md delete mode 100644 docs/registry_decisions_and_questions.md diff --git a/docs/model_registry_design.md b/docs/model_registry_design.md deleted file mode 100644 index 5b78659..0000000 --- a/docs/model_registry_design.md +++ /dev/null @@ -1,374 +0,0 @@ -# Model registry + unified `get_model` — design proposal - -Status: **implemented for LCLS** in `virtual_accelerator/registry/`. FACET-II entries are not -registered yet; §7 (screen PV prefixes) should land before they are. - -Goal: one entry point for every model, a registry of what exists and how to configure it, and -enough per-model metadata to stitch stages together safely. - -Scope: **any lattice element may be a start/end point**, with `CATHODE` as the canonical name for -models starting at the front of the machine. No element catalog is needed, because the only element -names that differ between engines are the gun and the RF cavity segments (see §1), so a small -per-model alias dict covers it. Screens are enumerated per model as a discovery aid and a typo check. - -Reference plane is the element **entrance** — see §5. - ---- - -## 1. Evidence: why a list of names is enough - -Checked against the real lattice files, not assumed. - -**Diagnostics already agree between Bmad and IMPACT, in both facilities.** The suspected -`YAG03`/`YAG3` mismatch does not exist. Confirmed screen-by-screen: - -| facility | diagnostic | Bmad | IMPACT-T | -|---|---|---|---| -| LCLS | YAG02, YAG03, OTR1, OTR2 | same | same | -| FACET | PR10241, PR10465, PR10471, PR10571 | same | same | - -Bmad aliases match `utils/cu_hxr_profmon_info.yaml` exactly (`OTR2[alias]=OTRS:IN20:571`, -`YAG03[alias]=YAGS:IN20:351`), so the profmon YAMLs already are a de-facto standard-name list — -this proposal just makes them addressable per model. - -Note that the *element names* agreeing is separate from the *PVs* agreeing — the latter is broken -for FACET, see §7. - -**The divergences that exist are all in non-diagnostic elements**, i.e. things that are never -handoff points: - -| element | Bmad | IMPACT-T | -|---|---|---| -| gun / cathode | `CATHODE` / `CATHODEF` | `GUN` / `GUNF` | -| L0A cavity | `L0A` (one lcavity) | `L0A_entrance`, `L0A_body_1`, `L0A_body_2`, `L0A_exit` | - -No alias dict is needed. The gun would have been the one entry (`CATHODE` → IMPACT `GUN`) but it is -unreachable: IMPACT has no start parameter to pass it to. Cavity segmentation is one-to-many and -cavities are not sensible handoff planes, so they are simply not registered as handoff points. - -Two facts that don't need machinery but should be written down: - -- IMPACT z and Bmad slice starts are both the element *entrance*; the lattice CSV's `SumL` is the - element *centre*. For a quad these differ by half the length (FACET `QA10361`: 4.3103893 vs - 4.4123893). This is why the entrance is the reference plane and the CSV's `SumL` is not used for it. -- FACET IMPACT geometry disagrees with the Bmad geometry by up to ~3.9 mm at `PR10571`, growing - downstream. LCLS agrees exactly. Real, but a physics caveat rather than something the registry - should police. - ---- - -## 2. Registry - -One table. Models only; no separate element registry. - -```python -@dataclass(frozen=True) -class ModelEntry: - name: str # "bmad_cu_hxr" - description: str # what models_available prints - facility: str # "lcls" | "facet2" - engine: str # "bmad" | "impact" | "surrogate" | "cheetah" - builder: str # "virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model" - extras: tuple[str, ...] # pip extras needed, e.g. ("bmad",) - params: dict[str, Any] # param name -> default, for validation + discovery - shared_params: frozenset[str] # must match across stages, e.g. {"n_particles"} - - handoff_points: tuple[str, ...] # suggested elements, ORDERED by lattice position - start_param: str | None # builder kwarg for start, None if fixed - end_param: str | None # builder kwarg for end, None if fixed - default_start: str | None - default_end: str | None -``` - -`start_param` / `end_param` replace the earlier `can_start_anywhere` flag: they carry the same -information (`None` means the extent is fixed) while also recording what each builder *calls* the -parameter, which differs between engines. - -`handoff_points` is a discovery aid and typo check, **not** a restriction — any lattice element may -be used as a start/end point. See §5 for how duplicate PVs at the handoff are resolved. - -Three choices worth calling out: - -**`builder` is a `"module:function"` string, not a callable.** The existing builders lazily import -`pytao` / `torch` / `impact` inside the function body via `import_optional`. A real callable in the -registry would force importing every engine module just to import the registry, breaking -`models_available` for anyone without all extras installed. - -**`handoff_points` is ordered by lattice position.** Ordering is then checkable by list index -instead of by metres, which is what lets the design drop positions entirely. - -**No data files, no generator script.** The lists are hand-maintained Python literals in -`registry/models.py`. They are short, they change only when a lattice model changes, and a diff is -readable. Revisit only if they start drifting from the lattice. - -``` -virtual_accelerator/registry/ -├── __init__.py # get_model, models_available, list_handoff_points -└── models.py # the ModelEntry table -``` - ---- - -## 3. Model catalog - -| registry name | facility | engine | builder | key params | suggested handoff points | -|---|---|---|---|---|---| -| `impact_cu_inj` | lcls | impact | `get_cu_inj_impact_model` | `n_particles=100`, `end_element="YAG03"` | YAG02, YAG03 | -| `bmad_cu_hxr` | lcls | bmad | `get_cu_hxr_bmad_model` | `start_element="OTR2"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | CATHODE, all 12 profmon screens, END | -| `surrogate_cu_inj` | lcls | surrogate | `get_cu_hxr_injector_surrogate_model` | `n_particles=1000` | OTR2 (fixed extent) | -| `cheetah_cu_hxr` | lcls | cheetah | `get_cu_hxr_cheetah_model` | `n_particles=1000` | — | - -Naming convention: `_`. - -**Not yet registered** — FACET-II, pending §7: - -| registry name | facility | engine | builder | key params | suggested handoff points | -|---|---|---|---|---|---| -| `impact_f2e_inj` | facet2 | impact | `get_facet_impact_model` | `n_particles=100`, `end_element="PR10571"` | PR10241, PR10465, PR10471, PR10571 | -| `bmad_f2_elec` | facet2 | bmad | `get_facet_bmad_model` | `start_element="L0AFEND"`, `end_element="END"`, `track_beam=False`, `custom_beam_path=None` | PR10241, PR10465, PR10471, PR10571, PR10711 | -| `surrogate_f2e_inj` | facet2 | surrogate | ⚠ **does not exist yet** | `n_particles=10000`, `surrogate_inputs="machine"` | PR10241 (fixed end) | - -Notes on the lists: - -- `impact_cu_inj` omits `OTRH1`/`OTRH2` (laser heater is unmodeled in the deck: `!!! Unmodeled: - Laser Heater from 9.076892 m to 10.690580 m`) and `OTR3`/`YAG01` (lines commented out — - `YAG01` is marked `!!! Broken:`). `OTR4` is past `stop_1` at z=16.5. -- `impact_f2e_inj` includes `PR10571` because the file actually loaded is - `ImpactT_template.in` (per `ImpactT.yaml`'s `input_file:` key). The checked-in `ImpactT.in` is a - stale truncated artifact stopping at z=12.0 and omitting it — reading that file gives the wrong - answer about available stop points. Worth an upstream issue. -- `bmad_f2_elec` defaults `start_element="L0AFEND"`, a marker rather than a screen. It exists in - both engines (Bmad superimposed marker, IMPACT write-beam element) so it is a legitimate handoff - plane. Since any element is now allowed, markers need no special treatment — they are simply listed - in `handoff_points` where useful. - -**Prerequisite:** `surrogate_f2e_inj` has no standalone builder — the `BeamOutputModel` is built -inline inside `get_facet_staged_model`. It needs extracting to -`get_facet_injector_surrogate_model()` to mirror `get_cu_hxr_injector_surrogate_model`, otherwise -it can't be a registry entry. - -`get_cu_hxr_staged_model` / `get_facet_staged_model` become thin back-compat wrappers over -`get_model([...])`. - ---- - -## 4. `get_model` - -```python -def get_model( - spec: str | Sequence[str], - *, - handoff_loc: str | Sequence[str] | None = None, - start_ele: str | None = None, - end_ele: str | None = None, - **kwargs, -) -> LUMEModel: -``` - -```python -# single -model = get_model("bmad_cu_hxr", end_ele="OTR4", track_beam=True) - -# staged, explicit handoff -model = get_model(["impact_cu_inj", "bmad_cu_hxr"], - handoff_loc="YAG03", end_ele="OTR4", n_particles=1000) - -# staged, handoff inferred from the surrogate's fixed end (OTR2) -model = get_model(["surrogate_cu_inj", "bmad_cu_hxr"], end_ele="OTR4", n_particles=10000) -``` - -Every example above is registered and working today, except that the `impact_cu_inj` stage needs -the IMPACT-T executable (`conda install -c conda-forge impact-t`) on top of the Python package. - -FACET-II is **not** registered yet, so this raises `KeyError` for now — see §3 and §7: - -```python -model = get_model(["impact_f2e_inj", "bmad_f2_elec"], handoff_loc="L0AFEND") # not yet -``` - -`start_ele` / `end_ele` on a staged call are the *overall* extent — first stage's start, last -stage's end. Interior extents come from `handoff_loc`. - -### Discovery - -```python -from virtual_accelerator.registry import models_available, list_handoff_points - -print(models_available) -# impact_cu_inj IMPACT-T LCLS injector, cathode -> OTR2 -# bmad_cu_hxr Bmad CU-HXR, gun -> undulator/dump -# ... - -list_handoff_points("bmad_cu_hxr") -# ('YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', -# 'OTR11', 'OTR12', 'OTR21', 'OTRDMP') -``` - -All 12 screens in `cu_hxr_profmon_info.yaml` were verified present in the `cu_hxr` lattice, in the -order shown. - -### Kwarg routing - -Because the registry declares each model's params, routing is a lookup, not signature -introspection: - -1. Param in `shared_params` (e.g. `n_particles`) → sent to every stage that declares it, and - the per-stage form is **rejected**: the beam flows through the stages, so differing values - would break a physical invariant rather than configure anything. -2. Declared by exactly one stage → routed there. -3. Declared by more than one stage and not shared → `ValueError` naming the candidates. The - builder spellings `end_element`/`start_element` are rejected flat, since they do not say - which stage they mean; use `end_ele`/`start_ele`, or qualify per stage. -4. `"."` always wins. -5. Matching no declared param → rejected with a suggestion, not silently forwarded. - -`track_beam=True` is forced on **every** stage that declares it, regardless of what the user -passes. An earlier draft said "non-terminal stages only", which is wrong and was caught by testing: -a non-final stage must *produce* `final_particles`, but a non-first stage must also *accept* -`initial_particles`, and `lume_bmad.model` raises `Cannot set initial_particles when track_type is -not 'beam'` otherwise. In a two-stage surrogate → Bmad chain the Bmad stage is terminal and still -needs it. - ---- - -## 5. Compatibility checking - -`StagedModel.validate_lume_model_instances` already checks mixin presence and duplicate variable -names — but only *after* both models are constructed, and `build_impact_model` calls -`impact.run()` during construction. So a duplicate-variable error costs a full IMPACT run before -it surfaces. The registry pre-validates from metadata alone, before instantiating anything. - -All checks are name-based. No positions involved. - -**C1 — same facility.** Staging `bmad_cu_hxr` onto `impact_f2e_inj` is rejected immediately. - -**C2 — handoff is a legal point in both stages.** `handoff_loc` must be in the upstream stage's -`diagnostics` and in the downstream stage's. Plus the downstream stage must have a `start_param`, -which is what makes `get_model(["bmad_cu_hxr", "impact_cu_inj"], ...)` fail with "IMPACT models can -only start at the cathode" instead of something inscrutable from inside `set_stop_location`. - -**C3 — duplicate PVs at the handoff are removed from the downstream stage.** This is the most -important rule in practice, and it was revised on 2026-09-02. - -Both stages include the handoff element, so both publish its PVs. IMPACT's `set_stop_location` -prunes to `s <= stop`, so a model stopped at `YAG03` *keeps* the screen and publishes its six -`YAGS:IN20:351:*` image PVs; Bmad sliced from `YAG03` publishes the same six. `StagedModel` then -rejects the pair as duplicates — after paying for a full IMPACT run. - -Resolution: compute the overlap after construction and unregister it from the **downstream** stage. -The upstream stage owns those PVs because it is the stage that actually tracks the beam to that -plane. `lume.actions` exposes `unregister_action_variable(name)` and `supported_variables` is a -property over `_action_variable_by_name`, so this is clean public API. Measured on a real Bmad model: -318 vars → 312 after removal, model still functional. - -This lets the handoff be named for the real element (`YAG03`) rather than the drift after it. An -earlier draft instead moved the downstream start to `element_after[handoff]` (`DL02A2`); that dict is -gone. It also explains `start_element="DL02A2"` in `examples/staged_example.ipynb` — a workaround for -this collision, no longer needed. - -**Safety rule:** a *writable* overlap raises rather than being dropped. Writable overlap means both -stages are driving the same magnet — extents overlapping rather than meeting at a plane — and -dropping it downstream would leave that stage tracking with a stale value. The real YAG03 overlap was -measured to be entirely read-only, so this does not fire in normal use. - -**Caveat:** the only staged path runnable without the IMPACT-T binary does not exercise this. -Measured overlap between `surrogate_cu_inj` and `bmad_cu_hxr` is **zero** — the surrogate publishes -`OTRS:IN20:571:XRMS`/`YRMS` while Bmad publishes `Image:*`/`RESOLUTION`/`X`. - -**C4 — beam handoff mechanics.** One place, in the registry, resolving an existing inconsistency: -`get_facet_staged_model` writes `final_particles` to a `NamedTemporaryFile` and passes it as -`custom_beam_path`, while `get_cu_hxr_staged_model` does neither and relies solely on the -`FinalParticlesMixIn` wiring in `StagedModel._set`. One of those is redundant or one is a latent -bug; the registry should own this in exactly one place. - -Also worth noting for C4: `get_facet_staged_model` hardcodes `t0=3.15391398e-09`, `p0c=6.3e06`, -`z0=0.9420843`. That `z0` is exactly `PR10241`'s s-position — these are handoff-plane quantities, -so if a second FACET handoff plane is ever used they will need to move somewhere per-plane rather -than staying inline. - ---- - -## 6. Deliberately deferred - -Cut from the previous draft, with the trigger for reconsidering each: - -| deferred | add it when | -|---|---| -| Element catalog generated from `lcls_elements.csv` | a handoff point is needed that isn't a diagnostic, or the hand-maintained lists start drifting from the lattice | -| s-positions on handoff points | we want to *numerically* verify a handoff rather than trust name equality | -| Tolerance-based position agreement check | ditto — this is where the FACET ~3.9 mm discrepancy would resurface | -| Physical-extent overlap check between stages | `StagedModel`'s duplicate-variable check proves insufficient in practice | -| PV names accepted as `start_ele` / `end_ele` | a control-room user asks for it; cheap to add later | - ---- - -## 7. Screen PV prefixes — resolved, and a bug it exposes - -**Resolved rule (supervisor, 2026-09-01): `PROF:` is correct for FACET VAs; `YAGS:`/`OTRS:` are -correct for LCLS / LCLS-II VAs.** - -This explains an asymmetry in the current code that otherwise looks arbitrary: -`get_facet_bmad_model` carries a `custom_aliases` dict while `get_cu_hxr_bmad_model` carries none. -FACET needs it precisely *because* the lattice aliases (`YAGS:/OTRS:IN10:*`) are wrong for VA -purposes; LCLS needs none because its lattice aliases are already right. - -### How each engine currently derives a screen's PV - -- **Bmad** — `bmad/variables.py:375`: `base_pv = tao.ele(screen_name).head.alias`, i.e. the Tao - element alias. `build_bmad_model` applies `custom_aliases` (factory.py:72-78) *before* - `get_variables` (factory.py:86), so the override ordering is correct. -- **IMPACT** — `impact/variables.py`: `alias_dict[element_name]`, where `alias_dict` is - `Element -> Control System Name` from `bmad/conversion/from_oracle/lcls_elements.csv`. There is - **no** `custom_aliases` equivalent on this path. -- **Neither** reads the `name:` field of `utils/*_profmon_info.yaml`. Bmad uses that file only for - `shape` and `pixel_size` (variables.py:379-382). So the `name:` field is currently dead data — - even though it holds the correct value in every case. - -### Consequence: FACET screen PVs are wrong today - -| screen | Bmad yields | IMPACT yields | correct (per rule) | -|---|---|---|---| -| PR10241 | `PROF:IN10:241` ✓ | `YAGS:IN10:241` ✗ | `PROF:IN10:241` | -| PR10465 | `OTRS:IN10:465` ✗ | `OTRS:IN10:465` ✗ | `PROF:IN10:465` | -| PR10471 | `OTRS:IN10:471` ✗ | `OTRS:IN10:471` ✗ | `PROF:IN10:471` | -| PR10571 | `PROF:IN10:571` ✓ | `OTRS:IN10:571` ✗ | `PROF:IN10:571` | -| PR10711 | `PROF:IN10:711` ✓ | `OTRS:IN10:711` ✗ | `PROF:IN10:711` | - -FACET Bmad is wrong for 2 of 5 screens (`custom_aliases` simply omits `PR10465`/`PR10471`); FACET -IMPACT is wrong for all 5. This directly affects `impact_f2e_inj` — the model about to be staged. - -LCLS is consistent across all three sources (lattice alias, CSV `Control System Name`, and -`cu_hxr_profmon_info.yaml`), so it is unaffected either way. - -### Proposed fix — no new data structure - -Make `utils/*_profmon_info.yaml` the single source of truth for screen PVs, and have **both** -engines read `screen_config[name]["name"]` instead of the lattice alias / CSV. Then: - -- Every FACET screen is fixed in both engines at once, because the YAML already holds the right - values for all five. -- LCLS is a **no-op** — its YAML values already equal its lattice aliases and CSV entries. -- `custom_aliases` in `models/facet2.py` loses its screen entries entirely, keeping only - `TCY10490 -> KLYS:LI10:51` (a TCAV, not a screen, and a separate question). -- The dead `name:` field becomes load-bearing, so it can no longer silently drift. - -This is deliberately *not* the element catalog returning: it reuses a file that already exists and -already has the answer. Worth doing as a small standalone PR **before** the registry work, since -the registry shouldn't inherit a known-wrong PV mapping. - -One fragility worth fixing alongside: `alias_dict[element_name]` in `impact/variables.py` is an -unguarded dict lookup, so any element absent from the CSV raises a bare `KeyError`. - ---- - -## 8. Open questions - -1. **Field name** — `diagnostics` or `handoff_points`? The FACET default start `L0AFEND` is a - marker, not a diagnostic, so the latter is more honest. -2. **`TCY10490 -> KLYS:LI10:51`** — the one non-screen entry in FACET's `custom_aliases`. The - lattice says `TCY10490[alias]=TCAV:IN10:490` and the CSV agrees. Is the `KLYS:` override - deliberate in the same way the `PROF:` ones were? Same question, different device class. -3. **`models/runners.py` CLI** — its four hardcoded `--model` choices become `get_model(args.model)`, - which is a strict improvement but changes the accepted values. -4. **`BmadModelSpec.database_relpath` is dead config** — declared at `bmad/factory.py:24`, never - read in `build_bmad_model`. Unrelated to this work; noted while reading. diff --git a/docs/registry_decisions_and_questions.md b/docs/registry_decisions_and_questions.md deleted file mode 100644 index acf8523..0000000 --- a/docs/registry_decisions_and_questions.md +++ /dev/null @@ -1,199 +0,0 @@ -# Model registry — design decisions and open questions - -Companion to `docs/model_registry_design.md`. LCLS registry is implemented and verified against -real models; FACET-II is not registered yet. - ---- - -## Part A — Design decisions - -### A1. Any element may be a start/end point; screens are enumerated as a hint - -**Revised 2026-09-02 per supervisor.** An earlier draft restricted handoffs to diagnostics. That is -lifted: any element in the lattice may be used, with `CATHODE` as the canonical name for models -starting at the front of the machine. - -Screens are still listed per model (`handoff_points`) for two reasons: discovery, and as a typo -check — screens are enumerated exhaustively, so a screen-shaped name absent from the list is -definitely wrong and is rejected early rather than failing deep inside Tao. Anything else passes -through to the engine. There are 3323 elements in `cu_hxr`, so exhaustive validation is not on the -table. - -`CATHODE` is listed only for `bmad_cu_hxr`, the one model whose start is configurable. The -injectors always begin at the cathode and cannot be told otherwise, so an alias mapping Bmad's -`CATHODE` to IMPACT's `GUN` was added and then removed again -- it was unreachable, since IMPACT -has no start parameter to pass it to. `element_aliases` is gone with it; diagnostics already agree -between the engines, so there was nothing left for it to carry. - -### A2. Reference plane is the element ENTRANCE - -**A midpoint default was considered and rejected on evidence.** Bmad's `-slice_lattice` begins at an -element's *entrance* and cannot begin at a midpoint: - -| element | L | entrance | centre | Bmad slice begins at | -|---|---|---|---|---| -| `QE01` | 0.108 | 8.440049 | 8.494049 | **8.440049** | -| `L0A` | 3.095 | 1.459000 | 3.006622 | **1.459000** | - -IMPACT's `impact.ele[name]["s"]` is also the entrance, though IMPACT *could* express a midpoint -since `impact.stop` is a float. So a midpoint default would be implementable in IMPACT but not in -Bmad, making the two engines silently disagree by half an element length — the exact class of bug -this design is meant to prevent. Entrance is the only plane both engines express identically. - -The separate issue of identical elements sitting at slightly different `s` in the two lattices -(FACET differs by up to ~3.9 mm, growing downstream; LCLS agrees exactly) is unaffected by this -choice and remains unpoliced — see A3. - -### A2b. Duplicate PVs at the handoff are unregistered from the downstream stage - -**Revised 2026-09-02 per supervisor.** An earlier draft moved the downstream start to the *next* -element (an `element_after` dict) to dodge the collision. That dict is now gone. - -Both stages include the handoff element, so both publish its PVs and `StagedModel` would reject the -pair as duplicates. Instead the overlap is computed after construction and removed from the -downstream stage, which lets the handoff be expressed as the real element name (`YAG03`, not -`DL02A2`). `lume.actions` provides `unregister_action_variable(name)` and `supported_variables` is a -property over `_action_variable_by_name`, so this is clean public API. Verified on a real Bmad model: -318 vars → 312 after removing the six `YAGS:IN20:351:*` PVs, model still functional. - -The upstream stage owns the handoff PVs because it is the stage that actually tracks the beam to -that plane. - -**One safety rule added:** a *writable* overlap raises instead of being dropped. Writable overlap -means both stages are driving the same magnet — the extents overlap rather than meeting at a plane — -and silently dropping it downstream would leave that stage tracking with a stale value. Measured -that the real YAG03 overlap is entirely read-only, so the rule does not fire in normal use. - -### A3. No positions stored — compatibility is name equality plus list-index ordering - -`handoff_points` lists are ordered by lattice position, so ordering is checkable by index without -storing metres. - -**Consequence:** the FACET IMPACT geometry disagrees with the Bmad geometry by up to ~3.9 mm at -`PR10571`, growing downstream (LCLS agrees exactly, to nine decimals). This is documented as a -physics caveat and deliberately **not** policed by the registry. Trigger for revisiting is recorded -in the design doc's "deliberately deferred" table. - -### A4. Builders referenced as `"module:function"` strings, resolved lazily - -The existing builders import `pytao` / `torch` / `impact` inside their function bodies via -`import_optional`. Holding real callables in the registry would force importing every engine module -just to import the registry. As strings, `models_available` and `list_handoff_points` work with **zero -optional dependencies installed** — verified. - -### A5. Registry declares each model's params; unknown kwargs are rejected - -Kwarg routing across stages is a table lookup, not signature introspection, so error messages can -name the candidate stages. Trade-off: adding a parameter to a builder means also adding it to the -registry. That cost buys real error messages and a truthful `models_available`. - -Routing rules: *shared* params (`n_particles`) go to every stage declaring them and cannot be set -per stage, since the beam flows through and differing values would break a physical invariant; a -param declared by exactly one stage routes there; declared by more than one is an error naming the -candidates; `"."` targets one stage. The builders' own `end_element` / -`start_element` are rejected flat in favour of `end_ele` / `start_ele`. - -### A6. Element names normalised to upper case at the API boundary - -Found by testing. Without it, `handoff_loc='yag03'` silently bypassed the A2b collision check and -started Bmad *at* the screen, resurfacing the duplicate-PV failure. Also fixes IMPACT, where -`impact.ele[...]` is a case-sensitive dict lookup. - -### A7. `track_beam=True` forced on every stage in a chain - -An earlier draft said non-terminal stages only. Wrong, and caught by testing: a non-final stage must -*produce* `final_particles`, but a non-first stage must also *accept* `initial_particles`, and -`lume_bmad` raises `Cannot set initial_particles when track_type is not 'beam'`. In a two-stage -surrogate → Bmad chain the Bmad stage is terminal and still needs it. - -### A8. Hand-maintained Python literals — no generated data files, no CI regenerator - -Per the "start simple" steer. The lists are short and change only when a lattice model changes, so a -diff is readable. Revisit if they drift from the lattice. - -### A9. Registry naming convention `_` - -`impact_cu_inj`, `bmad_cu_hxr`, `surrogate_cu_inj`, `cheetah_cu_hxr`. See Q6. - -### A10. Existing builders keep working; migration is deferred - -`get_cu_hxr_staged_model` etc. are untouched so far. Proven equivalent to the registry path: -identical variable sets, identical Bmad slice, and beam moments differing by 1.6e-05 relative — -less than the 4.6e-05 run-to-run variation of the existing builder against *itself* (distgen -sampling noise). Migration to thin wrappers is queued behind end-to-end verification of the -IMPACT → Bmad path. - ---- - -## Part B — Questions - -### Blocking now - -**Q1. `TCY10490 → KLYS:LI10:51` — deliberate, like the `PROF:` overrides?** -`models/facet2.py` overrides this alias, but the FACET lattice says `TCY10490[alias]=TCAV:IN10:490` -and the elements CSV agrees. This is the same shape of question as the screen-prefix one, but for a -TCAV rather than a screen. *Blocks the FACET PV fix, since that PR touches `custom_aliases`.* - -**Q2. ANSWERED 2026-09-02.** `PROF:` applies to both engines, and more importantly: *"the csv -file usually doesn't have correct values"*. So `lcls_elements.csv` is **not** authoritative for PVs. -That confirms the fix — make `utils/*_profmon_info.yaml` the single source of truth for screen PVs -for both Bmad and IMPACT. It already holds correct values for both facilities, in a `name:` field -that neither engine currently reads. LCLS is a no-op. Note this narrows the CSV's role to element -*names* and `SumL`; its `Control System Name` column should not be trusted. - -**Q3. ANSWERED 2026-09-02.** `element_after` is gone — replaced by unregistering overlapping -variables from the downstream stage (A2b). Simpler and more general. - -### Naming and API — cheap now, expensive later - -**Q4. Field name: `diagnostics` or `handoff_points`?** -`L0AFEND` is a marker, not a screen, so `diagnostics` is slightly dishonest (A1 caveat). - -**Q5. Which stage should own the handoff screen?** -I gave it to the **upstream** stage, on the reasoning that it physically images the screen as its -last element. The alternative is to have IMPACT stop just *before* the screen and let the -downstream Bmad model own it. That changes which model reports the handoff-plane measurement, so it -is a physics/operations call rather than a code one. - -**Q6. Are `impact_cu_inj` / `bmad_cu_hxr` / `surrogate_cu_inj` / `cheetah_cu_hxr` the names we want -users typing?** These become the public interface and are awkward to change once notebooks and -scripts use them. - -**Q7. Should `get_model` accept PVs as `start_ele` / `end_ele`** (e.g. `"OTRS:IN20:571"` as well as -`"OTR2"`)? Cheap to add; plausibly what a control-room user reaches for first. - -### Scope - -**Q8. `runners.py` CLI back-compatibility.** Routing it through the registry changes the accepted -`--model` values (`cu_hxr_bmad` → `bmad_cu_hxr`, etc.). Alias the old names, or make a clean break? - -**Q9. What else should be registered beyond LCLS cu_hxr and FACET-II?** -`examples/cheetah_diag0_model.ipynb` builds a diag0 Cheetah model inline with no factory function, -and a `models/sc_diag0.py` exists in stale build artifacts but not in current source. Is diag0 -in scope? Are LCLS-II / nc_sxr models expected? This determines whether the flat -`_` naming holds up. - -**Q10. Who owns keeping the diagnostics lists in sync with the lattice?** -A8 chose hand-maintained literals. If upstream renames or adds a screen, nothing detects it -automatically. Acceptable, or should there be a CI check that reconciles the lists against -`$LCLS_LATTICE`? - ---- - -## Verification status - -| path | status | -|---|---| -| Discovery / validation / routing / overlap logic, no extras installed | 51 unit tests, ~1 s | -| `bmad_cu_hxr` single model | built on real lattice, correct slice and PVs | -| `surrogate_cu_inj` -> `bmad_cu_hxr` staged | built, quad set, beam propagated, real OTR4 image | -| Equivalence with `get_cu_hxr_staged_model` | identical vars and slice; within sampling noise | -| Removing overlapping variables, on a real model | 318 -> 312 vars, model still functional | -| `impact_cu_inj` -> `bmad_cu_hxr` staged | **NOT verified** -- needs the IMPACT-T executable | - -**The last row matters more after this revision.** Removing overlapping variables is now the core -handoff mechanism, and the only staged path runnable here does not exercise it: measured overlap -between `surrogate_cu_inj` and `bmad_cu_hxr` is **zero**, because the surrogate publishes -`OTRS:IN20:571:XRMS`/`YRMS` while Bmad publishes `Image:*`/`RESOLUTION`/`X`. So the mechanism is -covered by unit tests and by a direct measurement on a Bmad model, but has never run inside a real -two-stage build. Closing that needs `conda install -c conda-forge impact-t`. From 25ba2ab3c83b88bbf4977058e12143115785f11a Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Fri, 4 Sep 2026 16:41:23 -0700 Subject: [PATCH 12/15] fixing docstring --- virtual_accelerator/registry/__init__.py | 185 +++++++++++++++-------- virtual_accelerator/registry/models.py | 86 ++++++----- 2 files changed, 174 insertions(+), 97 deletions(-) diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index aed0612..30091e3 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -40,7 +40,21 @@ def __repr__(self) -> str: def list_models(facility: str | None = None, engine: str | None = None) -> list[str]: - """Names of registered models, optionally filtered by facility or engine.""" + """ + Get the names of registered models, optionally filtered. + + Parameters + ---------- + facility : str, optional + Restrict to one facility, "lcls" or "facet2". Default is None, meaning all. + engine : str, optional + Restrict to one engine, e.g. "bmad". Default is None, meaning all. + + Returns + ------- + list[str] + Registry names, in registration order. + """ return [ name for name, entry in MODELS.items() @@ -50,44 +64,62 @@ def list_models(facility: str | None = None, engine: str | None = None) -> list[ def list_handoff_points(model_name: str) -> tuple[str, ...]: - """Suggested start/end/handoff elements, in lattice order. + """ + Get the suggested start, end and handoff elements for one model. - A discovery aid, not an exhaustive list -- any lattice element may be used. + Parameters + ---------- + model_name : str + Registry name. + + Returns + ------- + tuple[str, ...] + Element names in lattice order. A discovery aid rather than an exhaustive + list -- any element in the underlying lattice may be used. + + Raises + ------ + KeyError + If ``model_name`` is not registered. """ return _entry(model_name).handoff_points CATHODE = "CATHODE" -"""Start-of-machine marker. Never a valid handoff: nothing is upstream of it.""" def common_handoff_points(*model_names: str) -> tuple[str, ...]: - """Elements every named model can hand off at, in lattice order. - - ``CATHODE`` is always excluded -- it marks the front of the machine, so - nothing can hand over to a stage that begins there. - - This is the set intersection, not the union. A union would admit planes only - one stage can reach: ``impact_cu_inj`` stops by z=16.5 m and so cannot reach - ``OTR4`` at 17.80 m, but ``bmad_cu_hxr`` lists it, so the union would wrongly - accept ``handoff_loc="OTR4"`` for that pair. + """ + Get the elements every named model can hand off at, in lattice order. Parameters ---------- - *model_names + *model_names : str Two or more registry names. Returns ------- tuple[str, ...] Shared handoff elements, ordered by the first model's lattice order. + Empty if the models share none. - Examples - -------- - >>> common_handoff_points("impact_cu_inj", "bmad_cu_hxr") - ('YAG02', 'YAG03') - >>> common_handoff_points("surrogate_cu_inj", "bmad_cu_hxr") - ('OTR2',) + Raises + ------ + ValueError + If fewer than two model names are given. + KeyError + If any name is not registered. + + Notes + ----- + ``CATHODE`` is always excluded: it marks the front of the machine, so nothing + can hand over to a stage beginning there. + + This is the set intersection, not the union. A union would admit planes only + one stage can reach -- ``impact_cu_inj`` stops by z=16.5 m and so cannot reach + ``OTR4`` at 17.80 m, but ``bmad_cu_hxr`` lists it, so a union would wrongly + accept ``handoff_loc="OTR4"`` for that pair. """ if len(model_names) < 2: raise ValueError("Need at least two models to find common handoff points.") @@ -312,16 +344,42 @@ def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) - def _strip_overlapping_variables(upstream, downstream, upstream_name, downstream_name): - """Remove variables the downstream stage shares with the upstream stage. + """ + Remove variables the downstream stage shares with the upstream stage. + Parameters + ---------- + upstream : LUMEModel + Stage that tracks the beam to the handoff plane and so owns its PVs. + downstream : LUMEModel + Stage the duplicates are removed from. Must support + ``unregister_action_variable``. + upstream_name : str + Registry name of ``upstream``, used in error messages. + downstream_name : str + Registry name of ``downstream``, used in error messages. + + Returns + ------- + list[str] + Variable names removed from ``downstream``, empty if there was no overlap. + + Raises + ------ + ValueError + If any shared variable is writable. + TypeError + If ``downstream`` cannot unregister variables. + + Notes + ----- Both stages include the handoff element, so both publish its PVs and - ``StagedModel`` would reject the pair as duplicates. The upstream stage owns - them -- it is the stage that actually tracks the beam to that plane -- so they - are unregistered from the downstream stage. + ``StagedModel`` would reject the pair as duplicates. - A *writable* overlap means something different and worse: both stages would be - driving the same magnet, and dropping it downstream would silently leave that - stage tracking with a stale value. That is a slicing error, so it raises. + A writable overlap means something different and worse: both stages would be + driving the same magnet, so their extents overlap rather than meeting at a + plane, and dropping it downstream would leave that stage tracking a stale + value. """ from lume.actions import WritableActionMixin @@ -369,55 +427,62 @@ def get_model( end_ele: str | None = None, **kwargs: Any, ): - """Build a model, or a staged chain of models, by registry name. + """ + Build a model, or a staged chain of models, by registry name. Parameters ---------- - spec - A registry name, or an ordered list of names to stage together. - handoff_loc - Element where each consecutive pair hands the beam over. Inferred from the - upstream stage's standard end when omitted. Must be a shared handoff point - of both stages -- see :func:`common_handoff_points`. A list is required for - more than two stages. - start_ele, end_ele - Overall extent. For a staged model these apply to the first and last - stage respectively; interior extents come from ``handoff_loc``. + spec : str or list[str] + A registry name, or an ordered list of names to stage together, upstream + first. + handoff_loc : str or list[str], optional + Element where each consecutive pair hands the beam over. Must be a shared + handoff point of both stages, see ``common_handoff_points``. A list is + required for more than two stages. Default is None, meaning it is inferred + from the upstream stage's standard end. + start_ele : str, optional + Element to start tracking from. For a staged model this applies to the + first stage. Default is None, meaning the model's own default. + end_ele : str, optional + Element to stop tracking at. For a staged model this applies to the last + stage; interior extents come from ``handoff_loc``. Default is None, + meaning the model's own default. **kwargs - Builder parameters. For staged models, qualify an ambiguous parameter as - ``"."``. + Builder parameters. Params listed in a model's ``shared_params`` are sent + to every stage declaring them and cannot be set per stage. Others are + routed to the single stage declaring them, or qualified as + ``"."`` when more than one does. Returns ------- LUMEModel A single model, or a ``StagedModel`` wrapping the chain. + Raises + ------ + KeyError + If a name in ``spec`` is not registered. + ValueError + If the stages cannot be chained, the handoff is not shared by both, or a + kwarg cannot be routed unambiguously. + Notes ----- For staged chains this handles two things that are easy to get wrong by hand. - **Duplicate variables at the handoff are removed automatically.** Both stages + Duplicate variables at the handoff are removed automatically. Both stages include the handoff element, so both publish its PVs -- an IMPACT model stopped - at ``YAG03`` keeps the screen (it prunes to ``s <= stop``) and so does a Bmad - model sliced from ``YAG03``. ``StagedModel`` would reject the pair as - duplicates. The upstream stage owns them, because it is the stage that tracks - the beam to that plane, so they are unregistered from the downstream stage - before the chain is assembled. A *writable* overlap raises instead: that means - both stages drive the same magnet, so their extents overlap rather than meeting - at a plane, and dropping it downstream would leave that stage tracking a stale - value. + at YAG03 keeps the screen, since it prunes to ``s <= stop``, and so does a Bmad + model sliced from YAG03. ``StagedModel`` would reject the pair as duplicates. + The upstream stage owns them, because it is the stage that tracks the beam to + that plane, so they are unregistered from the downstream stage before the chain + is assembled. + + Beam tracking is forced on for every stage that supports it, since a non-final + stage must produce ``final_particles`` and a non-first stage must accept + ``initial_particles``. - **Beam tracking is forced on for every stage that supports it**, since a - non-final stage must produce ``final_particles`` and a non-first stage must - accept ``initial_particles``. - - Examples - -------- - >>> get_model("bmad_cu_hxr", end_ele="TD11") # doctest: +SKIP - >>> get_model(["surrogate_cu_inj", "bmad_cu_hxr"], # doctest: +SKIP - ... end_ele="OTR4", n_particles=500) - >>> get_model(["impact_cu_inj", "bmad_cu_hxr"], # doctest: +SKIP - ... handoff_loc="YAG03", end_ele="TD11") + See ``docs/model_registry_usage.md`` for worked examples. """ start_ele, end_ele = _normalize(start_ele), _normalize(end_ele) diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index 20db6d7..09645ed 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -1,7 +1,4 @@ -"""Registry of available virtual-accelerator models. - -Currently LCLS only; FACET-II entries are not registered yet. -""" +"""Registry of available virtual-accelerator models for LCLS and FACET-II.""" from dataclasses import dataclass from typing import Any @@ -9,11 +6,53 @@ @dataclass(frozen=True) class ModelEntry: - """Metadata describing one model and how to configure it. - - ``builder`` is a ``"module:function"`` string rather than a callable so that - importing this module does not import pytao / torch / impact. Discovery must - work with no optional dependencies installed. + """ + Metadata describing one model and how to configure it. + + Parameters + ---------- + name : str + Registry key, also used to qualify per-stage kwargs. + description : str + One-line summary shown by ``models_available``. + facility : str + "lcls" or "facet2". Models of different facilities cannot be staged. + engine : str + "bmad", "impact", "surrogate" or "cheetah". + builder : str + Builder function as a ``"module:function"`` string rather than a callable, + so that importing this module does not import pytao / torch / impact. + Discovery must work with no optional dependencies installed. + extras : tuple[str, ...] + Pip extras the builder needs, e.g. ``("bmad",)``. + params : dict[str, Any] + Configurable parameter name to default value. Also the allow-list against + which kwargs are validated. + handoff_points : tuple[str, ...] + Suggested start, end and handoff elements, in lattice order. A discovery + aid rather than a restriction: any element in the underlying lattice may + be used. Screens are enumerated exhaustively, so a screen-shaped name + absent from this tuple is treated as a typo and rejected; anything else + passes through to the engine. Positions refer to the entrance face of the + element, the only reference plane Bmad and IMPACT express identically. + start_param : str | None, optional + Builder kwarg controlling the start element. Default is None, meaning the + start is fixed and the model cannot be a downstream stage. + end_param : str | None, optional + Builder kwarg controlling the end element. Default is None, meaning the + end is not configurable. + default_start : str | None, optional + Standard start element. Default is None. + default_end : str | None, optional + Standard end element, used to infer the handoff when this model is + upstream in a chain. Default is None. + shared_params : frozenset[str], optional + Params that must hold the same value in every stage of a chain. The beam + flows through the stages, so a particle count differing between them is + physically meaningless. These are broadcast to every stage that declares + them, and the per-stage ``"."`` form is rejected for them -- + letting the values diverge would break the invariant rather than + configure anything. Default is empty. """ name: str @@ -22,44 +61,17 @@ class ModelEntry: engine: str builder: str extras: tuple[str, ...] - params: dict[str, Any] - """Configurable parameter name -> default. Also the allow-list for kwargs.""" - handoff_points: tuple[str, ...] - """Suggested start/end/handoff elements, in lattice order. - - A discovery aid, **not** a restriction: any element name in the underlying - lattice may be used. Screens are enumerated exhaustively, so a screen-shaped - name absent from this tuple is a typo and is rejected; anything else passes - through to the engine. - - Positions refer to the **entrance** face of the element. Bmad's - ``-slice_lattice`` begins at an element's entrance and cannot begin at a - midpoint, and IMPACT's ``impact.ele[name]["s"]`` is also the entrance, so the - entrance is the only reference plane both engines express identically. - """ - start_param: str | None = None - """Builder kwarg controlling the start element, or None if not configurable.""" - end_param: str | None = None - """Builder kwarg controlling the end element, or None if not configurable.""" - default_start: str | None = None default_end: str | None = None - shared_params: frozenset[str] = frozenset() - """Params that must hold the same value in every stage of a chain. - - The beam flows through the stages, so a particle count differing between them - is physically meaningless. These are broadcast to every stage that declares - them, and the per-stage ``"."`` form is rejected for them -- - letting the values diverge would break the invariant, not configure anything. - """ @property def configurable_extent(self) -> bool: + """Whether either end of the model's tracking range can be set.""" return self.start_param is not None or self.end_param is not None From 328118099ea7fd883d48bcf246ea60cd91b11c9d Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Tue, 8 Sep 2026 15:41:13 -0700 Subject: [PATCH 13/15] Fixes --- docs/model_registry_usage.md | 56 ++++++++++++++++-------- virtual_accelerator/impact/factory.py | 27 ++++++++++-- virtual_accelerator/models/cu_hxr.py | 5 ++- virtual_accelerator/models/facet2.py | 27 ++++++++---- virtual_accelerator/registry/__init__.py | 43 +++++++++++++++++- virtual_accelerator/registry/models.py | 12 ++++- 6 files changed, 134 insertions(+), 36 deletions(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index eb21ec8..29537c1 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -205,34 +205,52 @@ LCLS needs two handoff planes because its injector models end at different place neither can move. The NN surrogate predicts `OTRS:IN20:571` (OTR2) at 135 MeV and cannot produce a beam at YAG03, which sits before L0B at 64 MeV. -### Overlapping Variables Are Handled For You -Both stages include the handoff element, so both publish its PVs — an IMPACT model -stopped at `YAG03` keeps the screen (it prunes to `s <= stop`) and so does a Bmad model -sliced from `YAG03`. `StagedModel` would reject the pair as duplicates. +### The Handoff Element Belongs To The Downstream Stage +Tracking stops *at* the handoff plane without carrying on through the element, so the +upstream stage ends just before it and the downstream stage owns it. `get_model()` arranges +this; nothing is required of the caller. -`get_model()` resolves this automatically: the upstream stage owns those PVs, because it -is the stage that tracks the beam to that plane, so they are unregistered from the -downstream stage before the chain is assembled. Nothing is required of the caller. +The two engines express it differently: -The removal is surgical — only genuine collisions go. At YAG03 the IMPACT stage publishes -four PVs; the Bmad stage publishes those four plus `:X` and `:Y` centroid readbacks that -IMPACT does not provide. So the four move to IMPACT and the two Bmad-only ones stay: +| stage | how the exclusion is done | +|---|---| +| Bmad upstream | sliced to Tao's `"-1"`, the element before the handoff | +| IMPACT upstream | `include_end_element=False`, so the element on the stop plane is pruned | + +Neither changes where the beam stops. Every handoff point is zero-length, and IMPACT's stop +plane is already the element's entrance, so ending "before" the element and ending "at" it +are the same z. Only ownership of its PVs changes. + +The result is that the handoff element appears in exactly one stage: ```python ->>> m = get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="YAG03", end_ele="TD11") +>>> m = get_model(["impact_cu_inj", "bmad_cu_hxr"], handoff_loc="YAG03", end_ele="OTR4") >>> imp, bmad = m.lume_model_instances ->>> sorted(v for v in imp.supported_variables if "IN20:351" in v) -['YAGS:IN20:351:Image:ArrayData', 'YAGS:IN20:351:Image:ArraySize0_RBV', - 'YAGS:IN20:351:Image:ArraySize1_RBV', 'YAGS:IN20:351:RESOLUTION'] +>>> "YAG03" in imp.impact_model.simulator.ele +False +>>> len([v for v in imp.supported_variables if "IN20:351" in v]) +0 ->>> sorted(v for v in bmad.supported_variables if "IN20:351" in v) -['YAGS:IN20:351:X', 'YAGS:IN20:351:Y'] +>>> len([v for v in bmad.supported_variables if "IN20:351" in v]) +6 +>>> set(imp.supported_variables) & set(bmad.supported_variables) +set() ``` -A *writable* overlap raises instead of being dropped. That means both stages drive the -same magnet — their extents overlap rather than meeting at a plane — and dropping it -downstream would leave that stage tracking a stale value. +Note this applies only to the handoff. A `start_ele` or `end_ele` you ask for yourself stays +inclusive, so `end_ele="OTR4"` still gives you `OTR4_beam` and the OTR4 image PVs. + +#### If the extents overlap anyway +Because the stages meet at a plane rather than overlapping, there is normally nothing to +deduplicate. As a safeguard, any variables that do turn out to be shared are unregistered +from the downstream stage — `StagedModel` rejects duplicates outright, and this keeps the +failure from surfacing only after a full IMPACT run. + +A *writable* overlap raises instead of being dropped. That means both stages drive the same +magnet, so their extents genuinely overlap rather than meeting at a plane, and dropping it +downstream would leave that stage tracking a stale value. Check the handoff element if you +see it. ### Targeting One Stage With kwargs Parameters fall into two kinds, and which one it is decides how you pass it. diff --git a/virtual_accelerator/impact/factory.py b/virtual_accelerator/impact/factory.py index 3c52ecf..83dddc0 100644 --- a/virtual_accelerator/impact/factory.py +++ b/virtual_accelerator/impact/factory.py @@ -25,6 +25,8 @@ class ImpactModelSpec: impact_yaml_file: str = None numprocs: int = 1 space_charge: bool = False + include_stop_element: bool = True + custom_aliases: dict[str, str] | None = None def get_impact_and_distgen(spec: ImpactModelSpec): @@ -88,7 +90,9 @@ def get_actions_from_groups(impact: Impact, spec: ImpactModelSpec): return actions -def set_stop_location(impact: Impact, stop_location: str | float): +def set_stop_location( + impact: Impact, stop_location: str | float, include_stop_element: bool = True +): """ Set z stop location based on the beginning of the named element or a float value @@ -98,6 +102,12 @@ def set_stop_location(impact: Impact, stop_location: str | float): The impact model object. stop_location : str | float The stop location, either as the name of an element (str) or a float value representing the z position. + include_stop_element : bool, optional + Whether to keep the element sitting exactly on the stop plane. Default is + True. Pass False when handing the beam to a downstream model, so that the + element belongs to that model alone and the two do not both publish its + PVs. Tracking stops at the same z either way, since the stop plane is the + element's entrance. Returns: -------- @@ -118,9 +128,12 @@ def set_stop_location(impact: Impact, stop_location: str | float): impact.stop = stop_location_z # remove elements that are downstream of the stop location - impact.ele = {k: v for k, v in impact.ele.items() if v["s"] <= impact.stop} + def _keep(s: float) -> bool: + return s <= impact.stop if include_stop_element else s < impact.stop + + impact.ele = {k: v for k, v in impact.ele.items() if _keep(v["s"])} impact.input["lattice"] = [ - elem for elem in impact.lattice if elem.get("s", float("inf")) <= impact.stop + elem for elem in impact.lattice if _keep(elem.get("s", float("inf"))) ] return impact @@ -135,7 +148,9 @@ def build_impact_model(spec: ImpactModelSpec): impact.header["Bcurr"] = 1 if spec.space_charge else 0 if spec.stop_location is not None: - impact = set_stop_location(impact, spec.stop_location) + impact = set_stop_location( + impact, spec.stop_location, spec.include_stop_element + ) impact.run() @@ -146,9 +161,13 @@ def build_impact_model(spec: ImpactModelSpec): model = LUMEDistgenImpactModel.from_objects(distgen, impact) # register additional actions to lume model + # The elements CSV is not reliable for PV names, so a model may override the + # ones it cares about. Mirrors custom_aliases on the Bmad side. element_name_to_base_pv_mapping = get_element_name_to_base_pv_mapping( os.environ[spec.lattice_env_var] ) + if spec.custom_aliases: + element_name_to_base_pv_mapping.update(spec.custom_aliases) # get the screen configuration dictionary from the profmon config file config_path = Path(__file__).parent / ".." / "utils" / spec.profmon_config_filename diff --git a/virtual_accelerator/models/cu_hxr.py b/virtual_accelerator/models/cu_hxr.py index 74046d7..b0fbb23 100644 --- a/virtual_accelerator/models/cu_hxr.py +++ b/virtual_accelerator/models/cu_hxr.py @@ -183,7 +183,9 @@ def get_cu_hxr_cheetah_model(n_particles: int = 1000): return model -def get_cu_inj_impact_model(n_particles: int = 100, end_element="OTR2"): +def get_cu_inj_impact_model( + n_particles: int = 100, end_element="OTR2", include_end_element: bool = True +): from virtual_accelerator.impact.factory import ( ImpactModelSpec, build_impact_model, @@ -199,6 +201,7 @@ def get_cu_inj_impact_model(n_particles: int = 100, end_element="OTR2"): numprocs=1, space_charge=False, stop_location=end_element, + include_stop_element=include_end_element, ) model = build_impact_model(spec) diff --git a/virtual_accelerator/models/facet2.py b/virtual_accelerator/models/facet2.py index c8bece3..b675708 100644 --- a/virtual_accelerator/models/facet2.py +++ b/virtual_accelerator/models/facet2.py @@ -6,6 +6,18 @@ logger = logging.getLogger(__name__) +# The elements CSV and the lattice aliases both disagree with the PVs FACET VAs +# actually use, so these overrides are authoritative for both engines. The lattice +# says YAGS:/OTRS:IN10:* for the screens and TCAV:IN10:490 for the TCAV. +FACET_PV_OVERRIDES = { + "PR10241": "PROF:IN10:241", + "PR10465": "PROF:IN10:465", + "PR10471": "PROF:IN10:471", + "PR10571": "PROF:IN10:571", + "PR10711": "PROF:IN10:711", + "TCY10490": "KLYS:LI10:51", +} + IMPACT_GROUP_PV_MAPPING = { "group:L0AF_phase": {"pv": "KLYS:IN10:81:PDES", "element": "L0AF_entrance"}, "group:L0BF_phase": {"pv": "KLYS:IN10:41:PDES", "element": "L0BF_entrance"}, @@ -94,13 +106,6 @@ def get_facet_bmad_model( """ from virtual_accelerator.bmad.factory import BmadModelSpec, build_bmad_model - custom_aliases = { - "PR10241": "PROF:IN10:241", - "PR10571": "PROF:IN10:571", - "PR10711": "PROF:IN10:711", - "TCY10490": "KLYS:LI10:51", - } - spec = BmadModelSpec( feature="FACET-II Bmad model", lattice_env_var="FACET2_LATTICE", @@ -116,7 +121,7 @@ def get_facet_bmad_model( end_element=end_element, track_beam=track_beam, custom_beam_path=custom_beam_path, - custom_aliases=custom_aliases, + custom_aliases=FACET_PV_OVERRIDES, custom_tao_commands=[ "set bmad_com absolute_time_tracking=true", "set bmad_com lr_wakes_on=false", @@ -205,7 +210,9 @@ def get_facet_staged_model(n_particles=10000, surrogate_inputs="machine", **kwar return staged_model -def get_facet_impact_model(n_particles: int = 100, end_element="PR10571"): +def get_facet_impact_model( + n_particles: int = 100, end_element="PR10571", include_end_element: bool = True +): from virtual_accelerator.impact.factory import ( ImpactModelSpec, build_impact_model, @@ -221,6 +228,8 @@ def get_facet_impact_model(n_particles: int = 100, end_element="PR10571"): numprocs=1, space_charge=False, stop_location=end_element, + include_stop_element=include_end_element, + custom_aliases=FACET_PV_OVERRIDES, ) model = build_impact_model(spec) diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index 30091e3..dc41fbc 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -343,6 +343,39 @@ def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) - ) +def _exclusive_end(entry: ModelEntry, handoff: str) -> str: + """ + Get the end element for an upstream stage handing over at ``handoff``. + + Parameters + ---------- + entry : ModelEntry + The upstream stage. + handoff : str + Element where the beam is handed to the next stage. + + Returns + ------- + str + Element to pass as the upstream stage's end. For Bmad this is Tao's + ``"-1"`` offset form; other engines end at ``handoff`` itself. + + Notes + ----- + Tracking stops at the handoff plane without passing through the element, so + the downstream stage owns it and the two stages meet rather than overlap. + + Bmad's ``-slice_lattice`` accepts ``"OTR4-1"`` to mean the element before + OTR4, which also sidesteps naming the predecessor -- several of them are + duplicated in the lattice (both OTR1 and OTR3 follow an element called DE05) + and would otherwise need a ``##N`` index. + + For IMPACT the exclusion happens in ``set_stop_location``, which prunes + elements at or beyond the stop plane, so the name needs no adjustment here. + """ + return f"{handoff}-1" if entry.engine == "bmad" else handoff + + def _strip_overlapping_variables(upstream, downstream, upstream_name, downstream_name): """ Remove variables the downstream stage shares with the upstream stage. @@ -505,7 +538,9 @@ def get_model( stages = [] for i, entry in enumerate(entries): stage_start = start_ele if i == 0 else handoffs[i - 1] - stage_end = end_ele if i == len(entries) - 1 else handoffs[i] + stage_end = ( + end_ele if i == len(entries) - 1 else _exclusive_end(entry, handoffs[i]) + ) stage_kwargs = dict(routed[i]) # Every stage needs tracking on: a non-final stage has to produce @@ -514,6 +549,12 @@ def get_model( if "track_beam" in entry.params: stage_kwargs["track_beam"] = True + # An upstream stage stops at the handoff plane without keeping the element, + # so the downstream stage owns it. Bmad does this via the "-1" offset in + # _exclusive_end; IMPACT needs the flag because its prune is inclusive. + if i < len(entries) - 1 and "include_end_element" in entry.params: + stage_kwargs["include_end_element"] = False + stages.append( _build( entry, diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index 09645ed..bfec3c4 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -99,7 +99,11 @@ def configurable_extent(self) -> bool: engine="impact", builder="virtual_accelerator.models.cu_hxr:get_cu_inj_impact_model", extras=("impact",), - params={"n_particles": 100, "end_element": "YAG03"}, + params={ + "n_particles": 100, + "end_element": "YAG03", + "include_end_element": True, + }, # YAG01 and OTR3 exist in the deck but their lines are commented out; # OTR4 is past stop_1 at z=16.5. handoff_points=("YAG02", "YAG03"), @@ -160,7 +164,11 @@ def configurable_extent(self) -> bool: engine="impact", builder="virtual_accelerator.models.facet2:get_facet_impact_model", extras=("impact",), - params={"n_particles": 100, "end_element": "PR10241"}, + params={ + "n_particles": 100, + "end_element": "PR10241", + "include_end_element": True, + }, handoff_points=("PR10241",), end_param="end_element", default_end="PR10241", From 55d28641a34eae31699e0667428d6faf3b02c07e Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 10 Sep 2026 11:02:21 -0700 Subject: [PATCH 14/15] addressing comments --- docs/model_registry_usage.md | 8 ++++---- virtual_accelerator/registry/__init__.py | 14 +++++++------- virtual_accelerator/registry/models.py | 22 +++++++++++----------- 3 files changed, 22 insertions(+), 22 deletions(-) diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md index 29537c1..34a3183 100644 --- a/docs/model_registry_usage.md +++ b/docs/model_registry_usage.md @@ -36,13 +36,13 @@ surrogate_f2e_inj NN FACET-II injector surrogate, cathode -> PR10241 bmad_f2_elec Bmad FACET-II e- linac, injector handoff -> END ``` -Filter by facility or engine: +Filter by facility or simulator: ```python >>> list_models(facility="facet2") ['impact_f2e_inj', 'surrogate_f2e_inj', 'bmad_f2_elec'] ->>> list_models(engine="bmad") +>>> list_models(simulator="bmad") ['bmad_cu_hxr', 'bmad_f2_elec'] ``` @@ -56,7 +56,7 @@ Each model exposes a set of suggested handoff points — named locations where b impact_cu_inj ('YAG02', 'YAG03') bmad_cu_hxr ('CATHODE', 'YAG02', 'YAG03', 'OTRH1', 'OTRH2', 'OTR1', 'OTR2', 'OTR3', 'OTR4', 'OTR11', 'OTR12', 'OTR21', 'OTRDMP', 'END') surrogate_cu_inj ('OTR2',) -cheetah_cu_hxr () +cheetah_cu_hxr ('CATHODE', 'END') impact_f2e_inj ('PR10241',) surrogate_f2e_inj ('PR10241',) bmad_f2_elec ('CATHODEF', 'PR10241', 'L0AFEND', 'PR10465', 'PR10471', 'PR10571', 'PR10711', 'END') @@ -210,7 +210,7 @@ Tracking stops *at* the handoff plane without carrying on through the element, s upstream stage ends just before it and the downstream stage owns it. `get_model()` arranges this; nothing is required of the caller. -The two engines express it differently: +The two simulators express it differently: | stage | how the exclusion is done | |---|---| diff --git a/virtual_accelerator/registry/__init__.py b/virtual_accelerator/registry/__init__.py index dc41fbc..2a9872f 100644 --- a/virtual_accelerator/registry/__init__.py +++ b/virtual_accelerator/registry/__init__.py @@ -39,7 +39,7 @@ def __repr__(self) -> str: ) -def list_models(facility: str | None = None, engine: str | None = None) -> list[str]: +def list_models(facility: str | None = None, simulator: str | None = None) -> list[str]: """ Get the names of registered models, optionally filtered. @@ -47,8 +47,8 @@ def list_models(facility: str | None = None, engine: str | None = None) -> list[ ---------- facility : str, optional Restrict to one facility, "lcls" or "facet2". Default is None, meaning all. - engine : str, optional - Restrict to one engine, e.g. "bmad". Default is None, meaning all. + simulator : str, optional + Restrict to one simulator, e.g. "bmad". Default is None, meaning all. Returns ------- @@ -59,7 +59,7 @@ def list_models(facility: str | None = None, engine: str | None = None) -> list[ name for name, entry in MODELS.items() if (facility is None or entry.facility == facility) - and (engine is None or entry.engine == engine) + and (simulator is None or entry.simulator == simulator) ] @@ -325,7 +325,7 @@ def _validate_pair(upstream: ModelEntry, downstream: ModelEntry, handoff: str) - if downstream.start_param is None: reason = ( "IMPACT models can only start at the cathode" - if downstream.engine == "impact" + if downstream.simulator == "impact" else f"{downstream.name!r} has a fixed start" ) raise ValueError(f"{downstream.name!r} cannot be a downstream stage: {reason}.") @@ -358,7 +358,7 @@ def _exclusive_end(entry: ModelEntry, handoff: str) -> str: ------- str Element to pass as the upstream stage's end. For Bmad this is Tao's - ``"-1"`` offset form; other engines end at ``handoff`` itself. + ``"-1"`` offset form; other simulators end at ``handoff`` itself. Notes ----- @@ -373,7 +373,7 @@ def _exclusive_end(entry: ModelEntry, handoff: str) -> str: For IMPACT the exclusion happens in ``set_stop_location``, which prunes elements at or beyond the stop plane, so the name needs no adjustment here. """ - return f"{handoff}-1" if entry.engine == "bmad" else handoff + return f"{handoff}-1" if entry.simulator == "bmad" else handoff def _strip_overlapping_variables(upstream, downstream, upstream_name, downstream_name): diff --git a/virtual_accelerator/registry/models.py b/virtual_accelerator/registry/models.py index bfec3c4..1d7a3b9 100644 --- a/virtual_accelerator/registry/models.py +++ b/virtual_accelerator/registry/models.py @@ -17,7 +17,7 @@ class ModelEntry: One-line summary shown by ``models_available``. facility : str "lcls" or "facet2". Models of different facilities cannot be staged. - engine : str + simulator : str "bmad", "impact", "surrogate" or "cheetah". builder : str Builder function as a ``"module:function"`` string rather than a callable, @@ -33,7 +33,7 @@ class ModelEntry: aid rather than a restriction: any element in the underlying lattice may be used. Screens are enumerated exhaustively, so a screen-shaped name absent from this tuple is treated as a typo and rejected; anything else - passes through to the engine. Positions refer to the entrance face of the + passes through to the simulator. Positions refer to the entrance face of the element, the only reference plane Bmad and IMPACT express identically. start_param : str | None, optional Builder kwarg controlling the start element. Default is None, meaning the @@ -58,7 +58,7 @@ class ModelEntry: name: str description: str facility: str - engine: str + simulator: str builder: str extras: tuple[str, ...] params: dict[str, Any] @@ -96,7 +96,7 @@ def configurable_extent(self) -> bool: name="impact_cu_inj", description="IMPACT-T LCLS injector, cathode -> YAG03", facility="lcls", - engine="impact", + simulator="impact", builder="virtual_accelerator.models.cu_hxr:get_cu_inj_impact_model", extras=("impact",), params={ @@ -115,7 +115,7 @@ def configurable_extent(self) -> bool: name="bmad_cu_hxr", description="Bmad CU-HXR linac, injector handoff -> END", facility="lcls", - engine="bmad", + simulator="bmad", builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_bmad_model", extras=("bmad",), params={ @@ -136,7 +136,7 @@ def configurable_extent(self) -> bool: name="surrogate_cu_inj", description="NN LCLS injector surrogate, cathode -> OTR2", facility="lcls", - engine="surrogate", + simulator="surrogate", builder=( "virtual_accelerator.models.cu_hxr:get_cu_hxr_injector_surrogate_model" ), @@ -150,18 +150,18 @@ def configurable_extent(self) -> bool: name="cheetah_cu_hxr", description="Cheetah nc_hxr, cathode -> END", facility="lcls", - engine="cheetah", + simulator="cheetah", builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_cheetah_model", extras=("cheetah",), params={"n_particles": 1000}, - handoff_points=(), + handoff_points=("CATHODE", "END"), shared_params=frozenset({"n_particles"}), ), "impact_f2e_inj": ModelEntry( name="impact_f2e_inj", description="IMPACT-T FACET-II injector, cathode -> PR10241", facility="facet2", - engine="impact", + simulator="impact", builder="virtual_accelerator.models.facet2:get_facet_impact_model", extras=("impact",), params={ @@ -178,7 +178,7 @@ def configurable_extent(self) -> bool: name="surrogate_f2e_inj", description="NN FACET-II injector surrogate, cathode -> PR10241", facility="facet2", - engine="surrogate", + simulator="surrogate", builder="virtual_accelerator.models.facet2:get_facet_injector_surrogate_model", extras=("surrogate",), params={"n_particles": 10000, "surrogate_inputs": "machine"}, @@ -190,7 +190,7 @@ def configurable_extent(self) -> bool: name="bmad_f2_elec", description="Bmad FACET-II e- linac, injector handoff -> END", facility="facet2", - engine="bmad", + simulator="bmad", builder="virtual_accelerator.models.facet2:get_facet_bmad_model", extras=("bmad",), params={ From 0274b16435d04ae1ea45fca567abb1bbed1222f6 Mon Sep 17 00:00:00 2001 From: Gopika Bhardwaj Date: Thu, 10 Sep 2026 14:09:38 -0700 Subject: [PATCH 15/15] fixing tests --- virtual_accelerator/tests/test_registry.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/virtual_accelerator/tests/test_registry.py b/virtual_accelerator/tests/test_registry.py index 9f17f16..c5a11dd 100644 --- a/virtual_accelerator/tests/test_registry.py +++ b/virtual_accelerator/tests/test_registry.py @@ -28,7 +28,7 @@ def test_repr_is_aligned_table(self): assert len(text.splitlines()) == len(MODELS) def test_filter_by_engine_and_facility(self): - assert list_models(engine="bmad") == ["bmad_cu_hxr", "bmad_f2_elec"] + assert list_models(simulator="bmad") == ["bmad_cu_hxr", "bmad_f2_elec"] assert list_models(facility="facet2") == [ "impact_f2e_inj", "surrogate_f2e_inj", @@ -57,7 +57,7 @@ def test_cathode_is_only_listed_where_it_is_usable(self): # FACET's element is CATHODEF, not CATHODE. assert "CATHODE" in list_handoff_points("bmad_cu_hxr") assert "CATHODEF" in list_handoff_points("bmad_f2_elec") - for fixed in ("impact_cu_inj", "surrogate_cu_inj", "cheetah_cu_hxr"): + for fixed in ("impact_cu_inj", "surrogate_cu_inj"): assert not {"CATHODE", "CATHODEF"} & set(list_handoff_points(fixed)) def test_facet_handoff_is_restricted_to_pr10241(self): @@ -219,7 +219,7 @@ def test_ordered_by_lattice_position(self): ) def test_no_shared_points_gives_empty_tuple(self): - assert common_handoff_points("cheetah_cu_hxr", "bmad_cu_hxr") == () + assert common_handoff_points("impact_cu_inj", "surrogate_cu_inj") == () def test_cross_facility_pairs_share_nothing(self): assert common_handoff_points("impact_cu_inj", "bmad_f2_elec") == () @@ -322,11 +322,13 @@ def test_lowercase_bad_screen_is_still_rejected(self): get_model("impact_cu_inj", end_ele="otr99") def test_lowercase_valid_screen_is_accepted(self): - # Reaches the builder (and fails only because the extra is absent here), - # proving validation no longer rejects it. - with pytest.raises((ImportError, ValueError)) as excinfo: + # Validation must not reject a lowercase-but-valid screen. The builder may + # succeed or fail depending on environment (extras, lattice env vars); the + # only thing this test guards against is the "not an available" ValueError. + try: get_model("impact_cu_inj", end_ele="yag03") - assert "not an available" not in str(excinfo.value) + except Exception as exc: + assert "not an available" not in str(exc) def test_lowercase_handoff_normalises_before_resolution(self): entries = [MODELS["impact_cu_inj"], MODELS["bmad_cu_hxr"]]