Merge pull request #177 from yfukai/score_fix

Score fix

Author: yfukai

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "laptrack"
-version = "0.3.1"
+version = "0.3.2"
 description = "LapTrack"
 authors = ["Yohsuke Fukai <ysk@yfukai.net>"]
 license = "BSD-3-Clause"

src/laptrack/scores.py

@@ -109,23 +109,30 @@ def calc_scores(
         target_effectiveness = _calc_overlap_score(gt_edgess, pred_edgess)
 
         def get_children(m):
-            return [n for n in gt_tree.neighbors(m) if n[0] > m[0]]
+            return list(gt_tree.successors(m))
 
         dividing_nodes = [m for m in gt_tree.nodes() if len(get_children(m)) > 1]
         division_recovery_count = 0
+        total_count = 0
         for m in dividing_nodes:
             children = get_children(m)
-            if all(
-                [
-                    (n, m) in predicted_edges or (m, n) in predicted_edges
-                    for n in children
-                ]
-            ):
+
+            def check_in(edges):
+                return all([(n, m) in edges or (m, n) in edges for n in children])
+
+            excluded = check_in(exclude_true_edges)
+            if check_in(predicted_edges) and not excluded:
                 division_recovery_count += 1
-        division_recovery = division_recovery_count / len(dividing_nodes)
+            if not excluded:
+                total_count += 1
+        if total_count > 0:
+            division_recovery = division_recovery_count / total_count
+        else:
+            division_recovery = -1
 
         te = set(true_edges) - set(exclude_true_edges)
         pe = set(predicted_edges) - set(exclude_true_edges)
+        print(len(true_edges), len(exclude_true_edges), len(te))
         return {
             "union_ratio": len(te & pe) / len(te | pe),
             "true_ratio": len(te & pe) / len(te),

tests/test_scores.py

@@ -1,9 +1,11 @@
 import networkx as nx
+import pytest
 
 from laptrack.scores import calc_scores
 
 
-def test_scores() -> None:
+@pytest.fixture
+def test_trees():
     true_tree = nx.from_edgelist(
         [
             ((0, 0), (1, 0)),
@@ -20,8 +22,7 @@ def test_scores() -> None:
         ]
     )
 
-    pred_tree = nx.Graph()
-    pred_tree.add_edges_from(
+    pred_tree = nx.from_edgelist(
         [
             ((0, 0), (1, 0)),
             ((1, 0), (2, 0)),
@@ -34,7 +35,11 @@ def test_scores() -> None:
             ((4, 1), (5, 1)),
         ]
     )
+    return true_tree, pred_tree
 
+
+def test_scores(test_trees) -> None:
+    true_tree, pred_tree = test_trees
     score = {
         "union_ratio": 8 / 12,
         "true_ratio": 8 / 11,
@@ -43,5 +48,34 @@ def test_scores() -> None:
         "target_effectiveness": 6 / 11,
         "division_recovery": 1,
     }
-
     assert score == calc_scores(true_tree.edges, pred_tree.edges)
+
+
+def test_scores_exclude(test_trees) -> None:
+    true_tree, pred_tree = test_trees
+    exclude_edges = [((2, 0), (3, 0)), ((2, 0), (3, 1))]
+    score = {
+        "union_ratio": 6 / 10,
+        "true_ratio": 6 / 9,
+        "predicted_ratio": 6 / 7,
+        "track_purity": 7 / 9,
+        "target_effectiveness": 6 / 11,
+        "division_recovery": -1,
+    }
+    assert score == calc_scores(true_tree.edges, pred_tree.edges, exclude_edges)
+
+
+def test_scores_exclude2(test_trees) -> None:
+    true_tree, pred_tree = test_trees
+    exclude_edges = [
+        ((2, 0), (3, 0)),
+    ]
+    score = {
+        "union_ratio": 7 / 11,
+        "true_ratio": 7 / 10,
+        "predicted_ratio": 7 / 8,
+        "track_purity": 7 / 9,
+        "target_effectiveness": 6 / 11,
+        "division_recovery": 1,
+    }
+    assert score == calc_scores(true_tree.edges, pred_tree.edges, exclude_edges)