Replace Sankoff cost matrix with cardinality-composition decomposition
Replaced the legacy Sankoff parsimony pipeline with a new cardinality-composition decomposition that constructs row-normalized transition probability matrices symmetrized via geometric mean. This ensures reversibility, reduces free parameters from 240 to 120, and guarantees a zero diagonal. Tallies are now explicitly restricted to variable families to align with +ASC-corrected alignment populations. Additionally, fixed `--exclude-genome` handling to re-scan surviving sequences and drop newly monomorphic columns, preventing silent data corruption in downstream tree inference tools.
This commit is contained in:
Eric Coissac
@@ -0,0 +1,242 @@
|
||||
//! Cardinality × composition decomposition of the 16-state transition cost
|
||||
//! matrix — replaces `sankoff::build_cost_matrix`'s elementary-edge graph
|
||||
//! and its Floyd-Warshall shortest-path closure, which double-counts
|
||||
//! multi-hop transitions once IQ-TREE's own matrix exponential composes
|
||||
//! them again. See `docmd/theory/evolutionary_distances.md`, "`R` via
|
||||
//! `exp(-cost)` is wrong for a CTMC; cardinality/composition decomposition".
|
||||
//!
|
||||
//! Two small, directly-estimated first-order Markov models (diagonals
|
||||
//! included — "stay the same" is a real, calibrated outcome, not implicit):
|
||||
//! a 5-state cardinality model (`0..=4` members of a family) and a 4-state
|
||||
//! base-composition model (`A,C,G,T`, unambiguous/cardinality-1 loci only —
|
||||
//! composition bias is a substitution phenomenon, only meaningful when
|
||||
//! cardinality is conserved). Composed per pair of 16 states via a
|
||||
//! parsimony-style pairing of the non-shared elements (see
|
||||
//! `pairwise_transition_matrix`'s doc), then symmetrised — Sankoff
|
||||
//! parsimony's score is only independent of root placement (as required,
|
||||
//! since TNT/PhyG search *unrooted* trees) if the cost matrix is
|
||||
//! symmetric, the discrete-parsimony analogue of CTMC reversibility.
|
||||
|
||||
use crate::{BasePairTally, CardinalityTally};
|
||||
|
||||
/// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`),
|
||||
/// diagonal included ("stay at the same cardinality"), from
|
||||
/// [`CardinalityTally`]'s pooled, symmetric co-occurrence counts.
|
||||
pub fn cardinality_transition_probs(tally: &CardinalityTally) -> [[f64; 5]; 5] {
|
||||
let mut p = [[0.0f64; 5]; 5];
|
||||
for a in 0..5 {
|
||||
let row_sum: u64 = tally.counts[a].iter().sum();
|
||||
if row_sum == 0 {
|
||||
continue;
|
||||
}
|
||||
for b in 0..5 {
|
||||
p[a][b] = tally.counts[a][b] as f64 / row_sum as f64;
|
||||
}
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
/// Row-stochastic 4×4 base-composition transition probabilities
|
||||
/// (`A,C,G,T`), diagonal included ("stay the same base"), from
|
||||
/// [`BasePairTally`]'s pooled substitution (off-diagonal) and agreement
|
||||
/// (`same`, diagonal) counts — unambiguous, cardinality-1 loci only.
|
||||
pub fn composition_transition_probs(tally: &BasePairTally) -> [[f64; 4]; 4] {
|
||||
let mut p = [[0.0f64; 4]; 4];
|
||||
for a in 0..4 {
|
||||
let row_sum = tally.same[a] + (0..4).map(|b| tally.counts[a][b]).sum::<u64>();
|
||||
if row_sum == 0 {
|
||||
continue;
|
||||
}
|
||||
p[a][a] = tally.same[a] as f64 / row_sum as f64;
|
||||
for b in 0..4 {
|
||||
if a != b {
|
||||
p[a][b] = tally.counts[a][b] as f64 / row_sum as f64;
|
||||
}
|
||||
}
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
/// Enumerate every injective mapping from `small` (indices `0..small.len()`)
|
||||
/// into `large` (indices `0..large.len()`, `large.len() >= small.len()`),
|
||||
/// calling `visit` with each full permutation of chosen `large` indices.
|
||||
/// `small.len() <= 4` always here (at most 4 bases can be lost or gained),
|
||||
/// so a plain recursive enumeration (at most `4! = 24` calls) is simplest —
|
||||
/// not worth a dependency for.
|
||||
fn for_each_injection(small_len: usize, large_len: usize, visit: &mut dyn FnMut(&[usize])) {
|
||||
let mut chosen = vec![0usize; small_len];
|
||||
let mut used = vec![false; large_len];
|
||||
fn go(pos: usize, chosen: &mut [usize], used: &mut [bool], visit: &mut dyn FnMut(&[usize])) {
|
||||
if pos == chosen.len() {
|
||||
visit(chosen);
|
||||
return;
|
||||
}
|
||||
for l in 0..used.len() {
|
||||
if used[l] {
|
||||
continue;
|
||||
}
|
||||
used[l] = true;
|
||||
chosen[pos] = l;
|
||||
go(pos + 1, chosen, used, visit);
|
||||
used[l] = false;
|
||||
}
|
||||
}
|
||||
if small_len == 0 {
|
||||
visit(&[]);
|
||||
return;
|
||||
}
|
||||
go(0, &mut chosen, &mut used, visit);
|
||||
}
|
||||
|
||||
/// Parsimony-style pairing cost for the non-shared elements of a
|
||||
/// transition: pair every element of the smaller of `lost`/`gained` with
|
||||
/// some element of the larger one, choosing the pairing that minimises
|
||||
/// total `-ln(P_composition)` (maximises likelihood) — the discrete
|
||||
/// analogue of "prefer a substitution over an independent loss+gain".
|
||||
/// Leftover elements of the larger side (`|large|-|small|` of them) are
|
||||
/// *not* charged here — they're pure cardinality change, already priced by
|
||||
/// `P_cardinality(|A|→|B|)` in the caller.
|
||||
fn best_pairing_cost(lost: &[u8], gained: &[u8], p_comp: &[[f64; 4]; 4]) -> f64 {
|
||||
let (small, large) = if lost.len() <= gained.len() { (lost, gained) } else { (gained, lost) };
|
||||
if small.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let mut best = f64::INFINITY;
|
||||
for_each_injection(small.len(), large.len(), &mut |perm| {
|
||||
let mut cost = 0.0;
|
||||
for (i, &li) in perm.iter().enumerate() {
|
||||
let (x, y) = (small[i], large[li]);
|
||||
let p = p_comp[x as usize][y as usize];
|
||||
cost += if p > 0.0 { -p.ln() } else { f64::INFINITY };
|
||||
}
|
||||
if cost < best {
|
||||
best = cost;
|
||||
}
|
||||
});
|
||||
best
|
||||
}
|
||||
|
||||
/// The full 16-state transition cost matrix, replacing
|
||||
/// `sankoff::build_cost_matrix`. For every pair of states `A`, `B`
|
||||
/// (bitmasks, bit `0..4` = `A,C,G,T`): `shared = A∩B` contributes
|
||||
/// `P_composition(x→x)` per shared base (diagonal — "stayed"); the
|
||||
/// remaining `lost = A\B`, `gained = B\A` are parsimony-paired into
|
||||
/// substitutions via [`best_pairing_cost`], and whatever's left over after
|
||||
/// pairing (only possible on one side, since one side is fully consumed)
|
||||
/// is *pure* cardinality change, priced once via `P_cardinality(|A|→|B|)`
|
||||
/// — never chained through intermediate states. Row-normalised (`P(A→·)`
|
||||
/// sums to `1` over all `B`), converted to a cost via `-ln`, then
|
||||
/// symmetrised (`(cost(A,B)+cost(B,A))/2` — the row-normalised `P` is not
|
||||
/// symmetric in general, but a Sankoff parsimony cost must be, so the
|
||||
/// score is independent of where an unrooted tree gets rooted).
|
||||
pub fn pairwise_cost_matrix(p_card: &[[f64; 5]; 5], p_comp: &[[f64; 4]; 4]) -> [[f64; 16]; 16] {
|
||||
let mut raw = [[0.0f64; 16]; 16];
|
||||
for a in 0u8..16 {
|
||||
for b in 0u8..16 {
|
||||
let shared = a & b;
|
||||
let lost: Vec<u8> = (0..4).filter(|&i| a & (1 << i) != 0 && b & (1 << i) == 0).collect();
|
||||
let gained: Vec<u8> = (0..4).filter(|&i| b & (1 << i) != 0 && a & (1 << i) == 0).collect();
|
||||
|
||||
let mut log_p = 0.0; // accumulate ln(P), so 0.0 = probability 1
|
||||
let card_a = a.count_ones() as usize;
|
||||
let card_b = b.count_ones() as usize;
|
||||
let p_c = p_card[card_a][card_b];
|
||||
log_p += if p_c > 0.0 { p_c.ln() } else { f64::NEG_INFINITY };
|
||||
|
||||
for i in 0..4u8 {
|
||||
if shared & (1 << i) != 0 {
|
||||
let p = p_comp[i as usize][i as usize];
|
||||
log_p += if p > 0.0 { p.ln() } else { f64::NEG_INFINITY };
|
||||
}
|
||||
}
|
||||
log_p -= best_pairing_cost(&lost, &gained, p_comp);
|
||||
|
||||
raw[a as usize][b as usize] = log_p.exp();
|
||||
}
|
||||
}
|
||||
|
||||
// Row-normalise to a proper transition probability matrix.
|
||||
let mut p = [[0.0f64; 16]; 16];
|
||||
for a in 0..16 {
|
||||
let row_sum: f64 = raw[a].iter().sum();
|
||||
if row_sum > 0.0 {
|
||||
for b in 0..16 {
|
||||
p[a][b] = raw[a][b] / row_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cost = -ln(P), then symmetrise (see doc comment: parsimony on an
|
||||
// unrooted tree requires a symmetric cost matrix).
|
||||
let mut cost = [[0.0f64; 16]; 16];
|
||||
for a in 0..16 {
|
||||
for b in 0..16 {
|
||||
cost[a][b] = if p[a][b] > 0.0 { -p[a][b].ln() } else { f64::INFINITY };
|
||||
}
|
||||
}
|
||||
let mut sym = [[0.0f64; 16]; 16];
|
||||
for a in 0..16 {
|
||||
for b in 0..16 {
|
||||
sym[a][b] = if a == b { 0.0 } else { (cost[a][b] + cost[b][a]) / 2.0 };
|
||||
}
|
||||
}
|
||||
sym
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cardinality_probs_row_normalised() {
|
||||
let mut counts = [[0u64; 5]; 5];
|
||||
counts[0][0] = 10;
|
||||
counts[0][1] = 5;
|
||||
counts[1][0] = 5;
|
||||
counts[1][1] = 2;
|
||||
let tally = CardinalityTally { counts };
|
||||
let p = cardinality_transition_probs(&tally);
|
||||
assert!((p[0].iter().sum::<f64>() - 1.0).abs() < 1e-12);
|
||||
assert!((p[1].iter().sum::<f64>() - 1.0).abs() < 1e-12);
|
||||
assert!((p[0][0] - 10.0 / 15.0).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn composition_probs_include_diagonal() {
|
||||
let mut tally = BasePairTally { counts: [[0u64; 4]; 4], same: [0u64; 4] };
|
||||
tally.same[0] = 90; // A stays A 90 times
|
||||
tally.counts[0][1] = 10; // A -> C 10 times
|
||||
tally.counts[1][0] = 10;
|
||||
let p = composition_transition_probs(&tally);
|
||||
assert!((p[0][0] - 0.9).abs() < 1e-12);
|
||||
assert!((p[0][1] - 0.1).abs() < 1e-12);
|
||||
assert!((p[0].iter().sum::<f64>() - 1.0).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pairwise_cost_matrix_is_symmetric_and_zero_diagonal() {
|
||||
// A plausible-looking pair of models (not calibrated from real
|
||||
// data — just needs every state reachable with nonzero probability
|
||||
// for this structural test).
|
||||
let p_card = [
|
||||
[0.5, 0.3, 0.15, 0.04, 0.01],
|
||||
[0.3, 0.4, 0.2, 0.08, 0.02],
|
||||
[0.15, 0.2, 0.4, 0.2, 0.05],
|
||||
[0.04, 0.08, 0.2, 0.5, 0.18],
|
||||
[0.01, 0.02, 0.05, 0.18, 0.74],
|
||||
];
|
||||
let p_comp = [
|
||||
[0.7, 0.1, 0.15, 0.05],
|
||||
[0.1, 0.7, 0.05, 0.15],
|
||||
[0.15, 0.05, 0.7, 0.1],
|
||||
[0.05, 0.15, 0.1, 0.7],
|
||||
];
|
||||
let cost = pairwise_cost_matrix(&p_card, &p_comp);
|
||||
for a in 0..16 {
|
||||
assert_eq!(cost[a][a], 0.0);
|
||||
for b in 0..16 {
|
||||
assert!((cost[a][b] - cost[b][a]).abs() < 1e-9, "cost[{a}][{b}]={} cost[{b}][{a}]={}", cost[a][b], cost[b][a]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
pub mod error;
|
||||
pub mod meta;
|
||||
pub mod state;
|
||||
mod cardcomp;
|
||||
mod distance;
|
||||
mod dump;
|
||||
mod index;
|
||||
@@ -8,7 +9,6 @@ mod merge;
|
||||
mod numa;
|
||||
mod rebuild;
|
||||
mod reindex;
|
||||
mod sankoff;
|
||||
mod select;
|
||||
mod siblings;
|
||||
mod stats;
|
||||
@@ -20,8 +20,5 @@ pub use merge::MergeMode;
|
||||
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
|
||||
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
|
||||
pub use stats::IndexBitsPerKmer;
|
||||
pub use siblings::{BasePairTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
||||
pub use sankoff::{
|
||||
build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat, mean_substitution_cost,
|
||||
substitution_costs_from_tally, PHatEstimate, SankoffWeights,
|
||||
};
|
||||
pub use siblings::{BasePairTally, CardinalityTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
||||
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
|
||||
|
||||
@@ -1,402 +0,0 @@
|
||||
//! Sankoff parsimony cost matrix for the 16-state (powerset of `{A,C,G,T}`)
|
||||
//! family alphabet, and calibration of its parameters from real pairwise
|
||||
//! SNP/shared counts. See `docmd/theory/evolutionary_distances.md`,
|
||||
//! "Sankoff parsimony as the resolution of the 16-state model problem".
|
||||
|
||||
use crate::{BasePairTally, RawSnpDistanceOutput};
|
||||
|
||||
/// Transition base pairs in the project's fixed bit convention (see
|
||||
/// `siblings::central_base`: bit 0=A, 1=C, 2=G, 3=T). All other single-bit
|
||||
/// swaps are transversions.
|
||||
const TRANSITION_PAIRS: [(u8, u8); 2] = [(0, 2), (1, 3)]; // A<->G, C<->T
|
||||
|
||||
/// Tunable weights for the 16-state Sankoff cost matrix — see "A concrete
|
||||
/// Sankoff cost matrix for the 16-state alphabet" in the design doc.
|
||||
///
|
||||
/// Only two costs, not three: an earlier version had a separate `c_gl` for
|
||||
/// "gain/loss between two nonempty states" alongside `c_ctx` for "collapse
|
||||
/// to `∅`", but presence/absence tracking never observes "the flanks"
|
||||
/// independently of a whole 31-mer — losing one member of a family while
|
||||
/// others remain (`{A,C}->{A}`) and losing the last one (`{A}->∅`) are the
|
||||
/// same event: a specific, complete, homologous 31-mer that used to be
|
||||
/// observed no longer is. Both are exactly what `c_ctx` (see
|
||||
/// `c_ctx_from_p_hat`) prices, so it is used for every gain/loss uniformly
|
||||
/// — including a compound family losing several members at once, `∅`
|
||||
/// included, which costs `|X|*c_ctx` (no flat shortcut for `∅`: see
|
||||
/// `build_cost_matrix`'s docs for why an earlier version's flat special
|
||||
/// case there didn't hold up).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct SankoffWeights {
|
||||
/// Cost of a single-base substitution `a <-> b` (`a != b`), symmetric
|
||||
/// (`sub_cost[a][b] == sub_cost[b][a]`) and only meaningful off the
|
||||
/// diagonal. The full 6-category (AC, AG, AT, CG, CT, GT) resolution a
|
||||
/// symmetric cost matrix allows — see `substitution_costs_from_tally`
|
||||
/// for calibrating it from real data, or `SankoffWeights::ts_tv` for
|
||||
/// the coarser 2-category (transition/transversion) convenience
|
||||
/// constructor.
|
||||
pub sub_cost: [[f64; 4]; 4],
|
||||
/// Cost of losing or gaining one family member — the same constant
|
||||
/// everywhere a member is gained or lost, member count and target state
|
||||
/// (`∅` or not) included: see the struct docs.
|
||||
pub c_ctx: f64,
|
||||
}
|
||||
|
||||
impl SankoffWeights {
|
||||
/// Convenience constructor for the coarser 2-category
|
||||
/// (transition/transversion) substitution model: every transition pair
|
||||
/// (A<->G, C<->T) costs `c_ts`, every transversion pair costs `c_tv`.
|
||||
pub fn ts_tv(c_ts: f64, c_tv: f64, c_ctx: f64) -> Self {
|
||||
let mut sub_cost = [[0.0; 4]; 4];
|
||||
for a in 0..4usize {
|
||||
for b in 0..4usize {
|
||||
if a == b {
|
||||
continue;
|
||||
}
|
||||
let is_ts = TRANSITION_PAIRS.contains(&(a as u8, b as u8));
|
||||
sub_cost[a][b] = if is_ts { c_ts } else { c_tv };
|
||||
}
|
||||
}
|
||||
Self { sub_cost, c_ctx }
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the 16x16 Sankoff step-cost matrix, states indexed `0..=15` by
|
||||
/// their bitmask (matching `siblings::iupac_code`'s convention — state `0`
|
||||
/// is `∅`).
|
||||
///
|
||||
/// All 16 states, `∅` included, form a single graph: an edge of weight
|
||||
/// `c_ts`/`c_tv` between two states of equal cardinality differing by
|
||||
/// exactly one element (classified by whether the swapped bases form a
|
||||
/// transition or transversion pair), and an edge of weight `c_ctx` between
|
||||
/// a state and one that is its strict subset plus exactly one element
|
||||
/// (`∅` is a strict subset of every singleton state, so it connects
|
||||
/// directly to each of them this way — no special case). `cost(X, Y)` is
|
||||
/// the shortest-path distance in that graph (16 nodes — Floyd-Warshall,
|
||||
/// trivial at this size), uniformly, including `X <-> ∅`: a compound
|
||||
/// family losing several members at once, `∅` included, costs `|X|*c_ctx`
|
||||
/// via that many single-element steps — no cheaper flat alternative for
|
||||
/// `∅` specifically, since nothing distinguishes it from any other
|
||||
/// multi-element loss (see design doc — an earlier flat special case here
|
||||
/// was inconsistent with how every *other* multi-element transformation is
|
||||
/// already priced, and had no principled justification once that was
|
||||
/// noticed).
|
||||
pub fn build_cost_matrix(w: &SankoffWeights) -> [[f64; 16]; 16] {
|
||||
const INF: f64 = f64::INFINITY;
|
||||
let mut dist = [[INF; 16]; 16];
|
||||
for (i, row) in dist.iter_mut().enumerate() {
|
||||
row[i] = 0.0;
|
||||
}
|
||||
|
||||
for x in 0u8..16 {
|
||||
for y in (x + 1)..16 {
|
||||
let (xu, yu) = (x as usize, y as usize);
|
||||
let diff = x ^ y;
|
||||
let card_x = x.count_ones();
|
||||
let card_y = y.count_ones();
|
||||
let edge = if card_x == card_y && diff.count_ones() == 2 {
|
||||
// Exactly one element swapped: the bit only in x is the base
|
||||
// leaving, the bit only in y is the base entering.
|
||||
let a = (x & diff).trailing_zeros() as u8;
|
||||
let b = (y & diff).trailing_zeros() as u8;
|
||||
Some(w.sub_cost[a as usize][b as usize])
|
||||
} else if card_x.abs_diff(card_y) == 1 && (x & y) == x.min(y) {
|
||||
// True subset relationship: a pure gain/loss of one element.
|
||||
Some(w.c_ctx)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(cost) = edge {
|
||||
dist[xu][yu] = cost;
|
||||
dist[yu][xu] = cost;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for k in 0..16 {
|
||||
for i in 0..16 {
|
||||
if dist[i][k].is_infinite() {
|
||||
continue;
|
||||
}
|
||||
for j in 0..16 {
|
||||
let via = dist[i][k] + dist[k][j];
|
||||
if via < dist[i][j] {
|
||||
dist[i][j] = via;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
dist
|
||||
}
|
||||
|
||||
/// `p_hat` calibrated from real pairwise SNP/shared counts, restricted to
|
||||
/// pairs below `ratio_ceiling`, and its variance as a pooled Bernoulli
|
||||
/// proportion.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct PHatEstimate {
|
||||
pub p_hat: f64,
|
||||
/// `p_hat*(1-p_hat) / n_loci_total` — the pooled estimator's variance is
|
||||
/// governed by the total number of loci across all included pairs, not
|
||||
/// by any single pair's count (see design doc: this is why the
|
||||
/// exclusion criterion below is the per-pair *ratio*, not a per-pair
|
||||
/// minimum-count floor — a low-count pair barely moves the pool either
|
||||
/// way, but a saturated pair with plenty of loci would bias it).
|
||||
pub variance: f64,
|
||||
pub n_pairs_included: usize,
|
||||
pub n_loci_total: u64,
|
||||
}
|
||||
|
||||
/// Pool `snp`/`shared` counts across all genome pairs whose per-pair ratio
|
||||
/// `snp/(snp+shared)` is at most `ratio_ceiling`, then estimate
|
||||
/// `p_hat = sum(snp) / sum(snp+shared)` over the included pairs.
|
||||
///
|
||||
/// A pair with no eligible locus at all (`snp+shared == 0`) is always
|
||||
/// excluded (nothing to pool). `ratio_ceiling` should be well below 1.0 —
|
||||
/// pairs at or near saturation (e.g. cross-domain comparisons, where nearly
|
||||
/// every shared central-position family already differs) carry no
|
||||
/// information about `p_hat` and bias it upward if pooled in.
|
||||
pub fn calibrate_p_hat(raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> PHatEstimate {
|
||||
let n = raw.snp.nrows();
|
||||
let mut snp_sum: u64 = 0;
|
||||
let mut total_sum: u64 = 0;
|
||||
let mut n_pairs = 0usize;
|
||||
for i in 0..n {
|
||||
for j in (i + 1)..n {
|
||||
let snp = raw.snp[[i, j]];
|
||||
let shared = raw.shared[[i, j]];
|
||||
let total = snp + shared;
|
||||
if total == 0 {
|
||||
continue;
|
||||
}
|
||||
let ratio = snp as f64 / total as f64;
|
||||
if ratio > ratio_ceiling {
|
||||
continue;
|
||||
}
|
||||
snp_sum += snp;
|
||||
total_sum += total;
|
||||
n_pairs += 1;
|
||||
}
|
||||
}
|
||||
let p_hat = if total_sum > 0 { snp_sum as f64 / total_sum as f64 } else { 0.0 };
|
||||
let variance = if total_sum > 0 {
|
||||
p_hat * (1.0 - p_hat) / total_sum as f64
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
PHatEstimate { p_hat, variance, n_pairs_included: n_pairs, n_loci_total: total_sum }
|
||||
}
|
||||
|
||||
/// `c_ctx(p) = mean_sub_cost * [(2m*p) / (1 - (1-p)^(2m)) + p]` — see
|
||||
/// "Context, detectability, and a 3-way ordinal distance per pair" in the
|
||||
/// design doc for the bracketed term's derivation. `m` is the flank length
|
||||
/// on *each* side of the central position (`k = 2m+1`); pass `(k-1)/2`, not
|
||||
/// the project's minimizer-size parameter of the same name.
|
||||
///
|
||||
/// The bracketed term is `E[mutation count | at least one occurred]` — a
|
||||
/// *count* of mutations, not a cost. An earlier version used it directly as
|
||||
/// the cost, implicitly pricing every mutation at a flat `1` regardless of
|
||||
/// type. That stopped making sense once substitution costs were calibrated
|
||||
/// per base-pair type (`substitution_costs_from_tally`): transitions and
|
||||
/// transversions aren't equally likely (roughly 5-6x apart in this
|
||||
/// project's own data) nor equally costly, so "one mutation" isn't worth a
|
||||
/// flat unit — it's worth whatever the *empirical mix* of mutation types is
|
||||
/// worth on average. `mean_sub_cost` (see `mean_substitution_cost`) is that
|
||||
/// weighted average, turning the expected mutation *count* into an actual
|
||||
/// expected cost.
|
||||
pub fn c_ctx_from_p_hat(p_hat: f64, m: usize, mean_sub_cost: f64) -> f64 {
|
||||
if p_hat <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
let two_m = 2.0 * m as f64;
|
||||
let expected_mutations = (two_m * p_hat) / (1.0 - (1.0 - p_hat).powf(two_m)) + p_hat;
|
||||
expected_mutations * mean_sub_cost
|
||||
}
|
||||
|
||||
/// Mean substitution cost, weighted by each of the 6 base-pair categories'
|
||||
/// observed frequency in `tally` — the empirical average cost of "one
|
||||
/// mutation" under the calibrated substitution spectrum. Falls back to
|
||||
/// `1.0` (a flat per-mutation cost) if `tally` has no signal at all (no
|
||||
/// substitution ever observed), rather than dividing by zero.
|
||||
pub fn mean_substitution_cost(tally: &BasePairTally, sub_cost: &[[f64; 4]; 4]) -> f64 {
|
||||
let mut weighted_sum = 0.0;
|
||||
let mut total = 0u64;
|
||||
for a in 0..4 {
|
||||
for b in (a + 1)..4 {
|
||||
let count = tally.counts[a][b];
|
||||
weighted_sum += count as f64 * sub_cost[a][b];
|
||||
total += count;
|
||||
}
|
||||
}
|
||||
if total == 0 {
|
||||
return 1.0;
|
||||
}
|
||||
weighted_sum / total as f64
|
||||
}
|
||||
|
||||
/// Derive the 4x4 symmetric substitution cost table from an observed
|
||||
/// [`BasePairTally`]: `cost(a,b) = -ln(rate(a,b))`, normalised so the most
|
||||
/// frequently observed substitution type costs exactly `1.0` — the standard
|
||||
/// generalised-parsimony step-weighting heuristic (see design doc,
|
||||
/// "Transition/transversion refinement"), generalised here from 2
|
||||
/// categories (Ts/Tv) to the full 6-category resolution a symmetric matrix
|
||||
/// allows: a rarer substitution type is treated as less parsimonious to
|
||||
/// invoke, and therefore costs more.
|
||||
///
|
||||
/// A pair type never observed at all (`counts[a][b] == 0`) gets an infinite
|
||||
/// cost — parsimony should never spend a mutation on something with no
|
||||
/// empirical support in this data.
|
||||
pub fn substitution_costs_from_tally(tally: &BasePairTally) -> [[f64; 4]; 4] {
|
||||
let total: u64 = (0..4)
|
||||
.flat_map(|a| (a + 1..4).map(move |b| (a, b)))
|
||||
.map(|(a, b)| tally.counts[a][b])
|
||||
.sum();
|
||||
|
||||
let mut raw = [[0.0f64; 4]; 4];
|
||||
let mut min_cost = f64::INFINITY;
|
||||
for a in 0..4 {
|
||||
for b in (a + 1)..4 {
|
||||
let count = tally.counts[a][b];
|
||||
let cost = if count > 0 && total > 0 {
|
||||
-((count as f64 / total as f64).ln())
|
||||
} else {
|
||||
f64::INFINITY
|
||||
};
|
||||
raw[a][b] = cost;
|
||||
raw[b][a] = cost;
|
||||
if cost < min_cost {
|
||||
min_cost = cost;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut out = [[0.0f64; 4]; 4];
|
||||
for a in 0..4 {
|
||||
for b in 0..4 {
|
||||
if a != b {
|
||||
out[a][b] = raw[a][b] / min_cost;
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use ndarray::Array2;
|
||||
|
||||
fn weights() -> SankoffWeights {
|
||||
SankoffWeights::ts_tv(1.0, 2.0, 10.0)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_is_zero() {
|
||||
let m = build_cost_matrix(&weights());
|
||||
for row in m.iter().enumerate() {
|
||||
assert_eq!(m[row.0][row.0], 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_base_transition_costs_c_ts() {
|
||||
let m = build_cost_matrix(&weights());
|
||||
// A = 0b0001 (state 1), G = 0b0100 (state 4): A<->G is a transition.
|
||||
assert_eq!(m[1][4], 1.0);
|
||||
assert_eq!(m[4][1], 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_base_transversion_costs_c_tv() {
|
||||
let m = build_cost_matrix(&weights());
|
||||
// A = 0b0001 (state 1), C = 0b0010 (state 2): A<->C is a transversion.
|
||||
assert_eq!(m[1][2], 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gain_loss_costs_c_ctx() {
|
||||
let m = build_cost_matrix(&weights());
|
||||
// A = 0b0001 (state 1) -> {A,C} = 0b0011 (state 3): pure gain of C.
|
||||
assert_eq!(m[1][3], 10.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_set_costs_scale_with_cardinality() {
|
||||
let m = build_cost_matrix(&weights());
|
||||
assert_eq!(m[0][1], 10.0); // {A} -> ∅: one member, one c_ctx step.
|
||||
// N (all four) -> ∅: no flat shortcut — four single-element steps,
|
||||
// same as losing four members down to a nonempty state would cost.
|
||||
assert_eq!(m[0][15], 40.0);
|
||||
assert_eq!(m[0][0], 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matrix_is_symmetric() {
|
||||
let m = build_cost_matrix(&weights());
|
||||
for i in 0..16 {
|
||||
for j in 0..16 {
|
||||
assert_eq!(m[i][j], m[j][i], "asymmetry at ({i},{j})");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn calibrate_p_hat_excludes_saturated_pairs() {
|
||||
// 3 genomes: pair (0,1) informative (10% SNP), pair (0,2) saturated
|
||||
// (100% SNP) — must be excluded from the pooled estimate.
|
||||
let mut snp = Array2::<u64>::zeros((3, 3));
|
||||
let mut shared = Array2::<u64>::zeros((3, 3));
|
||||
snp[[0, 1]] = 10;
|
||||
shared[[0, 1]] = 90;
|
||||
snp[[1, 0]] = 10;
|
||||
shared[[1, 0]] = 90;
|
||||
snp[[0, 2]] = 50;
|
||||
shared[[0, 2]] = 0;
|
||||
snp[[2, 0]] = 50;
|
||||
shared[[2, 0]] = 0;
|
||||
let raw = RawSnpDistanceOutput { snp, shared };
|
||||
|
||||
let est = calibrate_p_hat(&raw, 0.5);
|
||||
assert_eq!(est.n_pairs_included, 1);
|
||||
assert_eq!(est.n_loci_total, 100);
|
||||
assert!((est.p_hat - 0.1).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn c_ctx_from_p_hat_zero_is_zero() {
|
||||
assert_eq!(c_ctx_from_p_hat(0.0, 15, 1.5), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn c_ctx_from_p_hat_matches_known_value() {
|
||||
// p=0.1, m=15: [(30*0.1)/(1-(0.9)^30) + 0.1] * mean_sub_cost
|
||||
let m = 15usize;
|
||||
let p = 0.1;
|
||||
let mean_sub_cost = 1.5;
|
||||
let expected = ((2.0 * m as f64 * p) / (1.0 - (1.0 - p).powf(2.0 * m as f64)) + p) * mean_sub_cost;
|
||||
assert!((c_ctx_from_p_hat(p, m, mean_sub_cost) - expected).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mean_substitution_cost_is_weighted_average() {
|
||||
let mut counts = [[0u64; 4]; 4];
|
||||
counts[0][1] = 3; // A/C
|
||||
counts[1][0] = 3;
|
||||
counts[0][2] = 1; // A/G
|
||||
counts[2][0] = 1;
|
||||
let tally = BasePairTally { counts };
|
||||
let mut sub_cost = [[0.0f64; 4]; 4];
|
||||
sub_cost[0][1] = 2.0;
|
||||
sub_cost[1][0] = 2.0;
|
||||
sub_cost[0][2] = 1.0;
|
||||
sub_cost[2][0] = 1.0;
|
||||
// (3*2.0 + 1*1.0) / 4 = 1.75
|
||||
assert!((mean_substitution_cost(&tally, &sub_cost) - 1.75).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mean_substitution_cost_falls_back_to_one_with_no_data() {
|
||||
let tally = BasePairTally { counts: [[0u64; 4]; 4] };
|
||||
let sub_cost = [[0.0f64; 4]; 4];
|
||||
assert_eq!(mean_substitution_cost(&tally, &sub_cost), 1.0);
|
||||
}
|
||||
}
|
||||
+235
-20
@@ -660,10 +660,18 @@ impl KmerIndex {
|
||||
/// (tallied once, at its minorant) of every layer of the already-built
|
||||
/// sibling annex, resolves each genome's single observed form (`None`
|
||||
/// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j,
|
||||
/// bi, bj)` for every genome pair `(i, j)` where both are unambiguous
|
||||
/// and single-copy (`bi == bj` means shared at that locus, `bi != bj`
|
||||
/// means a SNP). Layers are processed in parallel (rayon); each gets
|
||||
/// its own accumulator from `zero()`, combined pairwise via `combine`.
|
||||
/// bi, bj, variable)` for every genome pair `(i, j)` where both are
|
||||
/// unambiguous and single-copy (`bi == bj` means shared at that locus,
|
||||
/// `bi != bj` means a SNP). `variable` is the family's own
|
||||
/// `family_size() >= 2` (true if more than one member is observed
|
||||
/// *anywhere* in the family, i.e. it isn't fully invariant across the
|
||||
/// whole index) — `raw_snp_distance` ignores it (a fully-invariant
|
||||
/// family is still legitimately "shared"), but callers whose diagonal
|
||||
/// should only reflect genuine SNP-adjacent agreement, not the
|
||||
/// genome-wide invariant background, need it (see
|
||||
/// [`base_pair_tally`](Self::base_pair_tally)'s `same` field). Layers
|
||||
/// are processed in parallel (rayon); each gets its own accumulator
|
||||
/// from `zero()`, combined pairwise via `combine`.
|
||||
fn scan_family_pairs<Acc, F, C>(
|
||||
&self,
|
||||
label: &str,
|
||||
@@ -673,7 +681,7 @@ impl KmerIndex {
|
||||
) -> OKIResult<Acc>
|
||||
where
|
||||
Acc: Send,
|
||||
F: Fn(&mut Acc, usize, usize, u8, u8) + Sync,
|
||||
F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync,
|
||||
C: Fn(Acc, Acc) -> Acc,
|
||||
{
|
||||
let n_parts = self.n_partitions();
|
||||
@@ -751,6 +759,7 @@ impl KmerIndex {
|
||||
if !is_minorant(kmer, mask, k) {
|
||||
continue; // family tallied once, at its minorant
|
||||
}
|
||||
let variable = mask.family_size() >= 2;
|
||||
|
||||
single_form.clear();
|
||||
single_form.resize(n_cols, None);
|
||||
@@ -791,7 +800,7 @@ impl KmerIndex {
|
||||
continue;
|
||||
}
|
||||
let Some(bj) = single_form[j] else { continue };
|
||||
on_pair(&mut acc, i, j, bi, bj);
|
||||
on_pair(&mut acc, i, j, bi, bj, variable);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -816,7 +825,7 @@ impl KmerIndex {
|
||||
let (snp, shared) = self.scan_family_pairs(
|
||||
"raw_snp_distance",
|
||||
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
|
||||
|(snp, shared), i, j, bi, bj| {
|
||||
|(snp, shared), i, j, bi, bj, _variable| {
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
shared[[j, i]] += 1;
|
||||
@@ -837,9 +846,10 @@ impl KmerIndex {
|
||||
/// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG,
|
||||
/// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs
|
||||
/// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` —
|
||||
/// same saturation-exclusion discipline as `calibrate_p_hat`, for the
|
||||
/// same reason: a saturated pair's observed base-pair mix trends toward
|
||||
/// neutral base composition, not the true point-mutation spectrum.
|
||||
/// same saturation-exclusion discipline as
|
||||
/// [`cardinality_tally`](Self::cardinality_tally), for the same reason:
|
||||
/// a saturated pair's observed base-pair mix trends toward neutral base
|
||||
/// composition, not the true point-mutation spectrum.
|
||||
///
|
||||
/// A second full pass over the annex, sharing
|
||||
/// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by
|
||||
@@ -858,25 +868,38 @@ impl KmerIndex {
|
||||
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
|
||||
});
|
||||
|
||||
let counts = self.scan_family_pairs(
|
||||
let (counts, same) = self.scan_family_pairs(
|
||||
"base_pair_tally",
|
||||
|| [[0u64; 4]; 4],
|
||||
|counts, i, j, bi, bj| {
|
||||
if bi != bj && included[[i, j]] {
|
||||
|| ([[0u64; 4]; 4], [0u64; 4]),
|
||||
|(counts, same), i, j, bi, bj, variable| {
|
||||
if !included[[i, j]] {
|
||||
return;
|
||||
}
|
||||
if bi != bj {
|
||||
counts[bi as usize][bj as usize] += 1;
|
||||
counts[bj as usize][bi as usize] += 1;
|
||||
} else if variable {
|
||||
// Only count "stayed the same" from families that vary
|
||||
// *somewhere* in the index — a fully invariant family
|
||||
// (never varies anywhere) isn't a SNP-adjacent
|
||||
// agreement, it's genome-wide background, and would
|
||||
// otherwise swamp the diagonal (see
|
||||
// `docmd/theory/evolutionary_distances.md`, the
|
||||
// ascertainment-bias regression this was reverting).
|
||||
same[bi as usize] += 1;
|
||||
}
|
||||
},
|
||||
|mut total, partial| {
|
||||
|(mut counts, mut same), (partial_counts, partial_same)| {
|
||||
for a in 0..4 {
|
||||
same[a] += partial_same[a];
|
||||
for b in 0..4 {
|
||||
total[a][b] += partial[a][b];
|
||||
counts[a][b] += partial_counts[a][b];
|
||||
}
|
||||
}
|
||||
total
|
||||
(counts, same)
|
||||
},
|
||||
)?;
|
||||
Ok(BasePairTally { counts })
|
||||
Ok(BasePairTally { counts, same })
|
||||
}
|
||||
}
|
||||
|
||||
@@ -884,9 +907,15 @@ impl KmerIndex {
|
||||
pub struct BasePairTally {
|
||||
/// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled
|
||||
/// over included genome pairs, where the two genomes' single forms are
|
||||
/// `a` and `b` (0=A, 1=C, 2=G, 3=T). Diagonal always `0` (an `a == b`
|
||||
/// locus is `shared`, not tallied here).
|
||||
/// `a` and `b` (0=A, 1=C, 2=G, 3=T). Diagonal always `0` — an `a == b`
|
||||
/// locus is counted in `same`, not here.
|
||||
pub counts: [[u64; 4]; 4],
|
||||
/// `same[a]` = number of eligible loci, pooled over included genome
|
||||
/// pairs, where both genomes' single forms are `a` — the diagonal
|
||||
/// `counts` omits, needed to build a proper row-stochastic composition
|
||||
/// probability matrix (the "stay the same base" entries), not just the
|
||||
/// substitution-cost off-diagonal.
|
||||
pub same: [u64; 4],
|
||||
}
|
||||
|
||||
/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
|
||||
@@ -918,6 +947,192 @@ fn iupac_code(mask: u8) -> u8 {
|
||||
}
|
||||
}
|
||||
|
||||
/// See [`KmerIndex::cardinality_tally`].
|
||||
pub struct CardinalityTally {
|
||||
/// `counts[a][b] == counts[b][a]` = number of family sites, pooled over
|
||||
/// included genome pairs, where one genome's family cardinality
|
||||
/// (popcount of its presence mask, `0..=4`) is `a` and the other's is
|
||||
/// `b`. Diagonal is real data here (both genomes at the same
|
||||
/// cardinality), unlike [`BasePairTally::counts`].
|
||||
pub counts: [[u64; 5]; 5],
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Cardinality co-occurrence, pooled only over genome pairs whose
|
||||
/// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same
|
||||
/// saturation/no-data exclusion discipline as
|
||||
/// [`base_pair_tally`](Self::base_pair_tally). Unlike
|
||||
/// [`scan_family_pairs`](Self::scan_family_pairs) (which resolves each
|
||||
/// genome to a single form and silently drops any genome carrying more
|
||||
/// than one member of the family), this needs the *full* per-genome
|
||||
/// presence mask — a family member count of 2, 3 or 4 is exactly the
|
||||
/// signal being tallied, not noise to discard — so it re-implements the
|
||||
/// traversal rather than reusing that helper.
|
||||
///
|
||||
/// Restricted to variable families (`family_size() >= 2`), matching
|
||||
/// `snp_pseudo_alignment`'s own scope — briefly removed, then
|
||||
/// reinstated: without it, the diagonal is dominated by genome-wide
|
||||
/// invariant background (family_size()<2 loci vastly outnumber the
|
||||
/// ones that ever vary anywhere), which is inconsistent with the
|
||||
/// `+ASC`-corrected alignment this matrix is ultimately used with —
|
||||
/// `+ASC` exists specifically because the likelihood only ever sees
|
||||
/// variable sites, so a rate model calibrated mostly from invariant
|
||||
/// background sites doesn't describe the population it's applied to.
|
||||
/// Verified empirically: removing the filter measurably worsened a
|
||||
/// real IQ-TREE run (log-likelihood dropped, `NNI search needs
|
||||
/// unusual large number of steps to converge` warnings appeared) — see
|
||||
/// `docmd/theory/evolutionary_distances.md` for the full account.
|
||||
/// [`base_pair_tally`](Self::base_pair_tally)'s own diagonal (`same`)
|
||||
/// gets the matching restriction via `scan_family_pairs`'s new
|
||||
/// `variable` flag, rather than a `family_size()` check of its own (it
|
||||
/// doesn't have direct access to the family's mask).
|
||||
pub fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
|
||||
if i == j {
|
||||
return false;
|
||||
}
|
||||
let snp = raw.snp[[i, j]];
|
||||
let total = snp + raw.shared[[i, j]];
|
||||
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
|
||||
});
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
&self.root_path,
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
|
||||
|
||||
let mut layer_dirs = Vec::new();
|
||||
for part in 0..n_parts {
|
||||
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
continue;
|
||||
}
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||||
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||||
if !annex_path.exists() {
|
||||
return Err(OKIError::InvalidInput(format!(
|
||||
"no sibling annex at {} — run build_sibling_annex first",
|
||||
annex_path.display()
|
||||
)));
|
||||
}
|
||||
layer_dirs.push(layer_dir);
|
||||
}
|
||||
}
|
||||
|
||||
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
|
||||
let partials: Vec<[[u64; 5]; 5]> = layer_dirs
|
||||
.par_iter()
|
||||
.map(|layer_dir| -> OKIResult<[[u64; 5]; 5]> {
|
||||
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
|
||||
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
|
||||
|
||||
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
|
||||
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
|
||||
.map_err(OKIError::Partition)?;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
slot_kmer[slot] = Some(kmer);
|
||||
}
|
||||
}
|
||||
|
||||
let use_counts = with_counts && layer_dir.join("counts").exists();
|
||||
let mat = if use_counts {
|
||||
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
|
||||
} else {
|
||||
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
|
||||
};
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
|
||||
let mut counts = [[0u64; 5]; 5];
|
||||
let mut genome_mask: Vec<u8> = Vec::with_capacity(n_genomes);
|
||||
|
||||
for slot in 0..annex.len() {
|
||||
let Some(mask) = annex.get(slot) else { continue };
|
||||
let Some(kmer) = slot_kmer[slot] else { continue };
|
||||
if !is_minorant(kmer, mask, k) {
|
||||
continue; // family tallied once, at its minorant
|
||||
}
|
||||
if mask.family_size() < 2 {
|
||||
// Fully invariant family (never varies anywhere in
|
||||
// the index) — genome-wide background, not
|
||||
// SNP-adjacent signal; would otherwise swamp the
|
||||
// diagonal (`c=1/c=1` etc.), which needs to reflect
|
||||
// the same variable-families-only population the
|
||||
// `+ASC`-corrected alignment/likelihood actually
|
||||
// models. See `base_pair_tally`'s `variable` gate
|
||||
// on its own `same` diagonal for the matching fix.
|
||||
continue;
|
||||
}
|
||||
|
||||
genome_mask.clear();
|
||||
genome_mask.resize(n_genomes, 0);
|
||||
|
||||
for other in kmer.central_canonical_neighbors() {
|
||||
let base = central_base(other, k);
|
||||
if !mask.has(base) {
|
||||
continue;
|
||||
}
|
||||
let presence: Option<Vec<bool>> = if other == kmer {
|
||||
Some((0..n_cols).map(|g| mat.carries(g, slot)).collect())
|
||||
} else {
|
||||
let dest = partition_of(other, n_parts);
|
||||
cache.find_presence(dest, other, n_genomes)
|
||||
};
|
||||
let Some(presence) = presence else { continue };
|
||||
for (g, &present) in presence.iter().enumerate() {
|
||||
if present {
|
||||
genome_mask[g] |= 1 << base;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for i in 0..n_genomes {
|
||||
let card_i = genome_mask[i].count_ones() as usize;
|
||||
for j in (i + 1)..n_genomes {
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
let card_j = genome_mask[j].count_ones() as usize;
|
||||
counts[card_i][card_j] += 1;
|
||||
if card_i != card_j {
|
||||
counts[card_j][card_i] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pb.inc(1);
|
||||
Ok(counts)
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
pb.finish_and_clear();
|
||||
|
||||
let mut total = [[0u64; 5]; 5];
|
||||
for partial in partials {
|
||||
for a in 0..5 {
|
||||
for b in 0..5 {
|
||||
total[a][b] += partial[a][b];
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(CardinalityTally { counts: total })
|
||||
}
|
||||
}
|
||||
|
||||
/// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per
|
||||
/// genome, one column per variable family (`family_size() >= 2` — monomorphic
|
||||
/// families carry no signal and are skipped, unlike `raw_snp_distance`'s
|
||||
|
||||
@@ -10,9 +10,9 @@ use clap::Args;
|
||||
use kodama::{Method, linkage};
|
||||
use obifastwrite::{JsonVal, write_record};
|
||||
use obikindex::{
|
||||
DistanceMetric, KmerIndex, RawSnpDistanceOutput, SankoffWeights,
|
||||
SiblingAnnexStats, SnpAlignment, build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat,
|
||||
mean_substitution_cost, substitution_costs_from_tally,
|
||||
DistanceMetric, KmerIndex, RawSnpDistanceOutput,
|
||||
SiblingAnnexStats, SnpAlignment,
|
||||
cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix,
|
||||
};
|
||||
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
|
||||
use tracing::info;
|
||||
@@ -87,19 +87,20 @@ pub struct DistanceArgs {
|
||||
|
||||
/// Exclude a genome (by its exact label) from every computation below
|
||||
/// that reads the sibling annex — `--raw-snp-distance`/`--raw-snp-counts`,
|
||||
/// `--snp`, and `--sankoff` (and everything `--sankoff` implies:
|
||||
/// `p_hat`, `sub_cost`, `c_ctx`, the exported matrix/alignment,
|
||||
/// `--tnt`/`--phyg`/`--iqtree`). Repeatable. Does *not* affect the plain
|
||||
/// `--metric` distance matrix/NJ/UPGMA path (a different, unrelated
|
||||
/// computation). Applied by zeroing the excluded genome's row/column
|
||||
/// after `raw_snp_distance` runs (a pair with zero counts is already
|
||||
/// skipped by `calibrate_p_hat`/`base_pair_tally`, so this needs no
|
||||
/// change to the underlying traversal) and by dropping its row from
|
||||
/// `snp_pseudo_alignment`'s output — the annex is still built/scanned
|
||||
/// for the excluded genome too, just not used afterward. For a genome
|
||||
/// with almost no informative sites shared with anything else (see
|
||||
/// `docmd/theory/evolutionary_distances.md`, the IQ-TREE/Mash rogue-taxon
|
||||
/// discussion), its presence can otherwise silently bias `p_hat`/`R`.
|
||||
/// `--snp`, and `--sankoff` (and everything `--sankoff` implies: the
|
||||
/// cardinality/composition transition models, the exported
|
||||
/// matrix/alignment, `--tnt`/`--phyg`/`--iqtree`). Repeatable. Does
|
||||
/// *not* affect the plain `--metric` distance matrix/NJ/UPGMA path (a
|
||||
/// different, unrelated computation). Applied by zeroing the excluded
|
||||
/// genome's row/column after `raw_snp_distance` runs (a pair with zero
|
||||
/// counts is already skipped by `base_pair_tally`/`cardinality_tally`,
|
||||
/// so this needs no change to the underlying traversal) and by
|
||||
/// dropping its row from `snp_pseudo_alignment`'s output — the annex
|
||||
/// is still built/scanned for the excluded genome too, just not used
|
||||
/// afterward. For a genome with almost no informative sites shared
|
||||
/// with anything else (see `docmd/theory/evolutionary_distances.md`,
|
||||
/// the IQ-TREE/Mash rogue-taxon discussion), its presence can
|
||||
/// otherwise silently bias the transition models.
|
||||
#[arg(long = "exclude-genome", value_name = "LABEL")]
|
||||
pub exclude_genome: Vec<String>,
|
||||
|
||||
@@ -231,11 +232,12 @@ pub fn run(args: DistanceArgs) {
|
||||
|
||||
// ── Genome exclusion (`--exclude-genome`) ───────────────────────────────
|
||||
// Applied by zeroing a `RawSnpDistanceOutput`'s excluded rows/columns
|
||||
// (`zero_excluded_pairs`) — `calibrate_p_hat`/`base_pair_tally` already
|
||||
// skip any pair with zero total counts, so this needs no change to
|
||||
// `obikindex`'s traversal — and by dropping the excluded genome's row
|
||||
// from a `SnpAlignment` plus the matching label (`drop_excluded`),
|
||||
// since an all-`∅` row for an "excluded" genome would otherwise still
|
||||
// (`zero_excluded_pairs`) — `base_pair_tally`/`cardinality_tally`
|
||||
// already skip any pair with zero total counts, so this needs no
|
||||
// change to `obikindex`'s traversal — and by dropping the excluded
|
||||
// genome's row from a `SnpAlignment` plus the matching label
|
||||
// (`drop_excluded`), since an all-`∅` row for an "excluded" genome
|
||||
// would otherwise still
|
||||
// reach TNT/PhyG/IQ-TREE as a real (empty) taxon.
|
||||
let exclude_mask: Vec<bool> = {
|
||||
let mut mask = vec![false; n];
|
||||
@@ -263,15 +265,41 @@ pub fn run(args: DistanceArgs) {
|
||||
}
|
||||
}
|
||||
};
|
||||
// `snp_pseudo_alignment`'s "variable family" criterion
|
||||
// (`mask.family_size() >= 2`) is a property of the annex, computed
|
||||
// over *every* genome in the index — unaffected by `--exclude-genome`.
|
||||
// So dropping excluded rows alone can leave columns that are variable
|
||||
// only thanks to an excluded genome now monomorphic among the
|
||||
// survivors — silently wrong data for TNT/PhyG, and a hard failure
|
||||
// for IQ-TREE's `+ASC` (verified: excluding 2 taxa on the 20-genome
|
||||
// benchmark left 116,351 such columns). Re-check variability among the
|
||||
// *kept* genomes only, after dropping rows, and drop those columns too.
|
||||
let drop_excluded = |alignment: SnpAlignment| -> (SnpAlignment, Vec<String>) {
|
||||
let sequences = alignment.sequences.into_iter().enumerate()
|
||||
let mut sequences: Vec<Vec<u8>> = alignment.sequences.into_iter().enumerate()
|
||||
.filter(|(i, _)| !exclude_mask[*i])
|
||||
.map(|(_, seq)| seq)
|
||||
.collect();
|
||||
let kept_labels = labels.iter().enumerate()
|
||||
let kept_labels: Vec<String> = labels.iter().enumerate()
|
||||
.filter(|(i, _)| !exclude_mask[*i])
|
||||
.map(|(_, l)| l.clone())
|
||||
.collect();
|
||||
|
||||
if exclude_mask.iter().any(|&excluded| excluded) && !sequences.is_empty() {
|
||||
let n_sites = sequences[0].len();
|
||||
let keep_col: Vec<bool> = (0..n_sites)
|
||||
.map(|site| sequences.iter().any(|seq| seq[site] != sequences[0][site]))
|
||||
.collect();
|
||||
for seq in &mut sequences {
|
||||
let mut kept = Vec::with_capacity(seq.len());
|
||||
for (site, &b) in seq.iter().enumerate() {
|
||||
if keep_col[site] {
|
||||
kept.push(b);
|
||||
}
|
||||
}
|
||||
*seq = kept;
|
||||
}
|
||||
}
|
||||
|
||||
(SnpAlignment { sequences }, kept_labels)
|
||||
};
|
||||
|
||||
@@ -330,21 +358,20 @@ pub fn run(args: DistanceArgs) {
|
||||
std::process::exit(1);
|
||||
});
|
||||
zero_excluded_pairs(&mut raw);
|
||||
let estimate = calibrate_p_hat(&raw, args.sankoff_ratio_ceiling);
|
||||
let m = (idx.kmer_size() - 1) / 2;
|
||||
|
||||
let tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
|
||||
let base_tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
|
||||
eprintln!("error computing base-pair tally: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
let sub_cost = substitution_costs_from_tally(&tally);
|
||||
let mean_sub_cost = mean_substitution_cost(&tally, &sub_cost);
|
||||
let c_ctx = c_ctx_from_p_hat(estimate.p_hat, m, mean_sub_cost);
|
||||
|
||||
let weights = SankoffWeights { sub_cost, c_ctx };
|
||||
let matrix = build_cost_matrix(&weights);
|
||||
write_sankoff_matrix_csv(&matrix, &estimate, &weights, &args.output);
|
||||
write_sankoff_params(&estimate, &tally, &weights, args.sankoff_ratio_ceiling, m, mean_sub_cost, &args.output);
|
||||
let card_tally = idx.cardinality_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
|
||||
eprintln!("error computing cardinality tally: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
let p_card = cardinality_transition_probs(&card_tally);
|
||||
let p_comp = composition_transition_probs(&base_tally);
|
||||
let matrix = pairwise_cost_matrix(&p_card, &p_comp);
|
||||
write_sankoff_matrix_csv(&matrix, &args.output);
|
||||
write_sankoff_params(&card_tally, &p_card, &base_tally, &p_comp, args.sankoff_ratio_ceiling, &args.output);
|
||||
|
||||
let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP pseudo-alignment: {e}");
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::io::{BufWriter, Write};
|
||||
use std::path::PathBuf;
|
||||
|
||||
use obifastwrite::{JsonVal, write_record};
|
||||
use obikindex::{BasePairTally, PHatEstimate, SankoffWeights, SnpAlignment};
|
||||
use obikindex::{BasePairTally, CardinalityTally, SnpAlignment};
|
||||
use tracing::info;
|
||||
|
||||
// ── Sankoff pseudo-alignment → FASTA ────────────────────────────────────────
|
||||
@@ -65,19 +65,8 @@ pub(super) fn state_index_table() -> [u8; 128] {
|
||||
|
||||
pub(super) fn write_sankoff_matrix_csv(
|
||||
matrix: &[[f64; 16]; 16],
|
||||
estimate: &PHatEstimate,
|
||||
weights: &SankoffWeights,
|
||||
output: &Option<PathBuf>,
|
||||
) {
|
||||
info!(
|
||||
p_hat = format_args!("{:.6}", estimate.p_hat),
|
||||
variance = format_args!("{:.3e}", estimate.variance),
|
||||
n_pairs_included = estimate.n_pairs_included,
|
||||
n_loci_total = estimate.n_loci_total,
|
||||
c_ctx = format_args!("{:.4}", weights.c_ctx),
|
||||
"Sankoff matrix calibration"
|
||||
);
|
||||
|
||||
let path = output.as_ref()
|
||||
.map(|p| format!("{}_sankoff_matrix.csv", p.display()))
|
||||
.unwrap_or_else(|| "sankoff_matrix.csv".into());
|
||||
@@ -109,60 +98,64 @@ pub(super) fn write_sankoff_matrix_csv(
|
||||
// rather than re-parsing free text.
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct SankoffSubstitution {
|
||||
pair: String,
|
||||
struct CardinalityTransition {
|
||||
from: usize,
|
||||
to: usize,
|
||||
count: u64,
|
||||
cost: f64,
|
||||
probability: f64,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct CompositionTransition {
|
||||
from: char,
|
||||
to: char,
|
||||
count: u64,
|
||||
probability: f64,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct SankoffParamsReport {
|
||||
ratio_ceiling: f64,
|
||||
flank_length_m: usize,
|
||||
p_hat: f64,
|
||||
p_hat_variance: f64,
|
||||
n_pairs_included: usize,
|
||||
n_loci_total: u64,
|
||||
/// Weighted-average substitution cost — see `c_ctx_from_p_hat`'s docs:
|
||||
/// this is what turns the raw expected mutation *count* behind `c_ctx`
|
||||
/// into an actual cost (not every mutation is worth a flat `1`).
|
||||
mean_sub_cost: f64,
|
||||
c_ctx: f64,
|
||||
substitutions: Vec<SankoffSubstitution>,
|
||||
cardinality_transitions: Vec<CardinalityTransition>,
|
||||
composition_transitions: Vec<CompositionTransition>,
|
||||
}
|
||||
|
||||
pub(super) fn write_sankoff_params(
|
||||
estimate: &PHatEstimate,
|
||||
tally: &BasePairTally,
|
||||
weights: &SankoffWeights,
|
||||
card_tally: &CardinalityTally,
|
||||
p_card: &[[f64; 5]; 5],
|
||||
base_tally: &BasePairTally,
|
||||
p_comp: &[[f64; 4]; 4],
|
||||
ratio_ceiling: f64,
|
||||
m: usize,
|
||||
mean_sub_cost: f64,
|
||||
output: &Option<PathBuf>,
|
||||
) {
|
||||
const BASE_LETTER: [char; 4] = ['A', 'C', 'G', 'T'];
|
||||
|
||||
let mut substitutions = Vec::with_capacity(6);
|
||||
for a in 0..4 {
|
||||
for b in (a + 1)..4 {
|
||||
substitutions.push(SankoffSubstitution {
|
||||
pair: format!("{}/{}", BASE_LETTER[a], BASE_LETTER[b]),
|
||||
count: tally.counts[a][b],
|
||||
cost: weights.sub_cost[a][b],
|
||||
let mut cardinality_transitions = Vec::with_capacity(25);
|
||||
for a in 0..5 {
|
||||
for b in 0..5 {
|
||||
cardinality_transitions.push(CardinalityTransition {
|
||||
from: a,
|
||||
to: b,
|
||||
count: card_tally.counts[a][b],
|
||||
probability: p_card[a][b],
|
||||
});
|
||||
}
|
||||
}
|
||||
let report = SankoffParamsReport {
|
||||
ratio_ceiling,
|
||||
flank_length_m: m,
|
||||
p_hat: estimate.p_hat,
|
||||
p_hat_variance: estimate.variance,
|
||||
n_pairs_included: estimate.n_pairs_included,
|
||||
n_loci_total: estimate.n_loci_total,
|
||||
mean_sub_cost,
|
||||
c_ctx: weights.c_ctx,
|
||||
substitutions,
|
||||
};
|
||||
|
||||
let mut composition_transitions = Vec::with_capacity(16);
|
||||
for a in 0..4 {
|
||||
for b in 0..4 {
|
||||
let count = if a == b { base_tally.same[a] } else { base_tally.counts[a][b] };
|
||||
composition_transitions.push(CompositionTransition {
|
||||
from: BASE_LETTER[a],
|
||||
to: BASE_LETTER[b],
|
||||
count,
|
||||
probability: p_comp[a][b],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let report = SankoffParamsReport { ratio_ceiling, cardinality_transitions, composition_transitions };
|
||||
|
||||
let path = output.as_ref()
|
||||
.map(|p| format!("{}_sankoff_params.yaml", p.display()))
|
||||
@@ -183,15 +176,18 @@ pub(super) fn write_sankoff_params(
|
||||
/// closure* of the result (Floyd-Warshall over the 16 states again, on the
|
||||
/// now-integer values).
|
||||
///
|
||||
/// `matrix` is already a metric in its real-valued form (it's a
|
||||
/// shortest-path closure itself — see `obikindex::build_cost_matrix`), but
|
||||
/// rounding each cell independently can still break the triangle
|
||||
/// inequality: e.g. two real costs of `1.734` each round to `173`, summing
|
||||
/// to `346`, while their own real sum `3.468` rounds to `347` — TNT then
|
||||
/// reports "triangle inequality violated ... Fixed" and silently
|
||||
/// substitutes its own corrected value. Re-closing after rounding makes
|
||||
/// that correction explicit and reproducible here instead, rather than
|
||||
/// left implicit and tool-version-dependent.
|
||||
/// Unlike this project's earlier cost-matrix construction (a graph closed
|
||||
/// by shortest path, guaranteeing a metric by construction), `matrix` here
|
||||
/// comes from `obikindex::pairwise_cost_matrix`'s row-normalise-then-`-ln`
|
||||
/// composition, which gives no such guarantee — so this closure isn't only
|
||||
/// needed to correct integer-rounding artifacts (two real costs of `1.734`
|
||||
/// each round to `173`, summing to `346`, while their own real sum `3.468`
|
||||
/// rounds to `347` — TNT then reports "triangle inequality violated ...
|
||||
/// Fixed" and silently substitutes its own corrected value), it may also
|
||||
/// be the only thing making the *real-valued* matrix a metric in the first
|
||||
/// place. Re-closing after rounding makes both corrections explicit and
|
||||
/// reproducible here instead, rather than left implicit and
|
||||
/// tool-version-dependent.
|
||||
pub(super) fn scaled_metric_matrix(matrix: &[[f64; 16]; 16], cost_scale: f64) -> [[i64; 16]; 16] {
|
||||
let mut m = [[0i64; 16]; 16];
|
||||
for i in 0..16 {
|
||||
|
||||
Reference in New Issue
Block a user