diff --git a/.github/workflows/test-cheetah.yml b/.github/workflows/test-cheetah.yml index 304d169..72a8c17 100644 --- a/.github/workflows/test-cheetah.yml +++ b/.github/workflows/test-cheetah.yml @@ -54,4 +54,6 @@ jobs: FACET2_LATTICE: ${{ github.workspace }}/facet2-lattice run: > pytest virtual_accelerator/tests/test_optional_dependencies.py + virtual_accelerator/tests/test_cheetah_actions.py + virtual_accelerator/tests/test_cu_hxr.py -k Cheetah -v --tb=short -ra --strict-markers diff --git a/examples/cheetah_model_example.ipynb b/examples/cheetah_model_example.ipynb index df0b9ce..2e72aef 100644 --- a/examples/cheetah_model_example.ipynb +++ b/examples/cheetah_model_example.ipynb @@ -18,7 +18,7 @@ "metadata": {}, "outputs": [], "source": [ - "model = get_cu_hxr_cheetah_model()" + "model = get_cu_hxr_cheetah_model(n_particles=100000)" ] }, { @@ -60,13 +60,9 @@ "outputs": [], "source": [ "# get beam distribution on OTR screen\n", - "info = model.get(\n", - " [\n", - " \"OTRS:IN20:571:Image:ArrayData\",\n", - " ]\n", - ")\n", + "info = model.get(\"OTRS:IN20:571:Image:ArrayData\")\n", "fig, ax = plt.subplots()\n", - "ax.imshow(info[\"OTRS:IN20:571:Image:ArrayData\"])" + "ax.imshow(info)" ] }, { @@ -80,6 +76,14 @@ "model.reset()\n", "model.get(vars)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "187f625f", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -98,7 +102,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.14.3" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/pyproject.toml b/pyproject.toml index 56f43ef..25dd995 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "numpy", "scipy", "pyyaml", - "lume-base @ git+https://github.com/slaclab/lume-base", + "lume-base", ] [project.urls] @@ -37,12 +37,12 @@ Documentation = "https://github.com/slaclab/virtual-accelerator" [project.optional-dependencies] bmad = [ - "lume-bmad @ git+https://github.com/lume-science/lume-bmad", + "lume-bmad", ] cheetah = [ "torch", - "lume-cheetah @ git+https://github.com/lume-science/lume-cheetah", - "lume-torch @ git+https://github.com/lume-science/lume-torch", + "lume-cheetah", + "lume-torch", ] impact = [ "lume-impact @ git+https://github.com/ChristopherMayes/lume-impact" @@ -52,20 +52,20 @@ pva = [ ] surrogate = [ "torch", - "lume-torch @ git+https://github.com/lume-science/lume-torch", - "lume-cheetah @ git+https://github.com/lume-science/lume-cheetah", - "distgen @ git+https://github.com/ColwynGulliford/distgen", + "lume-torch", + "lume-cheetah", + "distgen", "facet2-inj-ml-model @ git+https://github.com/slaclab/facet2_inj_ml_model", "lcls-cu-inj-model @ git+https://github.com/slaclab/lcls_cu_injector_ml_model.git", ] all = [ "torch", - "lume-bmad @ git+https://github.com/lume-science/lume-bmad", - "lume-cheetah @ git+https://github.com/lume-science/lume-cheetah", + "lume-bmad", + "lume-cheetah", "lume-impact @ git+https://github.com/ChristopherMayes/lume-impact", - "lume-torch @ git+https://github.com/lume-science/lume-torch", + "lume-torch", "lume-pva @ git+https://github.com/lume-science/lume-pva", - "distgen @ git+https://github.com/ColwynGulliford/distgen", + "distgen", "facet2-inj-ml-model @ git+https://github.com/slaclab/facet2_inj_ml_model", "lcls-cu-inj-model @ git+https://github.com/slaclab/lcls_cu_injector_ml_model.git", ] diff --git a/virtual_accelerator/cheetah/actions.py b/virtual_accelerator/cheetah/actions.py new file mode 100644 index 0000000..ac6808f --- /dev/null +++ b/virtual_accelerator/cheetah/actions.py @@ -0,0 +1,316 @@ +"""Action-variable implementations for Cheetah-backed virtual accelerator PVs. + +This module provides the Cheetah action layer used by the virtual accelerator +package. Conversion and composite-device behavior live directly on action +classes. +""" + +from lume_cheetah.actions import ( + CheetahReadOnlyEnumVariable, + CheetahReadOnlyNDVariable, + CheetahReadOnlyScalarVariable, + CheetahWritableScalarVariable, +) +from lume_torch.variables import TorchScalarVariable +from lume.actions import ReadOnlyActionMixin +from lume.variables import EnumVariable + + +BCTRL_LIMIT = 100.0 + + +def get_magnetic_rigidity(energy): + """Calculate magnetic rigidity ($B\rho$) in kG-m for energy in eV.""" + return 33.356 * energy / 1e9 + + +class _ReadOnlyActionMixin(ReadOnlyActionMixin): + read_only: bool = True + + +class _ReadbackFromControlMixin(ReadOnlyActionMixin): + """Read-only mixin for readbacks that reuse control `_get` logic.""" + + read_only: bool = True + + def _get(self, simulator): + # Skip ReadOnlyActionMixin's abstract _get and delegate to the next base. + return super(ReadOnlyActionMixin, self)._get(simulator) + + def _set(self, simulator, value) -> None: + raise RuntimeError(f"{self.name} is read-only") + + +class QuadrupoleBCTRLVariable(CheetahWritableScalarVariable): + """Quadrupole control/desired magnetic strength (BCTRL/BDES) in kG.""" + + unit: str = "kG" + element_attribute: str = "k1" + + def _get(self, simulator): + element, energy = self._resolve_element_and_energy(simulator, self.element_name) + return ( + getattr(element, self.element_attribute) + * element.length + * get_magnetic_rigidity(energy) + ) + + def _set(self, simulator, value): + element, energy = self._resolve_element_and_energy(simulator, self.element_name) + new_k1 = value / get_magnetic_rigidity(energy) / element.length + setattr(element, self.element_attribute, new_k1) + + +class QuadrupoleBACTVariable(_ReadbackFromControlMixin, QuadrupoleBCTRLVariable): + """Quadrupole readback magnetic strength (BACT) in kG.""" + + +class SolenoidBCTRLVariable(CheetahWritableScalarVariable): + """Solenoid control/desired field variable in kG-equivalent units.""" + + unit: str = "kG" + element_attribute: str = "k" + + def _get(self, simulator): + element, energy = self._resolve_element_and_energy(simulator, self.element_name) + return getattr(element, self.element_attribute) * get_magnetic_rigidity(energy) + + def _set(self, simulator, value): + element, energy = self._resolve_element_and_energy(simulator, self.element_name) + new_k = value / get_magnetic_rigidity(energy) + setattr(element, self.element_attribute, new_k) + + +class SolenoidBACTVariable(_ReadbackFromControlMixin, SolenoidBCTRLVariable): + """Solenoid readback field variable in kG-equivalent units.""" + + +class SBendBCTRLVariable(CheetahWritableScalarVariable): + """SBend control/desired momentum-field representation in GeV/c.""" + + unit: str = "GeV/c" + + def _get(self, simulator): + element, _ = self._resolve_element_and_energy(simulator, self.element_name) + sbend = self._primary_element(element) + + if not all(hasattr(sbend, attr) for attr in ("g", "dg", "p0c")): + raise ValueError( + f"Element {self.element_name!r} does not expose sbend field attributes" + ) + + if sbend.g == 0: + return 0 + p = sbend.p0c * (1 + sbend.dg / sbend.g) + return p * 1e-9 + + def _set(self, simulator, value): + element, _ = self._resolve_element_and_energy(simulator, self.element_name) + sbend = self._primary_element(element) + + if not all(hasattr(sbend, attr) for attr in ("g", "p0c")): + raise ValueError( + f"Element {self.element_name!r} does not expose sbend field attributes" + ) + + dp = (value * 1e9 - sbend.p0c) / sbend.p0c + sbend.dg = dp * sbend.g + + +class SBendBACTVariable(_ReadbackFromControlMixin, SBendBCTRLVariable): + """SBend readback momentum-field representation in GeV/c.""" + + +class KickerBCTRLVariable(CheetahWritableScalarVariable): + """Horizontal/vertical corrector control variable (BCTRL/BDES) in kG.""" + + unit: str = "kG" + element_attribute: str = "angle" + + def _get(self, simulator): + element, energy = self._resolve_element_and_energy(simulator, self.element_name) + return getattr(element, self.element_attribute) * get_magnetic_rigidity(energy) + + def _set(self, simulator, value): + element, energy = self._resolve_element_and_energy(simulator, self.element_name) + new_angle = value / get_magnetic_rigidity(energy) + setattr(element, self.element_attribute, new_angle) + + +class KickerBACTVariable(_ReadbackFromControlMixin, KickerBCTRLVariable): + """Horizontal/vertical corrector readback variable (BACT) in kG.""" + + +class StatusVariable(TorchScalarVariable, _ReadOnlyActionMixin): + """Read-only device status scalar value.""" + + element_name: str + + def _get(self, simulator): + return 1.0 + + +class BminVariable(TorchScalarVariable, _ReadOnlyActionMixin): + """Read-only lower operating limit variable for magnet-like devices.""" + + element_name: str + + def _get(self, simulator): + return -BCTRL_LIMIT + + +class BmaxVariable(TorchScalarVariable, _ReadOnlyActionMixin): + """Read-only upper operating limit variable for magnet-like devices.""" + + element_name: str + + def _get(self, simulator): + return BCTRL_LIMIT + + +class ControlStateVariable(EnumVariable, _ReadOnlyActionMixin): + """Read-only high-level control state enum for magnet-like devices.""" + + element_name: str + + options: list[str] = ["Ready", "TRIM", "PERTURB", "BCON_TO_BDES", "BACT_TO_BDES"] + default_value: str = "Ready" + + def _get(self, simulator): + return "Ready" + + +class DummyEnumVariable(EnumVariable, _ReadOnlyActionMixin): + """Read-only placeholder enum used for PVs that exist for interface parity.""" + + element_name: str + + options: list[str] = ["0", "1"] + default_value: str = "0" + + def _get(self, simulator): + return "0" + + +class BPMXVariable(CheetahReadOnlyScalarVariable): + """Read-only BPM X readout in millimeters.""" + + unit: str = "mm" + element_attribute: str = "reading" + + def _get(self, simulator): + return super()._get(simulator)[0] + + +class BPMYVariable(CheetahReadOnlyScalarVariable): + """Read-only BPM Y readout in millimeters.""" + + unit: str = "mm" + element_attribute: str = "reading" + + def _get(self, simulator): + return super()._get(simulator)[1] + + +class BPMTMITDummyVariable(TorchScalarVariable, _ReadOnlyActionMixin): + """Read-only placeholder BPM intensity variable for interface parity.""" + + element_name: str + + unit: str = "arbitrary units" + + def _get(self, simulator): + return 1.0 + + +class CavityAREQVariable(CheetahWritableScalarVariable): + """Writable cavity amplitude request variable in MV.""" + + unit: str = "MV" + + def _get(self, simulator): + return self._get_direct_attribute(simulator, "voltage") + + def _set(self, simulator, value): + self._set_direct_attribute(simulator, "voltage", value) + + +class CavityAREQReadbackVariable(CheetahReadOnlyScalarVariable): + """Read-only cavity amplitude readback variable in MV.""" + + unit: str = "MV" + + def _get(self, simulator): + return self._get_direct_attribute(simulator, "voltage") + + +class CavityPREQVariable(CheetahWritableScalarVariable): + """Writable cavity phase request variable in degrees.""" + + unit: str = "degrees" + + def _get(self, simulator): + return self._get_direct_attribute(simulator, "phase") + + def _set(self, simulator, value): + self._set_direct_attribute(simulator, "phase", value) + + +class CavityPREQReadbackVariable(CheetahReadOnlyScalarVariable): + """Read-only cavity phase readback variable in degrees.""" + + unit: str = "degrees" + + def _get(self, simulator): + return self._get_direct_attribute(simulator, "phase") + + +class CavityMODECFGVariable(CheetahReadOnlyEnumVariable): + """Read-only cavity mode configuration enum.""" + + options: list[str] = ["Disable", "ACCEL", "STDBY", "ACCEL_STDBY"] + default_value: str = "ACCEL_STDBY" + + def _get(self, simulator): + return "ACCEL_STDBY" + + +class ScreenImageVariable(CheetahReadOnlyNDVariable): + """Read-only screen image array variable.""" + + element_attribute: str = "reading" + + def _get(self, simulator): + return super()._get(simulator).T * 65535 + + +class ScreenImageArraySizeVariable(CheetahReadOnlyScalarVariable): + """Read-only scalar for screen image dimension metadata values.""" + + element_attribute: str = "resolution" + index: int + + def _get(self, simulator): + return super()._get(simulator)[self.index] + + +class ScreenResolutionVariable(CheetahReadOnlyScalarVariable): + """Read-only scalar for screen pixel resolution in micrometers.""" + + element_attribute: str = "pixel_size" + unit: str = "um" + + def _get(self, simulator): + return super()._get(simulator)[0] * 1e6 + + +class ScreenPneumaticVariable(CheetahWritableScalarVariable): + """Writable scalar representing screen insertion/activation control.""" + + element_attribute: str = "is_active" + + def _get(self, simulator): + return 1.0 if bool(super()._get(simulator)) else 0.0 + + def _set(self, simulator, value): + super()._set(simulator, 1.0 if bool(value) else 0.0) diff --git a/virtual_accelerator/cheetah/transformer.py b/virtual_accelerator/cheetah/transformer.py deleted file mode 100644 index 5a81064..0000000 --- a/virtual_accelerator/cheetah/transformer.py +++ /dev/null @@ -1,108 +0,0 @@ -import torch - -from virtual_accelerator.cheetah.utils import access_cheetah_attribute -from lume_cheetah.transformer import CheetahTransformer - - -class SLACCheetahTransformer(CheetahTransformer): - """ - CheetahTransformer subclass for SLAC accelerator simulations. - - This class can be extended to include any necessary unit conversions or - special handling for SLAC-specific control variables and their mapping - to cheetah properties. - - Attributes - ---------- - control_name_to_cheetah : dict[str, str] - A dictionary mapping control variable names to cheetah elements - (e.g. {"QUAD:IN20:511 : "qe03"}) - #same something about how bctrl maps to k1, in utils.py - """ - - def __init__(self, control_name_to_cheetah: dict[str, str]): - """ - Initialize the transformer with a mapping from control variable names to - cheetah element names. - - Parameters - ---------- - control_name_to_cheetah : dict[str, str] - A dictionary mapping control variable names to cheetah elements - (e.g. {"QUAD:IN20:511" : "qe03"}) (cheetah names are converted to - lower case internally). - """ - self._control_name_to_cheetah = control_name_to_cheetah - - # convert mapping to lower case for cheetah element names - self._control_name_to_cheetah = { - k: v.lower() for k, v in self._control_name_to_cheetah.items() - } - - @property - def control_name_to_cheetah(self): - return self._control_name_to_cheetah - - def get_cheetah_property(self, simulator, control_variable_name): - """ - Get a property of a Cheetah element based on the control - variable name and return its value in EPICS units. - - Parameters - ---------- - simulator : CheetahSimulator - The simulator instance containing the segment and elements. - control_variable_name : str - The name of the control variable (e.g. "QUAD:IN20:511:BCTRL") - energy : float - The beam energy in eV, used for unit conversions if necessary. - """ - # get the last part after the last colon, which is the attribute name - control_name, attribute = ( - ":".join(control_variable_name.split(":", 3)[:3]), - control_variable_name.split(":", 3)[3], - ) - element_name = self.control_name_to_cheetah.get( - control_name - ) # mapping { "QUAD:IN20:511:BCTRL" : "QE03"} - if element_name is None: - raise ValueError( - f"No mapping found for control variable '{control_variable_name}'" - ) - - element = getattr(simulator.segment, element_name) - beam_energy_at_element = simulator.energies[element_name] - # due to getting beam energy this calc is very slow, maybe some list format should - # be passable for args. - if "STATCTRLSUB.T" in attribute: - return torch.tensor(1.0) - - return access_cheetah_attribute(element, attribute, beam_energy_at_element) - - def set_cheetah_property(self, simulator, control_variable_name, value): - """ - Set a property of a Cheetah element based on the control variable - name and value in EPICS units. - - Parameters - ---------- - simulator : CheetahSimulator - The simulator instance containing the segment and elements. - control_variable_name : str - The name of the control variable (e.g. "QUAD:IN20:511:BCTRL") - value : Any - The value to set for the corresponding cheetah property, in EPICS units. - """ - - control_name, attribute = control_variable_name.rsplit(":", 1) - element_name = self.control_name_to_cheetah.get(control_name) - if element_name is None: - raise ValueError( - f"No mapping found for control variable '{control_variable_name}'" - ) - - element = getattr(simulator.segment, element_name) - beam_energy_at_element = simulator.energies[element_name] - access_cheetah_attribute( - element, attribute, beam_energy_at_element, set_value=value - ) diff --git a/virtual_accelerator/cheetah/utils.py b/virtual_accelerator/cheetah/utils.py index 5f38ad3..edfb9e9 100644 --- a/virtual_accelerator/cheetah/utils.py +++ b/virtual_accelerator/cheetah/utils.py @@ -1,328 +1,27 @@ -import pandas as pd -import torch -import os -from pathlib import Path - - -class NoSetMethodError(Exception): - pass - - -class FieldAccessor: - """ - A class to access and set arbitrary attributes of Cheetah elements when logic - is more complex than simple attribute access (ie. nested attributes). - - This class is used to map process variable (PV) attributes to Cheetah element attributes. - It allows both getting and setting values of the attributes by providing a getter and setter function. - """ - - def __init__(self, getter, setter=None): - self.get = getter - self.set = setter - - def __call__(self, element, energy, value=None): - if value is None: - return self.get(element, energy) - else: - if self.set is None: - raise NoSetMethodError("Cannot set value for this attribute") - self.set(element, energy, value) - - -def get_magnetic_rigidity(energy): - """ - Calculate the magnetic rigidity (Bρ) in kG-m given the beam energy in eV. - """ - return 33.356 * energy / 1e9 - - -# define mappings for different element types - -BCTRL_LIMIT = 100.0 - -# -- include conversions for cheetah attributes to SLAC EPICS attributes -QUADRUPOLE_MAPPING = { - "BCTRL": FieldAccessor( - lambda e, energy: e.k1 * e.length * get_magnetic_rigidity(energy), - lambda e, energy, k1: setattr( - e, "k1", k1 / get_magnetic_rigidity(energy) / e.length - ), - ), - "BACT": FieldAccessor( - lambda e, energy: e.k1 * e.length * get_magnetic_rigidity(energy) - ), - "BMAX": FieldAccessor(lambda e, energy: BCTRL_LIMIT), - "BMIN": FieldAccessor(lambda e, energy: -BCTRL_LIMIT), - "BCTRL.DRVL": FieldAccessor(lambda e, energy: -BCTRL_LIMIT), - "BCTRL.DRVH": FieldAccessor(lambda e, energy: BCTRL_LIMIT), - "CTRL": FieldAccessor(lambda e, energy: "Ready"), - "BCON": FieldAccessor(lambda e, energy: 1.0), - "BDES": FieldAccessor( - lambda e, energy: e.k1 * e.length * get_magnetic_rigidity(energy) - ), -} - -SOLENOID_MAPPING = { - "BCTRL": FieldAccessor( - lambda e, energy: e.k * get_magnetic_rigidity(energy), - lambda e, energy, k: setattr(e, "k", k / (2 * get_magnetic_rigidity(energy))), - ), - "BACT": FieldAccessor(lambda e, energy: e.k * get_magnetic_rigidity(energy)), - "BMAX": FieldAccessor(lambda e, energy: BCTRL_LIMIT), - "BMIN": FieldAccessor(lambda e, energy: -BCTRL_LIMIT), - "BCTRL.DRVL": FieldAccessor(lambda e, energy: -BCTRL_LIMIT), - "BCTRL.DRVH": FieldAccessor(lambda e, energy: BCTRL_LIMIT), - "CTRL": FieldAccessor(lambda e, energy: "Ready"), - "BCON": FieldAccessor(lambda e, energy: 1.0), - "BDES": FieldAccessor(lambda e, energy: e.k * get_magnetic_rigidity(energy)), -} - -CORRECTOR_MAPPING = { - "BCTRL": FieldAccessor( - lambda e, energy: e.angle * get_magnetic_rigidity(energy), - lambda e, energy, a: setattr(e, "angle", a / get_magnetic_rigidity(energy)), - ), - "BACT": FieldAccessor(lambda e, energy: e.angle * get_magnetic_rigidity(energy)), - "BMAX": FieldAccessor(lambda e, energy: BCTRL_LIMIT), - "BMIN": FieldAccessor(lambda e, energy: -BCTRL_LIMIT), - "BCTRL.DRVL": FieldAccessor(lambda e, energy: -BCTRL_LIMIT), - "BCTRL.DRVH": FieldAccessor(lambda e, energy: BCTRL_LIMIT), - "CTRL": FieldAccessor(lambda e, energy: "Ready"), - "BCON": FieldAccessor(lambda e, energy: 1.0), - "BDES": FieldAccessor(lambda e, energy: e.angle * get_magnetic_rigidity(energy)), -} - -TRANSVERSE_DEFLECTING_CAVITY_MAPPING = { - "AREQ": "voltage", - "PREQ": "phase", - "AFBENB": FieldAccessor(lambda e, energy: 0.0), - "AFBST": FieldAccessor(lambda e, energy: 0.0), - "AMPL_W0CH0": FieldAccessor(lambda e, energy: 0.0), - "MODECFG": FieldAccessor(lambda e, energy: 0.0), - "PACT_AVGNT": FieldAccessor(lambda e, energy: 0.0), - "PFBENB": FieldAccessor(lambda e, energy: 0.0), - "PFBST": FieldAccessor(lambda e, energy: 0.0), - "RF_ENABLE": FieldAccessor(lambda e, energy: 1.0), -} - -BPM_MAPPING = { - "X": FieldAccessor(lambda e, energy: e.reading[0]), - "Y": FieldAccessor(lambda e, energy: e.reading[1]), - "XSCDT1H": FieldAccessor(lambda e, energy: e.reading[0]), - "YSCDT1H": FieldAccessor(lambda e, energy: e.reading[1]), - "TMIT": FieldAccessor(lambda e, energy: 1.0), -} - -# multiply image intensity by 16 bit number range (is similar to real machine?) -SCREEN_MAPPING = { - "Image:ArrayData": FieldAccessor(lambda e, energy: e.reading.T * 65535), - "PNEUMATIC": "is_active", - "Image:ArraySize1_RBV": FieldAccessor(lambda e, energy: e.resolution[0]), - "Image:ArraySize0_RBV": FieldAccessor(lambda e, energy: e.resolution[1]), - "RESOLUTION": FieldAccessor(lambda e, energy: e.pixel_size[0] * 1e6), - "IMAGE": FieldAccessor(lambda e, energy: e.reading.T * 65535), - "N_OF_ROW": FieldAccessor(lambda e, energy: e.resolution[0]), - "N_OF_COL": FieldAccessor(lambda e, energy: e.resolution[1]), -} - -# CHEETAH ELEMENT MAPPINGS -MAPPINGS = { - "Quadrupole": QUADRUPOLE_MAPPING, - "Solenoid": SOLENOID_MAPPING, - "HorizontalCorrector": CORRECTOR_MAPPING, - "VerticalCorrector": CORRECTOR_MAPPING, - "BPM": BPM_MAPPING, - "Screen": SCREEN_MAPPING, - "TransverseDeflectingCavity": TRANSVERSE_DEFLECTING_CAVITY_MAPPING, -} - -LCLS_ELEMENTS = os.path.join( - Path(__file__).parent.resolve(), - "lcls_elements.csv", -) - - -def handle_quadrupole_composite(elements, pv_attribute, energy, set_value): - """ - Handle composite quadrupole devices split into multiple subelements. - - Quadrupoles require special handling because their control system - variables (e.g., BCTRL, BDES) represent integrated magnetic strength. - - For a composite quadrupole: - - The total strength depends on the *sum of subelement lengths* - - All subelements share the same normalized strength (k1) - - Behavior: - - GET: - Returns the integrated strength using total length. - - SET: - Computes a single k1 using total length and applies it to all subelements. - - Args: - elements (list[Quadrupole]): - List of quadrupole subelements. - - pv_attribute (str): - Attribute name (e.g., 'BCTRL', 'BDES', 'BACT'). - - energy (torch.Tensor or float): - Beam energy for magnetic rigidity calculation. - - set_value (torch.Tensor or float or None): - Value to set, if performing a set operation. - - Returns: - torch.Tensor or float or None: - Attribute value if getting, otherwise None. - """ - total_length = sum(e.length for e in elements) - - if set_value is not None and pv_attribute in {"BCTRL", "BDES"}: - new_k1 = set_value / get_magnetic_rigidity(energy) / total_length +"""Cheetah utility helpers that are independent from action variable logic. - for e in elements: - e.k1 = new_k1 - return +The action conversion and PV-mapping layer now lives in +``virtual_accelerator.cheetah.actions``. This module keeps only static mapping +helpers used to load MAD/controls naming tables from CSV. +""" - if pv_attribute in {"BCTRL", "BACT", "BDES"}: - return elements[0].k1 * total_length * get_magnetic_rigidity(energy) - - # fallback behavior - return default_composite_handler(elements, pv_attribute, energy, set_value) - - -def default_composite_handler(elements, pv_attribute, energy, set_value): - """ - Default handler for composite devices. - Assumes all subelements share identical attributes. - - Behavior: - - GET: - Returns the attribute from the first subelement. - - SET: - Applies the value to all subelements. - - Args: - elements (list[object]): - List of subelements of the same type. - - pv_attribute (str): - Attribute name. - - energy (torch.Tensor or float): - Beam energy. - - set_value (torch.Tensor or float or None): - Value to set, if performing a set operation. - - Returns: - torch.Tensor or float or None: - Attribute value if getting, otherwise None. - """ - - if set_value is not None: - for e in elements: - access_cheetah_attribute(e, pv_attribute, energy, set_value) - return - - return access_cheetah_attribute(elements[0], pv_attribute, energy) - - -COMPOSITE_HANDLERS = { - "Quadrupole": handle_quadrupole_composite, - "TransverseDeflectingCavity": default_composite_handler, -} - - -def access_cheetah_attribute(element, pv_attribute, energy, set_value=None): - """ - - Return or set a Cheetah element attribute based on the PV attribute. - If `set_value` is provided, it sets the value of the Cheetah attribute. - - - This function supports both single elements and composite devices - represented as lists of subelements (e.g., split quadrupoles or TCAVs). - - For composite devices, behavior is delegated to a device-specific handler. - - Args: - element (element or list[element]): - A single Cheetah element or a list of subelements representing - a single logical device. - - pv_attribute (str): The process variable attribute to map. - - energy (float): The beam energy in eV. - - set_value (optional): If provided, sets the value of the Cheetah attribute. - - Returns: - value: The corresponding Cheetah attribute value if `set_value` is None, otherwise sets the value and returns None. - """ - - # implementing fix for quads, will rethink for tcavs - # simplest case, each subelement has sub_length = length/len(sub_elements) - # handling composite elements - - if isinstance(element, list): - if len(element) == 0: - raise ValueError("Cannot access attribute on empty element list") - element_type = type(element[0]).__name__ - if any(type(sub).__name__ != element_type for sub in element): - raise ValueError("All subelements in element list must have same type") - - handler = COMPOSITE_HANDLERS.get(element_type, default_composite_handler) - return handler(element, pv_attribute, energy, set_value) - - # handling for normal elements - - element_type = type(element).__name__ - if element_type not in MAPPINGS: - raise ValueError(f"Unsupported element type: {element_type}") - - mapping = MAPPINGS[element_type] - if pv_attribute not in mapping: - raise ValueError( - f"Unsupported PV attribute: {pv_attribute} for element type: {element_type}" - ) - - accessor = mapping[pv_attribute] - - # convert to tensor if the value is a float or int - if isinstance(set_value, (float, int)): - set_value = torch.tensor(set_value) +import os +from pathlib import Path - if isinstance(accessor, str): - if set_value is None: - return getattr(element, accessor) - else: - try: - setattr(element, accessor, set_value) - except NoSetMethodError as e: - raise ValueError( - f"Cannot set value for {pv_attribute} of element type {element_type}" - ) from e +import pandas as pd - elif isinstance(accessor, FieldAccessor): - try: - return accessor(element, energy, set_value) - except NoSetMethodError as e: - raise ValueError( - f"Cannot set value for {pv_attribute} of element type {element_type}" - ) from e +LCLS_ELEMENTS = os.path.join(Path(__file__).parent.resolve(), "lcls_elements.csv") def get_mad_control_mapping(fname: str | None = None): """ - Create a mapping from madnames to control names and device types - from a CSV file. - + Create a mapping from MAD element names to control-system names. - Args: - fname (str): Path to the CSV file containing the mapping. + Parameters + ---------- + fname : str | None + Optional path to a CSV file containing ``Element`` and + ``Control System Name`` columns. """ if fname is None: @@ -337,10 +36,13 @@ def get_mad_control_mapping(fname: str | None = None): def get_control_mad_mapping(fname: str | None = None): """ - Create a mapping from control system names to element names from a CSV file. + Create a mapping from control-system names to MAD element names. - Args: - fname (str): Path to the CSV file containing the mapping. + Parameters + ---------- + fname : str | None + Optional path to a CSV file containing ``Control System Name`` and + ``Element`` columns. """ diff --git a/virtual_accelerator/cheetah/variables.py b/virtual_accelerator/cheetah/variables.py index 0201224..24dfb59 100644 --- a/virtual_accelerator/cheetah/variables.py +++ b/virtual_accelerator/cheetah/variables.py @@ -1,11 +1,152 @@ +"""Variable factory utilities for action-based Cheetah model integration. + +This module builds action variables from a Cheetah lattice segment and SLAC PV +configuration mappings by resolving per-element variable classes from +``virtual_accelerator.cheetah.actions``. +""" + from typing import Any import warnings from lume.variables import Variable from virtual_accelerator.utils.variables import ( - get_variables_from_element_name, get_element_attr_mapping, - convert_to_torch_variables, ) +from virtual_accelerator.cheetah import actions as cheetah_actions + + +SKIPPED_ELEMENT_TYPES = { + "Drift", + "Marker", + "Cavity", + "Undulator", + "Dipole", + "Aperture", +} + +# Element-type aliases bridge Cheetah runtime names to SLAC config keys. +ELEMENT_TYPE_ALIASES = { + "TransverseDeflectingCavity": "Crab_Cavity", +} + +# Screen variables require an explicit mapping because the YAML configuration +# currently uses BMAD screen class names that do not apply to this action layer. +SCREEN_VARIABLE_CLASS_MAPPING = { + "Image:ArrayData": "ScreenImageVariable", + "PNEUMATIC": "ScreenPneumaticVariable", + "Image:ArraySize1_RBV": "ScreenImageArraySizeVariable", + "Image:ArraySize0_RBV": "ScreenImageArraySizeVariable", + "RESOLUTION": "ScreenResolutionVariable", + "IMAGE": "ScreenImageVariable", + "N_OF_ROW": "ScreenImageArraySizeVariable", + "N_OF_COL": "ScreenImageArraySizeVariable", +} + +SCREEN_ARRAY_SIZE_INDEX_BY_SUFFIX = { + "Image:ArraySize1_RBV": 0, + "N_OF_ROW": 0, + "Image:ArraySize0_RBV": 1, + "N_OF_COL": 1, +} + + +def _resolve_variable_class_name(var_spec: Any) -> str: + """Extract the variable class name from config entries. + + The SLAC mapping can provide either a string class name or a dictionary + containing ``variable_class`` metadata. + """ + if isinstance(var_spec, dict): + return var_spec["variable_class"] + return var_spec + + +def _resolve_element_variable_mapping( + element_type: str, + element_attr_mapping: dict[str, dict[str, Any]], +) -> dict[str, Any] | None: + """Resolve PV attribute-to-class mapping for a Cheetah element type. + + Screen mappings are handled separately to force action-layer classes. + All other element types are resolved through alias normalization and YAML + lookup. + """ + if element_type == "Screen": + return SCREEN_VARIABLE_CLASS_MAPPING + + mapping_key = ELEMENT_TYPE_ALIASES.get(element_type, element_type) + return element_attr_mapping.get(mapping_key) + + +def _resolve_control_name( + element_name: str, device_mapping: dict[str, str] +) -> str | None: + """Resolve control-system base PV from an element name. + + Supports direct element-name matches and split-element fallback by removing + ``#`` suffixes. + """ + normalized_name = element_name.upper() + if normalized_name in device_mapping: + return device_mapping[normalized_name] + + split_name = normalized_name.split("#", 1)[0] + return device_mapping.get(split_name) + + +def _instantiate_element_variables( + element_name: str, + control_name: str, + class_mapping: dict[str, Any], + image_shape: tuple[int, ...] | None = None, +) -> dict[str, Variable]: + """Instantiate mapped action variables for a single logical element. + + Parameters + ---------- + element_name : str + Logical element identifier used by action variables. + control_name : str + Base PV prefix used to build full variable names. + class_mapping : dict[str, Any] + Mapping of PV suffix to variable class descriptor. + image_shape : tuple[int, ...] | None, optional + Explicit NDVariable shape for screen image variables. + + Returns + ------- + dict[str, Variable] + Mapping of full PV names to instantiated action variables. + """ + element_variables: dict[str, Variable] = {} + + for attr, var_spec in class_mapping.items(): + var_class_name = _resolve_variable_class_name(var_spec) + var_class = getattr(cheetah_actions, var_class_name, None) + if var_class is None: + raise ValueError( + f"Unknown Cheetah variable class {var_class_name!r} for {element_name}.{attr}" + ) + + variable_name = f"{control_name}:{attr}" + init_kwargs = { + "name": variable_name, + "element_name": element_name, + } + + if issubclass(var_class, cheetah_actions.CheetahReadOnlyNDVariable): + init_kwargs["shape"] = image_shape or (1,) + + if var_class_name == "ScreenImageArraySizeVariable": + size_index = SCREEN_ARRAY_SIZE_INDEX_BY_SUFFIX.get(attr) + if size_index is None: + raise ValueError( + f"No screen array-size index configured for suffix {attr!r}" + ) + init_kwargs["index"] = size_index + + element_variables[variable_name] = var_class(**init_kwargs) + + return element_variables def get_variables_from_segment( @@ -43,33 +184,50 @@ def get_variables_from_segment( An example `device_mapping` might look like: device_mapping = {"QE03": "QUAD:IN20:511", "KLYS01": "KLYS:IN20:1001"} - See `get_variables_from_element_name` for details on the specification of `element_attr_mapping`. + Variable classes are resolved from `virtual_accelerator.cheetah.actions` + using the configured mapping for each element type. """ - from cheetah.accelerator import Screen - all_variables = {} + processed_control_names: set[str] = set() element_attr_mapping = element_attr_mapping or get_element_attr_mapping() for element in segment.elements: - if type(element).__name__ in ["Drift", "Marker", "Cavity"]: + element_type = type(element).__name__ + if element_type in SKIPPED_ELEMENT_TYPES: continue - elif element.name.upper() in device_mapping: - control_name = device_mapping[element.name.upper()] - else: + + control_name = _resolve_control_name(element.name, device_mapping) + if control_name is None: warnings.warn(f"Element {element.name} not found in device mapping") continue - element_variables = get_variables_from_element_name( - type(element).__name__, control_name, element_attr_mapping + if control_name in processed_control_names: + # Skip duplicate split-elements that map to one control device. + continue + processed_control_names.add(control_name) + + class_mapping = _resolve_element_variable_mapping( + element_type, element_attr_mapping ) + if class_mapping is None: + warnings.warn( + f"Element type {element_type} for {element.name} not found in variable mapping" + ) + continue - # if element type is a screen then modify the output variable - if isinstance(element, Screen): - element_variables[ - f"{control_name}:Image:ArrayData" - ].shape = element.resolution + base_element_name = element.name.split("#", 1)[0] + image_shape = None + if element_type == "Screen": + image_shape = tuple(element.resolution) + + element_variables = _instantiate_element_variables( + element_name=base_element_name, + control_name=control_name, + class_mapping=class_mapping, + image_shape=image_shape, + ) all_variables.update(element_variables) - return convert_to_torch_variables(all_variables) + return all_variables diff --git a/virtual_accelerator/models/cu_hxr.py b/virtual_accelerator/models/cu_hxr.py index b4063be..cff8778 100644 --- a/virtual_accelerator/models/cu_hxr.py +++ b/virtual_accelerator/models/cu_hxr.py @@ -1,4 +1,5 @@ -from virtual_accelerator.bmad.factory import BmadModelSpec, build_bmad_model +import os + from lume.staged_model import StagedModel @@ -26,6 +27,8 @@ def get_cu_hxr_bmad_model( Instance of the LUMEBmadModel for the CU_HXR lattice. """ + from virtual_accelerator.bmad.factory import BmadModelSpec, build_bmad_model + spec = BmadModelSpec( feature="CU HXR Bmad model", lattice_env_var="LCLS_LATTICE", @@ -85,3 +88,74 @@ def get_cu_hxr_staged_model(n_particles: int = 1000, **kwargs) -> StagedModel: staged_model = StagedModel([injector_surrogate, cu_hxr_bmad_model]) return staged_model + + +def get_cu_hxr_cheetah_model(n_particles: int = 1000): + """ + Get the LUMECheetahModel for the CU_HXR lattice. + + Returns + ------- + LUMECheetahModel + Instance of the LUMECheetahModel for the CU_HXR lattice. + """ + import torch + from cheetah.accelerator import Segment + from cheetah.particles import ParticleBeam + from lume_cheetah import LUMECheetahModel, CheetahSimulator + from virtual_accelerator.cheetah.utils import get_mad_control_mapping + from virtual_accelerator.cheetah.variables import get_variables_from_segment + + # Get path to beam distributions + # beam_dist = os.environ.get( + # 'BEAM_DISTRIBUTION', + # '/sdf/group/ad/sw/machine-learning/ + # Linac-Simulation-Server/simulation_server/beams' + # ) + # Create Cheetah particle Beam from file + + incoming_beam = ParticleBeam.from_twiss( + beta_x=torch.tensor(9.34), + alpha_x=torch.tensor(-1.6946), + emittance_x=torch.tensor(1e-7), + beta_y=torch.tensor(9.34), + alpha_y=torch.tensor(-1.6946), + emittance_y=torch.tensor(1e-7), + num_particles=n_particles, + energy=torch.tensor(90e6), + ) + incoming_beam.particle_charges = torch.tensor(1.0) + + # Get path to lattice files + lcls_lattice = os.environ.get("LCLS_LATTICE") + if lcls_lattice is None: + raise ValueError("LCLS_LATTICE environment variable must be set") + + # Create lattice from file + segment = Segment.from_lattice_json( + os.path.join(lcls_lattice, "cheetah/nc_hxr.json") + ) + + # Define the simulator using lattice and particle beam + simulator = CheetahSimulator( + segment=segment, + initial_beam_distribution=incoming_beam, + ) + + # get control system device to cheetah mapping + database_path = os.path.join( + lcls_lattice, "bmad/conversion/from_oracle/lcls_elements.csv" + ) + element_name_to_control_name = get_mad_control_mapping(database_path) + + # Get supported control system variables + # for the model + variables = get_variables_from_segment(segment, element_name_to_control_name) + + # Create model using action-based variable integration. + model = LUMECheetahModel( + simulator=simulator, + action_variables=list(variables.values()), + ) + + return model diff --git a/virtual_accelerator/tests/_bmad_model_test_utils.py b/virtual_accelerator/tests/_bmad_model_test_utils.py index 94bc1be..feff8b3 100644 --- a/virtual_accelerator/tests/_bmad_model_test_utils.py +++ b/virtual_accelerator/tests/_bmad_model_test_utils.py @@ -1,5 +1,6 @@ import math import os +from collections.abc import Iterable, Sequence from numbers import Real from pathlib import Path @@ -8,10 +9,200 @@ TEST_BEAM_PATH = os.path.join(Path(__file__).parent, "../bmad", "test_beam") +def _normalize_element_name(element_name: str) -> str: + """Return an element name without any split-index suffix. + + Parameters + ---------- + element_name : str + Raw element name from lattice metadata (for example ``Q1#2``). + + Returns + ------- + str + Base element name with any ``#...`` suffix removed. + """ + return element_name.split("#", 1)[0] + + +def _collect_elements_by_type( + element_names: Sequence[str], + element_types: Sequence[str], + requested_type: str, +) -> set[str]: + """Collect normalized lattice element names for a specific element type. + + Parameters + ---------- + element_names : Sequence[str] + Element names from a lattice/segment description. + element_types : Sequence[str] + Element type identifiers aligned positionally with ``element_names``. + requested_type : str + Element type to select. + + Returns + ------- + set[str] + Unique normalized element names for entries whose type matches + ``requested_type``. + + Raises + ------ + AssertionError + If ``element_names`` and ``element_types`` are not the same length. + """ + assert len(element_names) == len(element_types), ( + "Element-name and element-type sequences must have the same length." + ) + + return { + _normalize_element_name(element_name) + for element_name, element_type in zip(element_names, element_types) + if element_name not in ("BEGINNING", "END") and element_type == requested_type + } + + +def _assert_elements_have_pv_mapping_and_attrs( + model, + element_names: Iterable[str], + required_attrs: tuple[str, ...], + element_group_label: str, +) -> None: + """Assert PV mapping completeness and required suffixes for elements. + + Parameters + ---------- + model : LUMEModel + Model exposing ``supported_variables``. + element_names : Iterable[str] + Element names to validate. Split-element names are normalized. + required_attrs : tuple[str, ...] + PV suffixes that must exist for each element. + element_group_label : str + Label used in assertion messages to identify the validated element group. + + Raises + ------ + AssertionError + If any element is missing from element-name mapping metadata or if one + or more required PV suffixes are missing for any element. + """ + normalized_elements = sorted( + { + _normalize_element_name(element_name) + for element_name in element_names + if element_name not in ("BEGINNING", "END") + } + ) + + if not normalized_elements: + return + + pvs_by_element = get_pvs_by_element_name(model) + + missing_mapping = sorted( + element_name + for element_name in normalized_elements + if element_name not in pvs_by_element + ) + assert not missing_mapping, ( + f"{element_group_label} elements missing variables with element_name mapping: " + + ", ".join(missing_mapping) + ) + + missing_pvs = {} + for element_name in normalized_elements: + element_pvs = pvs_by_element[element_name] + absent_attrs = sorted( + attr + for attr in required_attrs + if not any(pv_name.endswith(f":{attr}") for pv_name in element_pvs) + ) + if absent_attrs: + missing_pvs[element_name] = absent_attrs + + assert not missing_pvs, ( + f"{element_group_label} PV attrs missing from model.supported_variables: " + + "; ".join( + f"{element}: {', '.join(attrs)}" for element, attrs in missing_pvs.items() + ) + ) + + +def _get_tao_lattice_element_metadata(model) -> tuple[list[str], list[str]]: + """Get element names and keys from a Tao lattice. + + Parameters + ---------- + model : LUMEBmadModel + Model exposing a Tao interface at ``model.tao``. + + Returns + ------- + tuple[list[str], list[str]] + Pair of lists ``(element_names, element_keys)`` from ``tao.lat_list``. + """ + return model.tao.lat_list("*", "ele.name"), model.tao.lat_list("*", "ele.key") + + +def _get_cheetah_segment_element_metadata(model) -> tuple[list[str], list[str]]: + """Get element names and runtime class names from a Cheetah segment. + + Parameters + ---------- + model : LUMECheetahModel + Model exposing ``model.simulator.segment.elements``. + + Returns + ------- + tuple[list[str], list[str]] + Pair ``(element_names, element_types)`` where ``element_types`` are + class names such as ``Quadrupole`` or ``Screen``. + + Raises + ------ + AssertionError + If required simulator/segment metadata is not available or an element + has an invalid ``name`` attribute. + """ + + elements = model.simulator.segment.elements + + element_names: list[str] = [] + element_types: list[str] = [] + + for element in elements: + element_name = getattr(element, "name", None) + assert isinstance(element_name, str) and element_name, ( + f"Cheetah segment element {element!r} does not expose a valid name." + ) + element_names.append(element_name) + element_types.append(type(element).__name__) + + return element_names, element_types + + def assert_bmad_model_initialization( get_model, required_control_variable: str | None = None, ) -> None: + """Assert that a BMAD model initializes and exposes writable controls. + + Parameters + ---------- + get_model : Callable[..., LUMEBmadModel] + Factory used to build the model under test. The helper passes + ``custom_beam_path=TEST_BEAM_PATH``. + required_control_variable : str | None, optional + If provided, this PV must be writable in ``model.supported_variables``. + + Raises + ------ + AssertionError + If no writable controls are found or if ``required_control_variable`` is + not writable. + """ model = get_model(custom_beam_path=TEST_BEAM_PATH) writable_control_variables = { @@ -26,6 +217,19 @@ def assert_bmad_model_initialization( def assert_bmad_model_twiss_outputs(get_model) -> None: + """Assert basic Twiss/readback output integrity for a BMAD model. + + Parameters + ---------- + get_model : Callable[..., LUMEBmadModel] + Factory used to build the model under test. + + Raises + ------ + AssertionError + If Twiss output lengths do not match lattice length or if expected + endpoint names are not present. + """ model = get_model(custom_beam_path=TEST_BEAM_PATH) outputs = model.get(["a.beta", "b.beta", "name"]) @@ -36,6 +240,18 @@ def assert_bmad_model_twiss_outputs(get_model) -> None: def assert_bmad_model_track_beam_custom_path(get_model) -> None: + """Assert that a BMAD model can initialize with ``track_beam=True``. + + Parameters + ---------- + get_model : Callable[..., LUMEBmadModel] + Factory used to build the model under test. + + Raises + ------ + AssertionError + If the returned model instance is ``None``. + """ # This test ensures shared track_beam setup works when custom_beam_path is given. model = get_model(track_beam=True, custom_beam_path=TEST_BEAM_PATH) assert model is not None @@ -44,6 +260,45 @@ def assert_bmad_model_track_beam_custom_path(get_model) -> None: def assert_magnet_pvs_match_tao_lattice( model, element_key: str, + excluded_elements: Iterable[str] = (), + element_attrs: tuple[str, ...] = ( + "BCTRL", + "BACT", + "BDES", + "BMIN", + "BMAX", + "STATCTRLSUB.T", + "CTRL", + ), +) -> None: + """Assert magnet PV coverage for a Tao-backed model. + + Parameters + ---------- + model : LUMEBmadModel + BMAD model exposing Tao lattice metadata. + element_key : str + Tao element key to validate (for example ``Quadrupole``). + excluded_elements : Iterable[str], optional + Element names to skip from validation. + element_attrs : tuple[str, ...], optional + PV suffixes that must exist for each validated element. + """ + element_names, element_keys = _get_tao_lattice_element_metadata(model) + assert_magnet_pvs_match_lattice_elements( + model=model, + element_key=element_key, + element_names=element_names, + element_keys=element_keys, + excluded_elements=excluded_elements, + element_attrs=element_attrs, + ) + + +def assert_magnet_pvs_match_cheetah_segment( + model, + element_key: str, + excluded_elements: Iterable[str] = (), element_attrs: tuple[str, ...] = ( "BCTRL", "BACT", @@ -54,72 +309,88 @@ def assert_magnet_pvs_match_tao_lattice( "CTRL", ), ) -> None: + """Assert magnet PV coverage for a Cheetah-backed model. + + Parameters + ---------- + model : LUMECheetahModel + Cheetah model exposing segment metadata. + element_key : str + Cheetah element class name to validate (for example ``Quadrupole``). + excluded_elements : Iterable[str], optional + Element names to skip from validation. + element_attrs : tuple[str, ...], optional + PV suffixes that must exist for each validated element. """ - Verify that all lattice elements of a given type have corresponding PV mappings. + element_names, element_keys = _get_cheetah_segment_element_metadata(model) + assert_magnet_pvs_match_lattice_elements( + model=model, + element_key=element_key, + element_names=element_names, + element_keys=element_keys, + excluded_elements=excluded_elements, + element_attrs=element_attrs, + ) - This utility checks: - 1. All elements of the specified type have PV prefix mappings. - 2. All expected PV attributes for those elements exist in model.supported_variables. + +def assert_magnet_pvs_match_lattice_elements( + model, + element_key: str, + element_names: Sequence[str], + element_keys: Sequence[str], + excluded_elements: Iterable[str] = (), + element_attrs: tuple[str, ...] = ( + "BCTRL", + "BACT", + "BDES", + "BMIN", + "BMAX", + "STATCTRLSUB.T", + "CTRL", + ), +) -> None: + """Assert magnet PV coverage from explicit lattice metadata sequences. Parameters ---------- model : LUMEModel - The LUME BMAD model instance to check. + Model exposing ``supported_variables``. element_key : str - The element type key to match (e.g., "Quadrupole", "HKicker", "VKicker"). - element_attrs : tuple[str, ...] - The PV attributes to check for each element (e.g., "BCTRL", "BACT", etc.). + Element type/class identifier to match. + element_names : Sequence[str] + Element names aligned positionally with ``element_keys``. + element_keys : Sequence[str] + Element type identifiers aligned positionally with ``element_names``. + excluded_elements : Iterable[str], optional + Element names to skip from validation after normalization. + element_attrs : tuple[str, ...], optional + PV suffixes required for each validated element. Raises ------ AssertionError - If any elements lack mappings or expected PVs are missing. + If metadata lengths do not match, if required element mappings are + missing, or if required PV suffixes are missing. """ - element_names = model.tao.lat_list("*", "ele.name") - element_keys = model.tao.lat_list("*", "ele.key") - - # Collect all unique element names of the specified type - matching_elements = { - element_name.split("#")[0] - for element_name, key in zip(element_names, element_keys) - if element_name not in ("BEGINNING", "END") and key == element_key + matching_elements = _collect_elements_by_type( + element_names=element_names, + element_types=element_keys, + requested_type=element_key, + ) + excluded_elements_normalized = { + _normalize_element_name(element_name) for element_name in excluded_elements } - - if not matching_elements: - # No elements of this type, test passes vacuously - return - - pvs_by_element = get_pvs_by_element_name(model) - - # Check that all elements have PV mappings - missing_mapping = sorted( + matching_elements = { element_name for element_name in matching_elements - if element_name not in pvs_by_element - ) - assert not missing_mapping, ( - f"{element_key} elements missing variables with element_name mapping: " - + ", ".join(missing_mapping) - ) - - # Check that all expected PV attributes exist in supported_variables - missing_pvs = {} - - for element_name in sorted(matching_elements): - element_pvs = pvs_by_element[element_name] - absent_pvs = sorted( - attr - for attr in element_attrs - if not any(pv_name.endswith(f":{attr}") for pv_name in element_pvs) - ) - if absent_pvs: - missing_pvs[element_name] = absent_pvs + if element_name not in excluded_elements_normalized + } - assert not missing_pvs, ( - f"{element_key} PV attrs missing from model.supported_variables: " - + "; ".join( - f"{element}: {', '.join(pvs)}" for element, pvs in missing_pvs.items() - ) + _assert_elements_have_pv_mapping_and_attrs( + model=model, + element_names=matching_elements, + required_attrs=element_attrs, + element_group_label=element_key, ) @@ -139,7 +410,7 @@ def assert_screen_image_pvs_in_supported_variables( Parameters ---------- model : LUMEModel - The LUME BMAD model instance to check. + Model exposing ``supported_variables``. screen_elements : tuple[str, ...] | list[str] | None Optional explicit list of lattice screen element names. If omitted, uses ``model.dump_locations``. @@ -155,38 +426,67 @@ def assert_screen_image_pvs_in_supported_variables( if screen_elements is None: screen_elements = tuple(getattr(model, "dump_locations", ()) or ()) - if not screen_elements: - return + _assert_elements_have_pv_mapping_and_attrs( + model=model, + element_names=screen_elements, + required_attrs=screen_attrs, + element_group_label="Screen", + ) - pvs_by_element = get_pvs_by_element_name(model) - missing_mapping = sorted( - element_name - for element_name in screen_elements - if element_name not in pvs_by_element - ) - assert not missing_mapping, ( - "Screen elements missing variables with element_name mapping: " - + ", ".join(missing_mapping) +def assert_screen_image_pvs_match_tao_dump_locations( + model, + screen_attrs: tuple[str, ...] = ( + "Image:ArrayData", + "Image:ArraySize1_RBV", + "Image:ArraySize0_RBV", + "RESOLUTION", + ), +) -> None: + """Assert screen image PV coverage using Tao ``dump_locations``. + + Parameters + ---------- + model : LUMEBmadModel + BMAD model that provides ``dump_locations``. + screen_attrs : tuple[str, ...], optional + PV suffixes required for each screen element. + """ + assert_screen_image_pvs_in_supported_variables( + model=model, + screen_elements=tuple(getattr(model, "dump_locations", ()) or ()), + screen_attrs=screen_attrs, ) - missing_image_pvs = {} - for element_name in screen_elements: - element_pvs = pvs_by_element[element_name] - absent_attrs = sorted( - attr - for attr in screen_attrs - if not any(pv_name.endswith(f":{attr}") for pv_name in element_pvs) - ) - if absent_attrs: - missing_image_pvs[element_name] = absent_attrs - assert not missing_image_pvs, ( - "Screen image PV attrs missing from model.supported_variables: " - + "; ".join( - f"{element}: {', '.join(attrs)}" - for element, attrs in missing_image_pvs.items() - ) +def assert_screen_image_pvs_match_cheetah_segment( + model, + screen_attrs: tuple[str, ...] = ( + "Image:ArrayData", + "Image:ArraySize1_RBV", + "Image:ArraySize0_RBV", + "RESOLUTION", + ), +) -> None: + """Assert screen image PV coverage for screen elements in a Cheetah segment. + + Parameters + ---------- + model : LUMECheetahModel + Cheetah model exposing segment elements. + screen_attrs : tuple[str, ...], optional + PV suffixes required for each screen element. + """ + element_names, element_types = _get_cheetah_segment_element_metadata(model) + screen_elements = _collect_elements_by_type( + element_names=element_names, + element_types=element_types, + requested_type="Screen", + ) + assert_screen_image_pvs_in_supported_variables( + model=model, + screen_elements=tuple(sorted(screen_elements)), + screen_attrs=screen_attrs, ) @@ -210,46 +510,70 @@ def assert_bpm_pvs_match_tao_lattice( If BPM lattice elements are missing PV prefix mappings, or if expected BPM PVs are missing from ``model.supported_variables``. """ - element_names = model.tao.lat_list("*", "ele.name") + element_names, _ = _get_tao_lattice_element_metadata(model) bpm_elements = { - element_name.split("#")[0] + _normalize_element_name(element_name) for element_name in element_names if element_name not in ("BEGINNING", "END") - and "BPM" in element_name.split("#")[0] + and "BPM" in _normalize_element_name(element_name) } - if not bpm_elements: - return + assert_bpm_pvs_match_elements( + model=model, + bpm_elements=bpm_elements, + bpm_attrs=bpm_attrs, + ) - pvs_by_element = get_pvs_by_element_name(model) - missing_mapping = sorted( - element_name - for element_name in bpm_elements - if element_name not in pvs_by_element +def assert_bpm_pvs_match_cheetah_segment( + model, + bpm_attrs: tuple[str, ...] = ("X", "Y", "TMIT"), +) -> None: + """Assert BPM PV coverage for BPM elements in a Cheetah segment. + + Parameters + ---------- + model : LUMECheetahModel + Cheetah model exposing segment elements. + bpm_attrs : tuple[str, ...], optional + PV suffixes required for each BPM element. + """ + element_names, element_types = _get_cheetah_segment_element_metadata(model) + bpm_elements = _collect_elements_by_type( + element_names=element_names, + element_types=element_types, + requested_type="BPM", ) - assert not missing_mapping, ( - "BPM elements missing variables with element_name mapping: " - + ", ".join(missing_mapping) + + assert_bpm_pvs_match_elements( + model=model, + bpm_elements=bpm_elements, + bpm_attrs=bpm_attrs, ) - missing_bpm_pvs = {} - for element_name in sorted(bpm_elements): - element_pvs = pvs_by_element[element_name] - absent_pvs = sorted( - attr - for attr in bpm_attrs - if not any(pv_name.endswith(f":{attr}") for pv_name in element_pvs) - ) - if absent_pvs: - missing_bpm_pvs[element_name] = absent_pvs +def assert_bpm_pvs_match_elements( + model, + bpm_elements: Iterable[str], + bpm_attrs: tuple[str, ...] = ("X", "Y", "TMIT"), +) -> None: + """Assert BPM PV coverage for an explicit BPM element name collection. - assert not missing_bpm_pvs, ( - "BPM PV attrs missing from model.supported_variables: " - + "; ".join( - f"{element}: {', '.join(pvs)}" for element, pvs in missing_bpm_pvs.items() - ) + Parameters + ---------- + model : LUMEModel + Model exposing ``supported_variables``. + bpm_elements : Iterable[str] + BPM element names to validate. + bpm_attrs : tuple[str, ...], optional + PV suffixes required for each BPM element. + """ + + _assert_elements_have_pv_mapping_and_attrs( + model=model, + element_names=bpm_elements, + required_attrs=bpm_attrs, + element_group_label="BPM", ) @@ -263,9 +587,16 @@ def assert_roundtrip_pv_get_set( ---------- model : LUMEModel The model instance under test. + + Notes + ----- + - Only writable scalar-like controls are exercised. + - ``track_type`` is excluded because it has model-level semantics that do + not fit this generic scalar roundtrip check. """ def assert_value_equal(pv_name, expected_value, actual_value) -> None: + """Assert equality with support for numeric tolerance and array values.""" if ( isinstance(expected_value, Real) and isinstance(actual_value, Real) @@ -313,23 +644,23 @@ def assert_value_equal(pv_name, expected_value, actual_value) -> None: "No writable variables found in model.supported_variables" ) - scalar_variable_cls = None - try: - from lume.variables import ScalarVariable as scalar_variable_cls - except ImportError: - pass - for pv_name in writable_supported_variables: - if pv_name != "track_type": - if scalar_variable_cls is not None and isinstance( - supported_variables[pv_name], scalar_variable_cls - ): - original_value = ( - model.get(pv_name) + 0.001 - ) # add small offset to ensure set does something - assert original_value != 0, ( - f"Original value for {pv_name} is zero, cannot perform roundtrip test with offset" - ) - model.set({pv_name: original_value}) - roundtrip_value = model.get(pv_name) - assert_value_equal(pv_name, original_value, roundtrip_value) + if pv_name == "track_type": + continue + + current_value = model.get(pv_name) + + if isinstance(current_value, bool): + target_value = not current_value + elif isinstance(current_value, Real): + if pv_name.endswith(":PNEUMATIC"): + target_value = 0.0 if float(current_value) >= 0.5 else 1.0 + else: + target_value = float(current_value) + 0.001 + else: + # Roundtrip helper currently targets scalar-like writable controls. + continue + + model.set({pv_name: target_value}) + roundtrip_value = model.get(pv_name) + assert_value_equal(pv_name, target_value, roundtrip_value) diff --git a/virtual_accelerator/tests/test_cheetah_actions.py b/virtual_accelerator/tests/test_cheetah_actions.py new file mode 100644 index 0000000..6091be4 --- /dev/null +++ b/virtual_accelerator/tests/test_cheetah_actions.py @@ -0,0 +1,99 @@ +import math + +import pytest + +from virtual_accelerator.tests.dependency_profiles import HAS_CHEETAH_DEPS + +pytestmark = [ + pytest.mark.requires_cheetah, +] + +if HAS_CHEETAH_DEPS: + import torch + from cheetah.accelerator import Drift, Quadrupole, Segment + from cheetah.particles import ParticleBeam + from lume_cheetah.simulator import CheetahSimulator + + from virtual_accelerator.cheetah.actions import ( + QuadrupoleBACTVariable, + QuadrupoleBCTRLVariable, + ) + from virtual_accelerator.cheetah.variables import get_variables_from_segment +else: + pytest.skip("requires cheetah optional dependencies", allow_module_level=True) + + +class TestCheetahActions: + @pytest.fixture + def simulator(self): + segment = Segment( + [ + Quadrupole( + name="Q1", + length=torch.tensor(0.5), + k1=torch.tensor(1.0), + ), + Drift(name="D1", length=torch.tensor(1.0)), + ] + ) + beam = ParticleBeam.from_twiss( + beta_x=torch.tensor(1.0), + beta_y=torch.tensor(1.0), + num_particles=256, + energy=torch.tensor(1e6), + ) + return CheetahSimulator(segment=segment, initial_beam_distribution=beam) + + def test_quadrupole_bctrl_roundtrip(self, simulator): + variable = QuadrupoleBCTRLVariable( + name="QUAD:IN20:511:BCTRL", + element_name="Q1", + element_attribute="k1", + ) + + initial_value = float(variable._get(simulator)) + expected_initial = 1.0 * 0.5 * 33.356 * 1e6 / 1e9 + assert math.isclose(initial_value, expected_initial, rel_tol=0.0, abs_tol=1e-8) + + variable._set(simulator, torch.tensor(0.05)) + roundtrip = float(variable._get(simulator)) + + assert math.isclose(roundtrip, 0.05, rel_tol=0.0, abs_tol=1e-8) + + def test_quadrupole_bact_is_read_only(self, simulator): + variable = QuadrupoleBACTVariable( + name="QUAD:IN20:511:BACT", + element_name="Q1", + element_attribute="k1", + ) + + with pytest.raises(RuntimeError): + variable._set(simulator, torch.tensor(0.1)) + + def test_get_variables_from_segment_instantiates_actions(self): + segment = Segment( + [ + Quadrupole( + name="Q1", + length=torch.tensor(0.5), + k1=torch.tensor(1.0), + ), + Drift(name="D1", length=torch.tensor(1.0)), + ] + ) + + variables = get_variables_from_segment( + segment, + device_mapping={"Q1": "QUAD:IN20:511"}, + element_attr_mapping={ + "Quadrupole": { + "BCTRL": "QuadrupoleBCTRLVariable", + "BACT": "QuadrupoleBACTVariable", + } + }, + ) + + assert "QUAD:IN20:511:BCTRL" in variables + assert "QUAD:IN20:511:BACT" in variables + assert isinstance(variables["QUAD:IN20:511:BCTRL"], QuadrupoleBCTRLVariable) + assert isinstance(variables["QUAD:IN20:511:BACT"], QuadrupoleBACTVariable) diff --git a/virtual_accelerator/tests/test_cu_hxr.py b/virtual_accelerator/tests/test_cu_hxr.py index 11c5631..f66e488 100644 --- a/virtual_accelerator/tests/test_cu_hxr.py +++ b/virtual_accelerator/tests/test_cu_hxr.py @@ -7,33 +7,43 @@ from lume.exceptions import ReadOnlyError from virtual_accelerator.tests.dependency_profiles import ( HAS_BMAD_DEPS, + HAS_CHEETAH_DEPS, HAS_LCLS_LATTICE, ) +from virtual_accelerator.models.cu_hxr import ( + get_cu_hxr_bmad_model, + get_cu_hxr_cheetah_model, +) from virtual_accelerator.tests._bmad_model_test_utils import ( + assert_bpm_pvs_match_cheetah_segment, + assert_screen_image_pvs_match_cheetah_segment, TEST_BEAM_PATH, assert_bpm_pvs_match_tao_lattice, assert_bmad_model_initialization, assert_bmad_model_twiss_outputs, assert_bmad_model_track_beam_custom_path, + assert_magnet_pvs_match_cheetah_segment, assert_magnet_pvs_match_tao_lattice, assert_roundtrip_pv_get_set, assert_screen_image_pvs_in_supported_variables, ) -pytestmark = [ - pytest.mark.requires_bmad, - pytest.mark.requires_lcls_lattice, -] - -if HAS_BMAD_DEPS and HAS_LCLS_LATTICE: - from virtual_accelerator.models.cu_hxr import get_cu_hxr_bmad_model -else: - pytest.skip( - "requires bmad optional dependencies and LCLS_LATTICE", - allow_module_level=True, - ) +CU_HXR_PROFMON_CONFIG_PATH = ( + Path(__file__).resolve().parents[1] / "utils" / "cu_hxr_profmon_info.yaml" +) +def _load_cu_hxr_screen_config(): + with CU_HXR_PROFMON_CONFIG_PATH.open("r", encoding="utf-8") as config_file: + return yaml.safe_load(config_file) + + +@pytest.mark.requires_bmad +@pytest.mark.requires_lcls_lattice +@pytest.mark.skipif( + not HAS_BMAD_DEPS or not HAS_LCLS_LATTICE, + reason="requires bmad optional dependencies and LCLS_LATTICE", +) class TestCUHXRBmad: def test_initialization(self): assert_bmad_model_initialization( @@ -103,11 +113,7 @@ def test_cu_hxr_screen(self): assert not (image == updated_image).all() def test_cu_hxr_screen_resolution_matches_yaml_and_expected_range(self): - config_path = ( - Path(__file__).resolve().parents[1] / "utils" / "cu_hxr_profmon_info.yaml" - ) - with config_path.open("r", encoding="utf-8") as config_file: - screen_config = yaml.safe_load(config_file) + screen_config = _load_cu_hxr_screen_config() otr4_config = screen_config["OTR4"] resolution_pv = f"{otr4_config['name']}:RESOLUTION" @@ -129,20 +135,13 @@ def test_cu_hxr_lcavity(self): ampl = model.get("KLYS:LI21:31:ENLD") assert np.isclose(ampl, enld + 5.0) - def test_quadrupole_pvs_match_tao_lattice(self): - model = get_cu_hxr_bmad_model() - assert_magnet_pvs_match_tao_lattice(model, "Quadrupole") - - def test_hkicker_pvs_match_tao_lattice(self): - model = get_cu_hxr_bmad_model() - assert_magnet_pvs_match_tao_lattice(model, "HKicker") - - def test_vkicker_pvs_match_tao_lattice(self): - model = get_cu_hxr_bmad_model() - assert_magnet_pvs_match_tao_lattice(model, "VKicker") + @pytest.mark.parametrize("element_type", ["Quadrupole", "HKicker", "VKicker"]) + def test_magnet_pvs_match_tao_lattice(self, element_type): + model = get_cu_hxr_bmad_model(custom_beam_path=TEST_BEAM_PATH) + assert_magnet_pvs_match_tao_lattice(model, element_type) def test_bpm_pvs_match_tao_lattice(self): - model = get_cu_hxr_bmad_model() + model = get_cu_hxr_bmad_model(custom_beam_path=TEST_BEAM_PATH) assert_bpm_pvs_match_tao_lattice(model) def test_roundtrip_pv_get_set(self): @@ -150,3 +149,110 @@ def test_roundtrip_pv_get_set(self): custom_beam_path=TEST_BEAM_PATH, end_element="OTR4" ) assert_roundtrip_pv_get_set(model) + + +class TestCUHXRCheetah: + pytestmark = [ + pytest.mark.requires_cheetah, + pytest.mark.requires_lcls_lattice, + pytest.mark.skipif( + not HAS_CHEETAH_DEPS or not HAS_LCLS_LATTICE, + reason="requires cheetah optional dependencies and LCLS_LATTICE", + ), + ] + + def test_initialization(self): + model = get_cu_hxr_cheetah_model() + writable_control_variables = { + name + for name, variable in model.supported_variables.items() + if not getattr(variable, "read_only", True) + } + + assert len(model.supported_variables) > 0 + assert len(writable_control_variables) > 0 + + # Smoke test that reading all variables works after initialization. + _ = model.get(list(model.supported_variables)) + + def test_bact_readback_is_not_writable(self): + model = get_cu_hxr_cheetah_model() + + bact_pv = next( + name for name in model.supported_variables if name.endswith(":BACT") + ) + + with pytest.raises(ReadOnlyError, match="is read-only"): + model.set({bact_pv: 0.0}) + + def test_cu_hxr_screen(self): + model = get_cu_hxr_cheetah_model() + + image_pv = next( + name + for name in model.supported_variables + if name.endswith(":Image:ArrayData") + ) + image = np.asarray(model.get(image_pv)) + assert image.ndim == 2 + assert image.size > 0 + + control_pv = next( + name + for name, variable in model.supported_variables.items() + if name.endswith(":BCTRL") and not getattr(variable, "read_only", True) + ) + current_value = float(model.get(control_pv)) + target_value = current_value + 0.001 + model.set({control_pv: target_value}) + assert np.isclose(float(model.get(control_pv)), target_value) + + updated_image = np.asarray(model.get(image_pv)) + assert updated_image.shape == image.shape + assert np.isfinite(updated_image).all() + + def test_cu_hxr_screen_resolution_matches_yaml_and_expected_range(self): + screen_config = _load_cu_hxr_screen_config() + + model = get_cu_hxr_cheetah_model() + + resolution_pv, expected_resolution = next( + ( + f"{config_entry['name']}:RESOLUTION", + float(config_entry["pixel_size"]), + ) + for config_entry in screen_config.values() + if f"{config_entry['name']}:RESOLUTION" in model.supported_variables + ) + + resolution = float(model.get(resolution_pv)) + assert np.isclose(resolution, expected_resolution) + assert 5.0 < resolution < 30.0 + + def test_screen_pvs_match_cheetah_segment(self): + model = get_cu_hxr_cheetah_model() + assert_screen_image_pvs_match_cheetah_segment(model) + + def test_bpm_pvs_have_expected_suffixes(self): + model = get_cu_hxr_cheetah_model() + assert_bpm_pvs_match_cheetah_segment(model) + + def test_roundtrip_pv_get_set(self): + model = get_cu_hxr_cheetah_model() + assert_roundtrip_pv_get_set(model) + + @pytest.mark.parametrize( + ("element_type", "excluded_elements"), + [ + ("Quadrupole", ()), + ("HorizontalCorrector", ("xcapm2",)), + ("VerticalCorrector", ("ycapm2",)), + ], + ) + def test_magnet_pvs_match_cheetah_segment(self, element_type, excluded_elements): + model = get_cu_hxr_cheetah_model() + assert_magnet_pvs_match_cheetah_segment( + model, + element_type, + excluded_elements=excluded_elements, + ) diff --git a/virtual_accelerator/tests/test_facet.py b/virtual_accelerator/tests/test_facet.py index 8419d99..9d2af25 100644 --- a/virtual_accelerator/tests/test_facet.py +++ b/virtual_accelerator/tests/test_facet.py @@ -173,21 +173,12 @@ def test_staged_model(self): ) assert screen_pv in staged_model.supported_variables - def test_quadrupole_pvs_match_tao_lattice(self): - model = get_facet_bmad_model(end_element="PR10711") - assert_magnet_pvs_match_tao_lattice(model, "Quadrupole") - - def test_hkicker_pvs_match_tao_lattice(self): - model = get_facet_bmad_model(end_element="PR10711") - assert_magnet_pvs_match_tao_lattice(model, "HKicker") - - def test_vkicker_pvs_match_tao_lattice(self): - model = get_facet_bmad_model(end_element="PR10711") - assert_magnet_pvs_match_tao_lattice(model, "VKicker") - - def test_sbend_pvs_match_tao_lattice(self): + @pytest.mark.parametrize( + "element_type", ["Quadrupole", "HKicker", "VKicker", "SBend"] + ) + def test_magnet_pvs_match_tao_lattice(self, element_type): model = get_facet_bmad_model(end_element="PR10711") - assert_magnet_pvs_match_tao_lattice(model, "SBend") + assert_magnet_pvs_match_tao_lattice(model, element_type) def test_bpm_pvs_match_tao_lattice(self): model = get_facet_bmad_model(end_element="PR10711") diff --git a/virtual_accelerator/tests/test_optional_dependencies.py b/virtual_accelerator/tests/test_optional_dependencies.py index 2b68553..9ab1b14 100644 --- a/virtual_accelerator/tests/test_optional_dependencies.py +++ b/virtual_accelerator/tests/test_optional_dependencies.py @@ -31,7 +31,7 @@ def fake_import_module(_name): with pytest.raises(ModuleNotFoundError): import_optional( - "virtual_accelerator.cheetah.transformer", + "virtual_accelerator.cheetah.actions", feature="test feature", extra="cheetah", )