Scarches Mapping

import scvi
import scanpy as sc
import numpy as np

print(f"scvi-tools version: {scvi.__version__}")

# Load saved reference model
# The model must have been trained with scvi-tools
reference_model = scvi.model.SCVI.load("reference_model/")

# Or load scANVI for label transfer
reference_model = scvi.model.SCANVI.load("reference_scanvi_model/")

# Check model info
print(f"Model type: {type(reference_model)}")
print(f"Training data shape: {reference_model.adata.shape}")

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

# CRITICAL: Match genes to reference
reference_genes = reference_model.adata.var_names
query_genes = adata_query.var_names

# Check overlap
common_genes = reference_genes.intersection(query_genes)
print(f"Reference genes: {len(reference_genes)}")
print(f"Query genes: {len(query_genes)}")
print(f"Overlap: {len(common_genes)}")

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

# Handle missing genes (filled with zeros automatically by prepare_query_anndata)

# Store raw counts
adata_query.layers["counts"] = adata_query.X.copy()

# Prepare query for mapping
# This aligns the query data structure to match the reference
scvi.model.SCVI.prepare_query_anndata(adata_query, reference_model)

# Create query model from reference
# This initializes with reference weights
query_model = scvi.model.SCVI.load_query_data(
    adata_query,
    reference_model
)

# The query model inherits:
# - Reference architecture
# - Reference encoder weights (frozen by default)
# - Decoder is fine-tuned for query

# Fine-tune the query model
# This adjusts decoder weights for query-specific effects
query_model.train(
    max_epochs=200,
    plan_kwargs={
        "weight_decay": 0.0  # Less regularization for fine-tuning
    }
)

# Check training
query_model.history['elbo_train'].plot()

# Get latent representation
# Query cells are embedded in same space as reference
adata_query.obsm["X_scVI"] = query_model.get_latent_representation()

# Visualize
sc.pp.neighbors(adata_query, use_rep="X_scVI")
sc.tl.umap(adata_query)
sc.pl.umap(adata_query, color=['cell_type', 'batch'])

# Reference must be scANVI model (trained with labels)
reference_scanvi = scvi.model.SCANVI.load("scanvi_reference/")

# Check available labels
print("Reference cell types:")
print(reference_scanvi.adata.obs['cell_type'].value_counts())

# Prepare query
adata_query.layers["counts"] = adata_query.X.copy()
adata_query = adata_query[:, reference_scanvi.adata.var_names].copy()

scvi.model.SCANVI.prepare_query_anndata(adata_query, reference_scanvi)

# Create query model
query_scanvi = scvi.model.SCANVI.load_query_data(
    adata_query,
    reference_scanvi
)

# Fine-tune
query_scanvi.train(
    max_epochs=100,
    plan_kwargs={"weight_decay": 0.0}
)

# Predict cell types
predictions = query_scanvi.predict()
adata_query.obs["predicted_cell_type"] = predictions

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

# Latent representation
adata_query.obsm["X_scANVI"] = query_scanvi.get_latent_representation()

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

# Distribution of predictions
print(adata_query.obs['predicted_cell_type'].value_counts())

# Confidence statistics
print(f"Mean confidence: {adata_query.obs['prediction_confidence'].mean():.3f}")
print(f"Low confidence (<0.5): {(adata_query.obs['prediction_confidence'] < 0.5).sum()}")

# Filter low-confidence predictions
high_conf = adata_query[adata_query.obs['prediction_confidence'] >= 0.7].copy()
print(f"High confidence cells: {len(high_conf)} ({len(high_conf)/len(adata_query)*100:.1f}%)")

# Load reference model
reference_model = scvi.model.SCVI.load("reference_model/")

# Get reference representation (before surgery)
adata_ref = reference_model.adata
adata_ref.obsm["X_scVI_before"] = reference_model.get_latent_representation()

# Add batch information
adata_ref.obs["dataset"] = "reference"
adata_query.obs["dataset"] = "query"

# Combine
adata_combined = sc.concat([adata_ref, adata_query])
adata_combined.layers["counts"] = adata_combined.X.copy()

# Option A: Use load_query_data (recommended)
scvi.model.SCVI.prepare_query_anndata(adata_query, reference_model)
extended_model = scvi.model.SCVI.load_query_data(adata_query, reference_model)
extended_model.train(max_epochs=200)

# Option B: Retrain with combined data (if query is large)
# This doesn't preserve reference exactly but may give better results
scvi.model.SCVI.setup_anndata(
    adata_combined,
    layer="counts",
    batch_key="dataset"
)
new_model = scvi.model.SCVI(adata_combined, n_latent=30)
new_model.train(max_epochs=200)

# Get latent representations
adata_ref.obsm["X_scVI"] = reference_model.get_latent_representation()
adata_query.obsm["X_scVI"] = query_model.get_latent_representation()

# Combine for visualization
adata_ref.obs["source"] = "reference"
adata_query.obs["source"] = "query"
adata_combined = sc.concat([adata_ref, adata_query])

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

# Visualize
import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

sc.pl.umap(adata_combined, color="source", ax=axes[0], show=False, title="Source")
sc.pl.umap(adata_combined, color="cell_type", ax=axes[1], show=False, title="Cell Type")
sc.pl.umap(adata_combined, color="batch", ax=axes[2], show=False, title="Batch")

plt.tight_layout()

from huggingface_hub import hf_hub_download

# Download model files
model_dir = hf_hub_download(
    repo_id="scvi-tools/model-name",  # Replace with actual repo
    filename="model.pt",
    local_dir="./downloaded_model/"
)

# Load model
atlas_model = scvi.model.SCANVI.load(model_dir)

# Many CellxGene datasets provide pre-trained models
# Check dataset documentation for model availability
# https://cellxgene.cziscience.com/

# Example workflow:
# 1. Download reference dataset and model
# 2. Load model: model = scvi.model.SCANVI.load("cellxgene_model/")
# 3. Map your query data using steps above

def map_query_to_reference(
    adata_query,
    reference_model_path,
    model_type="scanvi",
    max_epochs=100,
    confidence_threshold=0.5
):
    """
    Map query data to pre-trained reference model.

    Parameters
    ----------
    adata_query : AnnData
        Query data with raw counts
    reference_model_path : str
        Path to saved reference model
    model_type : str
        "scvi" or "scanvi"
    max_epochs : int
        Fine-tuning epochs
    confidence_threshold : float
        Minimum prediction confidence (for scANVI)

    Returns
    -------
    Mapped AnnData with predictions (if scANVI)
    """
    import scvi

    # Load reference
    if model_type == "scanvi":
        reference_model = scvi.model.SCANVI.load(reference_model_path)
        ModelClass = scvi.model.SCANVI
    else:
        reference_model = scvi.model.SCVI.load(reference_model_path)
        ModelClass = scvi.model.SCVI

    # Prepare query
    adata_query = adata_query.copy()
    adata_query = adata_query[:, reference_model.adata.var_names].copy()
    adata_query.layers["counts"] = adata_query.X.copy()

    # Map query
    ModelClass.prepare_query_anndata(adata_query, reference_model)
    query_model = ModelClass.load_query_data(adata_query, reference_model)

    # Fine-tune
    query_model.train(
        max_epochs=max_epochs,
        plan_kwargs={"weight_decay": 0.0}
    )

    # Get results
    rep_key = "X_scANVI" if model_type == "scanvi" else "X_scVI"
    adata_query.obsm[rep_key] = query_model.get_latent_representation()

    if model_type == "scanvi":
        adata_query.obs["predicted_cell_type"] = query_model.predict()
        soft = query_model.predict(soft=True)
        adata_query.obs["prediction_confidence"] = soft.max(axis=1)
        adata_query.obs["confident"] = adata_query.obs["prediction_confidence"] >= confidence_threshold

    # Compute UMAP
    sc.pp.neighbors(adata_query, use_rep=rep_key)
    sc.tl.umap(adata_query)

    return adata_query, query_model


# Usage
adata_mapped, model = map_query_to_reference(
    adata_query,
    "reference_scanvi_model/",
    model_type="scanvi"
)

# Visualize
sc.pl.umap(adata_mapped, color=['predicted_cell_type', 'prediction_confidence'])