kymotrackgroup: add multi-kymo dwell time analysis

lumicks/pylake/kymotracker/kymotrack.py

@@ -1064,6 +1064,76 @@ def _tracks_by_kymo(self):
             KymoTrackGroup([track for track in self if track._kymo == kymo]) for kymo in self._kymos
         ]
 
+    @staticmethod
+    def _extract_dwelltime_data_from_groups(groups, exclude_ambiguous_dwells):
+        """Compute data needed for dwelltime analysis from a dictionary of KymoTrackGroups.
+
+        This data consists of dwelltimes and observation limits per track. Note that dwelltimes of zero
+        are automatically dropped.
+
+        Parameters
+        ----------
+        groups : iterable of KymoTrackGroup
+            An iterable which provides a sequence of KymoTrackGroup. Note that each group can only
+            have one `Kymo` associated with it.
+        exclude_ambiguous_dwells : bool
+            Determines whether to exclude dwelltimes which are not exactly determined. If `True`,
+            tracks which start in the first frame or end in the last frame of the kymograph are not
+            used in the analysis, since the exact start/stop times of the binding event are not
+            definitively known.
+
+        Returns
+        -------
+        dwelltimes : numpy.ndarray
+            Dwelltimes
+        min_obs : numpy.ndarray
+            List of minimum observation times extracted from the kymos
+        max_obs : numpy.ndarray
+            List of maximum observation time
+        removed_zeros : bool
+            Whether zeroes were dropped
+
+        Raises
+        ------
+        ValueError
+            if one of the KymoTrackGroups has more than one `Kymo` associated with it
+        """
+        removed_zeros = False
+
+        def extract_dwelltime_data(group):
+            nonlocal removed_zeros
+
+            if len(group._kymos) > 1:
+                raise ValueError("This group has more than one Kymo associated with it.")
+
+            tracks = (
+                filter(KymoTrack._check_ends_are_defined, group)
+                if exclude_ambiguous_dwells
+                else group
+            )
+            dwelltimes_sec = np.array([track.seconds[-1] - track.seconds[0] for track in tracks])
+            nonzero_dwelltimes_sec = dwelltimes_sec[dwelltimes_sec > 0]
+            removed_zeros = removed_zeros or len(nonzero_dwelltimes_sec) != len(dwelltimes_sec)
+
+            # Gracefully handle empty groups
+            if nonzero_dwelltimes_sec.size == 0:
+                return np.empty((3, 0))
+
+            min_observation_time = np.min(nonzero_dwelltimes_sec)
+            max_observation_time = group[0]._image.shape[1] * group[0]._line_time_seconds
+
+            return np.vstack(
+                list(
+                    np.broadcast(nonzero_dwelltimes_sec, min_observation_time, max_observation_time)
+                )
+            ).T
+
+        dwelltimes, min_obs, max_obs = np.hstack(
+            [extract_dwelltime_data(tracks) for tracks in groups]
+        )
+
+        return dwelltimes, min_obs, max_obs, removed_zeros
+
     def fit_binding_times(
         self, n_components, *, exclude_ambiguous_dwells=True, tol=None, max_iter=None
     ):
@@ -1084,20 +1154,19 @@ def fit_binding_times(
             The maximum number of iterations to perform. This parameter is forwarded as the `maxiter` argument
             to `scipy.minimize(method="L-BFGS-B")`.
         """
-        self._validate_single_source("Dwelltime analysis")
+        if not len(self):
+            raise RuntimeError("No tracks available for analysis")
 
         if n_components not in (1, 2):
             raise ValueError(
                 "Only 1- and 2-component exponential distributions are currently supported."
             )
 
-        tracks = (
-            filter(KymoTrack._check_ends_are_defined, self) if exclude_ambiguous_dwells else self
+        dwelltimes, min_obs, max_obs, removed_zeros = self._extract_dwelltime_data_from_groups(
+            self._tracks_by_kymo(), exclude_ambiguous_dwells
         )
-        dwelltimes_sec = np.array([track.seconds[-1] - track.seconds[0] for track in tracks])
 
-        nonzero_dwelltimes_sec = dwelltimes_sec[dwelltimes_sec > 0]
-        if len(nonzero_dwelltimes_sec) != len(dwelltimes_sec):
+        if removed_zeros:
             warnings.warn(
                 RuntimeWarning(
                     "Some dwell times are zero. A dwell time of zero indicates that some of the "
@@ -1109,17 +1178,14 @@ def fit_binding_times(
                 stacklevel=2,
             )
 
-        if nonzero_dwelltimes_sec.size == 0:
+        if dwelltimes.size == 0:
             raise RuntimeError("No tracks available for analysis")
 
-        min_observation_time = np.min(nonzero_dwelltimes_sec)
-        max_observation_time = self[0]._image.shape[1] * self[0]._line_time_seconds
-
         return DwelltimeModel(
-            nonzero_dwelltimes_sec,
+            dwelltimes,
             n_components,
-            min_observation_time=min_observation_time,
-            max_observation_time=max_observation_time,
+            min_observation_time=min_obs,
+            max_observation_time=max_obs,
             tol=tol,
             max_iter=max_iter,
         )
@@ -1335,7 +1401,7 @@ def ensemble_diffusion(self, method, *, max_lag=None):
         RuntimeWarning
             if `method == "cve"` and the source kymographs do not have the same line times
             or pixel sizes. As a result, the localization variance and variance of the localization
-            variance are not be available. Estimates that are unavailable are returned as `np.nan`.
+            variance are not available. Estimates that are unavailable are returned as `np.nan`.
 
         References
         ----------

lumicks/pylake/kymotracker/tests/test_kymotrack.py

@@ -772,6 +772,96 @@ def test_fit_binding_times_empty():
         KymoTrackGroup([]).fit_binding_times(1)
 
 
+def test_multi_kymo_dwell():
+    kymos = [
+        _kymo_from_array(np.zeros((10, 10, 3)), "rgb", line_time_seconds=time, pixel_size_um=size)
+        for time, size in ((10e-4, 0.05), (10e-3, 0.1))
+    ]
+
+    k1 = KymoTrack(np.array([1, 2, 3]), np.array([1, 2, 3]), kymos[0], "red")
+    k2 = KymoTrack(np.array([2, 3, 4, 5]), np.array([2, 3, 4, 5]), kymos[0], "red")
+    k3 = KymoTrack(np.array([3, 4, 5]), np.array([3, 4, 5]), kymos[1], "red")
+    k4 = KymoTrack(np.array([4, 5, 6]), np.array([4, 5, 6]), kymos[1], "red")
+    k5 = KymoTrack(np.array([7]), np.array([7]), kymos[1], "red")
+
+    # Normal use case
+    dwell, obs_min, obs_max, removed_zeroes = KymoTrackGroup._extract_dwelltime_data_from_groups(
+        [KymoTrackGroup([k1, k2]), KymoTrackGroup([k3, k4])], False
+    )
+    np.testing.assert_allclose(dwell, [0.002, 0.003, 0.02, 0.02])
+    np.testing.assert_allclose(obs_min, [0.002, 0.002, 0.02, 0.02])
+    np.testing.assert_allclose(obs_max, [0.01, 0.01, 0.1, 0.1])
+    assert removed_zeroes is False
+
+    # Drop one "empty" dwell
+    dwell, obs_min, obs_max, removed_zeroes = KymoTrackGroup._extract_dwelltime_data_from_groups(
+        [KymoTrackGroup([k1, k2]), KymoTrackGroup([k3, k4, k5])], False
+    )
+    np.testing.assert_allclose(dwell, [0.002, 0.003, 0.02, 0.02])
+    np.testing.assert_allclose(obs_min, [0.002, 0.002, 0.02, 0.02])
+    np.testing.assert_allclose(obs_max, [0.01, 0.01, 0.1, 0.1])
+    assert removed_zeroes is True
+
+    # Test with empty group
+    dwell, obs_min, obs_max, removed_zeroes = KymoTrackGroup._extract_dwelltime_data_from_groups(
+        [KymoTrackGroup([k1, k2]), KymoTrackGroup([])], False
+    )
+    np.testing.assert_allclose(dwell, [0.002, 0.003])
+    np.testing.assert_allclose(obs_min, [0.002, 0.002])
+    np.testing.assert_allclose(obs_max, [0.01, 0.01])
+    assert removed_zeroes is False
+
+    # Test with group that is empty after filtering
+    dwell, obs_min, obs_max, removed_zeroes = KymoTrackGroup._extract_dwelltime_data_from_groups(
+        [KymoTrackGroup([k1, k2]), KymoTrackGroup([k5])], False
+    )
+    np.testing.assert_allclose(dwell, [0.002, 0.003])
+    np.testing.assert_allclose(obs_min, [0.002, 0.002])
+    np.testing.assert_allclose(obs_max, [0.01, 0.01])
+    assert removed_zeroes is True
+
+    dwell, obs_min, obs_max, removed_zeroes = KymoTrackGroup._extract_dwelltime_data_from_groups(
+        [KymoTrackGroup([k5]), KymoTrackGroup([k5])], False
+    )
+    np.testing.assert_allclose(dwell, [])
+    np.testing.assert_allclose(obs_min, [])
+    np.testing.assert_allclose(obs_max, [])
+    assert removed_zeroes is True
+
+
+def test_multi_kymo_dwelltime_analysis():
+    """This test tests only the happy path since others test more specific edge cases"""
+    np.random.seed(123)
+
+    shared_args = {
+        "image": np.empty((1, 5000, 3)),
+        "color_format": "rgb",
+        "start": np.int64(20e9),
+        "pixel_size_um": 2,
+    }
+
+    kymo1 = _kymo_from_array(**shared_args, line_time_seconds=0.01)
+    kymo2 = _kymo_from_array(**shared_args, line_time_seconds=0.02)
+    args = [np.array([1.0, 1.0]), kymo1, "red"]
+    group1 = KymoTrackGroup(
+        [
+            KymoTrack(np.array([1, int(round(t / kymo1.line_time_seconds)) + 1]), *args)
+            for t in np.random.exponential(2, size=(200,))
+        ],
+    )
+    args = [np.array([1.0, 1.0]), kymo2, "red"]
+    group2 = KymoTrackGroup(
+        [
+            KymoTrack(np.array([1, int(round(t / kymo2.line_time_seconds)) + 1]), *args)
+            for t in np.random.exponential(6, size=(200,))
+        ],
+    )
+
+    model = (group1 + group2).fit_binding_times(2, tol=1e-16)
+    np.testing.assert_allclose(model.lifetimes, (2.05754716, 6.89790037))
+    np.testing.assert_allclose(model.amplitudes, (0.61360417, 0.38639583))
+
+
 @pytest.mark.parametrize("method,max_lags", [("ols", 2), ("ols", None), ("gls", 2), ("gls", None)])
 def test_kymotrack_group_diffusion(blank_kymo, method, max_lags):
     """Tests whether we can call this function at the diffusion level"""

lumicks/pylake/kymotracker/tests/test_kymotrackgroup_sources.py

@@ -255,13 +255,6 @@ def test_multisource_not_implemented(kymos, coordinates):
             tracks, window=2, refine_missing_frames=False, overlap_strategy="ignore"
         )
 
-    with pytest.raises(
-        NotImplementedError,
-        match=(
-            r"Dwelltime analysis is not supported. This group contains tracks from 2 source kymographs."
-        ),
-    ):
-        tracks.fit_binding_times(n_components=1)
 
 def test_tracks_by_kymo(kymos, coordinates):
     time_indices, position_indices = coordinates
@@ -279,4 +272,3 @@ def make_tracks(kymo):
         assert len(tracks_group) == len(tracks_raw)
         for track_group, track_raw in zip(tracks_group, tracks_raw):
             id(track_group) == id(track_raw)
-