From e1637c91fe90038715fa5cac0ffbf2f6778fc6cb Mon Sep 17 00:00:00 2001
From: Joep Vanlier <joep.vanlier@gmail.com>
Date: Mon, 7 Oct 2024 15:52:14 +0200
Subject: [PATCH] colocalization: add method to classify binding order

---
 lumicks/pylake/kymotracker/colocalization.py  | 29 +++++++++++++++
 .../kymotracker/tests/test_colocalization.py  | 36 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 lumicks/pylake/kymotracker/colocalization.py

diff --git a/lumicks/pylake/kymotracker/colocalization.py b/lumicks/pylake/kymotracker/colocalization.py
new file mode 100644
index 000000000..4d3b49aa9
--- /dev/null
+++ b/lumicks/pylake/kymotracker/colocalization.py
@@ -0,0 +1,29 @@
+def classify_track(red, blue, timing_window=6):
+    """This function classifies the tracks into groups based on the order in which
+    the proteins bind"""
+
+    #  3 blue_start < red_start  blue_end < red_end
+    #  4 blue_start < red_start  blue_end = red_end
+    #  5 blue_start < red_start  blue_end > red_end
+
+    #  6 blue_start = red_start  blue_end < red_end
+    #  7 blue_start = red_start  blue_end = red_end
+    #  8 blue_start = red_start  blue_end > red_end
+
+    #  9 red_start < blue_start  blue_end < red_end
+    # 10 red_start < blue_start  blue_end = red_end
+    # 11 red_start < blue_start  blue_end > red_end
+
+    def classify(time, ref_time):
+        if ref_time - time > timing_window:
+            return 0
+        elif time - ref_time > timing_window:
+            return 2
+        else:
+            return 1
+
+    return (
+        3
+        + 3 * classify(blue.time_idx[0], red.time_idx[0])
+        + classify(blue.time_idx[-1], red.time_idx[-1])
+    )
diff --git a/lumicks/pylake/kymotracker/tests/test_colocalization.py b/lumicks/pylake/kymotracker/tests/test_colocalization.py
index 92845534a..cedc23ad7 100644
--- a/lumicks/pylake/kymotracker/tests/test_colocalization.py
+++ b/lumicks/pylake/kymotracker/tests/test_colocalization.py
@@ -1,6 +1,7 @@
 import pytest
 
 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,
     tracks_close,
@@ -56,6 +57,41 @@ def test_bbox_construction_test(
     assert bbox._valid == is_valid
 
 
+@pytest.mark.parametrize(
+    "red_idx, blue_idx, classification",
+    [
+        # Blue starts first
+        ([5, 13], [2, 10], 3),
+        ([5, 12], [2, 10], 4),
+        ([5, 10], [2, 10], 4),
+        ([5, 8], [2, 10], 4),
+        ([5, 7], [2, 10], 5),
+        # Starts at the same time
+        ([0, 13], [2, 10], 6),
+        ([2, 13], [2, 10], 6),
+        ([4, 13], [2, 10], 6),
+        ([5, 13], [2, 10], 3),  # Out of the time window for "same time"
+        ([2, 8], [2, 10], 7),
+        ([2, 10], [2, 10], 7),
+        ([2, 12], [2, 12], 7),
+        ([2, 7], [2, 10], 8),
+        # Red starts first
+        ([2, 13], [5, 10], 9),
+        ([2, 12], [5, 10], 10),
+        ([2, 10], [5, 10], 10),
+        ([2, 8], [5, 10], 10),
+        ([2, 7], [5, 10], 11),
+    ],
+)
+def test_track_colocalization_classifier(blank_kymo, red_idx, blue_idx, classification):
+    red = KymoTrack(red_idx, [1, 1], blank_kymo, 0, 0)
+    blue = KymoTrack(blue_idx, [1, 1], blank_kymo, 0, 0)
+
+    assert (
+        result := classify_track(red, blue, 2)
+    ) == classification, f"expected {classification}, got {result}"
+
+
 @pytest.mark.parametrize(
     "track1, track2, time_window, position_window, result",
     [