Skip to content

设计 ​

Rove 是一个 Rust crate。它为一个离散性状矩阵,在 Mk 模型下计算有根二叉树的似然。下面的决定是已经在代码中可见的那些。它们所铺垫的 SIMD 和快速路径工作在路线图上。

位点模式主序布局 ​

似然内核把时间花在把转移矩阵乘进部分似然上。BEAGLE 的 CPU 内核沿状态向量化,并针对四个状态做了特化,这适合核苷酸。形态学矩阵大多是二态和三态性状,所以这些内核会走带填充的通用路径。

Rove 反过来存放部分似然。对每个速率类别和每个状态,有一段连续的值,每个位点模式一个,填充到八个通道的倍数。内层循环沿模式运行,因此两个状态和六十一个状态的循环形状相同。

Layout 描述这种安排:状态数、速率类别数、模式数,以及填充后的步长。offset(category, state) 给出一段连续值的起点。

面向模型的快速路径 ​

Mk 有闭式解。一条分枝的转移矩阵是 P(t) = e·I + (1 − e)·π,其中 e 只取决于枝长,π 存放状态频率。把它乘进一个向量,每个模式只需 O(k) 而不是 O(k²)。

今天 model::Mk 生成完整矩阵,内核以通用方式做乘法。闭式解已经与 Lewis 的公式做过对照测试。在路线图上,基线测量和 Mkv 确认偏差校正排在直接 Mk 路径之前。

始终保留标量参考实现 ​

reference::Problem 是 Felsenstein 剪枝算法的朴素递归实现。它故意又慢又简单。等价性测试构建随机树,并把 LikelihoodCore 与它比较:用末端节点状态和用末端节点部分似然,有重缩放和无重缩放,相对容差为 1e-9。

以后加入的每条优化路径都必须通过同样的测试。参考实现会一直保留。

双缓冲与重缩放 ​

每个节点拥有两个部分似然缓冲区和两个转移矩阵缓冲区。在一次提议覆盖某个节点之前,set_node_partials_for_update 或 set_node_matrix_for_update 会把该节点切换到另一个缓冲区,这样先前的值得以保留。store 记录哪些缓冲区是当前的。restore 通过交换把那一组带回来,unstore 则通过复制。接受一次移动不花任何代价。这正是 MCMC 采样器要低成本地拒绝一次移动所需要的。

使用 Scaling::Always 时,每个内部节点的部分似然在计算完成后会除以其逐模式的最大值,并按节点保留该最大值的对数。对数似然再把这些缩放因子加回来。不做缩放,深的树会下溢。

NEON 优先,AVX2 其次 ​

布局的存在是为了让内层循环可以向量化。计划是先在 Apple silicon 上实现 NEON,再实现 AVX2,标量代码作为参考和回退保留。两者都还不存在;当前内核是标量 Rust。

单个 crate ​

Rove 不依赖 BEAST 或 BEAGLE。BEAST 是验证和基准测试的参考。BEAGLE C ABI 和通过 JNI 的 BEAST 2 LikelihoodCore 适配器是可能的兼容路径。两者都不是 Rust 核心所必需的,现在也都没有在构建。

非目标 ​

在大型核苷酸比对上与 BEAGLE 的 CUDA 后端竞争。那种情况请用 GPU。

用 Rust 实现的系统发育推断。MIT 或 Apache-2.0 许可。