opencode-knowledge-work 0.0.0-stage → 0.1.0

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

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""
+Differential expression analysis using scvi-tools models.
+
+Uses the trained model's differential_expression method which accounts
+for batch effects and uses the generative model for inference.
+
+Usage:
+    python differential_expression.py model_dir/ adata.h5ad output.csv --groupby leiden
+    python differential_expression.py model_dir/ adata.h5ad output.csv --groupby cell_type --group1 "T cells" --group2 "B cells"
+"""
+
+import argparse
+import os
+import sys
+
+
+def run_de_analysis(
+    model,
+    adata,
+    groupby,
+    group1=None,
+    group2=None,
+    n_genes=None
+):
+    """
+    Run differential expression analysis.
+
+    Parameters
+    ----------
+    model : scvi model
+        Trained model with differential_expression method
+    adata : AnnData
+        Data used for training
+    groupby : str
+        Column in obs to group by
+    group1 : str, optional
+        First group (if None, computes for all groups)
+    group2 : str, optional
+        Second group (rest if None)
+    n_genes : int, optional
+        Limit to top N genes per group
+
+    Returns
+    -------
+    DataFrame with DE results
+    """
+    import pandas as pd
+
+    if group1 is not None:
+        # Specific comparison
+        print(f"Comparing {group1} vs {group2 or 'rest'}...")
+        de_results = model.differential_expression(
+            groupby=groupby,
+            group1=group1,
+            group2=group2
+        )
+
+        # Add comparison info
+        de_results["comparison"] = f"{group1}_vs_{group2 or 'rest'}"
+
+    else:
+        # All pairwise or one-vs-rest
+        groups = adata.obs[groupby].unique()
+        print(f"Computing DE for {len(groups)} groups...")
+
+        all_results = []
+        for group in groups:
+            print(f"  Processing {group}...")
+            try:
+                de = model.differential_expression(
+                    groupby=groupby,
+                    group1=group
+                )
+                de["group"] = group
+                all_results.append(de)
+            except Exception as e:
+                print(f"  Warning: Failed for {group}: {e}")
+
+        de_results = pd.concat(all_results, ignore_index=False)
+
+    # Filter to significant
+    if "is_de_fdr_0.05" in de_results.columns:
+        n_sig = de_results["is_de_fdr_0.05"].sum()
+        print(f"Found {n_sig} significant DE genes (FDR < 0.05)")
+
+    # Limit to top genes if requested
+    if n_genes is not None and "lfc_mean" in de_results.columns:
+        if "group" in de_results.columns:
+            # Top N per group
+            de_results = de_results.groupby("group").apply(
+                lambda x: x.nlargest(n_genes, "lfc_mean")
+            ).reset_index(drop=True)
+        else:
+            de_results = de_results.nlargest(n_genes, "lfc_mean")
+
+    return de_results
+
+
+def plot_volcano(de_results, output_path, group_name=None):
+    """Create volcano plot of DE results."""
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    if "lfc_mean" not in de_results.columns:
+        print("Cannot create volcano plot: missing lfc_mean column")
+        return
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    # Get values
+    lfc = de_results["lfc_mean"].values
+    if "bayes_factor" in de_results.columns:
+        y_val = de_results["bayes_factor"].values
+        y_label = "Bayes Factor"
+    elif "proba_de" in de_results.columns:
+        y_val = -np.log10(1 - de_results["proba_de"].values + 1e-10)
+        y_label = "-log10(1 - P(DE))"
+    else:
+        y_val = np.ones(len(lfc))
+        y_label = ""
+
+    # Color by significance
+    if "is_de_fdr_0.05" in de_results.columns:
+        sig = de_results["is_de_fdr_0.05"].values
+        colors = ["red" if s else "gray" for s in sig]
+    else:
+        colors = "gray"
+
+    ax.scatter(lfc, y_val, c=colors, alpha=0.5, s=10)
+    ax.axvline(0, color="black", linestyle="--", alpha=0.5)
+    ax.set_xlabel("Log Fold Change")
+    ax.set_ylabel(y_label)
+
+    title = "Differential Expression"
+    if group_name:
+        title += f": {group_name}"
+    ax.set_title(title)
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"Volcano plot saved to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Differential expression with scvi-tools",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # DE for all clusters (one-vs-rest)
+    python differential_expression.py model/ adata.h5ad de_results.csv --groupby leiden
+
+    # Specific comparison
+    python differential_expression.py model/ adata.h5ad de_results.csv \\
+        --groupby cell_type --group1 "T cells" --group2 "B cells"
+
+    # Top 50 genes per cluster
+    python differential_expression.py model/ adata.h5ad de_results.csv \\
+        --groupby leiden --n-genes 50
+        """
+    )
+    parser.add_argument("model_dir", help="Directory containing saved model")
+    parser.add_argument("input", help="Input h5ad file (same as training)")
+    parser.add_argument("output", help="Output CSV file for DE results")
+    parser.add_argument("--groupby", required=True, help="Column to group by")
+    parser.add_argument("--group1", help="First group for comparison")
+    parser.add_argument("--group2", help="Second group (default: rest)")
+    parser.add_argument("--n-genes", type=int, help="Limit to top N genes per group")
+    parser.add_argument("--model-type", choices=["scvi", "scanvi", "totalvi"],
+                        default="scvi", help="Model type (default: scvi)")
+    parser.add_argument("--plot", action="store_true", help="Generate volcano plot")
+
+    args = parser.parse_args()
+
+    try:
+        import scvi
+        import scanpy as sc
+    except ImportError:
+        print("Error: scvi-tools and scanpy required")
+        sys.exit(1)
+
+    # Load data
+    print(f"Loading {args.input}...")
+    adata = sc.read_h5ad(args.input)
+
+    # Load model
+    print(f"Loading model from {args.model_dir}...")
+    if args.model_type == "scvi":
+        model = scvi.model.SCVI.load(args.model_dir, adata=adata)
+    elif args.model_type == "scanvi":
+        model = scvi.model.SCANVI.load(args.model_dir, adata=adata)
+    elif args.model_type == "totalvi":
+        model = scvi.model.TOTALVI.load(args.model_dir, adata=adata)
+
+    # Run DE
+    de_results = run_de_analysis(
+        model,
+        adata,
+        groupby=args.groupby,
+        group1=args.group1,
+        group2=args.group2,
+        n_genes=args.n_genes
+    )
+
+    # Save results
+    de_results.to_csv(args.output)
+    print(f"DE results saved to {args.output}")
+
+    # Plot
+    if args.plot:
+        plot_path = args.output.replace(".csv", "_volcano.png")
+        plot_volcano(de_results, plot_path, args.group1)
+
+    print("\nDone!")
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+"""
+Integrate multiple datasets using scvi-tools.
+
+Concatenates multiple h5ad files and runs batch correction with scVI or scANVI.
+
+Usage:
+    python integrate_datasets.py output_dir/ dataset1.h5ad dataset2.h5ad dataset3.h5ad
+    python integrate_datasets.py output_dir/ *.h5ad --batch-names study1,study2,study3
+"""
+
+import argparse
+import os
+import sys
+
+
+def integrate_datasets(
+    adatas,
+    batch_names=None,
+    labels_key=None,
+    n_top_genes=2000,
+    n_latent=30,
+    max_epochs=200
+):
+    """
+    Integrate multiple datasets.
+
+    Parameters
+    ----------
+    adatas : list of AnnData
+        Datasets to integrate
+    batch_names : list of str, optional
+        Names for each dataset (default: dataset_0, dataset_1, ...)
+    labels_key : str, optional
+        Cell type column (uses scANVI if provided)
+    n_top_genes : int
+        Number of HVGs
+    n_latent : int
+        Latent dimensions
+    max_epochs : int
+        Training epochs
+
+    Returns
+    -------
+    Integrated AnnData and trained model
+    """
+    import scvi
+    import scanpy as sc
+    import numpy as np
+
+    # Assign batch names
+    if batch_names is None:
+        batch_names = [f"dataset_{i}" for i in range(len(adatas))]
+
+    if len(batch_names) != len(adatas):
+        raise ValueError(f"Number of batch names ({len(batch_names)}) must match datasets ({len(adatas)})")
+
+    # Add batch labels
+    for adata, name in zip(adatas, batch_names):
+        adata.obs["batch"] = name
+        print(f"{name}: {adata.shape}")
+
+    # Find common genes
+    common_genes = set(adatas[0].var_names)
+    for adata in adatas[1:]:
+        common_genes = common_genes.intersection(adata.var_names)
+    common_genes = list(common_genes)
+    print(f"\nCommon genes: {len(common_genes)}")
+
+    # Subset to common genes
+    adatas = [adata[:, common_genes].copy() for adata in adatas]
+
+    # Concatenate
+    print("Concatenating datasets...")
+    adata = sc.concat(adatas, label="batch", keys=batch_names)
+    print(f"Combined: {adata.shape}")
+
+    # Store counts
+    adata.layers["counts"] = adata.X.copy()
+
+    # HVG selection
+    print(f"Selecting {n_top_genes} HVGs...")
+    sc.pp.highly_variable_genes(
+        adata,
+        n_top_genes=n_top_genes,
+        flavor="seurat_v3",
+        batch_key="batch",
+        layer="counts"
+    )
+    adata = adata[:, adata.var["highly_variable"]].copy()
+
+    # Train model
+    if labels_key is not None and labels_key in adata.obs.columns:
+        print(f"\nTraining scANVI with labels ({labels_key})...")
+
+        # First train scVI
+        scvi.model.SCVI.setup_anndata(adata, layer="counts", batch_key="batch")
+        scvi_model = scvi.model.SCVI(adata, n_latent=n_latent)
+        scvi_model.train(max_epochs=max_epochs, early_stopping=True)
+
+        # Then scANVI
+        model = scvi.model.SCANVI.from_scvi_model(
+            scvi_model,
+            labels_key=labels_key,
+            unlabeled_category="Unknown"
+        )
+        model.train(max_epochs=max_epochs // 4)
+
+        adata.obsm["X_scANVI"] = model.get_latent_representation()
+        rep_key = "X_scANVI"
+
+    else:
+        print("\nTraining scVI...")
+        scvi.model.SCVI.setup_anndata(adata, layer="counts", batch_key="batch")
+        model = scvi.model.SCVI(adata, n_latent=n_latent)
+        model.train(max_epochs=max_epochs, early_stopping=True)
+
+        adata.obsm["X_scVI"] = model.get_latent_representation()
+        rep_key = "X_scVI"
+
+    # Cluster
+    print("\nClustering...")
+    sc.pp.neighbors(adata, use_rep=rep_key)
+    sc.tl.umap(adata)
+    sc.tl.leiden(adata)
+
+    print(f"Found {adata.obs['leiden'].nunique()} clusters")
+
+    return adata, model
+
+
+def plot_integration(adata, output_dir, labels_key=None):
+    """Plot integration results."""
+    import scanpy as sc
+    import matplotlib.pyplot as plt
+
+    plots = [
+        ("batch", "By Batch"),
+        ("leiden", "Clusters")
+    ]
+
+    if labels_key is not None and labels_key in adata.obs.columns:
+        plots.append((labels_key, f"Cell Types ({labels_key})"))
+
+    if "predicted_cell_type" in adata.obs.columns:
+        plots.append(("predicted_cell_type", "Predicted Types"))
+
+    n_plots = len(plots)
+    fig, axes = plt.subplots(1, n_plots, figsize=(5 * n_plots, 4))
+    if n_plots == 1:
+        axes = [axes]
+
+    for ax, (color, title) in zip(axes, plots):
+        sc.pl.umap(adata, color=color, ax=ax, show=False, title=title)
+
+    plt.tight_layout()
+    plot_path = os.path.join(output_dir, "integration.png")
+    plt.savefig(plot_path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"Integration plot saved to {plot_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Integrate multiple datasets with scvi-tools",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Integrate multiple files
+    python integrate_datasets.py results/ data1.h5ad data2.h5ad data3.h5ad
+
+    # With custom batch names
+    python integrate_datasets.py results/ *.h5ad --batch-names ctrl,treat1,treat2
+
+    # With cell type labels (uses scANVI)
+    python integrate_datasets.py results/ *.h5ad --labels-key cell_type
+        """
+    )
+    parser.add_argument("output_dir", help="Output directory")
+    parser.add_argument("inputs", nargs="+", help="Input h5ad files")
+    parser.add_argument("--batch-names", help="Comma-separated batch names")
+    parser.add_argument("--labels-key", help="Cell type column (uses scANVI)")
+    parser.add_argument("--n-hvgs", type=int, default=2000, help="Number of HVGs (default: 2000)")
+    parser.add_argument("--n-latent", type=int, default=30, help="Latent dimensions (default: 30)")
+    parser.add_argument("--max-epochs", type=int, default=200, help="Max epochs (default: 200)")
+
+    args = parser.parse_args()
+
+    try:
+        import scvi
+        import scanpy as sc
+    except ImportError:
+        print("Error: scvi-tools and scanpy required")
+        sys.exit(1)
+
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Parse batch names
+    batch_names = None
+    if args.batch_names:
+        batch_names = args.batch_names.split(",")
+
+    # Load datasets
+    print("Loading datasets...")
+    adatas = []
+    for path in args.inputs:
+        print(f"  Loading {path}...")
+        adatas.append(sc.read_h5ad(path))
+
+    # Integrate
+    adata, model = integrate_datasets(
+        adatas,
+        batch_names=batch_names,
+        labels_key=args.labels_key,
+        n_top_genes=args.n_hvgs,
+        n_latent=args.n_latent,
+        max_epochs=args.max_epochs
+    )
+
+    # Save results
+    adata_path = os.path.join(args.output_dir, "integrated.h5ad")
+    adata.write_h5ad(adata_path)
+    print(f"\nIntegrated data saved to {adata_path}")
+
+    model_path = os.path.join(args.output_dir, "model")
+    model.save(model_path)
+    print(f"Model saved to {model_path}")
+
+    # Plot
+    plot_integration(adata, args.output_dir, args.labels_key)
+
+    print("\nDone!")
+
+
+if __name__ == "__main__":
+    main()