diff --git a/docs/model_registry_usage.md b/docs/model_registry_usage.md new file mode 100644 index 0000000..34a3183 --- /dev/null +++ b/docs/model_registry_usage.md @@ -0,0 +1,360 @@ +# 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) 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 +pip install git+https://github.com/slaclab/virtual-accelerator.git +``` + +### Quick Start +```python +from virtual_accelerator.registry import ( + get_model, + models_available, + list_models, + list_handoff_points, + common_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 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 simulator: + +```python +>>> list_models(facility="facet2") +['impact_f2e_inj', 'surrogate_f2e_inj', 'bmad_f2_elec'] + +>>> list_models(simulator="bmad") +['bmad_cu_hxr', 'bmad_f2_elec'] +``` + +### Handoff Points +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", "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 ('CATHODE', 'END') +impact_f2e_inj ('PR10241',) +surrogate_f2e_inj ('PR10241',) +bmad_f2_elec ('CATHODEF', 'PR10241', 'L0AFEND', 'PR10465', 'PR10471', 'PR10571', 'PR10711', 'END') +``` + +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 +— 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") +() + +>>> common_handoff_points("impact_f2e_inj", "bmad_f2_elec") +('PR10241',) +``` + +### 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", end_ele="TD11") + + +>>> get_model("impact_cu_inj", end_ele="YAG03") + +``` + +### Error: Unknown Model Name +Model IDs must be exact. Partial names are not supported: + +```python +>>> 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 +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: 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 — no `handoff_loc` needed, it is inferred from the +surrogate's fixed end (OTR2): +```python +>>> m = get_model(["surrogate_cu_inj", "bmad_cu_hxr"], 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 — hand off at YAG03: + +```python +>>> model = get_model( +... ["impact_cu_inj", "bmad_cu_hxr"], +... 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 +``` + +### 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: + +```python +>>> 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. +``` + +Standard chains: + +| Upstream | Downstream | Handoff | +|---|---|---| +| `impact_cu_inj` | `bmad_cu_hxr` | YAG03 | +| `surrogate_cu_inj` | `bmad_cu_hxr` | OTR2 (inferred) | + +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. + +### 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. + +The two simulators express it differently: + +| 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="OTR4") +>>> imp, bmad = m.lume_model_instances + +>>> "YAG03" in imp.impact_model.simulator.ele +False +>>> len([v for v in imp.supported_variables if "IN20:351" in v]) +0 + +>>> len([v for v in bmad.supported_variables if "IN20:351" in v]) +6 +>>> set(imp.supported_variables) & set(bmad.supported_variables) +set() +``` + +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. + +**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") +``` + +Or prefix with the model ID to set one stage: + +```python +>>> m = get_model( +... ["impact_cu_inj", "bmad_cu_hxr"], +... handoff_loc="YAG03", +... **{"bmad_cu_hxr.end_ele": "TD11"}, +... ) +``` + +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" +``` + +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: 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: + +```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"].shared_params +frozenset({'n_particles'}) +``` + +## API Reference +```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. + +```list_handoff_points(model_id: str) -> tuple[str, ...]``` + +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/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 eef6b5b..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", @@ -131,6 +136,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 +190,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") @@ -175,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, @@ -191,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 new file mode 100644 index 0000000..2a9872f --- /dev/null +++ b/virtual_accelerator/registry/__init__.py @@ -0,0 +1,576 @@ +"""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 +import logging +from typing import Any + +from virtual_accelerator.registry.models import MODELS, ModelEntry + +logger = logging.getLogger(__name__) + +__all__ = [ + "get_model", + "models_available", + "list_models", + "list_handoff_points", + "common_handoff_points", +] + + +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, simulator: str | None = None) -> list[str]: + """ + Get the names of registered models, optionally filtered. + + Parameters + ---------- + facility : str, optional + Restrict to one facility, "lcls" or "facet2". Default is None, meaning all. + simulator : str, optional + Restrict to one simulator, e.g. "bmad". Default is None, meaning all. + + Returns + ------- + list[str] + Registry names, in registration order. + """ + return [ + name + for name, entry in MODELS.items() + if (facility is None or entry.facility == facility) + and (simulator is None or entry.simulator == simulator) + ] + + +def list_handoff_points(model_name: str) -> tuple[str, ...]: + """ + Get the suggested start, end and handoff elements for one model. + + 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" + + +def common_handoff_points(*model_names: str) -> tuple[str, ...]: + """ + Get the elements every named model can hand off at, in lattice order. + + Parameters + ---------- + *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. + + 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.") + + 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] + 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 the registry's own ``handoff_points`` lookups would + silently miss. + """ + return name if name is None else name.upper() + + +def _check_element(entry: ModelEntry, name: str, role: str) -> None: + """Validate a start/end/handoff element. + + 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.handoff_points: + return + if name.startswith(("OTR", "YAG", "PR")): + raise ValueError( + f"{name!r} is not an available {role} screen for {entry.name!r}. " + f"Suggested points: {', '.join(entry.handoff_points)}" + ) + + +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)} + + # 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: + 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] + # 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 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}. " + f"Accepted: {', '.join(sorted(entries[index].params))}" + ) + 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}) + raise ValueError( + f"{key!r} is not a parameter of any stage. Accepted: {', '.join(known)}" + ) + + 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}). " + 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] = 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] = 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.simulator == "impact" + else f"{downstream.name!r} has a fixed start" + ) + raise ValueError(f"{downstream.name!r} cannot be a downstream stage: {reason}.") + + if handoff == CATHODE: + raise ValueError( + 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'}" + ) + + +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 simulators 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.simulator == "bmad" else handoff + + +def _strip_overlapping_variables(upstream, downstream, upstream_name, downstream_name): + """ + 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. + + 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 + + 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"{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( + 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 : 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. 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 + include the handoff element, so both publish its PVs -- an IMPACT model stopped + 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``. + + See ``docs/model_registry_usage.md`` for worked examples. + """ + 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 handoffs[i - 1] + 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 + # 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 + + # 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, + stage_kwargs, + stage_start if entry.start_param else None, + stage_end if entry.end_param else None, + ) + ) + + # 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 new file mode 100644 index 0000000..1d7a3b9 --- /dev/null +++ b/virtual_accelerator/registry/models.py @@ -0,0 +1,220 @@ +"""Registry of available virtual-accelerator models for LCLS and FACET-II.""" + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class ModelEntry: + """ + 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. + simulator : 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 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 + 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 + description: str + facility: str + simulator: str + builder: str + extras: tuple[str, ...] + params: dict[str, Any] + handoff_points: tuple[str, ...] + start_param: str | None = None + end_param: str | None = None + default_start: str | None = None + default_end: str | None = None + shared_params: frozenset[str] = frozenset() + + @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 + + +_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 -> YAG03", + facility="lcls", + simulator="impact", + builder="virtual_accelerator.models.cu_hxr:get_cu_inj_impact_model", + extras=("impact",), + 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"), + end_param="end_element", + default_end="YAG03", + shared_params=frozenset({"n_particles"}), + ), + "bmad_cu_hxr": ModelEntry( + name="bmad_cu_hxr", + description="Bmad CU-HXR linac, injector handoff -> END", + facility="lcls", + simulator="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, cathode -> OTR2", + facility="lcls", + simulator="surrogate", + builder=( + "virtual_accelerator.models.cu_hxr:get_cu_hxr_injector_surrogate_model" + ), + extras=("surrogate",), + params={"n_particles": 1000}, + handoff_points=("OTR2",), + default_end="OTR2", + shared_params=frozenset({"n_particles"}), + ), + "cheetah_cu_hxr": ModelEntry( + name="cheetah_cu_hxr", + description="Cheetah nc_hxr, cathode -> END", + facility="lcls", + simulator="cheetah", + builder="virtual_accelerator.models.cu_hxr:get_cu_hxr_cheetah_model", + extras=("cheetah",), + params={"n_particles": 1000}, + 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", + simulator="impact", + builder="virtual_accelerator.models.facet2:get_facet_impact_model", + extras=("impact",), + params={ + "n_particles": 100, + "end_element": "PR10241", + "include_end_element": True, + }, + 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", + simulator="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", + simulator="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 new file mode 100644 index 0000000..c5a11dd --- /dev/null +++ b/virtual_accelerator/tests/test_registry.py @@ -0,0 +1,336 @@ +"""Registry tests that need no engine dependencies or lattice checkouts.""" + +import pytest + +from lume.actions import WritableActionMixin + +from virtual_accelerator.registry import ( + _normalize, + _resolve_handoffs, + _route_kwargs, + _strip_overlapping_variables, + common_handoff_points, + get_model, + list_handoff_points, + 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(simulator="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") + assert diags.index("YAG02") < diags.index("YAG03") < diags.index("OTR2") + + 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 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): + # 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"): + 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: + @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_shared_params_are_declared(self, name): + entry = MODELS[name] + assert entry.shared_params <= set(entry.params) + + @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_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") + + 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"]) + + 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_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"]], {"n_particle": 5}) + + 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}] + + 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["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( + 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}) + + +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_explicit_handoff_is_used_verbatim(self): + 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): + # 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("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. + 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("impact_cu_inj", "surrogate_cu_inj") == () + + 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") + + 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.""" + + 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 handoff_points and + handoff_points 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_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 screen"): + get_model("impact_cu_inj", end_ele="otr99") + + def test_lowercase_valid_screen_is_accepted(self): + # 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") + 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"]] + handoffs = [_normalize(h) for h in _resolve_handoffs(entries, "yag03")] + assert handoffs == ["YAG03"]