From 7992a3fe7d48ca0b11b189ba9f6610b61fed0703 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 01:30:21 +0200 Subject: [PATCH 01/10] perf: numba kernels for harmony_integrate Replace the NumPy hot loops with parallel numba kernels that follow the input dtype. One fused kernel per clustering iteration, one exact correction path for any number of batch keys (block elimination of the largest covariate), fixed chunking so results do not depend on the thread count. Add seeding, stop_rule, penalty_combine, seeding_lloyd_steps, seeding_cells_per_cluster, recompute_counts and low_memory. Defaults are now seeding='afkmc2' and stop_rule='reference_record'; seeding='kmeans', stop_rule='reference' reproduces the previous output. --- docs/references.bib | 8 + src/scanpy/preprocessing/_harmony/__init__.py | 125 +- src/scanpy/preprocessing/_harmony/_kernels.py | 329 ++++++ src/scanpy/preprocessing/_harmony/core.py | 1007 ++++++++--------- tests/test_harmony.py | 685 ++++++++++- 5 files changed, 1560 insertions(+), 594 deletions(-) create mode 100644 src/scanpy/preprocessing/_harmony/_kernels.py diff --git a/docs/references.bib b/docs/references.bib index aa2261fe7b..ae203b5ea5 100644 --- a/docs/references.bib +++ b/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}, diff --git a/src/scanpy/preprocessing/_harmony/__init__.py b/src/scanpy/preprocessing/_harmony/__init__.py index 58e09f3939..de8a149530 100644 --- a/src/scanpy/preprocessing/_harmony/__init__.py +++ b/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 @@ -17,20 +16,6 @@ from ..._utils.random import RNGLike, SeedLike -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.""" - - def harmony_integrate( # noqa: PLR0913 adata: AnnData, key: str | Sequence[str], @@ -52,13 +37,23 @@ 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"] = "afkmc2", + 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, ) -> None: """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 + and follows the rapids-singlecell and harmony-pytorch packages. Multiple batch variables follow the per-covariate formulation described in the Harmony papers: each key is modeled separately instead of combining all keys into one joint category. As Harmony works by adjusting the principal components, @@ -105,15 +100,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: @@ -155,14 +153,48 @@ 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, as rapids-singlecell does; + ``"reference"`` records it only when the clustering converged, as harmonypy does. + ``"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. + ``"afkmc2"`` (default) 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), + as harmonypy does. + ``"kmeans++"`` is exact k-means++ seeding. + ``"paper_nearest"`` is the sampler of :cite:t:`Patikas2026`. + 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`. + 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. @@ -172,7 +204,7 @@ def harmony_integrate( # noqa: PLR0913 containing 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"}: @@ -181,17 +213,6 @@ def harmony_integrate( # noqa: PLR0913 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: @@ -238,27 +259,29 @@ def harmony_integrate( # noqa: PLR0913 ) raise ValueError(msg) - # Run Harmony - harmony = Harmony( + adata.obsm[adjusted_basis], _ = harmonize( + x, adata.obs, key, - 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) - - # Store result - adata.obsm[adjusted_basis] = harmony_out diff --git a/src/scanpy/preprocessing/_harmony/_kernels.py b/src/scanpy/preprocessing/_harmony/_kernels.py new file mode 100644 index 0000000000..9d7355cbc2 --- /dev/null +++ b/src/scanpy/preprocessing/_harmony/_kernels.py @@ -0,0 +1,329 @@ +"""Fused numba kernels for Harmony. Float arrays share one dtype (float32 or float64), C-contiguous. + +Layout: ``z_norm``/``x`` cells x d, ``r`` cells x K, ``o`` levels x K, +``y_t`` d x K (centroids transposed, so every inner loop runs over K and vectorises +without float reassociation), ``codes`` cells x C int32 with per-covariate offsets. +``n_chunks`` is the number of static cell chunks per parallel sweep; the driver passes a fixed value, +so results do not depend on the number of numba threads. They differ at rounding level between values of ``n_chunks``. +""" + +from __future__ import annotations + +import numpy as np +from fast_array_utils.numba import njit +from numba import prange + + +@njit +def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 + z_norm, + y_t, + r, + codes, + o, + r_sum, + pr_b, + theta, + perm, + bounds, + term, + stabilized, + n_chunks, + scatter, + sum_penalties, +): + """One clustering iteration: block-wise soft assignment, objective parts, next centroids. + + Mutates ``r``, ``o`` (levels x K) and ``r_sum`` (K,) in place. + Returns ``(dist, ent, o_fresh, y_t_next)``: ``dist = sum r*2(1-z.y)``, ``ent = sum r*log(r)`` + (float64), ``o`` rebuilt from the new ``r``, and the normalised centroids of the new ``r``. + + ``scatter`` selects how a block's counts are moved: ``False`` sums them in per-chunk accumulators in the + parallel sweeps (work per block ~ n_chunks * levels * K, best for many cells per level); ``True`` moves them + with serial loops over the block's cells and fills the penalty table only for levels that occur in the + block (work per block ~ cells * K, best for thousands of small levels). Same result up to rounding. + + ``sum_penalties`` selects how the penalty factors of a cell's several batch keys are combined: ``False`` + multiplies them (scanpy, rapids-singlecell, [M1 p.9]); ``True`` adds them, as the assignment rule is printed + in [V2 p.10]. With one key the two are the same. + """ + ft = z_norm.dtype.type + n_pc = z_norm.shape[1] + n_levels, n_clusters = o.shape + n_cov = codes.shape[1] + one = ft(1.0) + inv_cov = one / ft(n_cov) + combine_sum = sum_penalties and n_cov > 1 + tiny = ft(1e-38) + logpen = np.empty((n_levels, n_clusters), dtype=z_norm.dtype) + acc = np.empty( + (1 if scatter else n_chunks, n_levels, n_clusters), dtype=z_norm.dtype + ) + present = np.ones(n_levels, dtype=np.bool) # levels whose penalty row is needed + y_acc = np.zeros((n_chunks, n_pc, n_clusters), dtype=z_norm.dtype) + o_fresh = np.zeros((n_levels, n_clusters), dtype=z_norm.dtype) + dist = np.zeros(n_chunks, dtype=np.float64) + ent = np.zeros(n_chunks, dtype=np.float64) + for blk in range(bounds.size - 1): + lo = bounds[blk] + size = bounds[blk + 1] - lo + # 1. hold the block out of the counts: o -= sum of the block's old r, same for r_sum + if scatter: + present[:] = False + for j in range(lo, lo + size): + i = perm[j] + for c in range(n_cov): + b = codes[i, c] + present[b] = True + for k in range(n_clusters): + o[b, k] -= r[i, k] + for k in range(n_clusters): + r_sum[k] -= r[i, k] + else: + acc[:] = 0 + for t in prange(n_chunks): + for j in range( + lo + t * size // n_chunks, lo + (t + 1) * size // n_chunks + ): + i = perm[j] + for c in range(n_cov): + b = codes[i, c] + for k in range(n_clusters): + acc[t, b, k] += r[i, k] + for t in range(n_chunks): + for b in range(n_levels): + for k in range(n_clusters): + o[b, k] -= acc[t, b, k] + r_sum[k] -= acc[t, b, k] * inv_cov + # 2. log penalty table on held-out counts, e[b, k] = pr_b[b] * r_sum[k] + for b in range(n_levels): + if present[b]: + for k in range(n_clusters): + e = pr_b[b] * r_sum[k] + den = o[b, k] + e + one if stabilized else o[b, k] + one + logpen[b, k] = theta[b] * (np.log(e + one) - np.log(den)) + if combine_sum: # the table then holds the factors; cap the exponent so a huge e cannot overflow to inf + logpen[b, k] = np.exp(min(logpen[b, k], ft(80.0))) + # 3. new assignments; cells of a block are independent given the table + if not scatter: + acc[:] = 0 + for t in prange(n_chunks): + dots = np.empty(n_clusters, dtype=z_norm.dtype) + logit = np.empty(n_clusters, dtype=z_norm.dtype) + factor = np.empty(n_clusters, dtype=z_norm.dtype) + y_loc = y_acc[t] + dist_t = 0.0 + ent_t = 0.0 + for j in range(lo + t * size // n_chunks, lo + (t + 1) * size // n_chunks): + i = perm[j] + for k in range(n_clusters): + dots[k] = 0 + for dd in range(n_pc): + zv = z_norm[i, dd] + for k in range(n_clusters): + dots[k] += zv * y_t[dd, k] + for k in range(n_clusters): + logit[k] = term * (one - dots[k]) + if combine_sum: + for k in range(n_clusters): + factor[k] = 0 + for c in range(n_cov): + b = codes[i, c] + for k in range(n_clusters): + factor[k] += logpen[b, k] + for k in range(n_clusters): + logit[k] += np.log(max(factor[k], tiny)) + else: + for c in range(n_cov): + b = codes[i, c] + for k in range(n_clusters): + logit[k] += logpen[b, k] + top = logit[0] + for k in range(1, n_clusters): + top = max(top, logit[k]) + total = ft(0.0) + for k in range(n_clusters): + logit[k] -= top + p = np.exp(logit[k]) + r[i, k] = p + total += p + inv = one / total + cell_dist = ft(0.0) + cell_ent = ft(0.0) + for k in range(n_clusters): + rk = r[i, k] * inv + r[i, k] = rk + cell_dist += rk * (one - dots[k]) + cell_ent += rk * logit[k] # log r_k = logit_k - log(total) + dist_t += 2.0 * cell_dist + ent_t += cell_ent - np.log(total) + if not scatter: + for c in range(n_cov): + b = codes[i, c] + for k in range(n_clusters): + acc[t, b, k] += r[i, k] + for dd in range(n_pc): + zv = z_norm[i, dd] + for k in range(n_clusters): + y_loc[dd, k] += zv * r[i, k] + dist[t] += dist_t + ent[t] += ent_t + # 4. put the block back + if scatter: + for j in range(lo, lo + size): + i = perm[j] + for c in range(n_cov): + b = codes[i, c] + for k in range(n_clusters): + o[b, k] += r[i, k] + o_fresh[b, k] += r[i, k] + for k in range(n_clusters): + r_sum[k] += r[i, k] + else: + for t in range(n_chunks): + for b in range(n_levels): + for k in range(n_clusters): + o[b, k] += acc[t, b, k] + o_fresh[b, k] += acc[t, b, k] + r_sum[k] += acc[t, b, k] * inv_cov + # 5. centroids of the new r: y_k = normalise(sum_i r_ik z_i) + y_t_next = np.empty((n_pc, n_clusters), dtype=z_norm.dtype) + for k in range(n_clusters): + sq = ft(0.0) + for dd in range(n_pc): + v = ft(0.0) + for t in range(n_chunks): + v += y_acc[t, dd, k] + y_t_next[dd, k] = v + sq += v * v + scale = one / max(np.sqrt(sq), ft(1e-12)) + for dd in range(n_pc): + y_t_next[dd, k] *= scale + return dist.sum(), ent.sum(), o_fresh, y_t_next + + +@njit +def _joint_stats(x, r, joint, order, n_joint, n_chunks): # noqa: PLR0912, PLR0917 + """Sufficient statistics of the per-cluster regression, one row per joint batch category. + + ``joint[i]`` is the index of cell i's combination of batch levels (with one key: its batch) and ``order`` + lists the cells sorted by it. Returns ``s[j, k, :] = sum of r_ik * x_i`` and ``o[j, k] = sum of r_ik`` over + the cells of category j. + + The sorted cells are cut into ``n_chunks`` equal pieces. A category that lies strictly inside a piece is + written by that piece alone; only the first and last category of a piece can continue in a neighbour, and + those two go through a small per-piece buffer that is merged serially. Scratch is 2 * n_chunks * K * d floats + whatever the number of categories. + """ + n_cells, n_pc = x.shape + n_clusters = r.shape[1] + s = np.zeros((n_joint, n_clusters, n_pc), dtype=x.dtype) + o = np.zeros((n_joint, n_clusters), dtype=x.dtype) + s_edge = np.zeros((n_chunks, 2, n_clusters, n_pc), dtype=x.dtype) + o_edge = np.zeros((n_chunks, 2, n_clusters), dtype=x.dtype) + first = np.full(n_chunks, -1, dtype=np.int64) + last = np.full(n_chunks, -1, dtype=np.int64) + for t in range(n_chunks): + lo = t * n_cells // n_chunks + hi = (t + 1) * n_cells // n_chunks + if hi > lo: + first[t] = joint[order[lo]] + last[t] = joint[order[hi - 1]] + for t in prange(n_chunks): + for p in range(t * n_cells // n_chunks, (t + 1) * n_cells // n_chunks): + i = order[p] + j = joint[i] + if j == first[t] or j == last[t]: + side = 0 if j == first[t] else 1 + for k in range(n_clusters): + rk = r[i, k] + o_edge[t, side, k] += rk + for dd in range(n_pc): + s_edge[t, side, k, dd] += rk * x[i, dd] + else: + for k in range(n_clusters): + rk = r[i, k] + o[j, k] += rk + for dd in range(n_pc): + s[j, k, dd] += rk * x[i, dd] + for t in range(n_chunks): + if first[t] < 0: + continue + for side in range(2): + j = first[t] if side == 0 else last[t] + for k in range(n_clusters): + o[j, k] += o_edge[t, side, k] + for dd in range(n_pc): + s[j, k, dd] += s_edge[t, side, k, dd] + return s, o + + +@njit +def _apply_correction(x, r, joint, w_joint, z_hat, z_norm, n_chunks): # noqa: PLR0917 + """Apply the correction, normalise, and accumulate the next centroids. + + Writes ``z_hat = x_i - sum_k r_ik * w_joint[joint_i, k, :]`` and its row-normalised copy ``z_norm`` in place; + returns ``y_t`` (d x K), the normalised centroids of ``r`` on the new ``z_norm``. + + Each row i is written in full before it is read and no other row is touched, so passing the same array as both + ``z_hat`` and ``z_norm`` is safe and yields the normalised result (the low_memory path relies on this). + """ + ft = x.dtype.type + n_cells, n_pc = x.shape + n_clusters = r.shape[1] + one = ft(1.0) + y_acc = np.zeros((n_chunks, n_pc, n_clusters), dtype=x.dtype) + for t in prange(n_chunks): + y_loc = y_acc[t] + for i in range(t * n_cells // n_chunks, (t + 1) * n_cells // n_chunks): + j = joint[i] + for dd in range(n_pc): + z_hat[i, dd] = x[i, dd] + for k in range(n_clusters): + rk = r[i, k] + for dd in range(n_pc): + z_hat[i, dd] -= rk * w_joint[j, k, dd] + sq = ft(0.0) + for dd in range(n_pc): + sq += z_hat[i, dd] * z_hat[i, dd] + scale = one / max(np.sqrt(sq), ft(1e-12)) + for dd in range(n_pc): + zv = z_hat[i, dd] * scale + z_norm[i, dd] = zv + for k in range(n_clusters): + y_loc[dd, k] += zv * r[i, k] + y_t = np.empty((n_pc, n_clusters), dtype=x.dtype) + for k in range(n_clusters): + sq = ft(0.0) + for dd in range(n_pc): + v = ft(0.0) + for t in range(n_chunks): + v += y_acc[t, dd, k] + y_t[dd, k] = v + sq += v * v + scale = one / max(np.sqrt(sq), ft(1e-12)) + for dd in range(n_pc): + y_t[dd, k] *= scale + return y_t + + +@njit +def _apply_only(x, r, joint, w_joint, out, n_chunks): # noqa: PLR0917 + """``out[i] = x[i] - sum_k r_ik * w_joint[joint_i, k]``: the unnormalised corrected embedding, no centroids. + + The last step of the low-memory path, which reconstructs ``z_hat`` into the reused ``z_norm`` buffer once at the + end instead of keeping a separate ``z_hat`` array for every iteration (``z_hat`` is written but never read until + the return). Same per-cell arithmetic as the first half of ``_apply_correction``. + """ + n_cells, n_pc = x.shape + n_clusters = r.shape[1] + for t in prange(n_chunks): + for i in range(t * n_cells // n_chunks, (t + 1) * n_cells // n_chunks): + j = joint[i] + for dd in range(n_pc): + out[i, dd] = x[i, dd] + for k in range(n_clusters): + rk = r[i, k] + for dd in range(n_pc): + out[i, dd] -= rk * w_joint[j, k, dd] + return out diff --git a/src/scanpy/preprocessing/_harmony/core.py b/src/scanpy/preprocessing/_harmony/core.py index 02a37bf09c..d612124d8a 100644 --- a/src/scanpy/preprocessing/_harmony/core.py +++ b/src/scanpy/preprocessing/_harmony/core.py @@ -1,16 +1,19 @@ from __future__ import annotations -from dataclasses import KW_ONLY, InitVar, dataclass, field +from itertools import pairwise from typing import TYPE_CHECKING import numpy as np +from scipy import sparse from sklearn.cluster import KMeans from tqdm.auto import tqdm from ... import logging as log +from ..._compat import warn from ..._settings import settings from ..._settings.verbosity import Verbosity from ..._utils.random import _legacy_random_state +from ._kernels import _apply_correction, _apply_only, _cluster_pass, _joint_stats if TYPE_CHECKING: from collections.abc import Sequence @@ -21,222 +24,278 @@ from ..._utils.random import RNGLike, SeedLike -__all__ = ["_SUPPRESS_PENALTY", "Harmony", "_compute_lambda_kb"] +__all__ = ["_SUPPRESS_PENALTY", "_compute_lambda_kb", "harmonize"] _SUPPRESS_PENALTY = 1e30 - -@dataclass -class Harmony: - """Harmony batch correction algorithm. - - Parameters - ---------- - x - Data matrix (n_cells x d) - typically PCA embeddings. - batch_df - DataFrame containing batch information. - batch_key - Column name(s) in batch_df containing batch labels. +# stop_rule -> (window of the clustering test, record every round's objective, +# {flavor: (max_iter_clustering, tol_clustering, tol_harmony)}) +_REFERENCE = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (4, 1e-3, 1e-2)} +_PAPER = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (200, 1e-5, 1e-4)} +STOP_RULES = { + # the round objective is only recorded if clustering converged, as in harmonypy + "reference": (3, False, _REFERENCE), + # every round's objective is recorded, as in rapids-singlecell + "reference_record": (3, True, _REFERENCE), + "paper": (1, True, _PAPER), + "paper_window3": (3, True, _PAPER), +} +SEEDINGS = ("afkmc2", "kmeans", "kmeans++", "paper_nearest") +# Fixed number of static chunks per parallel sweep, so results do not depend on the thread count. +N_CHUNKS = 32 + + +def harmonize( # noqa: PLR0912, PLR0913, PLR0915 + x: NDArray[np.floating], + batch_df: pd.DataFrame, + batch_key: str | Sequence[str], + *, + flavor: str = "harmony2", + n_clusters: int | None = None, + sigma: float = 0.1, + theta: float | Sequence[float] = 2.0, + tau: int = 0, + ridge_lambda: float = 1.0, + alpha: float = 0.2, + batch_prune_threshold: float | None = 1e-5, + block_proportion: float = 0.05, + max_iter_harmony: int = 10, + max_iter_clustering: int | None = None, + tol_harmony: float | None = None, + tol_clustering: float | None = None, + stop_rule: str = "reference_record", + seeding: str = "afkmc2", + penalty_combine: str = "product", + seeding_lloyd_steps: int = 0, + seeding_cells_per_cluster: int | None = None, + recompute_counts: bool = True, + low_memory: bool = False, + rng: SeedLike | RNGLike | None = None, +) -> tuple[NDArray[np.floating], dict]: + """Run Harmony on ``x`` (cells x components) in the dtype of ``x``. + + Returns the corrected embedding and a dict with the Harmony objective per round + (``objectives``) and the clustering iterations per round (``iterations``). """ - - batch_df: InitVar[pd.DataFrame] - batch_key: InitVar[str | Sequence[str]] - _: KW_ONLY - theta: float | Sequence[float] - sigma: float - n_clusters: int | None - max_iter_harmony: int - max_iter_clustering: int - tol_harmony: float - tol_clustering: float - ridge_lambda: float - block_proportion: float - tau: int - rng: InitVar[SeedLike | RNGLike | None] - stabilized_penalty: bool = True - dynamic_lambda: bool = True - alpha: float = 0.2 - batch_prune_threshold: float | None = 1e-5 - - batch_codes: NDArray[np.int32] = field(init=False) - n_levels: NDArray[np.int32] = field(init=False) - n_batches: int = field(init=False) - n_covariates: int = field(init=False) - _rng: np.random.Generator = field(init=False) - - def __post_init__( - self, - batch_df: pd.DataFrame, - batch_key: str | Sequence[str], - rng: SeedLike | RNGLike | None, - ) -> None: - self._rng = np.random.default_rng(rng) - - if self.max_iter_harmony < 1: - msg = "max_iter_harmony must be >= 1" + stabilized = dynamic_lambda = flavor == "harmony2" + choices = dict( + stop_rule=tuple(STOP_RULES), + seeding=SEEDINGS, + penalty_combine=("product", "sum"), + ) + given = dict(stop_rule=stop_rule, seeding=seeding, penalty_combine=penalty_combine) + for name, allowed in choices.items(): + if given[name] not in allowed: + msg = f"{name} must be one of {allowed}, got {given[name]!r}." raise ValueError(msg) - - if self.dynamic_lambda: - if not np.isfinite(self.alpha) or self.alpha <= 0: - msg = f"alpha must be a finite positive number when dynamic_lambda=True, got {self.alpha}." - raise ValueError(msg) - if self.batch_prune_threshold is not None and not ( - 0 <= self.batch_prune_threshold <= 1 - ): - msg = f"batch_prune_threshold must be in [0, 1] or None, got {self.batch_prune_threshold}." - raise ValueError(msg) - elif not np.isfinite(self.ridge_lambda) or self.ridge_lambda <= 0: - msg = ( - "ridge_lambda must be a finite positive number when " - f"dynamic_lambda=False, got {self.ridge_lambda}." - ) + checks = { + "max_iter_harmony must be >= 1": max_iter_harmony >= 1, + "max_iter_clustering must be >= 1": max_iter_clustering is None + or max_iter_clustering >= 1, + f"sigma must be a finite positive number, got {sigma}.": np.isfinite(sigma) + and sigma > 0, + f"block_proportion must be in (0, 1], got {block_proportion}.": 0 + < block_proportion + <= 1, + f"alpha must be a finite positive number when flavor='harmony2', got {alpha}.": not dynamic_lambda + or (np.isfinite(alpha) and alpha > 0), + f"batch_prune_threshold must be in [0, 1] or None, got {batch_prune_threshold}.": not dynamic_lambda + or batch_prune_threshold is None + or 0 <= batch_prune_threshold <= 1, + f"ridge_lambda must be a finite positive number when flavor='harmony1', got {ridge_lambda}.": dynamic_lambda + or (np.isfinite(ridge_lambda) and ridge_lambda > 0), + f"seeding_lloyd_steps must be >= 0, got {seeding_lloyd_steps}.": seeding_lloyd_steps + >= 0, + f"seeding_cells_per_cluster must be >= 1 or None, got {seeding_cells_per_cluster}.": seeding_cells_per_cluster + is None + or seeding_cells_per_cluster >= 1, + } + for msg, ok in checks.items(): + if not ok: raise ValueError(msg) - - # Process batch keys - self.batch_codes, self.n_levels = _get_batch_codes(batch_df, batch_key) - self.n_batches = int(self.n_levels.sum()) - self.n_covariates = len(self.n_levels) - - def fit(self, x: np.ndarray) -> np.ndarray: - """Run Harmony. - - Returns - ------- - z_corr - Batch-corrected embedding matrix (n_cells x d). - """ - # Ensure input is C-contiguous float array (infer dtype from x) - x = np.ascontiguousarray(x) - n_cells = x.shape[0] - - # Normalize input for clustering - z_norm = _normalize_rows_l2(x) - - # Compute batch proportions - n_b = np.bincount(self.batch_codes.ravel(), minlength=self.n_batches).astype( - x.dtype + x = np.ascontiguousarray(x) + if x.ndim != 2: + msg = f"x must be a 2-dimensional array (cells x components), got {x.ndim} dimension(s)." + raise ValueError(msg) + if not np.isfinite(x).all(): + msg = "Input data contains non-finite values." + raise ValueError(msg) + ft = x.dtype.type + window, record_always, defaults = STOP_RULES[stop_rule] + # only differs from the product with two or more batch keys + sum_penalties = penalty_combine == "sum" + d_iter, d_tol_c, d_tol_h = defaults[flavor] + max_iter_clustering = d_iter if max_iter_clustering is None else max_iter_clustering + tol_clustering = d_tol_c if tol_clustering is None else tol_clustering + tol_harmony = d_tol_h if tol_harmony is None else tol_harmony + rng = np.random.default_rng(rng) + info = dict(objectives=[], iterations=[]) + + codes, n_levels = _get_batch_codes(batch_df, batch_key) + if codes.shape[0] != x.shape[0]: # the kernels index codes without bounds checks + msg = f"batch_df has {codes.shape[0]} rows but x has {x.shape[0]} cells." + raise ValueError(msg) + # all-zero cells have no direction on the unit sphere: leave them out and return them unchanged + x_all, keep = x, np.any(x != 0, axis=1) + if not keep.all(): + warn( + f"{(~keep).sum()} cells have an all-zero embedding; " + "they are excluded from Harmony and returned unchanged.", + UserWarning, ) - pr_b = (n_b / n_cells).reshape(-1, 1) - - theta_arr = _get_theta_array(self.theta, self.n_levels, x.dtype) - - # Set default n_clusters (needed before tau discounting) - if self.n_clusters is None: - n_clusters = int(min(100, n_cells / 30)) - n_clusters = max(n_clusters, 2) - else: - n_clusters = self.n_clusters - - # Apply tau discounting to theta - if self.tau > 0: - theta_arr = theta_arr * (1 - np.exp(-n_b / (n_clusters * self.tau)) ** 2) - - # Initialize centroids and state arrays - r, e, o, obj_init = _initialize_centroids( - z_norm, - self.batch_codes, - self.n_batches, - pr_b, - n_clusters=n_clusters, - sigma=self.sigma, - theta=theta_arr, - rng=self._rng, - stabilized_penalty=self.stabilized_penalty, + x, codes = np.ascontiguousarray(x[keep]), np.ascontiguousarray(codes[keep]) + n_cells = x.shape[0] + if n_cells == 0: + msg = "No cell has a non-zero embedding." + raise ValueError(msg) + n_batches = int(n_levels.sum()) + joint, levels = _joint_categories(codes) + joint_order = np.argsort(joint, kind="stable") + n_b = np.bincount(codes.ravel(), minlength=n_batches).astype(ft) + pr_b = n_b / ft(n_cells) + theta_b = _get_theta_array(theta, n_levels, x.dtype).ravel() + k = max(2, int(min(100, n_cells / 30))) if n_clusters is None else n_clusters + if tau > 0: + theta_b = (theta_b * (1 - np.exp(-n_b / (k * tau)) ** 2)).astype(ft) + + z_norm = _normalize_rows_l2(x) + # optional subsample for the parts of seeding that iterate over cells (KMeans, Lloyd polish); + # the other seeders read all cells and ignore the cap + z_seed = z_norm + uses_cap = seeding == "kmeans" or seeding_lloyd_steps > 0 + if ( + uses_cap + and seeding_cells_per_cluster is not None + and n_cells > seeding_cells_per_cluster * k + ): + cells = rng.choice(n_cells, seeding_cells_per_cluster * k, replace=False) + z_seed = z_norm[np.sort(cells)] + if seeding == "kmeans": + km = KMeans( + n_clusters=k, + random_state=_legacy_random_state(rng, always_state=True), + max_iter=25, ) - - # Main Harmony loop - objectives_harmony = [obj_init] - with tqdm( - range(self.max_iter_harmony), disable=settings.verbosity < Verbosity.info - ) as bar: - for i in bar: - r, e, o, obj = self._cluster( - z_norm, pr_b, r=r, e=e, o=o, theta=theta_arr + y = _normalize_rows_l2(km.fit(z_seed).cluster_centers_.astype(ft)) + elif seeding == "kmeans++": + y = _seed_kmeanspp(z_norm, k, rng) + elif seeding == "afkmc2": + y = _seed_afkmc2(z_norm, k, rng) + else: + y = _seed_paper_nearest(z_norm, k, rng) + if seeding_lloyd_steps: + y = _lloyd_spherical(z_seed, y, seeding_lloyd_steps) + + # initial r: one un-penalised pass over all cells as a single block + term = ft(-2.0 / sigma) + r = np.zeros((n_cells, k), dtype=ft) + o = np.zeros((n_batches, k), dtype=ft) + r_sum = np.zeros(k, dtype=ft) + dist, ent, _, y_t = _cluster_pass( + z_norm, + np.ascontiguousarray(y.T), + r, + codes, + o, + r_sum, + pr_b, + np.zeros(n_batches, dtype=ft), + np.arange(n_cells), + np.array([0, n_cells]), + term, + stabilized, + N_CHUNKS, + _use_scatter(n_batches, codes.shape[1], n_cells), + sum_penalties, + ) + e = pr_b[:, None] * r_sum[None, :] + objs_harmony = [ + dist + sigma * ent + _diversity(o, e, theta_b, sigma, stabilized=stabilized) + ] + + bounds = _block_bounds(n_cells, block_proportion) + scatter = _use_scatter(n_batches, codes.shape[1], int(bounds[1])) + # low_memory reuses the z_norm buffer as z_hat's scratch; otherwise z_hat is its own array + z_hat = z_norm if low_memory else np.empty_like(x) + with tqdm( + range(max_iter_harmony), disable=settings.verbosity < Verbosity.info + ) as bar: + for i in bar: + objs: list[float] = [] + done = False + for _ in range(max_iter_clustering): + perm = rng.permutation(n_cells) + # sort each block's cells so the kernels read z_norm/r/z_hat in ascending order; + # block membership is unchanged, only the order within a block differs + for lo, hi in pairwise(bounds): + perm[lo:hi].sort() + dist, ent, o_fresh, y_t = _cluster_pass( + z_norm, + y_t, + r, + codes, + o, + r_sum, + pr_b, + theta_b, + perm, + bounds, + term, + stabilized, + N_CHUNKS, + scatter, + sum_penalties, ) - if obj is not None: - objectives_harmony.append(obj) - - # Compute per-(k,b) ridge regularization - lambda_kb = _compute_lambda_kb( - e, - o=o, - n_b=n_b, - alpha=self.alpha, - threshold=self.batch_prune_threshold, - ridge_lambda=self.ridge_lambda, - dynamic_lambda=self.dynamic_lambda, + if recompute_counts: + o[:] = o_fresh + r_sum[:] = o_fresh[: n_levels[0]].sum(axis=0) + e = pr_b[:, None] * r_sum[None, :] + objs.append( + dist + + sigma * ent + + _diversity(o, e, theta_b, sigma, stabilized=stabilized) ) - - z_hat = self._correct(x, r, o, lambda_kb=lambda_kb) - z_norm = _normalize_rows_l2(z_hat) - if self._is_convergent(objectives_harmony, self.tol_harmony): - log.info(f"Harmony converged in {i + 1} iterations") + if done := _converged( + objs, tol_clustering, window=window, signed=window > 1 + ): break - else: - log.info( - f"Harmony did not converge after {self.max_iter_harmony} iterations." - ) - - return z_hat - - @staticmethod - def _is_convergent(objectives: list[float], tol: float) -> bool: - """Check Harmony convergence.""" - if len(objectives) < 2: - return False - obj_old = objectives[-2] - obj_new = objectives[-1] - return (obj_old - obj_new) < tol * abs(obj_old) - - def _cluster( - self, - z_norm: np.ndarray, - pr_b: np.ndarray, - *, - r: np.ndarray, - e: np.ndarray, - o: np.ndarray, - theta: np.ndarray, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray, float | None]: - """Perform clustering step. Modifies r, e, o in-place.""" - return _clustering( - z_norm, - self.batch_codes, - pr_b, - r=r, - e=e, - o=o, - theta=theta, - sigma=self.sigma, - max_iter=self.max_iter_clustering, - tol=self.tol_clustering, - block_proportion=self.block_proportion, - stabilized_penalty=self.stabilized_penalty, - rng=self._rng, - ) - - def _correct( - self, - x: np.ndarray, - r: np.ndarray, - o: np.ndarray, - *, - lambda_kb: np.ndarray, - ) -> np.ndarray: - """Perform correction step.""" - if self.n_covariates > 1: - return _correction_multi( - *(x, self.batch_codes, self.n_batches, r), - lambda_kb=lambda_kb, + info["iterations"].append(len(objs)) + if done or record_always: + objs_harmony.append(objs[-1]) + + lam = _compute_lambda_kb( + e, + o=o, + n_b=n_b, + alpha=alpha, + threshold=batch_prune_threshold, + ridge_lambda=ridge_lambda, + dynamic_lambda=dynamic_lambda, ) - - batch_codes = self.batch_codes[:, 0] - return _correction_fast( - *(x, batch_codes, self.n_batches, r, o), - lambda_kb=lambda_kb, - ) + s_joint, o_joint = _joint_stats( + x, r, joint, joint_order, levels.shape[0], N_CHUNKS + ) + w_joint = _solve_correction( + s_joint, o_joint, levels, n_levels, lam, _SUPPRESS_PENALTY + ) + # low_memory: write the normalised result into z_norm in place (z_hat aliases it) + y_t = _apply_correction(x, r, joint, w_joint, z_hat, z_norm, N_CHUNKS) + if _converged(objs_harmony, tol_harmony, window=1, signed=True): + log.info(f"Harmony converged in {i + 1} iterations") + break + else: + log.info(f"Harmony did not converge after {max_iter_harmony} iterations.") + info["objectives"] = objs_harmony + if ( + low_memory + ): # the shared buffer holds the normalised embedding; rebuild z_hat into it + z_hat = _apply_only(x, r, joint, w_joint, z_hat, N_CHUNKS) + if not keep.all(): + x_all = x_all.copy() + x_all[keep] = z_hat + z_hat = x_all + return z_hat, info def _get_batch_codes( @@ -310,178 +369,231 @@ def _normalize_rows_l2(x: np.ndarray) -> np.ndarray: return x / norms -def _normalize_rows_l1(r: np.ndarray) -> None: - """L1 normalize each row of r in-place (rows sum to 1).""" - row_sums = r.sum(axis=1, keepdims=True) - row_sums = np.maximum(row_sums, 1e-12) - r /= row_sums +def _converged(objs: list[float], tol: float, *, window: int, signed: bool) -> bool: + """Relative change of the summed last `window` objectives against the `window` before.""" + if len(objs) < window + 1: + return False + old, new = sum(objs[-window - 1 : -1]), sum(objs[-window:]) + diff = old - new + return (diff if signed else abs(diff)) < tol * abs(old) -def _initialize_centroids( - z_norm: np.ndarray, - batch_codes: np.ndarray, - n_batches: int, - pr_b: np.ndarray, - *, - n_clusters: int, - sigma: float, - theta: np.ndarray, - rng: np.random.Generator, - stabilized_penalty: bool = True, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]: - """Initialize cluster centroids using K-means.""" - kmeans = KMeans( - n_clusters=n_clusters, - random_state=_legacy_random_state(rng, always_state=True), - max_iter=25, - ) - kmeans.fit(z_norm) - - # Centroids - y = kmeans.cluster_centers_.copy() - y_norm = _normalize_rows_l2(y) - - # Compute soft cluster assignments r - term = -2.0 / sigma - r = _compute_r(z_norm, y_norm, term) - _normalize_rows_l1(r) - - # Initialize e (expected) and o (observed) - r_sum = r.sum(axis=0) - e = pr_b @ r_sum.reshape(1, -1) - # o[b, k] = sum of r[i, k] for cells in marginal category b - o = np.zeros((n_batches, n_clusters), dtype=z_norm.dtype) - for codes in batch_codes.T: - np.add.at(o, codes, r) - - # Compute initial objective - obj = _compute_objective( - y_norm, - z_norm, - r, - theta=theta, - sigma=sigma, - o=o, - e=e, - stabilized_penalty=stabilized_penalty, - ) +def _seed_paper_nearest( + z_norm: np.ndarray, k: int, rng: np.random.Generator +) -> np.ndarray: + """Seeding of :cite:t:`Patikas2026`. - return r, e, o, obj + k distinct cells; each new one is the cell with the highest ``log(u_j) / D_j``, ``u ~ U(0, 1)``, + which draws it with probability proportional to its Euclidean distance D_j to the nearest cell chosen so far. + """ + n = z_norm.shape[0] + chosen = [int(rng.integers(n))] + taken = np.zeros(n, dtype=bool) + dist = np.full(n, np.inf, dtype=z_norm.dtype) + while len(chosen) < k: + taken[chosen[-1]] = True + # unit rows: ||a - b|| = sqrt(2 - 2 a.b) + d = np.sqrt(np.maximum(2.0 - 2.0 * (z_norm @ z_norm[chosen[-1]]), 0.0)) + np.minimum(dist, d, out=dist) + with np.errstate(divide="ignore"): + s_j = np.log(rng.random(n)) / dist + s_j[taken] = -np.inf + chosen.append(int(np.argmax(s_j))) + return z_norm[chosen].copy() + + +def _seed_kmeanspp(z_norm: np.ndarray, k: int, rng: np.random.Generator) -> np.ndarray: + """k-means++ on unit vectors: each seed drawn with probability ∝ squared distance to the nearest seed.""" + n = z_norm.shape[0] + chosen = [int(rng.integers(n))] + closest = np.full(n, np.inf, dtype=z_norm.dtype) + while len(chosen) < k: + d2 = np.maximum(2.0 - 2.0 * (z_norm @ z_norm[chosen[-1]]), 0.0) + np.minimum(closest, d2, out=closest) + cum = np.cumsum(closest, dtype=np.float64) + # all remaining cells coincide with a seed: fall back to a uniform draw + pick = ( + np.searchsorted(cum, rng.random() * cum[-1], side="right") + if cum[-1] > 0 + else rng.integers(n) + ) + chosen.append(int(min(pick, n - 1))) + return z_norm[chosen].copy() -def _compute_r( - z: np.ndarray, - y: np.ndarray, - term: float, +def _seed_afkmc2( + z_norm: np.ndarray, k: int, rng: np.random.Generator, *, chain: int = 200 ) -> np.ndarray: - """Compute soft cluster assignments using NumPy dot.""" - dots = z @ y.T - return np.exp(term * (1.0 - dots)) + """AFK-MC² :cite:p:`Bachem2016`: k-means++ approximated by one Metropolis chain per seed. - -def _clustering( # noqa: PLR0913 - z_norm: NDArray[np.floating], - batch_codes: np.ndarray, - pr_b: np.ndarray, - *, - r: np.ndarray, - e: np.ndarray, - o: np.ndarray, - theta: np.ndarray, - sigma: float, - max_iter: int, - tol: float, - block_proportion: float, - rng: np.random.Generator, - stabilized_penalty: bool = True, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, float | None]: - """Run clustering iterations (modifies r, e, o in-place).""" - n_cells: int = z_norm.shape[0] - k: int = r.shape[1] - n_blocks = min(n_cells, 1 // block_proportion) - term = -2.0 / sigma - - objectives_clustering = [] - - # Pre-allocate work arrays - y = np.empty((k, z_norm.shape[1]), dtype=z_norm.dtype) - y_norm = np.empty_like(y) - - for _ in range(max_iter): - # Compute cluster centroids: y = r.T @ z_norm, then normalize - np.dot(r.T, z_norm, out=y) - norms = np.linalg.norm(y, axis=1, keepdims=True) - norms = np.maximum(norms, 1e-12) - np.divide(y, norms, out=y_norm) - - # Randomly shuffle cell indices - idx_list = rng.permutation(n_cells) - - # Process blocks. Every cell contributes once to each batch variable, - # while the expected counts depend only on its cluster assignment. - for cell_idx in np.array_split(idx_list, n_blocks): - cell_codes = batch_codes[cell_idx] - - # Remove old r contribution from all marginal categories and from e. - r_old = r[cell_idx] - r_old_sum = r_old.sum(axis=0) - for codes in cell_codes.T: - np.add.at(o, codes, -r_old) - e -= pr_b * r_old_sum - - # Compute new r values - dots = z_norm[cell_idx, :] @ y_norm.T - r_new = np.empty_like(dots) - - # Apply the product of the marginal penalties in log space. - # This both follows the Harmony formulation and avoids overflow - # when several variables have large finite penalty factors. - if stabilized_penalty: - # Harmony2: denominator is (O + E + 1) - log_penalty = theta.T * (np.log(e + 1.0) - np.log(o + e + 1.0)) - else: - # Harmony1: denominator is (O + 1) - log_penalty = theta.T * (np.log(e + 1.0) - np.log(o + 1.0)) - log_r_new = term * (1.0 - dots) - log_r_new += log_penalty[cell_codes].sum(axis=1) - log_r_new -= log_r_new.max(axis=1, keepdims=True) - np.exp(log_r_new, out=r_new) - - # Normalize rows to sum to 1 - row_sums = r_new.sum(axis=1, keepdims=True) - row_sums = np.maximum(row_sums, 1e-12) - r_new /= row_sums - - # Store back - r[cell_idx, :] = r_new - - # Add new r contribution to o and e - r_new_sum = r_new.sum(axis=0) - for codes in cell_codes.T: - np.add.at(o, codes, r_new) - e += pr_b * r_new_sum - - # Compute objective - obj = _compute_objective( - y_norm, - z_norm, - r, - theta=theta, - sigma=sigma, - o=o, - e=e, - stabilized_penalty=stabilized_penalty, + Proposal q = d²(x, first seed) / (2 Σ d²) + 1 / (2 n); one pass over the cells, then cost O(k² · chain · d). + """ + n = z_norm.shape[0] + first = int(rng.integers(n)) + d2 = np.maximum(2.0 - 2.0 * (z_norm @ z_norm[first]), 0.0).astype(np.float64) + total = d2.sum() + q = 0.5 * d2 / total + 0.5 / n if total > 0 else np.full(n, 1.0 / n) + candidates = rng.choice(n, size=(k - 1, chain), p=q) + uniforms = rng.random((k - 1, chain)) + seeds = np.empty((k, z_norm.shape[1]), dtype=z_norm.dtype) + seeds[0] = z_norm[first] + for c in range(1, k): + cand = z_norm[candidates[c - 1]] + # d² to the nearest seed so far + dist = np.maximum(2.0 - 2.0 * (cand @ seeds[:c].T), 0.0).min(axis=1) + weight = dist / q[candidates[c - 1]] + cur = 0 + for j in range(1, chain): + if weight[cur] == 0 or weight[j] > uniforms[c - 1, j] * weight[cur]: + cur = j + seeds[c] = cand[cur] + return seeds + + +def _lloyd_spherical(z_norm: np.ndarray, y: np.ndarray, steps: int) -> np.ndarray: + """Hard spherical k-means steps: assign each cell to its most similar centroid, re-normalise the sums.""" + cells = np.arange(z_norm.shape[0]) + for _ in range(steps): + labels = np.argmax(z_norm @ y.T, axis=1) + onehot = sparse.csr_matrix( # noqa: TID251 + (np.ones(cells.size, dtype=z_norm.dtype), (labels, cells)), + shape=(y.shape[0], cells.size), ) - objectives_clustering.append(obj) + sums = np.asarray(onehot @ z_norm) + norms = np.linalg.norm(sums, axis=1, keepdims=True) + # empty clusters keep their centroid + y = np.where(norms > 0, sums / np.maximum(norms, 1e-12), y).astype(y.dtype) + return y + + +def _joint_categories(codes: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Index of each cell's combination of batch levels, and the levels (J x C) of every combination.""" + levels, joint = np.unique(codes, axis=0, return_inverse=True) + return np.ascontiguousarray(joint.ravel().astype(np.int32)), levels + + +def _solve_correction( # noqa: PLR0915, PLR0917 + s_joint: np.ndarray, + o_joint: np.ndarray, + levels: np.ndarray, + n_levels: np.ndarray, + lam: np.ndarray, + sentinel: float, +) -> np.ndarray: + """Ridge coefficients of every cluster, returned as the total correction per joint category (J x K x d). - # Check convergence - if _is_convergent_clustering(objectives_clustering, tol): - obj = objectives_clustering[-1] - break - else: - obj = None + Per cluster the normal equations are ``(Phi* diag(r_k) Phi*^T + Lambda_k) W_k = Phi* diag(r_k) X^T`` + with an unpenalised intercept. The block of the covariate with the most levels is diagonal, + so it is eliminated in closed form and only the Schur complement over the intercept and the remaining + levels is solved; with one covariate that is a 1 x 1 system. ``lam[b, k] >= sentinel`` removes level b + from cluster k (its coefficient is 0). Runs in float64; nothing here has a cell dimension. + """ + n_joint, n_clusters, n_pc = s_joint.shape + offsets = np.concatenate([[0], np.cumsum(n_levels)]) + big = int(np.argmax(n_levels)) + n_big = int(n_levels[big]) + others = [c for c in range(len(n_levels)) if c != big] + big_of_joint = levels[:, big] - offsets[big] + small = np.concatenate( + [np.arange(offsets[c], offsets[c + 1]) for c in others] + + [np.empty(0, dtype=np.int64)] + ).astype(np.int64) + m = 1 + small.size # intercept and the levels that are not eliminated + row_of_level = np.full(int(offsets[-1]), -1) + row_of_level[small] = 1 + np.arange(small.size) + # rows[j] = the rows of the small system that joint category j belongs to: + # the intercept and one level per other key + rows = np.concatenate( + [np.zeros((n_joint, 1), dtype=np.int64), row_of_level[levels[:, others]]], + axis=1, + ) + + o = o_joint.astype(np.float64) # J x K + s = s_joint.astype(np.float64) # J x K x d + lam = lam.astype(np.float64) + pruned = lam >= sentinel + lam_big = lam[offsets[big] : offsets[big + 1]].T # K x n_big + pruned_big = pruned[offsets[big] : offsets[big + 1]].T + lam_small, pruned_small = lam[small].T, pruned[small].T # K x small + + # sums over joint categories; cluster axis last while accumulating, first afterwards + o_big = np.zeros((n_big, n_clusters)) + np.add.at(o_big, big_of_joint, o) + rhs_big = np.zeros((n_big, n_clusters, n_pc)) + np.add.at(rhs_big, big_of_joint, s) + gram = np.zeros((m, m, n_clusters)) + rhs = np.zeros((m, n_clusters, n_pc)) + coupling = np.zeros(( + m, + n_big, + n_clusters, + )) # mass shared by a small row and a big level + for a in range(rows.shape[1]): + np.add.at(rhs, rows[:, a], s) + np.add.at(coupling, (rows[:, a], big_of_joint), o) + for b in range(rows.shape[1]): + np.add.at(gram, (rows[:, a], rows[:, b]), o) + gram = gram.transpose(2, 0, 1).copy() + rhs = rhs.transpose(1, 0, 2).copy() + coupling = coupling.transpose(2, 0, 1).copy() + o_big, rhs_big = o_big.T, rhs_big.transpose(1, 0, 2) # K x n_big, K x n_big x d + + diag = o_big + lam_big + # 0 drops the level + d_inv = np.where(pruned_big | (diag <= 0), 0.0, 1.0 / np.where(diag > 0, diag, 1.0)) + idx = np.arange(1, m) + gram[:, idx, idx] += np.where(pruned_small, 0.0, lam_small) + # a pruned small level: empty row and column, unit diagonal after elimination + kk, ss = np.nonzero(pruned_small) + gram[kk, ss + 1, :] = 0 + gram[kk, :, ss + 1] = 0 + coupling[kk, ss + 1, :] = 0 + rhs[kk, ss + 1] = 0 + scaled = coupling * d_inv[:, None, :] # C D^-1, K x m x n_big + schur = gram - scaled @ coupling.transpose(0, 2, 1) + rhs -= scaled @ rhs_big + schur[kk, ss + 1, ss + 1] = 1 + empty = gram[:, 0, 0] <= 1e-12 # cluster without cells: no correction + schur[empty] = np.eye(m) + try: + w_small = np.linalg.solve(schur, rhs) # K x m x d + except np.linalg.LinAlgError: + # guard only: from `harmonize` every unpruned lambda is positive and the system is positive definite + w_small = np.stack([ + np.linalg.lstsq(a, b, rcond=None)[0] + for a, b in zip(schur, rhs, strict=True) + ]) + w_big = d_inv[:, :, None] * (rhs_big - coupling.transpose(0, 2, 1) @ w_small) + w_small[:, 0] = 0 # the intercept is not part of the correction + w_small[empty] = 0 + w_big[empty] = 0 + w_joint = w_big[:, big_of_joint] # K x J x d + for a in range(1, rows.shape[1]): + w_joint = w_joint + w_small[:, rows[:, a]] + return np.ascontiguousarray(w_joint.transpose(1, 0, 2), dtype=s_joint.dtype) - return r, e, o, obj + +def _correction_multi( + x: np.ndarray, + batch_codes: np.ndarray, + n_batches: int, + r: np.ndarray, + *, + lambda_kb: np.ndarray, +) -> np.ndarray: + """Ridge correction of ``x`` for any number of batch keys, as one Harmony round applies it.""" + # per-covariate level counts from the code ranges; an unused level has no mass, + # so where its range is cut does not change the result + ends = np.append(batch_codes.max(axis=0)[:-1] + 1, n_batches) + n_levels = np.diff(ends, prepend=0) + joint, levels = _joint_categories(batch_codes) + s_joint, o_joint = _joint_stats( + x, r, joint, np.argsort(joint, kind="stable"), levels.shape[0], N_CHUNKS + ) + w_joint = _solve_correction( + s_joint, o_joint, levels, n_levels, lambda_kb, _SUPPRESS_PENALTY + ) + return _apply_only(x, r, joint, w_joint, np.empty_like(x), N_CHUNKS) def _compute_lambda_kb( @@ -511,155 +623,26 @@ def _compute_lambda_kb( return lambda_kb -def _correction_multi( - x: np.ndarray, - batch_codes: np.ndarray, - n_batches: int, - r: np.ndarray, - *, - lambda_kb: np.ndarray, -) -> np.ndarray: - """Apply the exact ridge correction for a multi-covariate design.""" - _, d = x.shape - z = x.copy() - sentinel = x.dtype.type(_SUPPRESS_PENALTY) - - for k_idx, r_k in enumerate(r.T): - marginal_r = np.zeros(n_batches, dtype=x.dtype) - rhs = np.zeros((n_batches + 1, d), dtype=x.dtype) - weighted_x = r_k[:, np.newaxis] * x - rhs[0] = weighted_x.sum(axis=0) - for codes in batch_codes.T: - np.add.at(marginal_r, codes, r_k) - np.add.at(rhs, codes + 1, weighted_x) - - gram = np.zeros((n_batches + 1, n_batches + 1), dtype=x.dtype) - gram[0, 0] = r_k.sum() - gram[0, 1:] = marginal_r - gram[1:, 0] = marginal_r - gram[np.arange(1, n_batches + 1), np.arange(1, n_batches + 1)] = marginal_r - for left in range(batch_codes.shape[1] - 1): - left_codes = batch_codes[:, left] + 1 - for right in range(left + 1, batch_codes.shape[1]): - right_codes = batch_codes[:, right] + 1 - np.add.at(gram, (left_codes, right_codes), r_k) - np.add.at(gram, (right_codes, left_codes), r_k) - - active = lambda_kb[:, k_idx] < sentinel - retained = np.concatenate(([True], active)) - retained_idx = np.flatnonzero(retained) - gram[retained_idx[1:], retained_idx[1:]] += lambda_kb[active, k_idx] - - gram_reduced = gram[np.ix_(retained_idx, retained_idx)] - rhs_reduced = rhs[retained] - try: - if np.linalg.matrix_rank(gram_reduced) < gram_reduced.shape[0]: - raise np.linalg.LinAlgError - w_reduced = np.linalg.solve(gram_reduced, rhs_reduced) - except np.linalg.LinAlgError: - w_reduced = np.linalg.lstsq(gram_reduced, rhs_reduced, rcond=None)[0] - - w = np.zeros((n_batches + 1, d), dtype=x.dtype) - w[retained] = w_reduced - w[0] = 0 - correction = np.zeros_like(x) - for codes in batch_codes.T: - correction += w[codes + 1] - z -= r_k[:, np.newaxis] * correction - - return z - - -def _correction_fast( - x: np.ndarray, - batch_codes: np.ndarray, - n_batches: int, - r: np.ndarray, - o: np.ndarray, - *, - lambda_kb: np.ndarray, -) -> np.ndarray: - """Fast correction method using precomputed factors.""" - _, d = x.shape - - z = x.copy() - dtype = x.dtype - p = np.eye(n_batches + 1, dtype=dtype) - - # Compute all cluster/category right-hand sides together. - # This avoids rebuilding the same batch masks once per cluster and - # lets BLAS process the full responsibility matrix for each category. - phi_t_x = np.empty((r.shape[1], n_batches + 1, d), dtype=dtype) - phi_t_x[:, 0] = r.T @ x - for batch in range(n_batches): - mask = batch_codes == batch - phi_t_x[:, batch + 1] = r[mask].T @ x[mask] +def _use_scatter(n_batches: int, n_cov: int, block_size: int) -> bool: + """Whether the clustering kernel moves a block's counts by serial scatters instead of per-chunk accumulators. - for k_idx, (o_k, r_k) in enumerate(zip(o.T, r.T, strict=True)): - lam_k = lambda_kb[:, k_idx] - - factor = (1.0 / (o_k + lam_k)).astype(dtype) - c = np.sum(o_k) + np.sum(-factor * o_k**2) - c_inv = dtype.type(1.0 / c) - - p[0, 1:] = -factor * o_k - - p_t_b_inv = np.zeros((n_batches + 1, n_batches + 1), dtype=dtype) - p_t_b_inv[0, 0] = c_inv - p_t_b_inv[1:, 1:] = np.diag(factor) - p_t_b_inv[1:, 0] = p[0, 1:] * c_inv - - inv_mat = p_t_b_inv @ p - - w = inv_mat @ phi_t_x[k_idx] - w[0, :] = 0 + Accumulators cost about N_CHUNKS * levels * K per block, scatters about cells * keys * K. + Scatters win for thousands of small batches, accumulators on atlases with few large ones. + """ + return N_CHUNKS * n_batches > block_size * n_cov - w_batch = w[batch_codes + 1, :] - z -= r_k[:, np.newaxis] * w_batch - return z +def _block_bounds(n_cells: int, block_proportion: float) -> np.ndarray: + """Block boundaries within a permutation; the split of `np.array_split` into `1 // block_proportion` blocks.""" + n_blocks = int(min(n_cells, 1 // block_proportion)) + sizes = np.full(n_blocks, n_cells // n_blocks) + sizes[: n_cells % n_blocks] += 1 + return np.concatenate([[0], np.cumsum(sizes)]).astype(np.int64) -def _compute_objective( - y_norm: np.ndarray, - z_norm: np.ndarray, - r: np.ndarray, - *, - theta: np.ndarray, - sigma: float, - o: np.ndarray, - e: np.ndarray, - stabilized_penalty: bool = True, +def _diversity( + o: np.ndarray, e: np.ndarray, theta: np.ndarray, sigma: float, *, stabilized: bool ) -> float: - """Compute Harmony objective function.""" - zy = z_norm @ y_norm.T - kmeans_error = np.sum(r * 2.0 * (1.0 - zy)) - - r_row_sums = r.sum(axis=1, keepdims=True) - r_normalized = r / np.clip(r_row_sums, 1e-12, None) - entropy = sigma * np.sum(r_normalized * np.log(r_normalized + 1e-12)) - - if stabilized_penalty: - # Harmony2: numerator is (O + E + 1) - log_ratio = np.log((o + e + 1) / (e + 1)) - else: - # Harmony1: numerator is (O + 1) - log_ratio = np.log((o + 1) / (e + 1)) - diversity_penalty = sigma * np.sum(theta @ (o * log_ratio)) - - return kmeans_error + entropy + diversity_penalty - - -def _is_convergent_clustering( - objectives: list, - tol: float, - window_size: int = 3, -) -> bool: - """Check clustering convergence using window.""" - if len(objectives) < window_size + 1: - return False - - obj_old = sum(objectives[-window_size - 1 : -1]) - obj_new = sum(objectives[-window_size:]) - - return (obj_old - obj_new) < tol * abs(obj_old) + o, e = o.astype(np.float64), e.astype(np.float64) + ratio = (o + e + 1) / (e + 1) if stabilized else (o + 1) / (e + 1) + return sigma * float(np.sum(theta[:, None] * o * np.log(ratio))) diff --git a/tests/test_harmony.py b/tests/test_harmony.py index ada0be7e74..9268956376 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -1,8 +1,10 @@ from __future__ import annotations from functools import partial +from itertools import pairwise from typing import TYPE_CHECKING +import numba import numpy as np import pandas as pd import pooch @@ -11,6 +13,12 @@ from scipy.stats import pearsonr from scanpy.preprocessing import harmony_integrate +from scanpy.preprocessing._harmony import core +from scanpy.preprocessing._harmony._kernels import ( + _apply_correction, + _cluster_pass, + _joint_stats, +) from scanpy.preprocessing._harmony.core import ( _SUPPRESS_PENALTY, _compute_lambda_kb, @@ -60,6 +68,8 @@ rng=734, alpha=0.2, batch_prune_threshold=1e-5, + seeding="kmeans", # harmonypy seeds with k-means + stop_rule="reference", # and records a round's objective only if clustering converged ) @@ -242,35 +252,16 @@ def test_harmony_batch_keys_are_nonempty_and_complete() -> None: _get_batch_codes(obs, []) -@pytest.mark.parametrize("dtype", [np.float32, np.float64]) -def test_harmony_multikey_correction_matches_dense_design( - dtype: type[np.floating], -) -> None: - """Compare the optimized correction with a direct dense ridge-regression solve. - - Production avoids materializing the one-hot design and exploits its marginal - structure. This independent oracle explicitly builds that design and solves - the weighted, regularized least-squares problem, catching errors in category - indexing, cross-covariate terms, and suppressed penalties. - """ - rng = np.random.default_rng(734) - levels = np.array([2, 3], dtype=np.int32) - offsets = np.array([0, 2], dtype=np.int32) - n_batches = int(levels.sum()) - n_cells, n_pcs, n_clusters = 31, 5, 3 - - local_codes = np.column_stack([ - rng.integers(0, level, size=n_cells) for level in levels - ]).astype(np.int32) - codes = local_codes + offsets - x = rng.normal(size=(n_cells, n_pcs)).astype(dtype) - r = rng.random(size=(n_cells, n_clusters)).astype(dtype) - r /= r.sum(axis=1, keepdims=True) - lambda_kb = rng.uniform(0.2, 1.0, size=(n_batches, n_clusters)).astype(dtype) - lambda_kb[1, 1] = dtype(_SUPPRESS_PENALTY) - - result = _correction_multi(x, codes, n_batches, r, lambda_kb=lambda_kb) - +def _dense_correction( + x: np.ndarray, + codes: np.ndarray, + n_batches: int, + r: np.ndarray, + lambda_kb: np.ndarray, +) -> np.ndarray: + """Ridge correction by an explicit one-hot design and a least-squares solve per cluster.""" + dtype = x.dtype.type + n_cells, n_pcs = x.shape design = np.zeros((n_cells, n_batches + 1), dtype=dtype) design[:, 0] = 1 for covariate in range(codes.shape[1]): @@ -303,6 +294,38 @@ def test_harmony_multikey_correction_matches_dense_design( )[0] w[0] = 0 expected -= r_k[:, None] * (design @ w) + return expected + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_harmony_multikey_correction_matches_dense_design( + dtype: type[np.floating], +) -> None: + """Compare the optimized correction with a direct dense ridge-regression solve. + + Production avoids materializing the one-hot design and exploits its marginal + structure. This independent oracle explicitly builds that design and solves + the weighted, regularized least-squares problem, catching errors in category + indexing, cross-covariate terms, and suppressed penalties. + """ + rng = np.random.default_rng(734) + levels = np.array([2, 3], dtype=np.int32) + offsets = np.array([0, 2], dtype=np.int32) + n_batches = int(levels.sum()) + n_cells, n_pcs, n_clusters = 31, 5, 3 + + local_codes = np.column_stack([ + rng.integers(0, level, size=n_cells) for level in levels + ]).astype(np.int32) + codes = local_codes + offsets + x = rng.normal(size=(n_cells, n_pcs)).astype(dtype) + r = rng.random(size=(n_cells, n_clusters)).astype(dtype) + r /= r.sum(axis=1, keepdims=True) + lambda_kb = rng.uniform(0.2, 1.0, size=(n_batches, n_clusters)).astype(dtype) + lambda_kb[1, 1] = dtype(_SUPPRESS_PENALTY) + + result = _correction_multi(x, codes, n_batches, r, lambda_kb=lambda_kb) + expected = _dense_correction(x, codes, n_batches, r, lambda_kb) atol = 2e-5 if dtype == np.float32 else 1e-11 np.testing.assert_allclose(result, expected, atol=atol, rtol=atol) @@ -311,7 +334,7 @@ def test_harmony_multikey_correction_matches_dense_design( def test_harmony_multikey_correction_singular_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A system diagnosed as singular falls back to least squares.""" + """A system the solver rejects as singular falls back to least squares.""" rng = np.random.default_rng(734) x = rng.normal(size=(12, 3)) codes = np.column_stack(( @@ -323,7 +346,11 @@ def test_harmony_multikey_correction_singular_fallback( lambda_kb = np.ones((5, 2)) expected = _correction_multi(x, codes, 5, r, lambda_kb=lambda_kb) - monkeypatch.setattr(np.linalg, "matrix_rank", lambda _: 0) + + def singular(*_args, **_kwargs): + raise np.linalg.LinAlgError + + monkeypatch.setattr(np.linalg, "solve", singular) result = _correction_multi(x, codes, 5, r, lambda_kb=lambda_kb) np.testing.assert_allclose(result, expected, atol=1e-12, rtol=1e-12) @@ -393,6 +420,57 @@ def test_harmony2_multikey_reference( assert _get_measure(reference, result, "L2").max() < 0.1 +def _within_group_r2(x: np.ndarray, group: np.ndarray, covariate: np.ndarray) -> float: + """Share of the within-group variance of ``x`` explained by ``covariate``.""" + total = explained = 0.0 + for g in np.unique(group): + xg, cg = x[group == g], covariate[group == g] + total += ((xg - xg.mean(0)) ** 2).sum() + explained += sum( + (cg == c).sum() * ((xg[cg == c].mean(0) - xg.mean(0)) ** 2).sum() + for c in np.unique(cg) + ) + return explained / total + + +def test_harmony2_multikey_default_seeding(adata_harmonypy2_multikey: AnnData) -> None: + """The default seeding integrates at least as well as harmonypy on its multikey fixture. + + It starts further from the k-means fixed point than harmonypy, so it runs more rounds + and lands farther from the stored output: compare the biology, not the coordinates. + """ + from sklearn.cluster import KMeans + from sklearn.metrics import silhouette_score + + adata = adata_harmonypy2_multikey.copy() + x = adata.obsm["X_pca"] + group = KMeans(4, n_init=10, random_state=0).fit_predict( + x / np.linalg.norm(x, axis=1, keepdims=True) + ) + assert np.all( + np.bincount(group) == 20 + ) # four groups of 20, each balanced over batch x sex + reference = adata.obsm["harmony2_ref_nclust4"] + + harmony_integrate( + adata, + ["batch", "sex"], + theta=[2.0, 0.1], + n_clusters=4, + max_iter_clustering=4, + tol_clustering=1e-3, + tol_harmony=1e-2, + rng=734, + ) + + result = adata.obsm["X_pca_harmony"] + batch = adata.obs["batch"].to_numpy() + assert _within_group_r2(result, group, batch) <= _within_group_r2( + reference, group, batch + ) + assert silhouette_score(result, group) >= silhouette_score(reference, group) + + def test_harmony_custom_parameters() -> None: """Test Harmony with custom parameters.""" adata = pbmc68k_reduced() @@ -581,3 +659,548 @@ def test_harmony2_ircolitis_reference( result = adata.obsm["X_pca_harmony"] assert _get_measure(reference, result, "r").min() > 0.95 assert _get_measure(reference, result, "L2").max() < 0.1 + + +# Kernels and driver of the numba implementation + +_N_BATCH, _K, _SIGMA = 3, 5, 0.1 +_DTYPES = [np.float32, np.float64] + + +def _toy( + n: int = 600, d: int = 8, seed: int = 0, dtype: DTypeLike = np.float32 +) -> tuple[np.ndarray, pd.DataFrame]: + """Four cell types, three batches with a shift each, and a second key without effect.""" + rng = np.random.default_rng(seed) + cell_type, batch = rng.integers(4, size=n), rng.integers(_N_BATCH, size=n) + x = ( + 3 * rng.normal(size=(4, d))[cell_type] + + rng.normal(size=(_N_BATCH, d))[batch] + + 0.5 * rng.normal(size=(n, d)) + ) + obs = pd.DataFrame({ + "batch": pd.Categorical(batch.astype(str)), + "half": pd.Categorical((np.arange(n) % 2).astype(str)), + }) + return x.astype(dtype), obs + + +def _state(dtype: DTypeLike, keys: str | list[str] = "batch", seed: int = 1): + """Build a random soft assignment with consistent counts.""" + x, obs = _toy(dtype=dtype) + z = core._normalize_rows_l2(x) + codes, n_levels = _get_batch_codes(obs, keys) + r = np.random.default_rng(seed).random((len(x), _K)).astype(dtype) ** 8 + r /= r.sum(axis=1, keepdims=True) + o = np.zeros((int(n_levels.sum()), _K), dtype=dtype) + for covariate in codes.T: + np.add.at(o, covariate, r) + n_b = np.bincount(codes.ravel()).astype(dtype) + return x, z, codes, n_levels, r, o, n_b, n_b / dtype(len(x)) + + +def _bounds(n: int, n_blocks: int = 19) -> np.ndarray: + sizes = [len(a) for a in np.array_split(np.arange(n), n_blocks)] + return np.concatenate([[0], np.cumsum(sizes)]).astype(np.int64) + + +def _clustering_iteration( # noqa: PLR0917 + z, codes, pr_b, r, o, theta, perm, *, stabilized, sum_penalties +): + """One clustering iteration in float64 NumPy, block by block, with either combination rule.""" + r, o = r.astype(np.float64), o.astype(np.float64) + y = core._normalize_rows_l2(r.T @ z).astype(np.float64) + r_sum = r.sum(axis=0) + for idx in np.array_split(perm, 19): + for covariate in codes[idx].T: + np.add.at(o, covariate, -r[idx]) + r_sum -= r[idx].sum(axis=0) + e = pr_b[:, None] * r_sum + factor = ((e + 1) / (o + e + 1 if stabilized else o + 1)) ** theta[:, None] + per_key = factor[codes[idx]] # cells x keys x K + penalty = per_key.sum(axis=1) if sum_penalties else per_key.prod(axis=1) + new = np.exp(-2 / _SIGMA * (1 - z[idx] @ y.T)) * penalty + r[idx] = new / new.sum(axis=1, keepdims=True) + for covariate in codes[idx].T: + np.add.at(o, covariate, r[idx]) + r_sum += r[idx].sum(axis=0) + return r, o + + +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize( + ("keys", "stabilized", "n_chunks", "scatter", "sum_penalties"), + [ + ("batch", True, 32, False, False), + (["batch", "half"], False, 3, True, False), + ("batch", True, 3, False, False), + (["batch", "half"], False, 32, True, True), + (["batch", "half"], True, 32, False, True), + ], +) +def test_cluster_pass_matches_numpy( + dtype, keys, *, stabilized, n_chunks, scatter, sum_penalties +) -> None: + x, z, codes, _, r0, o0, _, pr_b = _state(dtype, keys) + theta = np.full(o0.shape[0], 2, dtype=dtype) + perm = np.random.default_rng(7).permutation(len(x)) + kw = dict(stabilized=stabilized, sum_penalties=sum_penalties) + r_ref, o_ref = _clustering_iteration(z, codes, pr_b, r0, o0, theta, perm, **kw) + + y = core._normalize_rows_l2(r0.T @ z) + r, o, r_sum = r0.copy(), o0.copy(), r0.sum(axis=0) + dist, ent, o_fresh, y_t = _cluster_pass( + *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b, theta, perm), + *(_bounds(len(x)), dtype(-2 / _SIGMA), stabilized, n_chunks, scatter), + sum_penalties, + ) + atol = 2e-5 if dtype == np.float32 else 1e-10 + np.testing.assert_allclose(r, r_ref, atol=atol) + np.testing.assert_allclose(o, o_ref, atol=50 * atol) + # objective parts, summed on the fly + assert dist == pytest.approx(np.sum(r_ref * 2 * (1 - z @ y.T)), rel=1e-4) + assert ent == pytest.approx(np.sum(r_ref * np.log(r_ref + 1e-300)), rel=1e-4) + # invariants + assert r.dtype == o.dtype == o_fresh.dtype == y_t.dtype == dtype + np.testing.assert_allclose(r.sum(axis=1), 1, atol=1e-5) + np.testing.assert_allclose(r_sum, r.sum(axis=0), rtol=1e-4) + exact = np.zeros_like(o) + for covariate in codes.T: + np.add.at(exact, covariate, r) + np.testing.assert_allclose(o_fresh, exact, rtol=1e-4, atol=1e-5) + np.testing.assert_allclose(y_t.T, core._normalize_rows_l2(r.T @ z), atol=1e-5) + if isinstance(keys, list): # the two rules really differ on this input + other, _ = _clustering_iteration( + *(z, codes, pr_b, r0, o0, theta, perm), + stabilized=stabilized, + sum_penalties=not sum_penalties, + ) + assert np.abs(r_ref - other).max() > 1e-3 + + +def test_cluster_pass_without_penalty_is_soft_kmeans() -> None: + x, z, codes, _, r0, _, _, pr_b = _state(np.float32) + y = core._normalize_rows_l2(r0.T @ z) + r = np.zeros_like(r0) + o, r_sum = np.zeros((_N_BATCH, _K), np.float32), np.zeros(_K, np.float32) + n = len(x) + _cluster_pass( + *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b), + *(np.zeros(_N_BATCH, np.float32), np.arange(n), np.array([0, n])), + *(np.float32(-20), True, 3, False, False), + ) + expected = np.exp(-20 * (1 - z @ y.T).astype(np.float64)) + expected /= expected.sum(axis=1, keepdims=True) + np.testing.assert_allclose(r, expected, atol=1e-5) + # sum_k o[b, k] = N_b + np.testing.assert_allclose(o.sum(axis=1), np.bincount(codes[:, 0]), rtol=1e-5) + + +def _random_design(rng: np.random.Generator, n: int, sizes: tuple[int, ...]): + """Random levels for several keys and a random soft assignment, in float64.""" + obs = pd.DataFrame({ + f"key{c}": pd.Categorical( + rng.integers(size, size=n).astype(str), + categories=[str(v) for v in range(size)], + ) + for c, size in enumerate(sizes) + }) + codes, n_levels = _get_batch_codes(obs, list(obs.columns)) + x = rng.normal(size=(n, 6)) + r = rng.random((n, _K)) ** 4 + r /= r.sum(axis=1, keepdims=True) + return codes, n_levels, x, r + + +@pytest.mark.parametrize("dtype", _DTYPES) +# big block first / last / a size-1 key / one key +@pytest.mark.parametrize("sizes", [(4, 9, 3), (4, 3, 9), (1, 8), (7,)]) +def test_correction_with_pruning_and_an_empty_cluster(dtype, sizes) -> None: + rng = np.random.default_rng(sum(sizes)) + codes, n_levels, x, r = _random_design(rng, 500, sizes) + r[:, 3] = 0 # an empty cluster + r /= r.sum(axis=1, keepdims=True) + n_batches = int(n_levels.sum()) + lam = rng.uniform(0.2, 2.0, size=(n_batches, _K)) + # pruned pairs in every block, several per cluster + lam[rng.random((n_batches, _K)) < 0.25] = _SUPPRESS_PENALTY + lam[:, 1] = _SUPPRESS_PENALTY # everything pruned in one cluster + lam[: n_levels[0], 2] = _SUPPRESS_PENALTY # a whole key pruned in another + + result = _correction_multi( + x.astype(dtype), codes, n_batches, r.astype(dtype), lambda_kb=lam.astype(dtype) + ) + expected = _dense_correction(x, codes, n_batches, r, lam) + assert result.dtype == dtype + np.testing.assert_allclose( + result, expected, atol=5e-5 if dtype == np.float32 else 1e-10 + ) + + +def test_correction_with_nested_keys_and_unused_levels() -> None: + rng = np.random.default_rng(1) + n = 600 + sample = rng.integers(12, size=n) + obs = pd.DataFrame({ + # donors 4 and 5 never occur + "donor": pd.Categorical( + (sample // 3).astype(str), categories=[str(v) for v in range(6)] + ), + "sample": pd.Categorical(sample.astype(str)), + }) + codes, _ = _get_batch_codes(obs, ["donor", "sample"]) + x = rng.normal(size=(n, 6)) + r = rng.random((n, _K)) ** 4 + r /= r.sum(axis=1, keepdims=True) + lam = np.full((18, _K), 0.7) + lam[5] = ( + _SUPPRESS_PENALTY # one unused level pruned, the other left with a finite penalty + ) + result = _correction_multi(x, codes, 18, r, lambda_kb=lam) + np.testing.assert_allclose( + result, _dense_correction(x, codes, 18, r, lam), atol=1e-10 + ) + + +@pytest.mark.parametrize("dtype", _DTYPES) +def test_apply_correction_normalises_and_returns_centroids(dtype) -> None: + x, _, codes, n_levels, r, o, n_b, pr_b = _state(dtype, ["batch", "half"]) + lam = _compute_lambda_kb( + pr_b[:, None] * r.sum(axis=0), + o=o, + n_b=n_b, + alpha=0.2, + threshold=1e-5, + ridge_lambda=1.0, + dynamic_lambda=True, + ) + joint, levels = core._joint_categories(codes) + s, o_joint = _joint_stats( + x, r, joint, np.argsort(joint, kind="stable"), len(levels), 3 + ) + w = core._solve_correction(s, o_joint, levels, n_levels, lam, _SUPPRESS_PENALTY) + z_hat, z_norm = np.empty_like(x), np.empty_like(x) + y_t = _apply_correction(x, r, joint, w, z_hat, z_norm, 3) + assert z_hat.dtype == z_norm.dtype == y_t.dtype == dtype + np.testing.assert_allclose( + z_hat, + _correction_multi(x, codes, int(n_levels.sum()), r, lambda_kb=lam), + atol=1e-5, + ) + np.testing.assert_allclose(np.linalg.norm(z_norm, axis=1), 1, atol=1e-5) + np.testing.assert_allclose(y_t.T, core._normalize_rows_l2(r.T @ z_norm), atol=1e-5) + + +@pytest.mark.parametrize("dtype", _DTYPES) +# normal, many categories across chunks, more chunks than cells (many empty) +@pytest.mark.parametrize(("n_chunks", "n_joint"), [(3, 7), (32, 300), (1000, 2)]) +def test_joint_stats_equal_direct_sums(dtype, n_chunks: int, n_joint: int) -> None: + rng = np.random.default_rng(n_chunks + n_joint) + n = 400 + x = rng.normal(size=(n, 5)).astype(dtype) + r = rng.random((n, _K)).astype(dtype) + joint = rng.integers(n_joint, size=n).astype(np.int32) + joint[:150] = 0 # one large category that spans many chunks + s, o = _joint_stats( + x, r, joint, np.argsort(joint, kind="stable"), n_joint, n_chunks + ) + s_ref, o_ref = np.zeros((n_joint, _K, 5)), np.zeros((n_joint, _K)) + np.add.at(o_ref, joint, r.astype(np.float64)) + np.add.at(s_ref, joint, r[:, :, None].astype(np.float64) * x[:, None, :]) + assert s.dtype == o.dtype == dtype + tol = 1e-4 if dtype == np.float32 else 1e-10 + np.testing.assert_allclose(o, o_ref, rtol=tol, atol=tol) + np.testing.assert_allclose(s, s_ref, rtol=tol, atol=tol) + + +def test_joint_categories() -> None: + rng = np.random.default_rng(0) + for n_keys, n_levels in ((1, 5), (3, 4), (16, 2), (40, 3)): + codes = rng.integers(n_levels, size=(3000, n_keys)) + n_levels * np.arange( + n_keys + ) + joint, levels = core._joint_categories(codes.astype(np.int32)) + assert joint.dtype == np.int32 + assert joint.flags.c_contiguous + np.testing.assert_array_equal(levels[joint], codes) + assert len(levels) == len(np.unique(codes, axis=0)) == joint.max() + 1 + + +@pytest.mark.parametrize("keys", ["batch", ["batch", "half"]]) +def test_harmonize_float32_matches_float64(keys) -> None: + x, obs = _toy(n=1500, dtype=np.float64) + kw = dict(seeding="kmeans++", rng=0) # the same seeds in both precisions + got64, info = core.harmonize(x, obs, keys, **kw) + got32, _ = core.harmonize(x.astype(np.float32), obs, keys, **kw) + assert got64.dtype == np.float64 + assert got32.dtype == np.float32 + assert np.abs(got64 - x).max() > 0.1 # it did correct something + assert len(info["iterations"]) <= 10 + r = [np.corrcoef(got32[:, j], got64[:, j])[0, 1] for j in range(x.shape[1])] + assert min(r) > 0.9999 + rel = np.linalg.norm(got32 - got64, axis=0) / np.linalg.norm(got64, axis=0) + assert rel.max() < 1e-3 + + +def test_harmonize_many_small_batches() -> None: + """A design the rule sends through the scatter path agrees across precisions.""" + assert not core._use_scatter(166, 1, 584_944 // 19) # atlas-like: accumulators + assert core._use_scatter( + 4000, 1, 20_000 // 19 + ) # thousands of tiny batches: scatters + assert core._use_scatter(11, 1, 10 * core.N_CHUNKS) + assert not core._use_scatter(10, 1, 10 * core.N_CHUNKS) + x, obs = _toy(n=1500, dtype=np.float64) + obs = obs.assign(well=pd.Categorical(np.arange(1500) % 250)) + assert core._use_scatter(250, 1, 1500 // 19) + got64, _ = core.harmonize(x, obs, "well", seeding="kmeans++", rng=0) + got32, _ = core.harmonize( + x.astype(np.float32), obs, "well", seeding="kmeans++", rng=0 + ) + rel = np.linalg.norm(got32 - got64, axis=0) / np.linalg.norm(got64, axis=0) + assert rel.max() < 1e-3 + + +def test_harmonize_stop_rules() -> None: + assert not core._converged([10.0], 1e-2, window=1, signed=False) + assert core._converged([10.0, 10.05], 1e-2, window=1, signed=False) + assert not core._converged([10.0, 11.0], 1e-2, window=1, signed=False) + # a rise passes the signed test + assert core._converged([10.0, 11.0], 1e-2, window=1, signed=True) + # needs window + 1 values + assert not core._converged([10.0, 9.0, 8.0], 1e-2, window=3, signed=True) + x, obs = _toy() + for rule in ("reference_record", "paper", "paper_window3"): + _, info = core.harmonize(x, obs, "batch", stop_rule=rule, rng=0) + objs = info["objectives"] + assert len(objs) == len(info["iterations"]) + 1 # every round recorded + assert max(info["iterations"]) <= 200 + tol = core.STOP_RULES[rule][2]["harmony2"][2] + passed = [old - new < tol * abs(old) for old, new in pairwise(objs)] + assert not any(passed[:-1]) # no earlier round met the stop test + # the last one did, or the cap ended the run + assert passed[-1] or len(info["iterations"]) == 10 + + +def test_harmonize_reference_rule_records_only_converged_rounds() -> None: + x, obs = _toy() + # a window of 3 can never pass with one iteration per round + kw = dict(max_iter_clustering=1, rng=0) + _, info = core.harmonize(x, obs, "batch", stop_rule="reference", **kw) + assert info["iterations"] == [1] * 10 + assert len(info["objectives"]) == 1 # only the initial objective + _, info = core.harmonize(x, obs, "batch", **kw) + assert len(info["objectives"]) == len(info["iterations"]) + 1 + + +@pytest.mark.parametrize("seeding", ["kmeans++", "afkmc2", "paper_nearest"]) +def test_harmonize_seedings_are_reproducible(seeding: str) -> None: + x, obs = _toy() + out = [core.harmonize(x, obs, "batch", seeding=seeding, rng=3)[0] for _ in range(2)] + np.testing.assert_array_equal(out[0], out[1]) + assert np.isfinite(out[0]).all() + + +def test_harmonize_edge_inputs() -> None: + x, obs = _toy() + for extra in (dict(tau=5), dict(stop_rule="paper_window3"), dict(n_clusters=2)): + assert np.isfinite(core.harmonize(x, obs, "batch", rng=0, **extra)[0]).all() + # a one-cell batch + obs1 = obs.assign(batch=pd.Categorical(["a"] * (len(x) - 1) + ["b"])) + assert np.isfinite(core.harmonize(x, obs1, "batch", rng=0)[0]).all() + # fewer cells per block than chunks + assert np.isfinite(core.harmonize(x[:40], obs[:40], "batch", rng=0)[0]).all() + + +def test_harmonize_does_not_depend_on_thread_count() -> None: + x, obs = _toy() + previous = numba.get_num_threads() + try: + out = [] + for n_threads in (1, 2): + numba.set_num_threads(n_threads) + # afkmc2 is bit-reproducible; scikit-learn's KMeans is not + out.append(core.harmonize(x, obs, "batch", rng=0)[0]) + finally: + numba.set_num_threads(previous) + np.testing.assert_array_equal(out[0], out[1]) + + +@pytest.mark.parametrize("keys", ["batch", ["batch", "half"]]) +def test_harmonize_low_memory_is_identical(keys) -> None: + x, obs = _toy(n=1500) + full = core.harmonize(x, obs, keys, rng=0)[0] + lean = core.harmonize(x, obs, keys, low_memory=True, rng=0)[0] + np.testing.assert_array_equal(full, lean) + + +def test_harmonize_leaves_out_all_zero_cells() -> None: + x, obs = _toy() + x0 = x.copy() + x0[[3, 77]] = 0 + with pytest.warns(UserWarning, match="2 cells have an all-zero embedding"): + out, _ = core.harmonize(x0, obs, "batch", rng=0) + keep = np.ones(len(x), dtype=bool) + keep[[3, 77]] = False + np.testing.assert_array_equal(out[~keep], 0) + # same as running without them + expected, _ = core.harmonize(x[keep], obs[keep], "batch", rng=0) + np.testing.assert_array_equal(out[keep], expected) + + +def test_harmonize_four_keys_and_a_level_emptied_by_the_zero_cell_filter() -> None: + x, obs = _toy() + obs = obs.assign( + third=pd.Categorical(np.arange(len(x)) % 3), + lone=pd.Categorical(["only"] + ["rest"] * (len(x) - 1)), + ) + x = x.copy() + x[0] = 0 # the one cell of level "only" is all-zero and gets filtered + with pytest.warns(UserWarning, match="1 cells have an all-zero embedding"): + out, _ = core.harmonize(x, obs, ["batch", "third", "half", "lone"], rng=0) + assert np.isfinite(out).all() + assert np.abs(out - x).max() > 0.1 + np.testing.assert_array_equal(out[0], 0) + + +def test_harmonize_tau_and_recompute_counts() -> None: + x, obs = _toy() + + def run(**extra): + return core.harmonize(x, obs, "batch", rng=0, **extra)[0] + + np.testing.assert_array_equal(run(tau=1e-9), run()) # discount factor 1 + np.testing.assert_allclose(run(tau=1e9), run(theta=0.0), atol=1e-4) # factor 0 + assert not np.array_equal(run(tau=10), run()) # in between it changes the result + incremental = run(recompute_counts=False) + assert not np.array_equal(incremental, run()) # the flag takes effect ... + np.testing.assert_allclose(incremental, run(), atol=1e-2) # ... at rounding level + + +def test_harmonize_penalty_combine() -> None: + x, obs = _toy() + one_key = [ + core.harmonize(x, obs, "batch", penalty_combine=rule, rng=0)[0] + for rule in ("product", "sum") + ] + np.testing.assert_array_equal(one_key[0], one_key[1]) # one key: nothing to combine + two_keys = [ + core.harmonize(x, obs, ["batch", "half"], penalty_combine=rule, rng=0)[0] + for rule in ("product", "sum") + ] + assert np.isfinite(two_keys[1]).all() + assert not np.array_equal(two_keys[0], two_keys[1]) + + +def test_harmonize_seeding_options() -> None: + x, obs = _toy() + kw = dict(seeding="afkmc2", seeding_lloyd_steps=3, rng=0) + plain = core.harmonize(x, obs, "batch", seeding="afkmc2", rng=0)[0] + polished = core.harmonize(x, obs, "batch", **kw)[0] + assert not np.array_equal(polished, plain) # the polish takes effect + # a cap above the number of cells changes nothing; a binding cap does + uncapped = core.harmonize(x, obs, "batch", seeding_cells_per_cluster=10**6, **kw)[0] + np.testing.assert_array_equal(uncapped, polished) + capped = core.harmonize(x, obs, "batch", seeding_cells_per_cluster=10, **kw)[0] + assert np.isfinite(capped).all() + assert not np.array_equal(capped, polished) + # subsampled KMeans + km = core.harmonize( + x, obs, "batch", seeding="kmeans", seeding_cells_per_cluster=10, rng=0 + ) + assert np.isfinite(km[0]).all() + # without Lloyd steps nothing uses the subsample, so a cap changes nothing + np.testing.assert_array_equal( + core.harmonize(x, obs, "batch", seeding_cells_per_cluster=10, rng=0)[0], + core.harmonize(x, obs, "batch", rng=0)[0], + ) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + (dict(alpha=float("nan")), "alpha"), + (dict(flavor="harmony1", ridge_lambda=0.0), "ridge_lambda"), + (dict(block_proportion=2.0), "block_proportion"), + (dict(sigma=0.0), "sigma"), + (dict(max_iter_harmony=0), "max_iter_harmony"), + (dict(seeding="nope"), "seeding must be one of"), + (dict(penalty_combine="mean"), "penalty_combine must be one of"), + (dict(stop_rule="never"), "stop_rule must be one of"), + ], +) +def test_harmonize_rejects_bad_parameters(kwargs: dict, match: str) -> None: + x, obs = _toy() + with pytest.raises(ValueError, match=match): + core.harmonize(x, obs, "batch", **kwargs) + + +def test_harmonize_rejects_bad_inputs() -> None: + x, obs = _toy() + with pytest.raises(ValueError, match="rows but x has"): + core.harmonize(x, obs[:500], "batch") + with pytest.raises(ValueError, match="rows but x has"): + core.harmonize(x[:500], obs, "batch") + with pytest.raises(ValueError, match="2-dimensional"): + core.harmonize(x[:, 0], obs, "batch") + x0 = x.copy() + x0[0, 0] = np.inf + with pytest.raises(ValueError, match="non-finite"): + core.harmonize(x0, obs, "batch") + with ( + pytest.raises(ValueError, match="non-zero embedding"), + pytest.warns(UserWarning, match="all-zero embedding"), + ): + core.harmonize(np.zeros_like(x), obs, "batch") + + +def test_block_bounds_match_array_split() -> None: + for n in (7, 40, 600, 1001): + n_blocks = int(min(n, 1 // 0.05)) + np.testing.assert_array_equal(core._block_bounds(n, 0.05), _bounds(n, n_blocks)) + + +def _closest_pair(y: np.ndarray) -> float: + d = np.linalg.norm(y[:, None] - y[None], axis=2) + return d[np.triu_indices(len(y), 1)].min() + + +# squared-distance weights (k-means++, AFK-MC²) spread seeds more than plain-distance weights +@pytest.mark.parametrize( + ("seeder", "factor"), + [("_seed_kmeanspp", 1.5), ("_seed_afkmc2", 1.5), ("_seed_paper_nearest", 1.2)], +) +def test_distance_weighted_seeders_spread_seeds(seeder: str, factor: float) -> None: + _, z, *_ = _state(np.float32) + gaps, uniform = [], [] + for seed in range(20): + rng = np.random.default_rng(seed) + y = getattr(core, seeder)(z, 10, rng) + assert y.shape == (10, z.shape[1]) + assert y.dtype == np.float32 + np.testing.assert_allclose(np.linalg.norm(y, axis=1), 1, atol=1e-5) + gaps.append(_closest_pair(y)) + uniform.append(_closest_pair(z[rng.choice(len(z), 10, replace=False)])) + # further apart than uniformly drawn cells + assert np.mean(gaps) > factor * np.mean(uniform) + + +def test_seeders_survive_identical_cells() -> None: + z = np.tile(np.array([[1.0, 0.0, 0.0]], dtype=np.float32), (50, 1)) + for seeder in (core._seed_kmeanspp, core._seed_afkmc2, core._seed_paper_nearest): + assert np.isfinite(seeder(z, 5, np.random.default_rng(0))).all() + + +@pytest.mark.parametrize("dtype", _DTYPES) +def test_lloyd_polish_lowers_the_kmeans_cost(dtype) -> None: + _, z, *_ = _state(dtype) + y = core._seed_afkmc2(z, _K, np.random.default_rng(0)) + + def cost(c): + return float(np.sum(1 - (z @ c.T).max(axis=1))) + + polished = core._lloyd_spherical(z, y, 3) + assert cost(polished) < cost(y) + assert cost(core._lloyd_spherical(z, polished, 1)) <= cost(polished) + 1e-4 + np.testing.assert_allclose(np.linalg.norm(polished, axis=1), 1, atol=1e-5) + assert polished.dtype == dtype From 02a88c33889f2f84c428b00433c07d8ed5e07f14 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 01:30:55 +0200 Subject: [PATCH 02/10] docs: release note for #4424 --- docs/release-notes/4424.feat.md | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 docs/release-notes/4424.feat.md diff --git a/docs/release-notes/4424.feat.md b/docs/release-notes/4424.feat.md new file mode 100644 index 0000000000..fee8285f19 --- /dev/null +++ b/docs/release-notes/4424.feat.md @@ -0,0 +1,2 @@ +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. +The default seeding is now `"afkmc2"` and the default stop rule `"reference_record"`; pass `seeding="kmeans", stop_rule="reference"` to reproduce previous results {smaller}`T Treis` From 254aee85670e212eae8754928b8e9c2c1ae33837 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 03:23:02 +0200 Subject: [PATCH 03/10] fix: validate n_clusters, dtype and flavor in harmonize n_clusters=0 reached the kernels with K=0 and returned the input unchanged; n_clusters above the number of cells ran with duplicate centroids. Other float types failed inside numba with an opaque typing error. --- src/scanpy/preprocessing/_harmony/core.py | 16 ++++++++++++++-- tests/test_harmony.py | 6 ++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/scanpy/preprocessing/_harmony/core.py b/src/scanpy/preprocessing/_harmony/core.py index d612124d8a..97c92f2c5b 100644 --- a/src/scanpy/preprocessing/_harmony/core.py +++ b/src/scanpy/preprocessing/_harmony/core.py @@ -34,7 +34,7 @@ _REFERENCE = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (4, 1e-3, 1e-2)} _PAPER = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (200, 1e-5, 1e-4)} STOP_RULES = { - # the round objective is only recorded if clustering converged, as in harmonypy + # the round objective is only recorded if clustering converged "reference": (3, False, _REFERENCE), # every round's objective is recorded, as in rapids-singlecell "reference_record": (3, True, _REFERENCE), @@ -80,11 +80,17 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 """ stabilized = dynamic_lambda = flavor == "harmony2" choices = dict( + flavor=("harmony1", "harmony2"), stop_rule=tuple(STOP_RULES), seeding=SEEDINGS, penalty_combine=("product", "sum"), ) - given = dict(stop_rule=stop_rule, seeding=seeding, penalty_combine=penalty_combine) + given = dict( + flavor=flavor, + stop_rule=stop_rule, + seeding=seeding, + penalty_combine=penalty_combine, + ) for name, allowed in choices.items(): if given[name] not in allowed: msg = f"{name} must be one of {allowed}, got {given[name]!r}." @@ -115,6 +121,9 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 if not ok: raise ValueError(msg) x = np.ascontiguousarray(x) + if x.dtype not in (np.float32, np.float64): + msg = f"dtype must be float32 or float64, got {x.dtype}." + raise ValueError(msg) if x.ndim != 2: msg = f"x must be a 2-dimensional array (cells x components), got {x.ndim} dimension(s)." raise ValueError(msg) @@ -156,6 +165,9 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 pr_b = n_b / ft(n_cells) theta_b = _get_theta_array(theta, n_levels, x.dtype).ravel() k = max(2, int(min(100, n_cells / 30))) if n_clusters is None else n_clusters + if not 1 <= k <= n_cells: + msg = f"n_clusters must be between 1 and the number of cells ({n_cells}), got {k}." + raise ValueError(msg) if tau > 0: theta_b = (theta_b * (1 - np.exp(-n_b / (k * tau)) ** 2)).astype(ft) diff --git a/tests/test_harmony.py b/tests/test_harmony.py index 9268956376..0b4f97c5f2 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -1127,6 +1127,9 @@ def test_harmonize_seeding_options() -> None: (dict(seeding="nope"), "seeding must be one of"), (dict(penalty_combine="mean"), "penalty_combine must be one of"), (dict(stop_rule="never"), "stop_rule must be one of"), + (dict(flavor="harmony3"), "flavor must be one of"), + (dict(n_clusters=0), "n_clusters must be between 1"), + (dict(n_clusters=601), "n_clusters must be between 1"), ], ) def test_harmonize_rejects_bad_parameters(kwargs: dict, match: str) -> None: @@ -1143,6 +1146,9 @@ def test_harmonize_rejects_bad_inputs() -> None: core.harmonize(x[:500], obs, "batch") with pytest.raises(ValueError, match="2-dimensional"): core.harmonize(x[:, 0], obs, "batch") + for dtype in (np.float16, np.int64): + with pytest.raises(ValueError, match="dtype must be float32 or float64"): + core.harmonize(x.astype(dtype), obs, "batch") x0 = x.copy() x0[0, 0] = np.inf with pytest.raises(ValueError, match="non-finite"): From b0f570b3ba24d7cd8192a23877c3e144ee24d12e Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 03:28:48 +0200 Subject: [PATCH 04/10] fix: thread-independent objective and exact summed penalties numba turned the final sums of the per-chunk objective parts into a parallel reduction, so float64 objectives (and with them the stop tests) could depend on the thread count. Sum them serially. penalty_combine='sum' capped the per-level factors at exp(80) and counted keys with theta=0 as a factor of 1. Combine the log factors with a shifted log-sum-exp over the keys with theta > 0 instead. --- src/scanpy/preprocessing/_harmony/_kernels.py | 36 ++++++----- tests/test_harmony.py | 59 ++++++++++++++++--- 2 files changed, 73 insertions(+), 22 deletions(-) diff --git a/src/scanpy/preprocessing/_harmony/_kernels.py b/src/scanpy/preprocessing/_harmony/_kernels.py index 9d7355cbc2..6f73c3ecca 100644 --- a/src/scanpy/preprocessing/_harmony/_kernels.py +++ b/src/scanpy/preprocessing/_harmony/_kernels.py @@ -44,8 +44,8 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 block (work per block ~ cells * K, best for thousands of small levels). Same result up to rounding. ``sum_penalties`` selects how the penalty factors of a cell's several batch keys are combined: ``False`` - multiplies them (scanpy, rapids-singlecell, [M1 p.9]); ``True`` adds them, as the assignment rule is printed - in [V2 p.10]. With one key the two are the same. + multiplies them :cite:p:`Korsunsky2019`; ``True`` adds them, as the assignment rule is printed in + :cite:t:`Patikas2026`, skipping levels with ``theta == 0``. With one key the two are the same. """ ft = z_norm.dtype.type n_pc = z_norm.shape[1] @@ -54,7 +54,7 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 one = ft(1.0) inv_cov = one / ft(n_cov) combine_sum = sum_penalties and n_cov > 1 - tiny = ft(1e-38) + neg_inf = ft(-np.inf) logpen = np.empty((n_levels, n_clusters), dtype=z_norm.dtype) acc = np.empty( (1 if scatter else n_chunks, n_levels, n_clusters), dtype=z_norm.dtype @@ -102,15 +102,12 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 e = pr_b[b] * r_sum[k] den = o[b, k] + e + one if stabilized else o[b, k] + one logpen[b, k] = theta[b] * (np.log(e + one) - np.log(den)) - if combine_sum: # the table then holds the factors; cap the exponent so a huge e cannot overflow to inf - logpen[b, k] = np.exp(min(logpen[b, k], ft(80.0))) # 3. new assignments; cells of a block are independent given the table if not scatter: acc[:] = 0 for t in prange(n_chunks): dots = np.empty(n_clusters, dtype=z_norm.dtype) logit = np.empty(n_clusters, dtype=z_norm.dtype) - factor = np.empty(n_clusters, dtype=z_norm.dtype) y_loc = y_acc[t] dist_t = 0.0 ent_t = 0.0 @@ -125,14 +122,20 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 for k in range(n_clusters): logit[k] = term * (one - dots[k]) if combine_sum: + # log of the summed factors of the keys with theta > 0, shifted by their maximum for k in range(n_clusters): - factor[k] = 0 - for c in range(n_cov): - b = codes[i, c] - for k in range(n_clusters): - factor[k] += logpen[b, k] - for k in range(n_clusters): - logit[k] += np.log(max(factor[k], tiny)) + top_pen = neg_inf + for c in range(n_cov): + b = codes[i, c] + if theta[b] != 0: + top_pen = max(top_pen, logpen[b, k]) + if top_pen > neg_inf: + summed = ft(0.0) + for c in range(n_cov): + b = codes[i, c] + if theta[b] != 0: + summed += np.exp(logpen[b, k] - top_pen) + logit[k] += top_pen + np.log(summed) else: for c in range(n_cov): b = codes[i, c] @@ -199,7 +202,12 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 scale = one / max(np.sqrt(sq), ft(1e-12)) for dd in range(n_pc): y_t_next[dd, k] *= scale - return dist.sum(), ent.sum(), o_fresh, y_t_next + # serial sums: a reduction over the parallel chunks would depend on the thread count + dist_total = ent_total = 0.0 + for t in range(n_chunks): + dist_total += dist[t] + ent_total += ent[t] + return dist_total, ent_total, o_fresh, y_t_next @njit diff --git a/tests/test_harmony.py b/tests/test_harmony.py index 0b4f97c5f2..aba54036da 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -716,10 +716,15 @@ def _clustering_iteration( # noqa: PLR0917 np.add.at(o, covariate, -r[idx]) r_sum -= r[idx].sum(axis=0) e = pr_b[:, None] * r_sum - factor = ((e + 1) / (o + e + 1 if stabilized else o + 1)) ** theta[:, None] - per_key = factor[codes[idx]] # cells x keys x K - penalty = per_key.sum(axis=1) if sum_penalties else per_key.prod(axis=1) - new = np.exp(-2 / _SIGMA * (1 - z[idx] @ y.T)) * penalty + den = o + e + 1 if stabilized else o + 1 + log_factor = theta[:, None] * (np.log(e + 1) - np.log(den)) + per_key = log_factor[codes[idx]] # cells x keys x K + if sum_penalties: + log_penalty = np.logaddexp.reduce(per_key, axis=1) + else: + log_penalty = per_key.sum(axis=1) + logit = -2 / _SIGMA * (1 - z[idx] @ y.T) + log_penalty + new = np.exp(logit - logit.max(axis=1, keepdims=True)) r[idx] = new / new.sum(axis=1, keepdims=True) for covariate in codes[idx].T: np.add.at(o, covariate, r[idx]) @@ -778,6 +783,30 @@ def test_cluster_pass_matches_numpy( assert np.abs(r_ref - other).max() > 1e-3 +@pytest.mark.parametrize("dtype", _DTYPES) +def test_cluster_pass_sum_rule_handles_large_penalties(dtype) -> None: + """Summed penalty factors far above exp(80) keep their order between clusters.""" + x, z, codes, _, r0, o0, _, pr_b = _state(dtype, ["batch", "half"]) + r0[codes[:, 0] == 0, :2] = 0 # batch 0 absent from clusters 0 and 1 + r0[:, 1] *= 0.3 # with different expected counts + r0 /= r0.sum(axis=1, keepdims=True) + o0[:] = 0 + for covariate in codes.T: + np.add.at(o0, covariate, r0) + theta = np.full(o0.shape[0], 30, dtype=dtype) + perm = np.random.default_rng(7).permutation(len(x)) + r_ref, _ = _clustering_iteration( + *(z, codes, pr_b, r0, o0, theta, perm), stabilized=False, sum_penalties=True + ) + y = core._normalize_rows_l2(r0.T @ z) + r, o, r_sum = r0.copy(), o0.copy(), r0.sum(axis=0) + _cluster_pass( + *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b, theta, perm), + *(_bounds(len(x)), dtype(-2 / _SIGMA), False, 32, False, True), + ) + np.testing.assert_allclose(r, r_ref, atol=5e-5 if dtype == np.float32 else 1e-10) + + def test_cluster_pass_without_penalty_is_soft_kmeans() -> None: x, z, codes, _, r0, _, _, pr_b = _state(np.float32) y = core._normalize_rows_l2(r0.T @ z) @@ -1012,18 +1041,23 @@ def test_harmonize_edge_inputs() -> None: assert np.isfinite(core.harmonize(x[:40], obs[:40], "batch", rng=0)[0]).all() -def test_harmonize_does_not_depend_on_thread_count() -> None: - x, obs = _toy() +@pytest.mark.skipif( + numba.config.NUMBA_NUM_THREADS < 2, reason="needs at least two numba threads" +) +@pytest.mark.parametrize("dtype", _DTYPES) +def test_harmonize_does_not_depend_on_thread_count(dtype) -> None: + x, obs = _toy(dtype=dtype) previous = numba.get_num_threads() try: out = [] for n_threads in (1, 2): numba.set_num_threads(n_threads) # afkmc2 is bit-reproducible; scikit-learn's KMeans is not - out.append(core.harmonize(x, obs, "batch", rng=0)[0]) + out.append(core.harmonize(x, obs, ["batch", "half"], rng=0)) finally: numba.set_num_threads(previous) - np.testing.assert_array_equal(out[0], out[1]) + np.testing.assert_array_equal(out[0][0], out[1][0]) + assert out[0][1]["objectives"] == out[1][1]["objectives"] @pytest.mark.parametrize("keys", ["batch", ["batch", "half"]]) @@ -1090,6 +1124,15 @@ def test_harmonize_penalty_combine() -> None: ] assert np.isfinite(two_keys[1]).all() assert not np.array_equal(two_keys[0], two_keys[1]) + # a key with theta = 0 is switched off under either rule, so with the other key alone + # the sum equals the product + off = [ + core.harmonize( + x, obs, ["batch", "half"], theta=[2.0, 0.0], penalty_combine=rule, rng=0 + )[0] + for rule in ("product", "sum") + ] + np.testing.assert_allclose(off[0], off[1], atol=1e-5) def test_harmonize_seeding_options() -> None: From f27095b8e13fab091c8043bf2c5bf38e4ba00dc6 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 03:33:31 +0200 Subject: [PATCH 05/10] fix: keep k-means seeding for harmony1 seeding now defaults to None: afkmc2 for harmony2, kmeans for harmony1, so flavor='harmony1' still runs the original initialisation. Document the paper_nearest reading and theta=0 under the sum rule, describe the stop rules by what they record, and state in the release note that all-zero cells are now returned unchanged. --- docs/release-notes/4424.feat.md | 3 ++- src/scanpy/preprocessing/_harmony/__init__.py | 25 +++++++++++-------- src/scanpy/preprocessing/_harmony/core.py | 11 +++++--- tests/test_harmony.py | 17 +++++++++++++ 4 files changed, 40 insertions(+), 16 deletions(-) diff --git a/docs/release-notes/4424.feat.md b/docs/release-notes/4424.feat.md index fee8285f19..f754c5b5fa 100644 --- a/docs/release-notes/4424.feat.md +++ b/docs/release-notes/4424.feat.md @@ -1,2 +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. -The default seeding is now `"afkmc2"` and the default stop rule `"reference_record"`; pass `seeding="kmeans", stop_rule="reference"` to reproduce previous results {smaller}`T Treis` +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` diff --git a/src/scanpy/preprocessing/_harmony/__init__.py b/src/scanpy/preprocessing/_harmony/__init__.py index de8a149530..2317ba9a5b 100644 --- a/src/scanpy/preprocessing/_harmony/__init__.py +++ b/src/scanpy/preprocessing/_harmony/__init__.py @@ -40,7 +40,7 @@ def harmony_integrate( # noqa: PLR0913 stop_rule: Literal[ "reference_record", "reference", "paper", "paper_window3" ] = "reference_record", - seeding: Literal["afkmc2", "kmeans", "kmeans++", "paper_nearest"] = "afkmc2", + 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", @@ -52,8 +52,7 @@ def harmony_integrate( # noqa: PLR0913 Integration adjusts a PCA embedding to reduce experiment-specific technical variation, allowing cells from different experiments to be compared. - This CPU implementation runs in parallel numba kernels - and follows 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 key is modeled separately instead of combining all keys into one joint category. As Harmony works by adjusting the principal components, @@ -82,7 +81,7 @@ def harmony_integrate( # noqa: PLR0913 adjusted_basis The name of the field in ``adata.obsm`` where the adjusted PCA table will be stored. 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. @@ -165,19 +164,22 @@ def harmony_integrate( # noqa: PLR0913 ``"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, as rapids-singlecell does; - ``"reference"`` records it only when the clustering converged, as harmonypy does. + ``"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. - ``"afkmc2"`` (default) approximates k-means++ seeding with Markov chains :cite:p:`Bachem2016` + 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), - as harmonypy does. + ``"kmeans"`` runs :class:`~sklearn.cluster.KMeans` (k-means++ and up to 25 Lloyd iterations). ``"kmeans++"`` is exact k-means++ seeding. - ``"paper_nearest"`` is the sampler of :cite:t:`Patikas2026`. + ``"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 @@ -186,7 +188,8 @@ def harmony_integrate( # noqa: PLR0913 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`. + ``"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. diff --git a/src/scanpy/preprocessing/_harmony/core.py b/src/scanpy/preprocessing/_harmony/core.py index 97c92f2c5b..a415e5bf19 100644 --- a/src/scanpy/preprocessing/_harmony/core.py +++ b/src/scanpy/preprocessing/_harmony/core.py @@ -65,7 +65,7 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 tol_harmony: float | None = None, tol_clustering: float | None = None, stop_rule: str = "reference_record", - seeding: str = "afkmc2", + seeding: str | None = None, penalty_combine: str = "product", seeding_lloyd_steps: int = 0, seeding_cells_per_cluster: int | None = None, @@ -79,6 +79,8 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 (``objectives``) and the clustering iterations per round (``iterations``). """ stabilized = dynamic_lambda = flavor == "harmony2" + if seeding is None: + seeding = "afkmc2" if flavor == "harmony2" else "kmeans" choices = dict( flavor=("harmony1", "harmony2"), stop_rule=tuple(STOP_RULES), @@ -393,10 +395,11 @@ def _converged(objs: list[float], tol: float, *, window: int, signed: bool) -> b def _seed_paper_nearest( z_norm: np.ndarray, k: int, rng: np.random.Generator ) -> np.ndarray: - """Seeding of :cite:t:`Patikas2026`. + """Seed as :cite:t:`Patikas2026` describe, read so that the seeds spread. - k distinct cells; each new one is the cell with the highest ``log(u_j) / D_j``, ``u ~ U(0, 1)``, - which draws it with probability proportional to its Euclidean distance D_j to the nearest cell chosen so far. + k distinct cells; each new one is the cell with the highest ``log(u_j) / D_j``, ``u ~ U(0, 1)`` + (the paper prints "lowest", which would pick the closest cells), which draws it with probability + proportional to its Euclidean distance D_j to the nearest cell chosen so far. """ n = z_norm.shape[0] chosen = [int(rng.integers(n))] diff --git a/tests/test_harmony.py b/tests/test_harmony.py index aba54036da..9d13138182 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -1030,6 +1030,23 @@ def test_harmonize_seedings_are_reproducible(seeding: str) -> None: assert np.isfinite(out[0]).all() +def test_harmonize_default_seeding_follows_flavor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class SeededError(Exception): + pass + + def afkmc2(*_args, **_kwargs): + raise SeededError + + monkeypatch.setattr(core, "_seed_afkmc2", afkmc2) + x, obs = _toy() + with pytest.raises(SeededError): + core.harmonize(x, obs, "batch", max_iter_harmony=1, rng=0) + # harmony1 keeps the k-means initialisation + core.harmonize(x, obs, "batch", flavor="harmony1", max_iter_harmony=1, rng=0) + + def test_harmonize_edge_inputs() -> None: x, obs = _toy() for extra in (dict(tau=5), dict(stop_rule="paper_window3"), dict(n_clusters=2)): From 3eb4b25100d1421960da174b3a57a12976815005 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 14:36:01 +0200 Subject: [PATCH 06/10] test: pin k-means seeding only where the fixture needs it Describe the default-seeding test by what it checks. --- tests/test_harmony.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/test_harmony.py b/tests/test_harmony.py index 9d13138182..2e6a5e5c98 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -68,8 +68,6 @@ rng=734, alpha=0.2, batch_prune_threshold=1e-5, - seeding="kmeans", # harmonypy seeds with k-means - stop_rule="reference", # and records a round's objective only if clustering converged ) @@ -413,6 +411,9 @@ def test_harmony2_multikey_reference( dtype=np.float64, n_clusters=n_clusters, max_iter_harmony=max_iter_harmony, + # these 80-cell outputs start from k-means centroids; + # test_harmony2_multikey_default_seeding covers the default + seeding="kmeans", ) result = adata.obsm["X_pca_harmony"] @@ -434,10 +435,10 @@ def _within_group_r2(x: np.ndarray, group: np.ndarray, covariate: np.ndarray) -> def test_harmony2_multikey_default_seeding(adata_harmonypy2_multikey: AnnData) -> None: - """The default seeding integrates at least as well as harmonypy on its multikey fixture. + """The default seeding integrates the multikey fixture at least as well as the stored output. - It starts further from the k-means fixed point than harmonypy, so it runs more rounds - and lands farther from the stored output: compare the biology, not the coordinates. + It starts further from the k-means fixed point, so it runs more rounds + and lands farther from the stored coordinates: compare the biology instead. """ from sklearn.cluster import KMeans from sklearn.metrics import silhouette_score From 0ab2bad4203581c2dada520e5eb3597b5a154316 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 14:46:21 +0200 Subject: [PATCH 07/10] perf: bound the memory of the multi-key correction With many joint batch categories (crossed keys), sum the regression statistics per level, one pass per key, instead of per joint category, and apply the coefficients per level. Solve the clusters eight at a time and accumulate the cross-key masses with bincount instead of np.add.at. --- src/scanpy/preprocessing/_harmony/_kernels.py | 22 +- src/scanpy/preprocessing/_harmony/core.py | 229 +++++++++++------- tests/test_harmony.py | 16 +- 3 files changed, 165 insertions(+), 102 deletions(-) diff --git a/src/scanpy/preprocessing/_harmony/_kernels.py b/src/scanpy/preprocessing/_harmony/_kernels.py index 6f73c3ecca..da533acfcd 100644 --- a/src/scanpy/preprocessing/_harmony/_kernels.py +++ b/src/scanpy/preprocessing/_harmony/_kernels.py @@ -267,10 +267,10 @@ def _joint_stats(x, r, joint, order, n_joint, n_chunks): # noqa: PLR0912, PLR09 @njit -def _apply_correction(x, r, joint, w_joint, z_hat, z_norm, n_chunks): # noqa: PLR0917 +def _apply_correction(x, r, cat, w, z_hat, z_norm, n_chunks): # noqa: PLR0912, PLR0917 """Apply the correction, normalise, and accumulate the next centroids. - Writes ``z_hat = x_i - sum_k r_ik * w_joint[joint_i, k, :]`` and its row-normalised copy ``z_norm`` in place; + Writes ``z_hat = x_i - sum_k r_ik * sum_c w[cat_ic, k, :]`` and its row-normalised copy ``z_norm`` in place; returns ``y_t`` (d x K), the normalised centroids of ``r`` on the new ``z_norm``. Each row i is written in full before it is read and no other row is touched, so passing the same array as both @@ -284,13 +284,14 @@ def _apply_correction(x, r, joint, w_joint, z_hat, z_norm, n_chunks): # noqa: P for t in prange(n_chunks): y_loc = y_acc[t] for i in range(t * n_cells // n_chunks, (t + 1) * n_cells // n_chunks): - j = joint[i] for dd in range(n_pc): z_hat[i, dd] = x[i, dd] for k in range(n_clusters): rk = r[i, k] - for dd in range(n_pc): - z_hat[i, dd] -= rk * w_joint[j, k, dd] + for c in range(cat.shape[1]): + j = cat[i, c] + for dd in range(n_pc): + z_hat[i, dd] -= rk * w[j, k, dd] sq = ft(0.0) for dd in range(n_pc): sq += z_hat[i, dd] * z_hat[i, dd] @@ -316,8 +317,8 @@ def _apply_correction(x, r, joint, w_joint, z_hat, z_norm, n_chunks): # noqa: P @njit -def _apply_only(x, r, joint, w_joint, out, n_chunks): # noqa: PLR0917 - """``out[i] = x[i] - sum_k r_ik * w_joint[joint_i, k]``: the unnormalised corrected embedding, no centroids. +def _apply_only(x, r, cat, w, out, n_chunks): # noqa: PLR0917 + """``out[i] = x[i] - sum_k r_ik * sum_c w[cat_ic, k]``: the unnormalised corrected embedding, no centroids. The last step of the low-memory path, which reconstructs ``z_hat`` into the reused ``z_norm`` buffer once at the end instead of keeping a separate ``z_hat`` array for every iteration (``z_hat`` is written but never read until @@ -327,11 +328,12 @@ def _apply_only(x, r, joint, w_joint, out, n_chunks): # noqa: PLR0917 n_clusters = r.shape[1] for t in prange(n_chunks): for i in range(t * n_cells // n_chunks, (t + 1) * n_cells // n_chunks): - j = joint[i] for dd in range(n_pc): out[i, dd] = x[i, dd] for k in range(n_clusters): rk = r[i, k] - for dd in range(n_pc): - out[i, dd] -= rk * w_joint[j, k, dd] + for c in range(cat.shape[1]): + j = cat[i, c] + for dd in range(n_pc): + out[i, dd] -= rk * w[j, k, dd] return out diff --git a/src/scanpy/preprocessing/_harmony/core.py b/src/scanpy/preprocessing/_harmony/core.py index a415e5bf19..8da8ea6a65 100644 --- a/src/scanpy/preprocessing/_harmony/core.py +++ b/src/scanpy/preprocessing/_harmony/core.py @@ -1,7 +1,7 @@ from __future__ import annotations from itertools import pairwise -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple import numpy as np from scipy import sparse @@ -44,6 +44,8 @@ SEEDINGS = ("afkmc2", "kmeans", "kmeans++", "paper_nearest") # Fixed number of static chunks per parallel sweep, so results do not depend on the thread count. N_CHUNKS = 32 +# clusters per batched solve of the correction; bounds its dense systems +_SOLVE_CHUNK = 8 def harmonize( # noqa: PLR0912, PLR0913, PLR0915 @@ -161,8 +163,7 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 msg = "No cell has a non-zero embedding." raise ValueError(msg) n_batches = int(n_levels.sum()) - joint, levels = _joint_categories(codes) - joint_order = np.argsort(joint, kind="stable") + plan = _correction_plan(codes, n_batches) n_b = np.bincount(codes.ravel(), minlength=n_batches).astype(ft) pr_b = n_b / ft(n_cells) theta_b = _get_theta_array(theta, n_levels, x.dtype).ravel() @@ -287,14 +288,9 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 ridge_lambda=ridge_lambda, dynamic_lambda=dynamic_lambda, ) - s_joint, o_joint = _joint_stats( - x, r, joint, joint_order, levels.shape[0], N_CHUNKS - ) - w_joint = _solve_correction( - s_joint, o_joint, levels, n_levels, lam, _SUPPRESS_PENALTY - ) + w = _correction_weights(x, r, plan, n_levels, lam) # low_memory: write the normalised result into z_norm in place (z_hat aliases it) - y_t = _apply_correction(x, r, joint, w_joint, z_hat, z_norm, N_CHUNKS) + y_t = _apply_correction(x, r, plan.cat, w, z_hat, z_norm, N_CHUNKS) if _converged(objs_harmony, tol_harmony, window=1, signed=True): log.info(f"Harmony converged in {i + 1} iterations") break @@ -304,7 +300,7 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 if ( low_memory ): # the shared buffer holds the normalised embedding; rebuild z_hat into it - z_hat = _apply_only(x, r, joint, w_joint, z_hat, N_CHUNKS) + z_hat = _apply_only(x, r, plan.cat, w, z_hat, N_CHUNKS) if not keep.all(): x_all = x_all.copy() x_all[keep] = z_hat @@ -487,26 +483,91 @@ def _joint_categories(codes: np.ndarray) -> tuple[np.ndarray, np.ndarray]: return np.ascontiguousarray(joint.ravel().astype(np.int32)), levels +class _Plan(NamedTuple): + levels: np.ndarray # J x C, the levels of every joint category + joint: np.ndarray # joint category of every cell + joint_order: np.ndarray + cols: list[np.ndarray] # category columns summed by the statistics passes + orders: list[np.ndarray] + cat: np.ndarray # n x C', the categories whose coefficients a cell receives + by_joint: bool + + +def _correction_plan(codes: np.ndarray, n_batches: int) -> _Plan: + """How the correction groups cells: by joint category when there are few, else by each key's level. + + Few joint categories (any single key, nested keys) take one statistics pass and one coefficient row per + category. Many (crossed keys) take one pass per key, so nothing larger than levels x K x d is held. + """ + joint, levels = _joint_categories(codes) + joint_order = np.argsort(joint, kind="stable") + # the cut-off keeps the per-category arrays within twice the per-level ones + by_joint = levels.shape[0] <= 2 * n_batches + cols = [joint] if by_joint else [np.ascontiguousarray(c) for c in codes.T] + orders = [joint_order] if by_joint else [np.argsort(c, kind="stable") for c in cols] + cat = np.ascontiguousarray(joint[:, None] if by_joint else codes) + return _Plan(levels, joint, joint_order, cols, orders, cat, by_joint) + + +def _correction_weights( + x: np.ndarray, + r: np.ndarray, + plan: _Plan, + n_levels: np.ndarray, + lam: np.ndarray, +) -> np.ndarray: + """Coefficients per category of ``plan`` (categories x K x d) for the apply kernels.""" + levels, joint, joint_order, cols, orders, _, by_joint = plan + n_joint, n_batches = levels.shape[0], int(n_levels.sum()) + if by_joint: + s_joint, o_joint = _joint_stats(x, r, joint, joint_order, n_joint, N_CHUNKS) + fold = sparse.csr_matrix( # noqa: TID251 + ( + np.ones(levels.size), + (levels.ravel(), np.repeat(np.arange(n_joint), levels.shape[1])), + ), + shape=(n_batches, n_joint), + ) + s = (fold @ s_joint.reshape(n_joint, -1).astype(np.float64)).reshape( + n_batches, *s_joint.shape[1:] + ) + else: + s = sum( + _joint_stats(x, r, c, order, n_batches, N_CHUNKS)[0] + for c, order in zip(cols, orders, strict=True) + ) + # masses per joint category: a pass without the PC dimension + no_pc = np.empty((x.shape[0], 0), dtype=x.dtype) + _, o_joint = _joint_stats(no_pc, r, joint, joint_order, n_joint, N_CHUNKS) + w = _solve_correction(s, o_joint, levels, n_levels, lam, _SUPPRESS_PENALTY) + if by_joint: + w = w[levels].sum(axis=1) + return np.ascontiguousarray(w, dtype=x.dtype) + + def _solve_correction( # noqa: PLR0915, PLR0917 - s_joint: np.ndarray, + s: np.ndarray, o_joint: np.ndarray, levels: np.ndarray, n_levels: np.ndarray, lam: np.ndarray, sentinel: float, ) -> np.ndarray: - """Ridge coefficients of every cluster, returned as the total correction per joint category (J x K x d). + """Ridge coefficients of every cluster per batch level (L x K x d, float64), without the intercept. + ``s[l, k] = sum of r_ik * x_i`` over the cells with level l; ``o_joint`` holds the masses per joint category. Per cluster the normal equations are ``(Phi* diag(r_k) Phi*^T + Lambda_k) W_k = Phi* diag(r_k) X^T`` with an unpenalised intercept. The block of the covariate with the most levels is diagonal, so it is eliminated in closed form and only the Schur complement over the intercept and the remaining levels is solved; with one covariate that is a 1 x 1 system. ``lam[b, k] >= sentinel`` removes level b - from cluster k (its coefficient is 0). Runs in float64; nothing here has a cell dimension. + from cluster k (its coefficient is 0). Runs in float64, ``_SOLVE_CHUNK`` clusters at a time, + so the dense systems stay at chunk x m x n_big. """ - n_joint, n_clusters, n_pc = s_joint.shape + n_clusters, n_pc = s.shape[1:] offsets = np.concatenate([[0], np.cumsum(n_levels)]) big = int(np.argmax(n_levels)) n_big = int(n_levels[big]) + big_lv = np.arange(offsets[big], offsets[big + 1]) others = [c for c in range(len(n_levels)) if c != big] big_of_joint = levels[:, big] - offsets[big] small = np.concatenate( @@ -519,73 +580,68 @@ def _solve_correction( # noqa: PLR0915, PLR0917 # rows[j] = the rows of the small system that joint category j belongs to: # the intercept and one level per other key rows = np.concatenate( - [np.zeros((n_joint, 1), dtype=np.int64), row_of_level[levels[:, others]]], + [np.zeros((len(levels), 1), dtype=np.int64), row_of_level[levels[:, others]]], axis=1, ) - o = o_joint.astype(np.float64) # J x K - s = s_joint.astype(np.float64) # J x K x d + o_all = o_joint.astype(np.float64) # J x K + s = s.astype(np.float64) # L x K x d lam = lam.astype(np.float64) - pruned = lam >= sentinel - lam_big = lam[offsets[big] : offsets[big + 1]].T # K x n_big - pruned_big = pruned[offsets[big] : offsets[big + 1]].T - lam_small, pruned_small = lam[small].T, pruned[small].T # K x small - - # sums over joint categories; cluster axis last while accumulating, first afterwards - o_big = np.zeros((n_big, n_clusters)) - np.add.at(o_big, big_of_joint, o) - rhs_big = np.zeros((n_big, n_clusters, n_pc)) - np.add.at(rhs_big, big_of_joint, s) - gram = np.zeros((m, m, n_clusters)) - rhs = np.zeros((m, n_clusters, n_pc)) - coupling = np.zeros(( - m, - n_big, - n_clusters, - )) # mass shared by a small row and a big level - for a in range(rows.shape[1]): - np.add.at(rhs, rows[:, a], s) - np.add.at(coupling, (rows[:, a], big_of_joint), o) - for b in range(rows.shape[1]): - np.add.at(gram, (rows[:, a], rows[:, b]), o) - gram = gram.transpose(2, 0, 1).copy() - rhs = rhs.transpose(1, 0, 2).copy() - coupling = coupling.transpose(2, 0, 1).copy() - o_big, rhs_big = o_big.T, rhs_big.transpose(1, 0, 2) # K x n_big, K x n_big x d - - diag = o_big + lam_big - # 0 drops the level - d_inv = np.where(pruned_big | (diag <= 0), 0.0, 1.0 / np.where(diag > 0, diag, 1.0)) - idx = np.arange(1, m) - gram[:, idx, idx] += np.where(pruned_small, 0.0, lam_small) - # a pruned small level: empty row and column, unit diagonal after elimination - kk, ss = np.nonzero(pruned_small) - gram[kk, ss + 1, :] = 0 - gram[kk, :, ss + 1] = 0 - coupling[kk, ss + 1, :] = 0 - rhs[kk, ss + 1] = 0 - scaled = coupling * d_inv[:, None, :] # C D^-1, K x m x n_big - schur = gram - scaled @ coupling.transpose(0, 2, 1) - rhs -= scaled @ rhs_big - schur[kk, ss + 1, ss + 1] = 1 - empty = gram[:, 0, 0] <= 1e-12 # cluster without cells: no correction - schur[empty] = np.eye(m) - try: - w_small = np.linalg.solve(schur, rhs) # K x m x d - except np.linalg.LinAlgError: - # guard only: from `harmonize` every unpruned lambda is positive and the system is positive definite - w_small = np.stack([ - np.linalg.lstsq(a, b, rcond=None)[0] - for a, b in zip(schur, rhs, strict=True) - ]) - w_big = d_inv[:, :, None] * (rhs_big - coupling.transpose(0, 2, 1) @ w_small) - w_small[:, 0] = 0 # the intercept is not part of the correction - w_small[empty] = 0 - w_big[empty] = 0 - w_joint = w_big[:, big_of_joint] # K x J x d - for a in range(1, rows.shape[1]): - w_joint = w_joint + w_small[:, rows[:, a]] - return np.ascontiguousarray(w_joint.transpose(1, 0, 2), dtype=s_joint.dtype) + w = np.zeros(s.shape) + for k0 in range(0, n_clusters, _SOLVE_CHUNK): + ks = slice(k0, min(k0 + _SOLVE_CHUNK, n_clusters)) + o = o_all[:, ks] # J x c + c = o.shape[1] + kl = np.arange(c) + pruned = lam[:, ks] >= sentinel + lam_big, pruned_big = lam[big_lv, ks].T, pruned[big_lv].T # c x n_big + lam_small, pruned_small = lam[small, ks].T, pruned[small].T # c x small + row_k = kl * m + rows[:, :, None] # J x C x c + o_big = _masses(kl * n_big + big_of_joint[:, None], o, (c, n_big)) + gram = _masses(row_k[:, :, None] * m + rows[:, None, :, None], o, (c, m, m)) + coupling = _masses( # mass shared by a small row and a big level + row_k * n_big + big_of_joint[:, None, None], o, (c, m, n_big) + ) + rhs_big = s[big_lv, ks].transpose(1, 0, 2) # c x n_big x d + rhs = np.empty((c, m, n_pc)) + rhs[:, 0] = rhs_big.sum(axis=1) + rhs[:, 1:] = s[small, ks].transpose(1, 0, 2) + + diag = o_big + lam_big + # 0 drops the level + d_inv = np.where( + pruned_big | (diag <= 0), 0.0, 1.0 / np.where(diag > 0, diag, 1.0) + ) + idx = np.arange(1, m) + gram[:, idx, idx] += np.where(pruned_small, 0.0, lam_small) + # a pruned small level: empty row and column, unit diagonal after elimination + kk, ss = np.nonzero(pruned_small) + gram[kk, ss + 1, :] = 0 + gram[kk, :, ss + 1] = 0 + coupling[kk, ss + 1, :] = 0 + rhs[kk, ss + 1] = 0 + scaled = coupling * d_inv[:, None, :] # C D^-1, c x m x n_big + schur = gram - scaled @ coupling.transpose(0, 2, 1) + rhs -= scaled @ rhs_big + schur[kk, ss + 1, ss + 1] = 1 + empty = gram[:, 0, 0] <= 1e-12 # cluster without cells: no correction + schur[empty] = np.eye(m) + try: + w_small = np.linalg.solve(schur, rhs) # c x m x d + except np.linalg.LinAlgError: + # guard only: from `harmonize` every unpruned lambda is positive and the system is positive definite + w_small = np.stack([ + np.linalg.lstsq(a, b, rcond=None)[0] + for a, b in zip(schur, rhs, strict=True) + ]) + w_big = d_inv[:, :, None] * (rhs_big - coupling.transpose(0, 2, 1) @ w_small) + w_small[empty] = 0 + w_big[empty] = 0 + w[big_lv, ks] = w_big.transpose(1, 0, 2) + w[small, ks] = w_small[:, 1:].transpose( + 1, 0, 2 + ) # the intercept is not part of the correction + return w def _correction_multi( @@ -601,14 +657,17 @@ def _correction_multi( # so where its range is cut does not change the result ends = np.append(batch_codes.max(axis=0)[:-1] + 1, n_batches) n_levels = np.diff(ends, prepend=0) - joint, levels = _joint_categories(batch_codes) - s_joint, o_joint = _joint_stats( - x, r, joint, np.argsort(joint, kind="stable"), levels.shape[0], N_CHUNKS - ) - w_joint = _solve_correction( - s_joint, o_joint, levels, n_levels, lambda_kb, _SUPPRESS_PENALTY + plan = _correction_plan(batch_codes, n_batches) + w = _correction_weights(x, r, plan, n_levels, lambda_kb) + return _apply_only(x, r, plan.cat, w, np.empty_like(x), N_CHUNKS) + + +def _masses(idx: np.ndarray, o: np.ndarray, shape: tuple[int, ...]) -> np.ndarray: + """Sum the joint masses ``o`` (J x c) into a dense array of ``shape`` at flat indices ``idx`` (J x ... x c).""" + weights = np.broadcast_to( + o.reshape(o.shape[0], *[1] * (idx.ndim - 2), -1), idx.shape ) - return _apply_only(x, r, joint, w_joint, np.empty_like(x), N_CHUNKS) + return np.bincount(idx.ravel(), weights.ravel(), int(np.prod(shape))).reshape(shape) def _compute_lambda_kb( diff --git a/tests/test_harmony.py b/tests/test_harmony.py index 2e6a5e5c98..ceb8d0c5f6 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -845,7 +845,12 @@ def _random_design(rng: np.random.Generator, n: int, sizes: tuple[int, ...]): @pytest.mark.parametrize("dtype", _DTYPES) # big block first / last / a size-1 key / one key @pytest.mark.parametrize("sizes", [(4, 9, 3), (4, 3, 9), (1, 8), (7,)]) -def test_correction_with_pruning_and_an_empty_cluster(dtype, sizes) -> None: +def test_correction_with_pruning_and_an_empty_cluster( + dtype, sizes, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + core, "_SOLVE_CHUNK", 2 + ) # the pruned and the empty cluster in different chunks rng = np.random.default_rng(sum(sizes)) codes, n_levels, x, r = _random_design(rng, 500, sizes) r[:, 3] = 0 # an empty cluster @@ -904,13 +909,10 @@ def test_apply_correction_normalises_and_returns_centroids(dtype) -> None: ridge_lambda=1.0, dynamic_lambda=True, ) - joint, levels = core._joint_categories(codes) - s, o_joint = _joint_stats( - x, r, joint, np.argsort(joint, kind="stable"), len(levels), 3 - ) - w = core._solve_correction(s, o_joint, levels, n_levels, lam, _SUPPRESS_PENALTY) + plan = core._correction_plan(codes, int(n_levels.sum())) + w = core._correction_weights(x, r, plan, n_levels, lam) z_hat, z_norm = np.empty_like(x), np.empty_like(x) - y_t = _apply_correction(x, r, joint, w, z_hat, z_norm, 3) + y_t = _apply_correction(x, r, plan.cat, w, z_hat, z_norm, 3) assert z_hat.dtype == z_norm.dtype == y_t.dtype == dtype np.testing.assert_allclose( z_hat, From fa1b82db91c6e82cd4a637473c8cb3cb84d89db1 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 14:54:52 +0200 Subject: [PATCH 08/10] perf: skip unused centroids and switch to count scatters later The centroids of the last allowed clustering iteration are replaced by the ones the correction computes, so that pass no longer accumulates them. The serial count scatters only beat the per-chunk accumulators at about four times the accumulators' work, so _use_scatter now requires that margin. --- src/scanpy/preprocessing/_harmony/_kernels.py | 25 ++++++++++++------- src/scanpy/preprocessing/_harmony/core.py | 13 ++++++---- tests/test_harmony.py | 24 ++++++++++++++---- 3 files changed, 43 insertions(+), 19 deletions(-) diff --git a/src/scanpy/preprocessing/_harmony/_kernels.py b/src/scanpy/preprocessing/_harmony/_kernels.py index da533acfcd..59213cde51 100644 --- a/src/scanpy/preprocessing/_harmony/_kernels.py +++ b/src/scanpy/preprocessing/_harmony/_kernels.py @@ -31,6 +31,7 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 n_chunks, scatter, sum_penalties, + want_y, ): """One clustering iteration: block-wise soft assignment, objective parts, next centroids. @@ -46,6 +47,9 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 ``sum_penalties`` selects how the penalty factors of a cell's several batch keys are combined: ``False`` multiplies them :cite:p:`Korsunsky2019`; ``True`` adds them, as the assignment rule is printed in :cite:t:`Patikas2026`, skipping levels with ``theta == 0``. With one key the two are the same. + + ``want_y=False`` skips the centroids and returns the input ``y_t``, for an iteration whose centroids are + never read. """ ft = z_norm.dtype.type n_pc = z_norm.shape[1] @@ -165,10 +169,11 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 b = codes[i, c] for k in range(n_clusters): acc[t, b, k] += r[i, k] - for dd in range(n_pc): - zv = z_norm[i, dd] - for k in range(n_clusters): - y_loc[dd, k] += zv * r[i, k] + if want_y: + for dd in range(n_pc): + zv = z_norm[i, dd] + for k in range(n_clusters): + y_loc[dd, k] += zv * r[i, k] dist[t] += dist_t ent[t] += ent_t # 4. put the block back @@ -189,6 +194,13 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 o[b, k] += acc[t, b, k] o_fresh[b, k] += acc[t, b, k] r_sum[k] += acc[t, b, k] * inv_cov + # serial sums: a reduction over the parallel chunks would depend on the thread count + dist_total = ent_total = 0.0 + for t in range(n_chunks): + dist_total += dist[t] + ent_total += ent[t] + if not want_y: + return dist_total, ent_total, o_fresh, y_t # 5. centroids of the new r: y_k = normalise(sum_i r_ik z_i) y_t_next = np.empty((n_pc, n_clusters), dtype=z_norm.dtype) for k in range(n_clusters): @@ -202,11 +214,6 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 scale = one / max(np.sqrt(sq), ft(1e-12)) for dd in range(n_pc): y_t_next[dd, k] *= scale - # serial sums: a reduction over the parallel chunks would depend on the thread count - dist_total = ent_total = 0.0 - for t in range(n_chunks): - dist_total += dist[t] - ent_total += ent[t] return dist_total, ent_total, o_fresh, y_t_next diff --git a/src/scanpy/preprocessing/_harmony/core.py b/src/scanpy/preprocessing/_harmony/core.py index 8da8ea6a65..741ceb1cb2 100644 --- a/src/scanpy/preprocessing/_harmony/core.py +++ b/src/scanpy/preprocessing/_harmony/core.py @@ -223,6 +223,7 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 N_CHUNKS, _use_scatter(n_batches, codes.shape[1], n_cells), sum_penalties, + True, # noqa: FBT003 ) e = pr_b[:, None] * r_sum[None, :] objs_harmony = [ @@ -239,9 +240,9 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 for i in bar: objs: list[float] = [] done = False - for _ in range(max_iter_clustering): + for it in range(max_iter_clustering): perm = rng.permutation(n_cells) - # sort each block's cells so the kernels read z_norm/r/z_hat in ascending order; + # sort each block's cells so the kernel reads z_norm and r in ascending order; # block membership is unchanged, only the order within a block differs for lo, hi in pairwise(bounds): perm[lo:hi].sort() @@ -261,6 +262,8 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 N_CHUNKS, scatter, sum_penalties, + # the correction recomputes the centroids after the last iteration + it < max_iter_clustering - 1, ) if recompute_counts: o[:] = o_fresh @@ -700,10 +703,10 @@ def _compute_lambda_kb( def _use_scatter(n_batches: int, n_cov: int, block_size: int) -> bool: """Whether the clustering kernel moves a block's counts by serial scatters instead of per-chunk accumulators. - Accumulators cost about N_CHUNKS * levels * K per block, scatters about cells * keys * K. - Scatters win for thousands of small batches, accumulators on atlases with few large ones. + Accumulators cost about N_CHUNKS * levels * K per block, scatters about cells * keys * K, but the scatters + are serial: they win only at about four times the accumulators' work (measured at 200k to 1M cells). """ - return N_CHUNKS * n_batches > block_size * n_cov + return N_CHUNKS * n_batches > 4 * block_size * n_cov def _block_bounds(n_cells: int, block_proportion: float) -> np.ndarray: diff --git a/tests/test_harmony.py b/tests/test_harmony.py index ceb8d0c5f6..b2ec2665c0 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -758,7 +758,7 @@ def test_cluster_pass_matches_numpy( dist, ent, o_fresh, y_t = _cluster_pass( *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b, theta, perm), *(_bounds(len(x)), dtype(-2 / _SIGMA), stabilized, n_chunks, scatter), - sum_penalties, + *(sum_penalties, True), ) atol = 2e-5 if dtype == np.float32 else 1e-10 np.testing.assert_allclose(r, r_ref, atol=atol) @@ -775,6 +775,20 @@ def test_cluster_pass_matches_numpy( np.add.at(exact, covariate, r) np.testing.assert_allclose(o_fresh, exact, rtol=1e-4, atol=1e-5) np.testing.assert_allclose(y_t.T, core._normalize_rows_l2(r.T @ z), atol=1e-5) + # without the centroids: same assignments, the input centroids back + y_in, r2, o2, r_sum2 = ( + np.ascontiguousarray(y.T), + r0.copy(), + o0.copy(), + r0.sum(axis=0), + ) + *_, y_out = _cluster_pass( + *(z, y_in, r2, codes, o2, r_sum2, pr_b, theta, perm), + *(_bounds(len(x)), dtype(-2 / _SIGMA), stabilized, n_chunks, scatter), + *(sum_penalties, False), + ) + np.testing.assert_array_equal(r2, r) + assert y_out is y_in if isinstance(keys, list): # the two rules really differ on this input other, _ = _clustering_iteration( *(z, codes, pr_b, r0, o0, theta, perm), @@ -803,7 +817,7 @@ def test_cluster_pass_sum_rule_handles_large_penalties(dtype) -> None: r, o, r_sum = r0.copy(), o0.copy(), r0.sum(axis=0) _cluster_pass( *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b, theta, perm), - *(_bounds(len(x)), dtype(-2 / _SIGMA), False, 32, False, True), + *(_bounds(len(x)), dtype(-2 / _SIGMA), False, 32, False, True, True), ) np.testing.assert_allclose(r, r_ref, atol=5e-5 if dtype == np.float32 else 1e-10) @@ -817,7 +831,7 @@ def test_cluster_pass_without_penalty_is_soft_kmeans() -> None: _cluster_pass( *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b), *(np.zeros(_N_BATCH, np.float32), np.arange(n), np.array([0, n])), - *(np.float32(-20), True, 3, False, False), + *(np.float32(-20), True, 3, False, False, True), ) expected = np.exp(-20 * (1 - z @ y.T).astype(np.float64)) expected /= expected.sum(axis=1, keepdims=True) @@ -980,8 +994,8 @@ def test_harmonize_many_small_batches() -> None: assert core._use_scatter( 4000, 1, 20_000 // 19 ) # thousands of tiny batches: scatters - assert core._use_scatter(11, 1, 10 * core.N_CHUNKS) - assert not core._use_scatter(10, 1, 10 * core.N_CHUNKS) + assert core._use_scatter(41, 1, 10 * core.N_CHUNKS) + assert not core._use_scatter(40, 1, 10 * core.N_CHUNKS) x, obs = _toy(n=1500, dtype=np.float64) obs = obs.assign(well=pd.Categorical(np.arange(1500) % 250)) assert core._use_scatter(250, 1, 1500 // 19) From 9a760698b4cbc1809a8c28d829f0c8d3fa4d65e9 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 14:54:52 +0200 Subject: [PATCH 09/10] test: cover the largest key first in the correction test --- tests/test_harmony.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_harmony.py b/tests/test_harmony.py index b2ec2665c0..30f63ad55a 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -857,8 +857,8 @@ def _random_design(rng: np.random.Generator, n: int, sizes: tuple[int, ...]): @pytest.mark.parametrize("dtype", _DTYPES) -# big block first / last / a size-1 key / one key -@pytest.mark.parametrize("sizes", [(4, 9, 3), (4, 3, 9), (1, 8), (7,)]) +# big block in the middle / first / last / a size-1 key / one key +@pytest.mark.parametrize("sizes", [(4, 9, 3), (9, 4, 3), (4, 3, 9), (1, 8), (7,)]) def test_correction_with_pruning_and_an_empty_cluster( dtype, sizes, monkeypatch: pytest.MonkeyPatch ) -> None: From 020898cfea7b8131f92ad76bc3e280c0142c1e39 Mon Sep 17 00:00:00 2001 From: anon Date: Sun, 18 Oct 2026 16:09:15 +0200 Subject: [PATCH 10/10] refactor: simplify the Harmony port Merge the two apply kernels and share the centroid step, use scikit-learn's k-means++, drop checks that harmonize already makes and the unreachable least-squares fallback, flatten the stop-rule table and the parameter checks, and move the test-only correction wrapper into the tests. --- src/scanpy/preprocessing/_harmony/__init__.py | 12 - src/scanpy/preprocessing/_harmony/_kernels.py | 89 +++----- src/scanpy/preprocessing/_harmony/core.py | 206 ++++++++---------- tests/test_harmony.py | 132 +++++------ 4 files changed, 162 insertions(+), 277 deletions(-) diff --git a/src/scanpy/preprocessing/_harmony/__init__.py b/src/scanpy/preprocessing/_harmony/__init__.py index 2317ba9a5b..9d2ed8d2c6 100644 --- a/src/scanpy/preprocessing/_harmony/__init__.py +++ b/src/scanpy/preprocessing/_harmony/__init__.py @@ -209,10 +209,6 @@ def harmony_integrate( # noqa: PLR0913 """ 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) @@ -254,14 +250,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) - adata.obsm[adjusted_basis], _ = harmonize( x, adata.obs, diff --git a/src/scanpy/preprocessing/_harmony/_kernels.py b/src/scanpy/preprocessing/_harmony/_kernels.py index 59213cde51..223fdc2295 100644 --- a/src/scanpy/preprocessing/_harmony/_kernels.py +++ b/src/scanpy/preprocessing/_harmony/_kernels.py @@ -12,6 +12,7 @@ import numpy as np from fast_array_utils.numba import njit from numba import prange +from numba.extending import register_jitable @njit @@ -39,17 +40,9 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 Returns ``(dist, ent, o_fresh, y_t_next)``: ``dist = sum r*2(1-z.y)``, ``ent = sum r*log(r)`` (float64), ``o`` rebuilt from the new ``r``, and the normalised centroids of the new ``r``. - ``scatter`` selects how a block's counts are moved: ``False`` sums them in per-chunk accumulators in the - parallel sweeps (work per block ~ n_chunks * levels * K, best for many cells per level); ``True`` moves them - with serial loops over the block's cells and fills the penalty table only for levels that occur in the - block (work per block ~ cells * K, best for thousands of small levels). Same result up to rounding. - - ``sum_penalties`` selects how the penalty factors of a cell's several batch keys are combined: ``False`` - multiplies them :cite:p:`Korsunsky2019`; ``True`` adds them, as the assignment rule is printed in - :cite:t:`Patikas2026`, skipping levels with ``theta == 0``. With one key the two are the same. - - ``want_y=False`` skips the centroids and returns the input ``y_t``, for an iteration whose centroids are - never read. + ``scatter``: move a block's counts by serial scatters instead of per-chunk accumulators (see ``_use_scatter``). + ``sum_penalties``: add the keys' penalty factors (over levels with ``theta != 0``) instead of multiplying them. + ``want_y=False``: skip the centroids and return the input ``y_t``. """ ft = z_norm.dtype.type n_pc = z_norm.shape[1] @@ -202,19 +195,27 @@ def _cluster_pass( # noqa: PLR0912, PLR0913, PLR0915, PLR0917 if not want_y: return dist_total, ent_total, o_fresh, y_t # 5. centroids of the new r: y_k = normalise(sum_i r_ik z_i) - y_t_next = np.empty((n_pc, n_clusters), dtype=z_norm.dtype) + return dist_total, ent_total, o_fresh, _centroids(y_acc) + + +@register_jitable +def _centroids(y_acc): + """Sum per-chunk centroid accumulators (chunks x d x K) serially and L2-normalise every column.""" + ft = y_acc.dtype.type + n_chunks, n_pc, n_clusters = y_acc.shape + y_t = np.empty((n_pc, n_clusters), dtype=y_acc.dtype) for k in range(n_clusters): sq = ft(0.0) for dd in range(n_pc): v = ft(0.0) for t in range(n_chunks): v += y_acc[t, dd, k] - y_t_next[dd, k] = v + y_t[dd, k] = v sq += v * v - scale = one / max(np.sqrt(sq), ft(1e-12)) + scale = ft(1.0) / max(np.sqrt(sq), ft(1e-12)) for dd in range(n_pc): - y_t_next[dd, k] *= scale - return dist_total, ent_total, o_fresh, y_t_next + y_t[dd, k] *= scale + return y_t @njit @@ -274,19 +275,16 @@ def _joint_stats(x, r, joint, order, n_joint, n_chunks): # noqa: PLR0912, PLR09 @njit -def _apply_correction(x, r, cat, w, z_hat, z_norm, n_chunks): # noqa: PLR0912, PLR0917 - """Apply the correction, normalise, and accumulate the next centroids. +def _apply_correction(x, r, cat, w, z_hat, z_norm, n_chunks, normalise): # noqa: PLR0917 + """Write ``z_hat = x_i - sum_k r_ik * sum_c w[cat_ic, k, :]`` in place. - Writes ``z_hat = x_i - sum_k r_ik * sum_c w[cat_ic, k, :]`` and its row-normalised copy ``z_norm`` in place; - returns ``y_t`` (d x K), the normalised centroids of ``r`` on the new ``z_norm``. - - Each row i is written in full before it is read and no other row is touched, so passing the same array as both - ``z_hat`` and ``z_norm`` is safe and yields the normalised result (the low_memory path relies on this). + With ``normalise``, also write its row-normalised copy into ``z_norm`` and return the normalised centroids + ``y_t`` (d x K) of ``r`` on it; otherwise return an empty array. Each row is written in full before it is read, + so ``z_hat`` and ``z_norm`` may be the same array (the low_memory path). """ ft = x.dtype.type n_cells, n_pc = x.shape n_clusters = r.shape[1] - one = ft(1.0) y_acc = np.zeros((n_chunks, n_pc, n_clusters), dtype=x.dtype) for t in prange(n_chunks): y_loc = y_acc[t] @@ -299,48 +297,17 @@ def _apply_correction(x, r, cat, w, z_hat, z_norm, n_chunks): # noqa: PLR0912, j = cat[i, c] for dd in range(n_pc): z_hat[i, dd] -= rk * w[j, k, dd] + if not normalise: + continue sq = ft(0.0) for dd in range(n_pc): sq += z_hat[i, dd] * z_hat[i, dd] - scale = one / max(np.sqrt(sq), ft(1e-12)) + scale = ft(1.0) / max(np.sqrt(sq), ft(1e-12)) for dd in range(n_pc): zv = z_hat[i, dd] * scale z_norm[i, dd] = zv for k in range(n_clusters): y_loc[dd, k] += zv * r[i, k] - y_t = np.empty((n_pc, n_clusters), dtype=x.dtype) - for k in range(n_clusters): - sq = ft(0.0) - for dd in range(n_pc): - v = ft(0.0) - for t in range(n_chunks): - v += y_acc[t, dd, k] - y_t[dd, k] = v - sq += v * v - scale = one / max(np.sqrt(sq), ft(1e-12)) - for dd in range(n_pc): - y_t[dd, k] *= scale - return y_t - - -@njit -def _apply_only(x, r, cat, w, out, n_chunks): # noqa: PLR0917 - """``out[i] = x[i] - sum_k r_ik * sum_c w[cat_ic, k]``: the unnormalised corrected embedding, no centroids. - - The last step of the low-memory path, which reconstructs ``z_hat`` into the reused ``z_norm`` buffer once at the - end instead of keeping a separate ``z_hat`` array for every iteration (``z_hat`` is written but never read until - the return). Same per-cell arithmetic as the first half of ``_apply_correction``. - """ - n_cells, n_pc = x.shape - n_clusters = r.shape[1] - for t in prange(n_chunks): - for i in range(t * n_cells // n_chunks, (t + 1) * n_cells // n_chunks): - for dd in range(n_pc): - out[i, dd] = x[i, dd] - for k in range(n_clusters): - rk = r[i, k] - for c in range(cat.shape[1]): - j = cat[i, c] - for dd in range(n_pc): - out[i, dd] -= rk * w[j, k, dd] - return out + if not normalise: + return np.empty((0, 0), dtype=x.dtype) + return _centroids(y_acc) diff --git a/src/scanpy/preprocessing/_harmony/core.py b/src/scanpy/preprocessing/_harmony/core.py index 741ceb1cb2..0e2564a101 100644 --- a/src/scanpy/preprocessing/_harmony/core.py +++ b/src/scanpy/preprocessing/_harmony/core.py @@ -5,7 +5,7 @@ import numpy as np from scipy import sparse -from sklearn.cluster import KMeans +from sklearn.cluster import KMeans, kmeans_plusplus from tqdm.auto import tqdm from ... import logging as log @@ -13,7 +13,7 @@ from ..._settings import settings from ..._settings.verbosity import Verbosity from ..._utils.random import _legacy_random_state -from ._kernels import _apply_correction, _apply_only, _cluster_pass, _joint_stats +from ._kernels import _apply_correction, _cluster_pass, _joint_stats if TYPE_CHECKING: from collections.abc import Sequence @@ -29,17 +29,16 @@ _SUPPRESS_PENALTY = 1e30 -# stop_rule -> (window of the clustering test, record every round's objective, -# {flavor: (max_iter_clustering, tol_clustering, tol_harmony)}) -_REFERENCE = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (4, 1e-3, 1e-2)} -_PAPER = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (200, 1e-5, 1e-4)} +# flavor -> (max_iter_clustering, tol_clustering, tol_harmony); the "paper" rules use harmony1's for both +_DEFAULTS = {"harmony1": (200, 1e-5, 1e-4), "harmony2": (4, 1e-3, 1e-2)} +# stop_rule -> (window of the clustering test, record every round's objective) STOP_RULES = { # the round objective is only recorded if clustering converged - "reference": (3, False, _REFERENCE), + "reference": (3, False), # every round's objective is recorded, as in rapids-singlecell - "reference_record": (3, True, _REFERENCE), - "paper": (1, True, _PAPER), - "paper_window3": (3, True, _PAPER), + "reference_record": (3, True), + "paper": (1, True), + "paper_window3": (3, True), } SEEDINGS = ("afkmc2", "kmeans", "kmeans++", "paper_nearest") # Fixed number of static chunks per parallel sweep, so results do not depend on the thread count. @@ -80,48 +79,53 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 Returns the corrected embedding and a dict with the Harmony objective per round (``objectives``) and the clustering iterations per round (``iterations``). """ - stabilized = dynamic_lambda = flavor == "harmony2" + stabilized = flavor == "harmony2" if seeding is None: - seeding = "afkmc2" if flavor == "harmony2" else "kmeans" - choices = dict( - flavor=("harmony1", "harmony2"), - stop_rule=tuple(STOP_RULES), - seeding=SEEDINGS, - penalty_combine=("product", "sum"), - ) - given = dict( - flavor=flavor, - stop_rule=stop_rule, - seeding=seeding, - penalty_combine=penalty_combine, - ) - for name, allowed in choices.items(): - if given[name] not in allowed: - msg = f"{name} must be one of {allowed}, got {given[name]!r}." + seeding = "afkmc2" if stabilized else "kmeans" + for name, value, allowed in ( + ("flavor", flavor, ("harmony1", "harmony2")), + ("stop_rule", stop_rule, tuple(STOP_RULES)), + ("seeding", seeding, SEEDINGS), + ("penalty_combine", penalty_combine, ("product", "sum")), + ): + if value not in allowed: + msg = f"{name} must be one of {allowed}, got {value!r}." raise ValueError(msg) - checks = { - "max_iter_harmony must be >= 1": max_iter_harmony >= 1, - "max_iter_clustering must be >= 1": max_iter_clustering is None - or max_iter_clustering >= 1, - f"sigma must be a finite positive number, got {sigma}.": np.isfinite(sigma) - and sigma > 0, - f"block_proportion must be in (0, 1], got {block_proportion}.": 0 - < block_proportion - <= 1, - f"alpha must be a finite positive number when flavor='harmony2', got {alpha}.": not dynamic_lambda - or (np.isfinite(alpha) and alpha > 0), - f"batch_prune_threshold must be in [0, 1] or None, got {batch_prune_threshold}.": not dynamic_lambda - or batch_prune_threshold is None - or 0 <= batch_prune_threshold <= 1, - f"ridge_lambda must be a finite positive number when flavor='harmony1', got {ridge_lambda}.": dynamic_lambda - or (np.isfinite(ridge_lambda) and ridge_lambda > 0), - f"seeding_lloyd_steps must be >= 0, got {seeding_lloyd_steps}.": seeding_lloyd_steps - >= 0, - f"seeding_cells_per_cluster must be >= 1 or None, got {seeding_cells_per_cluster}.": seeding_cells_per_cluster - is None - or seeding_cells_per_cluster >= 1, - } - for msg, ok in checks.items(): + finite_pos = lambda v: np.isfinite(v) and v > 0 # noqa: E731 + for ok, msg in ( + (max_iter_harmony >= 1, "max_iter_harmony must be >= 1"), + ( + max_iter_clustering is None or max_iter_clustering >= 1, + "max_iter_clustering must be >= 1", + ), + (finite_pos(sigma), f"sigma must be a finite positive number, got {sigma}."), + ( + 0 < block_proportion <= 1, + f"block_proportion must be in (0, 1], got {block_proportion}.", + ), + ( + not stabilized or finite_pos(alpha), + f"alpha must be a finite positive number when flavor='harmony2', got {alpha}.", + ), + ( + not stabilized + or batch_prune_threshold is None + or 0 <= batch_prune_threshold <= 1, + f"batch_prune_threshold must be in [0, 1] or None, got {batch_prune_threshold}.", + ), + ( + stabilized or finite_pos(ridge_lambda), + f"ridge_lambda must be a finite positive number when flavor='harmony1', got {ridge_lambda}.", + ), + ( + seeding_lloyd_steps >= 0, + f"seeding_lloyd_steps must be >= 0, got {seeding_lloyd_steps}.", + ), + ( + seeding_cells_per_cluster is None or seeding_cells_per_cluster >= 1, + f"seeding_cells_per_cluster must be >= 1 or None, got {seeding_cells_per_cluster}.", + ), + ): if not ok: raise ValueError(msg) x = np.ascontiguousarray(x) @@ -135,10 +139,11 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 msg = "Input data contains non-finite values." raise ValueError(msg) ft = x.dtype.type - window, record_always, defaults = STOP_RULES[stop_rule] - # only differs from the product with two or more batch keys + window, record_always = STOP_RULES[stop_rule] sum_penalties = penalty_combine == "sum" - d_iter, d_tol_c, d_tol_h = defaults[flavor] + d_iter, d_tol_c, d_tol_h = _DEFAULTS[ + "harmony1" if stop_rule.startswith("paper") else flavor + ] max_iter_clustering = d_iter if max_iter_clustering is None else max_iter_clustering tol_clustering = d_tol_c if tol_clustering is None else tol_clustering tol_harmony = d_tol_h if tol_harmony is None else tol_harmony @@ -232,7 +237,7 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 bounds = _block_bounds(n_cells, block_proportion) scatter = _use_scatter(n_batches, codes.shape[1], int(bounds[1])) - # low_memory reuses the z_norm buffer as z_hat's scratch; otherwise z_hat is its own array + # low_memory: z_hat shares z_norm's buffer and is rebuilt once at the end z_hat = z_norm if low_memory else np.empty_like(x) with tqdm( range(max_iter_harmony), disable=settings.verbosity < Verbosity.info @@ -289,21 +294,18 @@ def harmonize( # noqa: PLR0912, PLR0913, PLR0915 alpha=alpha, threshold=batch_prune_threshold, ridge_lambda=ridge_lambda, - dynamic_lambda=dynamic_lambda, + dynamic_lambda=stabilized, ) w = _correction_weights(x, r, plan, n_levels, lam) - # low_memory: write the normalised result into z_norm in place (z_hat aliases it) - y_t = _apply_correction(x, r, plan.cat, w, z_hat, z_norm, N_CHUNKS) + y_t = _apply_correction(x, r, plan.cat, w, z_hat, z_norm, N_CHUNKS, True) # noqa: FBT003 if _converged(objs_harmony, tol_harmony, window=1, signed=True): log.info(f"Harmony converged in {i + 1} iterations") break else: log.info(f"Harmony did not converge after {max_iter_harmony} iterations.") info["objectives"] = objs_harmony - if ( - low_memory - ): # the shared buffer holds the normalised embedding; rebuild z_hat into it - z_hat = _apply_only(x, r, plan.cat, w, z_hat, N_CHUNKS) + if low_memory: + _apply_correction(x, r, plan.cat, w, z_hat, z_norm, N_CHUNKS, False) # noqa: FBT003 if not keep.all(): x_all = x_all.copy() x_all[keep] = z_hat @@ -406,8 +408,7 @@ def _seed_paper_nearest( dist = np.full(n, np.inf, dtype=z_norm.dtype) while len(chosen) < k: taken[chosen[-1]] = True - # unit rows: ||a - b|| = sqrt(2 - 2 a.b) - d = np.sqrt(np.maximum(2.0 - 2.0 * (z_norm @ z_norm[chosen[-1]]), 0.0)) + d = np.sqrt(_sqdist(z_norm, z_norm[chosen[-1]])) np.minimum(dist, d, out=dist) with np.errstate(divide="ignore"): s_j = np.log(rng.random(n)) / dist @@ -417,22 +418,14 @@ def _seed_paper_nearest( def _seed_kmeanspp(z_norm: np.ndarray, k: int, rng: np.random.Generator) -> np.ndarray: - """k-means++ on unit vectors: each seed drawn with probability ∝ squared distance to the nearest seed.""" - n = z_norm.shape[0] - chosen = [int(rng.integers(n))] - closest = np.full(n, np.inf, dtype=z_norm.dtype) - while len(chosen) < k: - d2 = np.maximum(2.0 - 2.0 * (z_norm @ z_norm[chosen[-1]]), 0.0) - np.minimum(closest, d2, out=closest) - cum = np.cumsum(closest, dtype=np.float64) - # all remaining cells coincide with a seed: fall back to a uniform draw - pick = ( - np.searchsorted(cum, rng.random() * cum[-1], side="right") - if cum[-1] > 0 - else rng.integers(n) - ) - chosen.append(int(min(pick, n - 1))) - return z_norm[chosen].copy() + """k-means++ seeding, scikit-learn's with one candidate per step (the classic rule).""" + rs = _legacy_random_state(rng, always_state=True) + return z_norm[kmeans_plusplus(z_norm, k, random_state=rs, n_local_trials=1)[1]] + + +def _sqdist(z_norm: np.ndarray, y: np.ndarray) -> np.ndarray: + """Squared Euclidean distances between unit rows: ``2 - 2 a.b``, clipped at 0.""" + return np.maximum(2.0 - 2.0 * (z_norm @ y.T), 0.0) def _seed_afkmc2( @@ -444,7 +437,7 @@ def _seed_afkmc2( """ n = z_norm.shape[0] first = int(rng.integers(n)) - d2 = np.maximum(2.0 - 2.0 * (z_norm @ z_norm[first]), 0.0).astype(np.float64) + d2 = _sqdist(z_norm, z_norm[first]).astype(np.float64) total = d2.sum() q = 0.5 * d2 / total + 0.5 / n if total > 0 else np.full(n, 1.0 / n) candidates = rng.choice(n, size=(k - 1, chain), p=q) @@ -454,7 +447,7 @@ def _seed_afkmc2( for c in range(1, k): cand = z_norm[candidates[c - 1]] # d² to the nearest seed so far - dist = np.maximum(2.0 - 2.0 * (cand @ seeds[:c].T), 0.0).min(axis=1) + dist = _sqdist(cand, seeds[:c]).min(axis=1) weight = dist / q[candidates[c - 1]] cur = 0 for j in range(1, chain): @@ -474,18 +467,13 @@ def _lloyd_spherical(z_norm: np.ndarray, y: np.ndarray, steps: int) -> np.ndarra shape=(y.shape[0], cells.size), ) sums = np.asarray(onehot @ z_norm) - norms = np.linalg.norm(sums, axis=1, keepdims=True) # empty clusters keep their centroid - y = np.where(norms > 0, sums / np.maximum(norms, 1e-12), y).astype(y.dtype) + y = np.where( + sums.any(axis=1, keepdims=True), _normalize_rows_l2(sums), y + ).astype(y.dtype) return y -def _joint_categories(codes: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Index of each cell's combination of batch levels, and the levels (J x C) of every combination.""" - levels, joint = np.unique(codes, axis=0, return_inverse=True) - return np.ascontiguousarray(joint.ravel().astype(np.int32)), levels - - class _Plan(NamedTuple): levels: np.ndarray # J x C, the levels of every joint category joint: np.ndarray # joint category of every cell @@ -502,7 +490,8 @@ def _correction_plan(codes: np.ndarray, n_batches: int) -> _Plan: Few joint categories (any single key, nested keys) take one statistics pass and one coefficient row per category. Many (crossed keys) take one pass per key, so nothing larger than levels x K x d is held. """ - joint, levels = _joint_categories(codes) + levels, joint = np.unique(codes, axis=0, return_inverse=True) + joint = np.ascontiguousarray(joint.ravel(), dtype=np.int32) joint_order = np.argsort(joint, kind="stable") # the cut-off keeps the per-category arrays within twice the per-level ones by_joint = levels.shape[0] <= 2 * n_batches @@ -542,19 +531,18 @@ def _correction_weights( # masses per joint category: a pass without the PC dimension no_pc = np.empty((x.shape[0], 0), dtype=x.dtype) _, o_joint = _joint_stats(no_pc, r, joint, joint_order, n_joint, N_CHUNKS) - w = _solve_correction(s, o_joint, levels, n_levels, lam, _SUPPRESS_PENALTY) + w = _solve_correction(s, o_joint, levels, n_levels, lam) if by_joint: w = w[levels].sum(axis=1) return np.ascontiguousarray(w, dtype=x.dtype) -def _solve_correction( # noqa: PLR0915, PLR0917 +def _solve_correction( # noqa: PLR0915 s: np.ndarray, o_joint: np.ndarray, levels: np.ndarray, n_levels: np.ndarray, lam: np.ndarray, - sentinel: float, ) -> np.ndarray: """Ridge coefficients of every cluster per batch level (L x K x d, float64), without the intercept. @@ -562,7 +550,7 @@ def _solve_correction( # noqa: PLR0915, PLR0917 Per cluster the normal equations are ``(Phi* diag(r_k) Phi*^T + Lambda_k) W_k = Phi* diag(r_k) X^T`` with an unpenalised intercept. The block of the covariate with the most levels is diagonal, so it is eliminated in closed form and only the Schur complement over the intercept and the remaining - levels is solved; with one covariate that is a 1 x 1 system. ``lam[b, k] >= sentinel`` removes level b + levels is solved; with one covariate that is a 1 x 1 system. ``lam[b, k] >= _SUPPRESS_PENALTY`` removes level b from cluster k (its coefficient is 0). Runs in float64, ``_SOLVE_CHUNK`` clusters at a time, so the dense systems stay at chunk x m x n_big. """ @@ -596,7 +584,7 @@ def _solve_correction( # noqa: PLR0915, PLR0917 o = o_all[:, ks] # J x c c = o.shape[1] kl = np.arange(c) - pruned = lam[:, ks] >= sentinel + pruned = lam[:, ks] >= _SUPPRESS_PENALTY lam_big, pruned_big = lam[big_lv, ks].T, pruned[big_lv].T # c x n_big lam_small, pruned_small = lam[small, ks].T, pruned[small].T # c x small row_k = kl * m + rows[:, :, None] # J x C x c @@ -629,14 +617,8 @@ def _solve_correction( # noqa: PLR0915, PLR0917 schur[kk, ss + 1, ss + 1] = 1 empty = gram[:, 0, 0] <= 1e-12 # cluster without cells: no correction schur[empty] = np.eye(m) - try: - w_small = np.linalg.solve(schur, rhs) # c x m x d - except np.linalg.LinAlgError: - # guard only: from `harmonize` every unpruned lambda is positive and the system is positive definite - w_small = np.stack([ - np.linalg.lstsq(a, b, rcond=None)[0] - for a, b in zip(schur, rhs, strict=True) - ]) + # positive definite: every unpruned lambda is positive + w_small = np.linalg.solve(schur, rhs) # c x m x d w_big = d_inv[:, :, None] * (rhs_big - coupling.transpose(0, 2, 1) @ w_small) w_small[empty] = 0 w_big[empty] = 0 @@ -647,24 +629,6 @@ def _solve_correction( # noqa: PLR0915, PLR0917 return w -def _correction_multi( - x: np.ndarray, - batch_codes: np.ndarray, - n_batches: int, - r: np.ndarray, - *, - lambda_kb: np.ndarray, -) -> np.ndarray: - """Ridge correction of ``x`` for any number of batch keys, as one Harmony round applies it.""" - # per-covariate level counts from the code ranges; an unused level has no mass, - # so where its range is cut does not change the result - ends = np.append(batch_codes.max(axis=0)[:-1] + 1, n_batches) - n_levels = np.diff(ends, prepend=0) - plan = _correction_plan(batch_codes, n_batches) - w = _correction_weights(x, r, plan, n_levels, lambda_kb) - return _apply_only(x, r, plan.cat, w, np.empty_like(x), N_CHUNKS) - - def _masses(idx: np.ndarray, o: np.ndarray, shape: tuple[int, ...]) -> np.ndarray: """Sum the joint masses ``o`` (J x c) into a dense array of ``shape`` at flat indices ``idx`` (J x ... x c).""" weights = np.broadcast_to( diff --git a/tests/test_harmony.py b/tests/test_harmony.py index 30f63ad55a..371584865e 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -22,7 +22,6 @@ from scanpy.preprocessing._harmony.core import ( _SUPPRESS_PENALTY, _compute_lambda_kb, - _correction_multi, _get_batch_codes, _get_theta_array, ) @@ -295,6 +294,17 @@ def _dense_correction( return expected +def _correction_multi(x, codes, n_batches, r, *, lambda_kb): + """One Harmony round's correction through the production path.""" + # level counts from the code ranges; an unused level has no mass, so where its range ends does not matter + n_levels = np.diff(np.append(codes.max(axis=0)[:-1] + 1, n_batches), prepend=0) + plan = core._correction_plan(codes, n_batches) + w = core._correction_weights(x, r, plan, n_levels, lambda_kb) + out = np.empty_like(x) + _apply_correction(x, r, plan.cat, w, out, out, core.N_CHUNKS, False) # noqa: FBT003 + return out + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_harmony_multikey_correction_matches_dense_design( dtype: type[np.floating], @@ -329,31 +339,6 @@ def test_harmony_multikey_correction_matches_dense_design( np.testing.assert_allclose(result, expected, atol=atol, rtol=atol) -def test_harmony_multikey_correction_singular_fallback( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A system the solver rejects as singular falls back to least squares.""" - rng = np.random.default_rng(734) - x = rng.normal(size=(12, 3)) - codes = np.column_stack(( - np.resize([0, 1], len(x)), - np.resize([2, 3, 4], len(x)), - )).astype(np.int32) - r = rng.random(size=(len(x), 2)) - r /= r.sum(axis=1, keepdims=True) - lambda_kb = np.ones((5, 2)) - - expected = _correction_multi(x, codes, 5, r, lambda_kb=lambda_kb) - - def singular(*_args, **_kwargs): - raise np.linalg.LinAlgError - - monkeypatch.setattr(np.linalg, "solve", singular) - result = _correction_multi(x, codes, 5, r, lambda_kb=lambda_kb) - - np.testing.assert_allclose(result, expected, atol=1e-12, rtol=1e-12) - - @pytest.mark.parametrize("ridge_lambda", [0.0, -0.1, float("inf"), float("nan")]) @pytest.mark.parametrize("multikey", [False, True]) def test_harmony1_rejects_bad_ridge_lambda( @@ -665,7 +650,11 @@ def test_harmony2_ircolitis_reference( # Kernels and driver of the numba implementation _N_BATCH, _K, _SIGMA = 3, 5, 0.1 -_DTYPES = [np.float32, np.float64] + + +@pytest.fixture +def dtype(float_dtype): + return float_dtype def _toy( @@ -700,11 +689,6 @@ def _state(dtype: DTypeLike, keys: str | list[str] = "batch", seed: int = 1): return x, z, codes, n_levels, r, o, n_b, n_b / dtype(len(x)) -def _bounds(n: int, n_blocks: int = 19) -> np.ndarray: - sizes = [len(a) for a in np.array_split(np.arange(n), n_blocks)] - return np.concatenate([[0], np.cumsum(sizes)]).astype(np.int64) - - def _clustering_iteration( # noqa: PLR0917 z, codes, pr_b, r, o, theta, perm, *, stabilized, sum_penalties ): @@ -733,7 +717,6 @@ def _clustering_iteration( # noqa: PLR0917 return r, o -@pytest.mark.parametrize("dtype", _DTYPES) @pytest.mark.parametrize( ("keys", "stabilized", "n_chunks", "scatter", "sum_penalties"), [ @@ -757,7 +740,13 @@ def test_cluster_pass_matches_numpy( r, o, r_sum = r0.copy(), o0.copy(), r0.sum(axis=0) dist, ent, o_fresh, y_t = _cluster_pass( *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b, theta, perm), - *(_bounds(len(x)), dtype(-2 / _SIGMA), stabilized, n_chunks, scatter), + *( + core._block_bounds(len(x), 0.05), + dtype(-2 / _SIGMA), + stabilized, + n_chunks, + scatter, + ), *(sum_penalties, True), ) atol = 2e-5 if dtype == np.float32 else 1e-10 @@ -784,7 +773,13 @@ def test_cluster_pass_matches_numpy( ) *_, y_out = _cluster_pass( *(z, y_in, r2, codes, o2, r_sum2, pr_b, theta, perm), - *(_bounds(len(x)), dtype(-2 / _SIGMA), stabilized, n_chunks, scatter), + *( + core._block_bounds(len(x), 0.05), + dtype(-2 / _SIGMA), + stabilized, + n_chunks, + scatter, + ), *(sum_penalties, False), ) np.testing.assert_array_equal(r2, r) @@ -798,7 +793,6 @@ def test_cluster_pass_matches_numpy( assert np.abs(r_ref - other).max() > 1e-3 -@pytest.mark.parametrize("dtype", _DTYPES) def test_cluster_pass_sum_rule_handles_large_penalties(dtype) -> None: """Summed penalty factors far above exp(80) keep their order between clusters.""" x, z, codes, _, r0, o0, _, pr_b = _state(dtype, ["batch", "half"]) @@ -817,7 +811,15 @@ def test_cluster_pass_sum_rule_handles_large_penalties(dtype) -> None: r, o, r_sum = r0.copy(), o0.copy(), r0.sum(axis=0) _cluster_pass( *(z, np.ascontiguousarray(y.T), r, codes, o, r_sum, pr_b, theta, perm), - *(_bounds(len(x)), dtype(-2 / _SIGMA), False, 32, False, True, True), + *( + core._block_bounds(len(x), 0.05), + dtype(-2 / _SIGMA), + False, + 32, + False, + True, + True, + ), ) np.testing.assert_allclose(r, r_ref, atol=5e-5 if dtype == np.float32 else 1e-10) @@ -856,7 +858,6 @@ def _random_design(rng: np.random.Generator, n: int, sizes: tuple[int, ...]): return codes, n_levels, x, r -@pytest.mark.parametrize("dtype", _DTYPES) # big block in the middle / first / last / a size-1 key / one key @pytest.mark.parametrize("sizes", [(4, 9, 3), (9, 4, 3), (4, 3, 9), (1, 8), (7,)]) def test_correction_with_pruning_and_an_empty_cluster( @@ -911,7 +912,6 @@ def test_correction_with_nested_keys_and_unused_levels() -> None: ) -@pytest.mark.parametrize("dtype", _DTYPES) def test_apply_correction_normalises_and_returns_centroids(dtype) -> None: x, _, codes, n_levels, r, o, n_b, pr_b = _state(dtype, ["batch", "half"]) lam = _compute_lambda_kb( @@ -926,7 +926,7 @@ def test_apply_correction_normalises_and_returns_centroids(dtype) -> None: plan = core._correction_plan(codes, int(n_levels.sum())) w = core._correction_weights(x, r, plan, n_levels, lam) z_hat, z_norm = np.empty_like(x), np.empty_like(x) - y_t = _apply_correction(x, r, plan.cat, w, z_hat, z_norm, 3) + y_t = _apply_correction(x, r, plan.cat, w, z_hat, z_norm, 3, True) # noqa: FBT003 assert z_hat.dtype == z_norm.dtype == y_t.dtype == dtype np.testing.assert_allclose( z_hat, @@ -937,7 +937,6 @@ def test_apply_correction_normalises_and_returns_centroids(dtype) -> None: np.testing.assert_allclose(y_t.T, core._normalize_rows_l2(r.T @ z_norm), atol=1e-5) -@pytest.mark.parametrize("dtype", _DTYPES) # normal, many categories across chunks, more chunks than cells (many empty) @pytest.mark.parametrize(("n_chunks", "n_joint"), [(3, 7), (32, 300), (1000, 2)]) def test_joint_stats_equal_direct_sums(dtype, n_chunks: int, n_joint: int) -> None: @@ -959,22 +958,11 @@ def test_joint_stats_equal_direct_sums(dtype, n_chunks: int, n_joint: int) -> No np.testing.assert_allclose(s, s_ref, rtol=tol, atol=tol) -def test_joint_categories() -> None: - rng = np.random.default_rng(0) - for n_keys, n_levels in ((1, 5), (3, 4), (16, 2), (40, 3)): - codes = rng.integers(n_levels, size=(3000, n_keys)) + n_levels * np.arange( - n_keys - ) - joint, levels = core._joint_categories(codes.astype(np.int32)) - assert joint.dtype == np.int32 - assert joint.flags.c_contiguous - np.testing.assert_array_equal(levels[joint], codes) - assert len(levels) == len(np.unique(codes, axis=0)) == joint.max() + 1 - - -@pytest.mark.parametrize("keys", ["batch", ["batch", "half"]]) +# "well": 250 tiny batches, which take the scatter path +@pytest.mark.parametrize("keys", ["batch", ["batch", "half"], "well"]) def test_harmonize_float32_matches_float64(keys) -> None: x, obs = _toy(n=1500, dtype=np.float64) + obs = obs.assign(well=pd.Categorical(np.arange(1500) % 250)) kw = dict(seeding="kmeans++", rng=0) # the same seeds in both precisions got64, info = core.harmonize(x, obs, keys, **kw) got32, _ = core.harmonize(x.astype(np.float32), obs, keys, **kw) @@ -988,25 +976,6 @@ def test_harmonize_float32_matches_float64(keys) -> None: assert rel.max() < 1e-3 -def test_harmonize_many_small_batches() -> None: - """A design the rule sends through the scatter path agrees across precisions.""" - assert not core._use_scatter(166, 1, 584_944 // 19) # atlas-like: accumulators - assert core._use_scatter( - 4000, 1, 20_000 // 19 - ) # thousands of tiny batches: scatters - assert core._use_scatter(41, 1, 10 * core.N_CHUNKS) - assert not core._use_scatter(40, 1, 10 * core.N_CHUNKS) - x, obs = _toy(n=1500, dtype=np.float64) - obs = obs.assign(well=pd.Categorical(np.arange(1500) % 250)) - assert core._use_scatter(250, 1, 1500 // 19) - got64, _ = core.harmonize(x, obs, "well", seeding="kmeans++", rng=0) - got32, _ = core.harmonize( - x.astype(np.float32), obs, "well", seeding="kmeans++", rng=0 - ) - rel = np.linalg.norm(got32 - got64, axis=0) / np.linalg.norm(got64, axis=0) - assert rel.max() < 1e-3 - - def test_harmonize_stop_rules() -> None: assert not core._converged([10.0], 1e-2, window=1, signed=False) assert core._converged([10.0, 10.05], 1e-2, window=1, signed=False) @@ -1021,7 +990,7 @@ def test_harmonize_stop_rules() -> None: objs = info["objectives"] assert len(objs) == len(info["iterations"]) + 1 # every round recorded assert max(info["iterations"]) <= 200 - tol = core.STOP_RULES[rule][2]["harmony2"][2] + tol = core._DEFAULTS["harmony1" if rule.startswith("paper") else "harmony2"][2] passed = [old - new < tol * abs(old) for old, new in pairwise(objs)] assert not any(passed[:-1]) # no earlier round met the stop test # the last one did, or the cap ended the run @@ -1066,8 +1035,7 @@ def afkmc2(*_args, **_kwargs): def test_harmonize_edge_inputs() -> None: x, obs = _toy() - for extra in (dict(tau=5), dict(stop_rule="paper_window3"), dict(n_clusters=2)): - assert np.isfinite(core.harmonize(x, obs, "batch", rng=0, **extra)[0]).all() + assert np.isfinite(core.harmonize(x, obs, "batch", n_clusters=2, rng=0)[0]).all() # a one-cell batch obs1 = obs.assign(batch=pd.Categorical(["a"] * (len(x) - 1) + ["b"])) assert np.isfinite(core.harmonize(x, obs1, "batch", rng=0)[0]).all() @@ -1078,7 +1046,6 @@ def test_harmonize_edge_inputs() -> None: @pytest.mark.skipif( numba.config.NUMBA_NUM_THREADS < 2, reason="needs at least two numba threads" ) -@pytest.mark.parametrize("dtype", _DTYPES) def test_harmonize_does_not_depend_on_thread_count(dtype) -> None: x, obs = _toy(dtype=dtype) previous = numba.get_num_threads() @@ -1196,15 +1163,12 @@ def test_harmonize_seeding_options() -> None: @pytest.mark.parametrize( ("kwargs", "match"), [ - (dict(alpha=float("nan")), "alpha"), - (dict(flavor="harmony1", ridge_lambda=0.0), "ridge_lambda"), (dict(block_proportion=2.0), "block_proportion"), (dict(sigma=0.0), "sigma"), (dict(max_iter_harmony=0), "max_iter_harmony"), (dict(seeding="nope"), "seeding must be one of"), (dict(penalty_combine="mean"), "penalty_combine must be one of"), (dict(stop_rule="never"), "stop_rule must be one of"), - (dict(flavor="harmony3"), "flavor must be one of"), (dict(n_clusters=0), "n_clusters must be between 1"), (dict(n_clusters=601), "n_clusters must be between 1"), ], @@ -1240,7 +1204,10 @@ def test_harmonize_rejects_bad_inputs() -> None: def test_block_bounds_match_array_split() -> None: for n in (7, 40, 600, 1001): n_blocks = int(min(n, 1 // 0.05)) - np.testing.assert_array_equal(core._block_bounds(n, 0.05), _bounds(n, n_blocks)) + sizes = [len(a) for a in np.array_split(np.arange(n), n_blocks)] + np.testing.assert_array_equal( + core._block_bounds(n, 0.05), np.cumsum([0, *sizes]) + ) def _closest_pair(y: np.ndarray) -> float: @@ -1274,7 +1241,6 @@ def test_seeders_survive_identical_cells() -> None: assert np.isfinite(seeder(z, 5, np.random.default_rng(0))).all() -@pytest.mark.parametrize("dtype", _DTYPES) def test_lloyd_polish_lowers_the_kmeans_cost(dtype) -> None: _, z, *_ = _state(dtype) y = core._seed_afkmc2(z, _K, np.random.default_rng(0))