From 0ecc6ab29aa57101cd775d316c892cb8ed477ab8 Mon Sep 17 00:00:00 2001
From: Intron7 <sdicks@nvidia.com>
Date: Wed, 19 Aug 2026 15:28:31 +0200
Subject: [PATCH 1/4] add bbknn

Signed-off-by: Intron7 <sdicks@nvidia.com>
---
 docs/api/preprocessing.md            |   2 +
 docs/conf.py                         |   1 +
 src/scanpy/_utils/__init__.py        |   3 +-
 src/scanpy/external/pp/_bbknn.py     |   1 +
 src/scanpy/neighbors/_bbknn.py       | 383 +++++++++++++++++++++++++++
 src/scanpy/neighbors/_common.py      |  35 +++
 src/scanpy/neighbors/_types.py       |   4 +
 src/scanpy/preprocessing/__init__.py |   2 +
 tests/test_bbknn.py                  | 254 ++++++++++++++++++
 9 files changed, 684 insertions(+), 1 deletion(-)
 create mode 100644 src/scanpy/neighbors/_bbknn.py
 create mode 100644 tests/test_bbknn.py

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 <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
 
 ```
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/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/_bbknn.py b/src/scanpy/neighbors/_bbknn.py
new file mode 100644
index 0000000000..4157056413
--- /dev/null
+++ b/src/scanpy/neighbors/_bbknn.py
@@ -0,0 +1,383 @@
+"""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 .._settings import settings
+from .._utils import _doc_params, get_literal_vals
+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,
+)
+from ._connectivity import umap
+from ._doc import doc_n_pcs, doc_use_rep
+from ._types import _KnownTransformer
+
+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, _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,
+        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],
+    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],
+        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,
+    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``.
+
+    The heuristic for `transformer=None` matches the one of :func:`~scanpy.pp.neighbors`.
+    """
+    use_dense_distances = metric == "euclidean" and n_obs < 8192
+    shortcut = transformer == "sklearn" or (
+        transformer is None and (use_dense_distances or n_obs < 4096)
+    )
+    if shortcut:
+        from sklearn.neighbors import KNeighborsTransformer
+
+        return KNeighborsTransformer(
+            algorithm="brute",
+            n_jobs=settings.n_jobs,
+            n_neighbors=n_neighbors,
+            metric=metric,
+            metric_params=dict(metric_kwds),  # needs dict
+            # no random_state
+        ), True
+    if transformer is None or transformer == "pynndescent":
+        from pynndescent import PyNNDescentTransformer
+
+        kwds: dict[str, Any] = dict(
+            n_neighbors=n_neighbors,
+            metric=metric,
+            metric_kwds=dict(metric_kwds),  # needs to be cloneable
+            random_state=random_state,
+        )
+        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_obs**0.5 / 20.0)),
+                n_iters=max(5, round(np.log2(n_obs))),
+            )
+        return PyNNDescentTransformer(**kwds), False
+    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, False
+
+
+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..ad9183d806 100644
--- a/src/scanpy/neighbors/_common.py
+++ b/src/scanpy/neighbors/_common.py
@@ -216,3 +216,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..2663df021b 100644
--- a/src/scanpy/neighbors/_types.py
+++ b/src/scanpy/neighbors/_types.py
@@ -100,3 +100,7 @@ class NeighborsParams(TypedDict):
     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 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..ad040bd68b
--- /dev/null
+++ b/tests/test_bbknn.py
@@ -0,0 +1,254 @@
+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, _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
+
+
+N_PER_BATCH = [60, 40, 30]
+BATCHES = np.repeat(["a", "b", "c"], N_PER_BATCH)
+N_OBS = len(BATCHES)
+
+
+@pytest.fixture(scope="module")
+def rep() -> NDArray[np.float32]:
+    """Create a representation where each batch is shifted by a constant."""
+    rng = np.random.default_rng(0)
+    x = rng.normal(size=(N_OBS, 10)).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()))
+    adata.obsm["X_pca"] = rep.copy()
+    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_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("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),
+        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),
+        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 not adata.uns
+    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})

From 81fbb16116f4636e6e038aa0d175d9ead85a274b Mon Sep 17 00:00:00 2001
From: Intron7 <sdicks@nvidia.com>
Date: Wed, 19 Aug 2026 15:33:01 +0200
Subject: [PATCH 2/4] add release note

Signed-off-by: Intron7 <sdicks@nvidia.com>
---
 docs/release-notes/4306.feat.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 docs/release-notes/4306.feat.md

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`

From 6c996a8869571cc12dc19a1af9c343a743e54e02 Mon Sep 17 00:00:00 2001
From: Phil Schaf <flying-sheep@web.de>
Date: Thu, 20 Aug 2026 11:54:48 +0200
Subject: [PATCH 3/4] dedupe transformer creation

---
 benchmarks/benchmarks/preprocessing_log.py | 36 +++++++++
 hatch.toml                                 |  2 +-
 src/scanpy/neighbors/_bbknn.py             | 91 ++++++++++------------
 tests/test_bbknn.py                        | 76 +++++++++++++++---
 4 files changed, 143 insertions(+), 62 deletions(-)

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/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/neighbors/_bbknn.py b/src/scanpy/neighbors/_bbknn.py
index 4157056413..362d1001e4 100644
--- a/src/scanpy/neighbors/_bbknn.py
+++ b/src/scanpy/neighbors/_bbknn.py
@@ -12,21 +12,16 @@
 
 from .. import logging as logg
 from .._docs import doc_rng
-from .._settings import settings
-from .._utils import _doc_params, get_literal_vals
-from .._utils.random import (
-    _accepts_legacy_random_state,
-    _legacy_random_state,
-    _LegacyRng,
-)
+from .._utils import _doc_params
+from .._utils.random import _accepts_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
-from ._types import _KnownTransformer
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -36,8 +31,13 @@
     from numpy.typing import NDArray
 
     from .._compat import CSRBase
-    from .._utils.random import RNGLike, SeedLike, _LegacyRandom
-    from ._types import KnnTransformerLike, _Metric, _MetricFn
+    from .._utils.random import RNGLike, SeedLike
+    from ._types import (
+        KnnTransformerLike,
+        _KnownTransformer,
+        _Metric,
+        _MetricFn,
+    )
 
 
 @_doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep, rng=doc_rng)
@@ -168,6 +168,8 @@ def bbknn(  # noqa: PLR0913
     if batch_key not in adata.obs:
         msg = f"Batch key {batch_key!r} not found in `adata.obs`."
         raise KeyError(msg)
+    rng = np.random.default_rng(rng)
+
     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]):
@@ -182,11 +184,12 @@ def bbknn(  # noqa: PLR0913
         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),
+        rng=rng,
     )
     n_obs, n_neighbors = knn_indices.shape
     if trim is None:
@@ -238,11 +241,12 @@ def _compute_batch_balanced_knn(
     *,
     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,
+    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.
 
@@ -253,10 +257,11 @@ def _compute_batch_balanced_knn(
     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,
+        rng=rng,
     )
 
     n_obs = x.shape[0]
@@ -303,57 +308,39 @@ 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,
+    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``.
 
-    The heuristic for `transformer=None` matches the one of :func:`~scanpy.pp.neighbors`.
+    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.
     """
-    use_dense_distances = metric == "euclidean" and n_obs < 8192
     shortcut = transformer == "sklearn" or (
-        transformer is None and (use_dense_distances or n_obs < 4096)
-    )
-    if shortcut:
-        from sklearn.neighbors import KNeighborsTransformer
-
-        return KNeighborsTransformer(
-            algorithm="brute",
-            n_jobs=settings.n_jobs,
-            n_neighbors=n_neighbors,
-            metric=metric,
-            metric_params=dict(metric_kwds),  # needs dict
-            # no random_state
-        ), True
-    if transformer is None or transformer == "pynndescent":
-        from pynndescent import PyNNDescentTransformer
-
-        kwds: dict[str, Any] = dict(
-            n_neighbors=n_neighbors,
-            metric=metric,
-            metric_kwds=dict(metric_kwds),  # needs to be cloneable
-            random_state=random_state,
-        )
-        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_obs**0.5 / 20.0)),
-                n_iters=max(5, round(np.log2(n_obs))),
-            )
-        return PyNNDescentTransformer(**kwds), False
-    if isinstance(transformer, str):
-        msg = (
-            f"Unknown transformer: {transformer}. "
-            f"Try passing a class or one of {get_literal_vals(_KnownTransformer)}"
+        transformer is None
+        and (
+            max_batch_size < 4096
+            or (metric == "euclidean" and n_obs * max_batch_size < 10**9)
         )
-        raise ValueError(msg)
-    return transformer, False
+    )
+    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:
diff --git a/tests/test_bbknn.py b/tests/test_bbknn.py
index ad040bd68b..7c24c62a82 100644
--- a/tests/test_bbknn.py
+++ b/tests/test_bbknn.py
@@ -9,7 +9,11 @@
 from sklearn.neighbors import KNeighborsTransformer
 
 import scanpy as sc
-from scanpy.neighbors._bbknn import _compute_batch_balanced_knn, _trim
+from scanpy.neighbors._bbknn import (
+    _compute_batch_balanced_knn,
+    _handle_transformer,
+    _trim,
+)
 from testing.scanpy._pytest.params import ARRAY_TYPES_MEM
 
 if TYPE_CHECKING:
@@ -18,6 +22,7 @@
     from numpy.typing import NDArray
 
     from scanpy._compat import CSRBase
+    from scanpy.neighbors._types import _Metric
 
 
 N_PER_BATCH = [60, 40, 30]
@@ -25,11 +30,17 @@
 N_OBS = len(BATCHES)
 
 
-@pytest.fixture(scope="module")
-def rep() -> NDArray[np.float32]:
-    """Create a representation where each batch is shifted by a constant."""
+@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, 10)).astype(np.float32)
+    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
@@ -38,7 +49,8 @@ def rep() -> NDArray[np.float32]:
 @pytest.fixture
 def adata(rep: NDArray[np.float32]) -> AnnData:
     adata = AnnData(rep.copy(), obs=dict(batch=BATCHES.copy()))
-    adata.obsm["X_pca"] = rep.copy()
+    # fewer PCs than `.X` has columns, so the two differ
+    sc.pp.pca(adata, n_comps=5, key_added="pca")
     return adata
 
 
@@ -67,6 +79,20 @@ def test_bbknn(adata: AnnData) -> None:
     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")
@@ -107,6 +133,36 @@ def test_bbknn_transformer(adata: AnnData, transformer) -> None:
     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()))
@@ -165,11 +221,12 @@ def test_bbknn_knn_is_normalized(rep: NDArray[np.float32]) -> None:
         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,
+        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)
@@ -187,11 +244,12 @@ def test_bbknn_duplicate_cells() -> None:
         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,
+        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)
@@ -217,7 +275,7 @@ def test_bbknn_key_added(adata: AnnData) -> None:
 def test_bbknn_copy(adata: AnnData) -> None:
     copied = sc.pp.bbknn(adata, 3, batch_key="batch", copy=True)
     assert not adata.obsp
-    assert not adata.uns
+    assert "neighbors" not in adata.uns  # `.uns['pca']` is from the fixture
     assert set(copied.obsp) == {"distances", "connectivities"}
 
 

From b095f2cbc731caba486187b3440f86d5f3921eae Mon Sep 17 00:00:00 2001
From: "Philipp A." <flying-sheep@web.de>
Date: Fri, 21 Aug 2026 13:25:37 +0200
Subject: [PATCH 4/4] =?UTF-8?q?batch=5Fkey=E2=86=92batches?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/scanpy/external/pp/_bbknn.py |  3 +++
 src/scanpy/neighbors/_bbknn.py   | 33 ++++++++++++++++-------------
 src/scanpy/neighbors/_common.py  | 21 +++++++++++++++++--
 src/scanpy/neighbors/_types.py   | 14 +++++++------
 tests/test_bbknn.py              | 36 +++++++++++++++++---------------
 5 files changed, 68 insertions(+), 39 deletions(-)

diff --git a/src/scanpy/external/pp/_bbknn.py b/src/scanpy/external/pp/_bbknn.py
index 0c7f740e03..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,
diff --git a/src/scanpy/neighbors/_bbknn.py b/src/scanpy/neighbors/_bbknn.py
index 362d1001e4..fec094d945 100644
--- a/src/scanpy/neighbors/_bbknn.py
+++ b/src/scanpy/neighbors/_bbknn.py
@@ -13,10 +13,10 @@
 from .. import logging as logg
 from .._docs import doc_rng
 from .._utils import _doc_params
-from .._utils.random import _accepts_legacy_random_state, _LegacyRng
+from .._utils._doctests import doctest_needs
 from ._common import (
+    _get_bbknn_metadata,
     _get_indices_distances_from_rect_matrix,
-    _get_metadata,
     _get_sparse_matrix_from_indices_distances,
     _make_transformer,
 )
@@ -28,6 +28,7 @@
     from typing import Any
 
     from anndata import AnnData
+    from anndata.acc import AdRef
     from numpy.typing import NDArray
 
     from .._compat import CSRBase
@@ -40,14 +41,14 @@
     )
 
 
+@doctest_needs("anndata_acc")
 @_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",
+    batches: AdRef | str = "obs.batch",
     use_rep: str | None = None,
     transformer: KnnTransformerLike | _KnownTransformer | None = None,
     metric: _Metric | _MetricFn = "euclidean",
@@ -81,7 +82,7 @@ def bbknn(  # noqa: PLR0913
         which then serves as the basis for the construction of a symmetrical
         matrix of connectivities.
     {n_pcs}
-    batch_key
+    batches
         `adata.obs` column name discriminating between the batches.
     {use_rep}
     transformer
@@ -145,7 +146,7 @@ def bbknn(  # noqa: PLR0913
     >>> import scanpy as sc
     >>> adata = sc.datasets.pbmc68k_reduced()
     >>> adata.obs["batch"] = adata.obs["phase"]
-    >>> sc.pp.bbknn(adata, batch_key="batch")
+    >>> sc.pp.bbknn(adata, batches="obs.batch")
     >>> sc.tl.umap(adata)
 
     See Also
@@ -154,6 +155,8 @@ def bbknn(  # noqa: PLR0913
     :doc:`/how-to/knn-transformers`
 
     """
+    from anndata.acc import A, AdRef
+
     from ..tools._utils import _choose_representation
 
     start = logg.info("computing batch balanced neighbors")
@@ -162,16 +165,19 @@ def bbknn(  # noqa: PLR0913
     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 batch_key not in adata.obs:
-        msg = f"Batch key {batch_key!r} not found in `adata.obs`."
+    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)
 
-    batches = np.asarray(adata.obs[batch_key])
-    unique_batches, batch_sizes = np.unique(batches, return_counts=True)
+    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 = "
@@ -182,7 +188,7 @@ def bbknn(  # noqa: PLR0913
     x = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
     knn_indices, knn_distances = _compute_batch_balanced_knn(
         x,
-        batches=batches,
+        batches=batch_arr,
         unique_batches=unique_batches,
         batch_sizes=batch_sizes,
         neighbors_within_batch=neighbors_within_batch,
@@ -206,16 +212,15 @@ def bbknn(  # noqa: PLR0913
         connectivities = _trim(connectivities, trim)
     logg.debug("computed connectivities", time=start_connect)
 
-    key_added, neighbors_dict = _get_metadata(
+    key_added, neighbors_dict = _get_bbknn_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,
+        batches=A.to_json(batches),
         neighbors_within_batch=neighbors_within_batch,
         trim=trim,
     )
diff --git a/src/scanpy/neighbors/_common.py b/src/scanpy/neighbors/_common.py
index ad9183d806..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",
diff --git a/src/scanpy/neighbors/_types.py b/src/scanpy/neighbors/_types.py
index 2663df021b..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,7 +100,9 @@ class NeighborsParams(TypedDict):
     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]
+
+
+class BbknnParams(NeighborsParams):
+    batches: list[str | int | None]
+    neighbors_within_batch: int
+    trim: int
diff --git a/tests/test_bbknn.py b/tests/test_bbknn.py
index 7c24c62a82..e21d78e545 100644
--- a/tests/test_bbknn.py
+++ b/tests/test_bbknn.py
@@ -59,15 +59,14 @@ def n_per_row(a: CSRBase) -> NDArray[np.int64]:
 
 
 def test_bbknn(adata: AnnData) -> None:
-    assert sc.pp.bbknn(adata, 3, batch_key="batch") is 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",
-        random_state=0,
-        batch_key="batch",
+        batches=["obs", "batch"],
         neighbors_within_batch=3,
         trim=90,
     )
@@ -80,12 +79,12 @@ def test_bbknn(adata: AnnData) -> None:
 
 
 def test_bbknn_representation(adata: AnnData) -> None:
-    dists = sc.pp.bbknn(adata, 3, batch_key="batch", copy=True).obsp["distances"]
+    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, batch_key="batch", use_rep=rep, key_added=rep)
+        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()
@@ -95,7 +94,7 @@ def test_bbknn_representation(adata: AnnData) -> None:
 
 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")
+    sc.pp.bbknn(adata, 3, batches="obs.batch")
 
     dists = adata.obsp["distances"].tolil()
     for i, row in enumerate(dists.rows):
@@ -107,7 +106,7 @@ def test_bbknn_is_batch_balanced(adata: AnnData) -> None:
 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")
+    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()
@@ -129,7 +128,7 @@ def n_cross_batch(key: str) -> int:
     ],
 )
 def test_bbknn_transformer(adata: AnnData, transformer) -> None:
-    sc.pp.bbknn(adata, 3, batch_key="batch", transformer=transformer)
+    sc.pp.bbknn(adata, 3, batches="obs.batch", transformer=transformer)
     assert (n_per_row(adata.obsp["distances"]) == 8).all()
 
 
@@ -166,13 +165,13 @@ def test_bbknn_transformer_choice(
 @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")
+    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, batch_key="batch", trim=trim)
+    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)
@@ -181,7 +180,7 @@ def test_bbknn_trim(adata: AnnData, trim: int | None) -> None:
         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")
+    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
@@ -255,13 +254,13 @@ def test_bbknn_duplicate_cells() -> None:
     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")
+    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, batch_key="batch")
-    sc.pp.bbknn(adata, 3, batch_key="batch", key_added="test")
+    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"
@@ -273,7 +272,7 @@ def test_bbknn_key_added(adata: AnnData) -> None:
 
 
 def test_bbknn_copy(adata: AnnData) -> None:
-    copied = sc.pp.bbknn(adata, 3, batch_key="batch", copy=True)
+    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"}
@@ -283,7 +282,10 @@ def test_bbknn_copy(adata: AnnData) -> None:
     ("kwargs", "error", "pattern"),
     [
         pytest.param(
-            dict(batch_key="nope"), KeyError, r"Batch key 'nope' not found", id="key"
+            dict(batches="obs.nope"),
+            KeyError,
+            r"Batch key A.obs\['nope'\] not found",
+            id="key",
         ),
         pytest.param(
             dict(neighbors_within_batch=0),
@@ -309,4 +311,4 @@ 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})
+        sc.pp.bbknn(adata, **{"batches": "obs.batch", **kwargs})