Fix LSD computation and add tests for empty labels

src/lsd_lite/lsds.py

@@ -192,8 +192,13 @@ def get_lsds(
         descriptor = upsample(descriptor, df)
         label_descriptors.append(descriptor * (segmentation == label)[None, ...])
 
-    descriptors = np.sum(np.array(label_descriptors), axis=0)
-    np.clip(descriptors, 0.0, 1.0, out=descriptors)
+    if len(label_descriptors) == 0:
+        # No valid labels found, return zeros with expected shape
+        channels = 10 if dims == 3 else 6
+        descriptors = np.zeros((channels,) + segmentation.shape, dtype=np.float32)
+    else:
+        descriptors = np.sum(np.array(label_descriptors), axis=0)
+        np.clip(descriptors, 0.0, 1.0, out=descriptors)
 
     return descriptors
 

tests/conftest.py

@@ -17,6 +17,11 @@ def toy_labels() -> np.ndarray:
     labels[-1] = 2
     return labels
 
+@pytest.fixture(scope="session")
+def empty_labels() -> np.ndarray:
+    """All-background (zeros) 3D volume."""
+    labels = np.zeros((10, 10, 10), dtype=np.uint64)
+    return labels
 
 @pytest.fixture(scope="session")
 def real_labels() -> np.ndarray:

tests/test_affs.py

@@ -28,7 +28,6 @@
     "3d_long_range": [16, 56, 56, 32, 32, 32, 32, 32, 32],
 }
 
-
 @pytest.mark.parametrize("nb_key", NEIGHBORHOODS.keys())
 @pytest.mark.parametrize("dist_func", ["equality", "equality-no-bg"])
 def test_real_shapes(real_labels, nb_key, dist_func):
@@ -61,3 +60,20 @@ def test_no_background_dist_func(toy_labels, nb_key):
 def test_real_affs_non_empty(real_labels, nb_key, dist_func):
     affs = get_affs(real_labels, NEIGHBORHOODS[nb_key], dist=dist_func, pad=True)
     assert affs.any(), f"{nb_key} – {dist_func} produced all-zero affinities"
+
+
+@pytest.mark.parametrize("nb_key", NEIGHBORHOODS.keys())
+def test_empty_labels(empty_labels, nb_key):
+    """Test that get_affs handles empty (all-background) segmentations."""
+    neigh = NEIGHBORHOODS[nb_key]
+
+    # With equality-no-bg, background should not produce affinities
+    affs = get_affs(empty_labels, neigh, dist="equality-no-bg", pad=True)
+    assert affs.shape == (len(neigh), *empty_labels.shape)
+    assert not affs.any()
+
+    # With equality, background matches itself, so all True
+    affs_eq = get_affs(empty_labels, neigh, dist="equality", pad=True)
+    assert affs_eq.shape == (len(neigh), *empty_labels.shape)
+    assert affs_eq.all()
+

tests/test_lsds.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from lsd_lite import get_lsds
 
 
@@ -12,3 +14,13 @@ def test_lsds(real_labels):
 def test_sigmas(real_labels):
     for sigma in [5, 10, 15, 20]:
         _ = get_lsds(real_labels[0], sigma=sigma, downsample=2)
+
+
+def test_empty_labels(empty_labels):
+    """Test that get_lsds handles empty (all-background) segmentations."""
+    lsds = get_lsds(empty_labels, sigma=5, downsample=1)
+
+    # Should return zeros with correct shape
+    assert lsds.shape == (10,) + empty_labels.shape
+    assert np.all(lsds == 0)
+