Fix for spatial_neighbors breaks when transform='spectral' (#1028)

* fix

* attempt to add tests and add better documentation to the functions to describe what's happening

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add vibe coded tests prune later

* add tests for spectral transform

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* complete the merge conflict issue

* add fau and clean up code

* make loops parallel

* specify fast arrayutils dep

* Add fast_array_utils to project dependencies

* cache kernel

* forgot to save file bf commit

* Apply suggestion from @flying-sheep

* remove unused imports (idk why this didn't fail on linter)

* remove more unused imports

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* move to float32

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add tests for spectral transform

* add fau and clean up code

* specify fast arrayutils dep

* Add fast_array_utils to project dependencies

* Apply suggestion from @flying-sheep

* remove unused imports (idk why this didn't fail on linter)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* merge conflix fix on deps

* format imports again

* fix fau import

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Philipp A. <[email protected]>

Author: 3 people

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
   "dask[array]>=2021.2,<=2024.11.2",
   "dask-image>=0.5",
   "docrep>=0.3.1",
+  "fast-array-utils",
   "fsspec>=2021.11",
   "imagecodecs>=2025.8.2,<2026",
   "matplotlib>=3.3",

src/squidpy/gr/_build.py

@@ -13,11 +13,13 @@
 import pandas as pd
 from anndata import AnnData
 from anndata.utils import make_index_unique
-from numba import njit
+from fast_array_utils import stats as fau_stats
+from numba import njit, prange
 from scanpy import logging as logg
 from scipy.sparse import (
     SparseEfficiencyWarning,
     block_diag,
+    csr_array,
     csr_matrix,
     isspmatrix_csr,
     spmatrix,
@@ -28,9 +30,16 @@
 from sklearn.neighbors import NearestNeighbors
 from spatialdata import SpatialData
 from spatialdata._core.centroids import get_centroids
-from spatialdata._core.query.relational_query import get_element_instances, match_element_to_table
+from spatialdata._core.query.relational_query import (
+    get_element_instances,
+    match_element_to_table,
+)
 from spatialdata.models import get_table_keys
-from spatialdata.models.models import Labels2DModel, Labels3DModel, get_model
+from spatialdata.models.models import (
+    Labels2DModel,
+    Labels3DModel,
+    get_model,
+)
 
 from squidpy._constants._constants import CoordType, Transform
 from squidpy._constants._pkg_constants import Key
@@ -380,7 +389,7 @@ def _build_connectivity(
     if delaunay:
         tri = Delaunay(coords)
         indptr, indices = tri.vertex_neighbor_vertices
-        Adj = csr_matrix((np.ones_like(indices, dtype=np.float64), indices, indptr), shape=(N, N))
+        Adj = csr_matrix((np.ones_like(indices, dtype=np.float32), indices, indptr), shape=(N, N))
 
         if return_distance:
             # fmt: off
@@ -415,7 +424,7 @@ def _build_connectivity(
             col_indices = np.concatenate(col_indices)
 
         Adj = csr_matrix(
-            (np.ones_like(row_indices, dtype=np.float64), (row_indices, col_indices)),
+            (np.ones_like(row_indices, dtype=np.float32), (row_indices, col_indices)),
             shape=(N, N),
         )
         if return_distance:
@@ -431,28 +440,67 @@ def _build_connectivity(
 
 
 @njit
-def outer(indices: NDArrayA, indptr: NDArrayA, degrees: NDArrayA) -> NDArrayA:
-    res = np.empty_like(indices, dtype=np.float64)
-    start = 0
-    for i in range(len(indptr) - 1):
-        ixs = indices[indptr[i] : indptr[i + 1]]
-        res[start : start + len(ixs)] = degrees[i] * degrees[ixs]
-        start += len(ixs)
+def _csr_bilateral_diag_scale_helper(
+    mat: csr_array | csr_matrix,
+    degrees: NDArrayA,
+) -> NDArrayA:
+    """
+    Return an array F aligned with CSR non-zeros such that
+    F[k] = d[i] * data[k] * d[j] for the k-th non-zero (i, j) in CSR order.
+
+    Parameters
+    ----------
+
+    data : array of float
+        CSR `data` (non-zero values).
+    indices : array of int
+        CSR `indices` (column indices).
+    indptr : array of int
+        CSR `indptr` (row pointer).
+    degrees : array of float, shape (n,)
+        Diagonal scaling vector.
+
+    Returns
+    -------
+    array of float
+        Length equals len(data). Entry-wise factors d_i * d_j * data[k]
+    """
+
+    res = np.empty_like(mat.data, dtype=np.float32)
+    for i in prange(len(mat.indptr) - 1):
+        ixs = mat.indices[mat.indptr[i] : mat.indptr[i + 1]]
+        res[mat.indptr[i] : mat.indptr[i + 1]] = degrees[i] * degrees[ixs] * mat.data[mat.indptr[i] : mat.indptr[i + 1]]
 
     return res
 
 
+def symmetric_normalize_csr(adj: spmatrix) -> csr_matrix:
+    """
+    Return D^{-1/2} * A * D^{-1/2}, where D = diag(degrees(A)) and A = adj.
+
+
+    Parameters
+    ----------
+    adj : scipy.sparse.csr_matrix
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
+    """
+    degrees = np.squeeze(np.array(np.sqrt(1.0 / fau_stats.sum(adj, axis=0))))
+    if adj.shape[0] != len(degrees):
+        raise ValueError("len(degrees) must equal number of rows of adj")
+    res_data = _csr_bilateral_diag_scale_helper(adj, degrees)
+    return csr_matrix((res_data, adj.indices, adj.indptr), shape=adj.shape)
+
+
 def _transform_a_spectral(a: spmatrix) -> spmatrix:
     if not isspmatrix_csr(a):
         a = a.tocsr()
     if not a.nnz:
         return a
 
-    degrees = np.squeeze(np.array(np.sqrt(1.0 / a.sum(axis=0))))
-    a = a.multiply(outer(a.indices, a.indptr, degrees))
-    a.eliminate_zeros()
-
-    return a
+    return symmetric_normalize_csr(a)
 
 
 def _transform_a_cosine(a: spmatrix) -> spmatrix:

tests/graph/test_spatial_neighbors.py

@@ -49,7 +49,13 @@ def _adata_concat(adata1, adata2):
     # TODO: add edge cases
     # TODO(giovp): test with reshuffling
     @pytest.mark.parametrize(("n_rings", "n_neigh", "sum_dist"), [(1, 6, 0), (2, 18, 30), (3, 36, 84)])
-    def test_spatial_neighbors_visium(self, visium_adata: AnnData, n_rings: int, n_neigh: int, sum_dist: int):
+    def test_spatial_neighbors_visium(
+        self,
+        visium_adata: AnnData,
+        n_rings: int,
+        n_neigh: int,
+        sum_dist: int,
+    ):
         """
         check correctness of neighborhoods for visium coordinates
         """
@@ -355,3 +361,72 @@ def test_mask_graph(
                 negative_mask=True,
                 key_added=key_added,
             )
+
+    def test_spatial_neighbors_transform_mathematical_properties(self, non_visium_adata: AnnData):
+        """
+        Test mathematical properties of each transform.
+        """
+        # Test spectral transform properties
+        spatial_neighbors(non_visium_adata, delaunay=True, coord_type=None, transform="spectral")
+        adj_spectral = non_visium_adata.obsp[Key.obsp.spatial_conn()].toarray()
+
+        # Spectral transform should be symmetric
+        np.testing.assert_allclose(adj_spectral, adj_spectral.T, atol=1e-10)
+
+        # Spectral transform should have normalized rows (L2 norm <= 1)
+        row_norms = np.sqrt(np.sum(adj_spectral**2, axis=1))
+        np.testing.assert_array_less(row_norms, 1.0 + 1e-10)
+
+        # Test cosine transform properties
+        spatial_neighbors(non_visium_adata, delaunay=True, coord_type=None, transform="cosine")
+        adj_cosine = non_visium_adata.obsp[Key.obsp.spatial_conn()].toarray()
+
+        # Cosine transform should be symmetric
+        np.testing.assert_allclose(adj_cosine, adj_cosine.T, atol=1e-10)
+
+        # Cosine transform should have values in [-1, 1]
+        np.testing.assert_array_less(-1.0 - 1e-10, adj_cosine)
+        np.testing.assert_array_less(adj_cosine, 1.0 + 1e-10)
+
+        # Diagonal of cosine transform should be 1 (self-similarity)
+        np.testing.assert_allclose(np.diag(adj_cosine), 1.0, atol=1e-10)
+
+    def test_spatial_neighbors_transform_edge_cases(self, non_visium_adata: AnnData):
+        """
+        Test transforms with edge cases (empty graph, single node, etc.).
+        """
+        # Test with a very small dataset
+        small_adata = non_visium_adata[:5].copy()  # Only 5 points
+
+        # Test all transforms with small dataset
+        for transform in [None, "spectral", "cosine"]:
+            spatial_neighbors(small_adata, delaunay=True, coord_type=None, transform=transform)
+            assert Key.obsp.spatial_conn() in small_adata.obsp
+            assert Key.obsp.spatial_dist() in small_adata.obsp
+
+            # Verify transform parameter is saved
+            assert small_adata.uns[Key.uns.spatial_neighs()]["params"]["transform"] == transform
+
+    def test_spatial_neighbors_spectral_transform_properties(self, non_visium_adata: AnnData):
+        """
+        Test that spectral transform preserves nonzero pattern and normalizes rows to sum to 1.
+        """
+        # Apply spatial_neighbors without transform
+        spatial_neighbors(non_visium_adata, delaunay=True, coord_type=None, transform=None)
+        adj_no_transform = non_visium_adata.obsp[Key.obsp.spatial_conn()].copy()
+
+        # Apply spatial_neighbors with spectral transform
+        spatial_neighbors(non_visium_adata, delaunay=True, coord_type=None, transform="spectral")
+        adj_spectral = non_visium_adata.obsp[Key.obsp.spatial_conn()]
+
+        # Check that nonzero patterns are identical
+        np.testing.assert_array_equal(
+            adj_no_transform.nonzero(),
+            adj_spectral.nonzero(),
+            err_msg="Spectral transform should preserve the sparsity pattern",
+        )
+
+        w = np.linalg.eigvals(adj_spectral.toarray())
+        # Eigenvalues should be in range [-1, 1]
+        np.testing.assert_array_less(w, 1.0, err_msg="Eigenvalues should be <= 1")
+        np.testing.assert_array_less(-1.0, w, err_msg="Eigenvalues should be >= -1")