diff --git a/benchmarks/benchmarks/preprocessing_log.py b/benchmarks/benchmarks/preprocessing_log.py index 350bb66883..3243ac1ec1 100644 --- a/benchmarks/benchmarks/preprocessing_log.py +++ b/benchmarks/benchmarks/preprocessing_log.py @@ -67,6 +67,42 @@ def peakmem_scale(self, *_) -> None: sc.pp.scale(self.adata, max_value=10) +class NeighborsSuite: + """Benchmark neighbor graph construction. + + Both `pp.neighbors` and `pp.bbknn` pick an exact or approximate kNN backend + depending on how many observations they have to index, + so the small and the big dataset cover the two paths. + Both have a batch key, as `pp.bbknn` needs one. + """ + + params: tuple[list[Dataset]] = (["bmmc", "lung93k"],) + param_names = ("dataset",) + + def setup_cache(self) -> None: + """Without this caching, asv was running several processes which meant the data was repeatedly downloaded.""" + for dataset in self.params[0]: + adata, batch_key = get_dataset(dataset) + sc.pp.pca(adata) # so we time the kNN search, not the PCA + adata.uns["batch_key"] = batch_key + adata.write_zarr(f"{dataset}.zarr") + + def setup(self, dataset: Dataset) -> None: + self.adata = ad.read_zarr(f"{dataset}.zarr") + + def time_neighbors(self, *_) -> None: + sc.pp.neighbors(self.adata) + + def peakmem_neighbors(self, *_) -> None: + sc.pp.neighbors(self.adata) + + def time_bbknn(self, *_) -> None: + sc.pp.bbknn(self.adata, batch_key=self.adata.uns["batch_key"]) + + def peakmem_bbknn(self, *_) -> None: + sc.pp.bbknn(self.adata, batch_key=self.adata.uns["batch_key"]) + + class HVGSuite: # noqa: D101 params = (["seurat_v3", "cell_ranger", "seurat"], [True, False]) param_names = ("flavor", "use_dask") diff --git a/docs/api/preprocessing.md b/docs/api/preprocessing.md index 3b6c684deb..44261fc5aa 100644 --- a/docs/api/preprocessing.md +++ b/docs/api/preprocessing.md @@ -63,6 +63,7 @@ Note that a simple batch correction method is available via {func}`pp.regress_ou pp.harmony_integrate ``` +Batches can also be integrated at the level of the neighbor graph using {func}`pp.bbknn`. Also see {ref}`data integration tools ` and external {ref}`external data integration `. ## Doublet detection @@ -83,6 +84,7 @@ Also see {ref}`data integration tools ` and external {ref}`ext :nosignatures: :toctree: generated/ + pp.bbknn pp.neighbors ``` diff --git a/docs/conf.py b/docs/conf.py index 2b88b5b724..5920679ce6 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -178,6 +178,7 @@ array_support: dict[str, tuple[list[str], list[str]]] = { "experimental.pp.highly_variable_genes": (["np", "sp"], []), "get.aggregate": (["np", "sp", "da"], []), + "pp.bbknn": (["np", "sp"], []), "pp.calculate_qc_metrics": (["np", "sp", "da"], []), "pp.combat": (["np"], []), "pp.downsample_counts": (["np", "sp[csr]"], []), diff --git a/docs/release-notes/4306.feat.md b/docs/release-notes/4306.feat.md new file mode 100644 index 0000000000..c952efe4ba --- /dev/null +++ b/docs/release-notes/4306.feat.md @@ -0,0 +1 @@ +Add {func}`scanpy.pp.bbknn`, a native implementation of batch balanced kNN :cite:p:`Polanski2019` {smaller}`S Dicks` diff --git a/hatch.toml b/hatch.toml index 59baac1e9a..df3e9c1804 100644 --- a/hatch.toml +++ b/hatch.toml @@ -22,7 +22,7 @@ scripts.format-fix = "prek run -a --group=format" [envs.hatch-check-code] dependencies = [ "prek" ] -lint-check = "echo 'try `hatch check code --fix`'; false" +scripts.lint-check = "echo 'try `hatch check code --fix`'; false" scripts.lint-fix = "prek run -a --no-group=format" [envs.hatch-test] diff --git a/src/scanpy/_utils/__init__.py b/src/scanpy/_utils/__init__.py index 4ba294d12f..0ae8882ca1 100644 --- a/src/scanpy/_utils/__init__.py +++ b/src/scanpy/_utils/__init__.py @@ -47,7 +47,8 @@ from pandas._typing import Dtype as PdDtype from .._compat import CSRBase - from ..neighbors import NeighborsParams, RPForestDict + from ..neighbors import RPForestDict + from ..neighbors._types import NeighborsParams type _MemoryArray = NDArray | CSBase type _SupportedArray = _MemoryArray | DaskArray diff --git a/src/scanpy/external/pp/_bbknn.py b/src/scanpy/external/pp/_bbknn.py index 0a984fb59f..0c7f740e03 100644 --- a/src/scanpy/external/pp/_bbknn.py +++ b/src/scanpy/external/pp/_bbknn.py @@ -45,6 +45,7 @@ def bbknn( # noqa: PLR0913 This is just a wrapper of :func:`bbknn.bbknn`: up to date docstring, more information and bug reports there. + :func:`scanpy.pp.bbknn` implements the same algorithm without the extra dependency. Params ------ diff --git a/src/scanpy/neighbors/__init__.py b/src/scanpy/neighbors/__init__.py index bfae758eea..97a1fe1638 100644 --- a/src/scanpy/neighbors/__init__.py +++ b/src/scanpy/neighbors/__init__.py @@ -6,7 +6,7 @@ import sys from inspect import signature from textwrap import indent -from typing import TYPE_CHECKING, NamedTuple, TypedDict +from typing import TYPE_CHECKING, NamedTuple if sys.version_info < (3, 15): from types import MappingProxyType as frozendict # noqa: N813 @@ -21,7 +21,6 @@ from .._compat import CSBase, CSRBase, SpBase, pkg_version, warn from .._docs import doc_rng from .._keys import _EmbeddingKeys, _existing_preset_keys -from .._settings import settings from .._utils import NeighborsView, _doc_params, get_literal_vals from .._utils.random import ( _accepts_legacy_random_state, @@ -32,60 +31,35 @@ from ._common import ( _get_indices_distances_from_dense_matrix, _get_indices_distances_from_sparse_matrix, + _get_metadata, _get_sparse_matrix_from_indices_distances, + _make_transformer, ) from ._connectivity import umap from ._doc import doc_n_pcs, doc_use_rep -from ._types import _KnownTransformer, _Method +from ._types import KwdsForTransformer, _Method if TYPE_CHECKING: from collections.abc import Callable, Mapping, MutableMapping - from typing import Any, Literal, NotRequired, TypeAlias, Unpack + from typing import Any, Literal from anndata import AnnData from igraph import Graph from numpy.typing import NDArray - from .._utils.random import RNGLike, SeedLike, _LegacyRandom - from ._types import KnnTransformerLike, _Metric, _MetricFn - - # TODO: make `type` when https://github.com/sphinx-doc/sphinx/pull/13508 is released - RPForestDict: TypeAlias = Mapping[str, Mapping[str, np.ndarray]] # noqa: UP040 + from .._utils.random import RNGLike, SeedLike + from ._types import ( + KnnTransformerLike, + RPForestDict, + _KnownTransformer, + _Metric, + _MetricFn, + ) SCIPY_1_17 = pkg_version("scipy") >= Version("1.17") -class KwdsForTransformer(TypedDict): - """Keyword arguments passed to a _KnownTransformer. - - IMPORTANT: when changing the parameters set here, - update the “*ignored*” part in the parameter docs! - """ - - n_neighbors: int - metric: _Metric | _MetricFn - metric_params: Mapping[str, Any] - 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 | None - metric_kwds: NotRequired[Mapping[str, Any]] - use_rep: NotRequired[str] - n_pcs: NotRequired[int] - - @_doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep, rng=doc_rng) @_accepts_legacy_random_state(_DEFAULT_SEED := 0) def neighbors( # noqa: PLR0913 @@ -302,23 +276,6 @@ def neighbors( # noqa: PLR0913 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 @@ -725,13 +682,8 @@ def _handle_transformer( `method` will be coerced to `'gauss'`, `'umap'`, or `'jaccard'`. `transformer` is coerced from a str or instance to an instance class. - If `transformer` is `None` and there are few data points, - `transformer` will be set to a brute force - :class:`~sklearn.neighbors.KNeighborsTransformer`. - - If `transformer` is `None` and there are many data points, - `transformer` will be set like `umap` does (i.e. to a - ~`pynndescent.PyNNDescentTransformer` with custom `n_trees` and `n_iter`). + If `transformer` is `None`, it is chosen based on the number of data points, + see :func:`~scanpy.neighbors._common._make_transformer`. """ # legacy logic use_dense_distances = ( @@ -755,40 +707,16 @@ def _handle_transformer( # Coerce `transformer` to an instance if shortcut: - from sklearn.neighbors import KNeighborsTransformer - - assert transformer in {None, "sklearn"} n_neighbors = self._adata.n_obs - 1 if knn: # only obey n_neighbors arg if knn set n_neighbors = min(n_neighbors, kwds["n_neighbors"]) - transformer = KNeighborsTransformer( - algorithm="brute", - n_jobs=settings.n_jobs, - n_neighbors=n_neighbors, - metric=kwds["metric"], - metric_params=dict(kwds["metric_params"]), # needs dict - # no random_state - ) - elif transformer is None or transformer == "pynndescent": - from pynndescent import PyNNDescentTransformer - - kwds = kwds.copy() - kwds["metric_kwds"] = kwds.pop("metric_params") - if transformer is None: - # Use defaults from UMAP’s `nearest_neighbors` function - kwds.update( - n_jobs=settings.n_jobs, - n_trees=min(64, 5 + round((self._adata.n_obs) ** 0.5 / 20.0)), - n_iters=max(5, round(np.log2(self._adata.n_obs))), - ) - transformer = PyNNDescentTransformer(**kwds) - elif isinstance(transformer, str): - msg = ( - f"Unknown transformer: {transformer}. " - f"Try passing a class or one of {get_literal_vals(_KnownTransformer)}" - ) - raise ValueError(msg) - # else `transformer` is probably an instance + kwds = {**kwds, "n_neighbors": n_neighbors} + transformer = _make_transformer( + transformer, + shortcut=shortcut, + n_index=self._adata.n_obs, + **kwds, + ) return conn_method, transformer, shortcut def compute_transitions(self, *, density_normalize: bool = True) -> None: diff --git a/src/scanpy/neighbors/_bbknn.py b/src/scanpy/neighbors/_bbknn.py new file mode 100644 index 0000000000..a8aa269da1 --- /dev/null +++ b/src/scanpy/neighbors/_bbknn.py @@ -0,0 +1,372 @@ +"""Batch balanced k-nearest neighbors.""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +if sys.version_info < (3, 15): + from types import MappingProxyType as frozendict # noqa: N813 + +import numpy as np + +from .. import logging as logg +from .._docs import doc_rng +from .._utils import _doc_params +from .._utils.random import ( + _accepts_legacy_random_state, + _legacy_random_state, + _LegacyRng, +) +from ._common import ( + _get_indices_distances_from_rect_matrix, + _get_metadata, + _get_sparse_matrix_from_indices_distances, + _make_transformer, +) +from ._connectivity import umap +from ._doc import doc_n_pcs, doc_use_rep + +if TYPE_CHECKING: + from collections.abc import Mapping + from typing import Any + + from anndata import AnnData + from numpy.typing import NDArray + + from .._compat import CSRBase + from .._utils.random import RNGLike, SeedLike, _LegacyRandom + from ._types import ( + KnnTransformerLike, + _KnownTransformer, + _Metric, + _MetricFn, + ) + + +@_doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep, rng=doc_rng) +@_accepts_legacy_random_state(0) +def bbknn( # noqa: PLR0913 + adata: AnnData, + neighbors_within_batch: int = 3, + n_pcs: int | None = None, + *, + batch_key: str = "batch", + use_rep: str | None = None, + transformer: KnnTransformerLike | _KnownTransformer | None = None, + metric: _Metric | _MetricFn = "euclidean", + metric_kwds: Mapping[str, Any] = frozendict({}), + trim: int | None = None, + rng: SeedLike | RNGLike | None = None, + key_added: str | None = None, + copy: bool = False, +) -> AnnData | None: + """Compute a batch balanced neighborhood graph of observations :cite:p:`Polanski2019`. + + Batch balanced kNN alters the kNN procedure to identify each cell’s top neighbors + in each batch separately instead of the entire cell pool with no accounting for batch. + The nearest neighbors of each batch are then merged to create a final list of + neighbors for the cell, which aligns batches in a quick and lightweight manner. + + Use this as an alternative to :func:`~scanpy.pp.neighbors`: + it writes the same fields, so all downstream steps + (e.g. :func:`~scanpy.tl.umap` or :func:`~scanpy.tl.leiden`) work unchanged. + This CPU implementation is based on the rapids-singlecell package. + + .. array-support:: pp.bbknn + + Parameters + ---------- + adata + Annotated data matrix. + neighbors_within_batch + How many top neighbors to report for each batch. + The total number of neighbors is this number times the number of batches, + which then serves as the basis for the construction of a symmetrical + matrix of connectivities. + {n_pcs} + batch_key + `adata.obs` column name discriminating between the batches. + {use_rep} + transformer + kNN search backend following the API of + :class:`~sklearn.neighbors.KNeighborsTransformer`. + One index is built per batch and queried with all observations, + so its ``n_neighbors`` is ignored in favor of ``neighbors_within_batch``. + See :doc:`/how-to/knn-transformers` for more details. + Also accepts the following known options: + + `None` (the default) + Behavior depends on data size. + For small data, we will calculate exact kNN, otherwise we use + :class:`~pynndescent.pynndescent_.PyNNDescentTransformer` + `'pynndescent'` + :class:`~pynndescent.pynndescent_.PyNNDescentTransformer` + metric + A known metric’s name or a callable that returns a distance. + + *ignored if ``transformer`` is an instance.* + metric_kwds + Options for the metric. + + *ignored if ``transformer`` is an instance.* + trim + Trim each cell’s neighbors to these many top connectivities. + May help with population independence and improve the tidiness of clustering. + The lower the value, the more independent the individual populations, + at the cost of a more conserved batch effect. + If `None`, this is set to 10 times the total number of neighbors. + Set to 0 to skip trimming. + {rng} + + *ignored if ``transformer`` is an instance.* + key_added + If not specified, the neighbors data is stored in `.uns['neighbors']`, + 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[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' | f'{{key_added}}_distances']` : :class:`scipy.sparse.csr_matrix` (dtype `float`) + Distance matrix of the batch balanced nearest neighbors search. + Each row (cell) has ``neighbors_within_batch`` × ``n_batches`` - 1 non-zero entries: + its nearest neighbors in each batch, excluding the cell itself. + `adata.obsp['connectivities' | f'{{key_added}}_connectivities']` : :class:`scipy.sparse.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` + neighbors parameters. + + Examples + -------- + >>> import scanpy as sc + >>> adata = sc.datasets.pbmc68k_reduced() + >>> adata.obs["batch"] = adata.obs["phase"] + >>> sc.pp.bbknn(adata, batch_key="batch") + >>> sc.tl.umap(adata) + + See Also + -------- + :func:`~scanpy.pp.neighbors` + :doc:`/how-to/knn-transformers` + + """ + from ..tools._utils import _choose_representation + + start = logg.info("computing batch balanced neighbors") + + adata = adata.copy() if copy else adata + if adata.is_view: # we shouldn’t need this here... + adata._init_as_actual(adata.copy()) + + if neighbors_within_batch < 1: + msg = "`neighbors_within_batch` needs to be greater than 0." + raise ValueError(msg) + if batch_key not in adata.obs: + msg = f"Batch key {batch_key!r} not found in `adata.obs`." + raise KeyError(msg) + batches = np.asarray(adata.obs[batch_key]) + unique_batches, batch_sizes = np.unique(batches, return_counts=True) + if len(too_small := unique_batches[batch_sizes < neighbors_within_batch]): + msg = ( + f"Not all batches have at least `neighbors_within_batch = " + f"{neighbors_within_batch}` cells in them: {list(too_small)}." + ) + raise ValueError(msg) + + x = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs) + knn_indices, knn_distances = _compute_batch_balanced_knn( + x, + batches=batches, + unique_batches=unique_batches, + batch_sizes=batch_sizes, + neighbors_within_batch=neighbors_within_batch, + transformer=transformer, + metric=metric, + metric_kwds=metric_kwds, + random_state=_legacy_random_state(rng), + ) + n_obs, n_neighbors = knn_indices.shape + if trim is None: + trim = 10 * n_neighbors + + distances = _get_sparse_matrix_from_indices_distances( + knn_indices, knn_distances, keep_self=False + ) + start_connect = logg.debug("computed batch balanced neighbors", time=start) + connectivities = umap( + knn_indices, knn_distances, n_obs=n_obs, n_neighbors=n_neighbors + ) + if trim > 0: + connectivities = _trim(connectivities, trim) + logg.debug("computed connectivities", time=start_connect) + + key_added, neighbors_dict = _get_metadata( + key_added, + n_neighbors=n_neighbors, + method="umap", + metric=metric, + **(dict(random_state=rng.arg) if isinstance(rng, _LegacyRng) else {}), + **({} 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)), + batch_key=batch_key, + neighbors_within_batch=neighbors_within_batch, + trim=trim, + ) + adata.uns[key_added] = neighbors_dict + adata.obsp[neighbors_dict["distances_key"]] = distances + adata.obsp[neighbors_dict["connectivities_key"]] = connectivities + + logg.info( + " finished", + time=start, + deep=( + f"added to `.uns[{key_added!r}]`\n" + 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 _compute_batch_balanced_knn( + x: NDArray[np.float32 | np.float64] | CSRBase, + /, + *, + batches: NDArray[Any], + unique_batches: NDArray[Any], + batch_sizes: NDArray[np.int64], + neighbors_within_batch: int, + transformer: KnnTransformerLike | _KnownTransformer | None, + metric: _Metric | _MetricFn, + metric_kwds: Mapping[str, Any], + random_state: _LegacyRandom, +) -> tuple[NDArray[np.int64], NDArray[np.float32 | np.float64]]: + """Find the `neighbors_within_batch` nearest neighbors of each cell in each batch. + + Returns the merged indices and distances, sorted by distance within each row. + """ + from sklearn.base import clone + + proto, is_sklearn_shortcut = _handle_transformer( + transformer, + n_obs=x.shape[0], + max_batch_size=int(batch_sizes.max()), + n_neighbors=neighbors_within_batch, + metric=metric, + metric_kwds=metric_kwds, + random_state=random_state, + ) + + n_obs = x.shape[0] + n_neighbors = neighbors_within_batch * len(unique_batches) + knn_indices = np.empty((n_obs, n_neighbors), dtype=np.int64) + knn_distances = np.empty((n_obs, n_neighbors), dtype=np.float64) + obs_indices = np.arange(n_obs) + + for i, batch in enumerate(unique_batches): + mask = batches == batch + knn = clone(proto).fit(x[mask]) + d = ( + # ask for exactly `neighbors_within_batch`; `transform` would add one, + # which fails for batches that only have `neighbors_within_batch` cells + knn.kneighbors_graph(x, n_neighbors=neighbors_within_batch, mode="distance") + if is_sklearn_shortcut + else knn.transform(x) + ) + indices, distances = _get_indices_distances_from_rect_matrix( + d, neighbors_within_batch + ) + cols = slice(i * neighbors_within_batch, (i + 1) * neighbors_within_batch) + # the transformer’s indices are relative to the batch + knn_indices[:, cols] = obs_indices[mask][indices] + knn_distances[:, cols] = distances + logg.debug(f" computed neighbors within batch {batch!r}") + + # some backends report a tiny non-zero distance of a cell to itself, + # which `umap` would mistake for the radius of the cell’s local neighborhood + is_self = knn_indices == obs_indices[:, None] + knn_distances[is_self] = 0.0 + + # `umap` derives each cell’s local connectivity from its closest neighbors, + # so the merged rows need to be sorted by distance. + # Ties are broken in favor of the cell itself, which is dropped from `.obsp['distances']`. + order = np.lexsort((~is_self, knn_distances), axis=1) + return ( + np.take_along_axis(knn_indices, order, axis=1), + np.take_along_axis(knn_distances, order, axis=1), + ) + + +def _handle_transformer( + transformer: KnnTransformerLike | _KnownTransformer | None, + *, + n_obs: int, + max_batch_size: int, + n_neighbors: int, + metric: _Metric | _MetricFn, + metric_kwds: Mapping[str, Any], + random_state: _LegacyRandom, +) -> tuple[KnnTransformerLike, bool]: + """Coerce `transformer` to an instance to be cloned for each batch. + + Also returns whether it is a :class:`~sklearn.neighbors.KNeighborsTransformer` + we created ourselves, i.e. one we can query without going through ``transform``. + + Unlike :func:`~scanpy.pp.neighbors`, + we build one index per batch and query each with all `n_obs` observations, + so brute force costs ``n_obs × max_batch_size``, + while an approximate index’s cost is dominated by building it. + The cutoff is where the two met in benchmarks on ~50-dimensional data. + """ + shortcut = transformer == "sklearn" or ( + transformer is None + and ( + max_batch_size < 4096 + or (metric == "euclidean" and n_obs * max_batch_size < 10**9) + ) + ) + return _make_transformer( + transformer, + shortcut=shortcut, + n_index=max_batch_size, + n_neighbors=n_neighbors, + metric=metric, + metric_params=metric_kwds, + random_state=random_state, + ), shortcut + + +def _trim(connectivities: CSRBase, /, trim: int) -> CSRBase: + """Trim the graph in place to the `trim` strongest connections per cell. + + Following the reference implementation, an edge is dropped if its weight is + below the `trim`-th largest weight of *either* of the cells it connects, + which keeps the graph symmetric. + """ + n_nonzero = np.diff(connectivities.indptr) + if not (n_nonzero > trim).any(): + return connectivities + rows = np.repeat(np.arange(connectivities.shape[0]), n_nonzero) + # sort each row’s weights in descending order to find its `trim`-th largest one. + # rows with at most `trim` entries have no such weight and keep a cutoff of 0. + order = np.lexsort((-connectivities.data, rows)) + rank_in_row = np.arange(connectivities.nnz) - np.repeat( + connectivities.indptr[:-1], n_nonzero + ) + at_cutoff = rank_in_row == trim - 1 + cutoffs = np.zeros(connectivities.shape[0], dtype=connectivities.data.dtype) + cutoffs[rows[at_cutoff]] = connectivities.data[order][at_cutoff] + + keep_above = np.maximum(cutoffs[rows], cutoffs[connectivities.indices]) + connectivities.data[connectivities.data < keep_above] = 0 + connectivities.eliminate_zeros() + return connectivities diff --git a/src/scanpy/neighbors/_common.py b/src/scanpy/neighbors/_common.py index eb6e696b8e..6b8bdd899c 100644 --- a/src/scanpy/neighbors/_common.py +++ b/src/scanpy/neighbors/_common.py @@ -7,11 +7,83 @@ from scipy import sparse from .._compat import warn +from .._settings import settings +from .._utils import get_literal_vals +from ._types import NeighborsDict, _KnownTransformer if TYPE_CHECKING: + from typing import Unpack + from numpy.typing import NDArray from .._compat import CSRBase + from ._types import KnnTransformerLike, KwdsForTransformer, NeighborsParams + + +def _make_transformer( + transformer: KnnTransformerLike | _KnownTransformer | None, + /, + *, + shortcut: bool, + n_index: int, + **kwds: Unpack[KwdsForTransformer], +) -> KnnTransformerLike: + """Coerce `transformer` from `None` or a string to an instance. + + `shortcut` requests a brute force :class:`~sklearn.neighbors.KNeighborsTransformer`; + the caller decides that, as the trade-off depends on how the index is queried. + Otherwise, `transformer=None` is set up like `umap` does, i.e. to a + ~`pynndescent.PyNNDescentTransformer` with custom `n_trees` and `n_iters` + derived from `n_index`, the number of observations in the index. + """ + if shortcut: + from sklearn.neighbors import KNeighborsTransformer + + assert transformer in {None, "sklearn"} + return KNeighborsTransformer( + algorithm="brute", + n_jobs=settings.n_jobs, + n_neighbors=kwds["n_neighbors"], + metric=kwds["metric"], + metric_params=dict(kwds["metric_params"]), # needs dict + # no random_state + ) + if transformer is None or transformer == "pynndescent": + from pynndescent import PyNNDescentTransformer + + kwds["metric_kwds"] = dict(kwds.pop("metric_params")) # needs to be cloneable + if transformer is None: + # Use defaults from UMAP’s `nearest_neighbors` function + kwds.update( + n_jobs=settings.n_jobs, + n_trees=min(64, 5 + round(n_index**0.5 / 20.0)), + n_iters=max(5, round(np.log2(n_index))), + ) + return PyNNDescentTransformer(**kwds) + if isinstance(transformer, str): + msg = ( + f"Unknown transformer: {transformer}. " + f"Try passing a class or one of {get_literal_vals(_KnownTransformer)}" + ) + raise ValueError(msg) + return transformer # `transformer` is probably an instance + + +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, + ) def _has_self_column( @@ -141,3 +213,38 @@ def _ind_dist_shortcut( d.indices.reshape(n_obs, n_neighbors), d.data.reshape(n_obs, n_neighbors), ) + + +def _get_indices_distances_from_rect_matrix( + d: CSRBase, /, n_neighbors: int +) -> tuple[NDArray[np.int32 | np.int64], NDArray[np.float32 | np.float64]]: + """Get the `n_neighbors` nearest neighbors from a rectangular kNN distance matrix. + + In contrast to `_get_indices_distances_from_sparse_matrix`, + the columns of `d` index a subset of the observations the rows index, + so there is no self-column to take care of. + Rows are sorted by distance and truncated to `n_neighbors` entries. + """ + n_nonzero = np.diff(d.indptr) + if (n_too_few := int((n_nonzero < n_neighbors).sum())) > 0: + msg = ( + f"The transformer returned fewer than {n_neighbors} neighbors " + f"for {n_too_few} of {d.shape[0]} observations." + ) + raise ValueError(msg) + if is_constant(n_nonzero): + n_cols = int(n_nonzero[0]) + indices = d.indices.reshape(d.shape[0], n_cols) + distances = d.data.reshape(d.shape[0], n_cols) + else: # pad the rows to a common width, sorting the padding to the end + indices = np.zeros((d.shape[0], int(n_nonzero.max())), dtype=d.indices.dtype) + distances = np.full(indices.shape, np.inf, dtype=d.data.dtype) + rows = np.repeat(np.arange(d.shape[0]), n_nonzero) + cols = np.arange(d.nnz) - np.repeat(d.indptr[:-1], n_nonzero) + indices[rows, cols] = d.indices + distances[rows, cols] = d.data + order = np.argsort(distances, axis=1, kind="stable")[:, :n_neighbors] + return ( + np.take_along_axis(indices, order, axis=1), + np.take_along_axis(distances, order, axis=1), + ) diff --git a/src/scanpy/neighbors/_types.py b/src/scanpy/neighbors/_types.py index bfa0608c11..90e6eddb9a 100644 --- a/src/scanpy/neighbors/_types.py +++ b/src/scanpy/neighbors/_types.py @@ -1,17 +1,25 @@ from __future__ import annotations from collections.abc import Callable -from typing import TYPE_CHECKING, Literal, Protocol +from typing import TYPE_CHECKING, Literal, Protocol, TypedDict import numpy as np if TYPE_CHECKING: - from typing import Any, Self + from collections.abc import Mapping + from typing import Any, NotRequired, Self, TypeAlias from .._compat import CSRBase + from .._utils.random import _LegacyRandom + + # TODO: make `type` when https://github.com/sphinx-doc/sphinx/pull/13508 is released + RPForestDict: TypeAlias = Mapping[str, Mapping[str, np.ndarray]] # noqa: UP040 __all__ = [ "KnnTransformerLike", + "KwdsForTransformer", + "NeighborsDict", + "NeighborsParams", "_KnownTransformer", "_Method", "_Metric", @@ -62,3 +70,37 @@ def fit_transform(self, x, /, y: None = None) -> CSRBase: ... # from BaseEstimator def get_params(self, *, deep: bool = True) -> dict[str, Any]: ... def set_params(self, **params: Any) -> Self: ... + + +class KwdsForTransformer(TypedDict): + """Keyword arguments passed to a _KnownTransformer. + + IMPORTANT: when changing the parameters set here, + update the “*ignored*” part in the parameter docs! + """ + + n_neighbors: int + metric: _Metric | _MetricFn + metric_params: Mapping[str, Any] + random_state: _LegacyRandom + + +class NeighborsDict(TypedDict): + connectivities_key: str + distances_key: str + params: NeighborsParams + rp_forest: NotRequired[RPForestDict] + + +class NeighborsParams(TypedDict): + n_neighbors: int + method: _Method + random_state: _LegacyRandom + metric: _Metric | _MetricFn | None + metric_kwds: NotRequired[Mapping[str, Any]] + use_rep: NotRequired[str] + n_pcs: NotRequired[int] + # only set by `pp.bbknn` + batch_key: NotRequired[str] + neighbors_within_batch: NotRequired[int] + trim: NotRequired[int] diff --git a/src/scanpy/preprocessing/__init__.py b/src/scanpy/preprocessing/__init__.py index cb1aedb90c..d75adb8c8a 100644 --- a/src/scanpy/preprocessing/__init__.py +++ b/src/scanpy/preprocessing/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations from ..neighbors import neighbors +from ..neighbors._bbknn import bbknn from ._combat import combat from ._deprecated.sampling import subsample from ._harmony import harmony_integrate @@ -24,6 +25,7 @@ ) __all__ = [ + "bbknn", "calculate_qc_metrics", "combat", "downsample_counts", diff --git a/tests/test_bbknn.py b/tests/test_bbknn.py new file mode 100644 index 0000000000..223e10a312 --- /dev/null +++ b/tests/test_bbknn.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest +from anndata import AnnData +from scipy import sparse +from sklearn.neighbors import KNeighborsTransformer + +import scanpy as sc +from scanpy.neighbors._bbknn import ( + _compute_batch_balanced_knn, + _handle_transformer, + _trim, +) +from testing.scanpy._pytest.params import ARRAY_TYPES_MEM + +if TYPE_CHECKING: + from collections.abc import Callable + + from numpy.typing import NDArray + + from scanpy._compat import CSRBase + from scanpy.neighbors._types import _Metric + + +N_PER_BATCH = [60, 40, 30] +BATCHES = np.repeat(["a", "b", "c"], N_PER_BATCH) +N_OBS = len(BATCHES) + + +@pytest.fixture( + scope="module", params=[10, 2 * sc.settings.N_PCS], ids=["narrow", "wide"] +) +def rep(request: pytest.FixtureRequest) -> NDArray[np.float32]: + """Create a representation where each batch is shifted by a constant. + + Either narrow enough for `pp.bbknn` to use `.X`, + or wide enough that it uses the PCA. + """ + rng = np.random.default_rng(0) + x = rng.normal(size=(N_OBS, request.param)).astype(np.float32) + for i, batch in enumerate(np.unique(BATCHES)): + x[batch == BATCHES] += 5 * i + return x + + +@pytest.fixture +def adata(rep: NDArray[np.float32]) -> AnnData: + adata = AnnData(rep.copy(), obs=dict(batch=BATCHES.copy())) + # fewer PCs than `.X` has columns, so the two differ + sc.pp.pca(adata, n_comps=5, key_added="pca") + return adata + + +def n_per_row(a: CSRBase) -> NDArray[np.int64]: + return np.diff(a.indptr) + + +def test_bbknn(adata: AnnData) -> None: + assert sc.pp.bbknn(adata, 3, batch_key="batch") is None + + params = adata.uns["neighbors"]["params"] + assert params == dict( + n_neighbors=9, # 3 neighbors × 3 batches + method="umap", + metric="euclidean", + random_state=0, + batch_key="batch", + neighbors_within_batch=3, + trim=90, + ) + dists, conns = adata.obsp["distances"], adata.obsp["connectivities"] + assert dists.shape == conns.shape == (N_OBS, N_OBS) + # the cell itself is not part of its own neighborhood + assert (n_per_row(dists) == 8).all() + assert dists.diagonal().sum() == 0 + assert (conns != conns.T).nnz == 0 + + +def test_bbknn_representation(adata: AnnData) -> None: + dists = sc.pp.bbknn(adata, 3, batch_key="batch", copy=True).obsp["distances"] + # like `pp.neighbors`, we use the PCA – except for data narrower than `N_PCS` + used, unused = ("X", "pca") if adata.n_vars <= sc.settings.N_PCS else ("pca", "X") + + for rep in (used, unused): + sc.pp.bbknn(adata, 3, batch_key="batch", use_rep=rep, key_added=rep) + + np.testing.assert_allclose( + dists.toarray(), adata.obsp[f"{used}_distances"].toarray() + ) + assert (dists != adata.obsp[f"{unused}_distances"]).nnz > 0 + + +def test_bbknn_is_batch_balanced(adata: AnnData) -> None: + """Each cell has `neighbors_within_batch` neighbors in each batch, including itself.""" + sc.pp.bbknn(adata, 3, batch_key="batch") + + dists = adata.obsp["distances"].tolil() + for i, row in enumerate(dists.rows): + neighbors = np.asarray([*row, i]) # add back the cell itself + _, counts = np.unique(BATCHES[neighbors], return_counts=True) + assert (counts == 3).all() + + +def test_bbknn_connects_batches(adata: AnnData) -> None: + """Unlike `pp.neighbors`, `pp.bbknn` connects the (strongly separated) batches.""" + sc.pp.neighbors(adata, n_neighbors=9, key_added="knn") + sc.pp.bbknn(adata, 3, batch_key="batch", key_added="bbknn") + + def n_cross_batch(key: str) -> int: + i, j = adata.obsp[f"{key}_connectivities"].nonzero() + return int((BATCHES[i] != BATCHES[j]).sum()) + + assert n_cross_batch("knn") == 0 + assert n_cross_batch("bbknn") > 0 + + +@pytest.mark.parametrize( + "transformer", + [ + pytest.param(None, id="none"), + pytest.param("sklearn", id="sklearn"), + pytest.param("pynndescent", id="pynndescent"), + pytest.param( + KNeighborsTransformer(n_neighbors=3, algorithm="kd_tree"), id="instance" + ), + ], +) +def test_bbknn_transformer(adata: AnnData, transformer) -> None: + sc.pp.bbknn(adata, 3, batch_key="batch", transformer=transformer) + assert (n_per_row(adata.obsp["distances"]) == 8).all() + + +@pytest.mark.parametrize( + ("n_obs", "max_batch_size", "metric", "brute"), + [ + # brute force costs `n_obs` × index size, so both matter + pytest.param(20_000, 2_000, "euclidean", True, id="many_small_batches"), + pytest.param(100_000, 50_000, "euclidean", False, id="few_big_batches"), + pytest.param(300_000, 30_000, "euclidean", False, id="big_data"), + pytest.param(1_000, 500, "cosine", True, id="small_batch_other_metric"), + pytest.param(300_000, 30_000, "cosine", False, id="big_data_other_metric"), + ], +) +def test_bbknn_transformer_choice( + *, n_obs: int, max_batch_size: int, metric: _Metric, brute: bool +) -> None: + """`transformer=None` picks a backend based on how big the per-batch indices are.""" + from sklearn.neighbors import KNeighborsTransformer + + transformer, shortcut = _handle_transformer( + None, + n_obs=n_obs, + max_batch_size=max_batch_size, + n_neighbors=3, + metric=metric, + metric_kwds={}, + random_state=0, + ) + assert shortcut is brute + assert isinstance(transformer, KNeighborsTransformer) is brute + + +@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM) +def test_bbknn_array_types(rep: NDArray[np.float32], array_type: Callable) -> None: + adata = AnnData(array_type(np.abs(rep)), obs=dict(batch=BATCHES.copy())) + sc.pp.bbknn(adata, 3, 0, batch_key="batch") + assert (n_per_row(adata.obsp["distances"]) == 8).all() + + +@pytest.mark.parametrize("trim", [None, 0, 5, 12]) +def test_bbknn_trim(adata: AnnData, trim: int | None) -> None: + sc.pp.bbknn(adata, 3, batch_key="batch", trim=trim) + conns = adata.obsp["connectivities"] + + assert adata.uns["neighbors"]["params"]["trim"] == (90 if trim is None else trim) + if trim: + # ties are kept, so cells can end up with slightly more than `trim` neighbors + assert n_per_row(conns).max() >= trim + assert (conns != conns.T).nnz == 0 + # trimming only ever removes edges + sc.pp.bbknn(adata, 3, batch_key="batch", trim=0, key_added="untrimmed") + untrimmed = adata.obsp["untrimmed_connectivities"] + assert conns.nnz <= untrimmed.nnz + assert (conns != conns.multiply(untrimmed != 0)).nnz == 0 + + +def test_trim() -> None: + """`_trim` cuts each row at its `trim`-th largest value, but keeps the graph symmetric.""" + dense = [ + [0.0, 0.9, 0.8, 0.7], + [0.9, 0.0, 0.1, 0.0], + [0.8, 0.1, 0.0, 0.0], + [0.7, 0.0, 0.0, 0.0], + ] + conns = sparse.csr_matrix(dense) # noqa: TID251 + trimmed = _trim(conns.copy(), 2).toarray() + # row 0 keeps its top 2 (0.9, 0.8); 0.7 is dropped in both directions + np.testing.assert_allclose( + trimmed, + [ + [0.0, 0.9, 0.8, 0.0], + [0.9, 0.0, 0.1, 0.0], + [0.8, 0.1, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ], + ) + # rows with at most `trim` entries are left alone + np.testing.assert_allclose(_trim(conns.copy(), 4).toarray(), conns.toarray()) + + +def test_bbknn_knn_is_normalized(rep: NDArray[np.float32]) -> None: + """The merged neighbors are sorted, and each cell is its own first neighbor. + + kNN backends can report a tiny non-zero distance of a cell to itself, + which `umap` would mistake for the radius of that cell’s local neighborhood. + """ + knn_indices, knn_distances = _compute_batch_balanced_knn( + rep, + batches=BATCHES, + unique_batches=np.unique(BATCHES), + batch_sizes=np.asarray(N_PER_BATCH), + neighbors_within_batch=3, + transformer=None, + metric="euclidean", + metric_kwds={}, + random_state=0, + ) + np.testing.assert_array_equal(knn_indices[:, 0], np.arange(N_OBS)) + np.testing.assert_array_equal(knn_distances[:, 0], 0.0) + assert (np.diff(knn_distances, axis=1) >= 0).all() + + +def test_bbknn_duplicate_cells() -> None: + """Duplicate cells are at distance 0, but a cell is still its own first neighbor.""" + rng = np.random.default_rng(0) + x = rng.normal(size=(20, 5)) + x[10:] = x[:10] # each cell has an exact duplicate + batches = np.repeat(["a", "b"], 10) + + knn_indices, knn_distances = _compute_batch_balanced_knn( + x, + batches=batches, + unique_batches=np.unique(batches), + batch_sizes=np.asarray([10, 10]), + neighbors_within_batch=2, + transformer=None, + metric="euclidean", + metric_kwds={}, + random_state=0, + ) + np.testing.assert_array_equal(knn_indices[:, 0], np.arange(20)) + np.testing.assert_array_equal(knn_distances[:, 0], 0.0) + + adata = AnnData(x, obs=dict(batch=batches)) + sc.pp.bbknn(adata, 2, 0, batch_key="batch") + assert adata.obsp["distances"].diagonal().sum() == 0 + + +def test_bbknn_key_added(adata: AnnData) -> None: + sc.pp.bbknn(adata, 3, batch_key="batch") + sc.pp.bbknn(adata, 3, batch_key="batch", key_added="test") + + assert adata.uns["neighbors"]["params"] == adata.uns["test"]["params"] + assert adata.uns["test"]["distances_key"] == "test_distances" + assert adata.uns["test"]["connectivities_key"] == "test_connectivities" + for key in ("distances", "connectivities"): + np.testing.assert_allclose( + adata.obsp[key].toarray(), adata.obsp[f"test_{key}"].toarray() + ) + + +def test_bbknn_copy(adata: AnnData) -> None: + copied = sc.pp.bbknn(adata, 3, batch_key="batch", copy=True) + assert not adata.obsp + assert "neighbors" not in adata.uns # `.uns['pca']` is from the fixture + assert set(copied.obsp) == {"distances", "connectivities"} + + +@pytest.mark.parametrize( + ("kwargs", "error", "pattern"), + [ + pytest.param( + dict(batch_key="nope"), KeyError, r"Batch key 'nope' not found", id="key" + ), + pytest.param( + dict(neighbors_within_batch=0), + ValueError, + r"needs to be greater than 0", + id="n_neighbors", + ), + pytest.param( + dict(neighbors_within_batch=40), + ValueError, + r"Not all batches have at least .* \['c'\]", + id="batch_too_small", + ), + pytest.param( + dict(transformer="nope"), + ValueError, + r"Unknown transformer", + id="transformer", + ), + ], +) +def test_bbknn_errors( + adata: AnnData, kwargs: dict, error: type[Exception], pattern: str +) -> None: + with pytest.raises(error, match=pattern): + sc.pp.bbknn(adata, **{"batch_key": "batch", **kwargs})