add options to choose BPMs pairing for coupling calculation (#371)

Co-authored-by: Felix Soubelet <[email protected]>

Author: awegsche

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 #### 2022-04-25
 
+- Added:
+  - There is now an option, `coupling_pairing`, for the BPM pairing in coupling calculation, to let the user choose the number of BPMs instead of the usual "best candidate" method.
+
+#### 2022-04-25
+
 - Fixed:
   - Only perform index merging on the `NAME` column during coupling calculation. This solves an (at the moment) un-understood issue where some BPMs would have different `S` values in different files.
 

omc3/__init__.py

@@ -11,7 +11,7 @@
 __title__ = "omc3"
 __description__ = "An accelerator physics tools package for the OMC team at CERN."
 __url__ = "https://github.com/pylhc/omc3"
-__version__ = "0.2.6"
+__version__ = "0.2.7"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/hole_in_one.py

@@ -232,6 +232,13 @@ def hole_in_one_entrypoint(opt, rest):
         Flags: **--coupling_method**
         Choices: ``(0, 1, 2)``
         Default: ``2``
+      - **coupling_pairing**: Pairing mode for 2 BPM coupling method. If 0 is given, omc3 
+        will try to determine the best candidate. If a number n>=1 is given, then some BPMs are 
+        skipped and the n-th following BPM downstream is used for the pairing.
+
+        Flags: **--coupling_pairing**
+        Choices: ``(0, n>=1)``
+        Default: ``0``.
       - **nonlinear**: Calculate higher order RDTs or CRDT
 
         Flags: **--nonlinear**
@@ -508,6 +515,11 @@ def optics_params():
     params.add_parameter(name="coupling_method", type=int,
                          choices=(0, 1, 2), default=OPTICS_DEFAULTS["coupling_method"],
                          help="Analysis option for coupling: disabled, 1 BPM or 2 BPMs method")
+    params.add_parameter(name="coupling_pairing", type=int,
+                         default=OPTICS_DEFAULTS["coupling_pairing"],
+                         help="Pairing mode for 2 BPM coupling method. If 0 is given, omc3 will try to "
+                              "determine the best candidate. If a number n>=1 is given, then some BPMs are skipped "
+                              "and the n-th following BPM downstream is used for the pairing.")
     params.add_parameter(name="range_of_bpms", type=int,
                          choices=(5, 7, 9, 11, 13, 15),  default=OPTICS_DEFAULTS["range_of_bpms"],
                          help="Range of BPMs for beta from phase calculation")
@@ -553,6 +565,7 @@ def optics_params():
 
 OPTICS_DEFAULTS = {
         "coupling_method": 2,
+        "coupling_pairing": 0,
         "range_of_bpms": 11,
         "compensation": "model",
 }

omc3/optics_measurements/coupling.py

@@ -113,8 +113,8 @@ def calculate_coupling(
         )
 
     LOGGER.debug("Finding BPM pairs for momentum reconstruction")
-    bpm_pairs_x, deltas_x = _find_pair(phases_x)
-    bpm_pairs_y, deltas_y = _find_pair(phases_y)
+    bpm_pairs_x, deltas_x = _find_pair(phases_x, meas_input.coupling_pairing)
+    bpm_pairs_y, deltas_y = _find_pair(phases_y, meas_input.coupling_pairing)
 
     LOGGER.debug("Computing complex lines from spectra")
     A01: np.ndarray = 0.5 * _get_complex_line(
@@ -217,19 +217,37 @@ def compensate_rdts_ryoichi():
 
 # ----- Helpers ----- #
 
-# def _take_next(phases: tfs.TfsDataFrame, shift: int = 1):
-#     """
-#     Takes the following BPM for momentum reconstruction by a given shift.
-#
-#     Args:
-#         phases (tfs.TfsDataFrame): Dataframe matrix of phase advances, as calculated in phase.py.
-#         shift (int): ???
-#     """
-#     indices = np.roll(np.arange(phases.to_numpy().shape[0]), shift)
-#     return indices, phases.to_numpy()[np.arange(phases.to_numpy().shape[0]), indices] - 0.25
+def _find_pair(phases: tfs.TfsDataFrame, mode: int = 1):
+    """
+    Does the BPM pairing for coupling calculation.
+
+    Args:
+        mode (int): Value to determine the BPM pairing. If ``0`` is given,
+            tries to find the best candidate. If a value ``n>=1`` is given,
+            then takes the n-th following BPM downstream for the pairing.
+    """
+
+    if mode == 0:
+        return _find_candidate(phases)
+    else:
+        return _take_next(phases, mode)
+
+
+def _take_next(phases: tfs.TfsDataFrame, shift: int = 1):
+    """
+    Takes the following BPM for momentum reconstruction by a given shift.
+
+    Args:
+        phases (tfs.TfsDataFrame): Dataframe matrix of phase advances, as calculated in phase.py.
+        shift (int): Value to determine the BPM pairing. If ``0`` is given,
+           tries to find the best candidate. If a value ``n>=1`` is given,
+           then takes the n-th following BPM downstream for the pairing.
+   """
+    indices = np.roll(np.arange(phases.to_numpy().shape[0]), shift)
+    return indices, phases.to_numpy()[np.arange(phases.to_numpy().shape[0]), indices] - 0.25
 
 
-def _find_pair(phases: tfs.TfsDataFrame) -> Tuple[np.ndarray, np.ndarray]:
+def _find_candidate(phases: tfs.TfsDataFrame) -> Tuple[np.ndarray, np.ndarray]:
     """
     Finds the best candidate for momentum reconstruction.