perf: numba kernels for pp.harmony_integrate (#4424)

Author: timtreis

docs/references.bib

@@ -39,6 +39,14 @@ @article{Angerer2015
   pages     = {1241--1243},
 }
 
+@inproceedings{Bachem2016,
+  author    = {Bachem, Olivier and Lucic, Mario and Hassani, S. Hamed and Krause, Andreas},
+  title     = {Fast and Provably Good Seedings for k-Means},
+  booktitle = {Advances in Neural Information Processing Systems},
+  volume    = {29},
+  year      = {2016},
+}
+
 @article{Baron2016,
   author    = {Baron, Maayan and Veres, Adrian and Wolock, Samuel L. and Faust, Aubrey L. and Gaujoux, Renaud and Vetere, Amedeo and Ryu, Jennifer Hyoje and Wagner, Bridget K. and Shen-Orr, Shai S. and Klein, Allon M. and Melton, Douglas A. and Yanai, Itai},
   title     = {A Single-Cell Transcriptomic Map of the Human and Mouse Pancreas Reveals Inter- and Intra-cell Population Structure},

docs/release-notes/4424.perf.md

@@ -0,0 +1,3 @@
+Speed up {func}`~scanpy.pp.harmony_integrate` 6–13× with parallel numba kernels, and add the `seeding`, `stop_rule`, `penalty_combine`, `seeding_lloyd_steps`, `seeding_cells_per_cluster`, `recompute_counts` and `low_memory` parameters.
+For `flavor="harmony2"`, the default seeding is now `"afkmc2"`, and the default stop rule for both flavors is `"reference_record"`; pass `seeding="kmeans", stop_rule="reference"` for the previous behaviour.
+Cells with an all-zero embedding are now left out and returned unchanged {smaller}`T Treis`

src/scanpy/preprocessing/_harmony/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-import enum
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -24,20 +23,6 @@
     from ...get._types import MultiAcc
 
 
-class _Stopping(NamedTuple):
-    max_iter_clustering: int
-    tol_clustering: float
-    tol_harmony: float
-
-
-class Stopping(_Stopping, enum.Enum):
-    HARMONY2 = (4, 1e-3, 1e-2)
-    """Follows harmonypy 2.0.0 / harmony 2.0.5 (R) defaults."""
-
-    HARMONY1 = (200, 1e-5, 1e-4)
-    """Follows harmony-pytorch defaults."""
-
-
 @_doc_params(
     use=doc_use(
         "Which embedding to adjust. By default, the one :func:`~scanpy.pp.pca` wrote.",
@@ -70,13 +55,22 @@ def harmony_integrate(  # noqa: PLR0913
     batch_prune_threshold: float | None = 1e-5,
     correction_method: Literal["fast"] = "fast",
     block_proportion: float = 0.05,
+    stop_rule: Literal[
+        "reference_record", "reference", "paper", "paper_window3"
+    ] = "reference_record",
+    seeding: Literal["afkmc2", "kmeans", "kmeans++", "paper_nearest"] | None = None,
+    seeding_lloyd_steps: int = 0,
+    seeding_cells_per_cluster: int | None = None,
+    penalty_combine: Literal["product", "sum"] = "product",
+    recompute_counts: bool = True,
+    low_memory: bool = False,
     rng: SeedLike | RNGLike | None = None,
 ) -> np.ndarray | None:
     r"""Integrate different experiments using the Harmony algorithm :cite:p:`Korsunsky2019,Patikas2026`.
 
     Integration adjusts a PCA embedding to reduce experiment-specific technical variation,
     allowing cells from different experiments to be compared.
-    This CPU implementation is based on the rapids-singlecell and harmony-pytorch packages.
+    This CPU implementation runs in parallel numba kernels.
     Multiple batch variables follow the per-covariate formulation described in the Harmony papers:
     each batch variable is modeled separately instead of combining them all into one joint category.
     As Harmony works by adjusting the principal components,
@@ -105,7 +99,7 @@ def harmony_integrate(  # noqa: PLR0913
     {use}
     {out}
     dtype
-        The data type to use for Harmony computation.
+        The data type to use for Harmony computation, ``float32`` or ``float64``.
         If you use 32-bit you may experience numerical instability.
     flavor
         Which version of the Harmony algorithm to use.
@@ -123,15 +117,18 @@ def harmony_integrate(  # noqa: PLR0913
     max_iter_clustering
         Maximum iterations for the clustering step within each Harmony iteration.
         If :data:`None`, uses the value of the reference implementation for the
-        chosen ``flavor``: ``4`` for ``"harmony2"``, ``200`` for ``"harmony1"``.
+        chosen ``flavor``: ``4`` for ``"harmony2"``, ``200`` for ``"harmony1"``
+        (``200`` for both with the ``"paper"`` stop rules).
     tol_harmony
         Convergence tolerance for the Harmony objective function.
         The algorithm stops when the relative change in objective falls below this value.
-        If :data:`None`, ``1e-2`` for ``"harmony2"`` and ``1e-4`` for ``"harmony1"``.
+        If :data:`None`, ``1e-2`` for ``"harmony2"`` and ``1e-4`` for ``"harmony1"``
+        (``1e-4`` for both with the ``"paper"`` stop rules).
     tol_clustering
         Convergence tolerance for the clustering step within each
         Harmony iteration.
-        If :data:`None`, ``1e-3`` for ``"harmony2"`` and ``1e-5`` for ``"harmony1"``.
+        If :data:`None`, ``1e-3`` for ``"harmony2"`` and ``1e-5`` for ``"harmony1"``
+        (``1e-5`` for both with the ``"paper"`` stop rules).
     sigma
         Width of the soft-clustering kernel.
         Controls the entropy of cluster assignments:
@@ -173,14 +170,52 @@ def harmony_integrate(  # noqa: PLR0913
         Set to ``None`` to disable pruning.
     correction_method
         Method for the correction step.
-        ``"fast"`` uses a precomputed factorization that avoids the full inversion,
-        which can be faster for datasets with many batches.
-        Multiple keys automatically use the exact general-design solve,
-        because this optimization only applies to a single batch variable.
+        Only ``"fast"`` is supported: the ridge regression of every cluster is solved exactly,
+        eliminating the batch variable with the most levels in closed form,
+        for one or several batch keys.
     block_proportion
         Proportion of cells updated per clustering sub-iteration.
         Smaller values produce more stochastic updates.
         Larger values are faster but may converge to different solutions.
+    stop_rule
+        When the clustering and the outer Harmony loop stop.
+        ``"reference_record"`` (default) and ``"reference"`` use the iteration limits and tolerances
+        listed under ``max_iter_clustering``, ``tol_clustering`` and ``tol_harmony``,
+        and test clustering convergence over a window of 3 iterations.
+        ``"reference_record"`` records every round's objective for the outer test;
+        ``"reference"`` records it only when the clustering converged.
+        ``"paper"`` and ``"paper_window3"`` use the tolerances of :cite:t:`Korsunsky2019`
+        (200 clustering iterations, ``1e-5`` and ``1e-4``) for both flavors,
+        testing clustering convergence on the last iteration or a window of 3.
+    seeding
+        How the initial cluster centroids are chosen.
+        If :data:`None`, ``"afkmc2"`` for ``"harmony2"`` and ``"kmeans"`` for ``"harmony1"``.
+        ``"afkmc2"`` approximates k-means++ seeding with Markov chains :cite:p:`Bachem2016`
+        and needs a single pass over the cells.
+        ``"kmeans"`` runs :class:`~sklearn.cluster.KMeans` (k-means++ and up to 25 Lloyd iterations).
+        ``"kmeans++"`` is exact k-means++ seeding.
+        ``"paper_nearest"`` follows the sampler of :cite:t:`Patikas2026`,
+        taking the cell with the *highest* ``log(u) / D`` (the paper prints "lowest"),
+        where ``D`` is the distance to the nearest centroid chosen so far:
+        this draws cells with probability proportional to that distance.
+    seeding_lloyd_steps
+        Hard spherical k-means iterations applied to the seeds before Harmony starts.
+    seeding_cells_per_cluster
+        If set, ``seeding="kmeans"`` and the ``seeding_lloyd_steps`` run on a random subsample of
+        ``seeding_cells_per_cluster * n_clusters`` cells.
+    penalty_combine
+        How the diversity penalties of several batch keys are combined in the soft assignment.
+        ``"product"`` (default) multiplies the per-key penalty factors, as in :cite:t:`Korsunsky2019`;
+        ``"sum"`` adds them, as the assignment rule is printed in :cite:t:`Patikas2026`,
+        leaving out levels with ``theta == 0``.
+        Identical for a single key.
+    recompute_counts
+        Rebuild the per-batch cluster counts from the assignments after every clustering iteration.
+        If ``False``, the counts are only updated incrementally, which accumulates rounding error
+        in float32.
+    low_memory
+        Drop one cell x component buffer during the iterations
+        at the cost of one extra pass at the end; the result is identical.
     rng
         Random number generator or seed for deterministic behavior.
 
@@ -190,26 +225,11 @@ def harmony_integrate(  # noqa: PLR0913
     It contains the principal components adjusted by Harmony,
     such that different experiments are integrated.
     """
-    from .core import Harmony
+    from .core import harmonize
 
-    # Resolve flavor into internal flags
-    if flavor not in {"harmony1", "harmony2"}:
-        msg = f"flavor must be 'harmony1' or 'harmony2', got {flavor!r}."
-        raise ValueError(msg)
     if correction_method != "fast":
         msg = f"correction_method must be 'fast', got {correction_method!r}."
         raise ValueError(msg)
-    stabilized_penalty = flavor == "harmony2"
-    dynamic_lambda = flavor == "harmony2"
-
-    # Unset stopping rules follow the flavor's reference implementation
-    defaults = Stopping[flavor.upper()]
-    if max_iter_clustering is None:
-        max_iter_clustering = defaults.max_iter_clustering
-    if tol_clustering is None:
-        tol_clustering = defaults.tol_clustering
-    if tol_harmony is None:
-        tol_harmony = defaults.tol_harmony
 
     # Warn when flavor-incompatible parameters are explicitly set
     if flavor == "harmony2" and ridge_lambda != 1.0:
@@ -244,14 +264,6 @@ def harmony_integrate(  # noqa: PLR0913
         msg = f"Could not convert input of type {type(input_data).__name__} to NumPy array."
         raise TypeError(msg) from e
 
-    # Check for NaN values
-    if np.isnan(x).any():
-        msg = (
-            "Input data contains NaN values. Please handle these before "
-            "running harmony_integrate."
-        )
-        raise ValueError(msg)
-
     # Run Harmony
     batch_df = _obs_df(
         adata,
@@ -262,26 +274,32 @@ def harmony_integrate(  # noqa: PLR0913
             legacy_name="",
         ),
     )
-    harmony = Harmony(
+    harmony_out, _ = harmonize(
+        x,
         batch_df,
         list(batch_df.columns),
-        theta=theta,
-        sigma=sigma,
+        flavor=flavor,
         n_clusters=n_clusters,
+        sigma=sigma,
+        theta=theta,
+        tau=tau,
+        ridge_lambda=ridge_lambda,
+        alpha=alpha,
+        batch_prune_threshold=batch_prune_threshold,
+        block_proportion=block_proportion,
         max_iter_harmony=max_iter_harmony,
         max_iter_clustering=max_iter_clustering,
         tol_harmony=tol_harmony,
         tol_clustering=tol_clustering,
-        ridge_lambda=ridge_lambda,
-        block_proportion=block_proportion,
-        tau=tau,
+        stop_rule=stop_rule,
+        seeding=seeding,
+        penalty_combine=penalty_combine,
+        seeding_lloyd_steps=seeding_lloyd_steps,
+        seeding_cells_per_cluster=seeding_cells_per_cluster,
+        recompute_counts=recompute_counts,
+        low_memory=low_memory,
         rng=rng,
-        stabilized_penalty=stabilized_penalty,
-        dynamic_lambda=dynamic_lambda,
-        alpha=alpha,
-        batch_prune_threshold=batch_prune_threshold,
     )
-    harmony_out = harmony.fit(x)
 
     if dest is True:
         return harmony_out