STitch3D工程实践：从零构建空间转录组3D重建的深度学习流水线

最近在复现前沿的空间组学算法时，我常常感到一种割裂：论文里的数学公式优雅而抽象，但真要把代码跑起来，中间隔着一道名为“工程实现”的鸿沟。特别是像STitch3D这类结合了图神经网络、多模态数据整合的模型，理论理解是一回事，用PyTorch把它稳定、高效地实现出来，又是另一回事。这篇文章，就是想把我在复现STitch3D过程中的那些“坑”和“钥匙”分享出来，目标读者是那些不仅想读懂论文，更想亲手把算法“造”出来的生物信息工程师和算法开发者。我们会抛开纯理论的叙述，聚焦于数据如何流动、张量如何变换、以及损失函数如何被PyTorch优雅地表达。最终，你会获得一套可以直接在Colab上运行、并可被嵌入到你自己的分析流程中的模块化代码。

1. 环境搭建与数据预处理实战

在开始构建模型之前，一个可复现、隔离的环境是高效开发的基础。我强烈推荐使用Conda来管理依赖，它能很好地处理Python版本与科学计算包之间复杂的依赖关系。

# 创建并激活一个专门的conda环境
conda create -n stitch3d_env python=3.9 -y
conda activate stitch3d_env

# 安装核心依赖
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
pip install scanpy anndata squidpy pytorch-lightning
pip install ogb  # 用于图神经网络的一些工具
pip install paste-lab  # 用于切片对齐的PASTE算法

注意：PyTorch的版本需要与你的CUDA驱动版本匹配。上述命令适用于CUDA 11.7。如果你的环境不同，请访问PyTorch官网获取对应的安装命令。

数据预处理是模型成功的基石，对于STitch3D而言，这一步尤为关键，因为它涉及多切片对齐和3D图的构建。我们假设你已经有了多个空间转录组切片（例如10x Visium数据）的anndata对象和一个单细胞RNA-seq参考数据的anndata对象。

第一步：切片对齐与坐标初始化 STitch3D论文提到了ICP或PASTE算法。这里我们以PASTE为例，因为它专为空间转录组切片对齐设计。对齐的目标是为每个切片找到一个最优的刚性变换（旋转和平移），使它们在二维平面上对齐。

import paste as pst
import numpy as np

# 假设 slice_adata_list 是一个包含多个切片anndata对象的列表
# 每个anndata的.obsm['spatial'] 存储了原始的二维坐标
slice_adata_list = [slice1, slice2, slice3]

# 使用PASTE进行成对对齐，这里以顺序对齐为例
aligned_slices = []
aligned_slices.append(slice_adata_list[0]) # 第一片作为参考

for i in range(1, len(slice_adata_list)):
    # 执行PASTE对齐，pi是耦合矩阵，用于后续分析
    pi, aligned_coords = pst.pairwise_align(slice_adata_list[i-1], slice_adata_list[i])
    # 将当前切片的坐标更新为对齐后的坐标
    slice_adata_list[i].obsm['spatial_aligned'] = aligned_coords[1] # 取第二个切片对齐后的坐标
    aligned_slices.append(slice_adata_list[i])

第二步：构建3D坐标与邻接图 对齐后，我们需要赋予每个spot一个第三维（z轴）坐标，通常根据切片的物理顺序或已知的层间距来设定。然后，基于3D欧氏距离构建k近邻图或半径邻接图。

import torch
from torch_geometric.data import Data
from sklearn.neighbors import NearestNeighbors

def construct_3d_graph(aligned_slices, z_coords, radius_factor=1.1):
    """
    构建3D邻接图。
    aligned_slices: 对齐后的切片列表
    z_coords: 每个切片对应的z轴坐标列表
    radius_factor: 判定为邻居的距离阈值因子
    """
    all_coords_3d = []
    all_expressions = []

    for idx, adata in enumerate(aligned_slices):
        coords_2d = adata.obsm['spatial_aligned']
        # 为当前切片所有spot添加z坐标
        z = np.full((coords_2d.shape[0], 1), z_coords[idx])
        coords_3d = np.hstack([coords_2d, z])
        all_coords_3d.append(coords_3d)
        all_expressions.append(adata.X.toarray() if hasattr(adata.X, 'toarray') else adata.X)

    # 合并所有数据
    X_3d = np.vstack(all_coords_3d)  # (N_total, 3)
    Y = np.vstack(all_expressions)    # (N_total, G)

    # 计算距离阈值：通常取同一切片内最近邻距离的中位数乘以一个因子
    intra_distances = []
    start_idx = 0
    for coords in all_coords_3d:
        n_spots = coords.shape[0]
        if n_spots > 1:
            nbrs = NearestNeighbors(n_neighbors=2).fit(coords)
            distances, _ = nbrs.kneighbors(coords)
            intra_distances.extend(distances[:, 1]) # 到最近邻居的距离
        start_idx += n_spots
    radius = np.median(intra_distances) * radius_factor

    # 基于半径构建全局邻接关系
    nbrs = NearestNeighbors(radius=radius).fit(X_3d)
    distances, indices = nbrs.radius_neighbors(X_3d, return_distance=True)

    # 构建PyTorch Geometric需要的边索引 (2, num_edges)
    edge_index = []
    for i, neighbors in enumerate(indices):
        for j in neighbors:
            if i != j:  # 排除自环
                edge_index.append([i, j])
    edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()

    # 准备节点特征：这里先用基因表达数据，后续模型中会编码
    x = torch.tensor(Y, dtype=torch.float32)

    # 构建图数据对象
    graph_data = Data(x=x, edge_index=edge_index, pos=torch.tensor(X_3d, dtype=torch.float32))
    return graph_data, Y, X_3d

这个函数返回的graph_data对象可以直接喂给PyTorch Geometric的图神经网络。预处理阶段的核心思想是：将生物学的空间关系，转化为图数据结构中节点（spots）和边（邻接关系）的数学表达。

2. 核心模型：图注意力网络与细胞类型解卷积的实现

STitch3D模型的核心是一个图注意力网络（GAT），用于学习spot的低维隐表示（Z），以及一个多层感知机（MLP），用于从Z推断细胞类型比例。我们分模块构建。

2.1 图注意力编码器 我们实现一个多层GAT。每一层GAT都会聚合邻居节点的信息，并通过注意力机制赋予不同邻居不同的权重。

import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv

class GATEncoder(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, heads=4, dropout=0.2):
        super().__init__()
        self.conv1 = GATConv(in_channels, hidden_channels, heads=heads, dropout=dropout, concat=True)
        # 注意：第一层多头输出维度是 hidden_channels * heads
        self.conv2 = GATConv(hidden_channels * heads, out_channels, heads=1, dropout=dropout, concat=False)
        self.dropout = dropout

    def forward(self, x, edge_index):
        # x: (N, in_channels), edge_index: (2, E)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = F.elu(self.conv1(x, edge_index)) # 使用ELU激活函数
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv2(x, edge_index) # 最后一层不接激活函数，输出隐表示Z
        return x

2.2 细胞类型比例预测器 这个模块将隐表示Z映射到细胞类型比例（一个概率分布）。我们需要保证输出的每个spot的细胞类型比例之和为1，因此使用Softmax函数。

class CellTypeDeconvolution(nn.Module):
    def __init__(self, latent_dim, num_cell_types):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_cell_types)
        )

    def forward(self, z):
        # z: (N, latent_dim)
        logits = self.mlp(z)
        proportions = F.softmax(logits, dim=-1) # (N, num_cell_types)，每行和为1
        return proportions

2.3 批次效应校正模块 这是STitch3D的精华之一。它引入了两个可学习的偏差项：d_si（spot-level bias）和 b_sg（gene-level bias），来校正技术噪音。

class BatchEffectCorrector(nn.Module):
    def __init__(self, num_slices, num_genes, latent_dim, hidden_dim=64):
        super().__init__()
        self.num_slices = num_slices
        self.num_genes = num_genes

        # d_si: 每个切片-每个spot的偏差，用一个网络来生成，输入是z和切片标签
        self.d_net = nn.Sequential(
            nn.Linear(latent_dim + num_slices, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)  # 输出一个标量偏差
        )

        # b_sg: 每个切片-每个基因的偏差，作为直接可训练参数
        self.b_sg = nn.Parameter(torch.zeros(num_slices, num_genes))

    def forward(self, z, slice_labels, proportions, cell_type_profile):
        """
        z: (N, latent_dim)
        slice_labels: (N, ) 每个spot属于哪个切片（整数索引）
        proportions: (N, C) 细胞类型比例
        cell_type_profile: (C, G) 参考细胞类型的基因表达谱
        """
        N = z.size(0)
        # 将切片标签转为one-hot编码，用于与z拼接
        slice_onehot = F.one_hot(slice_labels, num_classes=self.num_slices).float() # (N, S)
        d_input = torch.cat([z, slice_onehot], dim=1) # (N, latent_dim+S)
        d_si = self.d_net(d_input).squeeze(-1)  # (N,)

        # 获取每个spot对应的切片-基因偏差
        b_sg_per_spot = self.b_sg[slice_labels]  # (N, G)

        # 重建基因表达： sum_c (比例_c * 表达谱_c) * exp(d_si + b_sg)
        # cell_type_profile: (C, G), proportions: (N, C)
        expected_expression = torch.matmul(proportions, cell_type_profile) # (N, G)
        # 加入偏差项，使用指数确保正值
        scale_factors = torch.exp(d_si.unsqueeze(-1) + b_sg_per_spot) # (N, G)
        reconstructed_expression = expected_expression * scale_factors

        # 还需要根据spot的总计数进行归一化（模拟原始数据的文库大小差异）
        # 这里简化处理，假设总计数信息已包含在预处理中或作为额外输入
        return reconstructed_expression, d_si, b_sg_per_spot

将以上三个核心模块组装起来，就构成了STitch3D的主干网络。

3. 损失函数设计与模型训练策略

STitch3D的损失函数是一个多任务组合，主要包含重建损失和正则化损失。理解每一部分的物理意义，对于调试模型至关重要。

3.1 重建损失 目标是让模型预测的基因表达（经过批次效应校正后）尽可能接近观测到的基因表达。由于基因表达数据通常是过离散的计数数据，我们采用负二项分布（Negative Binomial）损失，这比简单的均方误差（MSE）更符合生物学数据的特性。

import torch.distributions as dist

def negative_binomial_loss(observed, reconstructed, theta=10.0):
    """
    observed: 观测到的基因表达计数 (N, G)
    reconstructed: 模型重建的基因表达均值 (N, G)
    theta: 负二项分布的逆离散参数，越大分布越接近泊松
    """
    # 防止除零或log(0)
    epsilon = 1e-8
    reconstructed = reconstructed + epsilon
    # 负二项分布的参数化：均值mu = reconstructed，方差 = mu + mu^2/theta
    nb_dist = dist.NegativeBinomial(total_count=theta, logits=(reconstructed/theta).log() - torch.log1p(reconstructed/theta))
    # 计算负对数似然
    nll = -nb_dist.log_prob(observed).sum(dim=-1).mean()
    return nll

3.2 正则化损失 为了防止模型过度校正，抹除切片间真实的生物学差异，STitch3D引入了一个基于注意力的正则化器。其思想是：用一个切片特定的GAT去尝试从隐表示Z重建批次效应偏差d_si，如果d_si中包含了可被空间结构预测的生物学信息，那么这个重建就会很容易，损失就小。 我们鼓励d_si中只包含技术噪音，而非生物学信号，因此我们希望这个重建任务变得困难，即最大化这个重建损失（或最小化其负值作为正则项）。

class SliceSpecificGAT(nn.Module):
    """一个轻量级的GAT，用于正则化项"""
    def __init__(self, latent_dim, hidden_dim=32):
        super().__init__()
        self.conv1 = GATConv(latent_dim, hidden_dim, heads=2, concat=True)
        self.conv2 = GATConv(hidden_dim*2, 1, heads=1, concat=False) # 输出每个spot的d_si预测值

    def forward(self, z, edge_index):
        x = F.elu(self.conv1(z, edge_index))
        x = self.conv2(x, edge_index)
        return x.squeeze(-1)

def regularization_loss(z, d_si, edge_index, slice_labels):
    """
    计算正则化损失。
    对每个切片单独计算：用该切片内spot的z预测d_si，计算MSE，然后求和。
    """
    reg_loss = 0.0
    gat = SliceSpecificGAT(z.size(1))
    # 注意：在实际训练中，这个SliceSpecificGAT的参数也需要优化
    for s in torch.unique(slice_labels):
        mask = (slice_labels == s)
        z_s = z[mask]
        d_si_s = d_si[mask]
        # 需要获取该切片内部的子图edge_index
        # 这里简化处理，假设我们传递的是全局图，需要过滤出属于切片s的节点构成的边
        # 实际实现中需要更复杂的子图提取逻辑，此处省略细节
        # edge_index_s = subgraph(mask, edge_index)[0]
        # pred_d = gat(z_s, edge_index_s)
        # reg_loss += F.mse_loss(pred_d, d_si_s)
    # 简化返回一个占位值，实际训练需完善
    return torch.tensor(0.0, requires_grad=True)

3.3 总损失与训练循环 将上述损失组合，并加入权重超参数。

def total_loss(observed, reconstructed, z, d_si, edge_index, slice_labels, lambda_reg=0.1):
    loss_nb = negative_binomial_loss(observed, reconstructed)
    loss_reg = regularization_loss(z, d_si, edge_index, slice_labels)
    total = loss_nb + lambda_reg * loss_reg
    return total, loss_nb, loss_reg

训练时，我们使用PyTorch Lightning来组织代码，它能让训练循环更简洁，并方便集成早停、学习率调度等功能。

import pytorch_lightning as pl
from torch.optim import AdamW

class STitch3DLight(pl.LightningModule):
    def __init__(self, num_genes, num_cell_types, num_slices, latent_dim=64):
        super().__init__()
        self.save_hyperparameters()
        self.encoder = GATEncoder(num_genes, 256, latent_dim)
        self.deconv = CellTypeDeconvolution(latent_dim, num_cell_types)
        self.corrector = BatchEffectCorrector(num_slices, num_genes, latent_dim)
        # 假设 cell_type_profile 是已知的参考矩阵，作为模型的一个缓冲区
        self.register_buffer('cell_type_profile', torch.randn(num_cell_types, num_genes))

    def forward(self, data):
        x, edge_index, slice_labels = data.x, data.edge_index, data.slice_labels
        z = self.encoder(x, edge_index)
        proportions = self.deconv(z)
        recon_exp, d_si, b_sg = self.corrector(z, slice_labels, proportions, self.cell_type_profile)
        return z, proportions, recon_exp, d_si

    def training_step(self, batch, batch_idx):
        z, proportions, recon_exp, d_si = self(batch)
        total_loss, loss_nb, loss_reg = total_loss(
            batch.x, recon_exp, z, d_si, batch.edge_index, batch.slice_labels, lambda_reg=0.1
        )
        self.log('train_total_loss', total_loss)
        self.log('train_nb_loss', loss_nb)
        self.log('train_reg_loss', loss_reg)
        return total_loss

    def configure_optimizers(self):
        optimizer = AdamW(self.parameters(), lr=1e-3, weight_decay=1e-5)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5)
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'monitor': 'train_total_loss',
                'interval': 'epoch',
                'frequency': 1
            }
        }

# 初始化训练器并开始训练
trainer = pl.Trainer(max_epochs=100, accelerator='gpu', devices=1)
model = STitch3DLight(num_genes=2000, num_cell_types=8, num_slices=3)
trainer.fit(model, train_dataloader)

训练的关键在于监控三个损失的下降曲线。理想情况下，train_nb_loss和train_total_loss应稳步下降，而train_reg_loss可能保持在一个相对稳定的非零值，这表明正则化项在起作用。

4. 结果解析、可视化与下游分析

模型训练完成后，我们得到了每个spot的低维隐表示Z和细胞类型比例。这些结果是下游生物学分析的起点。

4.1 隐空间可视化 使用UMAP或t-SNE将高维的Z降维到2D进行可视化，可以观察spots在隐空间中的聚类情况，这通常对应不同的空间功能域。

import umap
import matplotlib.pyplot as plt

# 提取隐表示
model.eval()
with torch.no_grad():
    z, proportions, _, _ = model(graph_data) # graph_data包含所有信息

z_np = z.cpu().numpy()

# 使用UMAP降维
reducer = umap.UMAP(n_components=2, random_state=42)
z_2d = reducer.fit_transform(z_np)

# 根据切片来源着色
plt.figure(figsize=(10, 8))
scatter = plt.scatter(z_2d[:, 0], z_2d[:, 1], c=graph_data.slice_labels, cmap='tab20', s=5, alpha=0.7)
plt.colorbar(scatter, label='Slice Index')
plt.title('UMAP of Latent Representation Z (colored by slice)')
plt.xlabel('UMAP1')
plt.ylabel('UMAP2')
plt.show()

如果批次效应校正有效，不同切片的spots在隐空间中应该混合在一起，而不是按切片形成明显的独立簇。

4.2 细胞类型比例空间映射 这是STitch3D最具价值的结果之一。我们可以将每个spot的主导细胞类型或其比例映射回原始的3D空间坐标上。

# 获取每个spot的主要细胞类型
dominant_cell_type = torch.argmax(proportions, dim=1).cpu().numpy() # (N,)

# 假设我们有原始的3D坐标 X_3d (N, 3)
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')
scatter = ax.scatter(X_3d[:, 0], X_3d[:, 1], X_3d[:, 2],
                     c=dominant_cell_type, cmap='Set3', s=20, alpha=0.8)
plt.colorbar(scatter, ax=ax, label='Dominant Cell Type')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z (Slice Depth)')
ax.set_title('3D Spatial Distribution of Dominant Cell Types')
plt.show()

4.3 与下游分析工具衔接 得到的proportions可以作为一个新的anndata对象的obsm层，方便使用Scanpy等标准单细胞分析流程进行后续分析，如差异比例分析、轨迹推断等。

import scanpy as sc

# 创建一个包含隐表示和细胞类型比例的anndata对象
combined_adata = sc.AnnData(X=graph_data.x.numpy()) # 原始表达矩阵
combined_adata.obsm['X_stitch3d'] = z_np          # 隐表示
combined_adata.obsm['cell_type_proportions'] = proportions.cpu().numpy()
combined_adata.obsm['spatial_3d'] = X_3d          # 3D坐标

# 使用隐表示进行Leiden聚类
sc.pp.neighbors(combined_adata, use_rep='X_stitch3d')
sc.tl.leiden(combined_adata, resolution=0.5)
# 现在 combined_adata.obs['leiden'] 包含了基于隐表示的空间域聚类结果

在整个实现过程中，最耗费时间的往往不是模型本身，而是数据的准备、图的构建以及损失函数的调试。我建议在跑完整数据之前，先用一个极小的模拟数据集验证整个流水线是否通畅，从数据加载、图构建、模型前向传播、损失计算到反向传播，确保每一步都没有维度错误或数值问题。当隐空间可视化显示出有意义的结构，并且重建损失能够收敛时，你会获得将一篇复杂论文转化为可运行代码的巨大满足感。