Fix Numba layout handling in GrapPA features

Author: francois-drielsma

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## [0.14.1] - 2026-06-27
+
+### Fixed
+- **GrapPA inference stability**: Normalize indexed cluster coordinate views before Numba-compiled distance, endpoint, node-feature, and edge-feature helpers so full-chain GrapPA inference no longer fails intermittently on arbitrary-layout arrays in batch jobs.
+- **Regression coverage**: Add focused tests for arbitrary-layout Numba callers in distance helpers and GrapPA cluster/node/edge feature extraction.
+
 ## [0.14.0] - 2026-06-26
 
 ### Added

src/spine/math/distance.py

@@ -194,6 +194,7 @@ def pdist(
     """
     # Check on the input
     assert x.shape[1] == 3, "Only supports 3D points for now."
+    x = np.ascontiguousarray(x)
 
     # Dispatch (faster this way than dispatching at each distance call)
     if metric_id == MINKOWSKI:
@@ -293,6 +294,8 @@ def cdist(
     """
     # Check on the input
     assert x1.shape[1] == 3 and x2.shape[1] == 3, "Only supports 3D points for now."
+    x1 = np.ascontiguousarray(x1)
+    x2 = np.ascontiguousarray(x2)
 
     # Dispatch (faster this way than dispatching at each distance call)
     if metric_id == MINKOWSKI:
@@ -395,6 +398,8 @@ def farthest_pair(
     float
         Distance between the two points
     """
+    x = np.ascontiguousarray(x)
+
     # To save time, if Euclidean distance is used, use its square
     is_euclidean = False
     if metric_id == EUCLIDEAN:
@@ -461,6 +466,9 @@ def closest_pair_legacy(
     set switch after each closest-point update. New code should use
     :func:`closest_pair`.
     """
+    x1 = np.ascontiguousarray(x1)
+    x2 = np.ascontiguousarray(x2)
+
     # To save time, if Euclidean distance is used, use its square
     is_euclidean = False
     if metric_id == EUCLIDEAN:
@@ -566,6 +574,9 @@ def closest_pair(
     float
         Distance between the two points
     """
+    x1 = np.ascontiguousarray(x1)
+    x2 = np.ascontiguousarray(x2)
+
     # To save time, if Euclidean distance is used, use its square
     is_euclidean = False
     if metric_id == EUCLIDEAN:

src/spine/utils/gnn/cluster.py

@@ -846,7 +846,7 @@ def _get_cluster_features_base(
     for k in nb.prange(len(clusts)):
         # Get list of voxels in the cluster
         clust = clusts[ids[k]]
-        x = data[clust][:, COORD_COLS_LO:COORD_COLS_HI]
+        x = np.ascontiguousarray(data[clust][:, COORD_COLS_LO:COORD_COLS_HI])
 
         # Get cluster center
         center = sm.mean(x, 0)

src/spine/version.py

@@ -1,3 +1,3 @@
 """Module which stores the current software version."""
 
-__version__ = "0.14.0"
+__version__ = "0.14.1"

test/test_math/test_distance_regression.py

@@ -1,5 +1,6 @@
 """Regression tests for distance helpers."""
 
+import numba as nb
 import numpy as np
 import pytest
 
@@ -17,6 +18,13 @@
     [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.0, 0.0]],
     dtype=np.float32,
 )
+SQEUCLIDEAN_METRIC = METRICS["sqeuclidean"]
+
+
+@nb.njit(nb.float32[:, :](nb.float32[:, :]), cache=True)
+def _pdist_sqeuclidean_arbitrary_layout(points):
+    """Compile pdist from an arbitrary-layout caller, as PPN does."""
+    return pdist(points, SQEUCLIDEAN_METRIC)
 
 
 def test_get_metric_id_dispatches_minkowski_aliases_and_errors():
@@ -58,6 +66,17 @@ def test_pdist_dispatches_all_metrics_and_errors():
         pdist(POINTS, np.int64(99))
 
 
+def test_pdist_accepts_arbitrary_layout_from_jitted_callers():
+    """Nested Numba callers may pass arbitrary-layout row views."""
+    distances = _pdist_sqeuclidean_arbitrary_layout(POINTS)
+
+    np.testing.assert_allclose(
+        distances,
+        [[0.0, 1.0, 4.0], [1.0, 0.0, 5.0], [4.0, 5.0, 0.0]],
+        atol=1e-5,
+    )
+
+
 def test_cdist_dispatches_all_metrics_and_errors():
     """Cross-distance matrix should support every metric enumerator."""
     x1 = POINTS[:2]

test/test_utils/test_gnn_cluster.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from spine.utils.gnn.cluster import cluster_dedx
+from spine.utils.gnn.cluster import cluster_dedx, get_cluster_features_base
 
 
 def test_cluster_dedx_accepts_mixed_coordinate_dtypes():
@@ -17,3 +17,28 @@ def test_cluster_dedx_accepts_mixed_coordinate_dtypes():
     dedx = cluster_dedx(voxels, values, start, 5.0, True)
 
     assert dedx == np.float32(3.0)
+
+
+def test_cluster_features_base_accepts_indexed_float32_coordinates():
+    """Indexed cluster coordinate views should compile through Numba helpers."""
+    data = np.array(
+        [
+            [0.0, 0.0, 0.0, 1.0],
+            [1.0, 0.0, 0.0, 2.0],
+            [0.0, 1.0, 0.0, 3.0],
+            [2.0, 0.0, 0.0, 4.0],
+            [2.0, 1.0, 0.0, 5.0],
+            [3.0, 0.0, 0.0, 6.0],
+        ],
+        dtype=np.float32,
+    )
+    clusts = [
+        np.array([0, 2, 1], dtype=np.int64),
+        np.array([3, 5, 4], dtype=np.int64),
+    ]
+
+    feats = get_cluster_features_base(data, clusts)
+
+    assert feats.shape == (2, 16)
+    np.testing.assert_allclose(feats[0, :3], [1.0 / 3.0, 0.0, 2.0])
+    np.testing.assert_allclose(feats[1, :3], [1.0 / 3.0, 0.0, 5.0])

test/test_utils/test_gnn_network.py

@@ -6,7 +6,7 @@
 import numpy as np
 import pytest
 
-from spine.utils.gnn.network import inter_cluster_distance
+from spine.utils.gnn.network import get_cluster_edge_features, inter_cluster_distance
 
 GRAPH_BASE_PATH = (
     Path(__file__).resolve().parents[2]
@@ -93,3 +93,34 @@ def test_inter_cluster_distance_can_use_legacy_iterative_closest_pair():
     assert fixed_index[0, 1] == 11
     assert np.isclose(legacy_dist[0, 1], 0.7099366)
     assert legacy_index[0, 1] == 10
+
+
+def test_cluster_edge_features_accept_indexed_float32_coordinates_legacy():
+    """Indexed cluster coordinate views should compile in legacy CPA mode."""
+    data = np.array(
+        [
+            [0.0, 0.0, 0.0, 1.0],
+            [0.0, 1.0, 0.0, 2.0],
+            [0.0, 0.0, 1.0, 3.0],
+            [0.0, 4.0, 0.0, 4.0],
+            [0.0, 4.0, 1.0, 5.0],
+            [0.0, 5.0, 0.0, 6.0],
+        ],
+        dtype=np.float32,
+    )
+    clusts = [
+        np.array([0, 2, 1], dtype=np.int64),
+        np.array([3, 5, 4], dtype=np.int64),
+    ]
+    edge_index = np.array([[0, 1]], dtype=np.int64)
+
+    feats = get_cluster_edge_features(
+        data,
+        clusts,
+        edge_index,
+        iterative=False,
+        use_legacy_distance=True,
+    )
+
+    assert feats.shape == (1, 19)
+    assert feats.dtype == np.float32