Use cluster labels for truth point lookup

Author: francois-drielsma

src/spine/model/full_chain.py

@@ -783,7 +783,8 @@ def run_part_aggregation(self, data, clust_label=None, coord_label=None):
                     data,
                     fragments,
                     fragment_shapes,
-                    coord_label,
+                    clust_label=clust_label,
+                    coord_label=coord_label,
                     aggregate_shapes=True,
                     shape_use_primary=use_primary[name],
                     retain_primaries=use_primary[name],
@@ -895,7 +896,8 @@ def run_inter_aggregation(self, data, clust_label=None, coord_label=None):
                 particles,
                 particle_shapes,
                 particle_primaries,
-                coord_label,
+                clust_label=clust_label,
+                coord_label=coord_label,
                 point_use_primary=True,
             )
 
@@ -1002,6 +1004,7 @@ def run_grappa(
         clusts,
         clust_shapes,
         clust_primaries=None,
+        clust_label=None,
         coord_label=None,
         aggregate_shapes=False,
         shape_use_primary=False,
@@ -1024,6 +1027,9 @@ def run_grappa(
             Semantic type of each of the clusters
         clust_primaries : IndexBatch
             List of primary fragments associated with each input cluster
+        clust_label : TensorBatch, optional
+            Tensor used to fetch truth labels when building points from
+            `coord_label`. Defaults to `data`.
         coord_label : TensorBatch, optional
             (N, 1 + D + 6) Array of label particle end points
         aggregate_shapes : bool, default False
@@ -1058,6 +1064,7 @@ def run_grappa(
             clusts,
             clust_shapes,
             clust_primaries,
+            clust_label,
             coord_label,
             point_use_primary,
         )
@@ -1209,6 +1216,7 @@ def prepare_grappa_input(
         clusts,
         clust_shapes,
         clust_primaries=None,
+        clust_label=None,
         coord_label=None,
         point_use_primaries=False,
     ):
@@ -1231,6 +1239,9 @@ def prepare_grappa_input(
             Semantic type of each of the clusters
         clust_primaries : IndexBatch, optional
             List of primary fragment within each cluster to aggregate
+        clust_label : TensorBatch, optional
+            Tensor used to fetch truth labels when building points from
+            `coord_label`. Defaults to `data`.
         coord_label : TensorBatch, optional
             (N, 1 + D + 6) Array of label particle end points
         point_use_primaries:
@@ -1274,9 +1285,13 @@ def prepare_grappa_input(
                     "Must provide either `ppn_points` or `coord_label` to add "
                     "points to the GrapPA input."
                 )
+                assert clust_label is not None, (
+                    "Must provide `clust_label` with `coord_label` to add "
+                    "label points to the GrapPA input."
+                )
                 point_clusts = clusts if point_use_primaries else ref_clusts
                 points = get_cluster_points_label_batch(
-                    data,
+                    clust_label,
                     coord_label,
                     point_clusts,
                     use_group=point_use_primaries,

src/spine/utils/gnn/cluster.py

@@ -1033,12 +1033,23 @@ def _get_cluster_points_label(
         if use_group:
             # Use the true aggregate particle identity directly.
             label_ids = np.unique(data[c, GROUP_COL]).astype(np.int64)
-            label = coord_label[label_ids[0]]
+            label_id = label_ids[0]
         else:
             # Use the first constituent particle in time.
             part_ids = np.unique(data[c, PART_COL]).astype(np.int64)
-            min_id = part_ids[np.argmin(coord_label[part_ids, COORD_TIME_COL])]
-            label = coord_label[min_id]
+            label_id = -1
+            min_time = np.inf
+            for part_id in part_ids:
+                if part_id < 0 or part_id >= len(coord_label):
+                    raise IndexError("Invalid label index for coord_label.")
+                time = coord_label[part_id, COORD_TIME_COL]
+                if time < min_time:
+                    min_time = time
+                    label_id = part_id
+
+        if label_id < 0 or label_id >= len(coord_label):
+            raise IndexError("Invalid label index for coord_label.")
+        label = coord_label[label_id]
 
         start = label[COORD_START_COLS_LO:COORD_START_COLS_HI]
         end = label[COORD_END_COLS_LO:COORD_END_COLS_HI]

test/test_model/test_full_chain.py

@@ -1,6 +1,7 @@
 from types import SimpleNamespace
 
 import numpy as np
+import pytest
 import torch
 
 import spine.model.full_chain as full_chain_mod
@@ -183,6 +184,9 @@ def test_prepare_grappa_input_uses_label_points_without_ppn(monkeypatch):
 
     model = SimpleNamespace(node_encoder=SimpleNamespace(add_points=True))
     data = TensorBatch(np.zeros((3, 4), dtype=np.float32), counts=np.array([3]))
+    clust_label = TensorBatch(
+        np.zeros((3, GROUP_COL + 1), dtype=np.float32), counts=np.array([3])
+    )
     clusts = IndexBatch(
         [np.array([0], dtype=np.int64), np.array([1, 2], dtype=np.int64)],
         spans=np.array([3]),
@@ -214,12 +218,38 @@ def fake_label_points(data_arg, coord_label_arg, clusts_arg, **kwargs):
         clusts,
         clust_shapes,
         clust_primaries=primaries,
+        clust_label=clust_label,
         coord_label=coord_label,
         point_use_primaries=True,
     )
 
     assert grappa_input["points"] is label_points
     assert "coord_label" not in grappa_input
     assert calls == [
-        (data, coord_label, clusts, {"use_group": True}),
+        (clust_label, coord_label, clusts, {"use_group": True}),
     ]
+
+
+def test_prepare_grappa_input_requires_clust_label_for_label_points():
+    full_chain = object.__new__(FullChain)
+    full_chain.result = {}
+
+    model = SimpleNamespace(node_encoder=SimpleNamespace(add_points=True))
+    data = TensorBatch(np.zeros((1, 4), dtype=np.float32), counts=np.array([1]))
+    clusts = IndexBatch(
+        [np.array([0], dtype=np.int64)],
+        spans=np.array([1]),
+        counts=np.array([1]),
+        single_counts=np.array([1]),
+    )
+    clust_shapes = TensorBatch(np.array([SHOWR_SHP]), counts=np.array([1]))
+    coord_label = TensorBatch(np.zeros((1, 9), dtype=np.float32), counts=np.array([1]))
+
+    with pytest.raises(AssertionError, match="clust_label"):
+        full_chain.prepare_grappa_input(
+            model,
+            data,
+            clusts,
+            clust_shapes,
+            coord_label=coord_label,
+        )

test/test_utils/test_gnn_cluster.py

@@ -1,6 +1,7 @@
 """Regression tests for GNN cluster utilities."""
 
 import numpy as np
+import pytest
 
 from spine.constants import GROUP_COL, PART_COL
 from spine.utils.gnn.cluster import (
@@ -78,3 +79,32 @@ def test_cluster_points_label_can_use_group_identity():
 
     np.testing.assert_allclose(points[0, :3], [20.0, 0.0, 0.0])
     np.testing.assert_allclose(points[0, 3:], [20.0, 0.0, 0.0])
+
+
+def test_cluster_points_label_rejects_invalid_particle_id():
+    data = np.zeros((1, PART_COL + 1), dtype=np.float32)
+    data[:, PART_COL] = 1
+    coord_label = np.zeros((1, 9), dtype=np.float32)
+
+    with pytest.raises(IndexError, match="Invalid label index"):
+        get_cluster_points_label(
+            data,
+            coord_label,
+            [np.array([0], dtype=np.int64)],
+            random_order=False,
+        )
+
+
+def test_cluster_points_label_rejects_invalid_group_id():
+    data = np.zeros((1, GROUP_COL + 1), dtype=np.float32)
+    data[:, GROUP_COL] = 1
+    coord_label = np.zeros((1, 9), dtype=np.float32)
+
+    with pytest.raises(IndexError, match="Invalid label index"):
+        get_cluster_points_label(
+            data,
+            coord_label,
+            [np.array([0], dtype=np.int64)],
+            random_order=False,
+            use_group=True,
+        )