Skills › Knowledge Work Plugins › scvi-tools › cluster_embed.py
cluster_embed.py
212 lines · bio-research/skills/scvi-tools/scripts/cluster_embed.py at 2574e29
1#!/usr/bin/env python3
2"""
3Cluster and embed data using scvi-tools latent representation.
5Computes neighbors, UMAP, and Leiden clustering on the latent space.
6Input should have latent representation from train_model.py.
8Usage:
9 python cluster_embed.py input.h5ad output_dir/
10 python cluster_embed.py input.h5ad output_dir/ --resolution 0.5 --use-rep X_scVI
11"""
13import argparse
14import os
15import sys
18def cluster_and_embed(
19 adata,
20 use_rep=None,
21 n_neighbors=15,
22 resolution=1.0,
23 min_dist=0.3
24):
25 """
26 Cluster and compute UMAP embedding.
28 Parameters
29 ----------
30 adata : AnnData
31 Data with latent representation in obsm
32 use_rep : str, optional
33 Key in obsm to use (auto-detects if None)
34 n_neighbors : int
35 Number of neighbors for graph
36 resolution : float
37 Leiden clustering resolution
38 min_dist : float
39 UMAP min_dist parameter
41 Returns
42 -------
43 AnnData with neighbors, UMAP, and leiden clustering
44 """
45 import scanpy as sc
47 # Auto-detect representation
48 if use_rep is None:
49 candidates = ["X_scANVI", "X_scVI", "X_totalVI", "X_PeakVI", "X_MultiVI"]
50 for key in candidates:
51 if key in adata.obsm:
52 use_rep = key
53 break
55 if use_rep is None:
56 # Fall back to PCA
57 if "X_pca" not in adata.obsm:
58 print("No scvi-tools embedding found, computing PCA...")
59 sc.pp.pca(adata)
60 use_rep = "X_pca"
62 print(f"Using representation: {use_rep}")
63 print(f"Embedding shape: {adata.obsm[use_rep].shape}")
65 # Compute neighbors
66 print(f"Computing neighbors (n={n_neighbors})...")
67 sc.pp.neighbors(adata, use_rep=use_rep, n_neighbors=n_neighbors)
69 # UMAP
70 print(f"Computing UMAP (min_dist={min_dist})...")
71 sc.tl.umap(adata, min_dist=min_dist)
73 # Leiden clustering
74 print(f"Computing Leiden clustering (resolution={resolution})...")
75 sc.tl.leiden(adata, resolution=resolution)
77 n_clusters = adata.obs['leiden'].nunique()
78 print(f"Found {n_clusters} clusters")
80 return adata
83def plot_results(adata, output_dir, batch_key=None, labels_key=None):
84 """Generate and save visualization plots."""
85 import scanpy as sc
86 import matplotlib.pyplot as plt
88 plots = []
90 # Always plot clusters
91 plots.append(("leiden", "Clusters"))
93 # Plot batch if available
94 if batch_key is not None and batch_key in adata.obs.columns:
95 plots.append((batch_key, f"Batch ({batch_key})"))
97 # Plot labels if available
98 if labels_key is not None and labels_key in adata.obs.columns:
99 plots.append((labels_key, f"Labels ({labels_key})"))
101 # Check for common columns
102 for col in adata.obs.columns:
103 if col not in [p[0] for p in plots]:
104 if 'cell' in col.lower() and 'type' in col.lower():
105 plots.append((col, col))
106 elif 'predict' in col.lower():
107 plots.append((col, col))
109 # Limit to 6 plots
110 plots = plots[:6]
112 # Create figure
113 n_plots = len(plots)
114 n_cols = min(3, n_plots)
115 n_rows = (n_plots + n_cols - 1) // n_cols
117 fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows))
118 if n_plots == 1:
119 axes = [axes]
120 else:
121 axes = axes.flatten()
123 for i, (color, title) in enumerate(plots):
124 try:
125 sc.pl.umap(adata, color=color, ax=axes[i], show=False, title=title)
126 except Exception as e:
127 axes[i].set_title(f"Could not plot {color}: {e}")
129 # Hide unused axes
130 for i in range(len(plots), len(axes)):
131 axes[i].set_visible(False)
133 plt.tight_layout()
135 plot_path = os.path.join(output_dir, "umap_clusters.png")
136 plt.savefig(plot_path, dpi=150, bbox_inches="tight")
137 plt.close()
138 print(f"UMAP plot saved to {plot_path}")
140 # Save cluster counts
141 cluster_counts = adata.obs['leiden'].value_counts().sort_index()
142 counts_path = os.path.join(output_dir, "cluster_counts.csv")
143 cluster_counts.to_csv(counts_path)
144 print(f"Cluster counts saved to {counts_path}")
147def main():
148 parser = argparse.ArgumentParser(
149 description="Cluster and embed using scvi-tools latent space",
150 formatter_class=argparse.RawDescriptionHelpFormatter,
151 epilog="""
152Examples:
153 # Basic clustering
154 python cluster_embed.py adata_trained.h5ad results/
156 # Custom resolution
157 python cluster_embed.py adata_trained.h5ad results/ --resolution 0.5
159 # Specify representation
160 python cluster_embed.py adata_trained.h5ad results/ --use-rep X_scANVI
162 # Include batch and label columns in plots
163 python cluster_embed.py adata_trained.h5ad results/ --batch-key batch --labels-key cell_type
164 """
165 )
166 parser.add_argument("input", help="Input h5ad file with latent representation")
167 parser.add_argument("output_dir", help="Output directory")
168 parser.add_argument("--use-rep", help="Representation key in obsm (auto-detects)")
169 parser.add_argument("--n-neighbors", type=int, default=15, help="Neighbors for graph (default: 15)")
170 parser.add_argument("--resolution", type=float, default=1.0, help="Leiden resolution (default: 1.0)")
171 parser.add_argument("--min-dist", type=float, default=0.3, help="UMAP min_dist (default: 0.3)")
172 parser.add_argument("--batch-key", help="Batch column for plotting")
173 parser.add_argument("--labels-key", help="Labels column for plotting")
175 args = parser.parse_args()
177 try:
178 import scanpy as sc
179 except ImportError:
180 print("Error: scanpy required. Install with: pip install scanpy")
181 sys.exit(1)
183 # Create output directory
184 os.makedirs(args.output_dir, exist_ok=True)
186 # Load data
187 print(f"Loading {args.input}...")
188 adata = sc.read_h5ad(args.input)
189 print(f"Data: {adata.shape}")
191 # Cluster and embed
192 adata = cluster_and_embed(
193 adata,
194 use_rep=args.use_rep,
195 n_neighbors=args.n_neighbors,
196 resolution=args.resolution,
197 min_dist=args.min_dist
198 )
200 # Save results
201 adata_path = os.path.join(args.output_dir, "adata_clustered.h5ad")
202 adata.write_h5ad(adata_path)
203 print(f"AnnData saved to {adata_path}")
205 # Plot
206 plot_results(adata, args.output_dir, args.batch_key, args.labels_key)
208 print("\nDone!")
211if __name__ == "__main__":
212 main()