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 774e56062b..1ac8181ba2 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 @@ -93,6 +94,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 fb76d2c373..5448ef2e61 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..8832d28ffa 100644 --- a/src/scanpy/external/pp/_bbknn.py +++ b/src/scanpy/external/pp/_bbknn.py @@ -2,6 +2,8 @@ from typing import TYPE_CHECKING +from scverse_misc import Deprecation, deprecated + from ..._utils._doctests import doctest_needs if TYPE_CHECKING: @@ -11,6 +13,7 @@ from sklearn.metrics import DistanceMetric +@deprecated(Deprecation("1.13.0", "Use :func:`scanpy.pp.bbknn` instead.")) @doctest_needs("bbknn") def bbknn( # noqa: PLR0913 adata: AnnData, @@ -45,6 +48,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/_bbknn.py b/src/scanpy/neighbors/_bbknn.py new file mode 100644 index 0000000000..fec094d945 --- /dev/null +++ b/src/scanpy/neighbors/_bbknn.py @@ -0,0 +1,375 @@ +"""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._doctests import doctest_needs +from ._common import ( + _get_bbknn_metadata, + _get_indices_distances_from_rect_matrix, + _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 anndata.acc import AdRef + from numpy.typing import NDArray + + from .._compat import CSRBase + from .._utils.random import RNGLike, SeedLike + from ._types import ( + KnnTransformerLike, + _KnownTransformer, + _Metric, + _MetricFn, + ) + + +@doctest_needs("anndata_acc") +@_doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep, rng=doc_rng) +def bbknn( # noqa: PLR0913 + adata: AnnData, + neighbors_within_batch: int = 3, + n_pcs: int | None = None, + *, + batches: AdRef | str = "obs.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} + batches + `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, batches="obs.batch") + >>> sc.tl.umap(adata) + + See Also + -------- + :func:`~scanpy.pp.neighbors` + :doc:`/how-to/knn-transformers` + + """ + from anndata.acc import A, AdRef + + 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 not isinstance(batches, AdRef): + batches = A.resolve(batches, vec=True) + + if neighbors_within_batch < 1: + msg = "`neighbors_within_batch` needs to be greater than 0." + raise ValueError(msg) + if batches not in adata: + msg = f"Batch key {batches!r} not found in `adata.obs`." + raise KeyError(msg) + rng = np.random.default_rng(rng) + + batch_arr = np.asarray(adata[batches]) + unique_batches, batch_sizes = np.unique(batch_arr, 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=batch_arr, + unique_batches=unique_batches, + batch_sizes=batch_sizes, + neighbors_within_batch=neighbors_within_batch, + transformer=transformer, + metric=metric, + metric_kwds=metric_kwds, + rng=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_bbknn_metadata( + key_added, + n_neighbors=n_neighbors, + method="umap", + 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)), + batches=A.to_json(batches), + 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], + rng: np.random.Generator, +) -> 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, + rng=rng, + ) + + 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], + rng: np.random.Generator, +) -> 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, + rng=rng, + ), 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 1893d74f59..1c2e8a5175 100644 --- a/src/scanpy/neighbors/_common.py +++ b/src/scanpy/neighbors/_common.py @@ -18,7 +18,12 @@ from numpy.typing import NDArray from .._compat import CSRBase - from ._types import KnnTransformerLike, KwdsForTransformer, NeighborsParams + from ._types import ( + BbknnParams, + KnnTransformerLike, + KwdsForTransformer, + NeighborsParams, + ) def _make_transformer( @@ -75,7 +80,19 @@ def _make_transformer( def _get_metadata( key_added: str | None, /, **params: Unpack[NeighborsParams] -) -> tuple[str, NeighborsDict]: +) -> tuple[str, NeighborsDict[NeighborsParams]]: + return _metadata(key_added, params) + + +def _get_bbknn_metadata( + key_added: str | None, /, **params: Unpack[BbknnParams] +) -> tuple[str, NeighborsDict[BbknnParams]]: + return _metadata(key_added, params) + + +def _metadata[P: NeighborsParams]( + key_added: str | None, params: P +) -> tuple[str, NeighborsDict[P]]: if key_added is None: return "neighbors", NeighborsDict( connectivities_key="connectivities", @@ -216,3 +233,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 e26a1153a0..a6dbc17505 100644 --- a/src/scanpy/neighbors/_types.py +++ b/src/scanpy/neighbors/_types.py @@ -85,10 +85,10 @@ class KwdsForTransformer(TypedDict): rng: NotRequired[np.random.Generator] -class NeighborsDict(TypedDict): +class NeighborsDict[P: NeighborsParams](TypedDict): connectivities_key: str distances_key: str - params: NeighborsParams + params: P rp_forest: NotRequired[RPForestDict] @@ -100,3 +100,9 @@ class NeighborsParams(TypedDict): metric_kwds: NotRequired[Mapping[str, Any]] use_rep: NotRequired[str] n_pcs: NotRequired[int] + + +class BbknnParams(NeighborsParams): + batches: list[str | int | None] + neighbors_within_batch: int + trim: int diff --git a/src/scanpy/preprocessing/__init__.py b/src/scanpy/preprocessing/__init__.py index 3e1480c3df..10641e730d 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 @@ -25,6 +26,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..e21d78e545 --- /dev/null +++ b/tests/test_bbknn.py @@ -0,0 +1,314 @@ +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, batches="obs.batch") is None + + params = adata.uns["neighbors"]["params"] + assert params == dict( + n_neighbors=9, # 3 neighbors × 3 batches + method="umap", + metric="euclidean", + batches=["obs", "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, batches="obs.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, batches="obs.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, batches="obs.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, batches="obs.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, batches="obs.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={}, + rng=np.random.default_rng(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, batches="obs.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, batches="obs.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, batches="obs.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={}, + rng=np.random.default_rng(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={}, + rng=np.random.default_rng(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, batches="obs.batch") + assert adata.obsp["distances"].diagonal().sum() == 0 + + +def test_bbknn_key_added(adata: AnnData) -> None: + sc.pp.bbknn(adata, 3, batches="obs.batch") + sc.pp.bbknn(adata, 3, batches="obs.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, batches="obs.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(batches="obs.nope"), + KeyError, + r"Batch key A.obs\['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, **{"batches": "obs.batch", **kwargs})