scArches 参考映射 — 将新数据映射到已有单细胞参考图谱¶

# 创建专用 conda 环境
conda create -n scarches python=3.10 -y
conda activate scarches

# 安装 scArches（通过 scvi-tools，scArches 已集成其中）
pip install scvi-tools

# 验证安装
python -c "import scvi; print(scvi.__version__)"

# 安装可选依赖
pip install scanpy gdown   # gdown 用于从 Google Drive 下载参考模型

import scvi                            # scVI-tools 包含 scArches 功能
import scanpy as sc                    # Scanpy 数据处理
import anndata as ad

# ── 核心概念 ──────────────────────────────────────────
# 参考模型（Reference）：在大规模参考数据集上预训练好的 SCVI/SCANVI 模型
# 查询数据（Query）：新测序的样本，要映射到参考图谱上
# 映射（Surgery）：只微调解码器，保持编码器大致不变，保留参考知识

# ── 步骤1：准备参考模型（如果没有现成的） ─────────────
# 在参考数据集上训练 SCVI 模型
adata_ref = sc.read_h5ad('reference_atlas.h5ad')  # 读取参考图谱数据

# 设置参考模型
scvi.model.SCVI.setup_anndata(
    adata_ref,
    layer='counts',                    # 原始计数层
    batch_key='study'                  # 批次信息列
)

# 训练参考模型
arches_params = dict(
    use_layer_norm="both",            # 使用层归一化（scArches 必须设置）
    use_batch_norm="none",            # 关闭 BatchNorm（scArches 必须设置）
    encode_covariates=True,           # 编码协变量（批次等）
    dropout_rate=0.2,                 # Dropout 率
    n_layers=2                        # 网络层数
)

vae_ref = scvi.model.SCVI(adata_ref, **arches_params)
vae_ref.train(max_epochs=500, early_stopping=True)
vae_ref.save('reference_model/', overwrite=True)  # 保存参考模型

import scvi
import scanpy as sc
import numpy as np
import matplotlib.pyplot as plt

# ── 完整参考映射流程 ──────────────────────────────────

# 1. 读取查询数据（新测序样本）
adata_query = sc.read_h5ad('new_sample.h5ad')
print(f"查询数据：{adata_query.n_obs} 个细胞，{adata_query.n_vars} 个基因")

# 2. 确保查询数据的基因集与参考一致
# 加载参考模型（只需要模型文件，不需要参考数据的所有细胞）
vae_ref = scvi.model.SCVI.load('reference_model/')

# 获取参考模型使用的基因列表
ref_genes = vae_ref.adata.var_names.tolist()
print(f"参考模型使用 {len(ref_genes)} 个基因")

# 过滤查询数据到相同基因（只保留参考中有的基因）
query_genes = adata_query.var_names.tolist()
common_genes = list(set(ref_genes) & set(query_genes))
adata_query_filtered = adata_query[:, common_genes].copy()
print(f"共同基因数：{len(common_genes)}（查询/{len(query_genes)}，参考/{len(ref_genes)}）")

# 对于参考中有但查询中没有的基因，补零
# scArches 会自动处理这个问题

# 3. 将查询数据准备好用于参考模型
scvi.model.SCVI.prepare_query_anndata(
    adata_query,              # 查询数据（会自动对齐基因）
    vae_ref                   # 参考模型
)
print("查询数据基因对齐完成！")

# 4. 用 Surgery（手术）将查询数据映射到参考模型
vae_query = scvi.model.SCVI.load_query_data(
    adata_query,              # 查询数据
    vae_ref,                  # 参考模型
    freeze_dropout=True,      # 冻结 dropout（减少过拟合）
    inplace_subset_query_vars=True  # 自动过滤查询数据基因
)

# 5. 微调模型（只更新查询专属参数）
vae_query.train(
    max_epochs=200,           # 查询数据通常只需要 100-200 轮
    plan_kwargs={'weight_decay': 0.0},  # 不加权重衰减（保留参考知识）
    early_stopping=True
)

# 6. 获取查询数据的潜在嵌入（用于 UMAP 可视化）
adata_query.obsm['X_scVI'] = vae_query.get_latent_representation()

# 7. 合并参考和查询数据进行可视化
adata_ref_copy = vae_ref.adata.copy()
adata_ref_copy.obsm['X_scVI'] = vae_ref.get_latent_representation()
adata_ref_copy.obs['dataset'] = 'Reference'
adata_query.obs['dataset'] = 'Query'

# 拼接两个数据集（只保留共有列）
import anndata as ad
adata_full = ad.concat(
    [adata_ref_copy, adata_query],
    join='inner'              # 只保留两者共有的基因
)

# 计算 UMAP（只用 Query 的嵌入，参考的嵌入保持不变）
sc.pp.neighbors(adata_full, use_rep='X_scVI')
sc.tl.umap(adata_full)
sc.pl.umap(
    adata_full,
    color=['dataset', 'cell_type'],   # 查看查询数据是否嵌入到正确位置
    ncols=2,
    title=['数据集来源', '细胞类型']
)

# 8. 细胞类型迁移注释
# 用参考数据的已知细胞类型标注来推断查询数据的细胞类型
# 方法：找每个查询细胞在参考中的最近邻，用 KNN 投票

from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelEncoder

# 准备参考数据的标签和嵌入
ref_embeddings = adata_ref_copy.obsm['X_scVI']
ref_labels = adata_ref_copy.obs['cell_type'].values

# 训练 KNN 分类器（k=15）
knn = KNeighborsClassifier(n_neighbors=15, metric='euclidean')
knn.fit(ref_embeddings, ref_labels)

# 预测查询数据的细胞类型
query_embeddings = adata_query.obsm['X_scVI']
predicted_labels = knn.predict(query_embeddings)
prediction_proba = knn.predict_proba(query_embeddings)

adata_query.obs['predicted_cell_type'] = predicted_labels
adata_query.obs['prediction_confidence'] = prediction_proba.max(axis=1)

# 可视化预测结果
sc.pl.umap(
    adata_query,
    color=['predicted_cell_type', 'prediction_confidence'],
    title=['预测细胞类型', '预测置信度']
)

# 9. 保存结果
adata_query.write_h5ad('query_mapped.h5ad')
print("参考映射完成！")
print(f"预测细胞类型分布：\n{adata_query.obs['predicted_cell_type'].value_counts()}")

# 安装
# pip install scvi-tools scanpy

# 参考映射三步走
scvi.model.SCVI.prepare_query_anndata(query, ref_model)
vae_q = scvi.model.SCVI.load_query_data(query, ref_model)
vae_q.train(max_epochs=200)
query.obsm['X_scVI'] = vae_q.get_latent_representation()

# 常用参考图谱（可直接下载预训练模型）
# 人类细胞图谱 HCA：https://cellxgene.cziscience.com/
# Tabula Sapiens：https://tabula-sapiens-portal.ds.czbiohub.org/