Repository navigation
Conversation
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.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4424 +/- ##
=======================================
Coverage ? 84.37%
=======================================
Files ? 135
Lines ? 12732
Branches ? 0
=======================================
Hits ? 10743
Misses ? 1989
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more.
|
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.
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.
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.
Describe the default-seeding test by what it checks.
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.
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.
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.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This replaces the NumPy hot loops of
sc.pp.harmony_integratewith parallel numba kernels. With the settings the current code uses, the output is the same up to rounding (float64: relative L2 per PC ≤ 2e-12). The new defaults are 6–8× faster in float64 and 7–13× in float32.What changes
_harmony/_kernels.py(new): one fused kernel per clustering iteration (soft assignment, held-out counts, objective parts and next centroids in a single pass over each block), and the correction statistics and apply step. Kernels take their dtype from the input, so float32 and float64 both compile._harmony/core.py: aharmonizedriver replaces theHarmonyclass. One correction path handles any number of batch keys: per cluster, the covariate with the most levels is eliminated in closed form and only the small Schur complement is solved (float64; it has no cell dimension)._correction_fastis gone;_correction_multikeeps its signature and now runs this path, so the existing dense-design test (1e-11 in float64) exercises the production code.New parameters (defaults first)
seeding"afkmc2","kmeans","kmeans++","paper_nearest""kmeans"is the current behaviourstop_rule"reference_record","reference","paper","paper_window3""reference"is the current behaviour (round objective recorded only when clustering converged);"reference_record"records every round, as rapids-singlecell doespenalty_combine"product","sum""sum"is the rule as printed in Patikas et al. 2026seeding_lloyd_steps0seeding_cells_per_clusterNonerecompute_countsTruelow_memoryFalseDefault change:
seeding="afkmc2"andstop_rule="reference_record". The coordinates change (per-PC r 0.965–0.985 against the current output), but integration quality does not (table below).seeding="kmeans", stop_rule="reference"reproduces the current output.Benchmarks
Apple M-series, 8 threads, one batch key,
dtype=float64unless noted, numba compiled outside the timing:Two batch keys on HLCA (
sample+dataset), where the current code spends most of its time in the multi-key correction: float32 304.8 s → 6.3 s (the new code takes 9.5 s in float64). On Linux (Intel Xeon 6248R and AMD EPYC 9745, OpenBLAS and MKL, float32, one key), HLCA takes 1.7–3.9 s at 32 threads. Memory scales linearly, about 1 GB per million cells at 50 PCs in float32.Integration quality, current → new defaults (float64; float32 agrees to 4 decimals):