Merge pull request #220 from python-accelerator-middle-layer/orm-saving

orm saving

Author: JeanLucPons

examples/ESRF_ORM_example/measure_ideal_ORM_and_disp.py

@@ -6,9 +6,7 @@
 from pyaml.accelerator import Accelerator
 
 parent_folder = Path(__file__).parent
-config_path = parent_folder.parent.parent.joinpath(
-    "tests", "config", "EBSOrbit.yaml"
-).resolve()
+config_path = parent_folder.parent.parent.joinpath("tests", "config", "EBSOrbit.yaml").resolve()
 sr = Accelerator.load(config_path)
 
 ebs = sr.design
@@ -19,18 +17,9 @@
 ebs.dispersion.measure()
 dispersion_data = ebs.dispersion.get()
 
-rf_response = (
-    dispersion_data["frequency_response_x"] + dispersion_data["frequency_response_y"]
-)
-
-ideal_ORM_data = {
-    "type": "pyaml.tuning_tools.response_matrix",
-    "matrix": orm_data["matrix"],
-    "input_names": orm_data["input_names"],
-    "output_names": orm_data["output_names"],
-    "input_planes": orm_data["input_planes"],
-    "output_planes": orm_data["output_planes"],
-    "rf_response": rf_response,
-}
-
-json.dump(ideal_ORM_data, open("ideal_orm_disp.json", "w"))
+rf_response = dispersion_data["frequency_response_x"] + dispersion_data["frequency_response_y"]
+
+orm_data["rf_response"] = rf_response
+orm_data["type"] = "pyaml.tuning_tools.orbit_response_matrix_data"
+
+json.dump(orm_data, open("ideal_orm_disp.json", "w"))

pyaml/tuning_tools/orbit_response_matrix.py

@@ -11,6 +11,8 @@
 from ..common.element import ElementConfigModel
 from ..external.pySC_interface import pySCInterface
 from .measurement_tool import MeasurementTool
+from .orbit_response_matrix_data import ConfigModel as OrbitResponseMatrixDataConfigModel
+from .orbit_response_matrix_data import OrbitResponseMatrixData
 
 logger = logging.getLogger(__name__)
 
@@ -82,11 +84,6 @@ def measure(
             hcorrector_names = element_holder.get_magnets(self.hcorr_array_name).names()
             vcorrector_names = element_holder.get_magnets(self.vcorr_array_name).names()
             corrector_names = hcorrector_names + vcorrector_names
-        else:
-            all_hcorrector_names = element_holder.get_magnets(self.hcorr_array_name).names()
-            all_vcorrector_names = element_holder.get_magnets(self.vcorr_array_name).names()
-            hcorrector_names = [corr for corr in corrector_names if corr in all_hcorrector_names]
-            vcorrector_names = [corr for corr in corrector_names if corr in all_vcorrector_names]
 
         generator = measure_ORM(
             interface=interface,
@@ -117,19 +114,36 @@ def measure(
             logger.warning("Measurement aborted! Settings have not been restored.")
             return
 
-        response_data = measurement.response_data  # contains also pre-processed data
-        bpm_names = element_holder.get_bpms(self.bpm_array_name).names()
-        # This is because we assume always dual-plane bpms now.
-        response_data.output_names = bpm_names * 2
-        self.latest_measurement = response_data.model_dump()
+        orm_data = self._pySC_response_data_to_ORMData(measurement.response_data.model_dump())
+        self.latest_measurement = orm_data.model_dump()
+
+    def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataConfigModel:
+        # all metadata is discarded here. Should we keep something?
 
-        input_planes = []
-        for corr in corrector_names:
-            if corr in hcorrector_names:
-                input_planes.append("H")
-            elif corr in vcorrector_names:
-                input_planes.append("V")
-        self.latest_measurement["input_planes"] = input_planes
+        element_holder = self._peer
+        all_hcorrector_names = element_holder.get_magnets(self.hcorr_array_name).names()
+        all_vcorrector_names = element_holder.get_magnets(self.vcorr_array_name).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 = element_holder.get_bpms(self.bpm_array_name).names()
+        # This is because we assume always dual-plane bpms now.
         len_b = len(bpm_names)
-        self.latest_measurement["output_planes"] = ["H"] * len_b + ["V"] * len_b
+        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 = OrbitResponseMatrixDataConfigModel(**orm_data_model)
+        return orm_data