Scvi Tools

scripts/model_utils.py

#!/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")