Skills › Knowledge Work Plugins › scvi-tools › model_utils.py
model_utils.py
634 lines · bio-research/skills/scvi-tools/scripts/model_utils.py at 2574e29
1#!/usr/bin/env python3
2"""
3Utility functions for scvi-tools model training and evaluation.
5Usage:
6 from model_utils import prepare_adata, train_scvi, evaluate_integration
7"""
9import numpy as np
10import scanpy as sc
11from typing import Optional, List, Dict, Tuple
12import warnings
15def get_mito_genes(adata) -> np.ndarray:
16 """
17 Identify mitochondrial genes for both human and mouse data.
19 Handles common prefixes:
20 - Human: MT- (e.g., MT-CO1, MT-ND1)
21 - Mouse: mt- or Mt- (e.g., mt-Co1, Mt-Nd1)
23 Returns
24 -------
25 Boolean array indicating mitochondrial genes
26 """
27 return (
28 adata.var_names.str.startswith('MT-') |
29 adata.var_names.str.startswith('mt-') |
30 adata.var_names.str.startswith('Mt-')
31 )
34def prepare_adata(
35 adata,
36 batch_key: Optional[str] = None,
37 n_top_genes: int = 2000,
38 min_genes: int = 200,
39 max_genes: int = 5000,
40 max_mito_pct: float = 20.0,
41 min_cells: int = 3,
42 copy: bool = True
43):
44 """
45 Prepare AnnData for scvi-tools models.
47 Parameters
48 ----------
49 adata : AnnData
50 Raw count data
51 batch_key : str, optional
52 Column for batch information
53 n_top_genes : int
54 Number of highly variable genes
55 min_genes : int
56 Minimum genes per cell
57 max_genes : int
58 Maximum genes per cell
59 max_mito_pct : float
60 Maximum mitochondrial percentage
61 min_cells : int
62 Minimum cells per gene
63 copy : bool
64 Return copy of data
66 Returns
67 -------
68 AnnData prepared for scvi-tools
69 """
70 if copy:
71 adata = adata.copy()
73 # Calculate QC metrics
74 adata.var['mt'] = get_mito_genes(adata)
75 sc.pp.calculate_qc_metrics(adata, qc_vars=['mt'], inplace=True)
77 # Filter cells
78 adata = adata[adata.obs['n_genes_by_counts'] >= min_genes].copy()
79 adata = adata[adata.obs['n_genes_by_counts'] <= max_genes].copy()
80 adata = adata[adata.obs['pct_counts_mt'] < max_mito_pct].copy()
82 # Filter genes
83 sc.pp.filter_genes(adata, min_cells=min_cells)
85 # Store raw counts
86 adata.layers["counts"] = adata.X.copy()
88 # HVG selection
89 if batch_key and batch_key in adata.obs.columns:
90 sc.pp.highly_variable_genes(
91 adata,
92 n_top_genes=n_top_genes,
93 flavor="seurat_v3",
94 batch_key=batch_key,
95 layer="counts"
96 )
97 else:
98 # Need to normalize for non-seurat_v3 flavor
99 sc.pp.normalize_total(adata, target_sum=1e4)
100 sc.pp.log1p(adata)
101 sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes)
102 # Restore counts to X
103 adata.X = adata.layers["counts"].copy()
105 # Subset to HVGs
106 adata = adata[:, adata.var['highly_variable']].copy()
108 print(f"Prepared AnnData: {adata.shape}")
109 if batch_key:
110 print(f"Batches: {adata.obs[batch_key].nunique()}")
112 return adata
115def train_scvi(
116 adata,
117 batch_key: Optional[str] = None,
118 labels_key: Optional[str] = None,
119 n_latent: int = 30,
120 n_layers: int = 2,
121 max_epochs: int = 200,
122 early_stopping: bool = True,
123 use_gpu: bool = True
124):
125 """
126 Train scVI or scANVI model.
128 Parameters
129 ----------
130 adata : AnnData
131 Prepared data with counts layer
132 batch_key : str, optional
133 Batch column
134 labels_key : str, optional
135 Cell type labels (uses scANVI if provided)
136 n_latent : int
137 Latent dimensions
138 n_layers : int
139 Encoder/decoder layers
140 max_epochs : int
141 Maximum training epochs
142 early_stopping : bool
143 Use early stopping
144 use_gpu : bool
145 Use GPU if available
147 Returns
148 -------
149 Trained model
150 """
151 import scvi
153 # Setup AnnData
154 scvi.model.SCVI.setup_anndata(
155 adata,
156 layer="counts",
157 batch_key=batch_key
158 )
160 if labels_key and labels_key in adata.obs.columns:
161 # Train scVI first
162 scvi_model = scvi.model.SCVI(
163 adata,
164 n_latent=n_latent,
165 n_layers=n_layers
166 )
167 scvi_model.train(
168 max_epochs=max_epochs,
169 early_stopping=early_stopping
170 )
172 # Initialize scANVI
173 model = scvi.model.SCANVI.from_scvi_model(
174 scvi_model,
175 labels_key=labels_key,
176 unlabeled_category="Unknown"
177 )
178 model.train(max_epochs=max_epochs // 4)
180 # Store representation
181 adata.obsm["X_scANVI"] = model.get_latent_representation()
182 else:
183 # Train scVI only
184 model = scvi.model.SCVI(
185 adata,
186 n_latent=n_latent,
187 n_layers=n_layers
188 )
189 model.train(
190 max_epochs=max_epochs,
191 early_stopping=early_stopping
192 )
194 # Store representation
195 adata.obsm["X_scVI"] = model.get_latent_representation()
197 return model
200def evaluate_integration(
201 adata,
202 batch_key: str,
203 label_key: str,
204 embedding_key: str = "X_scVI"
205) -> Dict[str, float]:
206 """
207 Evaluate integration quality using basic metrics.
209 Parameters
210 ----------
211 adata : AnnData
212 Integrated data
213 batch_key : str
214 Batch column
215 label_key : str
216 Cell type column
217 embedding_key : str
218 Key in obsm for embedding
220 Returns
221 -------
222 Dictionary of metrics
223 """
224 from sklearn.metrics import silhouette_score
225 from sklearn.neighbors import NearestNeighbors
227 X = adata.obsm[embedding_key]
228 batch = adata.obs[batch_key].values
229 labels = adata.obs[label_key].values
231 metrics = {}
233 # Silhouette scores
234 try:
235 # Cell type silhouette (higher = better separation)
236 metrics["silhouette_label"] = silhouette_score(X, labels)
238 # Batch silhouette (lower = better mixing)
239 metrics["silhouette_batch"] = silhouette_score(X, batch)
240 except Exception as e:
241 warnings.warn(f"Silhouette calculation failed: {e}")
243 # Batch mixing in neighbors
244 try:
245 nn = NearestNeighbors(n_neighbors=50)
246 nn.fit(X)
247 distances, indices = nn.kneighbors(X)
249 batch_mixing = []
250 for i in range(len(X)):
251 neighbor_batches = batch[indices[i]]
252 unique_batches = len(np.unique(neighbor_batches))
253 batch_mixing.append(unique_batches / len(np.unique(batch)))
255 metrics["batch_mixing"] = np.mean(batch_mixing)
256 except Exception as e:
257 warnings.warn(f"Batch mixing calculation failed: {e}")
259 return metrics
262def get_marker_genes(
263 model,
264 adata,
265 groupby: str,
266 n_genes: int = 10
267) -> Dict[str, List[str]]:
268 """
269 Get marker genes using scVI differential expression.
271 Parameters
272 ----------
273 model : scvi model
274 Trained scVI/scANVI model
275 adata : AnnData
276 Data used for training
277 groupby : str
278 Column to group cells by
279 n_genes : int
280 Number of top markers per group
282 Returns
283 -------
284 Dictionary of {group: [marker_genes]}
285 """
286 markers = {}
287 groups = adata.obs[groupby].unique()
289 for group in groups:
290 # Get DE results for this group vs rest
291 de_results = model.differential_expression(
292 groupby=groupby,
293 group1=group
294 )
296 # Filter and sort
297 de_sig = de_results[
298 (de_results["is_de_fdr_0.05"] == True) &
299 (de_results["lfc_mean"] > 0.5)
300 ].sort_values("lfc_mean", ascending=False)
302 markers[group] = de_sig.index[:n_genes].tolist()
304 return markers
307def plot_training_history(model, save_path: Optional[str] = None):
308 """
309 Plot model training history.
311 Parameters
312 ----------
313 model : scvi model
314 Trained model
315 save_path : str, optional
316 Path to save figure
317 """
318 import matplotlib.pyplot as plt
320 fig, axes = plt.subplots(1, 2, figsize=(12, 4))
322 # ELBO
323 if "elbo_train" in model.history:
324 axes[0].plot(model.history["elbo_train"], label="Train")
325 if "elbo_validation" in model.history:
326 axes[0].plot(model.history["elbo_validation"], label="Validation")
327 axes[0].set_xlabel("Epoch")
328 axes[0].set_ylabel("ELBO")
329 axes[0].legend()
330 axes[0].set_title("Training Loss")
332 # Reconstruction
333 if "reconstruction_loss_train" in model.history:
334 axes[1].plot(model.history["reconstruction_loss_train"], label="Train")
335 if "reconstruction_loss_validation" in model.history:
336 axes[1].plot(model.history["reconstruction_loss_validation"], label="Validation")
337 axes[1].set_xlabel("Epoch")
338 axes[1].set_ylabel("Reconstruction Loss")
339 axes[1].legend()
340 axes[1].set_title("Reconstruction Loss")
342 plt.tight_layout()
344 if save_path:
345 plt.savefig(save_path, dpi=150, bbox_inches="tight")
347 return fig
350def save_results(
351 model,
352 adata,
353 output_dir: str,
354 save_model: bool = True,
355 save_adata: bool = True,
356 plot_umap: bool = True
357):
358 """
359 Save model, processed data, and visualization.
361 Parameters
362 ----------
363 model : scvi model
364 Trained model
365 adata : AnnData
366 Processed data with latent representation
367 output_dir : str
368 Output directory path
369 save_model : bool
370 Save the trained model
371 save_adata : bool
372 Save the processed AnnData
373 plot_umap : bool
374 Generate and save UMAP plot
375 """
376 import os
377 import scanpy as sc
378 import matplotlib.pyplot as plt
380 os.makedirs(output_dir, exist_ok=True)
382 # Save model
383 if save_model:
384 model_path = os.path.join(output_dir, "model")
385 model.save(model_path)
386 print(f"Model saved to {model_path}")
388 # Save AnnData
389 if save_adata:
390 adata_path = os.path.join(output_dir, "adata_processed.h5ad")
391 adata.write(adata_path)
392 print(f"AnnData saved to {adata_path}")
394 # Generate UMAP if needed
395 if plot_umap:
396 # Determine which embedding to use
397 if "X_scANVI" in adata.obsm:
398 rep_key = "X_scANVI"
399 elif "X_scVI" in adata.obsm:
400 rep_key = "X_scVI"
401 else:
402 rep_key = None
404 if rep_key is not None:
405 # Compute neighbors and UMAP if not present
406 if "X_umap" not in adata.obsm:
407 sc.pp.neighbors(adata, use_rep=rep_key)
408 sc.tl.umap(adata)
410 # Plot
411 fig, axes = plt.subplots(1, 2, figsize=(12, 5))
413 # Plot by batch if available
414 batch_cols = [c for c in adata.obs.columns if 'batch' in c.lower()]
415 if batch_cols:
416 sc.pl.umap(adata, color=batch_cols[0], ax=axes[0], show=False, title="By Batch")
418 # Plot by cluster
419 if "leiden" not in adata.obs:
420 sc.tl.leiden(adata)
421 sc.pl.umap(adata, color="leiden", ax=axes[1], show=False, title="Clusters")
423 plt.tight_layout()
424 plot_path = os.path.join(output_dir, "umap.png")
425 plt.savefig(plot_path, dpi=150, bbox_inches="tight")
426 plt.close()
427 print(f"UMAP plot saved to {plot_path}")
430def auto_select_model(adata) -> str:
431 """
432 Suggest the best scvi-tools model based on available data.
434 Parameters
435 ----------
436 adata : AnnData
437 Data to analyze
439 Returns
440 -------
441 String with model recommendation and reasoning
442 """
443 suggestions = []
445 # Check for multi-modal data
446 if 'protein_expression' in adata.obsm:
447 suggestions.append({
448 'model': 'totalVI',
449 'reason': 'CITE-seq data detected (protein + RNA)',
450 'priority': 1
451 })
453 if 'spliced' in adata.layers and 'unspliced' in adata.layers:
454 suggestions.append({
455 'model': 'veloVI',
456 'reason': 'RNA velocity data detected (spliced + unspliced)',
457 'priority': 1
458 })
460 # Check for ATAC data indicators
461 if adata.n_vars > 100000: # Many peaks suggest ATAC
462 suggestions.append({
463 'model': 'PeakVI',
464 'reason': f'Large number of features ({adata.n_vars}) suggests ATAC-seq peaks',
465 'priority': 2
466 })
468 # Check for labels
469 label_cols = [c for c in adata.obs.columns if 'cell' in c.lower() or 'type' in c.lower() or 'label' in c.lower()]
470 has_labels = len(label_cols) > 0
472 # Check for batch info
473 batch_cols = [c for c in adata.obs.columns if 'batch' in c.lower() or 'sample' in c.lower()]
474 has_batch = len(batch_cols) > 0
476 if has_batch:
477 if has_labels:
478 suggestions.append({
479 'model': 'scANVI',
480 'reason': f'Batch info ({batch_cols[0]}) + labels ({label_cols[0]}) available',
481 'priority': 1
482 })
483 else:
484 suggestions.append({
485 'model': 'scVI',
486 'reason': f'Batch info ({batch_cols[0]}) available, no labels',
487 'priority': 1
488 })
489 else:
490 suggestions.append({
491 'model': 'scVI',
492 'reason': 'Standard scRNA-seq analysis',
493 'priority': 2
494 })
496 # Sort by priority
497 suggestions.sort(key=lambda x: x['priority'])
499 # Format output
500 lines = ["Recommended models (in order of priority):"]
501 for i, s in enumerate(suggestions, 1):
502 lines.append(f" {i}. {s['model']}: {s['reason']}")
504 return "\n".join(lines)
507def compare_integrations(
508 adata,
509 batch_key: str,
510 label_key: str,
511 embedding_keys: List[str] = None
512) -> Dict[str, Dict[str, float]]:
513 """
514 Compare multiple integration methods using standard metrics.
516 Parameters
517 ----------
518 adata : AnnData
519 Data with integration embeddings in obsm
520 batch_key : str
521 Batch column in obs
522 label_key : str
523 Cell type column in obs
524 embedding_keys : list, optional
525 Keys in obsm to compare (default: auto-detect)
527 Returns
528 -------
529 Dictionary of {embedding: {metric: value}}
530 """
531 from sklearn.metrics import silhouette_score
533 # Auto-detect embeddings
534 if embedding_keys is None:
535 embedding_keys = [k for k in adata.obsm.keys()
536 if k.startswith('X_') and 'umap' not in k.lower()]
538 results = {}
540 for key in embedding_keys:
541 if key not in adata.obsm:
542 continue
544 X = adata.obsm[key]
545 batch = adata.obs[batch_key].values
546 labels = adata.obs[label_key].values
548 metrics = {}
550 try:
551 # Silhouette scores
552 metrics["silhouette_label"] = silhouette_score(X, labels)
553 metrics["silhouette_batch"] = silhouette_score(X, batch)
555 # Combined score (higher label preservation, lower batch separation = better)
556 metrics["integration_score"] = metrics["silhouette_label"] - metrics["silhouette_batch"]
558 except Exception as e:
559 metrics["error"] = str(e)
561 results[key] = metrics
563 return results
566def quick_clustering(
567 adata,
568 use_rep: str = None,
569 resolution: float = 1.0,
570 n_neighbors: int = 15
571):
572 """
573 Quick clustering pipeline on latent representation.
575 Parameters
576 ----------
577 adata : AnnData
578 Data with latent representation
579 use_rep : str, optional
580 Key in obsm (auto-detects scVI/scANVI if not specified)
581 resolution : float
582 Leiden clustering resolution
583 n_neighbors : int
584 Number of neighbors for graph
586 Returns
587 -------
588 AnnData with neighbors, UMAP, and leiden clustering
589 """
590 import scanpy as sc
592 # Auto-detect representation
593 if use_rep is None:
594 if "X_scANVI" in adata.obsm:
595 use_rep = "X_scANVI"
596 elif "X_scVI" in adata.obsm:
597 use_rep = "X_scVI"
598 elif "X_totalVI" in adata.obsm:
599 use_rep = "X_totalVI"
600 elif "X_PeakVI" in adata.obsm:
601 use_rep = "X_PeakVI"
602 elif "X_MultiVI" in adata.obsm:
603 use_rep = "X_MultiVI"
604 else:
605 raise ValueError("No scvi-tools embedding found in obsm")
607 print(f"Using representation: {use_rep}")
609 # Compute neighbors
610 sc.pp.neighbors(adata, use_rep=use_rep, n_neighbors=n_neighbors)
612 # UMAP
613 sc.tl.umap(adata)
615 # Leiden clustering
616 sc.tl.leiden(adata, resolution=resolution)
618 print(f"Found {adata.obs['leiden'].nunique()} clusters")
620 return adata
623if __name__ == "__main__":
624 print("scvi-tools model utilities")
625 print("\nAvailable functions:")
626 print(" - prepare_adata: Standard data preparation (QC, HVG, layer setup)")
627 print(" - train_scvi: Train scVI or scANVI with sensible defaults")
628 print(" - evaluate_integration: Compute batch mixing and silhouette metrics")
629 print(" - get_marker_genes: Extract markers using scVI differential expression")
630 print(" - plot_training_history: Visualize training convergence")
631 print(" - save_results: Save model, data, and visualizations")
632 print(" - auto_select_model: Suggest best model for your data")
633 print(" - compare_integrations: Compare multiple integration embeddings")
634 print(" - quick_clustering: Quick clustering on latent representation")