Kdtrees

mypy.ini

@@ -16,3 +16,6 @@ ignore_missing_imports = True
 
 [mypy-motile.*]
 ignore_missing_imports = True
+
+[mypy-scipy.*]
+ignore_missing_imports = True

src/motile_toolbox/candidate_graph/graph_attributes.py

@@ -19,5 +19,4 @@ class EdgeAttr(Enum):
     implementations of commonly used ones, listed here.
     """
 
-    DISTANCE = "distance"
     IOU = "iou"

src/motile_toolbox/candidate_graph/utils.py

@@ -1,10 +1,11 @@
 import logging
 import math
-from typing import Any
+from typing import Any, Iterable
 
 import networkx as nx
 import numpy as np
 from skimage.measure import regionprops
+from scipy.spatial import KDTree
 from tqdm import tqdm
 
 from .graph_attributes import EdgeAttr, NodeAttr
@@ -99,6 +100,11 @@ def _compute_node_frame_dict(cand_graph: nx.DiGraph) -> dict[int, list[Any]]:
     return node_frame_dict
 
 
+def create_kdtree(cand_graph: nx.DiGraph, node_ids: Iterable[Any]) -> KDTree:
+    positions = [cand_graph.nodes[node][NodeAttr.POS.value] for node in node_ids]
+    return KDTree(positions)
+
+
 def add_cand_edges(
     cand_graph: nx.DiGraph,
     max_edge_distance: float,
@@ -122,15 +128,20 @@ def add_cand_edges(
         node_frame_dict = _compute_node_frame_dict(cand_graph)
 
     frames = sorted(node_frame_dict.keys())
+    prev_node_ids = node_frame_dict[frames[0]]
+    prev_kdtree = create_kdtree(cand_graph, prev_node_ids)
     for frame in tqdm(frames):
         if frame + 1 not in node_frame_dict:
             continue
-        next_nodes = node_frame_dict[frame + 1]
-        next_locs = [cand_graph.nodes[n][NodeAttr.POS.value] for n in next_nodes]
-        for node in node_frame_dict[frame]:
-            loc = cand_graph.nodes[node][NodeAttr.POS.value]
-            for next_id, next_loc in zip(next_nodes, next_locs):
-                dist = math.dist(next_loc, loc)
-                if dist <= max_edge_distance:
-                    attrs = {EdgeAttr.DISTANCE.value: dist}
-                    cand_graph.add_edge(node, next_id, **attrs)
+        next_node_ids = node_frame_dict[frame + 1]
+        next_kdtree = create_kdtree(cand_graph, next_node_ids)
+
+        matched_indices = prev_kdtree.query_ball_tree(next_kdtree, max_edge_distance)
+
+        for prev_node_id, next_node_indices in zip(prev_node_ids, matched_indices):
+            for next_node_index in next_node_indices:
+                next_node_id = next_node_ids[next_node_index]
+                cand_graph.add_edge(prev_node_id, next_node_id)
+
+        prev_node_ids = next_node_ids
+        prev_kdtree = next_kdtree

tests/conftest.py

@@ -90,8 +90,8 @@ def graph_2d():
         ),
     ]
     edges = [
-        ("0_1", "1_1", {EdgeAttr.DISTANCE.value: 42.43, EdgeAttr.IOU.value: 0.0}),
-        ("0_1", "1_2", {EdgeAttr.DISTANCE.value: 11.18, EdgeAttr.IOU.value: 0.395}),
+        ("0_1", "1_1", {EdgeAttr.IOU.value: 0.0}),
+        ("0_1", "1_2", {EdgeAttr.IOU.value: 0.395}),
     ]
     graph.add_nodes_from(nodes)
     graph.add_edges_from(edges)
@@ -159,25 +159,25 @@ def multi_hypothesis_graph_2d():
     ]
 
     edges = [
-        ("0_0_1", "1_0_1", {EdgeAttr.DISTANCE.value: 42.426, EdgeAttr.IOU.value: 0.0}),
-        ("0_0_1", "1_1_1", {EdgeAttr.DISTANCE.value: 43.011, EdgeAttr.IOU.value: 0.0}),
+        ("0_0_1", "1_0_1", {EdgeAttr.IOU.value: 0.0}),
+        ("0_0_1", "1_1_1", {EdgeAttr.IOU.value: 0.0}),
         (
             "0_0_1",
             "1_0_2",
-            {EdgeAttr.DISTANCE.value: 11.180, EdgeAttr.IOU.value: 0.3931},
+            {EdgeAttr.IOU.value: 0.3931},
         ),
         (
             "0_0_1",
             "1_1_2",
-            {EdgeAttr.DISTANCE.value: 11.180, EdgeAttr.IOU.value: 0.4768},
+            {EdgeAttr.IOU.value: 0.4768},
         ),
-        ("0_1_1", "1_0_1", {EdgeAttr.DISTANCE.value: 43.011, EdgeAttr.IOU.value: 0.0}),
-        ("0_1_1", "1_1_1", {EdgeAttr.DISTANCE.value: 42.426, EdgeAttr.IOU.value: 0.0}),
-        ("0_1_1", "1_0_2", {EdgeAttr.DISTANCE.value: 15.0, EdgeAttr.IOU.value: 0.2402}),
+        ("0_1_1", "1_0_1", {EdgeAttr.IOU.value: 0.0}),
+        ("0_1_1", "1_1_1", {EdgeAttr.IOU.value: 0.0}),
+        ("0_1_1", "1_0_2", {EdgeAttr.IOU.value: 0.2402}),
         (
             "0_1_1",
             "1_1_2",
-            {EdgeAttr.DISTANCE.value: 11.180, EdgeAttr.IOU.value: 0.3931},
+            {EdgeAttr.IOU.value: 0.3931},
         ),
     ]
     graph.add_nodes_from(nodes)
@@ -276,9 +276,9 @@ def graph_3d():
     ]
     edges = [
         # math.dist([50, 50], [20, 80])
-        ("0_1", "1_1", {EdgeAttr.DISTANCE.value: 42.43}),
+        ("0_1", "1_1"),
         # math.dist([50, 50], [60, 45])
-        ("0_1", "1_2", {EdgeAttr.DISTANCE.value: 11.18}),
+        ("0_1", "1_2"),
     ]
     graph.add_nodes_from(nodes)
     graph.add_edges_from(edges)
@@ -336,12 +336,12 @@ def multi_hypothesis_graph_3d():
         ),
     ]
     edges = [
-        ("0_0_1", "1_0_1", {EdgeAttr.DISTANCE.value: 42.4264}),
-        ("0_0_1", "1_0_2", {EdgeAttr.DISTANCE.value: 11.1803}),
-        ("0_1_1", "1_0_1", {EdgeAttr.DISTANCE.value: 35.3553}),
-        ("0_1_1", "1_0_2", {EdgeAttr.DISTANCE.value: 18.0277}),
-        ("0_0_1", "1_1_1", {EdgeAttr.DISTANCE.value: 40.3112}),
-        ("0_1_1", "1_1_1", {EdgeAttr.DISTANCE.value: 33.5410}),
+        ("0_0_1", "1_0_1"),
+        ("0_0_1", "1_0_2"),
+        ("0_1_1", "1_0_1"),
+        ("0_1_1", "1_0_2"),
+        ("0_0_1", "1_1_1"),
+        ("0_1_1", "1_1_1"),
     ]
     graph.add_nodes_from(nodes)
     graph.add_edges_from(edges)

tests/test_candidate_graph/test_compute_graph.py

@@ -17,10 +17,6 @@ def test_graph_from_segmentation_2d(segmentation_2d, graph_2d):
         assert Counter(cand_graph.nodes[node]) == Counter(graph_2d.nodes[node])
     for edge in cand_graph.edges:
         print(cand_graph.edges[edge])
-        assert (
-            pytest.approx(cand_graph.edges[edge][EdgeAttr.DISTANCE.value], abs=0.01)
-            == graph_2d.edges[edge][EdgeAttr.DISTANCE.value]
-        )
         assert (
             pytest.approx(cand_graph.edges[edge][EdgeAttr.IOU.value], abs=0.01)
             == graph_2d.edges[edge][EdgeAttr.IOU.value]
@@ -33,9 +29,6 @@ def test_graph_from_segmentation_2d(segmentation_2d, graph_2d):
     )
     assert Counter(list(cand_graph.nodes)) == Counter(["0_1", "1_1", "1_2"])
     assert Counter(list(cand_graph.edges)) == Counter([("0_1", "1_2")])
-    assert cand_graph.edges[("0_1", "1_2")][EdgeAttr.DISTANCE.value] == pytest.approx(
-        11.18, abs=0.01
-    )
 
 
 def test_graph_from_segmentation_3d(segmentation_3d, graph_3d):
@@ -72,10 +65,6 @@ def test_graph_from_multi_segmentation_2d(
             multi_hypothesis_graph_2d.nodes[node]
         )
     for edge in cand_graph.edges:
-        assert (
-            pytest.approx(cand_graph.edges[edge][EdgeAttr.DISTANCE.value], abs=0.01)
-            == multi_hypothesis_graph_2d.edges[edge][EdgeAttr.DISTANCE.value]
-        )
         assert (
             pytest.approx(cand_graph.edges[edge][EdgeAttr.IOU.value], abs=0.01)
             == multi_hypothesis_graph_2d.edges[edge][EdgeAttr.IOU.value]

tests/test_candidate_graph/test_utils.py

@@ -75,20 +75,13 @@ def test_add_cand_edges_2d(graph_2d):
     cand_graph = nx.create_empty_copy(graph_2d)
     add_cand_edges(cand_graph, max_edge_distance=50)
     assert Counter(list(cand_graph.edges)) == Counter(list(graph_2d.edges))
-    for edge in cand_graph.edges:
-        assert (
-            pytest.approx(cand_graph.edges[edge][EdgeAttr.DISTANCE.value], abs=0.01)
-            == graph_2d.edges[edge][EdgeAttr.DISTANCE.value]
-        )
 
 
 def test_add_cand_edges_3d(graph_3d):
     cand_graph = nx.create_empty_copy(graph_3d)
     add_cand_edges(cand_graph, max_edge_distance=15)
     graph_3d.remove_edge("0_1", "1_1")
     assert Counter(list(cand_graph.edges)) == Counter(list(graph_3d.edges))
-    for edge in cand_graph.edges:
-        assert pytest.approx(cand_graph.edges[edge], abs=0.01) == graph_3d.edges[edge]
 
 
 def test_get_node_id():