diff --git a/examples/other_examples/OphydDevice_examples/tune_ophyd_device.py b/examples/other_examples/OphydDevice_examples/tune_ophyd_device.py index f98c19225..bb22a5884 100644 --- a/examples/other_examples/OphydDevice_examples/tune_ophyd_device.py +++ b/examples/other_examples/OphydDevice_examples/tune_ophyd_device.py @@ -179,6 +179,6 @@ def unit(self) -> str: # Main code sr = Accelerator.from_dict(acc_config) -tm1 = sr.live.get_betatron_tune_monitor("MY_TUNE_MONITOR") +tm1 = sr.live.diagnostic.get("MY_TUNE_MONITOR") print(tm1.tune.get()) print(tm1.frequency.get()) diff --git a/examples/use_cases/01-tune_correction.ipynb b/examples/use_cases/01-tune_correction.ipynb index cf6e3471b..4e20ff1f9 100644 --- a/examples/use_cases/01-tune_correction.ipynb +++ b/examples/use_cases/01-tune_correction.ipynb @@ -207,7 +207,7 @@ } ], "source": [ - "tune_monitor = SR.get_betatron_tune_monitor(\"BETATRON_TUNE\")\n", + "tune_monitor = SR.diagnostic.get(\"BETATRON_TUNE\")\n", "print(f\"Current tune: {tune_monitor.tune.get()}\")\n", "tune_monitor # string representation" ] diff --git a/examples/use_cases/01-tune_correction.py b/examples/use_cases/01-tune_correction.py index fc84e9293..1525b993a 100644 --- a/examples/use_cases/01-tune_correction.py +++ b/examples/use_cases/01-tune_correction.py @@ -80,7 +80,7 @@ # In[4]: -tune_monitor = SR.get_betatron_tune_monitor("BETATRON_TUNE") +tune_monitor = SR.diagnostic.get("BETATRON_TUNE") print(f"Current tune: {tune_monitor.tune.get()}") print(tune_monitor) # string representation diff --git a/examples/use_cases/02-chromaticity_measurement.ipynb b/examples/use_cases/02-chromaticity_measurement.ipynb index 9ac1afca7..05ea57a2c 100644 --- a/examples/use_cases/02-chromaticity_measurement.ipynb +++ b/examples/use_cases/02-chromaticity_measurement.ipynb @@ -267,7 +267,7 @@ } ], "source": [ - "chroma_monitor = SR.get_chromaticity_monitor(\"CHROMATICITY_MONITOR\")\n", + "chroma_monitor = SR.tool.get(\"CHROMATICITY_MONITOR\")\n", "\n", "chroma_monitor.measure(\n", " callback=chroma_callback,\n", diff --git a/examples/use_cases/02-chromaticity_measurement.py b/examples/use_cases/02-chromaticity_measurement.py index 966b23ef3..320aa0b94 100644 --- a/examples/use_cases/02-chromaticity_measurement.py +++ b/examples/use_cases/02-chromaticity_measurement.py @@ -109,7 +109,7 @@ def chroma_callback(action: int, cb_data: dict): # In[17]: -chroma_monitor = SR.get_chromaticity_monitor("CHROMATICITY_MONITOR") +chroma_monitor = SR.tool.get("CHROMATICITY_MONITOR") chroma_monitor.measure( callback=chroma_callback, diff --git a/pyaml/arrays/element.py b/pyaml/arrays/element.py index 6483f03de..4b93276a0 100644 --- a/pyaml/arrays/element.py +++ b/pyaml/arrays/element.py @@ -65,7 +65,7 @@ def fill_array(self, holder: ElementHolder): .. code-block:: python >>> elt_cfg.fill_array(sr.design) - >>> names = sr.design.get_elements("MyArray").names() + >>> names = sr.design.get("MyArray").names() >>> print(names) ['BPM_C04-01', 'SH1A-C04-H'] diff --git a/pyaml/arrays/element_array.py b/pyaml/arrays/element_array.py index 26c7b65b5..797fc9197 100644 --- a/pyaml/arrays/element_array.py +++ b/pyaml/arrays/element_array.py @@ -61,7 +61,7 @@ class ElementArray(list[Element]): .. code-block:: python >>> sr = Accelerator.load("acc.yaml") - >>> elements = sr.design.get_elements("QuadForTune") + >>> elements = sr.design.get("QuadForTune") """ def __init__(self, array_name: str, elements: list[Element], use_aggregator=True): @@ -282,7 +282,7 @@ def __and__(self, other: object): .. code-block:: python - >>> cell1 = sr.live.get_elements("C01") + >>> cell1 = sr.live.get("C01") >>> sexts = sr.live.get_magnets("SEXT") >>> cell1_sext = cell1 & sexts diff --git a/pyaml/common/holders/diagnostic_holder.py b/pyaml/common/holders/diagnostic_holder.py index b78eefbf2..fa156c684 100644 --- a/pyaml/common/holders/diagnostic_holder.py +++ b/pyaml/common/holders/diagnostic_holder.py @@ -35,13 +35,6 @@ class DiagnosticHolder: 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 @@ -144,7 +137,7 @@ def betatron_tune(self) -> BetatronTuneMonitor: :class:`~pyaml.diagnostics.tune_monitor.BetatronTuneMonitor`. """ name = "BETATRON_TUNE" - obj = self._peer.get_betatron_tune_monitor(name) + obj = self.get(name) if not isinstance(obj, BetatronTuneMonitor): raise PyAMLException(f"{name}: BetatronTuneMonitor expected but got {type(obj).__name__}") return obj diff --git a/pyaml/common/holders/element_holder.py b/pyaml/common/holders/element_holder.py index 21d303ac2..f5880dd2b 100644 --- a/pyaml/common/holders/element_holder.py +++ b/pyaml/common/holders/element_holder.py @@ -11,7 +11,6 @@ from ...magnet.serialized_magnet import SerializedMagnets from ...rf.rf_plant import RFPlant from ...rf.rf_transmitter import RFTransmitter -from ...tuning_tools.chromaticity_monitor import ChromaticityMonitor from ..abstract_aggregator import ScalarAggregator from ..element import Element from ..exception import PyAMLException @@ -31,15 +30,7 @@ if TYPE_CHECKING: from ...accelerator import Accelerator from ...configuration.unbound_element import UnboundElement - from ...tuning_tools.bba import BBA - from ...tuning_tools.chromaticity import Chromaticity - from ...tuning_tools.chromaticity_response_matrix import ChromaticityResponseMatrix - from ...tuning_tools.dispersion import Dispersion from ...tuning_tools.measurement_tool import MeasurementTool - 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 ...tuning_tools.tuning_tool import TuningTool @@ -85,36 +76,12 @@ class ElementHolder(metaclass=ABCMeta): Create and register a generic element array. add_element(element) Add an element to the global element store. - get_element(name) - Return a named element from the global store. - get_elements(name) - Return a named generic element array. - get_all_elements() - Return all registered elements in insertion order. - get_betatron_tune_monitor(name) - Return a named betatron tune monitor. + get(name=None) + Return a named element array, or every registered element when no name is given. add_betatron_tune_monitor(tune_monitor) Add a betatron tune monitor to the diagnostics store. add_tool(tool) Add a tuning or measurement tool to the tool store. - get_chromaticity_monitor(name) - Return a named chromaticity monitor. - get_chromaticity_tuning(name) - Return a named chromaticity tuning tool. - get_crm_tuning(name) - Return a named chromaticity response-matrix tool. - get_tune_tuning(name) - Return a named tune correction tool. - get_trm_tuning(name) - Return a named tune response-matrix tool. - get_orbit_tuning(name) - Return a named orbit correction tool. - get_orm_tuning(name) - Return a named orbit response-matrix tool. - get_bba(name) - Return a named beam-based alignment tool. - get_dispersion_tuning(name) - Return a named dispersion tuning tool. """ def __init__(self): @@ -211,7 +178,7 @@ def tool(self) -> ToolHolder: def post_init(self): """Run post-initialization hooks for every stored element.""" - for e in self.get_all_elements(): + for e in self._ALL.values(): e.post_init() def fill_device(self, elements: list[Element]): @@ -418,13 +385,27 @@ def _get(self, what, name, array) -> Element: return array[name] # Generic elements - def get(self) -> ElementArray: - """Return all registered elements in insertion order. + def get(self, name: str | None = None) -> ElementArray: + """Return a named element array, or every registered element when no name is given. + + Parameters + ---------- + name : str, optional + Name of the element array to look up, as declared in the configuration. + When omitted, every registered element is returned instead. Returns ------- ElementArray - New unnamed container sharing the registered element references. + The element array registered under ``name``, regardless of its concrete + family (magnet, BPM, combined-function magnet, serialized-magnet, or + generic element array), or a new unnamed container of every registered + element, in insertion order, when ``name`` is omitted. + + Raises + ------ + PyAMLException + If ``name`` is given and no array is registered under it. Notes ----- @@ -436,8 +417,11 @@ def get(self) -> ElementArray: -------- >>> elements = sr.live.get() >>> names = elements.names() + >>> cell08 = sr.live.get("CELL08") """ - return ElementArray("", list(self._ALL.values())) + if name is None: + return ElementArray("", list(self._ALL.values())) + return self._get_array(name) @overload def __getitem__(self, key: int) -> Element: ... @@ -545,7 +529,7 @@ def fill_element_array(self, arrayName: str, elementNames: list[str]): self._fill_array( arrayName, elementNames, - self.get_element, + self._get_element, ElementArray, self._ELEMENT_ARRAYS, ) @@ -561,60 +545,14 @@ def add_element(self, element: Element): """ self._ALL[element.get_name()] = element - def get_element(self, name: str) -> Element: + def _get_element(self, name: str) -> Element: """ - Return a named element from the global store. - - Parameters - ---------- - name : str - Name of the element to look up, as declared in the configuration. - - Returns - ------- - Element - The element registered under ``name``. + Generic single-element resolver used internally to build element arrays. """ return self._get("Element", name, self._ALL) - def get_elements(self, name: str) -> ElementArray: - """ - Return a named generic element array. - - Parameters - ---------- - name : str - Name of the element array to look up, as declared in the configuration. - - Returns - ------- - ElementArray - The element array registered under ``name``. - """ - return self._get("Element array", name, self._ELEMENT_ARRAYS) - - def get_all_elements(self) -> list[Element]: - """Return all registered elements in insertion order.""" - return [value for key, value in self._ALL.items()] - # Tune monitor - def get_betatron_tune_monitor(self, name: str) -> BetatronTuneMonitor: - """ - Return a named betatron tune monitor. - - Parameters - ---------- - name : str - Name of the betatron tune monitor to look up, as declared in the configuration. - - Returns - ------- - BetatronTuneMonitor - The betatron tune monitor registered under ``name``. - """ - return self._get("Diagnostic", name, self._DIAG) - def add_betatron_tune_monitor(self, tune_monitor: Element): """ Add a betatron tune monitor to the diagnostics store. @@ -639,161 +577,6 @@ def add_tool(self, tool: Element): """ self._add(self._TOOLS, tool) - # ---- Chromaticity ------------------------------------------------- - - def get_chromaticity_monitor(self, name: str) -> ChromaticityMonitor: - """ - Return a named chromaticity monitor. - - Parameters - ---------- - name : str - Name of the chromaticity monitor to look up, as declared in the configuration. - - Returns - ------- - ChromaticityMonitor - The chromaticity monitor registered under ``name``. - """ - obj = self._get("Chromaticity monitor", name, self._TOOLS) - return obj - - def get_chromaticity_tuning(self, name: str) -> "Chromaticity": - """ - Return a named chromaticity tuning tool. - - Parameters - ---------- - name : str - Name of the chromaticity tuning tool to look up, as declared in the configuration. - - Returns - ------- - 'Chromaticity' - The chromaticity tuning tool registered under ``name``. - """ - return self._get("Chromaticity tool", name, self._TOOLS) - - def get_crm_tuning(self, name: str) -> "ChromaticityResponseMatrix": - """ - Return a named chromaticity response-matrix tool. - - Parameters - ---------- - name : str - Name of the chromaticity response-matrix tool to look up, as declared in the configuration. - - Returns - ------- - 'ChromaticityResponseMatrix' - The chromaticity response-matrix tool registered under ``name``. - """ - return self._get("ChromaticityResponseMatrix tool", name, self._TOOLS) - - # ---- Tune --------------------------------------------------------- - - def get_tune_tuning(self, name: str) -> "Tune": - """ - Return a named tune correction tool. - - Parameters - ---------- - name : str - Name of the tune correction tool to look up, as declared in the configuration. - - Returns - ------- - 'Tune' - The tune correction tool registered under ``name``. - """ - return self._get("Tune tuning tool", name, self._TOOLS) - - def get_trm_tuning(self, name: str) -> "TuneResponseMatrix": - """ - Return a named tune response-matrix tool. - - Parameters - ---------- - name : str - Name of the tune response-matrix tool to look up, as declared in the configuration. - - Returns - ------- - 'TuneResponseMatrix' - The tune response-matrix tool registered under ``name``. - """ - return self._get("TuneResponseMatrix tool", name, self._TOOLS) - - # ---- Orbit -------------------------------------------------------- - - def get_orbit_tuning(self, name: str) -> "Orbit": - """ - Return a named orbit correction tool. - - Parameters - ---------- - name : str - Name of the orbit correction tool to look up, as declared in the configuration. - - Returns - ------- - 'Orbit' - The orbit correction tool registered under ``name``. - """ - return self._get("Orbit tuning tool", name, self._TOOLS) - - def get_orm_tuning(self, name: str) -> "OrbitResponseMatrix": - """ - Return a named orbit response-matrix tool. - - Parameters - ---------- - name : str - Name of the orbit response-matrix tool to look up, as declared in the configuration. - - Returns - ------- - 'OrbitResponseMatrix' - The orbit response-matrix tool registered under ``name``. - """ - return self._get("OrbitResponseMatrix tool", name, self._TOOLS) - - # ---- BBA -------------------------------------------------------- - - def get_bba(self, name: str) -> "BBA": - """ - Return a named beam-based alignment tool. - - Parameters - ---------- - name : str - Name of the beam-based alignment tool to look up, as declared in the configuration. - - Returns - ------- - 'BBA' - The beam-based alignment tool registered under ``name``. - """ - return self._get("BBA tool", name, self._TOOLS) - - # ---- Dispersive orbit -------------------------------------------- - - def get_dispersion_tuning(self, name: str) -> "Dispersion": - """ - Return a named dispersion tuning tool. - - Parameters - ---------- - name : str - Name of the dispersion tuning tool to look up, as declared in the configuration. - - Returns - ------- - 'Dispersion' - The dispersion tuning tool registered under ``name``. - """ - return self._get("Dispersion tool", name, self._TOOLS) - def _get_array(self, name: str): """ Generic array resolver used by YellowPages. @@ -864,7 +647,7 @@ def _set_energy(self, E: float): Energy in eV """ # Needed by energy dependant element (i.e. magnet coil current calculation) - for m in self.get_all_elements(): + for m in self._ALL.values(): m.set_energy(E) def _set_mcf(self, alphac: float): @@ -877,7 +660,7 @@ def _set_mcf(self, alphac: float): Moment compaction factor """ # Needed by some off energy dependant element (i.e. chromaticty tools) - for m in self.get_all_elements(): + for m in self._ALL.values(): m.set_mcf(alphac) def _set_harmonic(self, h: int): @@ -889,5 +672,5 @@ def _set_harmonic(self, h: int): h : int Harmonic number """ - for m in self.get_all_elements(): + for m in self._ALL.values(): m.set_harmonic(h) diff --git a/pyaml/tuning_tools/bba.py b/pyaml/tuning_tools/bba.py index d5cc32df2..05eb9b8e2 100644 --- a/pyaml/tuning_tools/bba.py +++ b/pyaml/tuning_tools/bba.py @@ -206,7 +206,7 @@ def measure( sr = Accelerator.load("tests/config/EBSOrbit.yaml") SR = sr.design - bba = SR.get_bba("BBA-BPM_C04-04") + bba = SR.tool.get("BBA-BPM_C04-04") # Add a misalignement SR.get_bpm("BPM_C04-04").offset.set([20e-6,-15e-6]) diff --git a/pyaml/tuning_tools/bba2.py b/pyaml/tuning_tools/bba2.py index 055446875..1e34d80d3 100644 --- a/pyaml/tuning_tools/bba2.py +++ b/pyaml/tuning_tools/bba2.py @@ -274,7 +274,7 @@ def quadrupole(self) -> "Magnet": def tune_correction(self) -> "Tune": """Return the tune-correction tool used during alignment.""" self.check_peer() - return self.peer.get_tune_tuning(self.tune_correction_name) + return self.peer.tool.get(self.tune_correction_name) @staticmethod def _x_intercept(x, k, n): @@ -360,8 +360,8 @@ def _init_responses( sth = design.magnet.get(steererhname) stv = design.magnet.get(steerervname) orbit = design.diagnostic.bpms.get(bpmname).positions - tune_design = design.get_tune_tuning(tunename) - tune_live = self._peer.get_tune_tuning(tunename) + tune_design = design.tool.get(tunename) + tune_live = self._peer.tool.get(tunename) # Get tune from live and adjust the model to improve quad response phase tune0 = tune_design.readback() @@ -585,7 +585,7 @@ def measure( sr = Accelerator.load("tests/config/EBSOrbit.yaml") SR = sr.live - bba = SR.get_bba("BBA2-BPM_C04-04") + bba = SR.tool.get("BBA2-BPM_C04-04") # Add a misalignement SR.get_bpm("BPM_C04-04").offset.set([200e-6,-150e-6]) diff --git a/pyaml/tuning_tools/chromaticity.py b/pyaml/tuning_tools/chromaticity.py index 94d475eea..d43c46726 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/pyaml/tuning_tools/chromaticity.py @@ -135,7 +135,7 @@ def load(self, load_path: Path): def chromaticity_monitor(self) -> "ChromaticityMonitor": """Return the chromaticity monitor used for readback.""" self.check_peer() - return self.peer.get_chromaticity_monitor(self._chromaticity_monitor_name) + return self.peer.tool.get(self._chromaticity_monitor_name) @property def sextupoles(self) -> "MagnetArray": diff --git a/pyaml/tuning_tools/chromaticity_monitor.py b/pyaml/tuning_tools/chromaticity_monitor.py index f9bfa5257..c03e284f0 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/pyaml/tuning_tools/chromaticity_monitor.py @@ -201,7 +201,7 @@ def __init__( def tune_monitor(self) -> "BetatronTuneMonitor": """Return the betatron tune monitor used for measurements.""" self.check_peer() - return self.peer.get_betatron_tune_monitor(self.betatron_tune_name) + return self.peer.diagnostic.get(self.betatron_tune_name) @property def rf_plant(self) -> "RFPlant": diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index 4d739e0e3..23d712ed4 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -115,7 +115,7 @@ def sextupoles(self) -> "MagnetArray": def chromaticity_monitor(self) -> "ChromaticityMonitor": """Return the chromaticity monitor used for the measurement.""" self.check_peer() - return self.peer.get_chromaticity_monitor(self.chromaticity_name) + return self.peer.tool.get(self.chromaticity_name) def measure( self, diff --git a/pyaml/tuning_tools/tune.py b/pyaml/tuning_tools/tune.py index 3c336d31f..18811bba3 100644 --- a/pyaml/tuning_tools/tune.py +++ b/pyaml/tuning_tools/tune.py @@ -153,7 +153,7 @@ def response_matrix(self) -> ResponseMatrixData | None: def tune_monitor(self) -> "BetatronTuneMonitor": """Return the betatron tune monitor used by the correction.""" self.check_peer() - return self.peer.get_betatron_tune_monitor(self.betatron_tune_name) + return self.peer.diagnostic.get(self.betatron_tune_name) @property def quadrupoles(self) -> "MagnetArray": diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 380a4105f..fe65e0d14 100644 --- a/pyaml/tuning_tools/tune_response_matrix.py +++ b/pyaml/tuning_tools/tune_response_matrix.py @@ -154,7 +154,7 @@ def quadrupoles(self) -> "MagnetArray": def tune_monitor(self) -> "BetatronTuneMonitor": """Return the betatron tune monitor used for the measurement.""" self.check_peer() - return self.peer.get_betatron_tune_monitor(self.betatron_tune_name) + return self.peer.diagnostic.get(self.betatron_tune_name) def measure( self, diff --git a/tests/arrays/test_array_selection_types.py b/tests/arrays/test_array_selection_types.py index 11e651815..8dce804ef 100644 --- a/tests/arrays/test_array_selection_types.py +++ b/tests/arrays/test_array_selection_types.py @@ -44,7 +44,7 @@ def test_magnet_holder_selection_uses_the_same_typing(design): def test_mixed_array_stays_generic_until_only_magnets_are_selected(design): - mixed = design.get_elements("ElArray") + mixed = design.get("ElArray") assert type(mixed[:]) is ElementArray assert mixed[:].names() == mixed.names() diff --git a/tests/arrays/test_arrays.py b/tests/arrays/test_arrays.py index 62f9c2024..5bd02a3d6 100644 --- a/tests/arrays/test_arrays.py +++ b/tests/arrays/test_arrays.py @@ -164,10 +164,10 @@ def test_arrays(install_test_package): assert np.abs(pos[1][1] - 7.4265634524358045e-06) < 1e-10 # Radom array - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") # Create an array that contains all elements - allElts = ElementArray("AllElements", sr.design.get_all_elements()) + allElts = ElementArray("AllElements", sr.design.get()) assert len(allElts) == 11 # Create an array that contains all elements @@ -210,7 +210,7 @@ def test_arrays(install_test_package): # Test dynamic arrays sr: Accelerator = Accelerator.load("tests/config/EBSOrbit.yaml", include_locations=False) - ae = ElementArray("All", sr.design.get_all_elements()) + ae = ElementArray("All", sr.design.get()) acfm = ElementArray("AllCFM", sr.design.combined_function_magnet.all(), use_aggregator=False) bpmC5 = ae["BPM*"][10:20] # All BPM C5 diff --git a/tests/arrays/test_arrays_ops.py b/tests/arrays/test_arrays_ops.py index 5fd706e2c..fd8c961fc 100644 --- a/tests/arrays/test_arrays_ops.py +++ b/tests/arrays/test_arrays_ops.py @@ -43,7 +43,7 @@ def test_element_array_and_mask_filters_and_is_autotyped_list_mask( sr.design.get_lattice().disable_6d() # "ElArray" is a mixed ElementArray in the dummy config (see existing tests) - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") assert isinstance(elts, ElementArray) assert len(elts) > 0 @@ -71,7 +71,7 @@ def test_element_array_and_mask_filters_and_is_autotyped_numpy_mask( sr = accelerator_from_fragments(*sr_configuration_fragments) sr.design.get_lattice().disable_6d() - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") assert len(elts) > 0 mask_list = [isinstance(e, Magnet) for e in elts] @@ -97,7 +97,7 @@ def test_element_array_sub_mask_removes_true_inverse_of_and( sr = accelerator_from_fragments(*sr_configuration_fragments) sr.design.get_lattice().disable_6d() - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") assert len(elts) > 0 # Keep only magnets with '& mask' @@ -134,7 +134,7 @@ def test_element_array_mask_length_mismatch_raises_for_and_and_sub( sr = accelerator_from_fragments(*sr_configuration_fragments) sr.design.get_lattice().disable_6d() - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") assert len(elts) > 0 bad_mask = [True] * (len(elts) - 1) @@ -159,7 +159,7 @@ def test_mask_by_type_returns_correct_boolean_mask( sr = accelerator_from_fragments(*sr_configuration_fragments) sr.design.get_lattice().disable_6d() - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") mask = elts.mask_by_type(Magnet) assert isinstance(mask, list) @@ -182,7 +182,7 @@ def test_filter_by_type_returns_autotyped_array(install_test_package, accelerato sr = accelerator_from_fragments(*sr_configuration_fragments) sr.design.get_lattice().disable_6d() - elts = sr.design.get_elements("ElArray") + elts = sr.design.get("ElArray") filtered = elts.of_type(Magnet) if len(filtered) == 0: diff --git a/tests/common/test_element_holder_collection.py b/tests/common/test_element_holder_collection.py index 3fdfc88d6..a90a15235 100644 --- a/tests/common/test_element_holder_collection.py +++ b/tests/common/test_element_holder_collection.py @@ -22,7 +22,7 @@ def test_exact_name_returns_the_element_or_raises(holder): def test_get_and_full_slice_keep_registration_order(holder): - names = [element.get_name() for element in holder.get_all_elements()] + names = [element.get_name() for element in holder.get()] assert type(holder.get()) is ElementArray assert type(holder[:]) is ElementArray @@ -79,7 +79,7 @@ def test_colons_are_part_of_names(holder): def test_magnet_subclasses_share_a_typed_array_in_either_order(holder): horizontal = holder.magnet.get("SH1A-C01-H") vertical = holder.magnet.get("SH1A-C01-V") - start = holder.get_all_elements().index(horizontal) + start = holder.get().index(horizontal) assert type(horizontal) is not type(vertical) assert isinstance(holder["SH1A-C01-[HV]"], MagnetArray) @@ -97,7 +97,7 @@ def test_mixed_selection_returns_a_generic_array(holder): def test_indices_and_slices_follow_insertion_order(holder): - registered = holder.get_all_elements() + registered = holder.get() assert holder[0] is registered[0] assert holder[-1] is registered[-1] @@ -119,7 +119,7 @@ def test_empty_holder_returns_empty_collections(ebs_lattice_file): def test_invalid_indices_and_keys_raise_clear_errors(holder): - size = len(holder.get_all_elements()) + size = len(holder.get()) with pytest.raises(IndexError): holder[size] @@ -132,14 +132,21 @@ def test_invalid_indices_and_keys_raise_clear_errors(holder): def test_selection_intersects_with_a_configured_family(holder): - selected = holder["SH1A-C0?-H"] & holder.get_elements("ElArray") + selected = holder["SH1A-C0?-H"] & holder.get("ElArray") assert isinstance(selected, MagnetArray) assert selected.names() == ["SH1A-C02-H"] - assert holder.get_element("SH1A-C02-H") is holder["SH1A-C02-H"] - assert holder.get_all_elements() == list(holder.get()) + + +def test_get_resolves_a_named_array_regardless_of_its_concrete_family(holder): + """`get(name)` must not be limited to generic element arrays, unlike the old `get_elements`.""" + assert isinstance(holder.get("HCORR"), MagnetArray) + assert holder.get("HCORR").names() == ["SH1A-C01-H", "SH1A-C02-H"] + assert isinstance(holder.get("BPMS"), BPMArray) + assert holder.get("BPMS").names() == ["BPM_C04-01", "BPM_C04-02"] + assert holder.get("ElArray").names() == ["BPM_C04-01", "BPM_C04-02", "SH1A-C01-V", "SH1A-C02-H"] with pytest.raises(PyAMLException): - holder.get_element("UNKNOWN") + holder.get("UNKNOWN") def test_existing_array_field_filters_still_work(holder): @@ -204,7 +211,7 @@ def test_disjoint_regex_and_list_selections_produce_a_plain_list(holder): def test_selection_combined_with_a_configured_family_via_regex_and_list(holder): """`re:` and list-of-patterns selections intersect with a configured family, same as wildcards.""" - el_array = holder.get_elements("ElArray") # BPM_C04-01, BPM_C04-02, SH1A-C01-V, SH1A-C02-H + el_array = holder.get("ElArray") # BPM_C04-01, BPM_C04-02, SH1A-C01-V, SH1A-C02-H by_regex = holder["re:^SH1A-C0[12]-H$"] & el_array assert isinstance(by_regex, MagnetArray) @@ -272,7 +279,7 @@ def test_getitem_supports_an_exclusion_in_a_list(holder): def test_getitem_lone_exclusion_pattern_means_everything_except(holder): - all_names = holder.get_all_elements() + all_names = holder.get() selected = holder["~QF1A-C01"] assert "QF1A-C01" not in selected.names() diff --git a/tests/diagnostics/test_diagnostic_accessors.py b/tests/diagnostics/test_diagnostic_accessors.py index b4c6f58b1..0046d975e 100644 --- a/tests/diagnostics/test_diagnostic_accessors.py +++ b/tests/diagnostics/test_diagnostic_accessors.py @@ -14,7 +14,7 @@ def test_diagnostic_get_returns_named_monitor(): include_locations=False, ).design - assert design.diagnostic.get("BETATRON_TUNE") is design.get_betatron_tune_monitor("BETATRON_TUNE") + assert design.diagnostic.get("BETATRON_TUNE") is design._DIAG["BETATRON_TUNE"] def test_diagnostic_get_with_no_name_returns_all_configured_diagnostics(): @@ -25,7 +25,7 @@ def test_diagnostic_get_with_no_name_returns_all_configured_diagnostics(): ).design all_diagnostics = design.diagnostic.get() - assert design.get_betatron_tune_monitor("BETATRON_TUNE") in all_diagnostics + assert design.diagnostic.get("BETATRON_TUNE") in all_diagnostics def test_diagnostic_betatron_tune_returns_default_monitor(): @@ -35,7 +35,7 @@ def test_diagnostic_betatron_tune_returns_default_monitor(): include_locations=False, ).design - assert design.diagnostic.betatron_tune is design.get_betatron_tune_monitor("BETATRON_TUNE") + assert design.diagnostic.betatron_tune is design.diagnostic.get("BETATRON_TUNE") def test_diagnostic_raises_when_default_missing(ebs_lattice_file): @@ -128,7 +128,7 @@ def test_diagnostic_getitem_list_selection_ignores_request_order(): ignore_external=True, include_locations=False, ).design - spare_monitor = copy.copy(design.get_betatron_tune_monitor("BETATRON_TUNE")) + spare_monitor = copy.copy(design.diagnostic.get("BETATRON_TUNE")) spare_monitor._name = "SPARE_BETATRON_TUNE" design._DIAG["SPARE_BETATRON_TUNE"] = spare_monitor diff --git a/tests/integration/test_dt4acc_twin_smoke.py b/tests/integration/test_dt4acc_twin_smoke.py index 44220e096..1e8916279 100644 --- a/tests/integration/test_dt4acc_twin_smoke.py +++ b/tests/integration/test_dt4acc_twin_smoke.py @@ -135,8 +135,8 @@ def deactivated_test_orbit_correction(config_key: str): accelerator = _build_accelerator(config_key) control_mode = accelerator.live 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 = control_mode.tool.get("DEFAULT_ORBIT_RESPONSE_MATRIX") + orbit_correction = control_mode.tool.get("DEFAULT_ORBIT_CORRECTION") orbit_response_matrix.measure() ormdata = orbit_response_matrix.get() orbit_response_matrix.save("orm.json") @@ -179,7 +179,7 @@ def deactivated_test_chromaticity_measurement(config_key: str): accelerator = _build_accelerator(config_key) control_mode = accelerator.live - chromaticity_measurement = control_mode.get_chromaticity_monitor("DEFAULT_CHROMATICITY_MEASUREMENT") + chromaticity_measurement = control_mode.tool.get("DEFAULT_CHROMATICITY_MEASUREMENT") def chroma_callback(action: int, cb_data: dict): if action == Action.MEASURE: diff --git a/tests/magnet/test_serialized_magnets.py b/tests/magnet/test_serialized_magnets.py index 1e95ca706..b362f7b77 100644 --- a/tests/magnet/test_serialized_magnets.py +++ b/tests/magnet/test_serialized_magnets.py @@ -27,11 +27,11 @@ def test_config_load(sr_file): sr: Accelerator = Accelerator.load(sr_file, include_locations=False, ignore_external=True) assert sr is not None magnets = [ - sr.design.get_element("QF8B-C04"), - sr.design.get_element("QF8B-C04"), - sr.design.get_element("QD5D-C04"), - sr.design.get_element("QF6D-C04"), - sr.design.get_element("QF4D-C04"), + sr.design["QF8B-C04"], + sr.design["QF8B-C04"], + sr.design["QD5D-C04"], + sr.design["QF6D-C04"], + sr.design["QF4D-C04"], ] assert None not in [magnets] @@ -115,7 +115,7 @@ def test_tune(sr_file): assert len(m.get_magnets()) == m.get_nb_magnets() quadForTuneDesign = sr.design.serialized_magnets.get("QForTune") - tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") + tune_monitor = sr.design.diagnostic.get("BETATRON_TUNE") # Build tune response matrix tunemat = np.zeros((len(quadForTuneDesign), 2)) diff --git a/tests/test_accelerator_load.py b/tests/test_accelerator_load.py index 0bc649ce1..aa5df3efc 100644 --- a/tests/test_accelerator_load.py +++ b/tests/test_accelerator_load.py @@ -9,10 +9,10 @@ def test_peer(): sr = Accelerator.load("tests/config/tune_monitor.yaml") - tm = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") + tm = sr.design.diagnostic.get("BETATRON_TUNE") assert isinstance(tm.peer.peer, Accelerator) assert isinstance(tm.peer, ElementHolder) - tm = sr.live.get_betatron_tune_monitor("BETATRON_TUNE") + tm = sr.live.diagnostic.get("BETATRON_TUNE") assert isinstance(tm.peer.peer, Accelerator) assert isinstance(tm.peer, ElementHolder) @@ -160,5 +160,5 @@ def test_config_dict(): assert sr.live.dconfig()["prefix"] == "VA:" assert sr.live.dconfig()["info"]["param1"] == "Param1 value" assert sr.live.dconfig()["info"]["param2"] == 12345.0 - assert isinstance(sr.live.get_element("MY_ELEMENT"), MyElement) - assert sr.live.get_element("MY_ELEMENT")._cfg.device_h == "TUNEZR:rdH" + assert isinstance(sr.live["MY_ELEMENT"], MyElement) + assert sr.live["MY_ELEMENT"]._cfg.device_h == "TUNEZR:rdH" diff --git a/tests/test_load_conf_with_code.py b/tests/test_load_conf_with_code.py index 25a98a8dc..41771dbfa 100644 --- a/tests/test_load_conf_with_code.py +++ b/tests/test_load_conf_with_code.py @@ -14,7 +14,7 @@ def test_load_conf_with_code(): 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.live[:].names() == [element.get_name() for element in sr.live.get()] 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 f42471031..22f60c144 100644 --- a/tests/tuning_tools/test_bba.py +++ b/tests/tuning_tools/test_bba.py @@ -12,13 +12,13 @@ def test_bba(): 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") + bba = SR.tool.get("BBA-BPM_C04-04") bba.measure() assert np.abs(bba.h_offset() - 20e-6) < 1e-6 assert np.abs(bba.v_offset() + 15e-6) < 1e-6 # BBA (model dependant method) - bba = SR.get_bba("BBA2-BPM_C04-04") + bba = SR.tool.get("BBA2-BPM_C04-04") bba.minicycle_sleep_time = 0 bba.measure() assert np.abs(bba.h_offset() - 20e-6) < 1e-6 diff --git a/tests/tuning_tools/test_chromaticity_monitor.py b/tests/tuning_tools/test_chromaticity_monitor.py index 3f7671d9f..4f9c7e115 100644 --- a/tests/tuning_tools/test_chromaticity_monitor.py +++ b/tests/tuning_tools/test_chromaticity_monitor.py @@ -13,7 +13,7 @@ def test_simulator_chromaticity_monitor(): sr: Accelerator = Accelerator.load("tests/config/EBSOrbit.yaml", ignore_external=True) sr.design.get_lattice().enable_6d() chromaAT = sr.design.get_lattice().get_chrom()[:-1] - chromaticity_monitor = sr.design.get_chromaticity_monitor("CHROMATICITY_MONITOR") + chromaticity_monitor = sr.design.tool.get("CHROMATICITY_MONITOR") assert chromaticity_monitor.chromaticity.unit() == "1" assert chromaticity_monitor.dispersion.unit() == "m" chromaticity_monitor.measure(fit_dispersion=True, callback=callback) @@ -43,7 +43,7 @@ def test_simulator_chromaticity_monitor(): ) def test_controlsystem_chromaticity_monitor(install_test_package): sr: Accelerator = Accelerator.load("tests/config/EBSOrbit.yaml") - chromaticity_monitor = sr.live.get_chromaticity_monitor("CHROMATICITY_MONITOR") + chromaticity_monitor = sr.live.tool.get("CHROMATICITY_MONITOR") assert chromaticity_monitor.chromaticity.get() is None chromaticity_monitor.measure( do_plot=False, diff --git a/tests/tuning_tools/test_tool_accessors.py b/tests/tuning_tools/test_tool_accessors.py index a5dd2500f..f582040ac 100644 --- a/tests/tuning_tools/test_tool_accessors.py +++ b/tests/tuning_tools/test_tool_accessors.py @@ -13,14 +13,14 @@ def test_tuning_tools_expose_configured_elements(): ) design = sr.design - tune_monitor = design.get_betatron_tune_monitor("BETATRON_TUNE") + tune_monitor = design.diagnostic.get("BETATRON_TUNE") quadrupoles = design.magnets.get("QForTune") assert design.tool.tune.tune_monitor is tune_monitor assert design.tool.tune.quadrupoles is quadrupoles assert design.tool.trm.tune_monitor is tune_monitor assert design.tool.trm.quadrupoles is quadrupoles - chromaticity_monitor = design.get_chromaticity_monitor("CHROMATICITY_MONITOR") + chromaticity_monitor = design.tool.get("CHROMATICITY_MONITOR") sextupoles = design.magnets.get("Sext") assert design.tool.chromaticity.chromaticity_monitor is chromaticity_monitor assert design.tool.chromaticity.sextupoles is sextupoles @@ -47,14 +47,14 @@ def test_tuning_tools_expose_configured_elements(): assert design.tool.orm.vcorrectors is vcorrectors for name in ("BBA-BPM_C04-04", "BBA2-BPM_C04-04"): - bba = design.get_bba(name) + bba = design.tool.get(name) assert bba.bpms is bpms 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.tool.tune + assert design.tool.get("BBA2-BPM_C04-04").tune_correction is design.tool.tune def test_tool_get_returns_named_tool(): diff --git a/tests/tuning_tools/test_tune.py b/tests/tuning_tools/test_tune.py index a114230cf..fe23a48e1 100644 --- a/tests/tuning_tools/test_tune.py +++ b/tests/tuning_tools/test_tune.py @@ -14,7 +14,7 @@ def test_tune(): sr.design.get_lattice().disable_6d() quadForTuneDesign = sr.design.magnets.get("QForTune") - tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") + tune_monitor = sr.design.diagnostic.get("BETATRON_TUNE") # Build tune response matrix tune = tune_monitor.tune.get() print(tune) diff --git a/tests/tuning_tools/test_tune_hardware.py b/tests/tuning_tools/test_tune_hardware.py index 8e83197df..e48ce2657 100644 --- a/tests/tuning_tools/test_tune_hardware.py +++ b/tests/tuning_tools/test_tune_hardware.py @@ -12,7 +12,7 @@ def test_tune(): # Build tune response matrix (hardware units) - tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") + tune_monitor = sr.design.diagnostic.get("BETATRON_TUNE") tune = tune_monitor.tune.get() print(tune) tunemat = np.zeros((len(quadForTuneDesign), 2)) diff --git a/tests/tuning_tools/test_tune_monitor.py b/tests/tuning_tools/test_tune_monitor.py index 893f096ec..66694c370 100644 --- a/tests/tuning_tools/test_tune_monitor.py +++ b/tests/tuning_tools/test_tune_monitor.py @@ -11,7 +11,7 @@ def test_simulator_tune_monitor( ignore_external=True, ) sr.design.get_lattice().disable_6d() - tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") + tune_monitor = sr.design.diagnostic.get("BETATRON_TUNE") assert tune_monitor.tune.get()[0] == sr.design.get_lattice().get_tune()[0] assert tune_monitor.tune.get()[1] == sr.design.get_lattice().get_tune()[1] assert np.abs(tune_monitor.frequency.get()[0] - 56834.22592393) < 1e-6 @@ -31,6 +31,6 @@ def test_controlsystem_tune_monitor( sr = accelerator_from_fragments( *tune_monitor_configuration_fragments, ) - tune_monitor = sr.live.get_betatron_tune_monitor("BETATRON_TUNE") + tune_monitor = sr.live.diagnostic.get("BETATRON_TUNE") assert tune_monitor.tune.get()[0] == 0.0 assert tune_monitor.tune.get()[1] == 0.0