Fix ESRF examples

examples/ESRF_ORM_example/correct_orbit.py

@@ -1,36 +1,35 @@
-import json
-import logging
+import time
 from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
 
 from pyaml.accelerator import Accelerator
-from pyaml.tuning_tools.orbit import ConfigModel as Orbit_ConfigModel
-from pyaml.tuning_tools.orbit import Orbit
 
 parent_folder = Path(__file__).parent
 pyaml_folder = parent_folder.parent.parent
 config_path = pyaml_folder.joinpath("tests/config/EBSOrbit.yaml").resolve()
 sr = Accelerator.load(config_path)
+#ebs = sr.live
 ebs = sr.design
 
 ## get reference
 ref_h, ref_v = ebs.get_bpms("BPM").positions.get().T
 reference = np.concatenate((ref_h, ref_v))
 ########################################################
 
-
-## generate some orbit
-std_kick = 1e-6
 hcorr = ebs.get_magnets("HCorr")
 vcorr = ebs.get_magnets("VCorr")
 bpms = ebs.get_bpms("BPM")
 
-np.random.seed(1)
-# mangle orbit
-hcorr.strengths.set(hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr)))
-vcorr.strengths.set(vcorr.strengths.get() + std_kick * np.random.normal(size=len(vcorr)))
+if ebs == sr.design:
+    ## generate some orbit
+    std_kick = 1e-6
+
+    np.random.seed(1)
+    # mangle orbit
+    hcorr.strengths.set(hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr)))
+    vcorr.strengths.set(vcorr.strengths.get() + std_kick * np.random.normal(size=len(vcorr)))
 
 h0 = hcorr.strengths.get()
 v0 = vcorr.strengths.get()
@@ -42,11 +41,13 @@
 
 ebs.orbit.set_virtual_weight(1000)
 ## Correct the orbit
-ebs.orbit.correct(reference=reference)
+ebs.orbit.correct(reference=None, rf=True)
 # ebs.orbit.correct(plane="H")
 # ebs.orbit.correct(plane="V")
 ########################################################
 
+time.sleep(5)
+
 ## inspect orbit correction
 positions_ac = bpms.positions.get()
 std_ac = np.std(positions_ac, axis=0)

examples/ESRF_ORM_example/measure_dispersion.py

@@ -3,9 +3,7 @@
 import matplotlib.pyplot as plt
 
 from pyaml.accelerator import Accelerator
-from pyaml.common.constants import ACTION_APPLY, ACTION_MEASURE, ACTION_RESTORE
-from pyaml.tuning_tools.dispersion import ConfigModel as Disp_ConfigModel
-from pyaml.tuning_tools.dispersion import Dispersion
+from pyaml.common.constants import Action
 
 parent_folder = Path(__file__).parent
 config_path = parent_folder.parent.parent.joinpath("tests", "config", "EBSOrbit.yaml").resolve()
@@ -14,11 +12,11 @@
 
 
 def callback(action: int, callback_data) -> bool:
-    if action == ACTION_APPLY:
-        print("Changing RF frequency.")
-    elif action == ACTION_MEASURE:
+    if action == Action.APPLY:
+        print("Changing RF frequency")
+    elif action == Action.MEASURE:
         print("Reading orbit.")
-    elif action == ACTION_RESTORE:
+    elif action == Action.RESTORE:
         print("Restoring RF frequency.")
     return True
 

examples/ESRF_ORM_example/measure_ideal_ORM_and_disp.py

@@ -20,6 +20,5 @@
 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"))
+ebs.orm.save("ideal_orm_disp.json")

examples/ESRF_ORM_example/measure_reduced_ORM.py

@@ -4,9 +4,7 @@
 import numpy as np
 
 from pyaml.accelerator import Accelerator
-from pyaml.common.constants import ACTION_RESTORE
-from pyaml.tuning_tools.orbit_response_matrix import ConfigModel as ORM_ConfigModel
-from pyaml.tuning_tools.orbit_response_matrix import OrbitResponseMatrix
+from pyaml.common.constants import Action
 
 # disable printing during ORM measurement to illustrate callback.
 logger = logging.getLogger("pyaml.tuning_tools.orbit_response_matrix").setLevel(logging.WARNING)
@@ -30,18 +28,20 @@
 
 
 def callback(action: int, cdata):
-    if action == ACTION_RESTORE:
-        i = cdata.last_number
-        n_bpms = cdata.raw_up.shape[0] // 2
-        n_correctors = cdata.raw_up.shape[1]
-        corrector = cdata.last_input
-        response = (cdata.raw_up[:, i] - cdata.raw_down[:, i]) / cdata.inputs_delta[i]
-        std_x_resp = np.std(response[:n_bpms])
-        std_y_resp = np.std(response[n_bpms:])
-        print(
-            f"[{i}/{n_correctors}], Measured response of {corrector}: "
-            f"r.m.s H.: {std_x_resp:.2f} mm/mrad, r.m.s. V: {std_y_resp:.2f} mm/mrad"
-        )
+    if action == Action.RESTORE:
+        if "response_data" in cdata:
+            rdata = cdata["response_data"]
+            i = rdata.last_number
+            n_bpms = rdata.raw_up.shape[0] // 2
+            n_correctors = rdata.raw_up.shape[1]
+            corrector = rdata.last_input
+            response = (rdata.raw_up[:, i] - rdata.raw_down[:, i]) / rdata.inputs_delta[i]
+            std_x_resp = np.std(response[:n_bpms])
+            std_y_resp = np.std(response[n_bpms:])
+            print(
+                f"[{i}/{n_correctors}], Measured response of {corrector}: "
+                f"r.m.s H.: {std_x_resp:.2f} mm/mrad, r.m.s. V: {std_y_resp:.2f} mm/mrad"
+            )
     return True
 
 

examples/ORM_Tango_server/ORM.py

@@ -15,7 +15,7 @@
 from pyaml.accelerator import Accelerator
 
 # Additional import
-from pyaml.common.constants import ACTION_RESTORE
+from pyaml.common.constants import Action
 
 __all__ = ["ORM", "main"]
 
@@ -81,26 +81,28 @@ def delete_device(self):
 
     # Orbit callback
     def orbit_callback(self, action: int, cdata):
-        if action == ACTION_RESTORE:
-            i = cdata.last_number
-            n_bpms = cdata.raw_up.shape[0] // 2
-            n_correctors = cdata.raw_up.shape[1]
-            corrector = cdata.last_input
-            response = cdata.raw_up[:, i] - cdata.raw_down[:, i]
-            response = response / cdata.inputs_delta[i]
-            std_x_resp = np.std(response[:n_bpms])
-            std_y_resp = np.std(response[n_bpms:])
-            self.progress_data[2 * i : 2 * i + 2] = [std_x_resp, std_y_resp]
-            self.set_status(
-                f"[{i}/{n_correctors}], Measured response of {corrector}: "
-                f"r.m.s H.: {std_x_resp:.2f} um/mrad, "
-                f"r.n.s V: {std_y_resp:.2f} um/mrad"
-            )
+        if action == Action.RESTORE:
+            if "response_data" in cdata:
+                rdata = cdata["response_data"]
+                i = rdata.last_number
+                n_bpms = rdata.raw_up.shape[0] // 2
+                n_correctors = rdata.raw_up.shape[1]
+                corrector = rdata.last_input
+                response = rdata.raw_up[:, i] - rdata.raw_down[:, i]
+                response = response / rdata.inputs_delta[i]
+                std_x_resp = np.std(response[:n_bpms])
+                std_y_resp = np.std(response[n_bpms:])
+                self.progress_data[2 * i : 2 * i + 2] = [std_x_resp, std_y_resp]
+                self.set_status(
+                    f"[{i}/{n_correctors}], Measured response of {corrector}: "
+                    f"r.m.s H.: {std_x_resp:.2f} um/mrad, "
+                    f"r.n.s V: {std_y_resp:.2f} um/mrad"
+                )
         return True
 
     def orm_run(self):
         # On design sleep 10ms to allow thread schedule
-        self.SR.design.orm.measure(callback=orbit_callback, set_wait_time=0.01)
+        self.SR.design.orm.measure(callback=orbit_callback, sleep_between_meas=0.01)
         self.orm_data = self.SR.design.orm.get()
         self.set_status(f"Ready to scan: {self.ConfigFileName}")
         self.set_state(DevState.ON)

pyaml/tuning_tools/dispersion.py

@@ -65,28 +65,46 @@ def measure(
             skip_save=True,
         )
 
-        self._register_callback(callback)
-        self._init_measure()
-
         aborted = False
-        for code, measurement in generator:
-            callback_data = measurement.dispersion_data  # to be defined better
-            if code is DispersionCode.AFTER_SET:
-                if not self.send_callback(Action.APPLY, callback_data):
-                    if aborted:
+        idx = 0
+        err = None
+        try:
+            self._register_callback(callback)
+            self._init_measure()
+            for code, measurement in generator:
+                callback_data = {"idx": idx, "dispersion_data": measurement.dispersion_data}
+                if code is DispersionCode.AFTER_SET:
+                    if not self.send_callback(Action.APPLY, callback_data):
+                        if aborted:
+                            break
+                elif code is DispersionCode.AFTER_GET:
+                    if not self.send_callback(Action.MEASURE, callback_data):
+                        aborted = True
+                        break
+                elif code is DispersionCode.AFTER_RESTORE:
+                    if not self.send_callback(Action.RESTORE, callback_data):
+                        aborted = True
                         break
-            elif code is DispersionCode.AFTER_GET:
-                if not self.send_callback(Action.MEASURE, callback_data):
-                    aborted = True
-                    break
-            elif code is DispersionCode.AFTER_RESTORE:
-                if not self.send_callback(Action.RESTORE, callback_data):
-                    aborted = True
-                    break
+                idx += 1
+        except Exception as ex:
+            err = ex
+        except KeyboardInterrupt as ex:
+            aborted = True
+        finally:
+            # Restore RF
+            # TODO
+            self.send_callback(
+                Action.RESTORE,
+                {"idx": idx},
+                raiseException=False,
+            )
+
+        if err is not None:
+            raise (err)
 
         if aborted:
-            logger.warning("Measurement aborted! Settings have not been restored.")
-            return
+            logger.warning(f"{self.get_name()} : measurement aborted (settings not restored)")
+            return False
 
         dispersion_data = measurement.dispersion_data
         # contains also pre-processed data
@@ -96,5 +114,7 @@ def measure(
         # ).names()
         self.latest_measurement.update(dispersion_data.model_dump())
 
+        return True
+
     def get(self):
         return self.latest_measurement

pyaml/tuning_tools/orbit_response_matrix.py

@@ -8,7 +8,6 @@
 from pySC.apps.codes import ResponseCode
 
 from ..common.constants import Action
-from ..common.element import ElementConfigModel
 from ..external.pySC_interface import pySCInterface
 from .measurement_tool import MeasurementTool, MeasurementToolConfigModel
 from .orbit_response_matrix_data import ConfigModel as OrbitResponseMatrixDataConfigModel
@@ -126,7 +125,7 @@ def measure(
             self._register_callback(callback)
             self._init_measure()
             for code, measurement in generator:
-                callback_data = measurement.response_data  # to be defined better
+                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:
@@ -159,6 +158,8 @@ def measure(
         self.latest_measurement.update(orm_data.model_dump())
         self.latest_measurement["type"] = "pyaml.tuning_tools.orbit_response_matrix_data"
 
+        return True
+
     def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataConfigModel:
         # all metadata is discarded here. Should we keep something?