scverse / scverse/rapids-singlecell

[FEA] gr.cluster_autok: CellCharter-compatible GPU ClusterAutoK

Open
#775 0 comments 0 reactions 0 assignees View on GitHub

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 k values (with
    CellCharter's ±1 k halo)
  • adjacent-k Fowlkes–Mallows stability across independent runs, mirrored
    into CellCharter's layout
  • best_k / peaks selection from the stability curve
  • writes labels to adata.obs and plot-ready metadata to adata.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 at best_k
  • adata.uns[key_added]best_k, peaks, reported_k, fitted_k,
    stability_mean, stability_std, params (including init)

Enough to plot k vs. mean stability (± SD) without keeping fitted GMM
models in uns.

Out of scope for v1
  • Multi-k predict(k=…) without refitting (CellCharter's class-based
    feature — our function writes best_k labels only; a different k needs a
    new call)
  • A plotting helper (pl.autok_stability or 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.pycluster_autok() +
    internal _ClusterAutoK / _ClusterAutoKGMM engine
  • tests/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.mdgr.cluster_autok entry
  • docs/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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.