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/docs/release-notes/4424.feat.md b/docs/release-notes/4424.feat.md new file mode 100644 index 0000000000..f754c5b5fa --- /dev/null +++ b/docs/release-notes/4424.feat.md @@ -0,0 +1,3 @@ +Speed up {func}`~scanpy.pp.harmony_integrate` 6–13× with parallel numba kernels, and add the `seeding`, `stop_rule`, `penalty_combine`, `seeding_lloyd_steps`, `seeding_cells_per_cluster`, `recompute_counts` and `low_memory` parameters. +For `flavor="harmony2"`, the default seeding is now `"afkmc2"`, and the default stop rule for both flavors is `"reference_record"`; pass `seeding="kmeans", stop_rule="reference"` for the previous behaviour. +Cells with an all-zero embedding are now left out and returned unchanged {smaller}`T Treis` diff --git a/src/scanpy/preprocessing/_harmony/__init__.py b/src/scanpy/preprocessing/_harmony/__init__.py index 58e09f3939..9d2ed8d2c6 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,22 @@ def harmony_integrate( # noqa: PLR0913 batch_prune_threshold: float | None = 1e-5, correction_method: Literal["fast"] = "fast", block_proportion: float = 0.05, + stop_rule: Literal[ + "reference_record", "reference", "paper", "paper_window3" + ] = "reference_record", + seeding: Literal["afkmc2", "kmeans", "kmeans++", "paper_nearest"] | None = None, + seeding_lloyd_steps: int = 0, + seeding_cells_per_cluster: int | None = None, + penalty_combine: Literal["product", "sum"] = "product", + recompute_counts: bool = True, + low_memory: bool = False, rng: SeedLike | RNGLike | None = None, ) -> 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. 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, @@ -87,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. @@ -105,15 +99,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 +152,52 @@ def harmony_integrate( # noqa: PLR0913 Set to ``None`` to disable pruning. correction_method Method for the correction step. - ``"fast"`` uses a precomputed factorization that avoids the full inversion, - which can be faster for datasets with many batches. - Multiple keys automatically use the exact general-design solve, - because this optimization only applies to a single batch variable. + Only ``"fast"`` is supported: the ridge regression of every cluster is solved exactly, + eliminating the batch variable with the most levels in closed form, + for one or several batch keys. block_proportion Proportion of cells updated per clustering sub-iteration. Smaller values produce more stochastic updates. Larger values are faster but may converge to different solutions. + stop_rule + When the clustering and the outer Harmony loop stop. + ``"reference_record"`` (default) and ``"reference"`` use the iteration limits and tolerances + listed under ``max_iter_clustering``, ``tol_clustering`` and ``tol_harmony``, + and test clustering convergence over a window of 3 iterations. + ``"reference_record"`` records every round's objective for the outer test; + ``"reference"`` records it only when the clustering converged. + ``"paper"`` and ``"paper_window3"`` use the tolerances of :cite:t:`Korsunsky2019` + (200 clustering iterations, ``1e-5`` and ``1e-4``) for both flavors, + testing clustering convergence on the last iteration or a window of 3. + seeding + How the initial cluster centroids are chosen. + If :data:`None`, ``"afkmc2"`` for ``"harmony2"`` and ``"kmeans"`` for ``"harmony1"``. + ``"afkmc2"`` approximates k-means++ seeding with Markov chains :cite:p:`Bachem2016` + and needs a single pass over the cells. + ``"kmeans"`` runs :class:`~sklearn.cluster.KMeans` (k-means++ and up to 25 Lloyd iterations). + ``"kmeans++"`` is exact k-means++ seeding. + ``"paper_nearest"`` follows the sampler of :cite:t:`Patikas2026`, + taking the cell with the *highest* ``log(u) / D`` (the paper prints "lowest"), + where ``D`` is the distance to the nearest centroid chosen so far: + this draws cells with probability proportional to that distance. + seeding_lloyd_steps + Hard spherical k-means iterations applied to the seeds before Harmony starts. + seeding_cells_per_cluster + If set, ``seeding="kmeans"`` and the ``seeding_lloyd_steps`` run on a random subsample of + ``seeding_cells_per_cluster * n_clusters`` cells. + penalty_combine + How the diversity penalties of several batch keys are combined in the soft assignment. + ``"product"`` (default) multiplies the per-key penalty factors, as in :cite:t:`Korsunsky2019`; + ``"sum"`` adds them, as the assignment rule is printed in :cite:t:`Patikas2026`, + leaving out levels with ``theta == 0``. + Identical for a single key. + recompute_counts + Rebuild the per-batch cluster counts from the assignments after every clustering iteration. + If ``False``, the counts are only updated incrementally, which accumulates rounding error + in float32. + low_memory + Drop one cell x component buffer during the iterations + at the cost of one extra pass at the end; the result is identical. rng Random number generator or seed for deterministic behavior. @@ -172,26 +207,11 @@ 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"}: - msg = f"flavor must be 'harmony1' or 'harmony2', got {flavor!r}." - raise ValueError(msg) if correction_method != "fast": msg = f"correction_method must be 'fast', got {correction_method!r}." raise ValueError(msg) - stabilized_penalty = flavor == "harmony2" - dynamic_lambda = flavor == "harmony2" - - # Unset stopping rules follow the flavor's reference implementation - defaults = Stopping[flavor.upper()] - if max_iter_clustering is None: - max_iter_clustering = defaults.max_iter_clustering - if tol_clustering is None: - tol_clustering = defaults.tol_clustering - if tol_harmony is None: - tol_harmony = defaults.tol_harmony # Warn when flavor-incompatible parameters are explicitly set if flavor == "harmony2" and ridge_lambda != 1.0: @@ -230,35 +250,29 @@ def harmony_integrate( # noqa: PLR0913 msg = f"Could not convert input of type {type(input_data).__name__} to NumPy array." raise TypeError(msg) from e - # Check for NaN values - if np.isnan(x).any(): - msg = ( - "Input data contains NaN values. Please handle these before " - "running harmony_integrate." - ) - raise ValueError(msg) - - # Run Harmony - 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..223fdc2295 --- /dev/null +++ b/src/scanpy/preprocessing/_harmony/_kernels.py @@ -0,0 +1,313 @@ +"""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 +from numba.extending import register_jitable + + +@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, + want_y, +): + """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``: 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] + 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 + 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 + ) + 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)) + # 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) + 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: + # log of the summed factors of the keys with theta > 0, shifted by their maximum + for k in range(n_clusters): + 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] + 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] + 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 + 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 + # 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) + 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[dd, k] = v + sq += v * v + scale = ft(1.0) / max(np.sqrt(sq), ft(1e-12)) + for dd in range(n_pc): + y_t[dd, k] *= scale + return y_t + + +@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, 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. + + 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] + 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): + for dd in range(n_pc): + z_hat[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): + 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 = 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] + 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 02a37bf09c..0e2564a101 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 typing import TYPE_CHECKING +from itertools import pairwise +from typing import TYPE_CHECKING, NamedTuple import numpy as np -from sklearn.cluster import KMeans +from scipy import sparse +from sklearn.cluster import KMeans, kmeans_plusplus 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, _cluster_pass, _joint_stats if TYPE_CHECKING: from collections.abc import Sequence @@ -21,222 +24,293 @@ 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. +# 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), + # every round's objective is recorded, as in rapids-singlecell + "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. +N_CHUNKS = 32 +# clusters per batched solve of the correction; bounds its dense systems +_SOLVE_CHUNK = 8 + + +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 | None = None, + 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 = flavor == "harmony2" + if seeding is None: + 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) - - 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}." - ) + 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) - - # 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.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) + if not np.isfinite(x).all(): + msg = "Input data contains non-finite values." + raise ValueError(msg) + ft = x.dtype.type + window, record_always = STOP_RULES[stop_rule] + sum_penalties = penalty_combine == "sum" + 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 + 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()) + 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() + 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) + + 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, + True, # noqa: FBT003 + ) + 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: 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 + ) as bar: + for i in bar: + objs: list[float] = [] + done = False + for it in range(max_iter_clustering): + perm = rng.permutation(n_cells) + # 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() + 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, + # the correction recomputes the centroids after the last iteration + it < max_iter_clustering - 1, ) - 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=stabilized, ) - - batch_codes = self.batch_codes[:, 0] - return _correction_fast( - *(x, batch_codes, self.n_batches, r, o), - lambda_kb=lambda_kb, - ) + w = _correction_weights(x, r, plan, n_levels, lam) + 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: + _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 + z_hat = x_all + return z_hat, info def _get_batch_codes( @@ -310,178 +384,257 @@ 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: + """Seed as :cite:t:`Patikas2026` describe, read so that the seeds spread. - 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)`` + (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))] + taken = np.zeros(n, dtype=bool) + dist = np.full(n, np.inf, dtype=z_norm.dtype) + while len(chosen) < k: + taken[chosen[-1]] = True + 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 + 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++ 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( + z_norm: np.ndarray, k: int, rng: np.random.Generator, *, chain: int = 200 +) -> np.ndarray: + """AFK-MC² :cite:p:`Bachem2016`: k-means++ approximated by one Metropolis chain per seed. + 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 = _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) + 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 = _sqdist(cand, seeds[:c]).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), + ) + sums = np.asarray(onehot @ z_norm) + # empty clusters keep their centroid + y = np.where( + sums.any(axis=1, keepdims=True), _normalize_rows_l2(sums), y + ).astype(y.dtype) + return y -def _compute_r( - z: np.ndarray, - y: np.ndarray, - term: float, -) -> np.ndarray: - """Compute soft cluster assignments using NumPy dot.""" - dots = z @ y.T - return np.exp(term * (1.0 - dots)) +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 _clustering( # noqa: PLR0913 - z_norm: NDArray[np.floating], - batch_codes: np.ndarray, - pr_b: np.ndarray, - *, + +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. + """ + 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 + 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, - 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, + 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:] ) - objectives_clustering.append(obj) - - # Check convergence - if _is_convergent_clustering(objectives_clustering, tol): - obj = objectives_clustering[-1] - break else: - obj = None + 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) + if by_joint: + w = w[levels].sum(axis=1) + return np.ascontiguousarray(w, dtype=x.dtype) + + +def _solve_correction( # noqa: PLR0915 + s: np.ndarray, + o_joint: np.ndarray, + levels: np.ndarray, + n_levels: np.ndarray, + lam: np.ndarray, +) -> np.ndarray: + """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] >= _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. + """ + 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( + [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((len(levels), 1), dtype=np.int64), row_of_level[levels[:, others]]], + axis=1, + ) - return r, e, o, obj + 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) + 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] >= _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 + 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) + # 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 + 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 _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 np.bincount(idx.ravel(), weights.ravel(), int(np.prod(shape))).reshape(shape) def _compute_lambda_kb( @@ -511,155 +664,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] - - 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 +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. - 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, 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 > 4 * 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..371584865e 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,10 +13,15 @@ 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, - _correction_multi, _get_batch_codes, _get_theta_array, ) @@ -242,35 +249,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,30 +291,52 @@ def test_harmony_multikey_correction_matches_dense_design( )[0] w[0] = 0 expected -= r_k[:, None] * (design @ w) + return expected - atol = 2e-5 if dtype == np.float32 else 1e-11 - np.testing.assert_allclose(result, expected, atol=atol, rtol=atol) +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 -def test_harmony_multikey_correction_singular_fallback( - monkeypatch: pytest.MonkeyPatch, + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_harmony_multikey_correction_matches_dense_design( + dtype: type[np.floating], ) -> None: - """A system diagnosed as singular falls back to least squares.""" + """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) - 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)) + 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 = np.ones((5, 2)) + lambda_kb = rng.uniform(0.2, 1.0, size=(n_batches, n_clusters)).astype(dtype) + lambda_kb[1, 1] = dtype(_SUPPRESS_PENALTY) - expected = _correction_multi(x, codes, 5, r, lambda_kb=lambda_kb) - monkeypatch.setattr(np.linalg, "matrix_rank", lambda _: 0) - result = _correction_multi(x, codes, 5, r, lambda_kb=lambda_kb) + result = _correction_multi(x, codes, n_batches, r, lambda_kb=lambda_kb) + expected = _dense_correction(x, codes, n_batches, r, lambda_kb) - np.testing.assert_allclose(result, expected, atol=1e-12, rtol=1e-12) + atol = 2e-5 if dtype == np.float32 else 1e-11 + np.testing.assert_allclose(result, expected, atol=atol, rtol=atol) @pytest.mark.parametrize("ridge_lambda", [0.0, -0.1, float("inf"), float("nan")]) @@ -386,6 +396,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"] @@ -393,6 +406,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 the multikey fixture at least as well as the stored output. + + 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 + + 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 +645,611 @@ 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 + + +@pytest.fixture +def dtype(float_dtype): + return float_dtype + + +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 _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 + 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]) + r_sum += r[idx].sum(axis=0) + return r, o + + +@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), + *( + 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 + 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) + # 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), + *( + core._block_bounds(len(x), 0.05), + 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), + stabilized=stabilized, + sum_penalties=not sum_penalties, + ) + assert np.abs(r_ref - other).max() > 1e-3 + + +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), + *( + 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) + + +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, True), + ) + 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 + + +# 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: + 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 + 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 + ) + + +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, + ) + 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, True) # noqa: FBT003 + 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) + + +# 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) + + +# "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) + 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_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._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 + 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_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() + 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() + # fewer cells per block than chunks + assert np.isfinite(core.harmonize(x[:40], obs[:40], "batch", rng=0)[0]).all() + + +@pytest.mark.skipif( + numba.config.NUMBA_NUM_THREADS < 2, reason="needs at least two numba threads" +) +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", "half"], rng=0)) + finally: + numba.set_num_threads(previous) + 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"]]) +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]) + # 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: + 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(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(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: + 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") + 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"): + 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)) + 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: + 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() + + +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