Label Transfer

import scvi
import scanpy as sc

# Load reference atlas
adata_ref = sc.read_h5ad("reference_atlas.h5ad")

# Check annotations
print(f"Reference cells: {adata_ref.n_obs}")
print(f"Cell types: {adata_ref.obs['cell_type'].nunique()}")
print(adata_ref.obs['cell_type'].value_counts())

# Ensure raw counts
adata_ref.layers["counts"] = adata_ref.raw.X.copy() if adata_ref.raw else adata_ref.X.copy()

# HVG selection
sc.pp.highly_variable_genes(
    adata_ref,
    n_top_genes=3000,
    flavor="seurat_v3",
    batch_key="batch" if "batch" in adata_ref.obs else None,
    layer="counts"
)
adata_ref = adata_ref[:, adata_ref.var["highly_variable"]].copy()

# First train scVI (unlabeled)
scvi.model.SCVI.setup_anndata(
    adata_ref,
    layer="counts",
    batch_key="batch"
)

scvi_ref = scvi.model.SCVI(adata_ref, n_latent=30)
scvi_ref.train(max_epochs=200)

# Initialize scANVI from scVI
scanvi_ref = scvi.model.SCANVI.from_scvi_model(
    scvi_ref,
    labels_key="cell_type",
    unlabeled_category="Unknown"
)

# Train scANVI
scanvi_ref.train(max_epochs=50)

# Save for later use
scanvi_ref.save("scanvi_reference_model/")

# Load query data
adata_query = sc.read_h5ad("query_data.h5ad")

# CRITICAL: Use same genes as reference
common_genes = adata_ref.var_names.intersection(adata_query.var_names)
print(f"Common genes: {len(common_genes)}")

# Subset query to reference genes
adata_query = adata_query[:, adata_ref.var_names].copy()

# Handle missing genes (set to 0)
missing_genes = set(adata_ref.var_names) - set(adata_query.var_names)
if missing_genes:
    # Add missing genes with zero expression
    import numpy as np
    from scipy.sparse import csr_matrix
    
    zero_matrix = csr_matrix((adata_query.n_obs, len(missing_genes)))
    # ... concat and reorder to match reference
    
# Store counts
adata_query.layers["counts"] = adata_query.X.copy()

# Prepare query data for mapping
scvi.model.SCANVI.prepare_query_anndata(adata_query, scanvi_ref)

# Create query model from reference
scanvi_query = scvi.model.SCANVI.load_query_data(
    adata_query,
    scanvi_ref
)

# Fine-tune on query (optional but recommended)
scanvi_query.train(
    max_epochs=100,
    plan_kwargs={"weight_decay": 0.0}
)

# Get predictions
adata_query.obs["predicted_cell_type"] = scanvi_query.predict()

# Get prediction probabilities
soft_predictions = scanvi_query.predict(soft=True)
adata_query.obs["prediction_score"] = soft_predictions.max(axis=1)

# Confidence scores
print(f"Mean prediction confidence: {adata_query.obs['prediction_score'].mean():.3f}")

# Low confidence predictions
low_conf = adata_query.obs['prediction_score'] < 0.5
print(f"Low confidence cells: {low_conf.sum()} ({low_conf.mean()*100:.1f}%)")

# Visualize
sc.pp.neighbors(adata_query, use_rep="X_scANVI")
sc.tl.umap(adata_query)
sc.pl.umap(adata_query, color=['predicted_cell_type', 'prediction_score'])

# scvi-tools maintains models on HuggingFace
# Check: https://huggingface.co/scvi-tools

# Example: Load pre-trained model
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="scvi-tools/example-model",
    filename="model.pt"
)

# Load model
model = scvi.model.SCANVI.load(model_path, adata=adata_query)

# Many atlases provide pre-trained models
# Example workflow with CellTypist-style model

# Download reference model
# model = scvi.model.SCANVI.load("atlas_model/", adata=adata_query)

# Load existing reference model
scanvi_ref = scvi.model.SCANVI.load("reference_model/")

# Surgery: prepare for query integration
scanvi_ref.freeze_layers()

# Map query data
scvi.model.SCANVI.prepare_query_anndata(adata_query, scanvi_ref)
scanvi_query = scvi.model.SCANVI.load_query_data(adata_query, scanvi_ref)

# Train only query-specific parameters
scanvi_query.train(
    max_epochs=200,
    plan_kwargs={"weight_decay": 0.0}
)

# Concatenate for joint visualization
adata_ref.obs["dataset"] = "reference"
adata_query.obs["dataset"] = "query"

# Get latent representations
adata_ref.obsm["X_scANVI"] = scanvi_ref.get_latent_representation()
adata_query.obsm["X_scANVI"] = scanvi_query.get_latent_representation()

# Combine
adata_combined = sc.concat([adata_ref, adata_query])

# Compute combined UMAP
sc.pp.neighbors(adata_combined, use_rep="X_scANVI")
sc.tl.umap(adata_combined)

# Plot
sc.pl.umap(
    adata_combined,
    color=["dataset", "cell_type", "predicted_cell_type"],
    ncols=2
)

# Filter predictions by confidence
confidence_threshold = 0.7

high_conf = adata_query[adata_query.obs['prediction_score'] >= confidence_threshold].copy()
low_conf = adata_query[adata_query.obs['prediction_score'] < confidence_threshold].copy()

print(f"High confidence: {len(high_conf)} ({len(high_conf)/len(adata_query)*100:.1f}%)")
print(f"Low confidence: {len(low_conf)} ({len(low_conf)/len(adata_query)*100:.1f}%)")

# Validate predictions with known markers
markers = {
    'T cells': ['CD3D', 'CD3E'],
    'B cells': ['CD19', 'MS4A1'],
    'Monocytes': ['CD14', 'LYZ']
}

for ct, genes in markers.items():
    ct_cells = adata_query[adata_query.obs['predicted_cell_type'] == ct]
    if len(ct_cells) > 0:
        for gene in genes:
            if gene in adata_query.var_names:
                expr = ct_cells[:, gene].X.mean()
                print(f"{ct} - {gene}: {expr:.3f}")

def transfer_labels(
    adata_ref,
    adata_query,
    cell_type_key="cell_type",
    batch_key=None,
    n_top_genes=3000,
    confidence_threshold=0.5
):
    """
    Transfer cell type labels from reference to query.
    
    Parameters
    ----------
    adata_ref : AnnData
        Annotated reference data
    adata_query : AnnData
        Unannotated query data
    cell_type_key : str
        Column with cell type annotations in reference
    batch_key : str, optional
        Batch column
    n_top_genes : int
        Number of HVGs
    confidence_threshold : float
        Minimum confidence for predictions
        
    Returns
    -------
    AnnData with predictions
    """
    import scvi
    import scanpy as sc
    
    # Prepare reference
    adata_ref = adata_ref.copy()
    adata_ref.layers["counts"] = adata_ref.X.copy()
    
    sc.pp.highly_variable_genes(
        adata_ref,
        n_top_genes=n_top_genes,
        flavor="seurat_v3",
        batch_key=batch_key,
        layer="counts"
    )
    adata_ref = adata_ref[:, adata_ref.var["highly_variable"]].copy()
    
    # Train reference model
    scvi.model.SCVI.setup_anndata(adata_ref, layer="counts", batch_key=batch_key)
    scvi_ref = scvi.model.SCVI(adata_ref, n_latent=30)
    scvi_ref.train(max_epochs=200)
    
    scanvi_ref = scvi.model.SCANVI.from_scvi_model(
        scvi_ref,
        labels_key=cell_type_key,
        unlabeled_category="Unknown"
    )
    scanvi_ref.train(max_epochs=50)
    
    # Prepare query
    adata_query = adata_query[:, adata_ref.var_names].copy()
    adata_query.layers["counts"] = adata_query.X.copy()
    
    # Map query
    scvi.model.SCANVI.prepare_query_anndata(adata_query, scanvi_ref)
    scanvi_query = scvi.model.SCANVI.load_query_data(adata_query, scanvi_ref)
    scanvi_query.train(max_epochs=100, plan_kwargs={"weight_decay": 0.0})
    
    # Get predictions
    adata_query.obs["predicted_cell_type"] = scanvi_query.predict()
    soft = scanvi_query.predict(soft=True)
    adata_query.obs["prediction_score"] = soft.max(axis=1)
    
    # Mark low confidence
    adata_query.obs["confident_prediction"] = adata_query.obs["prediction_score"] >= confidence_threshold
    
    # Add latent representation
    adata_query.obsm["X_scANVI"] = scanvi_query.get_latent_representation()
    
    return adata_query, scanvi_ref, scanvi_query

# Usage
adata_annotated, ref_model, query_model = transfer_labels(
    adata_ref,
    adata_query,
    cell_type_key="cell_type"
)

# Visualize
sc.pp.neighbors(adata_annotated, use_rep="X_scANVI")
sc.tl.umap(adata_annotated)
sc.pl.umap(adata_annotated, color=['predicted_cell_type', 'prediction_score'])