atlas.tl.CellRankExtension

atlas.tl.CellRankExtension#

class atlas.tl.CellRankExtension(mudata)#

Trajectory inference using the CellRank Pseudotime Kernel on MuData objects.

This class provides a high-level interface to run Psseudotime-based trajectory inference on multimodal single-cell data stored in a MuData. It operates directly on the input object and stores all intermediate and final results within it.

Parameters:

mudata (MuData) – Annotated multimodal data object containing precomputed neighborhood graphs and embeddings.

Notes

This implementation is adapted from the original CellRank algorithm [WLK+24], with modifications to operate on MuData objects.

Attributes table#

cluster_key

Key in mudata.obs storing cell-wise categories.

fate_key

Key in mudata.obsm storing cell-fate probabilities.

kernel

Underlying PseudotimeKernel object.

mudata

Underlying MuData object.

time_key

Key in mudata.obs storing cell-wise time.

Methods table#

compute_kernel([connectivity_key, time_key, ...])

Compute a pseudotime-based transition kernel.

run([n_components, initial_states, ...])

Run trajectory inference using GPCCA on a precomputed kernel.

Attributes#

CellRankExtension.cluster_key#

Key in mudata.obs storing cell-wise categories.

CellRankExtension.fate_key#

Key in mudata.obsm storing cell-fate probabilities.

CellRankExtension.kernel#

Underlying PseudotimeKernel object.

CellRankExtension.mudata#

Underlying MuData object.

Represents multimodal single-cell data containing the weighted nearest neighbor graph.

CellRankExtension.time_key#

Key in mudata.obs storing cell-wise time.

Methods#

CellRankExtension.compute_kernel(connectivity_key='wnn_connectivities', time_key='pseudotime', cluster_key=None, backward=False, threshold_scheme='hard', frac_to_keep=0.3, b=10.0, nu=0.5, n_jobs=-1, backend='loky')#

Compute a pseudotime-based transition kernel.

This function constructs a cellrank.kernels.PseudotimeKernel using a precomputed connectivity graph and pseudotime stored in the underlying MuData object. The kernel is used to model directed transitions between cells based on their pseudotemporal ordering.

Parameters:

  • connectivity_key (str (default: 'wnn_connectivities')) – Key in mudata.obsp where the connectivity matrix corresponding to the WNN is stored.

  • time_key (str (default: 'pseudotime')) – Key in mudata.obs where pseudotime values are stored.

  • cluster_key (str | None (default: None)) – Optional key in mudata.obs containing cell annotations. If provided, clusters are treated as categorical groups.

  • backward (bool (default: False)) – Whether to compute the backward process instead of the forward pseudotime dynamics. Refer to [WLK+24] for more information.

  • threshold_scheme (Literal['soft', 'hard'] (default: 'hard')) – Scheme used to sparsify the transition matrix. One of "soft" or "hard".

  • frac_to_keep (float (default: 0.3)) – Fraction of transitions to retain when using thresholding.

  • b (float (default: 10.0)) – Parameter controlling the steepness of the logistic function used in soft thresholding.

  • nu (float (default: 0.5)) – Parameter controlling the width of the kernel.

  • n_jobs (int (default: -1)) – Number of parallel jobs used for computation.

  • backend (str (default: 'loky')) – Which backend to use for multiprocessing (‘loky’, ‘multiprocessing’, ‘threading’).

Return type:

None

Returns:

Updates the MuData object:

Raises:

  • KeyError – If connectivity_key is not present in mudata.obsp.

  • KeyError – If cluster_key is provided but not present in mudata.obs.

Notes

This method creates a temporary AnnData object to interface with CellRank, while preserving the original MuData structure.

CellRankExtension.run(n_components=20, initial_states=None, terminal_states=None, initial_distribution=None, method='krylov', sorting_strategy='LR', eigengap_weight=1.0, verbose=None, n_cells=30, n_states=None, allow_overlap=False, alpha=1.0, stability_threshold=0.96, n_terminal_states=None, n_initial_states=1, solver='gmres', use_petsc=True, n_jobs=-1, tol=1e-06, preconditioner=None, backend='loky')#

Run trajectory inference using GPCCA on a precomputed kernel.

This method applies the Generalized Perron Cluster Cluster Analysis (GPCCA) algorithm to identify macrostates, infer initial and terminal states, and compute fate probabilities. Results are stored directly in the underlying MuData object.

Depending on the provided inputs, the method either:

  • uses user-defined initial and terminal states, or

  • automatically infers macrostates and predicts initial and terminal states.

Parameters:

  • n_components (int (default: 20)) – Number of Schur vectors to compute for the coarse-graining step.

  • initial_states (dict[str, Sequence[str]] | None (default: None)) – Mapping of initial states to cell identifiers. If None, initial states are inferred automatically.

  • terminal_states (dict[str, Sequence[str]] | None (default: None)) – Mapping of terminal states to cell identifiers. If None, terminal states are inferred automatically.

  • initial_distribution (ndarray | None (default: None)) – Optional initial distribution over cells used for the Schur decomposition.

  • method (Literal['krylov', 'brandts'] (default: 'krylov')) – Method used to compute the Schur decomposition.

  • sorting_strategy (Literal['LM', 'LR'] (default: 'LR')) – Strategy to sort eigenvalues (e.g. largest magnitude "LM" or largest real part "LR").

  • eigengap_weight (float (default: 1.0)) – Weight controlling the eigengap heuristic used during coarse-graining. Refer to cellrank.estimators.GPCCA.compute_schur() for more information.

  • verbose (bool | None (default: None)) – Whether to print progress information.

  • n_cells (int (default: 30)) – Number of cells to sample when identifying representative states.

  • n_states (int | Sequence[int] | None (default: None)) – Number of macrostates to compute. If None, it is determined automatically according to minChi and crispness, refer to [LBK+22] and cellrank.estimators.GPCCA.compute_macrostates() for more information.

  • allow_overlap (bool (default: False)) – Whether cells can belong to multiple states.

  • alpha (float (default: 1.0)) – Parameter controlling terminal state prediction. Refer to cellrank.estimators.GPCCA.predict_terminal_states() and cellrank.estimators.GPCCA.predict_initial_states().

  • stability_threshold (float (default: 0.96)) – Threshold used to identify stable terminal states.

  • n_terminal_states (int | None (default: None)) – Number of terminal states to predict when not provided.

  • n_initial_states (int | None (default: 1)) – Number of initial states to predict when not provided.

  • solver (Literal['direct', 'gmres', 'lgmres', 'bicgstab', 'gcrotmk'] (default: 'gmres')) – Linear solver used for fate probability computation.

  • use_petsc (bool (default: True)) – Whether to use petsc4py or scipy for solving linear systems.

  • n_jobs (int (default: -1)) – Number of parallel jobs for the iterative solver.

  • tol (float (default: 1e-06)) – Convergence tolerance for the iterative solver.

  • preconditioner (str | None (default: None)) – Optional preconditioner for the solver. See cellrank.estimators.GPCCA.compute_fate_probabilities().

  • backend (str (default: 'loky')) – Which backend to use for multiprocessing (‘loky’, ‘multiprocessing’, ‘threading’).

Return type:

None

Returns:

Updates the original MuData object:

  • Fate probabilities stored as a pandas.DataFrame in .obsm["fate_probabilities"].

  • Initial states stored in .uns["initial_states"].

  • Terminal states stored in .uns["terminal_states"].

  • Intermediate states (if inferred) stored in .uns["intermediate_states"].

  • State-specific colors stored in .uns["fate_state_colors"].

  • Entropy measures (e.g. Shannon entropy, KL divergence) stored in .obs.

Raises:

AttributeError – If the kernel has not been computed prior to calling this method.

Notes

This implementation relies on cellrank.estimators.GPCCA to perform coarse-graining of the Markov chain and infer lineage relationships.

If both initial_states and terminal_states are provided, no automatic state inference is performed. Otherwise, macrostates and lineage-driving states are inferred from the data.