Analyse rawdata using oscillation midpoint in FBBA.

Author: MichaelStubbings

src/dls_bba/fbba.py

@@ -1,6 +1,6 @@
 import logging as log
 from math import ceil
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 import numpy as np
 from cothread import Sleep
@@ -28,7 +28,7 @@ def __init__(self, machine: Machine) -> None:
         """
         super().__init__(machine)
 
-    def run(self, components_pair: List[Components]) -> RawData:
+    def run(self, components_pair: list[Components]) -> RawData:
         """The Fast BBA Process.
 
         Args:
@@ -37,8 +37,8 @@ def run(self, components_pair: List[Components]) -> RawData:
         Returns:
             The RawData object
         """
-        rawdata: Dict[str, Any] = {}
-        metadata: Dict[str, Any] = {}
+        rawdata: dict[str, Any] = {}
+        metadata: dict[str, Any] = {}
         config = self._machine.config.get_settings()
         metadata.update(config)
         metadata["method"] = "FastBBA"
@@ -52,7 +52,7 @@ def run(self, components_pair: List[Components]) -> RawData:
         for components in components_pair:
             log.debug(f"Component: {components}")
             for quadrupole, quad_name in zip(
-                components.quadrupoles, components.quadrupoles_names
+                components.quadrupoles, components.quadrupoles_names, strict=True
             ):
                 log.debug(f"Quad: {quad_name} of {components.quadrupoles_names}")
                 (
@@ -143,9 +143,9 @@ def select_data(
         self,
         data: np.ndarray,
         axis: str,
-        exc_data: Tuple[Excitation, Excitation],
+        exc_data: tuple[Excitation, Excitation],
         decimated: bool,
-    ) -> List[np.ndarray]:
+    ) -> list[np.ndarray]:
         """Extract FA data that covers the excitations exc_high and exc_low.
 
         The input data array should cover the full length of both excitations.
@@ -165,15 +165,12 @@ def select_data(
 
         exc_high, exc_low = exc_data
         # Note: array data must include the timestamps.
-        log.debug("Raw data shape: {}".format(data.shape))
+        log.debug(f"Raw data shape: {data.shape}")
+        log.debug(f"Timestamp range in raw data: {data[0, 0, 0]} - {data[-1, 0, 0]}")
+        log.debug(f"Excitation length: {exc_high.count}")
         log.debug(
-            "Timestamp range in raw data: {} - {}".format(data[0, 0, 0], data[-1, 0, 0])
-        )
-        log.debug("Excitation length: {}".format(exc_high.count))
-        log.debug(
-            "Trailing data to crop: {}.".format(
-                data[-1, 0, 0] - (exc_low.start_time + exc_low.count)
-            )
+            f"Trailing data to crop: "
+            f"{data[-1, 0, 0] - (exc_low.start_time + exc_low.count)}."
         )
         assert exc_high.count == exc_low.count, "Excitations different lengths"
         # Extract timestamps from data
@@ -186,7 +183,7 @@ def select_data(
         length = ceil(exc_high.count / 10) if decimated else exc_high.count
         high_data = data[high_start : high_start + length, :, plane]
         low_data = data[low_start : low_start + length, :, plane]
-        log.debug("Selected data shape: {} {}".format(high_data.shape, low_data.shape))
+        log.debug(f"Selected data shape: {high_data.shape} {low_data.shape}")
         assert high_data.shape == low_data.shape
         return [high_data, low_data]
 
@@ -260,8 +257,8 @@ def analyse(self, rawdata: RawData) -> Results:
         bpm_index = bpm_number - np.sum(
             enabled_bpms[:bpm_number] == False  # noqa false positive
         )
-        results: Dict[str, List[float]] = {}
-        plotting: Dict[str, Dict[str, np.ndarray]] = {}
+        results: dict[str, list[float]] = {}
+        plotting: dict[str, dict[str, np.ndarray]] = {}
 
         quad_names = []
         for key in data.keys():
@@ -289,9 +286,13 @@ def analyse(self, rawdata: RawData) -> Results:
                 good = q_diff.std(0) > q_diff.std(0).max() / 2
                 q_diff_good = q_diff[:, good]
 
-                # Use a single fit operation, then transform with the straight line equation
+                # Use a single fit operation, then transform with straight line equation
+                oscillation_midpoint = (
+                    q_high_clean[:, bpm_index] + q_low_clean[:, bpm_index]
+                ) / 2
+                oscillation_size = q_diff_good
                 fit = np.polynomial.polynomial.polyfit(
-                    q_high_clean[:, bpm_index], q_diff_good, 1
+                    oscillation_midpoint, oscillation_size, 1
                 )
                 p = np.array([1 / fit[1], -fit[0] / fit[1]]).T
 
@@ -302,8 +303,8 @@ def analyse(self, rawdata: RawData) -> Results:
 
                 # plotting data
                 plotting[key] = {
-                    "x": q_high_clean[:, bpm_index] / NM_TO_MM_UNIT_CONV,
-                    "y": q_diff_good / NM_TO_MM_UNIT_CONV,
+                    "x": oscillation_midpoint / NM_TO_MM_UNIT_CONV,
+                    "y": oscillation_size / NM_TO_MM_UNIT_CONV,
                 }
 
         offsets = self.create_offsets_dict(results, metadata)

src/dls_bba/simfbba.py

@@ -1,6 +1,6 @@
 import logging as log
 from math import ceil
-from typing import Any, Dict, List
+from typing import Any
 
 import numpy as np
 from cothread import Sleep
@@ -29,7 +29,7 @@ def __init__(self, machine: Machine) -> None:
         """
         super().__init__(machine)
 
-    def run(self, components_pair: List[Components]) -> RawData:
+    def run(self, components_pair: list[Components]) -> RawData:
         """The Simultaneous Fast BBA Process.
 
         Args:
@@ -38,8 +38,8 @@ def run(self, components_pair: List[Components]) -> RawData:
         Returns:
             The RawData object.
         """
-        rawdata: Dict[str, Any] = {}
-        metadata: Dict[str, Any] = {}
+        rawdata: dict[str, Any] = {}
+        metadata: dict[str, Any] = {}
         config = self._machine.config.get_settings()
         metadata.update(config)
         metadata["method"] = "SimFastBBA"
@@ -51,12 +51,15 @@ def run(self, components_pair: List[Components]) -> RawData:
 
         log.info(f"BPM: {components_pair[0].bpm_name}")
         for quadrupole, quad_name in zip(
-            components_pair[0].quadrupoles, components_pair[0].quadrupoles_names
+            components_pair[0].quadrupoles,
+            components_pair[0].quadrupoles_names,
+            strict=True,
         ):
             log.debug(f"BPM: {components_pair[0].bpm_name}")
             log.debug(f"Quad: {quad_name} of {components_pair[0].quadrupoles_names}")
             log.debug(
-                f"Corrector1: {components_pair[0].corrector_name}, Corrector2: {components_pair[1].corrector_name}"
+                f"Corrector1: {components_pair[0].corrector_name}, "
+                f"Corrector2: {components_pair[1].corrector_name}"
             )
             (
                 quad_start,
@@ -184,9 +187,9 @@ def select_data(
         self,
         data: np.ndarray,
         plane_index: int,
-        exc_data: List[Excitation],
+        exc_data: list[Excitation],
         decimated: bool,
-    ) -> List[np.ndarray]:
+    ) -> list[np.ndarray]:
         """Extract FA data that covers the excitations exc_high and exc_low.
 
         The input data array should cover the full length of both excitations.
@@ -201,15 +204,12 @@ def select_data(
         """
         # Note: array data must include the timestamps.
         exc_high, exc_low = exc_data
-        log.debug("Raw data shape: {}".format(data.shape))
+        log.debug(f"Raw data shape: {data.shape}")
+        log.debug(f"Timestamp range in raw data: {data[0, 0, 0]} - {data[-1, 0, 0]}")
+        log.debug(f"Excitation length: {exc_high.count}")
         log.debug(
-            "Timestamp range in raw data: {} - {}".format(data[0, 0, 0], data[-1, 0, 0])
-        )
-        log.debug("Excitation length: {}".format(exc_high.count))
-        log.debug(
-            "Trailing data to crop: {}.".format(
-                data[-1, 0, 0] - (exc_low.start_time + exc_low.count)
-            )
+            f"Trailing data to crop: "
+            f"{data[-1, 0, 0] - (exc_low.start_time + exc_low.count)}."
         )
         assert exc_high.count == exc_low.count, "Excitations different lengths"
         # Extract timestamps from data
@@ -222,7 +222,7 @@ def select_data(
         length = ceil(exc_high.count / 10) if decimated else exc_high.count
         high_data = data[high_start : high_start + length, :, plane_index]
         low_data = data[low_start : low_start + length, :, plane_index]
-        log.debug("Selected data shape: {} {}".format(high_data.shape, low_data.shape))
+        log.debug(f"Selected data shape: {high_data.shape} {low_data.shape}")
         assert high_data.shape == low_data.shape
         return [high_data, low_data]
 
@@ -295,8 +295,8 @@ def analyse(self, rawdata: RawData) -> Results:
         bpm_index = bpm_number - np.sum(
             enabled_bpms[:bpm_number] == False  # noqa false positive
         )
-        results: Dict[str, List[float]] = {}
-        plotting: Dict[str, Dict[str, np.ndarray]] = {}
+        results: dict[str, list[float]] = {}
+        plotting: dict[str, dict[str, np.ndarray]] = {}
 
         quad_names = []
         for key in data.keys():
@@ -323,9 +323,13 @@ def analyse(self, rawdata: RawData) -> Results:
                 good = q_diff.std(0) > q_diff.std(0).max() / 2
                 q_diff_good = q_diff[:, good]
 
-                # Use a single fit operation, then transform with the straight line equation
+                # Use a single fit operation, then transform with straight line equation
+                oscillation_midpoint = (
+                    q_high_clean[:, bpm_index] + q_low_clean[:, bpm_index]
+                ) / 2
+                oscillation_size = q_diff_good
                 fit = np.polynomial.polynomial.polyfit(
-                    q_high_clean[:, bpm_index], q_diff_good, 1
+                    oscillation_midpoint, oscillation_size, 1
                 )
                 p = np.array([1 / fit[1], -fit[0] / fit[1]]).T
 
@@ -335,8 +339,8 @@ def analyse(self, rawdata: RawData) -> Results:
 
                 # plotting data
                 plotting[key] = {
-                    "x": q_high_clean[:, bpm_index] / NM_TO_MM_UNIT_CONV,
-                    "y": q_diff_good / NM_TO_MM_UNIT_CONV,
+                    "x": oscillation_midpoint / NM_TO_MM_UNIT_CONV,
+                    "y": oscillation_size / NM_TO_MM_UNIT_CONV,
                 }
 
         offsets = self.create_offsets_dict(results, metadata)