feat: add distances parameter to sc.pp.neighbors (#3627)

Co-authored-by: Philipp A. <flying-sheep@web.de>

Author: amalia-k510

docs/release-notes/3627.feat.md

@@ -0,0 +1 @@
+Added `distances` parameter to {func}`scanpy.pp.neighbors`, allowing to compute graphs from a precomputed distance matrix {smaller}`A. Karesh`

src/scanpy/metrics/_metrics.py

@@ -60,7 +60,7 @@ def confusion_matrix(
     orig, new = pd.Series(orig), pd.Series(new)
     assert len(orig) == len(new)
 
-    unique_labels = pd.unique(np.concatenate((orig.values, new.values)))
+    unique_labels = pd.unique(np.concatenate((orig.to_numpy(), new.to_numpy())))
 
     # Compute
     mtx = _confusion_matrix(orig, new, labels=unique_labels)

src/scanpy/neighbors/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import contextlib
+from inspect import signature
 from textwrap import indent
 from types import MappingProxyType
 from typing import TYPE_CHECKING, NamedTuple, TypedDict
@@ -19,18 +20,21 @@
 from .._utils import NeighborsView, _doc_params, get_literal_vals
 from . import _connectivity
 from ._common import (
+    _get_indices_distances_from_dense_matrix,
     _get_indices_distances_from_sparse_matrix,
     _get_sparse_matrix_from_indices_distances,
 )
+from ._connectivity import umap
 from ._doc import doc_n_pcs, doc_use_rep
 from ._types import _KnownTransformer, _Method
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping, MutableMapping
-    from typing import Any, Literal, NotRequired, TypeAlias
+    from typing import Any, Literal, NotRequired, TypeAlias, Unpack
 
     from anndata import AnnData
     from igraph import Graph
+    from numpy.typing import NDArray
 
     from .._utils.random import _LegacyRandom
     from ._types import KnnTransformerLike, _Metric, _MetricFn
@@ -56,11 +60,18 @@ class KwdsForTransformer(TypedDict):
     random_state: _LegacyRandom
 
 
+class NeighborsDict(TypedDict):  # noqa: D101
+    connectivities_key: str
+    distances_key: str
+    params: NeighborsParams
+    rp_forest: NotRequired[RPForestDict]
+
+
 class NeighborsParams(TypedDict):  # noqa: D101
     n_neighbors: int
     method: _Method
     random_state: _LegacyRandom
-    metric: _Metric | _MetricFn
+    metric: _Metric | _MetricFn | None
     metric_kwds: NotRequired[Mapping[str, Any]]
     use_rep: NotRequired[str]
     n_pcs: NotRequired[int]
@@ -72,11 +83,12 @@ def neighbors(  # noqa: PLR0913
     n_neighbors: int = 15,
     n_pcs: int | None = None,
     *,
+    distances: np.ndarray | SpBase | None = None,
     use_rep: str | None = None,
     knn: bool = True,
     method: _Method = "umap",
     transformer: KnnTransformerLike | _KnownTransformer | None = None,
-    metric: _Metric | _MetricFn = "euclidean",
+    metric: _Metric | _MetricFn | None = None,
     metric_kwds: Mapping[str, Any] = MappingProxyType({}),
     random_state: _LegacyRandom = 0,
     key_added: str | None = None,
@@ -138,6 +150,8 @@ def neighbors(  # noqa: PLR0913
                Use :func:`rapids_singlecell.pp.neighbors` instead.
     metric
         A known metric’s name or a callable that returns a distance.
+        If `distances` is given, this parameter is simply stored in `.uns` (see below),
+        otherwise defaults to `'euclidean'`.
 
         *ignored if ``transformer`` is an instance.*
     metric_kwds
@@ -153,18 +167,18 @@ def neighbors(  # noqa: PLR0913
         distances and connectivities are stored in `.obsp['distances']` and
         `.obsp['connectivities']` respectively.
         If specified, the neighbors data is added to .uns[key_added],
-        distances are stored in `.obsp[key_added+'_distances']` and
-        connectivities in `.obsp[key_added+'_connectivities']`.
+        distances are stored in `.obsp[f'{{key_added}}_distances']` and
+        connectivities in `.obsp[f'{{key_added}}_connectivities']`.
     copy
         Return a copy instead of writing to adata.
 
     Returns
     -------
     Returns `None` if `copy=False`, else returns an `AnnData` object. Sets the following fields:
 
-    `adata.obsp['distances' | key_added+'_distances']` : :class:`scipy.sparse.csr_matrix` (dtype `float`)
+    `adata.obsp['distances' | f'{{key_added}}_distances']` : :class:`scipy.sparse.csr_matrix` (dtype `float`)
         Distance matrix of the nearest neighbors search. Each row (cell) has `n_neighbors`-1 non-zero entries. These are the distances to their `n_neighbors`-1 nearest neighbors (excluding the cell itself).
-    `adata.obsp['connectivities' | key_added+'_connectivities']` : :class:`scipy.sparse._csr.csr_matrix` (dtype `float`)
+    `adata.obsp['connectivities' | f'{{key_added}}_connectivities']` : :class:`scipy.sparse._csr.csr_matrix` (dtype `float`)
         Weighted adjacency matrix of the neighborhood graph of data
         points. Weights should be interpreted as connectivities.
     `adata.uns['neighbors' | key_added]` : :class:`dict`
@@ -189,68 +203,108 @@ def neighbors(  # noqa: PLR0913
     :doc:`/how-to/knn-transformers`
 
     """
-    start = logg.info("computing neighbors")
-    adata = adata.copy() if copy else adata
-    if adata.is_view:  # we shouldn't need this here...
-        adata._init_as_actual(adata.copy())
-    neighbors = Neighbors(adata)
-    neighbors.compute_neighbors(
-        n_neighbors,
-        n_pcs=n_pcs,
-        use_rep=use_rep,
-        knn=knn,
+    if distances is None:
+        if metric is None:
+            metric = "euclidean"
+        start = logg.info("computing neighbors")
+        adata = adata.copy() if copy else adata
+        if adata.is_view:  # we shouldn't need this here...
+            adata._init_as_actual(adata.copy())
+        neighbors_ = Neighbors(adata)
+        neighbors_.compute_neighbors(
+            n_neighbors,
+            n_pcs=n_pcs,
+            use_rep=use_rep,
+            knn=knn,
+            method=method,
+            transformer=transformer,
+            metric=metric,
+            metric_kwds=metric_kwds,
+            random_state=random_state,
+        )
+    else:
+        params = locals()
+        if ignored := {
+            p.name
+            for p in signature(neighbors).parameters.values()
+            if p.name in {"use_rep", "knn", "n_pcs", "metric_kwds", "random_state"}
+            if params[p.name] != p.default
+        }:
+            warn(
+                f"Parameter(s) ignored if `distances` is given: {ignored}",
+                UserWarning,
+            )
+            random_state = 0
+        if callable(metric):
+            msg = "`metric` must be a string if `distances` is given."
+            raise TypeError(msg)
+        start = logg.info("computing connectivities")
+        # if a precomputed distance matrix is provided, skip the PCA and distance computation
+        if isinstance(distances, SpBase):
+            if TYPE_CHECKING:
+                from scipy.sparse._base import _spbase
+
+                assert isinstance(distances, _spbase)
+            distances = distances.tocsr(copy=True)
+            distances.setdiag(0)
+            distances.eliminate_zeros()
+        else:
+            distances = np.asarray(distances)
+            np.fill_diagonal(distances, 0)
+
+        neighbors_ = Neighbors(adata)
+        neighbors_.n_neighbors = n_neighbors
+        neighbors_.knn = True
+        neighbors_._distances = distances
+        neighbors_._connectivities = neighbors_._compute_connectivites(method)
+
+    key_added, neighbors_dict = _get_metadata(
+        key_added,
+        n_neighbors=neighbors_.n_neighbors,
         method=method,
-        transformer=transformer,
-        metric=metric,
-        metric_kwds=metric_kwds,
         random_state=random_state,
+        metric=metric,
+        **({} if not metric_kwds else dict(metric_kwds=metric_kwds)),
+        **({} if use_rep is None else dict(use_rep=use_rep)),
+        **({} if n_pcs is None else dict(n_pcs=n_pcs)),
     )
 
-    if key_added is None:
-        key_added = "neighbors"
-        conns_key = "connectivities"
-        dists_key = "distances"
-    else:
-        conns_key = f"{key_added}_connectivities"
-        dists_key = f"{key_added}_distances"
-
-    adata.uns[key_added] = {}
-
-    neighbors_dict = adata.uns[key_added]
+    if neighbors_.rp_forest is not None:
+        neighbors_dict["rp_forest"] = neighbors_.rp_forest
 
-    neighbors_dict["connectivities_key"] = conns_key
-    neighbors_dict["distances_key"] = dists_key
+    adata.uns[key_added] = neighbors_dict
+    adata.obsp[neighbors_dict["distances_key"]] = neighbors_.distances
+    adata.obsp[neighbors_dict["connectivities_key"]] = neighbors_.connectivities
 
-    neighbors_dict["params"] = NeighborsParams(
-        n_neighbors=neighbors.n_neighbors,
-        method=method,
-        random_state=random_state,
-        metric=metric,
-    )
-    if metric_kwds:
-        neighbors_dict["params"]["metric_kwds"] = metric_kwds
-    if use_rep is not None:
-        neighbors_dict["params"]["use_rep"] = use_rep
-    if n_pcs is not None:
-        neighbors_dict["params"]["n_pcs"] = n_pcs
-
-    adata.obsp[dists_key] = neighbors.distances
-    adata.obsp[conns_key] = neighbors.connectivities
-
-    if neighbors.rp_forest is not None:
-        neighbors_dict["rp_forest"] = neighbors.rp_forest
     logg.info(
         "    finished",
         time=start,
         deep=(
             f"added to `.uns[{key_added!r}]`\n"
-            f"    `.obsp[{dists_key!r}]`, distances for each pair of neighbors\n"
-            f"    `.obsp[{conns_key!r}]`, weighted adjacency matrix"
+            f"    `.obsp[{neighbors_dict['distances_key']!r}]`, distances for each pair of neighbors\n"
+            f"    `.obsp[{neighbors_dict['connectivities_key']!r}]`, weighted adjacency matrix"
         ),
     )
     return adata if copy else None
 
 
+def _get_metadata(
+    key_added: str | None,
+    **params: Unpack[NeighborsParams],
+) -> tuple[str, NeighborsDict]:
+    if key_added is None:
+        return "neighbors", NeighborsDict(
+            connectivities_key="connectivities",
+            distances_key="distances",
+            params=params,
+        )
+    return key_added, NeighborsDict(
+        connectivities_key=f"{key_added}_connectivities",
+        distances_key=f"{key_added}_distances",
+        params=params,
+    )
+
+
 class FlatTree(NamedTuple):  # noqa: D101
     hyperplanes: None
     offsets: None
@@ -358,7 +412,7 @@ class Neighbors:
     n_dcs
         Number of diffusion components to use.
     neighbors_key
-        Where to look in `.uns` and `.obsp` for neighbors data
+        Where to look in `.uns` and `.obsp` for neighbors data.
 
     """
 
@@ -518,7 +572,7 @@ def to_igraph(self) -> Graph:
         return _utils.get_igraph_from_adjacency(self.connectivities)
 
     @_doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep)
-    def compute_neighbors(  # noqa: PLR0912
+    def compute_neighbors(
         self,
         n_neighbors: int = 30,
         n_pcs: int | None = None,
@@ -601,26 +655,11 @@ def compute_neighbors(  # noqa: PLR0912
                     self._rp_forest = _make_forest_dict(index)
         start_connect = logg.debug("computed neighbors", time=start_neighbors)
 
-        if method == "umap":
-            self._connectivities = _connectivity.umap(
-                knn_indices,
-                knn_distances,
-                n_obs=self._adata.shape[0],
-                n_neighbors=self.n_neighbors,
-            )
-        elif method == "gauss":
-            self._connectivities = _connectivity.gauss(
-                self._distances, self.n_neighbors, knn=self.knn
-            )
-        elif method == "jaccard":
-            self._connectivities = _connectivity.jaccard(
-                knn_indices,
-                n_obs=self._adata.shape[0],
-                n_neighbors=self.n_neighbors,
-            )
-        elif method is not None:
-            msg = f"{method!r} should have been coerced in _handle_transform_args"
-            raise AssertionError(msg)
+        self._connectivities = (
+            None
+            if method is None
+            else self._compute_connectivites(method, (knn_indices, knn_distances))
+        )
         self._number_connected_components = 1
         if isinstance(self._connectivities, CSBase):
             from scipy.sparse.csgraph import connected_components
@@ -630,6 +669,43 @@ def compute_neighbors(  # noqa: PLR0912
         if method is not None:
             logg.debug("computed connectivities", time=start_connect)
 
+    def _compute_connectivites(
+        self,
+        method: _Method,
+        knn_ind_dist: (
+            tuple[NDArray[np.int32 | np.int64], NDArray[np.float32 | np.float64]] | None
+        ) = None,
+    ) -> CSRBase | NDArray[np.float32 | np.float64] | None:
+        def get_knn():
+            if knn_ind_dist is not None:
+                return knn_ind_dist
+            if isinstance(self._distances, CSBase):
+                return _get_indices_distances_from_sparse_matrix(
+                    self._distances.tocsr(), self.n_neighbors
+                )
+            assert self._distances is not None
+            return _get_indices_distances_from_dense_matrix(
+                self._distances, self.n_neighbors
+            )
+
+        if method == "umap":
+            knn_indices, knn_distances = get_knn()
+            return umap(
+                knn_indices,
+                knn_distances,
+                n_obs=self._adata.n_obs,
+                n_neighbors=self.n_neighbors,
+            )
+        if method == "gauss":
+            return _connectivity.gauss(self._distances, self.n_neighbors, knn=self.knn)
+        if method == "jaccard":
+            knn_indices, _ = get_knn()
+            return _connectivity.jaccard(
+                knn_indices, n_obs=self._adata.n_obs, n_neighbors=self.n_neighbors
+            )
+        msg = f"Method {method} not implemented."
+        raise NotImplementedError(msg)
+
     def _handle_transformer(
         self,
         method: _Method | Literal["gauss"] | None,

src/scanpy/neighbors/_connectivity.py

@@ -14,7 +14,7 @@
 )
 
 
-def gauss[D: (NDArray[np.float32], CSRBase)](  # noqa: PLR0912
+def gauss[D: (NDArray[np.float32 | np.float64], CSRBase)](  # noqa: PLR0912
     distances: D, n_neighbors: int, *, knn: bool
 ) -> D:
     """Derive gaussian connectivities between data points from their distances.