#!/usr/bin/env python3 """ Utility functions for scvi-tools model training and evaluation. Usage: from model_utils import prepare_adata, train_scvi, evaluate_integration """ import numpy as np import scanpy as sc from typing import Optional, List, Dict, Tuple import warnings def get_mito_genes(adata) -> np.ndarray: """ Identify mitochondrial genes for both human and mouse data. Handles common prefixes: - Human: MT- (e.g., MT-CO1, MT-ND1) - Mouse: mt- or Mt- (e.g., mt-Co1, Mt-Nd1) Returns ------- Boolean array indicating mitochondrial genes """ return ( adata.var_names.str.startswith('MT-') | adata.var_names.str.startswith('mt-') | adata.var_names.str.startswith('Mt-') ) def prepare_adata( adata, batch_key: Optional[str] = None, n_top_genes: int = 2000, min_genes: int = 200, max_genes: int = 5000, max_mito_pct: float = 20.0, min_cells: int = 3, copy: bool = True ): """ Prepare AnnData for scvi-tools models. Parameters ---------- adata : AnnData Raw count data batch_key : str, optional Column for batch information n_top_genes : int Number of highly variable genes min_genes : int Minimum genes per cell max_genes : int Maximum genes per cell max_mito_pct : float Maximum mitochondrial percentage min_cells : int Minimum cells per gene copy : bool Return copy of data Returns ------- AnnData prepared for scvi-tools """ if copy: adata = adata.copy() # Calculate QC metrics adata.var['mt'] = get_mito_genes(adata) sc.pp.calculate_qc_metrics(adata, qc_vars=['mt'], inplace=True) # Filter cells adata = adata[adata.obs['n_genes_by_counts'] >= min_genes].copy() adata = adata[adata.obs['n_genes_by_counts'] <= max_genes].copy() adata = adata[adata.obs['pct_counts_mt'] < max_mito_pct].copy() # Filter genes sc.pp.filter_genes(adata, min_cells=min_cells) # Store raw counts adata.layers["counts"] = adata.X.copy() # HVG selection if batch_key and batch_key in adata.obs.columns: sc.pp.highly_variable_genes( adata, n_top_genes=n_top_genes, flavor="seurat_v3", batch_key=batch_key, layer="counts" ) else: # Need to normalize for non-seurat_v3 flavor sc.pp.normalize_total(adata, target_sum=1e4) sc.pp.log1p(adata) sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes) # Restore counts to X adata.X = adata.layers["counts"].copy() # Subset to HVGs adata = adata[:, adata.var['highly_variable']].copy() print(f"Prepared AnnData: {adata.shape}") if batch_key: print(f"Batches: {adata.obs[batch_key].nunique()}") return adata def train_scvi( adata, batch_key: Optional[str] = None, labels_key: Optional[str] = None, n_latent: int = 30, n_layers: int = 2, max_epochs: int = 200, early_stopping: bool = True, use_gpu: bool = True ): """ Train scVI or scANVI model. Parameters ---------- adata : AnnData Prepared data with counts layer batch_key : str, optional Batch column labels_key : str, optional Cell type labels (uses scANVI if provided) n_latent : int Latent dimensions n_layers : int Encoder/decoder layers max_epochs : int Maximum training epochs early_stopping : bool Use early stopping use_gpu : bool Use GPU if available Returns ------- Trained model """ import scvi # Setup AnnData scvi.model.SCVI.setup_anndata( adata, layer="counts", batch_key=batch_key ) if labels_key and labels_key in adata.obs.columns: # Train scVI first scvi_model = scvi.model.SCVI( adata, n_latent=n_latent, n_layers=n_layers ) scvi_model.train( max_epochs=max_epochs, early_stopping=early_stopping ) # Initialize scANVI model = scvi.model.SCANVI.from_scvi_model( scvi_model, labels_key=labels_key, unlabeled_category="Unknown" ) model.train(max_epochs=max_epochs // 4) # Store representation adata.obsm["X_scANVI"] = model.get_latent_representation() else: # Train scVI only model = scvi.model.SCVI( adata, n_latent=n_latent, n_layers=n_layers ) model.train( max_epochs=max_epochs, early_stopping=early_stopping ) # Store representation adata.obsm["X_scVI"] = model.get_latent_representation() return model def evaluate_integration( adata, batch_key: str, label_key: str, embedding_key: str = "X_scVI" ) -> Dict[str, float]: """ Evaluate integration quality using basic metrics. Parameters ---------- adata : AnnData Integrated data batch_key : str Batch column label_key : str Cell type column embedding_key : str Key in obsm for embedding Returns ------- Dictionary of metrics """ from sklearn.metrics import silhouette_score from sklearn.neighbors import NearestNeighbors X = adata.obsm[embedding_key] batch = adata.obs[batch_key].values labels = adata.obs[label_key].values metrics = {} # Silhouette scores try: # Cell type silhouette (higher = better separation) metrics["silhouette_label"] = silhouette_score(X, labels) # Batch silhouette (lower = better mixing) metrics["silhouette_batch"] = silhouette_score(X, batch) except Exception as e: warnings.warn(f"Silhouette calculation failed: {e}") # Batch mixing in neighbors try: nn = NearestNeighbors(n_neighbors=50) nn.fit(X) distances, indices = nn.kneighbors(X) batch_mixing = [] for i in range(len(X)): neighbor_batches = batch[indices[i]] unique_batches = len(np.unique(neighbor_batches)) batch_mixing.append(unique_batches / len(np.unique(batch))) metrics["batch_mixing"] = np.mean(batch_mixing) except Exception as e: warnings.warn(f"Batch mixing calculation failed: {e}") return metrics def get_marker_genes( model, adata, groupby: str, n_genes: int = 10 ) -> Dict[str, List[str]]: """ Get marker genes using scVI differential expression. Parameters ---------- model : scvi model Trained scVI/scANVI model adata : AnnData Data used for training groupby : str Column to group cells by n_genes : int Number of top markers per group Returns ------- Dictionary of {group: [marker_genes]} """ markers = {} groups = adata.obs[groupby].unique() for group in groups: # Get DE results for this group vs rest de_results = model.differential_expression( groupby=groupby, group1=group ) # Filter and sort de_sig = de_results[ (de_results["is_de_fdr_0.05"] == True) & (de_results["lfc_mean"] > 0.5) ].sort_values("lfc_mean", ascending=False) markers[group] = de_sig.index[:n_genes].tolist() return markers def plot_training_history(model, save_path: Optional[str] = None): """ Plot model training history. Parameters ---------- model : scvi model Trained model save_path : str, optional Path to save figure """ import matplotlib.pyplot as plt fig, axes = plt.subplots(1, 2, figsize=(12, 4)) # ELBO if "elbo_train" in model.history: axes[0].plot(model.history["elbo_train"], label="Train") if "elbo_validation" in model.history: axes[0].plot(model.history["elbo_validation"], label="Validation") axes[0].set_xlabel("Epoch") axes[0].set_ylabel("ELBO") axes[0].legend() axes[0].set_title("Training Loss") # Reconstruction if "reconstruction_loss_train" in model.history: axes[1].plot(model.history["reconstruction_loss_train"], label="Train") if "reconstruction_loss_validation" in model.history: axes[1].plot(model.history["reconstruction_loss_validation"], label="Validation") axes[1].set_xlabel("Epoch") axes[1].set_ylabel("Reconstruction Loss") axes[1].legend() axes[1].set_title("Reconstruction Loss") plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches="tight") return fig def save_results( model, adata, output_dir: str, save_model: bool = True, save_adata: bool = True, plot_umap: bool = True ): """ Save model, processed data, and visualization. Parameters ---------- model : scvi model Trained model adata : AnnData Processed data with latent representation output_dir : str Output directory path save_model : bool Save the trained model save_adata : bool Save the processed AnnData plot_umap : bool Generate and save UMAP plot """ import os import scanpy as sc import matplotlib.pyplot as plt os.makedirs(output_dir, exist_ok=True) # Save model if save_model: model_path = os.path.join(output_dir, "model") model.save(model_path) print(f"Model saved to {model_path}") # Save AnnData if save_adata: adata_path = os.path.join(output_dir, "adata_processed.h5ad") adata.write(adata_path) print(f"AnnData saved to {adata_path}") # Generate UMAP if needed if plot_umap: # Determine which embedding to use if "X_scANVI" in adata.obsm: rep_key = "X_scANVI" elif "X_scVI" in adata.obsm: rep_key = "X_scVI" else: rep_key = None if rep_key is not None: # Compute neighbors and UMAP if not present if "X_umap" not in adata.obsm: sc.pp.neighbors(adata, use_rep=rep_key) sc.tl.umap(adata) # Plot fig, axes = plt.subplots(1, 2, figsize=(12, 5)) # Plot by batch if available batch_cols = [c for c in adata.obs.columns if 'batch' in c.lower()] if batch_cols: sc.pl.umap(adata, color=batch_cols[0], ax=axes[0], show=False, title="By Batch") # Plot by cluster if "leiden" not in adata.obs: sc.tl.leiden(adata) sc.pl.umap(adata, color="leiden", ax=axes[1], show=False, title="Clusters") plt.tight_layout() plot_path = os.path.join(output_dir, "umap.png") plt.savefig(plot_path, dpi=150, bbox_inches="tight") plt.close() print(f"UMAP plot saved to {plot_path}") def auto_select_model(adata) -> str: """ Suggest the best scvi-tools model based on available data. Parameters ---------- adata : AnnData Data to analyze Returns ------- String with model recommendation and reasoning """ suggestions = [] # Check for multi-modal data if 'protein_expression' in adata.obsm: suggestions.append({ 'model': 'totalVI', 'reason': 'CITE-seq data detected (protein + RNA)', 'priority': 1 }) if 'spliced' in adata.layers and 'unspliced' in adata.layers: suggestions.append({ 'model': 'veloVI', 'reason': 'RNA velocity data detected (spliced + unspliced)', 'priority': 1 }) # Check for ATAC data indicators if adata.n_vars > 100000: # Many peaks suggest ATAC suggestions.append({ 'model': 'PeakVI', 'reason': f'Large number of features ({adata.n_vars}) suggests ATAC-seq peaks', 'priority': 2 }) # Check for labels label_cols = [c for c in adata.obs.columns if 'cell' in c.lower() or 'type' in c.lower() or 'label' in c.lower()] has_labels = len(label_cols) > 0 # Check for batch info batch_cols = [c for c in adata.obs.columns if 'batch' in c.lower() or 'sample' in c.lower()] has_batch = len(batch_cols) > 0 if has_batch: if has_labels: suggestions.append({ 'model': 'scANVI', 'reason': f'Batch info ({batch_cols[0]}) + labels ({label_cols[0]}) available', 'priority': 1 }) else: suggestions.append({ 'model': 'scVI', 'reason': f'Batch info ({batch_cols[0]}) available, no labels', 'priority': 1 }) else: suggestions.append({ 'model': 'scVI', 'reason': 'Standard scRNA-seq analysis', 'priority': 2 }) # Sort by priority suggestions.sort(key=lambda x: x['priority']) # Format output lines = ["Recommended models (in order of priority):"] for i, s in enumerate(suggestions, 1): lines.append(f" {i}. {s['model']}: {s['reason']}") return "\n".join(lines) def compare_integrations( adata, batch_key: str, label_key: str, embedding_keys: List[str] = None ) -> Dict[str, Dict[str, float]]: """ Compare multiple integration methods using standard metrics. Parameters ---------- adata : AnnData Data with integration embeddings in obsm batch_key : str Batch column in obs label_key : str Cell type column in obs embedding_keys : list, optional Keys in obsm to compare (default: auto-detect) Returns ------- Dictionary of {embedding: {metric: value}} """ from sklearn.metrics import silhouette_score # Auto-detect embeddings if embedding_keys is None: embedding_keys = [k for k in adata.obsm.keys() if k.startswith('X_') and 'umap' not in k.lower()] results = {} for key in embedding_keys: if key not in adata.obsm: continue X = adata.obsm[key] batch = adata.obs[batch_key].values labels = adata.obs[label_key].values metrics = {} try: # Silhouette scores metrics["silhouette_label"] = silhouette_score(X, labels) metrics["silhouette_batch"] = silhouette_score(X, batch) # Combined score (higher label preservation, lower batch separation = better) metrics["integration_score"] = metrics["silhouette_label"] - metrics["silhouette_batch"] except Exception as e: metrics["error"] = str(e) results[key] = metrics return results def quick_clustering( adata, use_rep: str = None, resolution: float = 1.0, n_neighbors: int = 15 ): """ Quick clustering pipeline on latent representation. Parameters ---------- adata : AnnData Data with latent representation use_rep : str, optional Key in obsm (auto-detects scVI/scANVI if not specified) resolution : float Leiden clustering resolution n_neighbors : int Number of neighbors for graph Returns ------- AnnData with neighbors, UMAP, and leiden clustering """ import scanpy as sc # Auto-detect representation if use_rep is None: if "X_scANVI" in adata.obsm: use_rep = "X_scANVI" elif "X_scVI" in adata.obsm: use_rep = "X_scVI" elif "X_totalVI" in adata.obsm: use_rep = "X_totalVI" elif "X_PeakVI" in adata.obsm: use_rep = "X_PeakVI" elif "X_MultiVI" in adata.obsm: use_rep = "X_MultiVI" else: raise ValueError("No scvi-tools embedding found in obsm") print(f"Using representation: {use_rep}") # Compute neighbors sc.pp.neighbors(adata, use_rep=use_rep, n_neighbors=n_neighbors) # UMAP sc.tl.umap(adata) # Leiden clustering sc.tl.leiden(adata, resolution=resolution) print(f"Found {adata.obs['leiden'].nunique()} clusters") return adata if __name__ == "__main__": print("scvi-tools model utilities") print("\nAvailable functions:") print(" - prepare_adata: Standard data preparation (QC, HVG, layer setup)") print(" - train_scvi: Train scVI or scANVI with sensible defaults") print(" - evaluate_integration: Compute batch mixing and silhouette metrics") print(" - get_marker_genes: Extract markers using scVI differential expression") print(" - plot_training_history: Visualize training convergence") print(" - save_results: Save model, data, and visualizations") print(" - auto_select_model: Suggest best model for your data") print(" - compare_integrations: Compare multiple integration embeddings") print(" - quick_clustering: Quick clustering on latent representation")