#!/usr/bin/env python3 """ Cluster and embed data using scvi-tools latent representation. Computes neighbors, UMAP, and Leiden clustering on the latent space. Input should have latent representation from train_model.py. Usage: python cluster_embed.py input.h5ad output_dir/ python cluster_embed.py input.h5ad output_dir/ --resolution 0.5 --use-rep X_scVI """ import argparse import os import sys def cluster_and_embed( adata, use_rep=None, n_neighbors=15, resolution=1.0, min_dist=0.3 ): """ Cluster and compute UMAP embedding. Parameters ---------- adata : AnnData Data with latent representation in obsm use_rep : str, optional Key in obsm to use (auto-detects if None) n_neighbors : int Number of neighbors for graph resolution : float Leiden clustering resolution min_dist : float UMAP min_dist parameter Returns ------- AnnData with neighbors, UMAP, and leiden clustering """ import scanpy as sc # Auto-detect representation if use_rep is None: candidates = ["X_scANVI", "X_scVI", "X_totalVI", "X_PeakVI", "X_MultiVI"] for key in candidates: if key in adata.obsm: use_rep = key break if use_rep is None: # Fall back to PCA if "X_pca" not in adata.obsm: print("No scvi-tools embedding found, computing PCA...") sc.pp.pca(adata) use_rep = "X_pca" print(f"Using representation: {use_rep}") print(f"Embedding shape: {adata.obsm[use_rep].shape}") # Compute neighbors print(f"Computing neighbors (n={n_neighbors})...") sc.pp.neighbors(adata, use_rep=use_rep, n_neighbors=n_neighbors) # UMAP print(f"Computing UMAP (min_dist={min_dist})...") sc.tl.umap(adata, min_dist=min_dist) # Leiden clustering print(f"Computing Leiden clustering (resolution={resolution})...") sc.tl.leiden(adata, resolution=resolution) n_clusters = adata.obs['leiden'].nunique() print(f"Found {n_clusters} clusters") return adata def plot_results(adata, output_dir, batch_key=None, labels_key=None): """Generate and save visualization plots.""" import scanpy as sc import matplotlib.pyplot as plt plots = [] # Always plot clusters plots.append(("leiden", "Clusters")) # Plot batch if available if batch_key is not None and batch_key in adata.obs.columns: plots.append((batch_key, f"Batch ({batch_key})")) # Plot labels if available if labels_key is not None and labels_key in adata.obs.columns: plots.append((labels_key, f"Labels ({labels_key})")) # Check for common columns for col in adata.obs.columns: if col not in [p[0] for p in plots]: if 'cell' in col.lower() and 'type' in col.lower(): plots.append((col, col)) elif 'predict' in col.lower(): plots.append((col, col)) # Limit to 6 plots plots = plots[:6] # Create figure n_plots = len(plots) n_cols = min(3, n_plots) n_rows = (n_plots + n_cols - 1) // n_cols fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows)) if n_plots == 1: axes = [axes] else: axes = axes.flatten() for i, (color, title) in enumerate(plots): try: sc.pl.umap(adata, color=color, ax=axes[i], show=False, title=title) except Exception as e: axes[i].set_title(f"Could not plot {color}: {e}") # Hide unused axes for i in range(len(plots), len(axes)): axes[i].set_visible(False) plt.tight_layout() plot_path = os.path.join(output_dir, "umap_clusters.png") plt.savefig(plot_path, dpi=150, bbox_inches="tight") plt.close() print(f"UMAP plot saved to {plot_path}") # Save cluster counts cluster_counts = adata.obs['leiden'].value_counts().sort_index() counts_path = os.path.join(output_dir, "cluster_counts.csv") cluster_counts.to_csv(counts_path) print(f"Cluster counts saved to {counts_path}") def main(): parser = argparse.ArgumentParser( description="Cluster and embed using scvi-tools latent space", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Basic clustering python cluster_embed.py adata_trained.h5ad results/ # Custom resolution python cluster_embed.py adata_trained.h5ad results/ --resolution 0.5 # Specify representation python cluster_embed.py adata_trained.h5ad results/ --use-rep X_scANVI # Include batch and label columns in plots python cluster_embed.py adata_trained.h5ad results/ --batch-key batch --labels-key cell_type """ ) parser.add_argument("input", help="Input h5ad file with latent representation") parser.add_argument("output_dir", help="Output directory") parser.add_argument("--use-rep", help="Representation key in obsm (auto-detects)") parser.add_argument("--n-neighbors", type=int, default=15, help="Neighbors for graph (default: 15)") parser.add_argument("--resolution", type=float, default=1.0, help="Leiden resolution (default: 1.0)") parser.add_argument("--min-dist", type=float, default=0.3, help="UMAP min_dist (default: 0.3)") parser.add_argument("--batch-key", help="Batch column for plotting") parser.add_argument("--labels-key", help="Labels column for plotting") args = parser.parse_args() try: import scanpy as sc except ImportError: print("Error: scanpy required. Install with: pip install scanpy") sys.exit(1) # Create output directory os.makedirs(args.output_dir, exist_ok=True) # Load data print(f"Loading {args.input}...") adata = sc.read_h5ad(args.input) print(f"Data: {adata.shape}") # Cluster and embed adata = cluster_and_embed( adata, use_rep=args.use_rep, n_neighbors=args.n_neighbors, resolution=args.resolution, min_dist=args.min_dist ) # Save results adata_path = os.path.join(args.output_dir, "adata_clustered.h5ad") adata.write_h5ad(adata_path) print(f"AnnData saved to {adata_path}") # Plot plot_results(adata, args.output_dir, args.batch_key, args.labels_key) print("\nDone!") if __name__ == "__main__": main()