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
2026-08-13 16:58:12 +02:00
parent c26623fa00
commit 0c86ea0385
9 changed files with 901 additions and 549 deletions
+130 -9
View File
@@ -1094,15 +1094,136 @@ identified. Gives 120 parameters, but derived from two small,
well-estimated pieces (a 5×5 cardinality model, a 4×4 composition model)
rather than fit or smoothed independently per pair.
**Status: designed, not implemented.** Would replace `build_cost_matrix`
(`obikindex/src/sankoff.rs`) and the single `c_ctx` scalar/parameter
entirely; `sub_cost`'s own calibration is untouched. Not yet decided
whether the "clean" parsimony-graph cost (`build_cost_matrix`'s current
output, still needed for `--tnt`/`--phyg`) should be replaced too, kept
as a separate simpler approximation, or derived as a special case of this
same decomposition (it likely can be — the same `shared`/`lost`/`gained`
logic, with the cardinality model reduced back to a single scalar `c_ctx`
substituted in, recovers exactly the current graph).
### Implemented (2026-08-12): `pairwise_cost_matrix` replaces `build_cost_matrix` entirely
New module `obikindex/src/cardcomp.rs`, replacing `sankoff::build_cost_matrix`
and the `c_ctx`/`SankoffWeights`/`PHatEstimate`/`calibrate_p_hat`/
`c_ctx_from_p_hat`/`mean_substitution_cost`/`substitution_costs_from_tally`
machinery it depended on outright — not kept in parallel as a fallback
(all now unreferenced outside their own tests; `sankoff.rs` itself is a
pending removal, not yet done).
**New primitives, `obikindex/src/siblings.rs`:**
- `BasePairTally` gained a `same: [u64; 4]` field (diagonal — "both
genomes at the same single base", pooled from
[`base_pair_tally`](obikindex::KmerIndex::base_pair_tally)'s existing
traversal, extended to also tally the `bi == bj` case it previously
discarded).
- `CardinalityTally { counts: [[u64; 5]; 5] }` and
`KmerIndex::cardinality_tally`, a new traversal (same shape as
`snp_pseudo_alignment`'s — needs full per-genome presence masks, not
`scan_family_pairs`'s single-resolved-form view, since cardinality 2-4
is exactly the signal being tallied, not noise to drop). Same
saturated/no-data pair exclusion as `base_pair_tally`. Restricted to
variable families (`family_size() >= 2`), matching
`snp_pseudo_alignment`'s own scope. Verified against the 20-genome
benchmark: counts match the earlier hand-rolled Python analysis exactly
(e.g. `c=0/c=0: 117,158,166`, `c=0/c=1: 13,707,223` — the same numbers
this whole investigation started from).
**`cardcomp.rs`:**
- `cardinality_transition_probs`/`composition_transition_probs`: row-
normalise the two tallies into proper transition probability matrices,
diagonal included ("stay the same" is a real, calibrated outcome).
- `pairwise_cost_matrix`: for every pair of the 16 states, `shared = A∩B`
contributes `∏ P_composition(x→x)`; `lost = A\B`, `gained = B\A` are
parsimony-paired (`best_pairing_cost`, brute-force over the ≤4!
injections — small enough that hand-rolling beats a dependency) into
substitution events on `P_composition`, minimising total `-ln`; the
cardinality-difference leftover is priced once via
`P_cardinality(|A|→|B|)`, never chained. Row-normalised, `-ln`'d, then
**symmetrised**: `cost_sym(A,B) = (cost(A,B)+cost(B,A))/2` — equivalent
to taking the *geometric* mean of the two raw probabilities
(`-ln(√(P(A,B)·P(B,A))) = (-ln P(A,B) - ln P(B,A))/2`), not their
arithmetic mean. Required, not just convenient for IQ-TREE's
lower-triangular file format: Sankoff parsimony's score is independent
of where an *unrooted* tree (what TNT/PhyG actually search over) gets
rooted only if the cost matrix is symmetric — the discrete-parsimony
analogue of CTMC reversibility, established by direct reasoning, not
assumed. Bonus of the same decision: 120 free parameters instead of the
240 a fully asymmetric matrix would need.
**Verified on the 20-genome benchmark**: resulting matrix symmetric
(checked numerically, zero asymmetric cells), zero diagonal, no NaN/Inf.
`--tnt` output still loads into TNT with no triangle-inequality warning
(`scaled_metric_matrix`'s rounding-metric-closure step still needed and
still applied — nothing in the new construction guarantees the *rounded
integer* matrix stays a metric, even though the real-valued one is exact
by construction here, unlike the old Floyd-Warshall-closed matrix which
needed it for a different reason). IQ-TREE loads the new model file and
reports the same `π` as before (only `R` changed).
### Two consistency bugs found and fixed post-implementation (2026-08-13)
**`--exclude-genome` didn't drop columns that become monomorphic once the
excluded genome(s) are gone.** `snp_pseudo_alignment`'s "variable family"
test (`family_size() >= 2`) is a property of the annex computed over
*every* genome in the index — unaffected by the CLI-level exclusion, which
only dropped the excluded genome's *row*. A family variable only because
of the excluded genome stayed in the alignment as a now-constant column —
silently wrong data for TNT/PhyG, and a hard failure for IQ-TREE's `+ASC`
(verified: excluding 2 taxa on the benchmark left 116,351 such columns —
matches the manual `+ASC` failures hit earlier in this same investigation,
before `--exclude-genome` existed). Fixed in `drop_excluded`
(`obikmer/src/cmd/distance/mod.rs`): after dropping excluded rows,
re-scan each column among the *surviving* sequences and drop any that are
now constant. Verified: 908,723 → 792,372 sites after excluding 2 taxa,
zero monomorphic columns remain, `π` recomputed from the corrected
alignment matches an independent recount exactly. The compact-alphabet
renumbering (`iqtree::compact_alphabet`) needed no equivalent fix — it
already recomputes which of the 16 states occur fresh on every call, from
whatever alignment it's actually handed, so a symbol disappearing (e.g.
excluding every genome that carries `N`) is already handled correctly;
verified directly (excluded 13 genomes to force `N` out: "15 of 16 states"
reported, correctly-shaped model file).
**`cardinality_tally`'s `family_size() >= 2` filter looked inconsistent
with `base_pair_tally` — removing it was tried, and was wrong; reverted.**
`cardinality_tally` (modelled after `snp_pseudo_alignment`) had the
filter; `base_pair_tally` didn't (it visits every family via
`scan_family_pairs` unconditionally, folding fully-invariant loci into its
own `same` diagonal). Read as `cardinality_tally` under-counting its
diagonal relative to `base_pair_tally`, and — independently — as another
angle on the `--exclude-genome` drift (`family_size()` being global-only
meant a family kept here post-exclusion could differ from what the
now-correctly-filtered alignment kept). First fix tried: drop
`cardinality_tally`'s filter entirely, matching `base_pair_tally`'s
whole-annex scope.
**That fix was empirically wrong, confirmed by a real IQ-TREE run, not
just a hunch.** Log-likelihood dropped from the earlier correct run's
`-8,364,671`/`-8,371,082` to `-9,170,228` (worse fit, not better), with
repeated `NNI search needs unusual large number of steps (20) to
converge!` warnings — and the completed run's **total tree length came
out at 67.644**, roughly 30× the earlier correct runs' ~2.0, i.e. branches
blowing up/saturating. Root cause, only clear in hindsight: `+ASC`
("ascertainment bias correction") exists specifically because the
likelihood only ever sees *variable* sites — the alignment fed to
IQ-TREE, by construction, contains not one invariant column. Calibrating
`R` from a population overwhelmingly dominated by genome-wide invariant
background (family_size()<2 loci outnumber the ~908k variable ones by
orders of magnitude) describes a completely different population than the
one `+ASC` and the alignment actually model — the "consistency" argument
for matching `base_pair_tally`'s scope was real, but pointed the wrong
way: `base_pair_tally`'s own unrestricted `same` diagonal turned out to
have the *identical* latent bug (only unmasked once its diagonal existed
at all, which happened earlier the same day when `same` was added), not a
correct baseline to match `cardinality_tally` to.
**Final fix**: restored `cardinality_tally`'s `family_size() >= 2` filter,
and gave `base_pair_tally`'s `same` diagonal the equivalent restriction —
`scan_family_pairs` (shared with `raw_snp_distance`, which legitimately
*does* want fully-invariant families counted as `shared`) now passes an
extra `variable: bool` (the family's own `family_size() >= 2`) to its
`on_pair` callback; `base_pair_tally` only increments `same` when
`variable` is true, `raw_snp_distance`'s callback ignores the new
argument. Both tallies now describe the same variable-families-only
population the `+ASC`-corrected alignment does. Verified: calibration
counts back to their original values exactly (`c=0/c=0`: 117,158,166,
matching the pre-regression run bit for bit), and a full IQ-TREE rerun
converged normally — log-likelihood `-8,389,106.273` (same order as the
two earlier correct runs), **total tree length 2.044** (was 67.644), no
NNI convergence warnings.
## Heterozygosity, ploidy, and consensus-assembly inputs
+16
View File
@@ -0,0 +1,16 @@
0.222641
0.222347 0.005197
0.206041 0.017442 0.019145
0.222349 0.024185 0.003581 0.013346
0.208691 0.017667 0.002921 0.000012 0.019168
0.205784 0.013974 0.019121 0.000085 0.018901 0.000018
0.027849 0.001121 0.001230 0.000109 0.001216 0.000109 0.000118
0.222640 0.004745 0.023685 0.013831 0.005020 0.002986 0.013813 0.000889
0.206337 0.017467 0.012967 0.000076 0.013365 0.000016 0.000059 0.000079 0.017590
0.208662 0.003028 0.019388 0.000017 0.002764 0.000003 0.000018 0.000019 0.017788 0.000017
0.027866 0.001122 0.001231 0.000109 0.000858 0.000016 0.000083 0.000009 0.001129 0.000100 0.000111
0.206081 0.013994 0.012951 0.000059 0.018929 0.000017 0.000082 0.000085 0.017568 0.000078 0.000012 0.000076
0.027872 0.001122 0.000833 0.000074 0.001217 0.000109 0.000080 0.000042 0.001130 0.000100 0.000017 0.000006 0.000108
0.027846 0.000899 0.001230 0.000087 0.001216 0.000019 0.000118 0.000009 0.001129 0.000080 0.000111 0.000044 0.000108 0.000009
0.009724 0.000172 0.000189 0.000044 0.000187 0.000044 0.000048 0.000047 0.000174 0.000041 0.000045 0.000043 0.000044 0.000043 0.000047
0.904811 0.012806 0.013642 0.007640 0.013786 0.007933 0.004196 0.000140 0.012382 0.006685 0.007847 0.000195 0.007527 0.000202 0.000134 0.000076
+160 -20
View File
@@ -1,27 +1,167 @@
ratio_ceiling: 0.5
flank_length_m: 15
p_hat: 0.010713940022064144
p_hat_variance: 3.2110559785689563e-10
n_pairs_included: 158
n_loci_total: 33008305
mean_sub_cost: 1.4764371767846869
c_ctx: 1.7343663443925388
substitutions:
- pair: A/C
cardinality_transitions:
- from: 0
to: 0
count: 117158166
probability: 0.8249649657257814
- from: 0
to: 1
count: 13707223
probability: 0.09651891232567299
- from: 0
to: 2
count: 10934100
probability: 0.07699206755884405
- from: 0
to: 3
count: 193315
probability: 0.0013612205430842902
- from: 0
to: 4
count: 23125
probability: 0.00016283384661730445
- from: 1
to: 0
count: 13707223
probability: 0.8980473700324103
- from: 1
to: 1
count: 1043692
probability: 0.06837890181868832
- from: 1
to: 2
count: 507953
probability: 0.033279232106318904
- from: 1
to: 3
count: 4270
probability: 0.0002797548613631216
- from: 1
to: 4
count: 225
probability: 0.000014741181219368235
- from: 2
to: 0
count: 10934100
probability: 0.9551033885158036
- from: 2
to: 1
count: 507953
probability: 0.04437014765794788
- from: 2
to: 2
count: 5241
probability: 0.00045780602511512846
- from: 2
to: 3
count: 690
probability: 0.000060272115498843476
- from: 2
to: 4
count: 96
probability: 8.385685634621701e-6
- from: 3
to: 0
count: 193315
probability: 0.9743748708410829
- from: 3
to: 1
count: 4270
probability: 0.021522285898618442
- from: 3
to: 2
count: 690
probability: 0.0034778401100812
- from: 3
to: 3
count: 98
probability: 0.0004939541025912429
- from: 3
to: 4
count: 26
probability: 0.0001310490476262481
- from: 4
to: 0
count: 23125
probability: 0.9846291407647109
- from: 4
to: 1
count: 225
probability: 0.009580175423656646
- from: 4
to: 2
count: 96
probability: 0.004087541514093502
- from: 4
to: 3
count: 26
probability: 0.0011070424934003236
- from: 4
to: 4
count: 14
probability: 0.0005960998041386358
composition_transitions:
- from: 'A'
to: 'A'
count: 162260
probability: 0.47062167539692207
- from: 'A'
to: 'C'
count: 27810
cost: 2.5439415226130238
- pair: A/G
probability: 0.08066059899413536
- from: 'A'
to: 'G'
count: 130153
cost: 1.0
- pair: A/T
probability: 0.3774979842101294
- from: 'A'
to: 'T'
count: 24555
cost: 2.6684722241428322
- pair: C/G
probability: 0.07121974139881315
- from: 'C'
to: 'A'
count: 27810
probability: 0.07790832534920075
- from: 'C'
to: 'C'
count: 184633
probability: 0.5172401234879174
- from: 'C'
to: 'G'
count: 19637
cost: 2.892062911764701
- pair: C/T
probability: 0.05501207424963161
- from: 'C'
to: 'T'
count: 124878
cost: 1.0413902163542994
- pair: G/T
probability: 0.3498394769132503
- from: 'G'
to: 'A'
count: 130153
probability: 0.36057158212891627
- from: 'G'
to: 'C'
count: 19637
probability: 0.054401697680925745
- from: 'G'
to: 'G'
count: 184557
probability: 0.5112906308956874
- from: 'G'
to: 'T'
count: 26616
cost: 2.5878424665170052
probability: 0.07373608929447062
- from: 'T'
to: 'A'
count: 24555
probability: 0.07337692220342934
- from: 'T'
to: 'C'
count: 124878
probability: 0.3731689387464813
- from: 'T'
to: 'G'
count: 26616
probability: 0.07953574267426085
- from: 'T'
to: 'T'
count: 158593
probability: 0.4739183963758285
+242
View File
@@ -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]);
}
}
}
}
+3 -6
View File
@@ -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};
-402
View File
@@ -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
View File
@@ -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
+61 -34
View File
@@ -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}");
+54 -58
View File
@@ -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 {