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
@@ -93,6 +94,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

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