Fix label fragmentation shape lookup

Author: francois-drielsma

src/spine/model/full_chain.py

@@ -712,7 +712,9 @@ def run_fragmentation(self, data, clust_label=None):
                 clust_label is not None
             ), "Must provide `clust_label` to use it for fragmentation."
             fragments = form_clusters_batch(clust_label.to_numpy(), column=CLUST_COL)
-            fragment_shapes = get_cluster_label_batch(clust_label, fragments)
+            fragment_shapes = get_cluster_label_batch(
+                clust_label, fragments, column=SHAPE_COL
+            )
 
         if fragments is not None:
             self.result["fragment_clusts"] = fragments

test/test_model/test_full_chain.py

@@ -4,7 +4,15 @@
 import torch
 
 import spine.model.full_chain as full_chain_mod
-from spine.constants import GHOST_SHP, GROUP_COL, PRGRP_COL, SHAPE_COL, SHOWR_SHP
+from spine.constants import (
+    CLUST_COL,
+    GHOST_SHP,
+    GROUP_COL,
+    PRGRP_COL,
+    SHAPE_COL,
+    SHOWR_SHP,
+    TRACK_SHP,
+)
 from spine.data import IndexBatch, TensorBatch
 from spine.model.full_chain import FullChain, FullChainLoss
 
@@ -112,6 +120,29 @@ def test_group_labels_accepts_shape_restriction_without_model():
     assert group_primaries is groups
 
 
+def test_label_fragmentation_reads_shapes_from_shape_column():
+    full_chain = object.__new__(FullChain)
+    full_chain.fragmentation = "label"
+    full_chain.result = {}
+
+    data = TensorBatch(np.zeros((4, 1), dtype=np.float32), counts=np.array([4]))
+    clust_label_array = np.zeros((4, CLUST_COL + 2), dtype=np.float32)
+    clust_label_array[:, CLUST_COL] = np.array([10, 10, 20, 20], dtype=np.float32)
+    clust_label_array[:, SHAPE_COL] = np.array(
+        [SHOWR_SHP, SHOWR_SHP, TRACK_SHP, TRACK_SHP], dtype=np.float32
+    )
+    clust_label = TensorBatch(clust_label_array, counts=np.array([4]))
+
+    full_chain.run_fragmentation(data, clust_label)
+
+    fragments = full_chain.result["fragment_clusts"]
+    fragment_shapes = full_chain.result["fragment_shapes"]
+
+    assert fragments.counts.tolist() == [2]
+    assert [f.tolist() for f in fragments.index_list] == [[0, 1], [2, 3]]
+    assert fragment_shapes.tensor.tolist() == [SHOWR_SHP, TRACK_SHP]
+
+
 def test_prepare_grappa_input_uses_label_points_without_ppn(monkeypatch):
     full_chain = object.__new__(FullChain)
     full_chain.result = {}