Batch Correction Sysvi

import scvi
import scanpy as sc
import numpy as np

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

                    ┌─────────────────┐
Input counts ──────►│    Encoder      │
                    │                 │
System info ───────►│  (conditioned)  │
                    └────────┬────────┘
                             │
                    ┌────────▼────────┐
                    │   Latent z      │
                    │  (biological)   │
                    └────────┬────────┘
                             │
                    ┌────────▼────────┐
System info ───────►│    Decoder      │
                    │  (conditioned)  │
                    └────────┬────────┘
                             │
                    Reconstructed counts

# Load datasets from different systems
adata1 = sc.read_h5ad("10x_data.h5ad")
adata2 = sc.read_h5ad("smartseq_data.h5ad")

# Add system labels
adata1.obs["system"] = "10x"
adata2.obs["system"] = "Smart-seq2"

# Add batch labels (within system)
# e.g., different samples within each technology

# Concatenate
adata = sc.concat([adata1, adata2])

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

# Select HVGs considering both batch and system
sc.pp.highly_variable_genes(
    adata,
    n_top_genes=4000,  # More genes for cross-system
    flavor="seurat_v3",
    batch_key="system",  # Consider system for HVG
    layer="counts"
)

# Optionally: ensure overlap between systems
# Check HVGs are expressed in both systems
adata = adata[:, adata.var["highly_variable"]].copy()

# Setup AnnData
# Note: sysVI may be accessed differently depending on version
# Check scvi-tools documentation for current API

scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="sample",           # Within-system batches
    categorical_covariate_keys=["system"]  # System-level covariate
)

# For true sysVI (if available in your version)
# scvi.model.SysVI.setup_anndata(...)

# Create model with system awareness
model = scvi.model.SCVI(
    adata,
    n_latent=30,
    n_layers=2,
    gene_likelihood="nb"
)

# Train
model.train(max_epochs=300)

# Get latent representation
adata.obsm["X_integrated"] = model.get_latent_representation()

# Clustering and visualization
sc.pp.neighbors(adata, use_rep="X_integrated")
sc.tl.umap(adata)
sc.tl.leiden(adata)

# Check integration
sc.pl.umap(adata, color=["system", "leiden", "cell_type"])

import scanpy.external as sce

# First run PCA
sc.pp.pca(adata)

# Apply Harmony for initial alignment
sce.pp.harmony_integrate(adata, key="system")

# Then train scVI on Harmony-corrected embedding
# Or use Harmony representation directly

# Include system as categorical covariate
scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="sample",
    categorical_covariate_keys=["system", "technology_version"]
)

model = scvi.model.SCVI(adata, n_latent=30)
model.train()

# Train separate models
scvi.model.SCVI.setup_anndata(adata1, layer="counts", batch_key="sample")
model1 = scvi.model.SCVI(adata1)
model1.train()

scvi.model.SCVI.setup_anndata(adata2, layer="counts", batch_key="sample")
model2 = scvi.model.SCVI(adata2)
model2.train()

# Get latent spaces
adata1.obsm["X_scVI"] = model1.get_latent_representation()
adata2.obsm["X_scVI"] = model2.get_latent_representation()

# Align with CCA or Harmony
# ... additional alignment step

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

# Color by system
sc.pl.umap(adata, color="system", ax=axes[0], show=False, title="By System")

# Color by cell type
sc.pl.umap(adata, color="cell_type", ax=axes[1], show=False, title="By Cell Type")

# Color by expression of marker
sc.pl.umap(adata, color="CD3D", ax=axes[2], show=False, title="CD3D Expression")

plt.tight_layout()

# Using scib-metrics
from scib_metrics.benchmark import Benchmarker

bm = Benchmarker(
    adata,
    batch_key="system",
    label_key="cell_type",
    embedding_obsm_keys=["X_integrated"]
)

bm.benchmark()

# Key metrics:
# - Batch mixing (ASW_batch, Graph connectivity)
# - Bio conservation (NMI, ARI, ASW_label)

# Local Inverse Simpson's Index
from scib_metrics import lisi

# Batch LISI (higher = better mixing)
batch_lisi = lisi.ilisi_graph(
    adata,
    batch_key="system",
    use_rep="X_integrated"
)

# Cell type LISI (lower = better preservation)
ct_lisi = lisi.clisi_graph(
    adata,
    label_key="cell_type", 
    use_rep="X_integrated"
)

print(f"Batch LISI: {batch_lisi.mean():.3f}")
print(f"Cell type LISI: {ct_lisi.mean():.3f}")

# Find common genes
common_genes = adata1.var_names.intersection(adata2.var_names)
print(f"Common genes: {len(common_genes)}")

# If too few, use gene mapping
# Or impute missing genes

# Add depth as continuous covariate
adata.obs["log_counts"] = np.log1p(adata.obs["total_counts"])

scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="sample",
    continuous_covariate_keys=["log_counts"]
)

# Check cell type distribution per system
import pandas as pd

ct_dist = pd.crosstab(adata.obs["system"], adata.obs["cell_type"], normalize="index")
print(ct_dist)

# If very unbalanced, consider:
# 1. Subsample to balance
# 2. Use scANVI with labels to preserve rare types

def integrate_cross_system(
    adatas: dict,
    system_key: str = "system",
    batch_key: str = "batch",
    cell_type_key: str = "cell_type",
    n_top_genes: int = 4000,
    n_latent: int = 30
):
    """
    Integrate datasets from different technological systems.
    
    Parameters
    ----------
    adatas : dict
        Dictionary of {system_name: AnnData}
    system_key : str
        Key for system annotation
    batch_key : str
        Key for within-system batch
    cell_type_key : str
        Key for cell type labels (optional)
    n_top_genes : int
        Number of HVGs
    n_latent : int
        Latent dimensions
        
    Returns
    -------
    Integrated AnnData with model
    """
    import scvi
    import scanpy as sc
    
    # Add system labels and concatenate
    for system_name, adata in adatas.items():
        adata.obs[system_key] = system_name
    
    adata = sc.concat(list(adatas.values()))
    
    # Find common genes
    for name, ad in adatas.items():
        if name == list(adatas.keys())[0]:
            common_genes = set(ad.var_names)
        else:
            common_genes = common_genes.intersection(ad.var_names)
    
    adata = adata[:, list(common_genes)].copy()
    print(f"Common genes: {len(common_genes)}")
    
    # Store counts
    adata.layers["counts"] = adata.X.copy()
    
    # HVG selection
    sc.pp.highly_variable_genes(
        adata,
        n_top_genes=n_top_genes,
        flavor="seurat_v3",
        batch_key=system_key,
        layer="counts"
    )
    adata = adata[:, adata.var["highly_variable"]].copy()
    
    # Setup with system as covariate
    scvi.model.SCVI.setup_anndata(
        adata,
        layer="counts",
        batch_key=batch_key if batch_key in adata.obs else None,
        categorical_covariate_keys=[system_key]
    )
    
    # Train
    model = scvi.model.SCVI(adata, n_latent=n_latent, n_layers=2)
    model.train(max_epochs=300, early_stopping=True)
    
    # Get representation
    adata.obsm["X_integrated"] = model.get_latent_representation()
    
    # Clustering
    sc.pp.neighbors(adata, use_rep="X_integrated")
    sc.tl.umap(adata)
    sc.tl.leiden(adata)
    
    return adata, model

# Usage
adatas = {
    "10x_v3": sc.read_h5ad("10x_v3_data.h5ad"),
    "Smart-seq2": sc.read_h5ad("smartseq_data.h5ad"),
    "Drop-seq": sc.read_h5ad("dropseq_data.h5ad")
}

adata_integrated, model = integrate_cross_system(adatas)

# Visualize
sc.pl.umap(adata_integrated, color=["system", "leiden"])