Restore part-based truth point selection

Author: francois-drielsma

src/spine/model/full_chain.py

@@ -1265,6 +1265,7 @@ def prepare_grappa_input(
         grappa_input["data"] = data
         grappa_input["clusts"] = clusts
         grappa_input["shapes"] = clust_shapes
+
         # Get the particle end points, if requested
         if hasattr(model.node_encoder, "add_points") and model.node_encoder.add_points:
             # Fetch the cluster list to use to get points
@@ -1289,12 +1290,10 @@ def prepare_grappa_input(
                     "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(
                     clust_label,
                     coord_label,
-                    point_clusts,
-                    use_group=point_use_primaries,
+                    ref_clusts,
                 )
 
             grappa_input["points"] = points

src/spine/utils/gnn/cluster.py

@@ -171,9 +171,7 @@ def get_cluster_closest_primary_label_batch(data, coord_label, clusts, primary_i
     return TensorBatch(labels, clusts.counts)
 
 
-def get_cluster_points_label_batch(
-    data, coord_label, clusts, random_order=True, use_group=False
-):
+def get_cluster_points_label_batch(data, coord_label, clusts, random_order=True):
     """Batched version of :func:`get_cluster_points_label`
 
     Parameters
@@ -187,11 +185,6 @@ def get_cluster_points_label_batch(
     random_order : bool, default True
         If `True`, randomize the order in which the start en end points of
         a track are stored in the output
-    use_group : bool, default False
-        If `True`, use the true group ID to fetch the particle end point
-        labels. Otherwise, use particle IDs and select the earliest particle
-        in time.
-
     Returns
     -------
     np.ndarray
@@ -210,7 +203,6 @@ def get_cluster_points_label_batch(
             coord_label[b],
             clusts[b],
             random_order,
-            use_group,
         )
 
     return TensorBatch(points, clusts.counts, coord_cols=points.shape[1])
@@ -982,9 +974,7 @@ def _get_cluster_features_extended(
     keep_torch=True,
     ref_arg="data",
 )
-def get_cluster_points_label(
-    data, coord_label, clusts, random_order=True, use_group=False
-):
+def get_cluster_points_label(data, coord_label, clusts, random_order=True):
     """Gets label points for each cluster.
 
     Returns start point of primary shower fragment twice if shower, delta or
@@ -1002,11 +992,6 @@ def get_cluster_points_label(
     random_order : bool, default True
         If `True`, randomize the order in which the start en end points of
         a track are stored in the output
-    use_group : bool, default False
-        If `True`, use the true group ID to fetch the particle end point
-        labels. Otherwise, use particle IDs and select the earliest particle
-        in time.
-
     Returns
     -------
     np.ndarray
@@ -1015,7 +1000,7 @@ def get_cluster_points_label(
     if len(clusts) == 0:
         return np.empty((0, 6), dtype=data.dtype)
 
-    return _get_cluster_points_label(data, coord_label, clusts, random_order, use_group)
+    return _get_cluster_points_label(data, coord_label, clusts, random_order)
 
 
 @nb.njit(cache=True)
@@ -1024,28 +1009,22 @@ def _get_cluster_points_label(
     coord_label: nb.float64[:, :],
     clusts: nb.types.List(nb.int64[:]),
     random_order: nb.boolean = True,
-    use_group: nb.boolean = False,
 ) -> nb.float64[:, :]:
 
     # Get start and end points (one and the same for all but track class)
     points = np.empty((len(clusts), 6), dtype=data.dtype)
     for i, c in enumerate(clusts):
-        if use_group:
-            # Use the true aggregate particle identity directly.
-            label_ids = np.unique(data[c, GROUP_COL]).astype(np.int64)
-            label_id = label_ids[0]
-        else:
-            # Use the first constituent particle in time.
-            part_ids = np.unique(data[c, PART_COL]).astype(np.int64)
-            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
+        # Use the first constituent particle in time.
+        part_ids = np.unique(data[c, PART_COL]).astype(np.int64)
+        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.")

test/test_model/test_full_chain.py

@@ -226,7 +226,7 @@ def fake_label_points(data_arg, coord_label_arg, clusts_arg, **kwargs):
     assert grappa_input["points"] is label_points
     assert "coord_label" not in grappa_input
     assert calls == [
-        (clust_label, coord_label, clusts, {"use_group": True}),
+        (clust_label, coord_label, primaries, {}),
     ]
 
 

test/test_utils/test_gnn_cluster.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from spine.constants import GROUP_COL, PART_COL
+from spine.constants import PART_COL
 from spine.utils.gnn.cluster import (
     cluster_dedx,
     get_cluster_features_base,
@@ -50,37 +50,6 @@ def test_cluster_features_base_accepts_indexed_float32_coordinates():
     np.testing.assert_allclose(feats[1, :3], [1.0 / 3.0, 0.0, 5.0])
 
 
-def test_cluster_points_label_can_use_group_identity():
-    data = np.zeros((2, GROUP_COL + 1), dtype=np.float32)
-    data[:, 1:4] = np.array(
-        [
-            [10.0, 0.0, 0.0],
-            [20.0, 0.0, 0.0],
-        ],
-        dtype=np.float32,
-    )
-    data[:, PART_COL] = [0, 1]
-    data[:, GROUP_COL] = 1
-    coord_label = np.array(
-        [
-            [0.0, 10.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 0.0],
-            [0.0, 20.0, 0.0, 0.0, 20.0, 0.0, 0.0, 5.0, 0.0],
-        ],
-        dtype=np.float32,
-    )
-
-    points = get_cluster_points_label(
-        data,
-        coord_label,
-        [np.array([0, 1], dtype=np.int64)],
-        random_order=False,
-        use_group=True,
-    )
-
-    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
@@ -93,18 +62,3 @@ def test_cluster_points_label_rejects_invalid_particle_id():
             [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,
-        )