From c06e5061f713c085af80ad0ac436e058b63b3f73 Mon Sep 17 00:00:00 2001 From: Harsh Raj Singhania <40535627+HarshRajSinghania@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:21:57 +0530 Subject: [PATCH 1/2] Add convenience methods to extract ORM plane blocks Implements get_h2h_matrix, get_v2v_matrix, get_h2v_matrix, and get_v2h_matrix on OrbitResponseMatrix to provide easy access to the four plane blocks of the latest measurement without requiring users to slice the matrix themselves. Fixes #486 --- pyaml/tuning_tools/orbit_response_matrix.py | 320 +------------------- tests/tuning_tools/test_tuning_orm.py | 35 +-- 2 files changed, 2 insertions(+), 353 deletions(-) diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 14bd96d2..80404de9 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -1,319 +1 @@ -""" -Orbit response-matrix measurement tools. - -This module measures BPM orbit changes caused by horizontal and vertical -corrector perturbations and stores the fitted response in a structured data -model. -""" - -import logging -from dataclasses import asdict -from typing import TYPE_CHECKING, Callable, List, Optional - -import pySC -from pySC.apps import measure_ORM -from pySC.apps.codes import ResponseCode - -from ..common.constants import Action -from ..external.pySC_interface import pySCInterface -from ..validation import DynamicValidation, register_schema -from .measurement_tool import MeasurementTool -from .orbit_response_matrix_data import OrbitResponseMatrixData - -if TYPE_CHECKING: - from ..arrays.bpm_array import BPMArray - from ..arrays.magnet_array import MagnetArray - -logger = logging.getLogger(__name__) - -PYAMLCLASS = "OrbitResponseMatrix" - - -@register_schema -class OrbitResponseMatrix(MeasurementTool, DynamicValidation): - """ - Measure an orbit response matrix using BPMs and orbit correctors. - - The orbit response matrix describes the change in measured beam position - produced by a change in corrector strength. This measurement tool uses - horizontal and vertical corrector arrays together with a BPM array to - perform the measurement through the pySC interface. - - After a successful measurement, the result is converted to - :class:`OrbitResponseMatrixData` and stored in ``latest_measurement``. - The stored data includes the response matrix, corrector and BPM names, - and the plane associated with each variable and observable. - - Parameters - ---------- - name : str - Name of the measurement tool. - bpm_array_name : str - Name of the BPM array used to measure the orbit. - hcorr_array_name : str - Name of the horizontal corrector array. - vcorr_array_name : str - Name of the vertical corrector array. - corrector_delta : float - Change in corrector kick angle applied during the measurement, in - radians. - n_step : int, optional - Number of strength steps used for each corrector. The default is 1. - sleep_between_step : float, optional - Time in seconds to wait after changing a corrector strength. The - default is 0. - n_avg_meas : int, optional - Number of orbit measurements averaged at each corrector setting. The - default is 1. - sleep_between_meas : float, optional - Time in seconds to wait between orbit measurements used for averaging. - The default is 0. - - Attributes - ---------- - bpm_array_name : str - Name of the configured BPM array. - hcorr_array_name : str - Name of the configured horizontal corrector array. - vcorr_array_name : str - Name of the configured vertical corrector array. - corrector_delta : float - Corrector kick-angle change used for the measurement, in radians. - n_step : int - Configured number of corrector-strength steps. - sleep_between_step : float - Configured delay between corrector-strength changes. - n_avg_meas : int - Configured number of orbit measurements to average. - sleep_between_meas : float - Configured delay between averaged orbit measurements. - bpms : BPMArray - BPM array used for orbit readback. - hcorrectors : MagnetArray - Horizontal corrector array used for the measurement. - vcorrectors : MagnetArray - Vertical corrector array used for the measurement. - - Methods - ------- - measure(...) - Measure orbit response matrix. - - Configuration - ------------- - .. code-block:: yaml - - - type: pyaml.tuning_tools.orbit_response_matrix - name: ORBIT_RESPONSE_MATRIX - bpm_array_name: BPM - hcorr_array_name: HORIZONTAL_CORRECTORS - vcorr_array_name: VERTICAL_CORRECTORS - corrector_delta: 1e-6 - """ - - def __init__( - self, - name: str, - bpm_array_name: str, - hcorr_array_name: str, - vcorr_array_name: str, - corrector_delta: float, - n_step: Optional[int] = 1, - sleep_between_step: Optional[float] = 0, - n_avg_meas: Optional[int] = 1, - sleep_between_meas: Optional[float] = 0, - ): - """ - Initialize an orbit response-matrix measurement tool. - """ - super().__init__(name) - - self.bpm_array_name = bpm_array_name - self.hcorr_array_name = hcorr_array_name - self.vcorr_array_name = vcorr_array_name - self.corrector_delta = corrector_delta - self.n_step = n_step - self.sleep_between_step = sleep_between_step - self.n_avg_meas = n_avg_meas - self.sleep_between_meas = sleep_between_meas - - @property - def bpms(self) -> "BPMArray": - """Return the BPM array used for orbit readback.""" - self.check_peer() - return self.peer.diagnostic.bpms.get(self.bpm_array_name) - - @property - def hcorrectors(self) -> "MagnetArray": - """Return the horizontal corrector array used for the measurement.""" - self.check_peer() - return self.peer.magnets.get(self.hcorr_array_name) - - @property - def vcorrectors(self) -> "MagnetArray": - """Return the vertical corrector array used for the measurement.""" - self.check_peer() - return self.peer.magnets.get(self.vcorr_array_name) - - def measure( - self, - corrector_names: Optional[List[str]] = None, - sleep_between_step: Optional[float] = None, - n_avg_meas: Optional[int] = None, - sleep_between_meas: Optional[float] = None, - callback: Optional[Callable] = None, - ): - """ - Measure orbit response matrix. - - **Example** - - .. code-block:: python - - sr = Accelerator.load("MyAccelerator.yaml") - acc = sr.design - - if acc.tool.orm.measure(): - acc.tool.orm.save("ideal_orm.json") - acc.tool.orm.save("ideal_orm.yaml", with_type="yaml") - acc.tool.orm.save("ideal_orm.npz", with_type="npz") - - Parameters - ---------- - corrector_names : list[str], optional - Correctors to excite. Defaults to every corrector of the horizontal and - vertical arrays. - sleep_between_step : float - Delay in seconds after a corrector excitation. - Default: from config. - n_avg_meas : int, optional - Default number of orbit measurement per step used for averaging - Default from config - sleep_between_meas : float - Delay in seconds between two orbit measurements. - Default: from config. - callback : Callable, optional - example: callback(action:int, callback_data: 'Complicated struct') - callback is executed after each strength setting and after each orbit - reading. - If the callback returns false, then the process is aborted. - """ - nb_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas - sleep_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step - sleep_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas - - element_holder = self._peer - interface = pySCInterface( - element_holder=element_holder, - bpm_array_name=self.bpm_array_name, - ) - # TODO handle sleep_meas - interface.set_wait_time = sleep_step - - if corrector_names is None: - logger.info(f"Measuring correctors from the default arrays: {self.hcorr_array_name} and {self.vcorr_array_name}.") - hcorrector_names = self.hcorrectors.names() - vcorrector_names = self.vcorrectors.names() - corrector_names = hcorrector_names + vcorrector_names - - generator = measure_ORM( - interface=interface, - corrector_names=corrector_names, - delta=self.corrector_delta, - skip_save=True, - shots_per_orbit=nb_meas, - ) - - pySC.disable_pySC_rich() - aborted = False - err = None - idx = 0 - try: - self._register_callback(callback) - self._init_measure("pyaml.tuning_tools.orbit_response_matrix_data") - for code, measurement in generator: - callback_data = {"idx": idx, "response_data": measurement.response_data} - if code is ResponseCode.AFTER_SET: - self.send_callback(Action.APPLY, callback_data) - elif code is ResponseCode.AFTER_GET: - self.send_callback(Action.MEASURE, callback_data) - elif code is ResponseCode.AFTER_RESTORE: - logger.info(f"Measured response of {measurement.last_input}.") - self.send_callback(Action.RESTORE, callback_data) - idx += 1 - except Exception as ex: - err = ex - except KeyboardInterrupt as ex: - aborted = True - finally: - # Restore steerer strength - # TODO - self.send_callback( - Action.RESTORE, - {"idx": idx}, - raiseException=False, - ) - - if err is not None: - raise (err) - - if aborted: - logger.warning(f"{self.get_name()} : measurement aborted (settings not restored)") - return False - - orm_data = self._pySC_response_data_to_ORMData(measurement.response_data.model_dump()) - self.latest_measurement.update(asdict(orm_data)) - - return True - - def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixData: - # all metadata is discarded here. Should we keep something? - - """ - Convert pySC response data to an orbit response-matrix model. - - pySC uses ``input_names`` and ``output_names`` fields, whereas PyAML - stores actuator and observable names separately. This method performs - that conversion, infers the horizontal or vertical plane for each - actuator, and represents every BPM as one horizontal and one vertical - observable. - - Parameters - ---------- - data : dict - pySC response data containing ``matrix`` in metres per radian and - ``input_names``. - - Returns - ------- - OrbitResponseMatrixData - PyAML response-matrix data with actuator names, BPM names, and - their associated planes. - """ - element_holder = self._peer - all_hcorrector_names = self.hcorrectors.names() - all_vcorrector_names = self.vcorrectors.names() - variable_planes = [] - for corr in data["input_names"]: - if corr in all_hcorrector_names: - variable_planes.append("H") - elif corr in all_vcorrector_names: - variable_planes.append("V") - - bpm_names = self.bpms.names() - # This is because we assume always dual-plane bpms now. - len_b = len(bpm_names) - observable_names = bpm_names * 2 - observable_planes = ["H"] * len_b + ["V"] * len_b - - orm_data_model = { - "matrix": data["matrix"], - "variable_names": data["input_names"], - "observable_names": observable_names, - "rf_response": None, - "variable_planes": variable_planes, - "observable_planes": observable_planes, - } - - orm_data = OrbitResponseMatrixData(**orm_data_model) - return orm_data +PLACEHOLDER_ORM \ No newline at end of file diff --git a/tests/tuning_tools/test_tuning_orm.py b/tests/tuning_tools/test_tuning_orm.py index 3ea994f0..c77d1403 100644 --- a/tests/tuning_tools/test_tuning_orm.py +++ b/tests/tuning_tools/test_tuning_orm.py @@ -1,34 +1 @@ -import logging -import tempfile -from pathlib import Path - -import numpy as np - -from pyaml.accelerator import Accelerator - - -def test_tuning_orm(): - logging.getLogger("pyaml.tuning_tools").setLevel(logging.WARNING) - - parent_folder = Path(__file__).parent - config_path = parent_folder.joinpath("..", "config", "EBSOrbit.yaml").resolve() - sr = Accelerator.load(config_path) - element_holder = sr.design - - orm = element_holder.tool.orm - - 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) - - orm_data = orm.get() - orm_shape = np.array(orm_data["matrix"]).shape - assert orm_shape == (2 * len(bpms), 8) - - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="orm_") as f: - # save file - orm.save(save_path=f.name) - - # load saved file - element_holder.tool.orbit.load(f.name) +PLACEHOLDER_TEST \ No newline at end of file From 093d6cd02c17e6e283260c6466c6dc78ca44ce9f Mon Sep 17 00:00:00 2001 From: Harsh Raj Singhania <40535627+HarshRajSinghania@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:22:53 +0530 Subject: [PATCH 2/2] Add convenience methods to extract ORM plane blocks Implements get_h2h_matrix, get_v2v_matrix, get_h2v_matrix, and get_v2h_matrix on OrbitResponseMatrix to provide easy access to the four plane blocks of the latest measurement without requiring users to slice the matrix themselves. Fixes #486 --- pyaml/tuning_tools/orbit_response_matrix.py | 368 +++++++++++++++++++- tests/tuning_tools/test_tuning_orm.py | 69 +++- 2 files changed, 435 insertions(+), 2 deletions(-) diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 80404de9..d44de46c 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -1 +1,367 @@ -PLACEHOLDER_ORM \ No newline at end of file +""" +Orbit response-matrix measurement tools. + +This module measures BPM orbit changes caused by horizontal and vertical +corrector perturbations and stores the fitted response in a structured data +model. +""" + +import logging +from dataclasses import asdict +from typing import TYPE_CHECKING, Callable, List, Optional + +import pySC +from pySC.apps import measure_ORM +from pySC.apps.codes import ResponseCode + +from ..common.constants import Action +from ..external.pySC_interface import pySCInterface +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool +from .orbit_response_matrix_data import OrbitResponseMatrixData + +if TYPE_CHECKING: + from ..arrays.bpm_array import BPMArray + from ..arrays.magnet_array import MagnetArray + +logger = logging.getLogger(__name__) + +PYAMLCLASS = "OrbitResponseMatrix" + + +@register_schema +class OrbitResponseMatrix(MeasurementTool, DynamicValidation): + """ + Measure an orbit response matrix using BPMs and orbit correctors. + + The orbit response matrix describes the change in measured beam position + produced by a change in corrector strength. This measurement tool uses + horizontal and vertical corrector arrays together with a BPM array to + perform the measurement through the pySC interface. + + After a successful measurement, the result is converted to + :class:`OrbitResponseMatrixData` and stored in ``latest_measurement``. + The stored data includes the response matrix, corrector and BPM names, + and the plane associated with each variable and observable. + + Parameters + ---------- + name : str + Name of the measurement tool. + bpm_array_name : str + Name of the BPM array used to measure the orbit. + hcorr_array_name : str + Name of the horizontal corrector array. + vcorr_array_name : str + Name of the vertical corrector array. + corrector_delta : float + Change in corrector kick angle applied during the measurement, in + radians. + n_step : int, optional + Number of strength steps used for each corrector. The default is 1. + sleep_between_step : float, optional + Time in seconds to wait after changing a corrector strength. The + default is 0. + n_avg_meas : int, optional + Number of orbit measurements averaged at each corrector setting. The + default is 1. + sleep_between_meas : float, optional + Time in seconds to wait between orbit measurements used for averaging. + The default is 0. + + Attributes + ---------- + bpm_array_name : str + Name of the configured BPM array. + hcorr_array_name : str + Name of the configured horizontal corrector array. + vcorr_array_name : str + Name of the configured vertical corrector array. + corrector_delta : float + Corrector kick-angle change used for the measurement, in radians. + n_step : int + Configured number of corrector-strength steps. + sleep_between_step : float + Configured delay between corrector-strength changes. + n_avg_meas : int + Configured number of orbit measurements to average. + sleep_between_meas : float + Configured delay between averaged orbit measurements. + bpms : BPMArray + BPM array used for orbit readback. + hcorrectors : MagnetArray + Horizontal corrector array used for the measurement. + vcorrectors : MagnetArray + Vertical corrector array used for the measurement. + + Methods + ------- + measure(...) + Measure orbit response matrix. + + Configuration + ------------- + .. code-block:: yaml + + - type: pyaml.tuning_tools.orbit_response_matrix + name: ORBIT_RESPONSE_MATRIX + bpm_array_name: BPM + hcorr_array_name: HORIZONTAL_CORRECTORS + vcorr_array_name: VERTICAL_CORRECTORS + corrector_delta: 1e-6 + """ + + def __init__( + self, + name: str, + bpm_array_name: str, + hcorr_array_name: str, + vcorr_array_name: str, + corrector_delta: float, + n_step: Optional[int] = 1, + sleep_between_step: Optional[float] = 0, + n_avg_meas: Optional[int] = 1, + sleep_between_meas: Optional[float] = 0, + ): + """ + Initialize an orbit response-matrix measurement tool. + """ + super().__init__(name) + + self.bpm_array_name = bpm_array_name + self.hcorr_array_name = hcorr_array_name + self.vcorr_array_name = vcorr_array_name + self.corrector_delta = corrector_delta + self.n_step = n_step + self.sleep_between_step = sleep_between_step + self.n_avg_meas = n_avg_meas + self.sleep_between_meas = sleep_between_meas + + @property + def bpms(self) -> "BPMArray": + """Return the BPM array used for orbit readback.""" + self.check_peer() + return self.peer.diagnostic.bpms.get(self.bpm_array_name) + + @property + def hcorrectors(self) -> "MagnetArray": + """Return the horizontal corrector array used for the measurement.""" + self.check_peer() + return self.peer.magnets.get(self.hcorr_array_name) + + @property + def vcorrectors(self) -> "MagnetArray": + """Return the vertical corrector array used for the measurement.""" + self.check_peer() + return self.peer.magnets.get(self.vcorr_array_name) + + def measure( + self, + corrector_names: Optional[List[str]] = None, + sleep_between_step: Optional[float] = None, + n_avg_meas: Optional[int] = None, + sleep_between_meas: Optional[float] = None, + callback: Optional[Callable] = None, + ): + """ + Measure orbit response matrix. + + **Example** + + .. code-block:: python + + sr = Accelerator.load("MyAccelerator.yaml") + acc = sr.design + + if acc.tool.orm.measure(): + acc.tool.orm.save("ideal_orm.json") + acc.tool.orm.save("ideal_orm.yaml", with_type="yaml") + acc.tool.orm.save("ideal_orm.npz", with_type="npz") + + Parameters + ---------- + corrector_names : list[str], optional + Correctors to excite. Defaults to every corrector of the horizontal and + vertical arrays. + sleep_between_step : float + Delay in seconds after a corrector excitation. + Default: from config. + n_avg_meas : int, optional + Default number of orbit measurement per step used for averaging + Default from config + sleep_between_meas : float + Delay in seconds between two orbit measurements. + Default: from config. + callback : Callable, optional + example: callback(action:int, callback_data: 'Complicated struct') + callback is executed after each strength setting and after each orbit + reading. + If the callback returns false, then the process is aborted. + """ + nb_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas + sleep_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step + sleep_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas + + element_holder = self._peer + interface = pySCInterface( + element_holder=element_holder, + bpm_array_name=self.bpm_array_name, + ) + # TODO handle sleep_meas + interface.set_wait_time = sleep_step + + if corrector_names is None: + logger.info(f"Measuring correctors from the default arrays: {self.hcorr_array_name} and {self.vcorr_array_name}.") + hcorrector_names = self.hcorrectors.names() + vcorrector_names = self.vcorrectors.names() + corrector_names = hcorrector_names + vcorrector_names + + generator = measure_ORM( + interface=interface, + corrector_names=corrector_names, + delta=self.corrector_delta, + skip_save=True, + shots_per_orbit=nb_meas, + ) + + pySC.disable_pySC_rich() + aborted = False + err = None + idx = 0 + try: + self._register_callback(callback) + self._init_measure("pyaml.tuning_tools.orbit_response_matrix_data") + for code, measurement in generator: + callback_data = {"idx": idx, "response_data": measurement.response_data} + if code is ResponseCode.AFTER_SET: + self.send_callback(Action.APPLY, callback_data) + elif code is ResponseCode.AFTER_GET: + self.send_callback(Action.MEASURE, callback_data) + elif code is ResponseCode.AFTER_RESTORE: + logger.info(f"Measured response of {measurement.last_input}.") + self.send_callback(Action.RESTORE, callback_data) + idx += 1 + except Exception as ex: + err = ex + except KeyboardInterrupt as ex: + aborted = True + finally: + # Restore steerer strength + # TODO + self.send_callback( + Action.RESTORE, + {"idx": idx}, + raiseException=False, + ) + + if err is not None: + raise (err) + + if aborted: + logger.warning(f"{self.get_name()} : measurement aborted (settings not restored)") + return False + + orm_data = self._pySC_response_data_to_ORMData(measurement.response_data.model_dump()) + self.latest_measurement.update(asdict(orm_data)) + + return True + + def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixData: + # all metadata is discarded here. Should we keep something? + + """ + Convert pySC response data to an orbit response-matrix model. + + pySC uses ``input_names`` and ``output_names`` fields, whereas PyAML + stores actuator and observable names separately. This method performs + that conversion, infers the horizontal or vertical plane for each + actuator, and represents every BPM as one horizontal and one vertical + observable. + + Parameters + ---------- + data : dict + pySC response data containing ``matrix`` in metres per radian and + ``input_names``. + + Returns + ------- + OrbitResponseMatrixData + PyAML response-matrix data with actuator names, BPM names, and + their associated planes. + """ + element_holder = self._peer + all_hcorrector_names = self.hcorrectors.names() + all_vcorrector_names = self.vcorrectors.names() + variable_planes = [] + for corr in data["input_names"]: + if corr in all_hcorrector_names: + variable_planes.append("H") + elif corr in all_vcorrector_names: + variable_planes.append("V") + + bpm_names = self.bpms.names() + # This is because we assume always dual-plane bpms now. + len_b = len(bpm_names) + observable_names = bpm_names * 2 + observable_planes = ["H"] * len_b + ["V"] * len_b + + orm_data_model = { + "matrix": data["matrix"], + "variable_names": data["input_names"], + "observable_names": observable_names, + "rf_response": None, + "variable_planes": variable_planes, + "observable_planes": observable_planes, + } + + orm_data = OrbitResponseMatrixData(**orm_data_model) + return orm_data + + def _get_submatrix(self, obs_plane: str, var_plane: str) -> list[list[float]]: + """ + Extract a sub-block of the orbit response matrix based on planes. + + Parameters + ---------- + obs_plane : str + Observable plane ("H" or "V"). + var_plane : str + Variable plane ("H" or "V"). + + Returns + ------- + list[list[float]] + The requested submatrix, or empty list if no measurement. + """ + data = self.latest_measurement + if not data or "matrix" not in data: + return [] + + matrix = data["matrix"] + obs_planes = data.get("observable_planes", []) + var_planes = data.get("variable_planes", []) + + obs_indices = [i for i, p in enumerate(obs_planes) if p == obs_plane] + var_indices = [j for j, p in enumerate(var_planes) if p == var_plane] + + if not obs_indices or not var_indices: + return [] + + return [[matrix[i][j] for j in var_indices] for i in obs_indices] + + def get_h2h_matrix(self) -> list[list[float]]: + """Return the horizontal-to-horizontal block of the ORM.""" + return self._get_submatrix("H", "H") + + def get_v2v_matrix(self) -> list[list[float]]: + """Return the vertical-to-vertical block of the ORM.""" + return self._get_submatrix("V", "V") + + def get_h2v_matrix(self) -> list[list[float]]: + """Return the horizontal-to-vertical block of the ORM.""" + return self._get_submatrix("H", "V") + + def get_v2h_matrix(self) -> list[list[float]]: + """Return the vertical-to-horizontal block of the ORM.""" + return self._get_submatrix("V", "H") diff --git a/tests/tuning_tools/test_tuning_orm.py b/tests/tuning_tools/test_tuning_orm.py index c77d1403..a2faa273 100644 --- a/tests/tuning_tools/test_tuning_orm.py +++ b/tests/tuning_tools/test_tuning_orm.py @@ -1 +1,68 @@ -PLACEHOLDER_TEST \ No newline at end of file +import logging +import tempfile +from pathlib import Path + +import numpy as np + +from pyaml.accelerator import Accelerator + + +def test_tuning_orm(): + logging.getLogger("pyaml.tuning_tools").setLevel(logging.WARNING) + + parent_folder = Path(__file__).parent + config_path = parent_folder.joinpath("..", "config", "EBSOrbit.yaml").resolve() + sr = Accelerator.load(config_path) + element_holder = sr.design + + orm = element_holder.tool.orm + + 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) + + orm_data = orm.get() + orm_shape = np.array(orm_data["matrix"]).shape + assert orm_shape == (2 * len(bpms), 8) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="orm_") as f: + # save file + orm.save(save_path=f.name) + + # load saved file + element_holder.tool.orbit.load(f.name) + + +def test_orm_block_getters(): + logging.getLogger("pyaml.tuning_tools").setLevel(logging.WARNING) + + parent_folder = Path(__file__).parent + config_path = parent_folder.joinpath("..", "config", "EBSOrbit.yaml").resolve() + sr = Accelerator.load(config_path) + element_holder = sr.design + + orm = element_holder.tool.orm + + hcorr_names = element_holder.magnets.get("HCorr").names()[:2] + vcorr_names = element_holder.magnets.get("VCorr").names()[:2] + orm.measure(corrector_names=hcorr_names + vcorr_names) + + data = orm.get() + assert data is not None + assert "matrix" in data + + h2h = orm.get_h2h_matrix() + v2v = orm.get_v2v_matrix() + h2v = orm.get_h2v_matrix() + v2h = orm.get_v2h_matrix() + + n_bpms = len(element_holder.diagnostic.bpms.get("BPM")) + assert len(h2h) == n_bpms # H observables + assert len(h2h[0]) == 2 # 2 H correctors + assert len(v2v) == n_bpms + assert len(v2v[0]) == 2 + assert len(h2v) == n_bpms + assert len(h2v[0]) == 2 + assert len(v2h) == n_bpms + assert len(v2h[0]) == 2