从单细胞转录组数据中重构生长与动态轨迹.

0. TL; DR

本文介绍的moscot (multi-omics single-cell optimal transport) 是一个为单细胞基因组学设计的、可扩展的OT计算框架，它系统性地解决了上述所有问题。moscot支持跨所有应用的多模态数据，通过工程和算法创新（如low-rank近似）实现了对百万级细胞图谱的高效分析，并提供了一个统一、易用的API。

作者通过一系列应用展示了moscot的强大能力：

  1. 时间轨迹：成功重构了包含170万个细胞、跨越20个时间点的小鼠胚胎发育轨迹。
  2. 空间映射：将多模态单细胞数据（CITE-seq）映射到小鼠肝脏空间转录组上，丰富了空间数据的生物学信息；并成功对齐了多个小鼠大脑的冠状切片。
  3. 时空动力学：提出了一种全新的moscot.spatiotemporal方法，结合时间和空间信息，揭示了小鼠胚胎发生的时空动态。
  4. 新生物学发现：在一个全新的小鼠胰腺发育多组学数据集上，解析了delta细胞和epsilon细胞的谱系关系，并预测了NEUROD2是epsilon祖细胞的关键调控因子，这一发现通过人诱导多能干细胞的KO实验得到了验证。

1. 背景介绍

单细胞基因组学技术，如单细胞RNA测序（scRNA-seq）和空间转录组学，极大地增进了我们对细胞分化和组织结构的理解。然而，这些技术通常是破坏性的，我们只能获得细胞在某个时间点或某个空间位置的“快照”，而无法追踪同一个细胞的完整生命历程。因此，如何将这些离散的“快照”重新对齐，以恢复细胞的动态和空间背景，成为一个核心的计算挑战。

最优传输（Optimal Transport, OT）作为一个强大的数学框架，已被证明在解决这类映射问题中非常有效。例如，Waddington-OT (WOT)被用于描绘细胞重编程过程，novoSpaRc被用于将单细胞数据映射到空间位置，PASTE则用于对齐不同的空间转录组切片。

尽管OT方法潜力巨大，但其在单细胞领域的应用仍面临三个关键挑战：

  1. 单模态限制：大多数现有工具仅为单模态数据（如仅RNA表达）设计，无法整合利用日益普及的多模态数据（如RNA+ATAC或RNA+蛋白质）。
  2. 可扩展性瓶颈：标准OT算法的时间和内存复杂度随细胞数量呈二次方（甚至三次方）增长，这使其难以应用于包含数百万细胞的生物学图谱。
  3. 实现不统一：现有的OT工具基于不同的实现，API异构，导致用户难以在不同问题间调整或组合使用，也为开发者扩展新功能带来了障碍。

为了克服这些限制，作者开发了moscot，一个统一、可扩展且支持多模态的OT计算框架，旨在释放OT在时间、空间和时空单细胞分析中的全部潜力。

2. Moscot 框架

Moscot的核心是将生物学上的细胞映射和对齐问题，转化为数学上的最优传输（OT）问题，并使用一套一致且高效的算法进行求解。其基本流程是：输入非配对的数据集（如不同时间点或空间切片的数据），结合先验生物学知识（如细胞增殖率），求解一个OT问题，最终生成一个概率性的“耦合矩阵”，该矩阵描述了不同数据集中细胞之间的对应关系。

2.1 Moscot使用的三种最优传输模型

为了应对不同的生物学问题，moscot内置了三种OT模型，它们在如何关联细胞分布上有所不同：

Wasserstein-type (W-type) OT

这是最基础的OT形式，用于比较具有相同特征的两个细胞群（例如，两个时间点的scRNA-seq数据）。它旨在找到一个耦合矩阵 $P$ ，使得在将细胞从一个分布“运输”到另一个分布时，总的“运输成本”最小。在实际计算中，作者采用了熵正则化的OT问题：

\[P^* = \text{argmin}_{P \in U(a,b)} \langle P, C \rangle - \epsilon H(P)\]

其中$C$是成本矩阵，$C_{ij}$表示将细胞$i$移动到细胞$j$的成本（通常是它们在基因表达空间中的距离）。$U(a,b)$是所有可能的耦合矩阵的集合，这些矩阵需要满足边际约束，即行和与列和分别等于预设的边际分布$a$和$b$。$H(P)$是熵正则项，它使得问题更易于求解，并能产生更平滑、非稀疏的解。$\epsilon > 0$是正则化强度。该问题可以通过高效的Sinkhorn算法求解。

为了更贴近生物学现实，moscot还引入了两个重要扩展：

Gromov–Wasserstein-type (GW-type) OT

当两个细胞群的特征空间不同，但其内部结构相似时（例如，对齐两个仅有空间坐标的组织切片），GW-type OT非常有用。它不直接比较细胞特征，而是比较细胞间的距离关系矩阵，寻找能最好地保持这种内部几何结构的映射。

Fused Gromov–Wasserstein-type (FGW-type) OT

这是一种混合模型，结合了W-type和GW-type的优点。它适用于细胞群具有部分共享特征和各自内部结构的情况。例如，在将单细胞RNA-seq数据映射到空间转录组数据时，我们可以同时最小化共享基因表达的差异（W-type部分）和最大化基因表达空间结构与物理空间结构的对应关系（GW-type部分）。其目标函数形式如下：

\[P^* = \text{argmin}_{P \in U(a,b)} \alpha \sum L(C_X, C_Y)P \otimes P + (1-\alpha) \langle P, C \rangle - \epsilon H(P)\]

第一项是GW-type成本，衡量两个空间内部结构（由距离矩阵$C_X$和$C_Y$定义）的差异。第二项是W-type成本，衡量跨空间特征（由成本矩阵$C$定义）的差异。$\alpha \in [0, 1]$是平衡这两个成本的权重。

2.2 Moscot的三大创新

Moscot通过三大设计原则克服了现有OT工具的局限性：

  1. 多模态 (Multimodality)：moscot通过在计算成本/距离时使用多模态数据的联合表示来实现这一点。例如，对于RNA和ATAC数据，作者可以将其分别投影到各自的低维空间，然后将这两个空间拼接起来，在这个联合空间中计算细胞间的距离。这使得OT的求解过程能够同时利用来自多个模态的信息。
  2. 可扩展性 (Scalability)：作者通过工程和方法论两方面的创新来解决性能瓶颈。
    • 工程优化：moscot基于JAX框架下的OTT库，支持即时编译（just-in-time compilation）、GPU加速和成本矩阵的在线计算（on-the-fly evaluation）。在线计算避免了在内存中存储巨大的成本矩阵（$N \times M$），使内存复杂度从二次方降为线性，从而可以在GPU上处理更大的数据集。
    • 方法论创新：对于图谱级别的超大数据集，作者引入了对耦合矩阵$P$的low-rank约束。这一近似方法可以将W-type、GW-type和FGW-type OT问题的时间和内存复杂度都降低到线性级别，从而实现了对百万级细胞数据集的分析。
  3. 统一框架 (Unified Framework)：moscot为时间、空间和时空问题提供了一致的API，并与scverse生态系统（如AnnData和scanpy）无缝集成。这使得用户可以方便地在不同分析任务之间切换，并极大地简化了新方法的开发和应用。

3. 实验分析

作者通过四个复杂的生物学应用，全面展示了moscot的性能和多功能性。

实验一：重构图谱规模的小鼠胚胎发育轨迹

作者将moscot.time应用于一个包含近170万个细胞、横跨20个时间点的小鼠胚胎发育图谱。

实验二：空间样本的映射与对齐

作者展示了moscot在处理空间转录组学数据上的两个核心应用：空间映射和空间对齐。

实验三：描绘小鼠时空发育图谱

作者引入了moscot.spatiotemporal，一种结合了时间和空间信息的全新轨迹推断方法，并将其应用于小鼠胚胎发生的MOSTA时空图谱。

实验四：解析小鼠胰腺发育的谱系关系

作者利用moscot分析了一个新生成的胰腺发育多组学（snRNA + ATAC）数据集，以解析delta和epsilon细胞的复杂谱系。