scverse / scverse/rapids-singlecell
[FEA] gr.cluster_autok: CellCharter-compatible GPU ClusterAutoK
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 394
- Forks
- 52
- Avg merge
- 22h 5m
- Merged PRs (30d)
- 20
Description
Summary
Add a function that runs ClusterAutoK on GPU:
- repeated full-covariance GMM fits across candidate
kvalues (with
CellCharter's ±1khalo) - adjacent-
kFowlkes–Mallows stability across independent runs, mirrored
into CellCharter's layout best_k/peaksselection from the stability curve- writes labels to
adata.obsand plot-ready metadata toadata.uns
This mirrors CellCharter's semantics (cellcharter.tl.ClusterAutoK), not
its class-based, PyTorch-Lightning-backed API.
It builds directly on the private full-covariance GMM added in #644 (thanks
@Intron7 — noticed while researching this that
you already built a "minimal GMM to do the cellcharter workflow" there,
which is exactly the primitive AutoK needs). Since #758 also just split
calculate_niche into flavor-specific functions, this seemed like a natural
moment to propose the AutoK companion for the CellCharter flavor.
Placement: proposing this under gr (alongside
calculate_niche_cellcharter) rather than tl, since it shares that
function's private squidpy_gpu._gmm dependency and CellCharter provenance,
and the two are commonly chained on a niche embedding. Open to tl instead
if you'd rather keep gr scoped to spatial-graph-specific tools — AutoK
itself doesn't require spatial coordinates, just an embedding.
Proposed API
rsc.gr.cluster_autok(
adata,
use_rep="X_cellcharter",
n_clusters=(2, 10),
max_runs=10,
convergence_tol=0.0,
rng=0,
init="sklearn_kmeans",
key_added="cluster_autok",
copy=False,
)
rng follows the numpy.random.Generator-based convention #761 introduced
across the scanpy-aligned APIs (each of the max_runs repetitions draws an
independent seed from it). copy follows the existing tl.kmeans/tl.leiden
convention (the gr.calculate_niche* family's inplace is scoped to that
module per #758, not project-wide as far as I can tell — happy to be
corrected, or to switch to inplace if that's actually the intended
direction for gr specifically).
Written fields
adata.obs[key_added]— labels atbest_kadata.uns[key_added]—best_k,peaks,reported_k,fitted_k,
stability_mean,stability_std,params(includinginit)
Enough to plot k vs. mean stability (± SD) without keeping fitted GMM
models in uns.
Out of scope for v1
- Multi-
kpredict(k=…)without refitting (CellCharter's class-based
feature — our function writesbest_klabels only; a differentkneeds a
new call) - A plotting helper (
pl.autok_stabilityor similar) — happy to follow up
once the data shape is settled - Replacing the private GMM internals with a new public estimator (see below)
Default init
sklearn_kmeans (CellCharter parity): sklearn k-means++ seeding, then
CUDA EM. Alternatives: kmeans (native cuML KMeans seeding, fully on-device,
faster) and random_from_data.
On a real spatial dataset (ATERA, PCA50 embedding, 10 runs, k=2..10):
sklearn_kmeans recovers the CPU CellCharter best_k=9 (stability curve
correlation 0.972 vs. CPU); native kmeans seeding selects best_k=7
(correlation 0.430) — worth documenting init as a first-class,
consequential choice rather than an incidental kwarg.
Implementation note / ask
The prototype uses the private rapids_singlecell.squidpy_gpu._gmm
functions (_initialize_parameters, _run_em, _precision_cholesky,
_choose_e_step, _e_step) because AutoK's "keep the lowest-NLL fit per k"
selection needs the fitted parameters and per-fit negative log-likelihood —
gmm_fit_predict only returns labels.
Question for maintainers: is depending on _gmm's private functions
acceptable for an initial merge (as gr.calculate_niche_cellcharter already
implicitly does via _gmm_cluster), or would you rather _gmm grow a small
public fit surface first (e.g. a GaussianMixture-like object exposing
weights_/means_/covariances_/NLL, or an nll=True flag on
gmm_fit_predict)? Either works — just want to build on whichever
shape you'd prefer to maintain long-term.
Prototype
A working prototype (function + engine + tests) exists on a local branch
against current main, exercised end-to-end against the real
squidpy_gpu._gmm CUDA kernels (not just synthetic mocks):
src/rapids_singlecell/squidpy_gpu/_cluster_autok.py—cluster_autok()+
internal_ClusterAutoK/_ClusterAutoKGMMenginetests/test_cluster_autok.py— unit tests for the halo/mirroring helpers,
schema/behavior tests on synthetic blobs, and an end-to-end smoke test on
scanpy.datasets.pbmc68k_reduced()docs/api/squidpy_gpu.md—gr.cluster_autokentrydocs/release-notes/0.17.0.md— Features stub
Happy to open a draft PR once the API shape below is roughly confirmed —
or sooner, if you'd rather iterate on the shape in code.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/rapids_singlecell/squidpy_gpu/_cluster_autok.py and the private squidpy_gpu._gmm functions it uses. Run tests/test_cluster_autok.py, including the synthetic and pbmc68k_reduced smoke tests, and review docs/api/squidpy_gpu.md plus the release-notes stub. Done means the API shape is agreed, the tests pass, and the documented fields and behavior match the CellCharter-compatible proposal.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, scikit-learn
- Domain
- bioinformatics, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100