scvi-tools MCP Server

from pydantic import Field from typing import Optional, List, Union, Dict, Any from .base import JSONParsingModel # ==================== Setup ==================== class SetupAnndataModel(JSONParsingModel): """Setup AnnData for scvi-tools models.""" layer: Optional[str] = Field( default=None, description="Key in adata.layers for count data. If None, uses adata.X" ) batch_key: Optional[str] = Field( default=None, description="Key in adata.obs for batch information" ) labels_key: Optional[str] = Field( default=None, description="Key in adata.obs for cell type labels (for SCANVI)" ) categorical_covariate_keys: Optional[List[str]] = Field( default=None, description="Keys in adata.obs for categorical covariates" ) continuous_covariate_keys: Optional[List[str]] = Field( default=None, description="Keys in adata.obs for continuous covariates" ) # ==================== SCVI ==================== class SCVISetupModel(SetupAnndataModel): """Setup AnnData for SCVI model.""" pass class SCVICreateModel(JSONParsingModel): """Create SCVI model.""" n_hidden: int = Field( default=128, description="Number of nodes per hidden layer" ) n_latent: int = Field( default=10, description="Dimensionality of the latent space" ) n_layers: int = Field( default=1, description="Number of hidden layers" ) dropout_rate: float = Field( default=0.1, description="Dropout rate for neural networks" ) gene_likelihood: str = Field( default="zinb", description="Gene likelihood: 'zinb', 'nb', 'poisson'" ) class SCVITrainModel(JSONParsingModel): """Train SCVI model.""" max_epochs: int = Field( default=400, description="Maximum number of training epochs" ) batch_size: int = Field( default=128, description="Minibatch size for training" ) early_stopping: bool = Field( default=True, description="Enable early stopping" ) early_stopping_patience: int = Field( default=45, description="Number of epochs with no improvement after which training stops" ) class SCVIGetLatentModel(JSONParsingModel): """Get latent representation from SCVI.""" give_mean: bool = Field( default=True, description="Return mean of latent distribution" ) save_key: str = Field( default="X_scVI", description="Key in adata.obsm to save latent representation" ) class SCVIGetNormalizedModel(JSONParsingModel): """Get normalized expression from SCVI.""" library_size: float = Field( default=1e4, description="Library size to use for normalization" ) save_layer: str = Field( default="scvi_normalized", description="Layer in adata.layers to save normalized expression" ) class SCVIDifferentialExpressionModel(JSONParsingModel): """Perform differential expression analysis with SCVI.""" groupby: str = Field( description="Key in adata.obs for grouping" ) group1: Optional[List[str]] = Field( default=None, description="First group of cells. If None, all groups are compared" ) group2: Optional[str] = Field( default=None, description="Second group to compare. If None, compare to rest" ) mode: str = Field( default="change", description="Mode: 'vanilla' or 'change'" ) delta: float = Field( default=0.25, description="Specific case of region inducing differential expression" ) batch_correction: bool = Field( default=False, description="Correct for batch effects" ) class SCVISaveModel(JSONParsingModel): """Save SCVI model.""" dir_path: str = Field( description="Path to directory where to save the model" ) overwrite: bool = Field( default=False, description="Overwrite existing directory" ) class SCVILoadModel(JSONParsingModel): """Load SCVI model.""" dir_path: str = Field( description="Path to saved model directory" ) # ==================== SCANVI ==================== class SCANVISetupModel(SetupAnndataModel): """Setup AnnData for SCANVI model.""" unlabeled_category: str = Field( description="Value used for unlabeled cells in labels_key" ) class SCANVICreateModel(JSONParsingModel): """Create SCANVI model.""" unlabeled_category: str = Field( description="Value used for unlabeled cells" ) n_hidden: int = Field( default=128, description="Number of nodes per hidden layer" ) n_latent: int = Field( default=10, description="Dimensionality of the latent space" ) n_layers: int = Field( default=1, description="Number of hidden layers" ) dropout_rate: float = Field( default=0.1, description="Dropout rate" ) gene_likelihood: str = Field( default="zinb", description="Gene likelihood: 'zinb', 'nb', 'poisson'" ) class SCANVIFromSCVIModel(JSONParsingModel): """Create SCANVI from pretrained SCVI.""" unlabeled_category: str = Field( description="Value used for unlabeled cells" ) class SCANVIPredictModel(JSONParsingModel): """Predict cell types with SCANVI.""" save_key: str = Field( default="scanvi_predictions", description="Key in adata.obs to save predictions" ) soft: bool = Field( default=False, description="Return probabilities instead of predicted labels" ) # ==================== TOTALVI ==================== class TOTALVISetupModel(JSONParsingModel): """Setup AnnData for TOTALVI model.""" protein_expression_obsm_key: str = Field( description="Key in adata.obsm for protein expression" ) layer: Optional[str] = Field( default=None, description="Layer for RNA counts" ) batch_key: Optional[str] = Field( default=None, description="Batch key in adata.obs" ) categorical_covariate_keys: Optional[List[str]] = Field( default=None, description="Categorical covariate keys" ) continuous_covariate_keys: Optional[List[str]] = Field( default=None, description="Continuous covariate keys" ) class TOTALVICreateModel(JSONParsingModel): """Create TOTALVI model.""" n_latent: int = Field( default=20, description="Dimensionality of latent space" ) gene_likelihood: str = Field( default="nb", description="Gene likelihood: 'nb', 'zinb'" ) class TOTALVIGetProteinForegroundProbModel(JSONParsingModel): """Get protein foreground probability.""" save_key: str = Field( default="totalvi_protein_fg_prob", description="Key to save in adata.obsm" ) # ==================== PEAKVI ==================== class PEAKVISetupModel(JSONParsingModel): """Setup AnnData for PEAKVI model.""" batch_key: Optional[str] = Field( default=None, description="Batch key in adata.obs" ) layer: Optional[str] = Field( default=None, description="Layer for counts. If None, use adata.X" ) class PEAKVICreateModel(JSONParsingModel): """Create PEAKVI model.""" n_hidden: int = Field( default=128, description="Number of nodes per hidden layer" ) n_latent: int = Field( default=10, description="Dimensionality of latent space" ) n_layers_encoder: int = Field( default=2, description="Number of layers for encoder" ) n_layers_decoder: int = Field( default=2, description="Number of layers for decoder" ) class PEAKVIDifferentialAccessibilityModel(JSONParsingModel): """Differential accessibility analysis with PEAKVI.""" groupby: str = Field( description="Key in adata.obs for grouping" ) group1: Optional[List[str]] = Field( default=None, description="First group" ) group2: Optional[str] = Field( default=None, description="Second group" ) # ==================== Common Operations ==================== class GetELBOModel(JSONParsingModel): """Get ELBO (Evidence Lower Bound).""" pass class GetReconstructionErrorModel(JSONParsingModel): """Get reconstruction error.""" pass