Make spatial_neighbors faster

CHANGELOG.md

@@ -1,3 +1,8 @@
+## [1.0.3] - xxxx-xx-xx
+
+### Changed
+- Make `novae.spatial_neighbors` faster (#40)
+
 ## [1.0.2] - 2026-01-15
 
 ### Added

novae/utils/build.py

@@ -8,7 +8,6 @@
 from collections.abc import Iterable
 from enum import Enum
 from functools import partial
-from itertools import chain
 from typing import Literal, get_args
 
 import numpy as np
@@ -17,7 +16,6 @@
 from anndata.utils import make_index_unique
 from scipy.sparse import SparseEfficiencyWarning, block_diag, csr_matrix, spmatrix
 from scipy.spatial import Delaunay
-from sklearn.metrics.pairwise import euclidean_distances
 from sklearn.neighbors import NearestNeighbors
 
 from .._constants import Keys, Nums
@@ -247,14 +245,9 @@ def _build_connectivity(
         Adj = csr_matrix((np.ones_like(indices, dtype=np.float64), indices, indptr), shape=(N, N))
 
         if return_distance:
-            # fmt: off
-            dists = np.array(list(chain(*(
-                euclidean_distances(coords[indices[indptr[i] : indptr[i + 1]], :], coords[np.newaxis, i, :])
-                for i in range(N)
-                if len(indices[indptr[i] : indptr[i + 1]])
-            )))).squeeze()
+            rows = np.repeat(np.arange(N), np.diff(indptr))
+            dists = np.linalg.norm(coords[rows] - coords[indices], axis=1)
             Dst = csr_matrix((dists, indices, indptr), shape=(N, N))
-            # fmt: on
     else:
         r = 1 if radius is None else radius if isinstance(radius, (int, float)) else max(radius)
         tree = NearestNeighbors(n_neighbors=n_neighs, radius=r, metric="euclidean")

tests/test_neighbors.py

@@ -1,8 +1,13 @@
+from itertools import chain
+
 import anndata
 import numpy as np
 import pandas as pd
 import pytest
 from anndata import AnnData
+from scipy.sparse import csr_matrix
+from scipy.spatial import Delaunay
+from sklearn.metrics import euclidean_distances
 
 import novae
 from novae._constants import Keys
@@ -284,3 +289,34 @@ def test_change_n_hops():
 
     assert mean_adj_view2 / mean_adj_view > 1.5
     assert mean_adj_local2 / mean_adj_local > 1.5
+
+
+def test_new_distance_calculation() -> None:
+    coords = np.random.rand(40, 2)
+
+    N = coords.shape[0]
+
+    tri = Delaunay(coords)
+    indptr, indices = tri.vertex_neighbor_vertices
+
+    dists = np.array(
+        list(
+            chain(
+                *(
+                    euclidean_distances(coords[indices[indptr[i] : indptr[i + 1]], :], coords[np.newaxis, i, :])
+                    for i in range(N)
+                    if len(indices[indptr[i] : indptr[i + 1]])
+                )
+            )
+        )
+    ).squeeze()
+    Dst = csr_matrix((dists, indices, indptr), shape=(N, N))
+
+    rows = np.repeat(np.arange(N), np.diff(indptr))
+    dists = np.linalg.norm(coords[rows] - coords[indices], axis=1)
+    Dst2 = csr_matrix((dists, indices, indptr), shape=(N, N))
+
+    assert (Dst.indices == Dst2.indices).all()
+    assert (Dst.indptr == Dst2.indptr).all()
+
+    assert np.allclose(Dst.data, Dst2.data)