diff --git a/pySC/apps/bba.py b/pySC/apps/bba.py index fc77702..834b7d3 100644 --- a/pySC/apps/bba.py +++ b/pySC/apps/bba.py @@ -1,5 +1,5 @@ from pydantic import BaseModel, PrivateAttr, ConfigDict, model_validator -from typing import Optional, ClassVar, Literal +from typing import Optional, ClassVar, Literal, Union, Tuple import datetime import logging import numpy as np @@ -22,9 +22,9 @@ class BBAData(BaseModel, extra="forbid"): """ quadrupole: str bpm: str - corrector: str + corrector: Union[str, Tuple[str,str]] plane: str - dk0l: float # Corrector k0 (max) step + dk0l: Union[float, Tuple[float,float]] # Corrector k0 (max) step dk1l: float # Quadrupole k1 step n0: int # Number of steps in the corrector strength shots_per_orbit: int @@ -32,7 +32,7 @@ class BBAData(BaseModel, extra="forbid"): magnet_type: BBA_MagnetType = MagnetType.norm_quad bpm_number: int - initial_k0l: Optional[float] = None + initial_k0l: Optional[Union[float, Tuple[float,float]]] = None initial_k1: Optional[float] = None timestamp: Optional[float] = None original_save_path: Optional[str] = None @@ -66,6 +66,16 @@ def save(self, folder_to_save: Optional[Path] = None) -> Path: filename = Path(folder_to_save) / Path(f'BBA_{self.bpm}_{self.plane}_{time_str}.h5') self.original_save_path = str(filename.resolve()) dict_to_save = self.model_dump() + + if isinstance(dict_to_save["corrector"], tuple): + dict_to_save["corrector"] = np.array(dict_to_save["corrector"], dtype="S") + if isinstance(dict_to_save["dk0l"], tuple): + dict_to_save["dk0l"] = np.array(dict_to_save["dk0l"], dtype=float) + if isinstance(dict_to_save["initial_k0l"], tuple): + dict_to_save["initial_k0l"] = np.array(dict_to_save["initial_k0l"], dtype=float) + if isinstance(dict_to_save["magnet_type"], MagnetType): + dict_to_save["magnet_type"] = dict_to_save["magnet_type"].value + dict_to_h5(dict_to_save, filename) logger.info(f'Saved data to {filename} .') return filename @@ -120,6 +130,7 @@ class BBA_Measurement(BaseModel, extra="forbid"): H_data: Optional[BBAData] = None V_data: Optional[BBAData] = None + HV_data: Optional[BBAData] = None last_ios: NPARRAY = np.array([]) last_bpm_pos: float = 0 @@ -152,7 +163,12 @@ def __init__(self, **kwargs): corrector=self.v_corrector, dk0l=self.dk0l_y, dk1l=self.dk1l_y, n0=self.n0, shots_per_orbit=self.shots_per_orbit, bipolar=self.bipolar, magnet_type=self.magnet_type) - + if self.h_corrector is not None and self.v_corrector is not None: + self.HV_data = BBAData(plane='HV', bpm=self.bpm, quadrupole=self.quadrupole, bpm_number=self.bpm_number, + corrector=(self.h_corrector, self.v_corrector), dk0l=(self.dk0l_x, self.dk0l_y), + dk1l=max(self.dk1l_x, self.dk1l_y), + n0=self.n0, shots_per_orbit=self.shots_per_orbit, bipolar=self.bipolar, + magnet_type=self.magnet_type) def print_init(self): logger.debug("Measurement plan:") @@ -174,7 +190,7 @@ def print_init(self): def one_plane_loop(self, plane: str): assert self._interface is not None - assert plane in ['H', 'V'] + assert plane in ['H', 'V', 'HV'] interface = self._interface if plane == 'H': @@ -182,28 +198,56 @@ def one_plane_loop(self, plane: str): code = BBACode.HORIZONTAL ios_ready_code = BBACode.HORIZONTAL_IOS_READY code_done = BBACode.HORIZONTAL_DONE - else: + elif plane == 'V': logger.debug('Starting measurement in vertical plane.') code = BBACode.VERTICAL ios_ready_code = BBACode.VERTICAL_IOS_READY code_done = BBACode.VERTICAL_DONE + elif plane == 'HV': + logger.debug('Starting measurement in horizontal and vertical plane.') + code = BBACode.HORIZONTAL_VERTICAL + ios_ready_code = BBACode.HORIZONTAL_VERTICAL_IOS_READY + code_done = BBACode.HORIZONTAL_VERTICAL_DONE - corrector = self.h_corrector if plane == 'H' else self.v_corrector - initial_k0l = self.initial_h_k0l if plane == 'H' else self.initial_v_k0l - data = self.H_data if plane == 'H' else self.V_data + if plane == 'H': + corrector = self.h_corrector + initial_k0l = self.initial_h_k0l + data = self.H_data + elif plane == 'V': + corrector = self.v_corrector + initial_k0l = self.initial_v_k0l + data = self.V_data + elif plane == 'HV': + corrector = (self.h_corrector, self.v_corrector) + initial_k0l = (self.initial_h_k0l, self.initial_v_k0l) + data = self.HV_data logger.debug('Setting corrector to under first value (k0 - 1.2 dk0) for hysteresis') - interface.set(corrector, initial_k0l - 1.2 * data.dk0l) + if plane == 'HV': + interface.set(corrector[0], initial_k0l[0] - 1.2 * data.dk0l[0]) + interface.set(corrector[1], initial_k0l[1] - 1.2 * data.dk0l[1]) + else: + interface.set(corrector, initial_k0l - 1.2 * data.dk0l) yield code - k0_array = np.linspace(-data.dk0l, data.dk0l, self.n0) + initial_k0l + if plane == 'HV': + k0_array0 = np.linspace(-data.dk0l[0], data.dk0l[0], self.n0) + initial_k0l[0] + k0_array1 = np.linspace(-data.dk0l[1], data.dk0l[1], self.n0) + initial_k0l[1] + k0_array = [(k0_array0[i], k0_array1[i]) for i in range(self.n0)] + else: + k0_array = np.linspace(-data.dk0l, data.dk0l, self.n0) + initial_k0l get_orbit = self._interface.get_orbit for ii, k0_sp in enumerate(k0_array): # set next setpoint in corrector - logger.debug(f'{ii+1}/{self.n0} Stepping to next corrector setpoint: {k0_sp*1e6:+.1f} murad') - interface.set(corrector, k0_sp) + if plane == 'HV': + logger.debug(f'{ii+1}/{self.n0} Stepping to next corrector setpoint: ({k0_sp[0]*1e6:+.1f}, {k0_sp[1]*1e6:+.1f}) murad') + interface.set(corrector[0], k0_sp[0]) + interface.set(corrector[1], k0_sp[1]) + else: + logger.debug(f'{ii+1}/{self.n0} Stepping to next corrector setpoint: {k0_sp*1e6:+.1f} murad') + interface.set(corrector, k0_sp) yield code # correct vertical orbit? @@ -255,12 +299,23 @@ def one_plane_loop(self, plane: str): if self.live_ios: # TODO decide how to use n_downstream after writing an actual application for trajectory-based BBA. self.last_bpm_pos, self.last_ios = get_one_ios(data=data, ii=ii, n_downstream=None) - logger.debug(f" Position at BPM = {1e6*self.last_bpm_pos:.3f} μm, std(I.O.S.) = {1e6*np.std(self.last_ios):.3f} μm") + if plane == "HV": + bpm0 = 1e6*self.last_bpm_pos[0] + bpm1 = 1e6*self.last_bpm_pos[1] + ios0 = 1e6*np.std(self.last_ios[0]) + ios1 = 1e6*np.std(self.last_ios[1]) + logger.debug(f" Position at BPM = ({bpm0:.3f}, {bpm1:.3f}) μm, std(I.O.S.) = ({ios0:.3f}, {ios1:.3f}) μm") + else: + logger.debug(f" Position at BPM = {1e6*self.last_bpm_pos:.3f} μm, std(I.O.S.) = {1e6*np.std(self.last_ios):.3f} μm") yield ios_ready_code logger.debug("") # restore corrector to initial setpoint - interface.set(corrector, initial_k0l) + if plane == 'HV': + interface.set(corrector[0], initial_k0l[0]) + interface.set(corrector[1], initial_k0l[1]) + else: + interface.set(corrector, initial_k0l) #save data yield code_done @@ -286,6 +341,13 @@ def generate(self, interface: AbstractInterface, plane: Optional[str] = None, sk self.V_data.initial_k1 = self.initial_k1l self.V_data.timestamp = timestamp + if plane == "HV": + if self.h_corrector is None or self.v_corrector is None: + raise ValueError("BBA_Measurement.h_corrector or BBA_Measurement.v_corrector is not defined.") + self.HV_data.initial_k0l = (self.initial_h_k0l, self.initial_v_k0l) + self.HV_data.initial_k1 = self.initial_k1l + self.HV_data.timestamp = timestamp + self.print_init() if self.h_corrector is None: @@ -307,6 +369,10 @@ def generate(self, interface: AbstractInterface, plane: Optional[str] = None, sk for code in self.one_plane_loop('V'): yield code + if (plane == 'HV') and self.h_corrector is not None and self.v_corrector is not None: + for code in self.one_plane_loop('HV'): + yield code + yield BBACode.DONE # def run(self, generator=None): @@ -316,6 +382,8 @@ def generate(self, interface: AbstractInterface, plane: Optional[str] = None, sk # logger.debug(f' Got code: {code}') def prep_ios(data: BBAData, n_downstream: Optional[int] = None) -> tuple[np.ndarray, np.ndarray]: + if data.plane == 'HV': + raise Exception("prep_ios does not support HV plane.") ( bpm_number, bpm_position, induced_orbit_shift, start, k1_arr, all_x, all_y, ) = _prepare_data_for_ios_calculation(data=data, n_downstream=n_downstream) @@ -381,7 +449,7 @@ def calc_ios(ii: int, magnet_type: MagnetType, delta: float, k1_arr: list[float] induced_orbit_shift = np.polyfit(k1_arr, all_x[:,ii], 1)[0] * delta else: raise Exception(f"Unknown magnet type {magnet_type}.") - else: + elif plane == 'V': bpm_position = np.mean(all_y[:, ii, bpm_number - start]) if magnet_type in [MagnetType.skew_quad, MagnetType.norm_sext]: induced_orbit_shift = np.polyfit(k1_arr, all_x[:,ii], 1)[0] * delta @@ -389,6 +457,17 @@ def calc_ios(ii: int, magnet_type: MagnetType, delta: float, k1_arr: list[float] induced_orbit_shift = np.polyfit(k1_arr, all_y[:,ii], 1)[0] * delta else: raise Exception(f"Unknown magnet type {magnet_type}.") + elif plane == 'HV': + bpm_position = (np.mean(all_x[:, ii, bpm_number - start]), + np.mean(all_y[:, ii, bpm_number - start]) + ) + if magnet_type in [MagnetType.skew_quad]: + induced_orbit_shift = (np.polyfit(k1_arr, all_y[:,ii], 1)[0] * delta, np.polyfit(k1_arr, all_x[:,ii], 1)[0] * delta) + elif magnet_type in [MagnetType.norm_quad]: + induced_orbit_shift = (np.polyfit(k1_arr, all_x[:,ii], 1)[0] * delta, np.polyfit(k1_arr, all_y[:,ii], 1)[0] * delta) + else: + raise Exception(f"Magnet type {magnet_type} not supported in plane=HV measurement.") + return bpm_position, induced_orbit_shift def reject_bpm_outlier(induced_orbit_shift: np.ndarray, bpm_outlier_sigma: float) -> np.ndarray[bool]: @@ -412,32 +491,32 @@ def reject_center_outlier(center: np.ndarray, center_cutoff: float) -> np.ndarra return mask class BBAAnalysis(BaseModel): - offset: float - offset_error: float + offset: Union[float, Tuple[float, float]] + offset_error: Union[float, Tuple[float, float]] - quadratics: NPARRAY - slopes: NPARRAY - intercepts: NPARRAY - centers: NPARRAY + quadratics: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + slopes: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + intercepts: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + centers: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] - quadratics_err: NPARRAY - slopes_err: NPARRAY - intercepts_err: NPARRAY - centers_err: NPARRAY + quadratics_err: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + slopes_err: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + intercepts_err: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + centers_err: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] - induced_orbit_shift: NPARRAY - bpm_position: NPARRAY + induced_orbit_shift: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] + bpm_position: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] - mask_accepted: NPARRAY + mask_accepted: Union[NPARRAY, Tuple[NPARRAY, NPARRAY]] n_downstream: Optional[int] fit_order:int - rejected_outliers: int - rejected_slopes: int - rejected_centers: int - total_rejections: int = 0 + rejected_outliers: Union[int, Tuple[int, int]] + rejected_slopes: Union[int, Tuple[int, int]] + rejected_centers: Union[int, Tuple[int, int]] + total_rejections: Union[int, Tuple[int, int]] = 0 bpm_outlier_sigma: float slope_cutoff: float @@ -462,6 +541,40 @@ def analyze(cls, data: BBAData, n_downstream: Optional[int] = None, bpm_outlier_ if center_cutoff is None: center_cutoff = cls.default_center_cutoff + if data.plane == "HV": + try: + data.plane = "H" + resultH = cls.analyze(data=data, n_downstream=n_downstream, bpm_outlier_sigma=bpm_outlier_sigma, + slope_cutoff=slope_cutoff, center_cutoff=center_cutoff) + data.plane = "V" + resultV = cls.analyze(data=data, n_downstream=n_downstream, bpm_outlier_sigma=bpm_outlier_sigma, + slope_cutoff=slope_cutoff, center_cutoff=center_cutoff) + finally: + data.plane = "HV" + return BBAAnalysis(offset=(resultH.offset, resultV.offset), + offset_error=(resultH.offset_error, resultV.offset_error), + quadratics=(resultH.quadratics, resultV.quadratics), + slopes=(resultH.slopes, resultV.slopes), + intercepts=(resultH.intercepts, resultV.intercepts), + centers=(resultH.centers, resultV.centers), + quadratics_err=(resultH.quadratics_err, resultV.quadratics_err), + slopes_err=(resultH.slopes_err, resultV.slopes_err), + intercepts_err=(resultH.intercepts_err, resultV.intercepts_err), + centers_err=(resultH.centers_err, resultV.centers_err), + induced_orbit_shift=(resultH.induced_orbit_shift, resultV.induced_orbit_shift), + bpm_position=(resultH.bpm_position, resultV.bpm_position), + mask_accepted=(resultH.mask_accepted, resultV.mask_accepted), + n_downstream=n_downstream, + fit_order=resultH.fit_order, + rejected_outliers=(resultH.rejected_outliers, resultV.rejected_outliers), + rejected_slopes=(resultH.rejected_slopes, resultV.rejected_slopes), + rejected_centers=(resultH.rejected_centers, resultV.rejected_centers), + total_rejections=(resultH.total_rejections, resultV.total_rejections), + bpm_outlier_sigma=bpm_outlier_sigma, + slope_cutoff=slope_cutoff, + center_cutoff=center_cutoff, + ) + bpm_position, induced_orbit_shift = prep_ios(data=data, n_downstream=n_downstream) nanmask = ~np.isnan(bpm_position) diff --git a/pySC/apps/codes.py b/pySC/apps/codes.py index 9a2127b..e97e5cd 100644 --- a/pySC/apps/codes.py +++ b/pySC/apps/codes.py @@ -15,6 +15,9 @@ class BBACode(IntEnum): DONE = 7 HORIZONTAL_IOS_READY = 8 VERTICAL_IOS_READY = 9 + HORIZONTAL_VERTICAL = 10 + HORIZONTAL_VERTICAL_DONE = 11 + HORIZONTAL_VERTICAL_IOS_READY = 12 class ResponseCode(IntEnum): INITIALIZED = MeasurementCode.INITIALIZED.value diff --git a/pySC/apps/measurements.py b/pySC/apps/measurements.py index 9124d89..945e2f9 100644 --- a/pySC/apps/measurements.py +++ b/pySC/apps/measurements.py @@ -99,7 +99,7 @@ def measure_bba(interface: AbstractInterface, bpm_name, config: dict, shots_per_ live_ios=live_ios, ) - generator = measurement.generate(interface=interface, plane=plane) + generator = measurement.generate(interface=interface, plane=plane, skip_cycle=skip_cycle) # run measurement loop for code in generator: @@ -108,6 +108,8 @@ def measure_bba(interface: AbstractInterface, bpm_name, config: dict, shots_per_ measurement.H_data.save(folder_to_save=folder_to_save) if not skip_save and code is BBACode.VERTICAL_DONE: measurement.V_data.save(folder_to_save=folder_to_save) + if not skip_save and code is BBACode.HORIZONTAL_VERTICAL_DONE: + measurement.HV_data.save(folder_to_save=folder_to_save) yield code, measurement def measure_ORM(interface: AbstractInterface, corrector_names: list[str], delta: Union[float, list[float]], diff --git a/pySC/version.py b/pySC/version.py index a06ff4e..3a384fb 100644 --- a/pySC/version.py +++ b/pySC/version.py @@ -1 +1 @@ -__version__ = "1.5.3" +__version__ = "1.5.4" diff --git a/tests/apps/test_bba.py b/tests/apps/test_bba.py index 79a090d..9e573a7 100644 --- a/tests/apps/test_bba.py +++ b/tests/apps/test_bba.py @@ -83,6 +83,56 @@ def _make_bba_data(n0=7, n_bpms=10, bpm_number=3, plane='H', bipolar=True, return data +def _make_hv_bba_data(n0=7, n_bpms=10, bpm_number=3, bipolar=True, + dk0l_x=1e-4, dk0l_y=2e-4, dk1l=0.05, + offset_x=0.0, offset_y=0.0): + """Build synthetic normal-quadrupole HV BBA data with known offsets.""" + data = BBAData( + quadrupole='Q1', bpm='BPM3', corrector=('CH1', 'CV1'), plane='HV', + dk0l=(dk0l_x, dk0l_y), dk1l=dk1l, n0=n0, shots_per_orbit=1, + bipolar=bipolar, bpm_number=bpm_number, + ) + + x_steps = np.linspace(-dk0l_x, dk0l_x, n0) + y_steps = np.linspace(-dk0l_y, dk0l_y, n0) + rng = np.random.default_rng(321) + x_response = np.abs(1.0 + 0.3 * rng.standard_normal(n_bpms)) + y_response = np.abs(1.0 + 0.3 * rng.standard_normal(n_bpms)) + + for x_step, y_step in zip(x_steps, y_steps): + x_center = np.zeros(n_bpms) + y_center = np.zeros(n_bpms) + x_center[bpm_number] = x_step + offset_x + y_center[bpm_number] = y_step + offset_y + + x_ios = x_response * x_step + y_ios = y_response * y_step + + x_up = x_center + x_ios + y_up = y_center + y_ios + x_down = x_center.copy() + y_down = y_center.copy() + if bipolar: + x_down -= x_ios + y_down -= y_ios + + data.raw_bpm_x_center.append(list(x_center)) + data.raw_bpm_y_center.append(list(y_center)) + data.raw_bpm_x_up.append(list(x_up)) + data.raw_bpm_y_up.append(list(y_up)) + data.raw_bpm_x_down.append(list(x_down)) + data.raw_bpm_y_down.append(list(y_down)) + + data.raw_bpm_x_center_err.append(list(np.zeros(n_bpms))) + data.raw_bpm_y_center_err.append(list(np.zeros(n_bpms))) + data.raw_bpm_x_up_err.append(list(np.zeros(n_bpms))) + data.raw_bpm_y_up_err.append(list(np.zeros(n_bpms))) + data.raw_bpm_x_down_err.append(list(np.zeros(n_bpms))) + data.raw_bpm_y_down_err.append(list(np.zeros(n_bpms))) + + return data + + def _make_sextupole_bba_data(n0=9, n_bpms=10, bpm_number=3, plane='H', dk0l=1e-4, dk1l=0.05, offset=0.0): """Build synthetic normal-sextupole BBA data with a known quadratic center.""" @@ -196,6 +246,29 @@ def test_get_one_ios_matches_prep_ios(self, bipolar, n_downstream): assert one_bpm_position == pytest.approx(bpm_position[ii]) np.testing.assert_allclose(one_induced_orbit_shift, induced_orbit_shift[ii]) + def test_prep_ios_hv_raises_clear_error(self): + """HV data are analyzed as separate planes, not through prep_ios.""" + data = _make_hv_bba_data() + + with pytest.raises(Exception, match="does not support HV plane"): + prep_ios(data) + + assert data.plane == 'HV' + + def test_get_one_ios_hv_returns_both_planes(self): + """Single-step IOS for HV data returns separate H and V values.""" + data = _make_hv_bba_data(n0=5, n_bpms=8, bpm_number=2, + offset_x=0.001, offset_y=-0.002) + + bpm_position, induced_orbit_shift = get_one_ios(data, ii=2) + + assert len(bpm_position) == 2 + assert len(induced_orbit_shift) == 2 + assert induced_orbit_shift[0].shape == (8,) + assert induced_orbit_shift[1].shape == (8,) + assert np.all(np.isfinite(induced_orbit_shift[0])) + assert np.all(np.isfinite(induced_orbit_shift[1])) + class TestRejectBpmOutlier: def test_reject_bpm_outlier_no_outliers(self): @@ -323,6 +396,18 @@ def test_bba_analysis_total_rejections_stored(self): ) assert result.total_rejections == expected + def test_bba_analysis_hv_known_offsets(self): + """HV analysis returns one aggregate result with both planes.""" + data = _make_hv_bba_data(n0=9, offset_x=0.0015, offset_y=-0.002) + + result = BBAAnalysis.analyze(data) + + np.testing.assert_allclose(result.offset, (0.0015, -0.002), atol=1e-6) + assert len(result.bpm_position) == 2 + assert len(result.induced_orbit_shift) == 2 + assert len(result.mask_accepted) == 2 + assert data.plane == 'HV' + # =========================================================================== # Generator / measurement tests (need MockInterface) @@ -444,6 +529,49 @@ def test_bba_measurement_live_ios_yields_ready_codes(self, mock_interface): assert len(ready_codes) == meas.n0 assert code == BBACode.DONE + def test_bba_measurement_generate_hv_plane(self, mock_interface): + """Generator produces HV codes and collects HV data.""" + iface = mock_interface(n_bpms=10) + meas = self._make_measurement() + + codes = list(meas.generate(iface, plane='HV', skip_cycle=True)) + + assert BBACode.HYSTERESIS not in codes + assert BBACode.HORIZONTAL_VERTICAL in codes + assert BBACode.HORIZONTAL_VERTICAL_DONE in codes + assert BBACode.HORIZONTAL not in codes + assert BBACode.VERTICAL not in codes + assert codes[-1] == BBACode.DONE + assert len(meas.HV_data.raw_bpm_x_center) == meas.n0 + assert len(meas.HV_data.raw_bpm_y_center) == meas.n0 + assert meas.HV_data.initial_k0l == ( + pytest.approx(meas.initial_h_k0l), + pytest.approx(meas.initial_v_k0l), + ) + assert meas.HV_data.initial_k1 == pytest.approx(meas.initial_k1l) + assert meas.HV_data.timestamp is not None + + def test_bba_measurement_live_ios_hv_yields_ready_codes(self, mock_interface): + """live_ios=True yields one HV IOS-ready code per measured point.""" + iface = mock_interface(n_bpms=10) + meas = self._make_measurement() + meas.live_ios = True + + ready_codes = [] + for code in meas.generate(iface, plane='HV', skip_cycle=True): + if code == BBACode.HORIZONTAL_VERTICAL_IOS_READY: + ready_codes.append(code) + assert iface.get(meas.quadrupole) == pytest.approx(meas.initial_k1l) + assert len(meas.last_bpm_pos) == 2 + assert len(meas.last_ios) == 2 + assert meas.last_ios[0].shape == (10,) + assert meas.last_ios[1].shape == (10,) + assert np.all(np.isfinite(meas.last_ios[0])) + assert np.all(np.isfinite(meas.last_ios[1])) + + assert len(ready_codes) == meas.n0 + assert code == BBACode.DONE + def test_quad_is_skew_deprecated_argument_sets_magnet_type(self): """Deprecated quad_is_skew input remains backward compatible.""" meas = BBA_Measurement( diff --git a/tests/apps/test_measurements.py b/tests/apps/test_measurements.py index 9a88de9..95ff67d 100644 --- a/tests/apps/test_measurements.py +++ b/tests/apps/test_measurements.py @@ -6,6 +6,7 @@ from pySC.apps.measurements import orbit_correction, measure_bba, measure_ORM, measure_dispersion from pySC.apps.response_matrix import ResponseMatrix from pySC.apps.codes import BBACode, ResponseCode, DispersionCode +from pySC.utils.file_tools import h5_to_dict # --------------------------------------------------------------------------- @@ -147,6 +148,37 @@ def test_deprecated_quad_is_skew_config_sets_magnet_type(self, mock_interface): assert measurement.magnet_type == 'skew_quadrupole' assert measurement.H_data.magnet_type == 'skew_quadrupole' + def test_measure_bba_hv_saves_data(self, mock_interface, tmp_path): + """HV BBA measurement saves one HDF5 file with split corrector data.""" + iface = mock_interface(n_bpms=5) + config = _bba_config() + del config['QUAD_is_skew'] + config['magnet_type'] = 'normal_quadrupole' + + results = list(measure_bba( + iface, + 'BPM0', + config, + skip_save=False, + folder_to_save=tmp_path, + plane='HV', + skip_cycle=True, + )) + + codes = [code for code, _ in results] + files = list(tmp_path.glob('BBA_BPM0_HV_*.h5')) + + assert BBACode.HYSTERESIS not in codes + assert BBACode.HYSTERESIS_DONE not in codes + assert BBACode.HORIZONTAL_VERTICAL_DONE in codes + assert len(files) == 1 + + saved_data = h5_to_dict(files[0]) + assert saved_data['plane'] == 'HV' + assert tuple(saved_data['corrector']) == (b'CH0', b'CV0') + np.testing.assert_allclose(saved_data['dk0l'], [1e-4, 1e-4]) + np.testing.assert_allclose(saved_data['initial_k0l'], [0.0, 0.0]) + # --------------------------------------------------------------------------- # measure_ORM tests