Add Ophyd test using BESSY2 VA

examples/BESSY2_example/README.md

@@ -0,0 +1,30 @@
+Run the BESSY va:
+
+```bash
+apptainer run oras://registry.hzdr.de/digital-twins-for-accelerators/containers/pyat-softioc-digital-twin:v0-1-2-bessy.2475331
+```
+
+Run the epics tools:
+```bash
+apptainer run oras://registry.hzdr.de/digital-twins-for-accelerators/epics-tools:v0-1-0.2028728
+```
+
+In the epics tools:
+```bash
+epics tools ~/Desktop > pvlist
+GUID 0x4A709B99A866F1326CBBE22F version 2: tcp@[ 160.103.10.135:5075 ]
+epics tools ~/Desktop > pvlist 0x4A709B99A866F1326CBBE22F
+pons:CAVH1T8R:freq
+pons:CAVH2T8R:freq
+...
+```
+
+Edit your yaml config file and setup the control system prefix:
+```yaml
+controls:
+  - type: pyaml_cs_oa.controlsystem
+    prefix: "your_prefix:"
+    name: live
+```
+
+Run the exmaples

examples/BESSY2_example/bessy2-orbit.py

@@ -0,0 +1,89 @@
+import json
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyaml.accelerator import Accelerator
+from pyaml.tuning_tools.orbit_response_matrix import ConfigModel as ORM_ConfigModel
+from pyaml.tuning_tools.orbit_response_matrix import OrbitResponseMatrix
+
+# Load the configuration
+sr = Accelerator.load("BESSY2Orbit.yaml")
+
+
+# if the ORM is not present measure it
+if sr.design.orbit.response_matrix is None:
+    # Measure ORM on design or on live
+
+    # SR = sr.design
+    SR = sr.live
+
+    orm = OrbitResponseMatrix(
+        cfg=ORM_ConfigModel(
+            bpm_array_name="BPM",
+            hcorr_array_name="HCorr",
+            vcorr_array_name="VCorr",
+            corrector_delta=1e-6,
+        ),
+        element_holder=SR,  # Measurement target
+    )
+    orm.measure(set_wait_time=0.0 if SR == sr.design else 2.0)
+    orm_data = orm.get()
+    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"],
+        "inputs_plane": orm_data["inputs_plane"],
+    }
+    json.dump(ideal_ORM_data, open("ideal_orm.json", "w"))
+    # load the response on the live
+    sr.live.orbit.load_response_matrix("ideal_orm.json")
+
+# handle for live
+orbit = sr.live.get_bpms("BPM").positions
+hcorr = sr.live.get_magnets("HCorr")
+vcorr = sr.live.get_magnets("VCorr")
+
+# Mangle the orbit
+std_kick = 1e-6
+hcorr.strengths.set(std_kick * np.random.normal(size=len(hcorr)))
+vcorr.strengths.set(std_kick * np.random.normal(size=len(vcorr)))
+time.sleep(3)
+
+
+positions_bc = orbit.get()
+
+# Correct the orbit
+sr.live.orbit.correct()
+# sr.live.orbit.correct(plane="H")
+# sr.live.orbit.correct(plane="V",gain = 1.0/2.5)
+
+time.sleep(3)
+positions_ac = orbit.get()
+
+# Plot
+fig = plt.figure()
+ax1 = fig.add_subplot(311)
+ax2 = fig.add_subplot(312)
+ax3 = fig.add_subplot(313)
+ax1.plot(positions_bc[:, 0] * 1e6, label="Orbit before correction")
+ax2.plot(positions_bc[:, 1] * 1e6, label="Orbit before correction")
+ax1.plot(positions_ac[:, 0] * 1e6, label="Orbit after correction")
+ax2.plot(positions_ac[:, 1] * 1e6, label="Orbit after correction")
+
+ax3.plot(hcorr.strengths.get(), label="H Steerers")
+ax3.plot(vcorr.strengths.get(), label="V Steerers")
+
+ax1.set_ylabel("Horizontal pos. [μm]")
+ax2.set_ylabel("Vertical pos. [μm]")
+ax2.set_xlabel("BPM number")
+ax3.set_ylabel("Strength (rad)")
+ax3.set_xlabel("Steerer number")
+ax1.legend()
+ax2.legend()
+ax3.legend()
+fig.tight_layout()
+
+plt.show()

examples/BESSY2_example/bessy2-tune.py

@@ -0,0 +1,32 @@
+import time
+
+import numpy as np
+
+from pyaml.accelerator import Accelerator
+from pyaml.common.constants import ACTION_RESTORE
+from pyaml.magnet.magnet import Magnet
+
+
+def tune_callback(step: int, action: int, m: Magnet, dtune: np.array):
+    if action == ACTION_RESTORE:
+        # On action restore, the measured delta tune is passed as argument
+        print(f"Tune response: #{step} {m.get_name()} {dtune}")
+    return True
+
+
+sr = Accelerator.load("BESSY2Tune.yaml")
+
+# print(sr.live.get_magnets("QForTune").strengths.get())
+# print(sr.design.get_magnets("QForTune").strengths.get())
+# quit()
+
+sr.design.get_lattice().disable_6d()
+tune_adjust = sr.design.tune
+tune_adjust.response.measure(callback=tune_callback)
+tune_adjust.response.save_json("tunemat-bessy.json")
+
+sr.live.tune.response.load_json("tunemat-bessy.json")
+print(sr.live.tune.readback())
+sr.live.tune.set([0.83, 0.84], iter=2, wait_time=3)
+time.sleep(3)
+print(sr.live.tune.readback())

pyaml/control/abstract_impl.py

@@ -151,12 +151,18 @@ def set(self, value: NDArray[np.float64]):
         raise Exception("BPM are not writable")
 
     def get(self) -> NDArray[np.float64]:
-        # TODO read using DeviceAccessList
-        allValues = []
-        for i, d in enumerate(self._devs):
-            v = d.get()
-            allValues.extend(v[self._indices[i]])
-        return np.array(allValues)
+        if len(self._devs) == 1:
+            v = self._devs[0].get()
+            return v[self._indices[0]]
+        else:
+            # TODO read using DeviceAccessList
+            v0 = self._devs[0].get()[self._indices[0]]
+            v1 = self._devs[1].get()[self._indices[1]]
+            # Interleave
+            xy = np.zeros(v0.size + v1.size)
+            xy[0::2] = v0
+            xy[1::2] = v1
+            return xy
 
     def readback(self) -> np.array:
         return self.get()

pyaml/external/pySC_interface.py

@@ -1,3 +1,4 @@
+import time
 from typing import TYPE_CHECKING, Optional, Tuple
 
 import numpy as np
@@ -8,6 +9,8 @@
 
 
 class pySCInterface:
+    set_wait_time: float = 0
+
     def __init__(
         self,
         element_holder: "ElementHolder",
@@ -36,6 +39,7 @@ def get(self, name: str) -> float:
     def set(self, name: str, value: float) -> None:
         magnet = self.element_holder.get_magnet(name=name)
         magnet.strength.set(value=value)  # ideally set_and_wait but not implemented
+        time.sleep(self.set_wait_time)
         return
 
     def get_rf_main_frequency(self) -> float:

pyaml/tuning_tools/dispersion.py

@@ -44,12 +44,13 @@ def __init__(self, element_holder: ElementHolder, cfg: ConfigModel):
         self.frequency_delta = cfg.frequency_delta
         self.latest_measurement = None
 
-    def measure(self):
+    def measure(self, set_waiting_time: float = 0):
         interface = pySCInterface(
             element_holder=self.element_holder,
             bpm_array_name=self.bpm_array_name,
             rf_plant_name=self.rf_plant_name,
         )
+        interface.set_wait_time = set_waiting_time
 
         generator = measure_dispersion(
             interface=interface,

pyaml/tuning_tools/orbit_response_matrix.py

@@ -51,11 +51,14 @@ def __init__(self, element_holder: ElementHolder, cfg: ConfigModel):
         self.corrector_delta = cfg.corrector_delta
         self.latest_measurement = None
 
-    def measure(self, corrector_names: Optional[List[str]] = None):
+    def measure(
+        self, corrector_names: Optional[List[str]] = None, set_wait_time: float = 0
+    ):
         interface = pySCInterface(
             element_holder=self.element_holder,
             bpm_array_name=self.bpm_array_name,
         )
+        interface.set_wait_time = set_wait_time
 
         if corrector_names is None:
             logger.info(
@@ -69,6 +72,19 @@ def measure(self, corrector_names: Optional[List[str]] = None):
                 self.vcorr_array_name
             ).names()
             corrector_names = hcorrector_names + vcorrector_names
+        else:
+            all_hcorrector_names = self.element_holder.get_magnets(
+                self.hcorr_array_name
+            ).names()
+            all_vcorrector_names = self.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,
@@ -87,6 +103,9 @@ def measure(self, corrector_names: Optional[List[str]] = None):
             self.bpm_array_name
         ).names()
         self.latest_measurement = response_data.model_dump()
+        len_h = len(hcorrector_names)
+        len_v = len(vcorrector_names)
+        self.latest_measurement["inputs_plane"] = ["H"] * len_h + ["V"] * len_v
 
     def get(self):
         return self.latest_measurement

pyaml/tuning_tools/response_matrix.py

@@ -26,7 +26,8 @@ class ConfigModel(BaseModel):
     matrix: list[list[float]]
     input_names: Optional[list[str]]
     output_names: list[str]
-    rf_response: Optional[list[float]]
+    rf_response: Optional[list[float]] = None
+    inputs_plane: Optional[list[str]] = None
 
 
 class ResponseMatrix(object):