【单细胞整合避坑指南】数据集不平衡如何毁掉你的分析？Nature Biotech 方法实战教程

scvi.model.SCVI.setup_anndata(adata, batch_key="batch")
model = scvi.model.SCVI(adata)
model.train()

# 基础环境
pip install scanpy scvi-tools harmony-pytorch bbknn

# 安装文章作者开发的平衡聚类指标库
pip install balanced-clustering

import scanpy as sc
import pandas as pd
import numpy as np
from scipy.spatial.distance import cosine
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
from balanced_clustering import balanced_adjusted_rand_index, balanced_adjusted_mutual_info

# 假设 adata.obs 中有 'batch' 和 'cell_type' 两列
for batch in sorted(adata.obs['batch'].unique()):
    print(f"\n=== {batch} ===")
    props = adata[adata.obs['batch']==batch].obs['cell_type'].value_counts(normalize=True)
    print(props.map(lambda x: f"{x:.1%}"))

def compute_cell_type_distances(adata, cell_type_key='cell_type', n_pcs=20):
    """计算各细胞类型在PCA空间中的中心距离"""
    sc.tl.pca(adata, n_comps=n_pcs, use_highly_variable=True)
    
    centers = {}
    for ct in adata.obs[cell_type_key].unique():
        mask = adata.obs[cell_type_key] == ct
        centers[ct] = adata.obsm['X_pca'][mask].mean(axis=0)
    
    # 计算每个细胞类型的最近邻
    results = []
    for ct1 in centers:
        min_dist = 1.0
        nearest = ""
        for ct2 in centers:
            if ct1 != ct2:
                # 使用余弦距离，与文章一致
                d = cosine(centers[ct1], centers[ct2])
                if d < min_dist:
                    min_dist = d
                    nearest = ct2
        results.append({
            'cell_type': ct1,
            'nearest_neighbor': nearest,
            'min_center_distance': min_dist,
            'total_cells': np.sum(adata.obs[cell_type_key] == ct1)
        })
    
    return pd.DataFrame(results).sort_values('min_center_distance')

# 运行诊断
diagnosis_df = compute_cell_type_distances(adata)
print(diagnosis_df)

      cell_type nearest_neighbor  min_center_distance  total_cells
0        pDC           cDC2               0.12            45
1    CD8_TCM        CD8_TEM               0.18           320
2       cDC2           pDC               0.12          2800

# Seurat (R语言示例，Python 可用 seurat-via-scanpy 或 rpy2)
# 减少锚点数量，防止强行拉平
FindIntegrationAnchors(..., k.anchor = 5)

# Harmony
adata = sc.external.pp.harmony_integrate(adata, key='batch', theta=1)  # theta 默认2，降低则更保守

# scVI
SCVI.setup_anndata(adata, batch_key="batch")
model = SCVI(adata, n_latent=20)  # 增加 latent dim，保留更多信号

# 示例：将免疫细胞和上皮细胞分开整合
# 先对每个批次做粗分群（或利用已知标记），分成 major compartments
for compartment in ['immune', 'epithelial', 'stromal']:
    sub = adata[adata.obs['compartment'] == compartment].copy()
    # 对每个 compartment 单独整合
    scvi.model.SCVI.setup_anndata(sub, batch_key='batch')
    model = scvi.model.SCVI(sub)
    model.train()
    sub.obsm['X_integrated'] = model.get_latent_representation()
    # 保存结果，最后合并

# Seurat v4/v5 的 MapQuery 思路
# reference: 大样本/对照样本
# query1, query2, query3: 其他样本
# 分别将 query 映射到 reference，而不是全部混在一起做整合

# Day0+Day3 先整合，结果再与 Day7 整合
integrated_d0_d3 = integrate(day0, day3)
final = integrate(integrated_d0_d3, day7)

# 假设你已经有了整合后的聚类标签 'leiden' 和真实标签 'cell_type'
bARI = balanced_adjusted_rand_index(adata.obs['cell_type'], adata.obs['leiden'])
bAMI = balanced_adjusted_mutual_info(adata.obs['cell_type'], adata.obs['leiden'])

print(f"Balanced ARI: {bARI:.3f}")
print(f"Balanced AMI: {bAMI:.3f}")

from sklearn.metrics import adjusted_rand_score
ARI = adjusted_rand_score(adata.obs['cell_type'], adata.obs['leiden'])
print(f"ARI: {ARI:.3f} vs bARI: {bARI:.3f}")

def evaluate_integration_per_cell_type(adata, embedding_key='X_pca', label_key='cell_type'):
    """在整合后的嵌入空间上评估每个细胞类型的分类性能"""
    X = adata.obsm[embedding_key]
    y = adata.obs[label_key]
    
    # 分层抽样，确保训练/测试集中各细胞类型比例一致
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, stratify=y, random_state=42
    )
    
    knn = KNeighborsClassifier(n_neighbors=15)
    knn.fit(X_train, y_train)
    y_pred = knn.predict(X_test)
    
    # 逐个细胞类型输出 F1
    results = []
    for ct in sorted(y.unique()):
        f1 = f1_score(y_test, y_pred, labels=[ct], average=None, zero_division=0)[0]
        results.append({'cell_type': ct, 'F1_score': f1})
    
    return pd.DataFrame(results).sort_values('F1_score')

# 对整合后的数据运行（如 X_scVI, X_harmony 等）
results_df = evaluate_integration_per_cell_type(adata, embedding_key='X_scVI')
print(results_df)

# 错误做法：直接 resolution=1
# sc.tl.leiden(adata, resolution=1)

# 正确做法：先用最平衡的批次/对照样本确定最优 resolution
# 使得聚类数 ≈ 真实细胞类型数
for res in [0.2, 0.5, 1.0, 1.5]:
    sc.tl.leiden(adata, resolution=res)
    n_clusters = adata.obs['leiden'].nunique()
    print(f"resolution={res}: {n_clusters} clusters")

# 不要直接在整个整合后的数据上跑 rank_genes_groups 找全局 marker
# 而是按批次分别找 marker，然后取交集，过滤掉批次特异性基因

markers_batch1 = find_markers(adata[adata.obs['batch']=='batch1'])
markers_batch2 = find_markers(adata[adata.obs['batch']=='batch2'])
consensus_markers = intersect(markers_batch1, markers_batch2)

# 检查关键中间态在各批次中的数量
for batch in adata.obs['batch'].unique():
    print(f"\n{batch}:")
    print(adata[adata.obs['batch']==batch].obs['cell_type'].value_counts())