From b01daf9cde6120d5d24499216e81cef52b8e2739 Mon Sep 17 00:00:00 2001
From: Joep Vanlier <joep.vanlier@gmail.com>
Date: Fri, 24 Jan 2025 18:20:24 +0100
Subject: [PATCH] calibration: fallback for driving data

---
 lumicks/pylake/calibration.py                 |  8 +++-
 .../force_calibration/calibration_item.py     | 14 ++++++-
 lumicks/pylake/group.py                       | 42 ++++++++++++++-----
 3 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/lumicks/pylake/calibration.py b/lumicks/pylake/calibration.py
index 3d9caae31..9ba9c841d 100644
--- a/lumicks/pylake/calibration.py
+++ b/lumicks/pylake/calibration.py
@@ -90,7 +90,9 @@ def filter_calibration(self, start, stop):
         )
 
     @staticmethod
-    def from_field(hdf5, force_channel, force_slice=None) -> "ForceCalibrationList":
+    def from_field(
+        hdf5, force_channel, force_slice=None, driving_slice=None
+    ) -> "ForceCalibrationList":
         """Fetch force calibration data from the HDF5 file
 
         Parameters
@@ -99,6 +101,10 @@ def from_field(hdf5, force_channel, force_slice=None) -> "ForceCalibrationList":
             A Bluelake HDF5 file.
         force_channel : str
             Calibration field to access (e.g. "Force 1x").
+        force_slice : Slice
+            Slice with force calibration data
+        driving_slice : Slice
+            Slice that has driving input data
         """
 
         def make_slice(dset, field, y_label, title) -> Slice:
diff --git a/lumicks/pylake/force_calibration/calibration_item.py b/lumicks/pylake/force_calibration/calibration_item.py
index 31f9e2403..49087ea88 100644
--- a/lumicks/pylake/force_calibration/calibration_item.py
+++ b/lumicks/pylake/force_calibration/calibration_item.py
@@ -20,12 +20,14 @@ def __init__(
         sum_voltage=None,
         driving=None,
         force_slice=None,
+        driving_slice=None,
     ):
         super().__init__(dictionary)
         self._voltage = voltage
         self._sum_voltage = sum_voltage
         self._driving = driving
         self._force_slice = force_slice
+        self._driving_slice = driving_slice
 
     @property
     def voltage(self):
@@ -54,7 +56,17 @@ def sum_voltage(self):
     @property
     def driving(self):
         """Driving signal used for active calibration"""
-        return self._driving if self._driving is not None else empty_slice
+        if self._driving is not None:
+            return self._driving
+        elif (
+            self.active_calibration
+            and self._driving_slice
+            and self._driving_slice.start <= self.start
+            and self._driving_slice.stop >= self.stop
+        ):
+            return self._driving_slice[self.start : self.stop]
+
+        return empty_slice
 
     @staticmethod
     def _verify_full(method):
diff --git a/lumicks/pylake/group.py b/lumicks/pylake/group.py
index 232f44cf7..08f41182d 100644
--- a/lumicks/pylake/group.py
+++ b/lumicks/pylake/group.py
@@ -4,6 +4,36 @@
 from .calibration import ForceCalibrationList
 
 
+def process_calibrated_channel(cls, item, lk_file, item_name):
+    """Process a calibrated channel
+
+    Parameters
+    ----------
+    cls : TimeTags | TimeSeries | Continuous
+        Channel class
+    item : h5py.Dataset
+        Item to process
+    lk_file : lumicks.pylake.File
+        Pylake file handle
+    item_name : str
+        Name of the channel (e.g. "Force 1x")
+    """
+    force_slice = cls.from_dataset(
+        item,
+        calibration=ForceCalibrationList.from_field(lk_file.h5, item_name),
+    )
+    driving_channel = f"Nanostage position/{item_name[-1].upper()}"
+    driving_slice = lk_file[driving_channel] if driving_channel in lk_file else None
+
+    return cls.from_dataset(
+        item,
+        calibration=ForceCalibrationList.from_field(
+            lk_file.h5, item_name, force_slice, driving_slice
+        ),
+        y_label="Force (pN)",
+    )
+
+
 class Group:
     """A thin wrapper around an HDF5 group with a Bluelake-specific `__getitem__`
 
@@ -53,17 +83,7 @@ def __getitem__(self, item):
                 cls = channel_class(thing)
 
                 if item_type in ("Force HF", "Force LF"):
-                    force_slice = cls.from_dataset(
-                        thing,
-                        calibration=ForceCalibrationList.from_field(self._lk_file.h5, item_name),
-                    )
-                    return cls.from_dataset(
-                        thing,
-                        calibration=ForceCalibrationList.from_field(
-                            self._lk_file.h5, item_name, force_slice
-                        ),
-                        y_label="Force (pN)",
-                    )
+                    return process_calibrated_channel(cls, thing, self._lk_file, item_name)
 
                 return cls.from_dataset(thing)