From b2dcde9931041e0e42d227607b0f4bae68c04951 Mon Sep 17 00:00:00 2001 From: PONS Date: Mon, 20 Jul 2026 11:57:45 +0200 Subject: [PATCH] Added BBA plot --- pyaml/tuning_tools/bba.py | 115 ++++++++++++++++++++++++++++++++------ 1 file changed, 99 insertions(+), 16 deletions(-) diff --git a/pyaml/tuning_tools/bba.py b/pyaml/tuning_tools/bba.py index de9b1700..2e915ba1 100644 --- a/pyaml/tuning_tools/bba.py +++ b/pyaml/tuning_tools/bba.py @@ -1,6 +1,7 @@ import logging from typing import Callable, Optional +import matplotlib.pyplot as plt import numpy as np import pySC from pydantic import ConfigDict @@ -9,6 +10,7 @@ from pySC.apps.codes import BBACode from ..common.constants import Action +from ..common.exception import PyAMLException from ..external.pySC_interface import pySCInterface from .measurement_tool import MeasurementTool, MeasurementToolConfigModel @@ -69,6 +71,7 @@ def measure( n_avg_meas: Optional[int] = None, sleep_between_meas: Optional[float] = None, callback: Optional[Callable] = None, + plane: Optional[str] = None, ): """ Measure BBA. @@ -77,8 +80,17 @@ def measure( .. code-block:: python - sr = Accelerator.load("MyAccelerator.yaml") - TODO + sr = Accelerator.load("tests/config/EBSOrbit.yaml") + SR = sr.design + bba = SR.get_bba("BBA-BPM_C04-04") + + # Add a misalignement + SR.get_bpm("BPM_C04-04").offset.set([20e-6,-15e-6]) + + bba.measure(plane="H") + print(f"HOffset: {bba.h_offset()}") + + bba.plot_data("H") Parameters ---------- @@ -96,6 +108,8 @@ def measure( callback is executed after each strength setting and after each orbit reading. If the callback returns false, then the process is aborted. + plane: str, optional + Plane to perform ("H" or "V", None => both plane) """ nb_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step @@ -120,9 +134,11 @@ def measure( "QUAD_dk_H": self._cfg.hquad_delta, "VCORR_delta": self._cfg.vcorr_delta, "QUAD_dk_V": self._cfg.vquad_delta, - "QUAD_is_skew": False, + "magnet_type": "normal_quadrupole", } + # logging.getLogger("pySC.apps.measurements").setLevel(logging.DEBUG) + generator = measure_bba( interface=interface, bpm_name=self._cfg.bpm_name, @@ -131,28 +147,35 @@ def measure( n_corr_steps=self._cfg.n_step, bipolar=False, skip_save=True, + plane=plane, ) pySC.disable_pySC_rich() aborted = False err = None - idx = 0 - self._latest_measurement = {"HOffset": np.nan, "VOffset": np.nan, "HOffsetError": np.nan, "VOffsetError": np.nan} + hstep = 0 + vstep = 0 try: self._register_callback(callback) self._init_measure() + self.latest_measurement["HData"] = None + self.latest_measurement["VData"] = None for code, measurement_object in generator: - # print(f"Got code: {code.name}") + if code == BBACode.HORIZONTAL: + self.send_callback(Action.MEASURE, {"step": hstep, "plane": "H", "bba_data": measurement_object.H_data}) + hstep += 1 + if code == BBACode.VERTICAL: + self.send_callback(Action.MEASURE, {"step": vstep, "plane": "V", "bba_data": measurement_object.V_data}) + vstep += 1 + if code == BBACode.HORIZONTAL_DONE: result = BBAAnalysis.analyze(measurement_object.H_data) - self.latest_measurement["HOffset"] = result.offset - self.latest_measurement["HOffsetError"] = result.offset_error + self.latest_measurement["HData"] = result if code == BBACode.VERTICAL_DONE: result = BBAAnalysis.analyze(measurement_object.V_data) - self.latest_measurement["VOffset"] = result.offset - self.latest_measurement["VOffsetError"] = result.offset_error - idx += 1 + self.latest_measurement["VData"] = result + except Exception as ex: err = ex except KeyboardInterrupt as ex: @@ -162,7 +185,7 @@ def measure( # TODO self.send_callback( Action.RESTORE, - {"idx": idx}, + {}, raiseException=False, ) @@ -176,13 +199,73 @@ def measure( return True def h_offset(self) -> float: - return self.latest_measurement["HOffset"] + return self.latest_measurement["HData"].offset if self.latest_measurement["HData"] is not None else np.nan def h_offset_error(self) -> float: - return self.latest_measurement["HOffsetError"] + return self.latest_measurement["HData"].offset_error if self.latest_measurement["HData"] is not None else np.nan def v_offset(self) -> float: - return self.latest_measurement["VOffset"] + return self.latest_measurement["VData"].offset if self.latest_measurement["VData"] is not None else np.nan def v_offset_error(self) -> float: - return self.latest_measurement["VOffsetError"] + return self.latest_measurement["VData"].offset_error if self.latest_measurement["VData"] is not None else np.nan + + def plot_data(self, plane: str): + """ + Plot BBA data. + + Parameters + ---------- + plane: str + Plane to plot ("H" or "V") + """ + + planeData = plane + "Data" + if planeData not in self.latest_measurement or self.latest_measurement[planeData] is None: + raise PyAMLException("No BBA data to plot, please call measure() first") + + fig, axes = plt.subplot_mosaic(mosaic=[["A", "A", "S", "S"], ["A", "A", "C", "C"]]) + + bpm_pos = self.latest_measurement[planeData].bpm_position + ios = self.latest_measurement[planeData].induced_orbit_shift + slopes = self.latest_measurement[planeData].slopes + centers = self.latest_measurement[planeData].centers + final_mask = self.latest_measurement[planeData].mask_accepted + offset = self.latest_measurement[planeData].offset + + xp_min = np.min(bpm_pos) * 1e6 + xp_max = np.max(bpm_pos) * 1e6 + for kk in range(len(ios[0, :])): + p0 = np.polyfit(bpm_pos[:] * 1e6, ios[:, kk] * 1e6, 1) + xp = np.linspace(xp_min, xp_max, 10) + yp = p0[0] * xp + p0[1] + axes["A"].plot(xp, yp, "-", c="C0" if final_mask[kk] else "C1", alpha=0.3) + + for kk in range(len(bpm_pos[:])): + yy = ios[kk] * 1e6 + xx = np.ones_like(yy) * bpm_pos[kk] * 1e6 + axes["A"].plot(xx[final_mask], yy[final_mask], ".", c="C0") + axes["A"].plot(xx[~final_mask], yy[~final_mask], ".", c="C1") + + axes["A"].set_xlabel(f"BPM position [μm]\n{self._cfg.bpm_name} offset = {offset * 1e6:.3f} [μm]") + axes["A"].set_ylabel("Modulation [μm]") + axes["A"].grid() + + bpm_numbers = np.arange(len(final_mask)) + axes["S"].plot(bpm_numbers[final_mask], slopes[final_mask], ".", c="C0", label="") + axes["S"].plot(bpm_numbers[~final_mask], slopes[~final_mask], ".", c="C1") + axes["S"].set_xlabel("BPM number") + axes["S"].set_ylabel("Slope") + axes["S"].grid() + + axes["C"].plot(bpm_numbers[final_mask], centers[final_mask] * 1e6, ".", c="C0") + axes["C"].plot(bpm_numbers[~final_mask], centers[~final_mask] * 1e6, ".", c="C1", label="rejected") + axes["C"].set_xlabel("BPM number") + axes["C"].set_ylabel("Center [μm]") + axes["C"].legend() + axes["C"].grid() + + fig.tight_layout() + fig.canvas.manager.set_window_title(f"{plane} BBA {self._cfg.bpm_name}") + + plt.show()