dynamic valve calibration times

iblrig/gui/valve.py

@@ -47,10 +47,10 @@ def values(self, values: ValveValues):
 
     def update(self):
         self._points.setData(x=self.values.open_times_ms, y=self.values.volumes_ul)
-        if len(self.values.open_times_ms) < 2:
+        if len(self.values.open_times_ms) < 1:
             self._curve.setData(x=[], y=[])
         else:
-            time_range = list(np.linspace(self.values.open_times_ms[0], self.values.open_times_ms[-1], 100))
+            time_range = list(np.linspace(0, self.values.open_times_ms[-1], 100))
             self._curve.setData(x=time_range, y=self.values.ms2ul(time_range))
 
     def clear(self):
@@ -240,13 +240,21 @@ def clear_calibration(self):
         self._next_calibration_time = self.get_next_calibration_time()
 
     def get_next_calibration_time(self) -> float | None:
-        remaining_calibration_times = [
-            t for t in self.valve.new_calibration_open_times if t not in self.new_calibration.values.open_times_ms
-        ]
-        if len(remaining_calibration_times) > 0:
-            return max(remaining_calibration_times)
-        else:
+        if len(self.new_calibration.values.open_times_ms) == 0:
+            # we start with the longest opening time ...
+            return self.hw_settings.device_valve.WATER_CALIBRATION_MAX_OPEN_TIME_MS
+        elif (
+            # ... and stop, once we reached the defined lower volume threshold
+            min(self.new_calibration.values.volumes_ul)
+            <= self.hw_settings.device_valve.WATER_CALIBRATION_LOWER_VOLUME_THRESHOLD_UL
+        ):
             return None
+        else:
+            # calibration times are given by the previous time, multiplied by a reduction factor
+            return round(
+                min(self.new_calibration.values.open_times_ms)
+                * self.hw_settings.device_valve.WATER_CALIBRATION_OPEN_TIME_REDUCTION_FACTOR
+            )
 
     def initialize_scale(self, port: str) -> bool:
         if port is None:
@@ -346,7 +354,9 @@ def _on_tare_finished(self, success: bool):
 
     @QtCore.pyqtSlot()
     def calibrate(self):
-        n_samples = int(np.ceil(50 * max(self.valve.new_calibration_open_times) / self._next_calibration_time))
+        n_samples = int(
+            np.ceil(50 * self.hw_settings.device_valve.WATER_CALIBRATION_MAX_OPEN_TIME_MS / self._next_calibration_time)
+        )
         self.labelGuideText.setText(
             f'Getting {n_samples} samples for a valve opening time of {self._next_calibration_time} ms ...'
         )

iblrig/pydantic_definitions.py

@@ -124,7 +124,15 @@ class HardwareSettingsSound(BunchModel):
 class HardwareSettingsValve(BunchModel):
     WATER_CALIBRATION_DATE: date
     WATER_CALIBRATION_RANGE: list[PositiveFloat] = Field(min_items=2, max_items=2)  # type: ignore
-    WATER_CALIBRATION_N: PositiveInt = Field(ge=3, default=5)
+    WATER_CALIBRATION_MAX_OPEN_TIME_MS: PositiveInt = Field(
+        default=100, description='Longest opening time to start calibration with'
+    )
+    WATER_CALIBRATION_LOWER_VOLUME_THRESHOLD_UL: PositiveFloat = Field(
+        default=1.0, description='Lower volume threshold to reach in order to finish calibration'
+    )
+    WATER_CALIBRATION_OPEN_TIME_REDUCTION_FACTOR: PositiveFloat = Field(
+        default=0.66, lt=1, description='Factor to reduce opening time by with each iteration of the calibration'
+    )
     WATER_CALIBRATION_OPEN_TIMES: list[PositiveFloat] = Field(min_items=2)  # type: ignore
     WATER_CALIBRATION_WEIGHT_PERDROP: list[float] = Field(PositiveFloat, min_items=2)  # type: ignore
     FREE_REWARD_VOLUME_UL: PositiveFloat = 1.5

iblrig/valve.py

@@ -1,36 +1,88 @@
 import datetime
 import warnings
 from collections.abc import Sequence
 
 import numpy as np
 import scipy
 from numpy.polynomial import Polynomial
-from pydantic import PositiveFloat, validate_call
+from pydantic import PositiveFloat, validate_call, NonNegativeFloat
 
 from iblrig.pydantic_definitions import HardwareSettingsValve
 
 
 class ValveValues:
     _dtype = [('open_times_ms', float), ('weights_g', float)]
     _data: np.ndarray
     _polynomial: Polynomial
 
     def __init__(self, open_times_ms: Sequence[float], weights_g: Sequence[float]):
+        """
+        Initialize a ValveValues object.
+
+        Parameters
+        ----------
+        open_times_ms : Sequence[float]
+            Sequence of open times in milliseconds.
+        weights_g : Sequence[float]
+            Sequence of weights in grams corresponding to the open times.
+
+        Returns
+        -------
+        None
+        """
         self.clear_data()
         self.add_samples(open_times_ms, weights_g)
 
     @staticmethod
-    def _fcn(x: np.ndarray, a: float, b: float, c: float) -> np.ndarray:
-        return a + b * x + c * np.square(x)
+    def _fcn(x: np.ndarray, b: float, c: float) -> np.ndarray:
+        """
+        Function for fitting a quadratic curve.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Input data.
+        b : float
+            Coefficient for lineaer term.
+        c : float
+            Coefficient for quadratic term.
+
+        Returns
+        -------
+        np.ndarray
+            The result of the polynomial curve fitting.
+        """
+        return b * x + c * np.square(x)
 
     @validate_call
     def add_samples(self, open_times_ms: Sequence[PositiveFloat], weights_g: Sequence[PositiveFloat]):
+        """
+        Add samples of open times and weights to the data.
+
+        Parameters
+        ----------
+        open_times_ms : Sequence[PositiveFloat]
+            Sequence of open times in milliseconds.
+        weights_g : Sequence[PositiveFloat]
+            Sequence of weights in grams corresponding to the open times.
+
+        Returns
+        -------
+        None
+        """
         incoming = np.rec.fromarrays([open_times_ms, weights_g], dtype=self._dtype)
         self._data = np.append(self._data, incoming)
         self._data = np.sort(self._data)
         self._update_fit()
 
-    def clear_data(self):
+    def clear_data(self) -> None:
+        """
+        Clear all data stored in the object.
+
+        Returns
+        -------
+        None
+        """
         self._data = np.empty((0,), dtype=self._dtype)
         self._update_fit()
 
@@ -47,54 +99,127 @@
         return self._data['weights_g'] * 1e3
 
     def _update_fit(self) -> None:
+        """
+        Update the polynomial fit based on the data stored in the object.
+
+        Returns
+        -------
+        None
+        """
         if len(self._data) >= 2:
             with warnings.catch_warnings():
                 warnings.simplefilter('ignore')
                 try:
                     c, _ = scipy.optimize.curve_fit(
-                        self._fcn, self.open_times_ms, self.volumes_ul, bounds=([-np.inf, 0, 0], np.inf)
+                        self._fcn, self.open_times_ms, self.volumes_ul, bounds=([0, 0], np.inf)
                     )
                 except RuntimeError:
-                    c = [np.nan, np.nan, np.nan]
+                    c = [np.nan, np.nan]
         else:
-            c = [np.nan, np.nan, np.nan]
-        self._polynomial = Polynomial(coef=c)
+            c = [np.nan, np.nan]
+        self._polynomial = Polynomial(coef=np.append(0, c))
 
     @validate_call
-    def ul2ms(self, volume_ul: PositiveFloat) -> PositiveFloat:
+    def ul2ms(self, volume_ul: NonNegativeFloat) -> NonNegativeFloat:
+        """
+        Convert from volume to opening time.
+
+        Parameters
+        ----------
+        volume_ul : PositiveFloat
+            Volume in microliters.
+
+        Returns
+        -------
+        PositiveFloat
+            The corresponding opening time in milliseconds.
+        """
         return max((self._polynomial - volume_ul).roots())
 
     @validate_call
-    def ms2ul(self, volume_ul: PositiveFloat | list[PositiveFloat]) -> PositiveFloat | np.ndarray:
+    def ms2ul(self, volume_ul: NonNegativeFloat | list[NonNegativeFloat]) -> NonNegativeFloat | np.ndarray:
+        """
+        Convert from opening time to volume.
+
+        Parameters
+        ----------
+        volume_ul : PositiveFloat | list[PositiveFloat]
+            Opening time in milliseconds or a list of times in milliseconds.
+
+        Returns
+        -------
+        PositiveFloat | np.ndarray
+            The corresponding volume(s) in microliters.
+        """
         return self._polynomial(np.array(volume_ul))
 
 
 class Valve:
     def __init__(self, settings: HardwareSettingsValve):
+        """
+        Initialize a Valve object.
+
+        Parameters
+        ----------
+        settings : HardwareSettingsValve
+            The hardware settings for the valve.
+
+        Returns
+        -------
+        None
+        """
         self._settings = settings
         volumes_ul = settings.WATER_CALIBRATION_WEIGHT_PERDROP
         weights_g = [volume / 1e3 for volume in volumes_ul]
         self.values = ValveValues(settings.WATER_CALIBRATION_OPEN_TIMES, weights_g)
 
     @property
     def calibration_date(self) -> datetime.date:
+        """
+        Get the date of the valve's last calibration.
+
+        Returns
+        -------
+        datetime.date
+            The calibration date.
+        """
         return self._settings.WATER_CALIBRATION_DATE
 
     @property
-    def calibration_range(self) -> list[float, float]:
+    def calibration_range(self) -> list[float]:
+        """
+        Get the calibration range of the valve.
+
+        Returns
+        -------
+        np.ndarray
+            A list containing the minimum and maximum calibration values.
+        """
         return self._settings.WATER_CALIBRATION_RANGE
 
-    @property
-    def new_calibration_open_times(self) -> set[float]:
-        return set(np.linspace(self.calibration_range[0], self.calibration_range[1], self._settings.WATER_CALIBRATION_N))
-
     @property
     def free_reward_time(self) -> float:
-        return self._settings.FREE_REWARD_VOLUME_UL
+        """
+        Get the free reward time of the valve.
+
+        Returns
+        -------
+        float
+            The free reward time in seconds.
+        """
+        return self.values.ul2ms(self._settings.FREE_REWARD_VOLUME_UL) * 1000
 
     @property
     def settings(self) -> HardwareSettingsValve:
+        """
+        Get the current hardware settings of the valve.
+
+        Returns
+        -------
+        HardwareSettingsValve
+            The current hardware settings.
+        """
         settings = self._settings
-        settings.WATER_CALIBRATION_OPEN_TIMES = self.values.open_times_ms
-        settings.WATER_CALIBRATION_WEIGHT_PERDROP = self.values.volumes_ul
+        settings.WATER_CALIBRATION_OPEN_TIMES = list(self.values.open_times_ms)
+        settings.WATER_CALIBRATION_WEIGHT_PERDROP = list(self.values.volumes_us)
         return settings

settings/hardware_settings_template.yaml

@@ -28,6 +28,8 @@ device_microphone:
   BONSAI_WORKFLOW: devices/microphone/record_mic.bonsai
 device_valve:
   WATER_CALIBRATION_DATE: 2099-12-31
+  WATER_CALIBRATION_MAX_OPEN_TIME_MS: 100
+  WATER_CALIBRATION_LOWER_VOLUME_THRESHOLD_UL: 1.0
   WATER_CALIBRATION_OPEN_TIMES: [50, 100]
   WATER_CALIBRATION_RANGE: [40, 140]
   WATER_CALIBRATION_WEIGHT_PERDROP: [1.25, 2.75]