Monitors and actuators acces from tuning tools

pyaml/tuning_tools/bba.py

@@ -7,7 +7,7 @@
 """
 
 import logging
-from typing import Callable, Optional
+from typing import TYPE_CHECKING, Callable, Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -22,6 +22,11 @@
 from ..validation import DynamicValidation, register_schema
 from .measurement_tool import MeasurementTool
 
+if TYPE_CHECKING:
+    from ..arrays.bpm_array import BPMArray
+    from ..bpm.bpm import BPM
+    from ..magnet.magnet import Magnet
+
 logger = logging.getLogger(__name__)
 
 PYAMLCLASS = "BBA"
@@ -88,6 +93,19 @@ class BBA(MeasurementTool, DynamicValidation):
         Return the uncertainty of the vertical center offset.
     plot_data(plane)
         Plot BBA data.
+
+    Attributes
+    ----------
+    bpms : BPMArray
+        BPM array used for the measurement.
+    bpm : BPM
+        Reference BPM used for the alignment.
+    hcorrector : Magnet
+        Horizontal corrector used for the alignment.
+    vcorrector : Magnet
+        Vertical corrector used for the alignment.
+    quadrupole : Magnet
+        Quadrupole being aligned.
     """
 
     def __init__(
@@ -125,6 +143,36 @@ def __init__(
         self.n_avg_meas = n_avg_meas
         self.sleep_between_meas = sleep_between_meas
 
+    @property
+    def bpms(self) -> "BPMArray":
+        """Return the BPM array used for the measurement."""
+        self.check_peer()
+        return self.peer.bpms.get(self.bpm_array_name)
+
+    @property
+    def bpm(self) -> "BPM":
+        """Return the reference BPM used for the alignment."""
+        self.check_peer()
+        return self.peer.bpm.get(self.bpm_name)
+
+    @property
+    def hcorrector(self) -> "Magnet":
+        """Return the horizontal corrector used for the alignment."""
+        self.check_peer()
+        return self.peer.magnet.get(self.hcorr_name)
+
+    @property
+    def vcorrector(self) -> "Magnet":
+        """Return the vertical corrector used for the alignment."""
+        self.check_peer()
+        return self.peer.magnet.get(self.vcorr_name)
+
+    @property
+    def quadrupole(self) -> "Magnet":
+        """Return the quadrupole being aligned."""
+        self.check_peer()
+        return self.peer.magnet.get(self.quad_name)
+
     def measure(
         self,
         sleep_between_step: Optional[float] = None,
@@ -183,7 +231,7 @@ def measure(
         interface.set_wait_time = sleep_step
         interface.read_wait_time = sleep_meas
 
-        bpms_names = element_holder.bpms.get(self.bpm_array_name).names()
+        bpms_names = self.bpms.names()
 
         bba_pySC_config = {
             "number": bpms_names.index(self.bpm_name),

pyaml/tuning_tools/bba2.py

@@ -7,7 +7,7 @@
 
 import logging
 import time
-from typing import Callable, Optional
+from typing import TYPE_CHECKING, Callable, Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -17,6 +17,12 @@
 from ..validation import DynamicValidation, register_schema
 from .measurement_tool import MeasurementTool
 
+if TYPE_CHECKING:
+    from ..arrays.bpm_array import BPMArray
+    from ..bpm.bpm import BPM
+    from ..magnet.magnet import Magnet
+    from .tune import Tune
+
 logger = logging.getLogger(__name__)
 
 PYAMLCLASS = "BBA2"
@@ -162,6 +168,21 @@ class BBA2(MeasurementTool, DynamicValidation):
         Plot measured kicks and the fitted alignment response for one plane.
     plot_data()
         Plot BBA data.
+
+    Attributes
+    ----------
+    bpms
+        Return the BPM array used for the measurement.
+    bpm
+        Return the reference BPM used for the alignment.
+    hcorrector
+        Return the horizontal corrector used for the alignment.
+    vcorrector
+        Return the vertical corrector used for the alignment.
+    quadrupole
+        Return the quadrupole being aligned.
+    tune_correction
+        Return the tune-correction tool used during alignment.
     """
 
     def __init__(
@@ -203,6 +224,42 @@ def __init__(
         self.n_avg_meas = n_avg_meas
         self.sleep_between_meas = sleep_between_meas
 
+    @property
+    def bpms(self) -> "BPMArray":
+        """Return the BPM array used for the measurement."""
+        self.check_peer()
+        return self.peer.bpms.get(self.bpm_array_name)
+
+    @property
+    def bpm(self) -> "BPM":
+        """Return the reference BPM used for the alignment."""
+        self.check_peer()
+        return self.peer.bpm.get(self.bpm_name)
+
+    @property
+    def hcorrector(self) -> "Magnet":
+        """Return the horizontal corrector used for the alignment."""
+        self.check_peer()
+        return self.peer.magnet.get(self.hcorr_name)
+
+    @property
+    def vcorrector(self) -> "Magnet":
+        """Return the vertical corrector used for the alignment."""
+        self.check_peer()
+        return self.peer.magnet.get(self.vcorr_name)
+
+    @property
+    def quadrupole(self) -> "Magnet":
+        """Return the quadrupole being aligned."""
+        self.check_peer()
+        return self.peer.magnet.get(self.quad_name)
+
+    @property
+    def tune_correction(self) -> "Tune":
+        """Return the tune-correction tool used during alignment."""
+        self.check_peer()
+        return self.peer.get_tune_tuning(self.tune_correction_name)
+
     @staticmethod
     def _x_intercept(x, k, n):
         # Linear fit on last n points
@@ -549,10 +606,10 @@ def measure(
 
         # Device handles
         self.check_peer()
-        self._h_steer = self.peer.magnet.get(self.hcorr_name)
-        self._v_steer = self.peer.magnet.get(self.vcorr_name)
-        self._quad = self.peer.magnet.get(self.quad_name)
-        self._bpms = self.peer.bpms.get(self.bpm_array_name)
+        self._h_steer = self.hcorrector
+        self._v_steer = self.vcorrector
+        self._quad = self.quadrupole
+        self._bpms = self.bpms
         self._bpmi = self._bpms.names().index(self.bpm_name)
 
         # Initial values

pyaml/tuning_tools/chromaticity.py

@@ -53,6 +53,10 @@ class Chromaticity(TuningTool, DynamicValidation):
     ----------
     response_matrix
         Return the loaded chromaticity response matrix, if available.
+    chromaticity_monitor
+        Return the chromaticity monitor used for readback.
+    sextupoles
+        Return the sextupole array used for correction.
 
     Methods
     -------
@@ -118,14 +122,14 @@ def load(self, load_path: Path):
         self._correctionmat = np.linalg.pinv(self._response_matrix)
 
     @property
-    def _cm(self) -> "ChromaticityMonitor":
-        """Return the chromaticity monitor."""
+    def chromaticity_monitor(self) -> "ChromaticityMonitor":
+        """Return the chromaticity monitor used for readback."""
         self.check_peer()
         return self.peer.get_chromaticity_monitor(self._chromaticity_monitor_name)
 
     @property
-    def _sextu(self) -> "MagnetArray":
-        """Return the sextupole array."""
+    def sextupoles(self) -> "MagnetArray":
+        """Return the sextupole array used for correction."""
         self.check_peer()
         return self.peer.magnets.get(self.sextu_array_name)
 
@@ -135,8 +139,8 @@ def get(self):
 
     def readback(self):
         """Measure and return the current dimensionless horizontal and vertical chromaticity."""
-        self._cm.measure()
-        return self._cm.chromaticity.get()
+        self.chromaticity_monitor.measure()
+        return self.chromaticity_monitor.chromaticity.get()
 
     def set(self, chroma: np.array, iter: int = 1, wait_time: float = 0.0):
         """
@@ -194,8 +198,8 @@ def add(self, dchroma: np.array, wait_time: float = 0.0):
         wait_time : float
             Delay in seconds after changing sextupole strengths.
         """
-        strengths = self._sextu.strengths.get()
+        strengths = self.sextupoles.strengths.get()
         strengths += self.correct(dchroma)
-        self._sextu.strengths.set(strengths)
+        self.sextupoles.strengths.set(strengths)
         time.sleep(wait_time)
         self._setpoint += dchroma

pyaml/tuning_tools/chromaticity_monitor.py

@@ -9,6 +9,7 @@
 import logging
 from collections.abc import Callable
 from time import sleep
+from typing import TYPE_CHECKING
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -20,6 +21,11 @@
 from ..tuning_tools.measurement_tool import MeasurementTool
 from ..validation import DynamicValidation, register_schema
 
+if TYPE_CHECKING:
+    from ..arrays.bpm_array import BPMArray
+    from ..diagnostics.tune_monitor import BetatronTuneMonitor
+    from ..rf.rf_plant import RFPlant
+
 logger = logging.getLogger(__name__)
 
 PYAMLCLASS = "ChromaticityMonitor"
@@ -116,6 +122,12 @@ class ChromaticityMonitor(MeasurementTool, DynamicValidation):
         Get the chromaticity values.
     dispersion
         Get the dispersion values.
+    tune_monitor
+        Return the betatron tune monitor used for measurements.
+    rf_plant
+        Return the RF plant varied during measurements.
+    bpms
+        Return the optional BPM array used for dispersion measurements.
 
     Methods
     -------
@@ -173,6 +185,26 @@ def __init__(
         self._dispersion = RChromaDispArray(self, "dispersion", "m")
         self._alphac = None
 
+    @property
+    def tune_monitor(self) -> "BetatronTuneMonitor":
+        """Return the betatron tune monitor used for measurements."""
+        self.check_peer()
+        return self.peer.get_betatron_tune_monitor(self.betatron_tune_name)
+
+    @property
+    def rf_plant(self) -> "RFPlant":
+        """Return the RF plant varied during measurements."""
+        self.check_peer()
+        return self.peer.rf.get(self.rf_plant_name)
+
+    @property
+    def bpms(self) -> "BPMArray | None":
+        """Return the optional BPM array used for dispersion measurements."""
+        self.check_peer()
+        if self.bpm_array_name is None:
+            return None
+        return self.peer.bpms.get(self.bpm_array_name)
+
     @property
     def chromaticity(self) -> ReadFloatArray:
         """
@@ -306,14 +338,14 @@ def measure(
 
         # Get devices
         self.check_peer()
-        tm = self.peer.get_betatron_tune_monitor(self.betatron_tune_name)
-        rf = self.peer.rf.get(self.rf_plant_name)
+        tm = self.tune_monitor
+        rf = self.rf_plant
         bpms = None
         n_bpm = 0
         orbit = None
         if fit_dispersion and fit_disp_order is not None and self.bpm_array_name is not None:
             # For dispersion fit
-            bpms = self.peer.bpms.get(self.bpm_array_name)
+            bpms = self.bpms
             n_bpm = len(bpms)
 
         f0 = rf.frequency.get()

pyaml/tuning_tools/chromaticity_response_matrix.py

@@ -9,7 +9,7 @@
 import logging
 import time
 from dataclasses import asdict
-from typing import Callable, Optional
+from typing import TYPE_CHECKING, Callable, Optional
 
 import numpy as np
 
@@ -18,6 +18,10 @@
 from .measurement_tool import MeasurementTool
 from .response_matrix_data import ResponseMatrixData
 
+if TYPE_CHECKING:
+    from ..arrays.magnet_array import MagnetArray
+    from .chromaticity_monitor import ChromaticityMonitor
+
 logger = logging.getLogger(__name__)
 
 PYAMLCLASS = "ChromaticityResponseMatrix"
@@ -57,6 +61,13 @@ class ChromaticityResponseMatrix(MeasurementTool, DynamicValidation):
     -------
     measure(...)
         Measure the chromaticity response matrix.
+
+    Attributes
+    ----------
+    sextupoles : MagnetArray
+        Sextupole array used for the measurement.
+    chromaticity_monitor : ChromaticityMonitor
+        Chromaticity monitor used for the measurement.
     """
 
     def __init__(
@@ -84,6 +95,18 @@ def __init__(
         self.sleep_between_meas = sleep_between_meas
         self.aborted = False
 
+    @property
+    def sextupoles(self) -> "MagnetArray":
+        """Return the sextupole array used for the measurement."""
+        self.check_peer()
+        return self.peer.magnets.get(self.sextu_array_name)
+
+    @property
+    def chromaticity_monitor(self) -> "ChromaticityMonitor":
+        """Return the chromaticity monitor used for the measurement."""
+        self.check_peer()
+        return self.peer.get_chromaticity_monitor(self.chromaticity_name)
+
     def measure(
         self,
         sextu_delta: Optional[float] = None,
@@ -165,8 +188,8 @@ def callback(action: Action, data:dict):
         """
         # Get devices
         self.check_peer()
-        sextus = self._peer.magnets.get(self.sextu_array_name)
-        cm = self._peer.get_chromaticity_monitor(self.chromaticity_name)
+        sextus = self.sextupoles
+        cm = self.chromaticity_monitor
 
         self._register_callback(callback)
         self._init_measure("pyaml.tuning_tools.response_matrix_data")