diff --git a/pyaml/tuning_tools/bba.py b/pyaml/tuning_tools/bba.py index 15b170ef..ed0b8cc2 100644 --- a/pyaml/tuning_tools/bba.py +++ b/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), diff --git a/pyaml/tuning_tools/bba2.py b/pyaml/tuning_tools/bba2.py index 7217ae0d..b1b7bb07 100644 --- a/pyaml/tuning_tools/bba2.py +++ b/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 diff --git a/pyaml/tuning_tools/chromaticity.py b/pyaml/tuning_tools/chromaticity.py index f4efd782..e1aed264 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/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 diff --git a/pyaml/tuning_tools/chromaticity_monitor.py b/pyaml/tuning_tools/chromaticity_monitor.py index 05da69a5..c1da2507 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/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() diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index e077c3d7..6746f403 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/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") diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index 6eb5dd84..6bbc3844 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/pyaml/tuning_tools/dispersion.py @@ -7,7 +7,7 @@ """ import logging -from typing import Callable, Optional +from typing import TYPE_CHECKING, Callable, Optional from pySC.apps import measure_dispersion from pySC.apps.codes import DispersionCode @@ -17,6 +17,10 @@ from ..validation import DynamicValidation, register_schema from .measurement_tool import MeasurementTool +if TYPE_CHECKING: + from ..arrays.bpm_array import BPMArray + from ..rf.rf_plant import RFPlant + logger = logging.getLogger(__name__) PYAMLCLASS = "Dispersion" @@ -41,6 +45,10 @@ class Dispersion(MeasurementTool, DynamicValidation): Name of the RF plant whose frequency is varied. frequency_delta : float RF-frequency change applied during the measurement, in hertz. + bpms : BPMArray + BPM array used for orbit readback. + rf_plant : RFPlant + RF plant varied during the measurement. Attributes ---------- @@ -69,6 +77,18 @@ def __init__(self, name: str, bpm_array_name: str, rf_plant_name: str, frequency self.rf_plant_name = rf_plant_name self.frequency_delta = frequency_delta + @property + def bpms(self) -> "BPMArray": + """Return the BPM array used for orbit readback.""" + self.check_peer() + return self.peer.bpms.get(self.bpm_array_name) + + @property + def rf_plant(self) -> "RFPlant": + """Return the RF plant varied during the measurement.""" + self.check_peer() + return self.peer.rf.get(self.rf_plant_name) + def measure( self, set_waiting_time: float = 0, diff --git a/pyaml/tuning_tools/orbit.py b/pyaml/tuning_tools/orbit.py index ecb8afc7..3f2c8b39 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -14,6 +14,7 @@ from pySC import ResponseMatrix as pySC_ResponseMatrix from pySC.apps import orbit_correction +from ..arrays.bpm_array import BPMArray from ..arrays.magnet_array import MagnetArray from ..common.exception import PyAMLException from ..external.pySC_interface import pySCInterface @@ -62,6 +63,16 @@ class Orbit(TuningTool, DynamicValidation): ---------- response_matrix Return the response matrix if it has been loaded None otherwise + bpms + Return the BPM array used for orbit readback. + hcorrectors + Return the horizontal corrector array used for correction. + vcorrectors + Return the vertical corrector array used for correction. + correctors + Return the combined horizontal and vertical corrector array. + rf_plant + Return the optional RF plant used for correction. Methods ------- @@ -184,6 +195,46 @@ def response_matrix(self) -> OrbitResponseMatrixData | None: """ return self._response_matrix + @property + def bpms(self) -> BPMArray: + """Return the BPM array used for orbit readback.""" + self.check_peer() + return self.peer.bpms.get(self.bpm_array_name) + + @property + def hcorrectors(self) -> MagnetArray: + """Return the horizontal corrector array used for correction.""" + self.check_peer() + if self._hcorr is None: + return self.peer.magnets.get(self.hcorr_array_name) + return self._hcorr + + @property + def vcorrectors(self) -> MagnetArray: + """Return the vertical corrector array used for correction.""" + self.check_peer() + if self._vcorr is None: + return self.peer.magnets.get(self.vcorr_array_name) + return self._vcorr + + @property + def correctors(self) -> MagnetArray: + """Return the combined horizontal and vertical corrector array.""" + self.check_peer() + if self._hvcorr is None: + return MagnetArray("", [*self.hcorrectors, *self.vcorrectors]) + return self._hvcorr + + @property + def rf_plant(self) -> RFPlant | None: + """Return the optional RF plant used for orbit correction.""" + self.check_peer() + if self.rf_plant_name is None: + return None + if self._rf_plant is None: + return self.peer.rf.get(self.rf_plant_name) + return self._rf_plant + def correct( self, plane: Optional[Literal["H", "V"]] = None, @@ -291,7 +342,7 @@ def correct( # take care of rf trim rf_flag = rf and (plane is None or plane == "H") if rf_flag: - if self._rf_plant is None: + if self.rf_plant is None: raise PyAMLException("RF plant is not defined!") eff_gain_rf = gain_rf if gain_rf is not None else eff_gain_h ## pySC returns with an 'rf' entry into the dictionary if rf=True @@ -305,17 +356,17 @@ def correct( for trim in trims_v: trims_v[trim] *= eff_gain_v trims = {**trims_h, **trims_v} - corr_array = self._hvcorr + corr_array = self.correctors elif plane == "H": for trim in trims_h: trims_h[trim] *= eff_gain_h trims = trims_h - corr_array = self._hcorr + corr_array = self.hcorrectors elif plane == "V": for trim in trims_v: trims_v[trim] *= eff_gain_v trims = trims_v - corr_array = self._vcorr + corr_array = self.vcorrectors corrector_names = corr_array.names() corrector_to_index = {name: idx for idx, name in enumerate(corrector_names)} @@ -333,8 +384,8 @@ def correct( # send trims corr_array.strengths.set(data_to_send) if rf_flag: - rf_frequency = self._rf_plant.frequency.get() - self._rf_plant.frequency.set(rf_frequency + rf_trim) + rf_frequency = self.rf_plant.frequency.get() + self.rf_plant.frequency.set(rf_frequency + rf_trim) return diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index cb332df7..7b065b49 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -8,7 +8,7 @@ import logging from dataclasses import asdict -from typing import Callable, List, Optional +from typing import TYPE_CHECKING, Callable, List, Optional import pySC from pySC.apps import measure_ORM @@ -20,6 +20,10 @@ from .measurement_tool import MeasurementTool from .orbit_response_matrix_data import OrbitResponseMatrixData +if TYPE_CHECKING: + from ..arrays.bpm_array import BPMArray + from ..arrays.magnet_array import MagnetArray + logger = logging.getLogger(__name__) PYAMLCLASS = "OrbitResponseMatrix" @@ -83,6 +87,12 @@ class OrbitResponseMatrix(MeasurementTool, DynamicValidation): Configured number of orbit measurements to average. sleep_between_meas : float Configured delay between averaged orbit measurements. + bpms : BPMArray + BPM array used for orbit readback. + hcorrectors : MagnetArray + Horizontal corrector array used for the measurement. + vcorrectors : MagnetArray + Vertical corrector array used for the measurement. Methods ------- @@ -116,6 +126,24 @@ 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 orbit readback.""" + self.check_peer() + return self.peer.bpms.get(self.bpm_array_name) + + @property + def hcorrectors(self) -> "MagnetArray": + """Return the horizontal corrector array used for the measurement.""" + self.check_peer() + return self.peer.magnets.get(self.hcorr_array_name) + + @property + def vcorrectors(self) -> "MagnetArray": + """Return the vertical corrector array used for the measurement.""" + self.check_peer() + return self.peer.magnets.get(self.vcorr_array_name) + def measure( self, corrector_names: Optional[List[str]] = None, @@ -173,8 +201,8 @@ def measure( if corrector_names is None: logger.info(f"Measuring correctors from the default arrays: {self.hcorr_array_name} and {self.vcorr_array_name}.") - hcorrector_names = element_holder.magnets.get(self.hcorr_array_name).names() - vcorrector_names = element_holder.magnets.get(self.vcorr_array_name).names() + hcorrector_names = self.hcorrectors.names() + vcorrector_names = self.vcorrectors.names() corrector_names = hcorrector_names + vcorrector_names generator = measure_ORM( @@ -252,8 +280,8 @@ def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixData: their associated planes. """ element_holder = self._peer - all_hcorrector_names = element_holder.magnets.get(self.hcorr_array_name).names() - all_vcorrector_names = element_holder.magnets.get(self.vcorr_array_name).names() + all_hcorrector_names = self.hcorrectors.names() + all_vcorrector_names = self.vcorrectors.names() variable_planes = [] for corr in data["input_names"]: if corr in all_hcorrector_names: @@ -261,7 +289,7 @@ def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixData: elif corr in all_vcorrector_names: variable_planes.append("V") - bpm_names = element_holder.bpms.get(self.bpm_array_name).names() + bpm_names = self.bpms.names() # This is because we assume always dual-plane bpms now. len_b = len(bpm_names) observable_names = bpm_names * 2 diff --git a/pyaml/tuning_tools/tune.py b/pyaml/tuning_tools/tune.py index d918b313..068caf13 100644 --- a/pyaml/tuning_tools/tune.py +++ b/pyaml/tuning_tools/tune.py @@ -27,6 +27,7 @@ # Define the main class name for this module PYAMLCLASS = "Tune" +DEFAULT_BETATRON_TUNE_MONITOR = "BETATRON_TUNE" @register_schema @@ -50,13 +51,13 @@ class Tune(TuningTool, DynamicValidation): Name of the tuning tool. quad_array_name : str Name of the quadrupole array used to adjust the tune. - betatron_tune_name : str - Name of the betatron tune monitor used to measure the horizontal and - vertical tunes. response_matrix : str or ResponseMatrixData Tune response matrix or path to a file containing the response matrix. The matrix is expected to have one row for each tune plane and one column for each quadrupole. + betatron_tune_name : str, optional + Name of the betatron tune monitor used to measure the horizontal and + vertical tunes. The default is ``"BETATRON_TUNE"``. Attributes ---------- @@ -64,6 +65,10 @@ class Tune(TuningTool, DynamicValidation): Name of the configured quadrupole array. betatron_tune_name : str Name of the configured betatron tune monitor. + quadrupoles : MagnetArray + Quadrupole array used by the correction. + tune_monitor : BetatronTuneMonitor + Betatron tune monitor used by the correction. response_matrix : ResponseMatrixData or None Loaded tune response matrix. @@ -87,8 +92,8 @@ def __init__( self, name: str, quad_array_name: str, - betatron_tune_name: str, response_matrix: str | ResponseMatrixData, + betatron_tune_name: str = DEFAULT_BETATRON_TUNE_MONITOR, ): """ Initialize a betatron-tune correction tool. @@ -135,14 +140,14 @@ def response_matrix(self) -> ResponseMatrixData | None: return self._response_matrix @property - def _tm(self) -> "BetatronTuneMonitor": - """Return the betatron tune monitor.""" + def tune_monitor(self) -> "BetatronTuneMonitor": + """Return the betatron tune monitor used by the correction.""" self.check_peer() return self.peer.get_betatron_tune_monitor(self.betatron_tune_name) @property - def _quads(self) -> "MagnetArray": - """Return the quadrupole array.""" + def quadrupoles(self) -> "MagnetArray": + """Return the quadrupole array used by the correction.""" self.check_peer() return self.peer.magnets.get(self.quad_array_name) @@ -153,7 +158,7 @@ def get(self): def readback(self): """Return the current dimensionless horizontal and vertical betatron tune.""" self.check_peer() - return self._tm.tune.get() + return self.tune_monitor.tune.get() def set(self, tune: np.array, iter: int = 1, wait_time: float = 0.0): """ @@ -215,8 +220,8 @@ def add(self, dtune: np.array, wait_time: float = 0.0): """ if np.shape(dtune) != (2,): raise PyAMLException("Tune.add(): invalid input dtune dimension, (2,) expected") - strengths = self._quads.strengths.get() + strengths = self.quadrupoles.strengths.get() strengths += self.correct(dtune) - self._quads.strengths.set(strengths) + self.quadrupoles.strengths.set(strengths) sleep(wait_time) self._setpoint += dtune diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 11707ea3..023b86ca 100644 --- a/pyaml/tuning_tools/tune_response_matrix.py +++ b/pyaml/tuning_tools/tune_response_matrix.py @@ -9,7 +9,7 @@ import logging from dataclasses import asdict from time import sleep -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 ..diagnostics.tune_monitor import BetatronTuneMonitor + logger = logging.getLogger(__name__) PYAMLCLASS = "TuneResponseMatrix" @@ -85,6 +89,10 @@ class TuneResponseMatrix(MeasurementTool, DynamicValidation): Configured number of tune measurements to average. sleep_between_meas : float Configured delay between averaged tune measurements. + quadrupoles : MagnetArray + Quadrupole array used for the measurement. + tune_monitor : BetatronTuneMonitor + Betatron tune monitor used for the measurement. Methods ------- @@ -126,6 +134,18 @@ def __init__( self.n_avg_meas = n_avg_meas self.sleep_between_meas = sleep_between_meas + @property + def quadrupoles(self) -> "MagnetArray": + """Return the quadrupole array used for the measurement.""" + self.check_peer() + return self.peer.magnets.get(self.quad_array_name) + + @property + def tune_monitor(self) -> "BetatronTuneMonitor": + """Return the betatron tune monitor used for the measurement.""" + self.check_peer() + return self.peer.get_betatron_tune_monitor(self.betatron_tune_name) + def measure( self, quad_delta: Optional[float] = None, @@ -201,8 +221,8 @@ def callback(action: Action, data:dict): """ # Get devices self.check_peer() - quads = self._peer.magnets.get(self.quad_array_name) - tm = self._peer.get_betatron_tune_monitor(self.betatron_tune_name) + quads = self.quadrupoles + tm = self.tune_monitor tunemat = np.zeros((len(quads), 2)) initial_tune = tm.tune.get() diff --git a/tests/tuning_tools/test_tool_accessors.py b/tests/tuning_tools/test_tool_accessors.py new file mode 100644 index 00000000..1ccbf5df --- /dev/null +++ b/tests/tuning_tools/test_tool_accessors.py @@ -0,0 +1,53 @@ +from pyaml.accelerator import Accelerator + + +def test_tuning_tools_expose_configured_elements(): + sr = Accelerator.load( + "tests/config/EBSOrbit.yaml", + ignore_external=True, + include_locations=False, + ) + design = sr.design + + tune_monitor = design.get_betatron_tune_monitor("BETATRON_TUNE") + quadrupoles = design.magnets.get("QForTune") + assert design.tune.tune_monitor is tune_monitor + assert design.tune.quadrupoles is quadrupoles + assert design.trm.tune_monitor is tune_monitor + assert design.trm.quadrupoles is quadrupoles + + chromaticity_monitor = design.get_chromaticity_monitor("CHROMATICITY_MONITOR") + sextupoles = design.magnets.get("Sext") + assert design.chromaticity.chromaticity_monitor is chromaticity_monitor + assert design.chromaticity.sextupoles is sextupoles + assert design.crm.chromaticity_monitor is chromaticity_monitor + assert design.crm.sextupoles is sextupoles + + bpms = design.bpms.get("BPM") + rf_plant = design.rf.get("DEFAULT_RF_PLANT") + assert chromaticity_monitor.tune_monitor is tune_monitor + assert chromaticity_monitor.bpms is bpms + assert chromaticity_monitor.rf_plant is rf_plant + assert design.dispersion.bpms is bpms + assert design.dispersion.rf_plant is rf_plant + + hcorrectors = design.magnets.get("HCorr") + vcorrectors = design.magnets.get("VCorr") + assert design.orbit.bpms is bpms + assert design.orbit.hcorrectors is hcorrectors + assert design.orbit.vcorrectors is vcorrectors + assert design.orbit.correctors.names() == hcorrectors.names() + vcorrectors.names() + assert design.orbit.rf_plant is rf_plant + assert design.orm.bpms is bpms + assert design.orm.hcorrectors is hcorrectors + assert design.orm.vcorrectors is vcorrectors + + for name in ("BBA-BPM_C04-04", "BBA2-BPM_C04-04"): + bba = design.get_bba(name) + assert bba.bpms is bpms + assert bba.bpm is design.bpm.get("BPM_C04-04") + assert bba.hcorrector is design.magnet.get("SF2E-C02-H") + assert bba.vcorrector is design.magnet.get("SD1A-C26-V") + assert bba.quadrupole is design.magnet.get("QF6B-C04") + + assert design.get_bba("BBA2-BPM_C04-04").tune_correction is design.tune