perf: run tracking cross-correlation in parallel threads

Author: Master92

src/pybox/gui/tracking_plots.py

@@ -27,7 +27,7 @@ class _TrackingResult:
     time: np.ndarray
     gyros: list[np.ndarray]      # [roll, pitch, yaw]
     setpoints: list[np.ndarray]  # [roll, pitch, yaw]
-    delays: list[float]          # ms per axis
+    delays: list[float | None]   # ms per axis
 
 
 AXIS_NAMES = ["Roll", "Pitch", "Yaw"]
@@ -37,43 +37,28 @@ class _TrackingWorker(QThread):
     """Compute cross-correlation delays and downsample in a background thread."""
     finished = Signal(object)  # _TrackingResult or None
 
-    def __init__(self, entry: LogEntry, parent=None):
+    def __init__(self, entry: LogEntry, axis: int, setpoints: np.ndarray, gyros: np.ndarray, parent=None):
         super().__init__(parent)
         self._entry = entry
+        self._axis = axis
+        self._setpoints = setpoints
+        self._gyros = gyros
 
     def run(self):
         entry = self._entry
-        time_s, gyro_r, gyro_p, gyro_y = entry.gyro_arrays()
-        sp_r, sp_p, sp_y = entry.setpoint_arrays()
+        time_s, _, _, _ = entry.gyro_arrays()
 
         if len(time_s) == 0:
-            self.finished.emit(None)
+            self.finished.emit([self._axis, None])
             return
 
         mask = entry.time_mask()
         t_full = time_s[mask]
-        gyros_full = [gyro_r[mask], gyro_p[mask], gyro_y[mask]]
-        setpoints_full = [sp_r[mask], sp_p[mask], sp_y[mask]]
 
         dt = np.median(np.diff(t_full)) if len(t_full) > 2 else 0.001
-        delays = []
-        for axis in range(3):
-            delays.append(TrackingPlots._xcorr_delay(setpoints_full[axis], gyros_full[axis], dt))
-
-        # Downsample for plotting
-        t = t_full
-        gyros = gyros_full
-        setpoints = setpoints_full
-        max_pts = 20000
-        if len(t) > max_pts:
-            step = len(t) // max_pts
-            t = t[::step]
-            gyros = [g[::step] for g in gyros]
-            setpoints = [s[::step] for s in setpoints]
+        delays = TrackingPlots._xcorr_delay(self._setpoints, self._gyros, dt)
 
-        self.finished.emit(_TrackingResult(
-            time=t, gyros=gyros, setpoints=setpoints, delays=delays,
-        ))
+        self.finished.emit([self._axis, delays])
 
 
 class TrackingPlots(QWidget):
@@ -112,7 +97,8 @@ def __init__(self, parent=None):
         self._curves: dict[tuple[str, int], pg.PlotDataItem] = {}
         # key: ("gyro"|"sp", axis)
         self._cache_key: tuple | None = None
-        self._worker: _TrackingWorker | None = None
+        self._workers: list[_TrackingWorker | None] | None = []
+        self._results: _TrackingResult | None = None
 
         for i, name in enumerate(AXIS_NAMES):
             if i > 0:
@@ -133,18 +119,42 @@ def __init__(self, parent=None):
         for i in range(1, 3):
             self._plots[i].setXLink(self._plots[0])
 
+    def _prepare_data(self, entry: LogEntry) -> tuple[list[np.ndarray], list[np.ndarray]]:
+        mask = entry.time_mask()
+        time_s, gyro_r, gyro_p, gyro_y = entry.gyro_arrays()
+        sp_r, sp_p, sp_y = entry.setpoint_arrays()
+        t_full = time_s[mask]
+        gyros_full = [gyro_r[mask], gyro_p[mask], gyro_y[mask]]
+        setpoints_full = [sp_r[mask], sp_p[mask], sp_y[mask]]
+
+        # Downsample for plotting
+        t = t_full
+        gyros = gyros_full
+        setpoints = setpoints_full
+        max_pts = 20000
+        if len(t) > max_pts:
+            step = len(t) // max_pts
+            t = t[::step]
+            gyros = [g[::step] for g in gyros]
+            setpoints = [s[::step] for s in setpoints]
+
+        self._results = _TrackingResult(time=t, gyros=gyros, setpoints=setpoints, delays=[None, None, None])
+
+        return gyros_full, setpoints_full
+
     @staticmethod
     def _entry_key(entry: LogEntry) -> tuple:
         return (entry.file_path, entry.log_index, entry.time_start_s, entry.time_end_s)
 
     def update_plots(self, entry: LogEntry | None):
         """Display setpoint/gyro tracking for the given (selected) log entry."""
         # Cancel any running worker
-        if self._worker is not None and self._worker.isRunning():
-            self._worker.finished.disconnect()
-            self._worker.quit()
-            self._worker.wait(500)
-            self._worker = None
+        for worker in self._workers:
+            if worker is not None and worker.isRunning():
+                worker.finished.disconnect()
+                worker.quit()
+                worker.wait(500)
+        self._workers = []
 
         if entry is not None:
             key = self._entry_key(entry)
@@ -159,11 +169,14 @@ def update_plots(self, entry: LogEntry | None):
         if entry is None:
             return
 
+        gyros_full, setpoints_full = self._prepare_data(entry)
+
         # Show loading indicator and start background computation
         self._info_label.setText(self.tr("Computing tracking delays..."))
-        self._worker = _TrackingWorker(entry, parent=self)
-        self._worker.finished.connect(self._on_worker_finished)
-        self._worker.start()
+        for i in range(3):
+            self._workers.append(_TrackingWorker(entry, i, setpoints_full[i], gyros_full[i], parent=self))
+            self._workers[i].finished.connect(self._on_worker_finished)
+            self._workers[i].start()
 
     @staticmethod
     def _xcorr_delay(setpoint: np.ndarray, gyro: np.ndarray, dt: float) -> float:
@@ -183,28 +196,35 @@ def _xcorr_delay(setpoint: np.ndarray, gyro: np.ndarray, dt: float) -> float:
         lag = int(np.argmax(search))
         return lag * dt * 1000.0  # ms
 
-    def _on_worker_finished(self, result: _TrackingResult | None):
+    def _on_worker_finished(self, result: tuple[int, list[float] | None]):
         """Called on main thread when background computation is done."""
-        self._worker = None
-        if result is None:
+        axis = result[0]
+        self._workers[axis] = None
+        if result[1] is None:
             self._info_label.setText("")
             return
-        self._draw_result(result)
+        self._results.delays[axis] = result[1]
+        self._draw_result(self._results)
 
     def _draw_result(self, r: _TrackingResult):
         """Draw pre-computed tracking data (runs on main thread)."""
         axis_colors = ["#e6194b", "#3cb44b", "#4363d8"]  # R, G, B
         t_theme = current_theme()
 
         for axis in range(3):
-            if len(r.gyros[axis]) < 10:
+            # If this axis is already contained in self._curves, skip it
+            if self._curves.get(("gyro", axis)) is not None:
+                continue
+
+            if len(r.gyros[axis]) < 10 or r.delays[axis] is None:
                 continue
 
             color = axis_colors[axis]
 
             # Update title with delay
+            time_str = f"{r.delays[axis]:.1f} ms" if r.delays[axis] is not None else "N/A"
             self._plots[axis].setTitle(
-                f"{AXIS_NAMES[axis]}  –  delay: {r.delays[axis]:.1f} ms",
+                f"{AXIS_NAMES[axis]}  –  delay: {time_str}",
                 color=t_theme.plot_title_color, size="12pt",
             )
 
@@ -223,8 +243,9 @@ def _draw_result(self, r: _TrackingResult):
             self._curves[("gyro", axis)] = gyro_curve
 
         # Update info label
+        time_str = [f"{r.delays[axis]:.1f} ms" if r.delays [axis] is not None else "N/A" for axis in range(3)]
         self._info_label.setText(
-            f"Delay  –  Roll: {r.delays[0]:.1f} ms  |  Pitch: {r.delays[1]:.1f} ms  |  Yaw: {r.delays[2]:.1f} ms"
+            f"Delay  –  {AXIS_NAMES[0]}: {time_str[0]}  |  {AXIS_NAMES[1]}: {time_str[1]}  |  {AXIS_NAMES[2]}: {time_str[2]}"
         )
 
         # Auto-range