Added BBA plot

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()