colocalization: add function to compute track proximity

lumicks/pylake/kymotracker/detail/track_proximity.py

@@ -14,7 +14,7 @@ class BoundingBox:
     _valid: bool = True
 
     @staticmethod
-    def from_track(track, time_slack=1, coord_slack=1):
+    def from_track(track, time_slack=1, coord_slack=1.0):
         if len(track.time_idx) == 0:
             return BoundingBox(0, 0, 0, 0, False)
 
@@ -27,10 +27,10 @@ def from_track(track, time_slack=1, coord_slack=1):
         )
 
     def overlaps_with(self, other):
+        """Test whether this bounding box overlaps with another bounding box"""
         if not (self._valid and other._valid):
             return False
 
-        """Test whether this bounding box overlaps with another bounding box"""
         if self.time_max < other.time_min:
             return False
 
@@ -44,3 +44,82 @@ def overlaps_with(self, other):
             return False
 
         return True
+
+
+def tracks_close(track1, track2, position_window, time_window):
+    """Check whether two tracks overlap (with some tolerance)
+
+    Parameters
+    ----------
+    track1, track2 : lumicks.pylake.kymotracker.kymotrack.KymoTrack
+        Kymotracks to compare
+    position_window : float
+        How much distance is there allowed to be between two detections (in pixels).
+    time_window : int
+        How much temporal distance is there allowed to be between two detections (in frame indices).
+
+    Returns
+    -------
+    bool
+        True when the tracks are close enough in time and space.
+    """
+
+    def dist(coord1, coord2):
+        return np.abs(np.asarray(coord2) - np.asarray(coord1)[:, np.newaxis])
+
+    coord_dists = dist(track1.coordinate_idx, track2.coordinate_idx)
+    time_dists = dist(track1.time_idx, track2.time_idx)
+    if np.any(np.logical_and(coord_dists <= position_window, time_dists <= time_window)):
+        return True
+
+    return False
+
+
+def find_colocalized(
+    tracks1,
+    tracks2,
+    position_window,
+    time_window,
+    *,
+    interpolate=True,
+):
+    """Returns the list of tracks which have points in close proximity
+
+    Parameters
+    ----------
+    tracks1, tracks2 : lumicks.pylake.kymotracker.KymoTrack.KymoTrackGroup
+        Groups of tracks
+    position_window : float
+        Spatial distance in pixels above which track points are not considered close anymore.
+    time_window : int
+        Temporal distance in pixels above which track points are not considered close anymore.
+    interpolate : bool
+        Interpolate track before starting.
+
+    Returns
+    -------
+    list of tuple
+        Set of pairs of track indices in close enough proximity to be candidates for interaction.
+        Note that a track can appear in multiple pairs.
+    """
+    track_groups = [tracks1, tracks2]
+
+    bboxes = [
+        [BoundingBox.from_track(track, time_window, position_window) for track in tracks]
+        for tracks in track_groups
+    ]
+
+    if interpolate:
+        track_groups = [[t.interpolate() for t in tracks] for tracks in track_groups]
+
+    colocalizations = set()
+    for idx1, (track1, bbox1) in enumerate(zip(track_groups[0], bboxes[0])):
+        for idx2, (track2, bbox2) in enumerate(zip(track_groups[1], bboxes[1])):
+            if bbox1.overlaps_with(bbox2):
+                if tracks_close(track1, track2, position_window, time_window):
+                    colocalizations.add((idx1, idx2))
+
+    sorted_list = list(colocalizations)
+    sorted_list.sort()
+
+    return sorted_list

lumicks/pylake/kymotracker/tests/test_colocalization.py

@@ -1,8 +1,12 @@
 import pytest
 
-from lumicks.pylake.kymotracker.kymotrack import KymoTrack
+from lumicks.pylake.kymotracker.kymotrack import KymoTrack, KymoTrackGroup
 from lumicks.pylake.kymotracker.colocalization import classify_track
-from lumicks.pylake.kymotracker.detail.track_proximity import BoundingBox
+from lumicks.pylake.kymotracker.detail.track_proximity import (
+    BoundingBox,
+    tracks_close,
+    find_colocalized,
+)
 
 
 @pytest.mark.parametrize(
@@ -86,3 +90,50 @@ def test_track_colocalization_classifier(blank_kymo, red_idx, blue_idx, classifi
     assert (
         result := classify_track(red, blue, 2)
     ) == classification, f"expected {classification}, got {result}"
+
+
+@pytest.mark.parametrize(
+    "track1, track2, time_window, position_window, result",
+    [
+        (([1, 2, 3], [1.0, 1.0, 1.0]), ([1, 2, 3], [3.0, 3.0, 3.0]), 0, 2, True),
+        (([1, 2, 3], [1.0, 1.0, 1.0]), ([1, 2, 3], [3.0, 3.0, 3.0]), 0, 1.99, False),
+        (([1, 3, 5], [1.0, 1.0, 1.0]), ([2, 4, 6], [1.0, 1.0, 1.0]), 1, 1, True),
+        (([1, 3, 5], [1.0, 1.0, 1.0]), ([2, 4, 6], [1.0, 1.0, 1.0]), 0, 1, False),
+        (([1, 3], [1.0, 1.0]), ([2, 4, 6], [1.0, 1.0, 1.0]), 1, 1, True),
+        (([1, 3], [1.0, 1.0]), ([2, 4, 6], [1.0, 1.0, 1.0]), 0, 1, False),
+        (([], []), ([], []), 0, 1, False),
+    ],
+)
+def test_colocalization_interacting(
+    blank_kymo, track1, track2, time_window, position_window, result
+):
+    track1 = KymoTrack(*track1, blank_kymo, "red", 0)
+    track2 = KymoTrack(*track2, blank_kymo, "red", 0)
+    assert (
+        tracks_close(track1, track2, time_window=time_window, position_window=position_window)
+        == result
+    )
+
+
+def test_colocalization_find_colocalization(
+    blank_kymo,  # time_window, position_window, result
+):
+    track_coords1 = [
+        ([1, 2, 3], [1, 1, 1]),
+        ([10, 14], [1, 1]),
+    ]
+
+    track_coords2 = [
+        ([1, 2, 3], [3, 3, 3]),
+        ([11, 12, 13], [2, 2, 2]),  # Will only show up when colocalized
+    ]
+
+    tracks1 = KymoTrackGroup([KymoTrack(*tc, blank_kymo, "red", 0) for tc in track_coords1])
+    tracks2 = KymoTrackGroup([KymoTrack(*tc, blank_kymo, "red", 0) for tc in track_coords2])
+
+    # assert (
+    find_colocalized(
+        tracks1, tracks2, time_window=1, position_window=1
+    )  # , time_window=time_window, position_window=position_window)
+    #    == result
+    # )