feat: add bbknn

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")

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 <data-integration>` and external {ref}`external data integration <external-data-integration>`.
 
 ## Doublet detection
@@ -83,6 +84,7 @@ Also see {ref}`data integration tools <data-integration>` and external {ref}`ext
    :nosignatures:
    :toctree: generated/
 
+   pp.bbknn
    pp.neighbors
 
 ```

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]"], []),

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`

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]

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

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
     ------

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: