opencode-knowledge-work 0.0.0-stage → 0.1.1

package/content/bio-research/skills/bio-research-scvi-tools/scripts/model_utils.py

@@ -0,0 +1,634 @@
+#!/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")