Type checking in Element and add get/readback in tuning_tools.Tune

Author: JeanLucPons

pyaml/common/element.py

@@ -1,7 +1,12 @@
+from typing import TYPE_CHECKING
+
 from pydantic import BaseModel, ConfigDict
 
 from .exception import PyAMLException
 
+if TYPE_CHECKING:
+    from ..common.element_holder import ElementHolder
+
 
 def __pyaml_repr__(obj):
     """
@@ -39,7 +44,7 @@ class Element(object):
 
     def __init__(self, name: str):
         self.__name: str = name
-        self._peer = None  # Peer: ControlSystem, Simulator
+        self._peer: "ElementHolder" = None  # Peer: ControlSystem, Simulator
 
     def get_name(self):
         """

pyaml/tuning_tools/tune.py

@@ -57,7 +57,7 @@ def measure(self, callback: Callable | None = None):
         """
         quads = self.quads()  # Returns attached quad devices
         tunemat = np.zeros((len(quads), 2))
-        initial_tune = self.__parent.get()
+        initial_tune = self.__parent.readback()
         delta = self.__parent._cfg.delta  # TODO: handle delta array
         aborted = False
 
@@ -70,7 +70,7 @@ def measure(self, callback: Callable | None = None):
                 aborted = True
                 break
 
-            tune = self.__parent.get()
+            tune = self.__parent.readback()
             dq = tune - initial_tune
             tunemat[idx] = dq / delta
 
@@ -166,10 +166,17 @@ def __init__(self, cfg: ConfigModel):
         self._cfg = cfg
         self.__tm = None
         self.__tr: TuneResponse = None
+        self.__setpoint = np.array([np.nan, np.nan])
 
     def get(self):
         """
-        Return the betatron tune
+        Return the betatron tune setpoint
+        """
+        return self.__setpoint
+
+    def readback(self):
+        """
+        Return the betatron tune measurement
         """
         self.check_peer()
         if not self.__tm:
@@ -190,8 +197,9 @@ def set(self, tune: np.array, iter: int = 1, wait_time: float = 0.0):
         wait_time: float
             Time to wait in second between 2 iterations
         """
+        self.__setpoint = tune
         for _ in range(iter):
-            diff_tune = tune - self.get()
+            diff_tune = tune - self.readback()
             str = self.__tr.quads().strengths.get()
             str += self.__tr.correct(diff_tune)
             self.__tr.quads().strengths.set(str)

tests/test_tuning_tools.py

@@ -21,7 +21,7 @@ def test_tuning_tools():
     tune_adjust.response.save_json("tunemat.json")
     tune_adjust.response.load_json("tunemat.json")
     tune_adjust.set([0.17, 0.32], iter=2)
-    tune = tune_adjust.get()
+    tune = tune_adjust.readback()
     assert np.abs(tune[0] - 0.17) < 1e-5
     assert np.abs(tune[1] - 0.32) < 1e-5