Skills › Knowledge Work Plugins › scvi-tools › transfer_labels.py
transfer_labels.py
224 lines · bio-research/skills/scvi-tools/scripts/transfer_labels.py at 2574e29
1#!/usr/bin/env python3
2"""
3Transfer cell type labels from reference to query using scANVI.
5Maps query cells to a pre-trained reference model and predicts cell types.
7Usage:
8 python transfer_labels.py reference_model/ query.h5ad output_dir/
9 python transfer_labels.py reference_model/ query.h5ad output_dir/ --confidence 0.7
10"""
12import argparse
13import os
14import sys
17def transfer_labels(
18 reference_model,
19 adata_query,
20 max_epochs=100,
21 confidence_threshold=0.5
22):
23 """
24 Transfer labels from reference to query.
26 Parameters
27 ----------
28 reference_model : SCANVI model
29 Pre-trained scANVI model
30 adata_query : AnnData
31 Query data to annotate
32 max_epochs : int
33 Fine-tuning epochs
34 confidence_threshold : float
35 Minimum confidence for predictions
37 Returns
38 -------
39 AnnData with predictions
40 """
41 import scvi
42 import numpy as np
44 # Get reference genes
45 ref_genes = reference_model.adata.var_names
46 print(f"Reference genes: {len(ref_genes)}")
48 # Check gene overlap
49 query_genes = adata_query.var_names
50 common = ref_genes.intersection(query_genes)
51 print(f"Query genes: {len(query_genes)}")
52 print(f"Common genes: {len(common)} ({len(common)/len(ref_genes)*100:.1f}%)")
54 if len(common) < len(ref_genes) * 0.5:
55 print("Warning: Less than 50% gene overlap. Results may be unreliable.")
57 # Subset query to reference genes
58 # Missing genes will be filled with zeros
59 adata_query = adata_query[:, adata_query.var_names.isin(ref_genes)].copy()
61 # Ensure counts layer
62 if "counts" not in adata_query.layers:
63 adata_query.layers["counts"] = adata_query.X.copy()
65 # Prepare query for mapping
66 print("Preparing query data...")
67 scvi.model.SCANVI.prepare_query_anndata(adata_query, reference_model)
69 # Create query model
70 print("Creating query model...")
71 query_model = scvi.model.SCANVI.load_query_data(
72 adata_query,
73 reference_model
74 )
76 # Fine-tune
77 print(f"Fine-tuning ({max_epochs} epochs)...")
78 query_model.train(
79 max_epochs=max_epochs,
80 plan_kwargs={"weight_decay": 0.0}
81 )
83 # Get predictions
84 print("Getting predictions...")
85 predictions = query_model.predict()
86 soft_predictions = query_model.predict(soft=True)
88 adata_query.obs["predicted_cell_type"] = predictions
89 adata_query.obs["prediction_confidence"] = soft_predictions.max(axis=1)
90 adata_query.obs["confident_prediction"] = adata_query.obs["prediction_confidence"] >= confidence_threshold
92 # Get latent representation
93 adata_query.obsm["X_scANVI"] = query_model.get_latent_representation()
95 # Stats
96 n_confident = adata_query.obs["confident_prediction"].sum()
97 print(f"\nPrediction summary:")
98 print(f" Total cells: {adata_query.n_obs}")
99 print(f" Confident (>= {confidence_threshold}): {n_confident} ({n_confident/adata_query.n_obs*100:.1f}%)")
100 print(f" Mean confidence: {adata_query.obs['prediction_confidence'].mean():.3f}")
102 print("\nPredicted cell types:")
103 print(adata_query.obs["predicted_cell_type"].value_counts())
105 return adata_query, query_model
108def plot_predictions(adata, output_dir):
109 """Plot prediction results."""
110 import scanpy as sc
111 import matplotlib.pyplot as plt
113 # Compute UMAP if needed
114 if "X_umap" not in adata.obsm:
115 sc.pp.neighbors(adata, use_rep="X_scANVI")
116 sc.tl.umap(adata)
118 # Plot
119 fig, axes = plt.subplots(1, 3, figsize=(15, 4))
121 sc.pl.umap(adata, color="predicted_cell_type", ax=axes[0], show=False,
122 title="Predicted Cell Type")
123 sc.pl.umap(adata, color="prediction_confidence", ax=axes[1], show=False,
124 title="Prediction Confidence", cmap="viridis")
125 sc.pl.umap(adata, color="confident_prediction", ax=axes[2], show=False,
126 title="Confident Predictions")
128 plt.tight_layout()
129 plot_path = os.path.join(output_dir, "predictions.png")
130 plt.savefig(plot_path, dpi=150, bbox_inches="tight")
131 plt.close()
132 print(f"Prediction plot saved to {plot_path}")
135def main():
136 parser = argparse.ArgumentParser(
137 description="Transfer cell type labels using scANVI",
138 formatter_class=argparse.RawDescriptionHelpFormatter,
139 epilog="""
140Examples:
141 # Basic label transfer
142 python transfer_labels.py reference_model/ query.h5ad results/
144 # With confidence threshold
145 python transfer_labels.py reference_model/ query.h5ad results/ --confidence 0.7
147 # More fine-tuning
148 python transfer_labels.py reference_model/ query.h5ad results/ --max-epochs 200
149 """
150 )
151 parser.add_argument("model_dir", help="Directory containing reference scANVI model")
152 parser.add_argument("query", help="Query h5ad file to annotate")
153 parser.add_argument("output_dir", help="Output directory")
154 parser.add_argument("--reference-adata", help="Reference adata used for training (if not saved with model)")
155 parser.add_argument("--max-epochs", type=int, default=100,
156 help="Fine-tuning epochs (default: 100)")
157 parser.add_argument("--confidence", type=float, default=0.5,
158 help="Confidence threshold (default: 0.5)")
160 args = parser.parse_args()
162 try:
163 import scvi
164 import scanpy as sc
165 except ImportError:
166 print("Error: scvi-tools and scanpy required")
167 sys.exit(1)
169 # Create output directory
170 os.makedirs(args.output_dir, exist_ok=True)
172 # Load query data
173 print(f"Loading query data: {args.query}")
174 adata_query = sc.read_h5ad(args.query)
175 print(f"Query: {adata_query.shape}")
177 # Load reference model
178 print(f"Loading reference model: {args.model_dir}")
179 if args.reference_adata:
180 ref_adata = sc.read_h5ad(args.reference_adata)
181 reference_model = scvi.model.SCANVI.load(args.model_dir, adata=ref_adata)
182 else:
183 # Try loading without adata (works if model was saved with adata)
184 try:
185 reference_model = scvi.model.SCANVI.load(args.model_dir)
186 except ValueError as e:
187 if "no saved anndata" in str(e).lower():
188 print("Error: Model was saved without adata. Please provide --reference-adata")
189 sys.exit(1)
190 raise
191 print(f"Reference: {reference_model.adata.shape}")
193 # Transfer labels
194 adata_annotated, query_model = transfer_labels(
195 reference_model,
196 adata_query,
197 max_epochs=args.max_epochs,
198 confidence_threshold=args.confidence
199 )
201 # Save results
202 adata_path = os.path.join(args.output_dir, "query_annotated.h5ad")
203 adata_annotated.write_h5ad(adata_path)
204 print(f"\nAnnotated data saved to {adata_path}")
206 # Save query model
207 model_path = os.path.join(args.output_dir, "query_model")
208 query_model.save(model_path)
209 print(f"Query model saved to {model_path}")
211 # Save predictions CSV
212 pred_df = adata_annotated.obs[["predicted_cell_type", "prediction_confidence", "confident_prediction"]]
213 pred_path = os.path.join(args.output_dir, "predictions.csv")
214 pred_df.to_csv(pred_path)
215 print(f"Predictions saved to {pred_path}")
217 # Plot
218 plot_predictions(adata_annotated, args.output_dir)
220 print("\nDone!")
223if __name__ == "__main__":
224 main()