Merge pull request #335 from python-accelerator-middle-layer/schema-registry-tuning-tools

Changes for schema registry for tuning tools

Author: gubaidulinvadim

pyaml/common/holders/element_holder.py

@@ -19,7 +19,7 @@
 from ...magnet.serialized_magnet import SerializedMagnets
 from ...rf.rf_plant import RFPlant
 from ...rf.rf_transmitter import RFTransmitter
-from ...tuning_tools.chromaticity_monitor import ChomaticityMonitor
+from ...tuning_tools.chromaticity_monitor import ChromaticityMonitor
 from ..abstract_aggregator import ScalarAggregator
 from ..element import Element
 from ..exception import PyAMLException
@@ -254,8 +254,8 @@ def add_tool(self, tool: Element):
 
     # ---- Chromaticity -------------------------------------------------
 
-    def get_chromaticity_monitor(self, name: str) -> ChomaticityMonitor:
-        obj = self._get("Chomaticity monitor", name, self._TOOLS)
+    def get_chromaticity_monitor(self, name: str) -> ChromaticityMonitor:
+        obj = self._get("Chromaticity monitor", name, self._TOOLS)
         return obj
 
     def get_chromaticity_tuning(self, name: str) -> "Chromaticity":

pyaml/tuning_tools/chromaticity.py

@@ -3,7 +3,7 @@
 
 from .. import PyAMLException
 from ..validation import DynamicValidation, register_schema
-from .chromaticity_monitor import ChomaticityMonitor
+from .chromaticity_monitor import ChromaticityMonitor
 from .response_matrix_data import ResponseMatrixData
 from .tuning_tool import TuningTool
 
@@ -61,7 +61,7 @@ def __init__(
 
         # Invert matrix
         if self._response_matrix:
-            self._response_matrix = np.array(self._response_matrix._cfg.matrix)
+            self._response_matrix = np.array(self._response_matrix.matrix)
             self._correctionmat = np.linalg.pinv(self._response_matrix)
 
         # TODO: Initialise first setpoint
@@ -84,11 +84,11 @@ def load(self, load_path: Path):
             Filename of the :class:`~.ResponseMatrixData` to load
         """
         self._response_matrix = ResponseMatrixData.load(load_path)
-        self._response_matrix = np.array(self._response_matrix._cfg.matrix)
+        self._response_matrix = np.array(self._response_matrix.matrix)
         self._correctionmat = np.linalg.pinv(self._response_matrix)
 
     @property
-    def _cm(self) -> "ChomaticityMonitor":
+    def _cm(self) -> "ChromaticityMonitor":
         self.check_peer()
         return self.peer.get_chromaticity_monitor(self._chromaticity_monitor_name)
 

pyaml/tuning_tools/chromaticity_monitor.py

@@ -14,7 +14,7 @@
 
 logger = logging.getLogger(__name__)
 
-PYAMLCLASS = "ChomaticityMonitor"
+PYAMLCLASS = "ChromaticityMonitor"
 
 
 class RChromaDispArray(ReadFloatArray):
@@ -23,7 +23,7 @@ class RChromaDispArray(ReadFloatArray):
     Returns arrays of shape (fit_order,2) or None
     """
 
-    def __init__(self, parent: "ChomaticityMonitor", name: str, unit: str):
+    def __init__(self, parent: "ChromaticityMonitor", name: str, unit: str):
         self._parent = parent
         self._name = name
         self._unit = unit
@@ -40,7 +40,7 @@ def unit(self) -> str:
 
 
 @register_schema
-class ChomaticityMonitor(MeasurementTool, DynamicValidation):
+class ChromaticityMonitor(MeasurementTool, DynamicValidation):
     """
     Class providing access to a chromaticity monitor
     of a physical or simulated lattice. The monitor provides

pyaml/tuning_tools/chromaticity_response_matrix.py

@@ -1,13 +1,14 @@
 import logging
 import time
+from dataclasses import asdict
 from typing import Callable, Optional
 
 import numpy as np
 
 from ..common.constants import Action
 from ..validation import DynamicValidation, register_schema
 from .measurement_tool import MeasurementTool
-from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel
+from .response_matrix_data import ResponseMatrixData
 
 logger = logging.getLogger(__name__)
 
@@ -252,12 +253,12 @@ def callback(action: Action, data:dict):
             logger.warning(f"{self.get_name()} : measurement aborted")
             return False
 
-        mat = ResponseMatrixDataConfigModel(
+        mat = ResponseMatrixData(
             matrix=chromamat.T.tolist(),
             variable_names=sextus.names(),
             observable_names=[cm.get_name() + ".x", cm.get_name() + ".y"],
         )
-        self.latest_measurement.update(mat.model_dump())
+        self.latest_measurement.update(asdict(mat))
         self.latest_measurement["type"] = "pyaml.tuning_tools.response_matrix_data"
 
         return True

pyaml/tuning_tools/dispersion.py

@@ -1,50 +1,54 @@
 import logging
-from typing import Callable, Optional, Self
+from typing import Callable, Optional
 
-from pydantic import ConfigDict
 from pySC.apps import measure_dispersion
 from pySC.apps.codes import DispersionCode
 
 from ..common.constants import Action
-from ..common.element import ElementConfigModel
-from ..common.holders.element_holder import ElementHolder
 from ..external.pySC_interface import pySCInterface
+from ..validation import DynamicValidation, register_schema
 from .measurement_tool import MeasurementTool
 
 logger = logging.getLogger(__name__)
 
 PYAMLCLASS = "Dispersion"
 
 
-class ConfigModel(ElementConfigModel):
-    """
-    Configuration model for dispersion measurement
+@register_schema
+class Dispersion(MeasurementTool, DynamicValidation):
+    """Measure beam dispersion by changing the RF frequency.
+
+    The measurement uses a :class:`pySCInterface` to change the frequency of
+    an RF plant and acquire orbit data from a BPM array. Progress is reported
+    through the callback mechanism provided by :class:`MeasurementTool`.
 
     Parameters
     ----------
+    name : str
+        Name of the dispersion measurement tool.
     bpm_array_name : str
-        BPM array name
+        Name of the BPM array used to measure the orbit.
     rf_plant_name : str
-        RF plant name
+        Name of the RF plant whose frequency is varied.
     frequency_delta : float
-        Frequency delta for measurement
-    """
-
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
-
-    bpm_array_name: str
-    rf_plant_name: str
-    frequency_delta: float
+        RF-frequency change applied during the measurement.
 
+    Attributes
+    ----------
+    bpm_array_name : str
+        Name of the BPM array used for the measurement.
+    rf_plant_name : str
+        Name of the RF plant used for the measurement.
+    frequency_delta : float
+        RF-frequency change applied during the measurement.
+    """
 
-class Dispersion(MeasurementTool):
-    def __init__(self, cfg: ConfigModel):
-        super().__init__(cfg.name)
-        self._cfg = cfg
+    def __init__(self, name: str, bpm_array_name: str, rf_plant_name: str, frequency_delta: float):
+        super().__init__(name)
 
-        self.bpm_array_name = cfg.bpm_array_name
-        self.rf_plant_name = cfg.rf_plant_name
-        self.frequency_delta = cfg.frequency_delta
+        self.bpm_array_name = bpm_array_name
+        self.rf_plant_name = rf_plant_name
+        self.frequency_delta = frequency_delta
 
     def measure(
         self,

pyaml/tuning_tools/measurement_tool.py

@@ -17,31 +17,12 @@
 
 
 class MeasurementToolConfigModel(ElementConfigModel):
-    """
-    Measurement tool configuration model
-
-    Parameters
-    ----------
-    n_step: int, optional
-        Number of measurement step [-delta/n_step..delta/n_step]
-        Default 1
-    sleep_between_step: float, optional
-        Default sleep time after an actuator excitation
-        Default: 0
-    n_avg_meas : int, optional
-        Default number of measurement per step used for averaging
-        Default 1
-    sleep_between_meas: float, optional
-        Default sleep time between two measurments
-        Default: 0
-    """
-
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
 
-    n_step: Optional[int] = 1
-    sleep_between_step: Optional[float] = 0
-    n_avg_meas: Optional[int] = 1
-    sleep_between_meas: Optional[float] = 0
+    n_step: int = 10
+    sleep_between_step: float = 0
+    n_avg_meas: int = 1
+    sleep_between_meas: float = 0
 
 
 class MeasurementTool(Element, metaclass=ABCMeta):

pyaml/tuning_tools/orbit.py

@@ -1,25 +1,17 @@
 import logging
+from dataclasses import asdict
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Optional, Union
-
-try:
-    from typing import Self  # Python 3.11+
-except ImportError:
-    from typing_extensions import Self  # Python 3.10 and earlier
+from typing import Literal, Optional, Union
 
 import numpy as np
-from pydantic import ConfigDict
-
-if TYPE_CHECKING:
-    from ..common.holders.element_holder import ElementHolder
 from pySC import ResponseMatrix as pySC_ResponseMatrix
 from pySC.apps import orbit_correction
 
 from ..arrays.magnet_array import MagnetArray
-from ..common.element import Element, ElementConfigModel
 from ..common.exception import PyAMLException
 from ..external.pySC_interface import pySCInterface
 from ..rf.rf_plant import RFPlant
+from ..validation import DynamicValidation, register_schema
 from .orbit_response_matrix_data import OrbitResponseMatrixData
 from .tuning_tool import TuningTool
 
@@ -29,57 +21,56 @@
 PYAMLCLASS = "Orbit"
 
 
-class ConfigModel(ElementConfigModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
-
-    bpm_array_name: str
-    hcorr_array_name: str
-    vcorr_array_name: str
-    rf_plant_name: Optional[str] = None
-    singular_values: Optional[int] = None
-    singular_values_H: Optional[int] = None
-    singular_values_V: Optional[int] = None
-    virtual_target: float = 0
-    response_matrix: Union[str, OrbitResponseMatrixData]
-
+@register_schema
+class Orbit(TuningTool, DynamicValidation):
+    def __init__(
+        self,
+        name: str,
+        bpm_array_name: str,
+        hcorr_array_name: str,
+        vcorr_array_name: str,
+        response_matrix: Union[str, OrbitResponseMatrixData],
+        rf_plant_name: Optional[str] = None,
+        singular_values: Optional[int] = None,
+        singular_values_H: Optional[int] = None,
+        singular_values_V: Optional[int] = None,
+        virtual_target: float = 0,
+    ):
+        super().__init__(name)
 
-class Orbit(TuningTool):
-    def __init__(self, cfg: ConfigModel):
-        super().__init__(cfg.name)
-        self._cfg = cfg
-        self.bpm_array_name = cfg.bpm_array_name
-        self.hcorr_array_name = cfg.hcorr_array_name
-        self.vcorr_array_name = cfg.vcorr_array_name
+        self.bpm_array_name = bpm_array_name
+        self.hcorr_array_name = hcorr_array_name
+        self.vcorr_array_name = vcorr_array_name
         self._pySC_response_matrix = None
+        self.rf_plant_name = rf_plant_name
+        self.virtual_target = virtual_target
 
-        self.virtual_target = cfg.virtual_target
-
-        if cfg.singular_values is None:
-            if cfg.singular_values_H is None or cfg.singular_values_V is None:
+        if singular_values is None:
+            if singular_values_H is None or singular_values_V is None:
                 raise PyAMLException(
                     "Either `singular_values` or `singular_values_H` and `singular_values_V` must be provided."
                 )
-            self.singular_values_H = cfg.singular_values_H
-            self.singular_values_V = cfg.singular_values_V
+            self.singular_values_H = singular_values_H
+            self.singular_values_V = singular_values_V
         else:
-            if cfg.singular_values_H is not None or cfg.singular_values_V is not None:
+            if singular_values_H is not None or singular_values_V is not None:
                 raise PyAMLException(
                     "Either `singular_values` or `singular_values_H` and `singular_values_V` must be provided, not both."
                 )
-            self.singular_values_H = cfg.singular_values
-            self.singular_values_V = cfg.singular_values
+            self.singular_values_H = singular_values
+            self.singular_values_V = singular_values
 
         # If the configuration response matrix is a filename, load it
-        if type(cfg.response_matrix) is str:
+        if type(response_matrix) is str:
             try:
-                cfg.response_matrix = OrbitResponseMatrixData.load(cfg.response_matrix)
+                self._response_matrix = OrbitResponseMatrixData.load(response_matrix)
             except Exception as e:
-                logger.warning(f"Loading {cfg.response_matrix} failed {str(e)}")
-                cfg.response_matrix = None
+                logger.warning(f"Loading {response_matrix} failed {str(e)}")
+                self._response_matrix = None
 
         # Converts to self._pySC_response_matrix
-        if cfg.response_matrix:
-            self._set_response_matrix(cfg.response_matrix)
+        if self._response_matrix:
+            self._set_response_matrix(self._response_matrix)
 
         self._hcorr: MagnetArray = None
         self._vcorr: MagnetArray = None
@@ -95,24 +86,25 @@ def load(self, load_path: Path):
         load_path : Path
             Filename of the :class:`~.OrbitResponseMatrixData` to load
         """
-        self._cfg.response_matrix = OrbitResponseMatrixData.load(load_path)
-        self._set_response_matrix(self._cfg.response_matrix)
+        self._response_matrix = OrbitResponseMatrixData.load(load_path)
+        self._set_response_matrix(self.response_matrix)
 
     def _set_response_matrix(self, mat):
-        m = mat._cfg.model_dump()
+        m = asdict(mat)
         m["input_names"] = m.pop("variable_names")
         m["output_names"] = m.pop("observable_names")
         m["input_planes"] = m.pop("variable_planes")
         m["output_planes"] = m.pop("observable_planes")
-        self._cfg.response_matrix = mat
+        m.pop("type", None)
+        self._response_matrix = mat
         self._pySC_response_matrix = pySC_ResponseMatrix.model_validate(m)
 
     @property
     def response_matrix(self) -> OrbitResponseMatrixData | None:
         """
         Return the response matrix if it has been loaded None otherwise
         """
-        return self._cfg.response_matrix
+        return self._response_matrix
 
     def correct(
         self,
@@ -312,11 +304,11 @@ def get_rf_weight(self) -> float:
         return self._pySC_response_matrix.rf_weight
 
     def post_init(self):
-        self._hcorr = self.peer.magnets.get(self._cfg.hcorr_array_name)
-        self._vcorr = self.peer.magnets.get(self._cfg.vcorr_array_name)
+        self._hcorr = self.peer.magnets.get(self.hcorr_array_name)
+        self._vcorr = self.peer.magnets.get(self.vcorr_array_name)
         hvElts = []
         hvElts.extend(self._hcorr)
         hvElts.extend(self._vcorr)
         self._hvcorr = MagnetArray("", hvElts)
-        if self._cfg.rf_plant_name is not None:
-            self._rf_plant = self.peer.rf.get(self._cfg.rf_plant_name)
+        if self.rf_plant_name is not None:
+            self._rf_plant = self.peer.rf.get(self.rf_plant_name)