设计
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。
