Setting the file. One moment.
Model Utils · Scvi Tools · anthropics/knowledge-work-plugins · Skills Docs
Back to the top of the page Spatial Deconvolution
22
Validate Data
Tech Debt
62
Recruiting Pipeline
71
Vendor Check
125
Zoom Meeting SDK Web
88
Vendor Review
181
Create An Asset
Video Sdk/web
16 of 21 scripts/ model_utils.py
Python · 634 lines · 17 KB
15 def get_mito_genes (adata) -> np.ndarray:
16 """
17 Identify mitochondrial genes for both human and mouse data.
18
19 Handles common prefixes:
20 - Human: MT- (e.g., MT-CO1, MT-ND1)
21 - Mouse: mt- or Mt- (e.g., mt-Co1, Mt-Nd1)
22
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 )
32
33
34 def 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.
46
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
65
66 Returns
67 -------
68 AnnData prepared for scvi-tools
69 """
70 if copy:
71 adata = adata.copy()
72
73 # Calculate QC metrics
74 adata.var[ 'mt' ] = get_mito_genes(adata)
75 sc.pp.calculate_qc_metrics(adata, qc_vars = [ 'mt' ], inplace = True )
76
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()
81
82 # Filter genes
83 sc.pp.filter_genes(adata, min_cells = min_cells)
84
85 # Store raw counts
86 adata.layers[ "counts" ] = adata.X.copy()
87
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()
104
105 # Subset to HVGs
106 adata = adata[:, adata.var[ 'highly_variable' ]].copy()
107
108 print ( f "Prepared AnnData: { adata.shape } " )
109 if batch_key:
110 print ( f "Batches: { adata.obs[batch_key].nunique() } " )
111
112 return adata
113
114
115 def 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.
127
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
146
147 Returns
148 -------
149 Trained model
150 """
151 import scvi
152
153 # Setup AnnData
154 scvi.model. SCVI .setup_anndata(
155 adata,
156 layer = "counts" ,
157 batch_key = batch_key
158 )
159
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 )
171
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 )
179
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 )
193
194 # Store representation
195 adata.obsm[ "X_scVI" ] = model.get_latent_representation()
196
197 return model
198
199
200 def 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.
208
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
219
220 Returns
221 -------
222 Dictionary of metrics
223 """
224 from sklearn.metrics import silhouette_score
225 from sklearn.neighbors import NearestNeighbors
226
227 X = adata.obsm[embedding_key]
228 batch = adata.obs[batch_key].values
229 labels = adata.obs[label_key].values
230
231 metrics = {}
232
233 # Silhouette scores
234 try :
235 # Cell type silhouette (higher = better separation)
236 metrics[ "silhouette_label" ] = silhouette_score(X, labels)
237
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 } " )
242
243 # Batch mixing in neighbors
244 try :
245 nn = NearestNeighbors( n_neighbors = 50 )
246 nn.fit(X)
247 distances, indices = nn.kneighbors(X)
248
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)))
254
255 metrics[ "batch_mixing" ] = np.mean(batch_mixing)
256 except Exception as e:
257 warnings.warn( f "Batch mixing calculation failed: { e } " )
258
259 return metrics
260
261
262 def 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.
270
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
281
282 Returns
283 -------
284 Dictionary of {group: [marker_genes]}
285 """
286 markers = {}
287 groups = adata.obs[groupby].unique()
288
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 )
295
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 )
301
302 markers[group] = de_sig.index[:n_genes].tolist()
303
304 return markers
305
306
307 def plot_training_history (model, save_path: Optional[ str ] = None ):
308 """
309 Plot model training history.
310
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
319
320 fig, axes = plt.subplots( 1 , 2 , figsize = ( 12 , 4 ))
321
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" )
331
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" )
341
342 plt.tight_layout()
343
344 if save_path:
345 plt.savefig(save_path, dpi = 150 , bbox_inches = "tight" )
346
347 return fig
348
349
350 def 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.
360
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
379
380 os.makedirs(output_dir, exist_ok = True )
381
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 } " )
387
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 } " )
393
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
403
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)
409
410 # Plot
411 fig, axes = plt.subplots( 1 , 2 , figsize = ( 12 , 5 ))
412
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" )
417
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" )
422
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 } " )
428
429
430 def auto_select_model (adata) -> str :
431 """
432 Suggest the best scvi-tools model based on available data.
433
434 Parameters
435 ----------
436 adata : AnnData
437 Data to analyze
438
439 Returns
440 -------
441 String with model recommendation and reasoning
442 """
443 suggestions = []
444
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 })
452
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 })
459
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 })
467
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
471
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
475
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 })
495
496 # Sort by priority
497 suggestions.sort( key =lambda x: x[ 'priority' ])
498
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' ] } " )
503
504 return " \n " .join(lines)
505
506
507 def 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.
515
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)
526
527 Returns
528 -------
529 Dictionary of {embedding: {metric: value}}
530 """
531 from sklearn.metrics import silhouette_score
532
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()]
537
538 results = {}
539
540 for key in embedding_keys:
541 if key not in adata.obsm:
542 continue
543
544 X = adata.obsm[key]
545 batch = adata.obs[batch_key].values
546 labels = adata.obs[label_key].values
547
548 metrics = {}
549
550 try :
551 # Silhouette scores
552 metrics[ "silhouette_label" ] = silhouette_score(X, labels)
553 metrics[ "silhouette_batch" ] = silhouette_score(X, batch)
554
555 # Combined score (higher label preservation, lower batch separation = better)
556 metrics[ "integration_score" ] = metrics[ "silhouette_label" ] - metrics[ "silhouette_batch" ]
557
558 except Exception as e:
559 metrics[ "error" ] = str (e)
560
561 results[key] = metrics
562
563 return results
564
565
566 def 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.
574
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
585
586 Returns
587 -------
588 AnnData with neighbors, UMAP, and leiden clustering
589 """
590 import scanpy as sc
591
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" )
606
607 print ( f "Using representation: { use_rep } " )
608
609 # Compute neighbors
610 sc.pp.neighbors(adata, use_rep = use_rep, n_neighbors = n_neighbors)
611
612 # UMAP
613 sc.tl.umap(adata)
614
615 # Leiden clustering
616 sc.tl.leiden(adata, resolution = resolution)
617
618 print ( f "Found { adata.obs[ 'leiden' ].nunique() } clusters" )
619
620 return adata
621
622
623 if __name__ == "__main__" :
624 print ( "scvi-tools model utilities" )
625 print ( " \n Available 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" )