Improvement of MeasurementTool interface

pyaml/diagnostics/chromaticity_monitor.py

@@ -147,7 +147,16 @@ def measure(
         callback: callable = None,
     ):
         """
-        Main function for chromaticity measurment
+        Main function for chromaticity measurment.
+        :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement` contains:
+
+        .. code-block:: python
+
+            chromaticity:np.array # First order chromaticity, Array of [q'x,q'y]
+            dispersion:np.array # First order dispersion, [[dx'0,..,dx'n],[dy'0,..,dy'n]]
+            chromaticity_fit:np.array # Array of [[qx,qy],[q'x,q'y],[q''x,q''y],...]
+            dispersionx_fit:np.array # Array of [[dx0],..,[dxn],[dx'0],..,[dx'n],...]
+            dispersiony_fit:np.array # Array of [[dy0,..,dyn],[dy'0,..,dy'n],...]
 
         Parameters
         ----------
@@ -176,10 +185,22 @@ def measure(
         fit_dispersion : bool, optionnal
             Fit dispersion, [default: from config]
         do_plot : bool
-            Do you want to plot the fittinf results ?
+            Do you want to plot the fitting results ?
         callback: Callable, optional
             Callback is executed after each measurement or setting.
             If the callback return false, then the process is aborted.
+            callback_data dict contains:
+
+            .. code-block:: python
+
+              source:MeasurementTool # Tool that triggered the callback
+              step:int # The current step
+              avg_step:int # The current averaging step
+              rf:float # RF frequency used for the current step
+              tune:np.array # The measured tune (on Action.MEASURE)
+              orbit:np_array # The measured orbit, if fit_dispersion is True, (on Action.MEASURE)
+              dtune:np.array # The tune variation (on Action.RESTORE)
+
         """
         n_step = n_step if n_step is not None else self._cfg.n_step
         alphac = alphac if alphac is not None else self._alphac
@@ -198,7 +219,8 @@ def measure(
         if alphac is None:
             raise PyAMLException("Moment compaction factor is not defined")
 
-        self.register_callback(callback)
+        self._init_measure("")
+        self._register_callback(callback)
 
         # Get devices
         self.check_peer()
@@ -276,7 +298,8 @@ def measure(
 
     def fit(self, deltas, Q, order, orbit=None, fit_disp_order=None, do_plot=False):
         """
-        Compute chromaticity (and dispersion) from measurement data.
+        Compute chromaticity (and dispersion) from input data and update
+        :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement`.
 
         Parameters
         ----------
@@ -292,14 +315,8 @@ def fit(self, deltas, Q, order, orbit=None, fit_disp_order=None, do_plot=False):
         plot : bool, optional
             If True, plot the fit.
 
-        Returns
-        -------
-        dict
-            Dict with horizontal and veritical chromaticity and dispersion
-
         """
         chroma = np.polynomial.polynomial.polyfit(deltas, Q, order).T
-        self.latest_measurement = dict()
         self.latest_measurement["chromaticity_fit"] = chroma
         self.latest_measurement["chromaticity"] = chroma[:, 1]  # First order chroma
         if orbit is not None:

pyaml/tuning_tools/chromaticity_response_matrix.py

@@ -52,7 +52,14 @@ def measure(
         callback: Optional[Callable] = None,
     ):
         """
-        Measure tune response matrix
+        Measure chromaticity response matrix.
+        :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement` contains:
+
+        .. code-block:: python
+
+            matrix:list[list[float] # The response matrix
+            variable_names:list[str] # Variable names
+            observable_names:list[str] # Observables names
 
         **Example**
 
@@ -96,23 +103,25 @@ def callback(action: Action, data:dict):
             If the callback return false, then the scan is aborted and strength restored.
             callback_data dict contains:
 
-            .. code-block::
+            .. code-block:: python
 
-              source:str Object that triggered the source
-              step:int The current step
-              avg_step:int The current avg step
-              magnet:Magnet The magnet being excited
-              strength:Magnet strength
-              chroma:np.array The measured chroma (on Action.MEASURE)
-              dchroma:np.array The chroma variation (on Action.RESTORE)
+              source:MeasurementTool # Tool that triggered the callback
+              idx:int # The index in the element array being processed
+              step:int # The current step
+              avg_step:int # The current averaging step
+              magnet:str # The magnet being excited
+              strength:float # Magnet strength
+              chroma:np.array # The measured chroma (on Action.MEASURE)
+              dchroma:np.array # The chroma variation (on Action.RESTORE)
 
         """
         # Get devices
         self.check_peer()
         sextus = self._peer.get_magnets(self._cfg.sextu_array_name)
         cm = self._peer.get_chromaticity_monitor(self._cfg.chromaticity_name)
 
-        self.register_callback(callback)
+        self._register_callback(callback)
+        self._init_measure("pyaml.tuning_tools.response_matrix_data")
 
         chromamat = np.zeros((len(sextus), 2))
 
@@ -140,7 +149,9 @@ def callback(action: Action, data:dict):
                     # apply strength
                     m.strength.set(str + d)
 
-                    self.send_callback(Action.APPLY, {"step": qidx, "magnet": m.get_name(), "strength": float(str + d)})
+                    self.send_callback(
+                        Action.APPLY, {"idx": qidx, "step": step, "magnet": m.get_name(), "strength": float(str + d)}
+                    )
 
                     time.sleep(sleep_step)
 
@@ -153,7 +164,8 @@ def callback(action: Action, data:dict):
                         Qp[step] += chroma
 
                         self.send_callback(
-                            Action.MEASURE, {"step": qidx, "avg_step": avg, "magnet": m.get_name(), "chroma": chroma}
+                            Action.MEASURE,
+                            {"idx": qidx, "step": step, "avg_step": avg, "magnet": m.get_name(), "chroma": chroma},
                         )
 
                         if avg < nb_meas - 1:
@@ -171,16 +183,16 @@ def callback(action: Action, data:dict):
                 m.strength.set(str)
                 self.send_callback(
                     Action.RESTORE,
-                    {"step": qidx, "magnet": m.get_name(), "strength": float(str), "dchroma": chromamat[qidx]},
+                    {"idx": qidx, "magnet": m.get_name(), "strength": float(str), "dchroma": chromamat[qidx]},
                 )
 
         except Exception as ex:
             err = ex
         except KeyboardInterrupt as ex:
             aborted = True
         finally:
-            # Restore
-            m.strength.set(str)  # restore strength
+            # Restore strength
+            m.strength.set(str)
             self.send_callback(
                 Action.RESTORE,
                 {"step": qidx, "magnet": m.get_name(), "strength": float(str), "dchroma": chromamat[qidx]},
@@ -194,11 +206,11 @@ def callback(action: Action, data:dict):
             logger.warning(f"{self.get_name()} : measurement aborted")
             return False
 
-        self.latest_measurement = ResponseMatrixDataConfigModel(
+        mat = ResponseMatrixDataConfigModel(
             matrix=chromamat.T.tolist(),
             variable_names=sextus.names(),
             observable_names=[cm.get_name() + ".x", cm.get_name() + ".y"],
-        ).model_dump()
-        self.latest_measurement["type"] = "pyaml.tuning_tools.response_matrix_data"
+        )
+        self.latest_measurement.update(mat.model_dump())
 
         return True

pyaml/tuning_tools/dispersion.py

@@ -65,7 +65,8 @@ def measure(
             skip_save=True,
         )
 
-        self.register_callback(callback)
+        self._register_callback(callback)
+        self._init_measure()
 
         aborted = False
         for code, measurement in generator:
@@ -93,7 +94,7 @@ def measure(
         # dispersion_data.output_names = self.element_holder.get_bpms(
         #     self.bpm_array_name
         # ).names()
-        self.latest_measurement = dispersion_data.model_dump()
+        self.latest_measurement.update(dispersion_data.model_dump())
 
     def get(self):
         return self.latest_measurement

pyaml/tuning_tools/measurement_tool.py

@@ -50,10 +50,18 @@ class MeasurementTool(Element, metaclass=ABCMeta):
 
     def __init__(self, name):
         super().__init__(name)
-        self.latest_measurement: dict = None
+        self._latest_measurement: dict = None
+        """
+        """
         self._peer: "ElementHolder" = None  # Peer: ControlSystem or Simulator
         self._callback: Callable = None
 
+    def _init_measure(self, measurement_type: str | None = None):
+        # Initialize measurement data
+        self._latest_measurement = {}
+        if measurement_type is not None:
+            self._latest_measurement["type"] = measurement_type
+
     @abstractmethod
     def measure(self) -> bool:
         """
@@ -65,28 +73,30 @@ def measure(self) -> bool:
         """
         raise NotImplementedError()
 
-    def get(self) -> dict:
+    @property
+    def latest_measurement(self) -> dict:
         """
-        Return last measurement data
+        Return last measurement data, a dictionary containing last measurement data.
+        See sub class of MeasurementTool to get description.
 
         Returns
         -------
-        ResponseMatrixData
+        dict
             Return latest measurement or None
         """
-        return self.latest_measurement
+        return self._latest_measurement
 
-    # TODO: Abstract this method
-    def load(self, load_path: Path):
+    def get(self) -> dict:
         """
-        Load measurement data
+        Return last measurement data, a dictionary containing last measurement data.
+        See sub class of MeasurementTool to get description.
 
-        Parameters
-        ----------
-        load_path: Path
-            Matrix filename
+        Returns
+        -------
+        dict
+            Return latest measurement or None
         """
-        raise NotImplementedError()
+        return self._latest_measurement
 
     def save(self, save_path: Path, with_type: str = "json"):
         """
@@ -142,16 +152,16 @@ def callback(action: Action, data: dict):
         """
         ok = True
         if self._callback is not None:
-            # Add source
+            # Add source and peer
             cb_data["mode"] = f"{self.get_peer()}"
-            cb_data["source_name"] = f"{self.get_name()}"
+            cb_data["source"] = self
             ok = self._callback(action, cb_data)
         if not ok and raiseException:
             # Abort, same as ctrl+C
             raise KeyboardInterrupt
         return ok
 
-    def register_callback(self, callback: Callable):
+    def _register_callback(self, callback: Callable):
         self._callback = callback
 
     def attach(self, peer: "ElementHolder") -> Self:

pyaml/tuning_tools/orbit_response_matrix.py

@@ -123,9 +123,10 @@ def measure(
         pySC.disable_pySC_rich()
         aborted = False
         err = None
-        step = 0
+        idx = 0
         try:
-            self.register_callback(callback)
+            self._register_callback(callback)
+            self._init_measure()
             for code, measurement in generator:
                 callback_data = measurement.response_data  # to be defined better
                 if code is ResponseCode.AFTER_SET:
@@ -135,7 +136,7 @@ def measure(
                 elif code is ResponseCode.AFTER_RESTORE:
                     logger.info(f"Measured response of {measurement.last_input}.")
                     self.send_callback(Action.RESTORE, callback_data)
-                step += 1
+                idx += 1
         except Exception as ex:
             err = ex
         except KeyboardInterrupt as ex:
@@ -145,7 +146,7 @@ def measure(
             # TODO
             self.send_callback(
                 Action.RESTORE,
-                {"step": step},
+                {"idx": idx},
                 raiseException=False,
             )
 
@@ -157,7 +158,7 @@ def measure(
             return False
 
         orm_data = self._pySC_response_data_to_ORMData(measurement.response_data.model_dump())
-        self.latest_measurement = orm_data.model_dump()
+        self.latest_measurement.update(orm_data.model_dump())
 
     def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataConfigModel:
         # all metadata is discarded here. Should we keep something?

pyaml/tuning_tools/response_matrix_data.py

@@ -26,8 +26,6 @@ class ConfigModel(BaseModel):
         Observable names, basically the measurements
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
-
     matrix: list[list[float]]
     variable_names: Optional[list[str]]
     observable_names: list[str]