Fix for spatial_neighbors breaks when transform='spectral'

src/squidpy/gr/_build.py

@@ -439,7 +439,26 @@ def _build_connectivity(
 
 
 @njit
-def outer(indices: NDArrayA, indptr: NDArrayA, degrees: NDArrayA) -> NDArrayA:
+def _didj_factors_csr(indices: NDArrayA, indptr: NDArrayA, degrees: NDArrayA) -> NDArrayA:
+    """
+    Return an array F aligned with CSR nonzeros such that
+    F[k] = d[i] * d[j] for the k-th nonzero (i, j) in CSR order.
+
+    Parameters
+    ----------
+    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(indices). Entry-wise factors d_i * d_j.
+    """
+
     res = np.empty_like(indices, dtype=np.float64)
     start = 0
     for i in range(len(indptr) - 1):
@@ -450,17 +469,60 @@ def outer(indices: NDArrayA, indptr: NDArrayA, degrees: NDArrayA) -> NDArrayA:
     return res
 
 
+def csr_bilateral_diag_scale(A: spmatrix, degrees: NDArrayA) -> csr_matrix:
+    """
+    Return D * A * D where D = diag(d) and A is CSR.
+
+    For every stored nonzero (i, j):
+        (D A D)_{ij} = d[i] * A_{ij} * d[j].
+
+    Parameters
+    ----------
+    A : scipy.sparse.csr_matrix
+        Input sparse matrix. Converted to CSR if needed.
+    degrees : array_like, shape (n,)
+        Diagonal vector. len(degrees) must equal A.shape[0].
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
+        Matrix with same sparsity pattern as A.
+    """
+    if A.shape[0] != len(degrees):
+        raise ValueError("len(d) must equal number of rows of A")
+    factors = _didj_factors_csr(A.indices, A.indptr, degrees)
+    F = csr_matrix((factors, A.indices, A.indptr), shape=A.shape)
+    B = A.multiply(F)
+    B.eliminate_zeros()
+    return B
+
+
+def symmetric_normalize_csr(A: spmatrix) -> csr_matrix:
+    """
+    Return D^{-1/2} * A * D^{-1/2}, where D = diag(degrees).
+
+    degrees[i] = sum_j A_{ij}  (row degree).
+    Zeros in degree are left as-is to avoid division by zero.
+
+    Parameters
+    ----------
+    A : scipy.sparse.csr_matrix
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
+    """
+    degrees = np.squeeze(np.array(np.sqrt(1.0 / A.sum(axis=0))))
+    return csr_bilateral_diag_scale(A, degrees)
+
+
 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
         """
@@ -357,3 +363,48 @@ 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