diff --git a/examples/other_examples/ControlSystem/tango/main.py b/examples/other_examples/ControlSystem/tango/main.py index 8fe0782b..4a21dc35 100644 --- a/examples/other_examples/ControlSystem/tango/main.py +++ b/examples/other_examples/ControlSystem/tango/main.py @@ -3,7 +3,7 @@ sr = Accelerator.load("config.yaml") # print the BPM position -bpm = sr.live.bpm.get("BPM_C01-01") # bpm is a BPM +bpm = sr.live.diagnostic.bpm.get("BPM_C01-01") # bpm is a BPM print(bpm.positions.get()) # Direct access to control system diff --git a/examples/use_cases/03-orbit_correction.ipynb b/examples/use_cases/03-orbit_correction.ipynb index e476b7a0..58cad3cc 100644 --- a/examples/use_cases/03-orbit_correction.ipynb +++ b/examples/use_cases/03-orbit_correction.ipynb @@ -185,24 +185,11 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "6adb7436", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "BPMs: 320, H correctors: 288, V correctors: 288\n" - ] - } - ], - "source": [ - "bpms = SR.bpms.get(\"BPM\")\n", - "hcorr = SR.magnets.get(\"HCorr\")\n", - "vcorr = SR.magnets.get(\"VCorr\")\n", - "print(f\"BPMs: {len(bpms)}, H correctors: {len(hcorr)}, V correctors: {len(vcorr)}\")" - ] + "outputs": [], + "source": "bpms = SR.diagnostic.bpms.get(\"BPM\")\nhcorr = SR.magnets.get(\"HCorr\")\nvcorr = SR.magnets.get(\"VCorr\")\nprint(f\"BPMs: {len(bpms)}, H correctors: {len(hcorr)}, V correctors: {len(vcorr)}\")" }, { "cell_type": "markdown", diff --git a/examples/use_cases/03-orbit_correction.py b/examples/use_cases/03-orbit_correction.py index ceef8ec1..b7af039c 100644 --- a/examples/use_cases/03-orbit_correction.py +++ b/examples/use_cases/03-orbit_correction.py @@ -84,7 +84,7 @@ # In[4]: -bpms = SR.bpms.get("BPM") +bpms = SR.diagnostic.bpms.get("BPM") hcorr = SR.magnets.get("HCorr") vcorr = SR.magnets.get("VCorr") print(f"BPMs: {len(bpms)}, H correctors: {len(hcorr)}, V correctors: {len(vcorr)}") diff --git a/pyaml/arrays/bpm.py b/pyaml/arrays/bpm.py index 4ea84927..dd9d2996 100644 --- a/pyaml/arrays/bpm.py +++ b/pyaml/arrays/bpm.py @@ -88,4 +88,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate the :py:class:`.BPMArray` with. """ - holder.bpms.add(self._name, self._elements) + holder.diagnostic.bpms.add(self._name, self._elements) diff --git a/pyaml/common/holders/diagnostic_holder.py b/pyaml/common/holders/diagnostic_holder.py new file mode 100644 index 00000000..361a06ed --- /dev/null +++ b/pyaml/common/holders/diagnostic_holder.py @@ -0,0 +1,121 @@ +"""Holder interface for diagnostics.""" + +from typing import TYPE_CHECKING + +from ...arrays.element_array import ElementArray +from ...diagnostics.tune_monitor import BetatronTuneMonitor +from ..element import Element, __pyaml_repr__ +from ..exception import PyAMLException +from .sub_holders import BPMHolder, BPMsHolder + +if TYPE_CHECKING: + from .element_holder import ElementHolder + + +class DiagnosticHolder: + """ + Provide access to diagnostics. + + Parameters + ---------- + peer : 'ElementHolder' + Parent holder containing the diagnostic store. + + Attributes + ---------- + bpm, bpms + Single BPM by name, or a named BPM array. + betatron_tune + Betatron tune monitor configured as ``BETATRON_TUNE``, the default, validated + against its expected class. + + Methods + ------- + get(name=None) + Return a named diagnostic, or all configured diagnostics when no name is given. + + Notes + ----- + :meth:`ElementHolder.get_betatron_tune_monitor + ` stays + available as the named, untyped lookup. This holder adds the default-name, + type-validated convenience property. + + Examples + -------- + >>> default_tune_monitor = sr.live.diagnostic.betatron_tune + >>> measured_tune = default_tune_monitor.tune.get() + >>> spare_tune_monitor = sr.live.diagnostic.get("SPARE_BETATRON_TUNE_MONITOR") + >>> all_diagnostics = sr.live.diagnostic.get() + >>> bpm = sr.live.diagnostic.bpm.get("BPM01") + >>> bpms = sr.live.diagnostic.bpms.get("BPMS") + """ + + def __init__(self, peer: "ElementHolder"): + """ + Initialize a diagnostic holder for an element holder. + """ + self._peer = peer + self._bpm_holder = BPMHolder(peer) + self._bpms_holder = BPMsHolder(peer) + + @property + def bpm(self) -> BPMHolder: + """Return the bpm.""" + return self._bpm_holder + + @property + def bpms(self) -> BPMsHolder: + """Return the bpms.""" + return self._bpms_holder + + def get(self, name: str = None) -> "Element | ElementArray": + """ + Return a named diagnostic, or all configured diagnostics when no name is given. + + Parameters + ---------- + name : str, optional + Name of the diagnostic to look up, as declared in the configuration. When + omitted, every configured diagnostic is returned instead. + + Returns + ------- + Element or ElementArray + The diagnostic registered under ``name``, or an + :class:`~pyaml.arrays.element_array.ElementArray` holding every configured + diagnostic, in insertion order, when ``name`` is omitted. + + Raises + ------ + PyAMLException + If ``name`` is given and no diagnostic is registered under it. + """ + if name is None: + return ElementArray("", list(self._peer._DIAG.values())) + return self._peer._get_diagnostic(name) + + @property + def betatron_tune(self) -> BetatronTuneMonitor: + """ + Return the betatron tune monitor configured as ``BETATRON_TUNE``. + + Returns + ------- + BetatronTuneMonitor + Betatron tune monitor registered under ``BETATRON_TUNE``. + + Raises + ------ + PyAMLException + If no diagnostic is registered under ``BETATRON_TUNE``, or if it is not a + :class:`~pyaml.diagnostics.tune_monitor.BetatronTuneMonitor`. + """ + name = "BETATRON_TUNE" + obj = self._peer.get_betatron_tune_monitor(name) + if not isinstance(obj, BetatronTuneMonitor): + raise PyAMLException(f"{name}: BetatronTuneMonitor expected but got {type(obj).__name__}") + return obj + + def __repr__(self): + return __pyaml_repr__(self) diff --git a/pyaml/common/holders/element_holder.py b/pyaml/common/holders/element_holder.py index aa2b16dc..2a0fb27e 100644 --- a/pyaml/common/holders/element_holder.py +++ b/pyaml/common/holders/element_holder.py @@ -17,10 +17,9 @@ from ..abstract_aggregator import ScalarAggregator from ..element import Element from ..exception import PyAMLException +from .diagnostic_holder import DiagnosticHolder from .rf_holder import RFHolder from .sub_holders import ( - BPMHolder, - BPMsHolder, CombinedFunctionMagnetHolder, CombinedFunctionMagnetsHolder, MagnetHolder, @@ -28,6 +27,7 @@ SerializedMagnetHolder, SerializedMagnetsHolder, ) +from .tool_holder import ToolHolder if TYPE_CHECKING: from ...accelerator import Accelerator @@ -57,18 +57,21 @@ class ElementHolder(metaclass=ABCMeta): ---------- magnet, magnets Single magnet by name, or a named magnet array. - bpm, bpms - Single BPM by name, or a named BPM array. combined_function_magnet, combined_function_magnets Single combined-function magnet by name, or a named array. serialized_magnet, serialized_magnets Single serialized magnet group by name, or a named array. rf RF plant and transmitters of this mode. + diagnostic + Diagnostics of this mode, with typed default-name access. + tool + Tuning and measurement tools of this mode, with typed default-name access. tune, chromaticity, orbit, dispersion - Tuning tools attached to this mode, looked up by name. + Backward-compatible aliases for ``tool.tune``, ``tool.chromaticity``, ``tool.orbit`` + and ``tool.dispersion``. trm, crm, orm - Response-matrix measurement tools, looked up by name. + Backward-compatible aliases for ``tool.trm``, ``tool.crm`` and ``tool.orm``. Methods ------- @@ -156,9 +159,9 @@ def __init__(self): self._serialized_magnets_holder = SerializedMagnetsHolder(self) self._combined_function_magnet_holder = CombinedFunctionMagnetHolder(self) self._combined_function_magnets_holder = CombinedFunctionMagnetsHolder(self) - self._bpm_holder = BPMHolder(self) - self._bpms_holder = BPMsHolder(self) self._rf_holder = RFHolder(self) + self._diagnostic_holder = DiagnosticHolder(self) + self._tool_holder = ToolHolder(self) @property def peer(self) -> "Accelerator": @@ -198,19 +201,19 @@ def combined_function_magnets(self) -> CombinedFunctionMagnetsHolder: return self._combined_function_magnets_holder @property - def bpm(self) -> BPMHolder: - """Return the bpm.""" - return self._bpm_holder + def rf(self) -> RFHolder: + """Return the rf.""" + return self._rf_holder @property - def bpms(self) -> BPMsHolder: - """Return the bpms.""" - return self._bpms_holder + def diagnostic(self) -> DiagnosticHolder: + """Return the diagnostic.""" + return self._diagnostic_holder @property - def rf(self) -> RFHolder: - """Return the rf.""" - return self._rf_holder + def tool(self) -> ToolHolder: + """Return the tool.""" + return self._tool_holder def post_init(self): """Run post-initialization hooks for every stored element.""" diff --git a/pyaml/common/holders/generic_array_holder.py b/pyaml/common/holders/generic_array_holder.py index 0ce62441..4004a52a 100644 --- a/pyaml/common/holders/generic_array_holder.py +++ b/pyaml/common/holders/generic_array_holder.py @@ -42,6 +42,13 @@ class GenericArrayHolder(Generic[T, A]): Return a named array or a transient array of all elements. add(arrayName, elementNames) Create and register a named array from element selectors. + + Notes + ----- + A configured array is also reachable as an attribute when its name is a + valid Python identifier, e.g. ``holder.QuadForTune`` is equivalent to + ``holder.get("QuadForTune")``. Array names appear in ``dir(holder)`` so + interactive completion (IPython, Jupyter) discovers them. """ def __init__( @@ -117,5 +124,49 @@ def __getitem__(self, key): """ return self.get().__getitem__(key) + def __getattr__(self, name: str) -> A: + """ + Return a configured array through attribute access. + + Only called when normal attribute lookup fails, so it never shadows + :meth:`get`, :meth:`add`, or any other existing attribute. + + Parameters + ---------- + name : str + Configured array name. Must be a valid Python identifier. + + Returns + ------- + A + The array registered under ``name``. + + Raises + ------ + AttributeError + If ``name`` starts with an underscore, is not a valid Python + identifier, or does not match a configured array. + + Examples + -------- + >>> quad_family = sr.live.magnets.get("QuadForTune") + >>> same_quad_family = sr.live.magnets.QuadForTune + """ + if name.startswith("_") or not name.isidentifier() or name not in self._array_store: + raise AttributeError(f"'{type(self).__name__}' object has no array named '{name}'") + return self._array_store[name] + + def __dir__(self) -> list[str]: + """ + List attributes, including configured array names. + + Returns + ------- + list of str + Default attributes plus configured array names that are valid + Python identifiers, for interactive completion (IPython, Jupyter). + """ + return sorted(set(super().__dir__()) | {name for name in self._array_store if name.isidentifier()}) + def __repr__(self): return __pyaml_repr__(self) diff --git a/pyaml/common/holders/sub_holders.py b/pyaml/common/holders/sub_holders.py index 399485eb..f5e36b92 100644 --- a/pyaml/common/holders/sub_holders.py +++ b/pyaml/common/holders/sub_holders.py @@ -42,6 +42,19 @@ class MagnetsHolder(GenericArrayHolder[Magnet, MagnetArray]): ---------- peer : 'ElementHolder' Parent element holder that owns this specialized holder. + + Notes + ----- + A configured magnet family is also reachable as an attribute when its + name is a valid Python identifier, e.g. ``magnets.QuadForTune`` is + equivalent to ``magnets.get("QuadForTune")``. Family names appear in + ``dir(magnets)`` for interactive completion. + + Examples + -------- + >>> quad_family = sr.live.magnets.get("QuadForTune") + >>> same_quad_family = sr.live.magnets.QuadForTune + >>> combined_function_magnets = sr.live.magnets.get_cfm() """ def __init__(self, peer: "ElementHolder"): @@ -57,6 +70,22 @@ def __init__(self, peer: "ElementHolder"): "Magnet array", ) + def get_cfm(self) -> CombinedFunctionMagnetArray: + """ + Return all combined-function magnets. + + Returns + ------- + CombinedFunctionMagnetArray + New unnamed container with every registered combined-function + magnet. + + Examples + -------- + >>> combined_function_magnets = sr.live.magnets.get_cfm() + """ + return self._peer.combined_function_magnets.get() + class CombinedFunctionMagnetHolder(GenericElementHolder[CombinedFunctionMagnet]): """ @@ -83,6 +112,13 @@ class CombinedFunctionMagnetsHolder(GenericArrayHolder[CombinedFunctionMagnet, C ---------- peer : 'ElementHolder' Parent element holder that owns this specialized holder. + + Notes + ----- + A configured array is also reachable as an attribute when its name is a + valid Python identifier, e.g. ``combined_function_magnets.CFM`` is + equivalent to ``combined_function_magnets.get("CFM")``. Array names + appear in ``dir(combined_function_magnets)`` for interactive completion. """ def __init__(self, peer: "ElementHolder"): @@ -124,6 +160,13 @@ class SerializedMagnetsHolder(GenericArrayHolder[SerializedMagnets, SerializedMa ---------- peer : 'ElementHolder' Parent element holder that owns this specialized holder. + + Notes + ----- + A configured array is also reachable as an attribute when its name is a + valid Python identifier, e.g. ``serialized_magnets.QForTune`` is + equivalent to ``serialized_magnets.get("QForTune")``. Array names appear + in ``dir(serialized_magnets)`` for interactive completion. """ def __init__(self, peer: "ElementHolder"): @@ -165,17 +208,25 @@ class BPMsHolder(GenericArrayHolder[BPM, BPMArray]): ---------- peer : 'ElementHolder' Parent element holder that owns this specialized holder. + + Notes + ----- + A configured array is also reachable as an attribute when its name is a + valid Python identifier, e.g. ``bpms.BPMS`` is equivalent to + ``bpms.get("BPMS")``. Array names appear in ``dir(bpms)`` for interactive + completion. """ def __init__(self, peer: "ElementHolder"): """ Initialize the BPMsHolder. """ + bpm_holder = BPMHolder(peer) super().__init__( peer, peer._BPM_ARRAYS, - peer.bpm.all, - peer.bpm.get, + bpm_holder.all, + bpm_holder.get, BPMArray, "BPM array", ) diff --git a/pyaml/common/holders/tool_holder.py b/pyaml/common/holders/tool_holder.py new file mode 100644 index 00000000..be7d17a0 --- /dev/null +++ b/pyaml/common/holders/tool_holder.py @@ -0,0 +1,159 @@ +"""Holder interface for tuning and measurement tools.""" + +from typing import TYPE_CHECKING + +from ...arrays.element_array import ElementArray +from ..element import Element, __pyaml_repr__ +from ..exception import PyAMLException + +if TYPE_CHECKING: + from ...tuning_tools.chromaticity import Chromaticity + from ...tuning_tools.chromaticity_response_matrix import ChromaticityResponseMatrix + from ...tuning_tools.dispersion import Dispersion + from ...tuning_tools.orbit import Orbit + from ...tuning_tools.orbit_response_matrix import OrbitResponseMatrix + from ...tuning_tools.tune import Tune + from ...tuning_tools.tune_response_matrix import TuneResponseMatrix + from .element_holder import ElementHolder + + +class ToolHolder: + """ + Provide access to tuning and measurement tools. + + Parameters + ---------- + peer : 'ElementHolder' + Parent holder containing the tool store. + + Attributes + ---------- + tune, chromaticity, orbit, dispersion + Default tuning tool for each kind, resolved from its ``DEFAULT_...`` name and + validated against the expected class. + trm, crm, orm + Default response-matrix measurement tool for each kind, resolved the same way. + + Methods + ------- + get(name=None) + Return a named tool, or all configured tools when no name is given. + + Notes + ----- + :attr:`ElementHolder.tune `, + ``.trm``, ``.orbit``, ``.orm``, ``.chromaticity``, ``.crm`` and ``.dispersion`` are + backward-compatible aliases for the corresponding properties here. + + Examples + -------- + >>> tune_correction = sr.live.tool.tune + >>> tune_correction.set([0.31, 0.22]) + >>> tune_response = sr.live.tool.trm + >>> tune_response.measure() + >>> other_tool = sr.live.tool.get("OTHER_TUNE_CORRECTION") + >>> all_tools = sr.live.tool.get() + """ + + def __init__(self, peer: "ElementHolder"): + """ + Initialize a tool holder for an element holder. + """ + self._peer = peer + + def get(self, name: str = None) -> "Element | ElementArray": + """ + Return a named tool, or all configured tools when no name is given. + + Parameters + ---------- + name : str, optional + Name of the tuning or measurement tool to look up, as declared in the + configuration. When omitted, every configured tool is returned instead. + + Returns + ------- + Element or ElementArray + The tool registered under ``name``, or an :class:`~pyaml.arrays.element_array.ElementArray` + holding every configured tool, in insertion order, when ``name`` is omitted. + + Raises + ------ + PyAMLException + If ``name`` is given and no tool is registered under it. + """ + if name is None: + return ElementArray("", list(self._peer._TOOLS.values())) + return self._peer._get_tool(name) + + def _validate_type(self, name: str, obj: Element, expected_type: type) -> Element: + """ + Ensure a resolved default tool matches the type its accessor expects. + + Raises + ------ + PyAMLException + If ``obj`` is not an instance of ``expected_type``. + """ + if not isinstance(obj, expected_type): + raise PyAMLException(f"{name}: {expected_type.__name__} expected but got {type(obj).__name__}") + return obj + + @property + def chromaticity(self) -> "Chromaticity": + """Return the chromaticity tuning tool configured as ``DEFAULT_CHROMATICITY_CORRECTION``.""" + from ...tuning_tools.chromaticity import Chromaticity + + name = "DEFAULT_CHROMATICITY_CORRECTION" + return self._validate_type(name, self._peer.get_chromaticity_tuning(name), Chromaticity) + + @property + def crm(self) -> "ChromaticityResponseMatrix": + """Return the chromaticity response-matrix tool configured as ``DEFAULT_CHROMATICITY_RESPONSE_MATRIX``.""" + from ...tuning_tools.chromaticity_response_matrix import ChromaticityResponseMatrix + + name = "DEFAULT_CHROMATICITY_RESPONSE_MATRIX" + return self._validate_type(name, self._peer.get_crm_tuning(name), ChromaticityResponseMatrix) + + @property + def tune(self) -> "Tune": + """Return the tune correction tool configured as ``DEFAULT_TUNE_CORRECTION``.""" + from ...tuning_tools.tune import Tune + + name = "DEFAULT_TUNE_CORRECTION" + return self._validate_type(name, self._peer.get_tune_tuning(name), Tune) + + @property + def trm(self) -> "TuneResponseMatrix": + """Return the tune response-matrix tool configured as ``DEFAULT_TUNE_RESPONSE_MATRIX``.""" + from ...tuning_tools.tune_response_matrix import TuneResponseMatrix + + name = "DEFAULT_TUNE_RESPONSE_MATRIX" + return self._validate_type(name, self._peer.get_trm_tuning(name), TuneResponseMatrix) + + @property + def orbit(self) -> "Orbit": + """Return the orbit correction tool configured as ``DEFAULT_ORBIT_CORRECTION``.""" + from ...tuning_tools.orbit import Orbit + + name = "DEFAULT_ORBIT_CORRECTION" + return self._validate_type(name, self._peer.get_orbit_tuning(name), Orbit) + + @property + def orm(self) -> "OrbitResponseMatrix": + """Return the orbit response-matrix tool configured as ``DEFAULT_ORBIT_RESPONSE_MATRIX``.""" + from ...tuning_tools.orbit_response_matrix import OrbitResponseMatrix + + name = "DEFAULT_ORBIT_RESPONSE_MATRIX" + return self._validate_type(name, self._peer.get_orm_tuning(name), OrbitResponseMatrix) + + @property + def dispersion(self) -> "Dispersion": + """Return the dispersion tool configured as ``DEFAULT_DISPERSION``.""" + from ...tuning_tools.dispersion import Dispersion + + name = "DEFAULT_DISPERSION" + return self._validate_type(name, self._peer.get_dispersion_tuning(name), Dispersion) + + def __repr__(self): + return __pyaml_repr__(self) diff --git a/pyaml/control/controlsystem.py b/pyaml/control/controlsystem.py index 9199f98d..d2e25470 100644 --- a/pyaml/control/controlsystem.py +++ b/pyaml/control/controlsystem.py @@ -248,7 +248,7 @@ def _fill_bpm(self, bpm: BPM) -> None: positions = RBpmArray(position_devices[0], position_devices[1]) tilt = RWBpmTiltScalar(tilt_devices[0]) offsets = RWBpmOffsetArray(offset_devices[0], offset_devices[1]) - self.bpm.add(bpm.attach(self, positions, offsets, tilt)) + self.diagnostic.bpm.add(bpm.attach(self, positions, offsets, tilt)) def _fill_rf_plant(self, rf_plant: RFPlant) -> None: attached_transmitters: list[RFTransmitter] = [] diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index 3be5196e..b9e60493 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -56,7 +56,7 @@ def __init__( """ self.element_holder = element_holder - self.bpm_array = element_holder.bpms.get(bpm_array_name) + self.bpm_array = element_holder.diagnostic.bpms.get(bpm_array_name) self.rf_plant_name = rf_plant_name if rf_plant_name is not None: diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index 071d66a8..b4a0a0a6 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -291,7 +291,9 @@ def _fill_bpm(self, bpm: BPM) -> None: if len(bpm_elt.Offset) != 2: raise PyAMLException(f"BPM {bpm.get_name()} offset must be a 2-element array.") update_bpm_transform_matrix(bpm_elt) - self.bpm.add(bpm.attach(self, RBpmArray(bpm_elt, self.ring), RWBpmOffsetArray(bpm_elt), RWBpmTiltScalar(bpm_elt))) + self.diagnostic.bpm.add( + bpm.attach(self, RBpmArray(bpm_elt, self.ring), RWBpmOffsetArray(bpm_elt), RWBpmTiltScalar(bpm_elt)) + ) def _fill_rf_plant(self, rf_plant: RFPlant) -> None: if rf_plant.transmitters: diff --git a/pyaml/tuning_tools/bba.py b/pyaml/tuning_tools/bba.py index ed0b8cc2..51f53934 100644 --- a/pyaml/tuning_tools/bba.py +++ b/pyaml/tuning_tools/bba.py @@ -147,13 +147,13 @@ def __init__( def bpms(self) -> "BPMArray": """Return the BPM array used for the measurement.""" self.check_peer() - return self.peer.bpms.get(self.bpm_array_name) + return self.peer.diagnostic.bpms.get(self.bpm_array_name) @property def bpm(self) -> "BPM": """Return the reference BPM used for the alignment.""" self.check_peer() - return self.peer.bpm.get(self.bpm_name) + return self.peer.diagnostic.bpm.get(self.bpm_name) @property def hcorrector(self) -> "Magnet": diff --git a/pyaml/tuning_tools/bba2.py b/pyaml/tuning_tools/bba2.py index b1b7bb07..c265dbab 100644 --- a/pyaml/tuning_tools/bba2.py +++ b/pyaml/tuning_tools/bba2.py @@ -228,13 +228,13 @@ def __init__( def bpms(self) -> "BPMArray": """Return the BPM array used for the measurement.""" self.check_peer() - return self.peer.bpms.get(self.bpm_array_name) + return self.peer.diagnostic.bpms.get(self.bpm_array_name) @property def bpm(self) -> "BPM": """Return the reference BPM used for the alignment.""" self.check_peer() - return self.peer.bpm.get(self.bpm_name) + return self.peer.diagnostic.bpm.get(self.bpm_name) @property def hcorrector(self) -> "Magnet": @@ -343,7 +343,7 @@ def _init_responses( quad = design.magnet.get(quadname) sth = design.magnet.get(steererhname) stv = design.magnet.get(steerervname) - orbit = design.bpms.get(bpmname).positions + orbit = design.diagnostic.bpms.get(bpmname).positions tune_design = design.get_tune_tuning(tunename) tune_live = self._peer.get_tune_tuning(tunename) diff --git a/pyaml/tuning_tools/chromaticity_monitor.py b/pyaml/tuning_tools/chromaticity_monitor.py index c1da2507..fa59ee1a 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/pyaml/tuning_tools/chromaticity_monitor.py @@ -203,7 +203,7 @@ def bpms(self) -> "BPMArray | None": self.check_peer() if self.bpm_array_name is None: return None - return self.peer.bpms.get(self.bpm_array_name) + return self.peer.diagnostic.bpms.get(self.bpm_array_name) @property def chromaticity(self) -> ReadFloatArray: diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index 6bbc3844..a5ce10e2 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/pyaml/tuning_tools/dispersion.py @@ -81,7 +81,7 @@ def __init__(self, name: str, bpm_array_name: str, rf_plant_name: str, frequency def bpms(self) -> "BPMArray": """Return the BPM array used for orbit readback.""" self.check_peer() - return self.peer.bpms.get(self.bpm_array_name) + return self.peer.diagnostic.bpms.get(self.bpm_array_name) @property def rf_plant(self) -> "RFPlant": diff --git a/pyaml/tuning_tools/orbit.py b/pyaml/tuning_tools/orbit.py index 3f2c8b39..cb649f54 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -199,7 +199,7 @@ def response_matrix(self) -> OrbitResponseMatrixData | None: def bpms(self) -> BPMArray: """Return the BPM array used for orbit readback.""" self.check_peer() - return self.peer.bpms.get(self.bpm_array_name) + return self.peer.diagnostic.bpms.get(self.bpm_array_name) @property def hcorrectors(self) -> MagnetArray: diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 7b065b49..0d725c0f 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -130,7 +130,7 @@ def __init__( def bpms(self) -> "BPMArray": """Return the BPM array used for orbit readback.""" self.check_peer() - return self.peer.bpms.get(self.bpm_array_name) + return self.peer.diagnostic.bpms.get(self.bpm_array_name) @property def hcorrectors(self) -> "MagnetArray": diff --git a/tests/arrays/test_arrays.py b/tests/arrays/test_arrays.py index 7ae9b871..62f9c202 100644 --- a/tests/arrays/test_arrays.py +++ b/tests/arrays/test_arrays.py @@ -128,32 +128,32 @@ def test_arrays(install_test_package): # Test BPMs array # Using aggregator - pos = sr.design.bpms.get("BPMS").positions.get() + pos = sr.design.diagnostic.bpms.get("BPMS").positions.get() assert np.abs(pos[0][0] - 7.22262850488348e-05) < 1e-10 assert np.abs(pos[0][1] - 3.4291613955705856e-05) < 1e-10 assert np.abs(pos[1][0] + 1.1696152238807462e-04) < 1e-10 assert np.abs(pos[1][1] - 7.4265634524358045e-06) < 1e-10 # Using aggregator (h and v) - pos_h = sr.design.bpms.get("BPMS").h.get() - pos_v = sr.design.bpms.get("BPMS").v.get() + pos_h = sr.design.diagnostic.bpms.get("BPMS").h.get() + pos_v = sr.design.diagnostic.bpms.get("BPMS").v.get() assert np.all(np.isclose(pos[:, 0], pos_h, rtol=1e-15, atol=1e-15)) assert np.all(np.isclose(pos[:, 1], pos_v, rtol=1e-15, atol=1e-15)) # Test BPM transformation matrices - sr.design.bpm.get("BPM_C04-01").offset.set([0.1, 0.2]) - sr.design.bpm.get("BPM_C04-02").offset.set([0.3, 0.4]) - pos = sr.design.bpms.get("BPMS").positions.get() + sr.design.diagnostic.bpm.get("BPM_C04-01").offset.set([0.1, 0.2]) + sr.design.diagnostic.bpm.get("BPM_C04-02").offset.set([0.3, 0.4]) + pos = sr.design.diagnostic.bpms.get("BPMS").positions.get() assert np.abs(pos[0][0] - 7.22262850488348e-05 - 0.1) < 1e-10 assert np.abs(pos[0][1] - 3.4291613955705856e-05 - 0.2) < 1e-10 assert np.abs(pos[1][0] + 1.1696152238807462e-04 - 0.3) < 1e-10 assert np.abs(pos[1][1] - 7.4265634524358045e-06 - 0.4) < 1e-10 - sr.design.bpm.get("BPM_C04-01").offset.set([0.0, 0.0]) - sr.design.bpm.get("BPM_C04-02").offset.set([0.0, 0.0]) + sr.design.diagnostic.bpm.get("BPM_C04-01").offset.set([0.0, 0.0]) + sr.design.diagnostic.bpm.get("BPM_C04-02").offset.set([0.0, 0.0]) # No aggregator bpms = [] - for b in sr.design.bpms.get("BPMS"): + for b in sr.design.diagnostic.bpms.get("BPMS"): bpms.append(b) bpms = BPMArray("BPM_noagg", bpms, use_aggregator=False) @@ -175,7 +175,7 @@ def test_arrays(install_test_package): assert len(allMags) == 7 # Create an array that contains all BPM - allBpms = BPMArray("AllBPMs", sr.design.bpm.all()) + allBpms = BPMArray("AllBPMs", sr.design.diagnostic.bpm.all()) assert len(allBpms) == 2 cfm = sr.design.combined_function_magnets.get("CFM") @@ -204,7 +204,7 @@ def test_arrays(install_test_package): assert np.abs(strHVSQ[4] + 0.000017) < 1e-10 # V assert np.abs(strHVSQ[5] - 1e-6) < 1e-10 # SQ - bpmsLive = BPMArray("", sr.live.bpm.all()) + bpmsLive = BPMArray("", sr.live.diagnostic.bpm.all()) bpmsLive.positions.get() # Test dynamic arrays @@ -232,7 +232,7 @@ def test_arrays(install_test_package): emptyBPM = BPM(name="emptyBPM", elements=[]) emptyBPM.fill_array(sr.design) # Attach the array - v = sr.design.bpms.get("emptyBPM").positions.get() # Ensure good attach + v = sr.design.diagnostic.bpms.get("emptyBPM").positions.get() # Ensure good attach assert np.shape(v) == (0,) emptyCFM = CombinedFunctionMagnet(name="emptyCFM", elements=[]) diff --git a/tests/bpm/test_bpm.py b/tests/bpm/test_bpm.py index 93ac68ec..3bf8ee71 100644 --- a/tests/bpm/test_bpm.py +++ b/tests/bpm/test_bpm.py @@ -9,7 +9,7 @@ def test_simulator_bpm_tilt(): sr.design.get_lattice().disable_6d() sr.design.magnet.get("SH1A-C01-H").strength.set(10e-6) # Add orbit sr.design.magnet.get("SH1A-C01-V").strength.set(10e-6) # Add orbit - bpm = sr.design.bpm.get("BPM_C01-01") + bpm = sr.design.diagnostic.bpm.get("BPM_C01-01") assert np.allclose(bpm.positions.get(), np.array([5.90809968e-05, 2.24832853e-05])) assert bpm.tilt.get() == 0 alpha = np.pi / 3 @@ -30,7 +30,7 @@ def test_simulator_bpm_tilt(): def test_simulator_bpm_offset(): sr: Accelerator = Accelerator.load("tests/config/bpms.yaml", ignore_external=True) sr.design.get_lattice().disable_6d() - bpm = sr.design.bpm.get("BPM_C01-01") + bpm = sr.design.diagnostic.bpm.get("BPM_C01-01") assert bpm.offset.get()[0] == 0 assert bpm.offset.get()[1] == 0 @@ -48,8 +48,8 @@ def test_simulator_bpm_offset(): def test_simulator_bpm_position(install_test_package): sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") sr.design.get_lattice().disable_6d() - bpm = sr.design.bpm.get("BPM_C01-01") - bpm_simple = sr.live.bpm.get("BPM_C01-02") + bpm = sr.design.diagnostic.bpm.get("BPM_C01-01") + bpm_simple = sr.live.diagnostic.bpm.get("BPM_C01-02") assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) assert np.allclose(bpm_simple.positions.get(), np.array([0.0, 0.0])) @@ -58,9 +58,9 @@ def test_simulator_bpm_position(install_test_package): def test_simulator_bpm_position_with_bad_corrector_strength(): sr: Accelerator = Accelerator.load("tests/config/bpms.yaml", ignore_external=True) sr.design.get_lattice().disable_6d() - bpm1 = sr.design.bpm.get("BPM_C01-01") - bpm_simple = sr.design.bpm.get("BPM_C01-02") - bpm3 = sr.design.bpm.get("BPM_C01-03") + bpm1 = sr.design.diagnostic.bpm.get("BPM_C01-01") + bpm_simple = sr.design.diagnostic.bpm.get("BPM_C01-02") + bpm3 = sr.design.diagnostic.bpm.get("BPM_C01-03") sr.design.magnet.get("SH1A-C01-H").strength.set(-1e-6) sr.design.magnet.get("SH1A-C01-V").strength.set(-1e-6) diff --git a/tests/bpm/test_bpm_controlsystem.py b/tests/bpm/test_bpm_controlsystem.py index 13321ca1..8b534536 100644 --- a/tests/bpm/test_bpm_controlsystem.py +++ b/tests/bpm/test_bpm_controlsystem.py @@ -11,7 +11,7 @@ ) def test_controlsystem_bpm_tilt(install_test_package): sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.bpm.get("BPM_C01-01") + bpm = sr.live.diagnostic.bpm.get("BPM_C01-01") print(bpm.tilt.get()) assert bpm.tilt.get() == 0 @@ -26,7 +26,7 @@ def test_controlsystem_bpm_tilt(install_test_package): ) def test_controlsystem_bpm_offset(install_test_package): sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.bpm.get("BPM_C01-01") + bpm = sr.live.diagnostic.bpm.get("BPM_C01-01") assert bpm.offset.get()[0] == 0 assert bpm.offset.get()[1] == 0 @@ -43,8 +43,8 @@ def test_controlsystem_bpm_offset(install_test_package): ) def test_controlsystem_bpm_position(install_test_package): sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.bpm.get("BPM_C01-01") - bpm_simple = sr.live.bpm.get("BPM_C01-02") + bpm = sr.live.diagnostic.bpm.get("BPM_C01-01") + bpm_simple = sr.live.diagnostic.bpm.get("BPM_C01-02") assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) assert np.allclose(bpm_simple.positions.get(), np.array([0.0, 0.0])) @@ -91,6 +91,6 @@ def test_controlsystem_bpm_position_indexed(install_test_package): set_attribute("srdiag/bpm/c01-04/Position", [0.0, 1.0], unit="mm") sr: Accelerator = Accelerator.load("tests/config/bpms.yaml") - bpm = sr.live.bpm.get("BPM_C01-04") + bpm = sr.live.diagnostic.bpm.get("BPM_C01-04") assert np.allclose(bpm.positions.get(), np.array([0.0, 1.0])) diff --git a/tests/common/test_array_holder_navigation.py b/tests/common/test_array_holder_navigation.py new file mode 100644 index 00000000..7e630268 --- /dev/null +++ b/tests/common/test_array_holder_navigation.py @@ -0,0 +1,99 @@ +import pytest + +from pyaml.accelerator import Accelerator +from pyaml.arrays.cfm_magnet_array import CombinedFunctionMagnetArray + + +@pytest.fixture +def holder(accelerator_from_fragments, sr_configuration_fragments): + sr = accelerator_from_fragments(*sr_configuration_fragments) + sr.design.get_lattice().disable_6d() + return sr.design + + +@pytest.fixture +def serialized_holder(): + sr = Accelerator.load("tests/config/sr_serialized_magnets.yaml", include_locations=False, ignore_external=True) + return sr.design + + +def test_dynamic_attribute_returns_named_magnet_array(holder): + assert holder.magnets.HCORR is holder.magnets.get("HCORR") + assert holder.magnets.VCORR is holder.magnets.get("VCORR") + assert holder.magnets.HVCORR is holder.magnets.get("HVCORR") + + +def test_dynamic_attribute_returns_named_combined_function_magnet_array(holder): + assert holder.combined_function_magnets.CFM is holder.combined_function_magnets.get("CFM") + + +def test_dynamic_attribute_returns_named_bpm_array(holder): + assert holder.diagnostic.bpms.BPMS is holder.diagnostic.bpms.get("BPMS") + + +def test_dynamic_attribute_returns_named_serialized_magnet_array(serialized_holder): + assert serialized_holder.serialized_magnets.QForTune is serialized_holder.serialized_magnets.get("QForTune") + assert serialized_holder.serialized_magnets.series is serialized_holder.serialized_magnets.get("series") + + +def test_dynamic_attribute_unknown_name_raises_attribute_error(holder): + with pytest.raises(AttributeError): + _ = holder.magnets.UNKNOWN + + +def test_dynamic_attribute_does_not_shadow_existing_methods(holder): + holder.magnets._array_store["get"] = holder.magnets.get("HCORR") + holder.magnets._array_store["add"] = holder.magnets.get("HCORR") + + assert callable(holder.magnets.get) + assert callable(holder.magnets.add) + assert holder.magnets.get("HCORR").names() == ["SH1A-C01-H", "SH1A-C02-H"] + + +def test_dynamic_attribute_rejects_non_identifier_names(holder): + holder.magnets._array_store["not-an-id"] = holder.magnets.get("HCORR") + + with pytest.raises(AttributeError): + getattr(holder.magnets, "not-an-id") + + +def test_dir_includes_configured_array_names(holder): + assert {"HCORR", "VCORR", "HVCORR"} <= set(dir(holder.magnets)) + assert "CFM" in dir(holder.combined_function_magnets) + assert "BPMS" in dir(holder.diagnostic.bpms) + + +def test_dir_excludes_non_identifier_array_names(holder): + holder.magnets._array_store["not-an-id"] = holder.magnets.get("HCORR") + + assert "not-an-id" not in dir(holder.magnets) + + +def test_get_cfm_returns_all_combined_function_magnets(holder): + combined_function_magnets = holder.magnets.get_cfm() + + assert type(combined_function_magnets) is CombinedFunctionMagnetArray + assert combined_function_magnets.names() == holder.combined_function_magnets.get().names() + + +def test_issue_373_example(): + """Reproduce the #373 example verbatim, against real test lattice data.""" + sr = Accelerator.load("tests/config/EBSTune-patterns.yaml", ignore_external=True) + sr.design.get_lattice().disable_6d() + + all_magnets = sr.design.magnets[:] + quad_family = sr.design.magnets.get("QForTune") + same_quad_family = sr.design.magnets.QForTune + assert same_quad_family.names() == quad_family.names() + assert len(quad_family) == 124 + assert len(all_magnets) >= len(quad_family) + + combined_function_magnets = sr.design.magnets.get_cfm() + assert len(combined_function_magnets) == 0 + + matching_quadrupoles = sr.design.magnets["QF1*"] + assert matching_quadrupoles.names() == sr.design.magnets["QF1*"].names() + + one_magnet = sr.design.magnet.get("QF1E-C04") + one_magnet.strength.set(0.8) + assert one_magnet.strength.get() == pytest.approx(0.8) diff --git a/tests/common/test_element_holder_collection.py b/tests/common/test_element_holder_collection.py index 937d2da7..4a55fadf 100644 --- a/tests/common/test_element_holder_collection.py +++ b/tests/common/test_element_holder_collection.py @@ -16,7 +16,7 @@ def holder(accelerator_from_fragments, sr_configuration_fragments): def test_exact_name_returns_the_element_or_none(holder): - assert holder["BPM_C04-01"] is holder.bpm.get("BPM_C04-01") + assert holder["BPM_C04-01"] is holder.diagnostic.bpm.get("BPM_C04-01") assert holder["UNKNOWN"] is None @@ -47,7 +47,7 @@ def test_new_calls_reflect_registry_additions(holder): assert "EXTRA_BPM" not in previous.names() assert holder.get().names() == previous.names() + ["EXTRA_BPM"] - assert holder["EXTRA_BPM"] is holder.bpm.get("EXTRA_BPM") + assert holder["EXTRA_BPM"] is holder.diagnostic.bpm.get("EXTRA_BPM") def test_patterns_always_return_arrays(holder): @@ -70,7 +70,7 @@ def test_character_classes_in_name_patterns(holder): def test_colons_are_part_of_names(holder): holder.fill_device([BPM("CELL04:BPM01", lattice_names="list(BPM_C04-01)")]) - assert holder["CELL04:BPM01"] is holder.bpm.get("CELL04:BPM01") + assert holder["CELL04:BPM01"] is holder.diagnostic.bpm.get("CELL04:BPM01") assert holder["CELL04:BPM*"].names() == ["CELL04:BPM01"] assert holder["model_name:*"].names() == [] diff --git a/tests/common/test_errors.py b/tests/common/test_errors.py index 38ce79f5..4265a20c 100644 --- a/tests/common/test_errors.py +++ b/tests/common/test_errors.py @@ -43,7 +43,7 @@ def test_tune(install_test_package): assert "Magnet QF1A-C05XX not defined" in str(exc.value) with pytest.raises(PyAMLException) as exc: - m2 = sr.design.bpm.get("QF1A-C05XX") + m2 = sr.design.diagnostic.bpm.get("QF1A-C05XX") assert "BPM QF1A-C05XX not defined" in str(exc.value) diff --git a/tests/diagnostics/test_diagnostic_accessors.py b/tests/diagnostics/test_diagnostic_accessors.py new file mode 100644 index 00000000..2247db71 --- /dev/null +++ b/tests/diagnostics/test_diagnostic_accessors.py @@ -0,0 +1,81 @@ +import pytest + +from pyaml.accelerator import Accelerator +from pyaml.common.exception import PyAMLException +from pyaml.lattice.simulator import Simulator + + +def test_diagnostic_get_returns_named_monitor(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + assert design.diagnostic.get("BETATRON_TUNE") is design.get_betatron_tune_monitor("BETATRON_TUNE") + + +def test_diagnostic_get_with_no_name_returns_all_configured_diagnostics(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + all_diagnostics = design.diagnostic.get() + assert design.get_betatron_tune_monitor("BETATRON_TUNE") in all_diagnostics + + +def test_diagnostic_betatron_tune_returns_default_monitor(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + assert design.diagnostic.betatron_tune is design.get_betatron_tune_monitor("BETATRON_TUNE") + + +def test_diagnostic_raises_when_default_missing(ebs_lattice_file): + holder = Simulator(name="empty", lattice=str(ebs_lattice_file)) + + with pytest.raises(PyAMLException) as exc: + _ = holder.diagnostic.betatron_tune + assert "BETATRON_TUNE" in str(exc.value) + + +def test_diagnostic_raises_when_default_wrong_type(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + # Swap the registered default so it points to an object of the wrong type. + design._DIAG["BETATRON_TUNE"] = design.orbit + + with pytest.raises(PyAMLException) as exc: + _ = design.diagnostic.betatron_tune + assert "BETATRON_TUNE" in str(exc.value) + assert "BetatronTuneMonitor" in str(exc.value) + + +def test_diagnostic_bpm_returns_named_bpm(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + assert design.diagnostic.bpm.get("BPM_C04-04").get_name() == "BPM_C04-04" + + +def test_diagnostic_bpms_returns_named_array(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + bpms = design.diagnostic.bpms.get("BPM") + assert design.diagnostic.bpm.get("BPM_C04-04") in bpms diff --git a/tests/integration/test_dt4acc_twin_smoke.py b/tests/integration/test_dt4acc_twin_smoke.py index 549bdb4e..44220e09 100644 --- a/tests/integration/test_dt4acc_twin_smoke.py +++ b/tests/integration/test_dt4acc_twin_smoke.py @@ -134,7 +134,7 @@ def deactivated_test_orbit_correction(config_key: str): try: accelerator = _build_accelerator(config_key) control_mode = accelerator.live - bpms = control_mode.bpms.get("bpms") + bpms = control_mode.diagnostic.bpms.get("bpms") orbit_response_matrix = control_mode.get_orm_tuning("DEFAULT_ORBIT_RESPONSE_MATRIX") orbit_correction = control_mode.get_orbit_tuning("DEFAULT_ORBIT_CORRECTION") orbit_response_matrix.measure() diff --git a/tests/test_accelerator_load.py b/tests/test_accelerator_load.py index f30a68e3..0bc649ce 100644 --- a/tests/test_accelerator_load.py +++ b/tests/test_accelerator_load.py @@ -19,8 +19,8 @@ def test_peer(): def test_repr_is_informative_and_bounded(): sr = Accelerator.load("tests/config/EBSOrbit.yaml") - bpm = sr.design.bpm.get("BPM_C04-04") - bpms = sr.design.bpms.get("BPM") + bpm = sr.design.diagnostic.bpm.get("BPM_C04-04") + bpms = sr.design.diagnostic.bpms.get("BPM") assert repr(bpm) == ( "BPM(name='BPM_C04-04', lattice_names='BPM_C04-04', " @@ -46,7 +46,7 @@ def test_repr_options_limit_sequences(): set_repr_options(max_items=1) sr = Accelerator.load("tests/config/EBSOrbit.yaml") - bpms = sr.design.bpms.get("BPM") + bpms = sr.design.diagnostic.bpms.get("BPM") assert f"... +{len(bpms) - 1} more ..." in repr(bpms) finally: set_repr_options( diff --git a/tests/test_catalogs.py b/tests/test_catalogs.py index 828627aa..43627e1d 100644 --- a/tests/test_catalogs.py +++ b/tests/test_catalogs.py @@ -59,7 +59,7 @@ def test_inline_catalog_is_supported(install_test_package): } ) - bpm = sr.live.bpm.get("BPM_C02-01") + bpm = sr.live.diagnostic.bpm.get("BPM_C02-01") assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0])) @@ -184,7 +184,7 @@ def test_indexed_catalog_entry_extracts_scalar_from_vector_attribute(install_tes } ) - bpm = sr.live.bpm.get("BPM_TEST") + bpm = sr.live.diagnostic.bpm.get("BPM_TEST") positions = bpm.positions.get() assert np.isclose(positions[0], 1.5) assert np.isclose(positions[1], -0.3) diff --git a/tests/test_load_conf_with_code.py b/tests/test_load_conf_with_code.py index de476102..25a98a8d 100644 --- a/tests/test_load_conf_with_code.py +++ b/tests/test_load_conf_with_code.py @@ -8,11 +8,13 @@ def test_load_conf_with_code(): config_path = parent_folder.joinpath("config", "EBSOrbit.yaml").resolve() sr: Accelerator = Accelerator.load(config_path) - bpms = sr.live.bpms.get("BPM") + bpms = sr.live.diagnostic.bpms.get("BPM") assert bpms is not None assert len(bpms) == 320 assert sr.live[bpms[0].get_name()] is bpms[0] assert sr.live["BPM*"].names() == bpms.names() assert sr.live[:].names() == [element.get_name() for element in sr.live.get_all_elements()] - assert sr.design["BPM*"].names() == sr.design.bpms.get("BPM").names() + assert sr.design["BPM*"].names() == sr.design.diagnostic.bpms.get("BPM").names() + + assert sr.live.diagnostic.bpms.BPM is bpms diff --git a/tests/tuning_tools/test_bba.py b/tests/tuning_tools/test_bba.py index 330c8118..f4247103 100644 --- a/tests/tuning_tools/test_bba.py +++ b/tests/tuning_tools/test_bba.py @@ -9,7 +9,7 @@ def test_bba(): SR = sr.design # Add a misalignement - SR.bpm.get("BPM_C04-04").offset.set([20e-6, -15e-6]) + SR.diagnostic.bpm.get("BPM_C04-04").offset.set([20e-6, -15e-6]) # BBA (standard bow tie, model independant) bba = SR.get_bba("BBA-BPM_C04-04") diff --git a/tests/tuning_tools/test_tool_accessors.py b/tests/tuning_tools/test_tool_accessors.py index 1ccbf5df..d9a54e7b 100644 --- a/tests/tuning_tools/test_tool_accessors.py +++ b/tests/tuning_tools/test_tool_accessors.py @@ -1,4 +1,8 @@ +import pytest + from pyaml.accelerator import Accelerator +from pyaml.common.exception import PyAMLException +from pyaml.lattice.simulator import Simulator def test_tuning_tools_expose_configured_elements(): @@ -23,7 +27,7 @@ def test_tuning_tools_expose_configured_elements(): assert design.crm.chromaticity_monitor is chromaticity_monitor assert design.crm.sextupoles is sextupoles - bpms = design.bpms.get("BPM") + bpms = design.diagnostic.bpms.get("BPM") rf_plant = design.rf.get("DEFAULT_RF_PLANT") assert chromaticity_monitor.tune_monitor is tune_monitor assert chromaticity_monitor.bpms is bpms @@ -45,9 +49,77 @@ def test_tuning_tools_expose_configured_elements(): for name in ("BBA-BPM_C04-04", "BBA2-BPM_C04-04"): bba = design.get_bba(name) assert bba.bpms is bpms - assert bba.bpm is design.bpm.get("BPM_C04-04") + assert bba.bpm is design.diagnostic.bpm.get("BPM_C04-04") assert bba.hcorrector is design.magnet.get("SF2E-C02-H") assert bba.vcorrector is design.magnet.get("SD1A-C26-V") assert bba.quadrupole is design.magnet.get("QF6B-C04") assert design.get_bba("BBA2-BPM_C04-04").tune_correction is design.tune + + +def test_tool_get_returns_named_tool(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + assert design.tool.get("DEFAULT_TUNE_CORRECTION") is design.tune + assert design.tool.get("DEFAULT_ORBIT_RESPONSE_MATRIX") is design.orm + + +def test_tool_get_with_no_name_returns_all_configured_tools(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + all_tools = design.tool.get() + assert design.tune in all_tools + assert design.trm in all_tools + assert design.orbit in all_tools + assert design.orm in all_tools + assert design.chromaticity in all_tools + assert design.crm in all_tools + assert design.dispersion in all_tools + + +def test_tool_typed_properties_alias_existing_defaults(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + assert design.tool.tune is design.tune + assert design.tool.trm is design.trm + assert design.tool.orbit is design.orbit + assert design.tool.orm is design.orm + assert design.tool.chromaticity is design.chromaticity + assert design.tool.crm is design.crm + assert design.tool.dispersion is design.dispersion + + +def test_tool_raises_when_default_missing(ebs_lattice_file): + holder = Simulator(name="empty", lattice=str(ebs_lattice_file)) + + with pytest.raises(PyAMLException) as exc: + _ = holder.tool.tune + assert "DEFAULT_TUNE_CORRECTION" in str(exc.value) + + +def test_tool_raises_when_default_wrong_type(): + design = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ).design + + # Swap the registered default so it points to an object of the wrong type. + design._TOOLS["DEFAULT_TUNE_CORRECTION"] = design.orbit + + with pytest.raises(PyAMLException) as exc: + _ = design.tool.tune + assert "DEFAULT_TUNE_CORRECTION" in str(exc.value) + assert "Tune" in str(exc.value) diff --git a/tests/tuning_tools/test_tuning_dispersion.py b/tests/tuning_tools/test_tuning_dispersion.py index 16d2cd39..9cd3e66b 100644 --- a/tests/tuning_tools/test_tuning_dispersion.py +++ b/tests/tuning_tools/test_tuning_dispersion.py @@ -17,7 +17,7 @@ def test_tuning_orm(): dispersion.measure() dispersion_data = dispersion.get() - bpms = element_holder.bpms.get("BPM") + bpms = element_holder.diagnostic.bpms.get("BPM") assert len(dispersion_data["frequency_response_x"]) == len(bpms) assert len(dispersion_data["frequency_response_y"]) == len(bpms) diff --git a/tests/tuning_tools/test_tuning_orbit_correction.py b/tests/tuning_tools/test_tuning_orbit_correction.py index 4bf3b6a6..7cf8fb56 100644 --- a/tests/tuning_tools/test_tuning_orbit_correction.py +++ b/tests/tuning_tools/test_tuning_orbit_correction.py @@ -22,7 +22,7 @@ def test_tuning_orbit_correction(): std_kick = 1e-6 hcorr = element_holder.magnets.get("HCorr") vcorr = element_holder.magnets.get("VCorr") - bpms = element_holder.bpms.get("BPM") + bpms = element_holder.diagnostic.bpms.get("BPM") x, y = bpms.positions.get().T # get reference orbit reference = np.concat((x, y)) diff --git a/tests/tuning_tools/test_tuning_orm.py b/tests/tuning_tools/test_tuning_orm.py index 553bce16..f3652f0c 100644 --- a/tests/tuning_tools/test_tuning_orm.py +++ b/tests/tuning_tools/test_tuning_orm.py @@ -17,7 +17,7 @@ def test_tuning_orm(): orm = element_holder.orm - bpms = element_holder.bpms.get("BPM") + bpms = element_holder.diagnostic.bpms.get("BPM") hcorr_names = element_holder.magnets.get("HCorr").names()[:4] vcorr_names = element_holder.magnets.get("VCorr").names()[:4] orm.measure(corrector_names=hcorr_names + vcorr_names)