diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 14bd96d2..d44de46c 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -317,3 +317,51 @@ def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixData: 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 3ea994f0..a2faa273 100644 --- a/tests/tuning_tools/test_tuning_orm.py +++ b/tests/tuning_tools/test_tuning_orm.py @@ -32,3 +32,37 @@ def test_tuning_orm(): # 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