Add convenience methods to extract ORM plane blocks

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")

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