from pydantic import (
Field,
ValidationInfo,
computed_field,
field_validator,
model_validator, BaseModel
)
from typing import Optional, List, Dict, Union, Literal, Tuple, Any, Sequence
class GPCCAFitModel(BaseModel):
"""
Input schema for CellRank's GPCCA.fit method which prepares the estimator for terminal states prediction.
"""
kernel: Literal['pseudotime', 'cytotrace', 'velocity', 'connectivity', 'realtime'] = Field(
description="Type of kernel to use."
)
n_states: Optional[Union[int, List[int]]] = Field(
default=None,
description="Number of macrostates to compute. If a list, use the minChi criterion. If None, use the eigengap heuristic."
)
n_cells: Optional[int] = Field(
default=30,
description="Number of most likely cells from each macrostate to select."
)
cluster_key: Optional[str] = Field(
default=None,
description="If a key to cluster labels is given, names and colors of the states will be associated with the clusters."
)
# # Parameters for compute_schur method n_states is n_components in compute_schur
# n_components: int = Field(
# default=20,
# description="Number of Schur vectors to compute."
# )
initial_distribution: Optional[Any] = Field(
default=None,
description="Input distribution over all cells. If None, uniform distribution is used."
)
method: Literal['krylov', 'brandts'] = Field(
default='krylov',
description="Method for calculating the Schur vectors. 'krylov' is an iterative procedure for large, sparse matrices. 'brandts' is full sorted Schur decomposition of a dense matrix."
)
which: Literal['LR', 'LM'] = Field(
default='LR',
description="How to sort the eigenvalues. 'LR' - the largest real part. 'LM' - the largest magnitude."
)
alpha: float = Field(
default=1.0,
description="Used to compute the eigengap. alpha is the weight given to the deviation of an eigenvalue from one."
)
class GPCCAPredictInitialStatesModel(BaseModel):
"""
Input schema for CellRank's GPCCA.predict_initial_states method which computes initial states from macrostates.
"""
kernel: Literal['pseudotime', 'cytotrace', 'velocity', 'connectivity', 'realtime'] = Field(
...,
description="Type of kernel to use."
)
n_states: int = Field(
default=1,
description="Number of initial states to compute."
)
n_cells: int = Field(
default=30,
description="Number of most likely cells from each macrostate to select."
)
allow_overlap: bool = Field(
default=False,
description="Whether to allow overlapping names between initial and terminal states."
)
class GPCCAPredictTerminalStatesModel(BaseModel):
"""
Input schema for CellRank's GPCCA.predict_terminal_states method which automatically selects terminal states from macrostates.
"""
kernel: Literal['pseudotime', 'cytotrace', 'velocity', 'connectivity', 'realtime'] = Field(
description="Type of kernel to use."
)
method: Literal['stability', 'top_n', 'eigengap', 'eigengap_coarse'] = Field(
default='stability',
description="How to select the terminal states. 'eigengap' - select based on eigengap of transition_matrix. "
"'eigengap_coarse' - select based on eigengap of diagonal of coarse_T. "
"'top_n' - select top n_states based on probability of diagonal of coarse_T. "
"'stability' - select states with stability >= stability_threshold."
)
n_cells: int = Field(
default=30,
description="Number of most likely cells from each macrostate to select."
)
alpha: float = Field(
default=1,
description="Weight given to the deviation of an eigenvalue from one. Only used when method = 'eigengap' or method = 'eigengap_coarse'."
)
stability_threshold: float = Field(
default=0.96,
description="Threshold used when method = 'stability'."
)
n_states: Optional[int] = Field(
default=None,
description="Number of states used when method = 'top_n'."
)
allow_overlap: bool = Field(
default=False,
description="Whether to allow overlapping names between initial and terminal states."
)
class GPCCAComputeFateProbabilitiesModel(BaseModel):
"""
Input schema for CellRank's GPCCA.compute_fate_probabilities method which calculates the probability
of each cell being absorbed in any of the terminal states.
"""
kernel: Literal['pseudotime', 'cytotrace', 'velocity', 'connectivity', 'realtime'] = Field(
description="Type of kernel to use."
)
keys: Optional[List[str]] = Field(
default=None,
description="Terminal states for which to compute the fate probabilities. If None, use all states defined in terminal_states."
)
solver: Literal['direct', 'gmres', 'lgmres', 'bicgstab', 'gcrotmk'] = Field(
default='gmres',
description="Solver to use for the linear problem. Options are 'direct', 'gmres', 'lgmres', 'bicgstab' or 'gcrotmk' when use_petsc = False."
)
use_petsc: bool = Field(
default=True,
description="Whether to use solvers from petsc4py or scipy. Recommended for large problems."
)
n_jobs: Optional[int] = Field(
default=None,
description="Number of parallel jobs to use when using an iterative solver."
)
backend: Literal['loky', 'multiprocessing', 'threading'] = Field(
default='loky',
description="Which backend to use for multiprocessing."
)
tol: float = Field(
default=1e-06,
description="Convergence tolerance for the iterative solver."
)
preconditioner: Optional[str] = Field(
default=None,
description="Preconditioner to use, only available when use_petsc = True. We recommend the 'ilu' preconditioner for badly conditioned problems."
)
