Add bipolar quadrupole modulation option to the BBA tool

pyaml/tuning_tools/bba.py

@@ -79,6 +79,11 @@ class BBA(MeasurementTool, DynamicValidation):
         Number of BPM measurements to average at each step.
     sleep_between_meas : float, default=0
         Time in seconds to wait between individual BPM measurements.
+    bipolar : bool, default=False
+        Quadrupole modulation. If ``False`` (unipolar), the orbit is measured
+        at the initial strength and at ``+quad_delta`` for each corrector
+        step. If ``True`` (bipolar), it is also measured at ``-quad_delta``
+        and the induced orbit shift is fitted over the three strengths.
 
     Methods
     -------
@@ -112,7 +117,7 @@ class BBA(MeasurementTool, DynamicValidation):
     -------------
     .. code-block:: yaml
 
-        - type: pyaml.tuning_tools.bba
+        - class: pyaml.tuning_tools.bba.BBA
           name: BBA
           bpm_array_name: BPM
           bpm_name: BPM_NAME
@@ -123,6 +128,7 @@ class BBA(MeasurementTool, DynamicValidation):
           vcorr_delta: 1e-5
           hquad_delta: 5e-3
           vquad_delta: 5e-3
+          bipolar: true
     """
 
     def __init__(
@@ -141,6 +147,7 @@ def __init__(
         sleep_between_step: float = 0,
         n_avg_meas: int = 1,
         sleep_between_meas: float = 0,
+        bipolar: bool = False,
     ):
         """
         Initialize the BBA.
@@ -159,6 +166,7 @@ def __init__(
         self.sleep_between_step = sleep_between_step
         self.n_avg_meas = n_avg_meas
         self.sleep_between_meas = sleep_between_meas
+        self.bipolar = bipolar
 
     @property
     def bpms(self) -> "BPMArray":
@@ -197,6 +205,7 @@ def measure(
         sleep_between_meas: Optional[float] = None,
         callback: Optional[Callable] = None,
         plane: Optional[str] = None,
+        bipolar: Optional[bool] = None,
     ):
         """
         Measure BBA.
@@ -235,10 +244,15 @@ def measure(
             If the callback returns false, then the process is aborted.
         plane : str, optional
             Plane to perform ("H" or "V", None => both plane)
+        bipolar : bool, optional
+            Bipolar (+-quad_delta) or unipolar (+quad_delta) quadrupole
+            modulation.
+            Default: from config.
         """
         nb_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas
         sleep_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step
         sleep_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas
+        bipolar = bipolar if bipolar is not None else self.bipolar
 
         element_holder = self._peer
         interface = pySCInterface(
@@ -271,7 +285,7 @@ def measure(
             config=bba_pySC_config,
             shots_per_orbit=nb_meas,
             n_corr_steps=self.n_step,
-            bipolar=False,
+            bipolar=bipolar,
             skip_save=True,
             plane=plane,
             live_ios=True,

tests/tuning_tools/test_bba.py

@@ -25,3 +25,18 @@ def test_bba():
     assert np.abs(bba.v_offset() + 15e-6) < 1e-6
     assert np.abs(bba.h_offset_error() - 3.425616879733572e-07) < 1e-10
     assert np.abs(bba.v_offset_error() - 2.593196843945284e-07) < 1e-10
+
+
+@pytest.mark.parametrize("bipolar", [False, True])
+def test_bba_bipolar(bipolar):
+    sr = Accelerator.load("tests/config/EBSOrbit.yaml")
+    SR = sr.design
+    SR.diagnostic.bpm.get("BPM_C04-04").offset.set([20e-6, -15e-6])
+
+    bba = SR.tool.get("BBA-BPM_C04-04")
+    assert bba.bipolar is False
+    k1 = bba.quadrupole.strength.get()
+    bba.measure(bipolar=bipolar)
+    assert np.abs(bba.h_offset() - 20e-6) < 1e-6
+    assert np.abs(bba.v_offset() + 15e-6) < 1e-6
+    assert bba.quadrupole.strength.get() == pytest.approx(k1)