feat: add Sankoff parsimony model and directory locking
Implements a calibrated 16-state Sankoff substitution cost matrix and CLI pipeline for evolutionary distance computation, including empirical calibration via saturation-filtered SNP counts. Refactors the sibling scanning stage to use batched transforms for improved synchronization efficiency. Introduces an OS-level advisory directory lock across all index-modifying commands to prevent concurrent write corruption. Updates dependencies and exposes new Sankoff utilities in the public API.
This commit is contained in:
Eric Coissac
Generated
+136
-18
@@ -48,6 +48,15 @@ dependencies = [
|
||||
"as-slice",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aligned-vec"
|
||||
version = "0.6.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b"
|
||||
dependencies = [
|
||||
"equator",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "allocator-api2"
|
||||
version = "0.2.21"
|
||||
@@ -497,6 +506,18 @@ dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "console"
|
||||
version = "0.16.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c"
|
||||
dependencies = [
|
||||
"encode_unicode",
|
||||
"libc",
|
||||
"unicode-width",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.10.2"
|
||||
@@ -820,6 +841,26 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equator"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc"
|
||||
dependencies = [
|
||||
"equator-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equator-macro"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
@@ -907,6 +948,16 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs4"
|
||||
version = "0.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e8c6b3bd49c37d2aa3f3f2220233b29a7cd23f79d1fe70e5337d25fb390793de"
|
||||
dependencies = [
|
||||
"rustix 0.38.44",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs_at"
|
||||
version = "0.2.1"
|
||||
@@ -1228,13 +1279,26 @@ version = "0.17.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
|
||||
dependencies = [
|
||||
"console",
|
||||
"console 0.15.11",
|
||||
"number_prefix",
|
||||
"portable-atomic",
|
||||
"unicode-width",
|
||||
"web-time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indicatif"
|
||||
version = "0.18.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c"
|
||||
dependencies = [
|
||||
"console 0.16.4",
|
||||
"portable-atomic",
|
||||
"unicode-width",
|
||||
"unit-prefix",
|
||||
"web-time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "infer"
|
||||
version = "0.19.0"
|
||||
@@ -1326,12 +1390,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kodama"
|
||||
version = "0.2.3"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7a44f3a71a44fbf49ce38152db7dc9adf959d4fe5c29344cd1858bdbda8d9091"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
checksum = "f868feceed4703c842925c14232667c5f29c8059d0354154b2d06ec3bf9daf82"
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
@@ -1365,6 +1426,12 @@ dependencies = [
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.4.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.12.1"
|
||||
@@ -1695,7 +1762,7 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"crossbeam-channel",
|
||||
"hwlocality",
|
||||
"indicatif",
|
||||
"indicatif 0.17.11",
|
||||
"ndarray",
|
||||
"obicompactvec",
|
||||
"obikpartitionner",
|
||||
@@ -1720,7 +1787,7 @@ version = "1.1.44"
|
||||
dependencies = [
|
||||
"clap",
|
||||
"csv",
|
||||
"indicatif",
|
||||
"indicatif 0.18.6",
|
||||
"kodama",
|
||||
"obidebruinj",
|
||||
"obifastwrite",
|
||||
@@ -1737,7 +1804,9 @@ dependencies = [
|
||||
"obitaxonomy",
|
||||
"pprof",
|
||||
"rayon",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
"speedytree",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
@@ -1749,7 +1818,7 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"cacheline-ef",
|
||||
"epserde",
|
||||
"indicatif",
|
||||
"indicatif 0.17.11",
|
||||
"memmap2",
|
||||
"niffler 3.0.0",
|
||||
"obicompactvec",
|
||||
@@ -1853,7 +1922,7 @@ dependencies = [
|
||||
"memmap2",
|
||||
"niffler 3.0.0",
|
||||
"obikseq",
|
||||
"rustix",
|
||||
"rustix 1.1.4",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
@@ -1863,7 +1932,8 @@ dependencies = [
|
||||
name = "obisys"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"indicatif",
|
||||
"fs4",
|
||||
"indicatif 0.17.11",
|
||||
"libc",
|
||||
"sysinfo",
|
||||
"tracing",
|
||||
@@ -2017,10 +2087,11 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "pprof"
|
||||
version = "0.13.0"
|
||||
version = "0.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ef5c97c51bd34c7e742402e216abdeb44d415fbe6ae41d56b114723e953711cb"
|
||||
checksum = "38a01da47675efa7673b032bf8efd8214f1917d89685e07e395ab125ea42b187"
|
||||
dependencies = [
|
||||
"aligned-vec",
|
||||
"backtrace",
|
||||
"cfg-if",
|
||||
"findshlibs",
|
||||
@@ -2028,15 +2099,15 @@ dependencies = [
|
||||
"log",
|
||||
"nix 0.26.4",
|
||||
"once_cell",
|
||||
"parking_lot",
|
||||
"prost",
|
||||
"prost-build",
|
||||
"prost-derive",
|
||||
"sha2",
|
||||
"smallvec",
|
||||
"spin",
|
||||
"symbolic-demangle",
|
||||
"tempfile",
|
||||
"thiserror 1.0.69",
|
||||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2386,6 +2457,19 @@ dependencies = [
|
||||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "0.38.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys 0.4.15",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.4"
|
||||
@@ -2395,7 +2479,7 @@ dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"linux-raw-sys 0.12.1",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
@@ -2521,6 +2605,19 @@ dependencies = [
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_yaml"
|
||||
version = "0.9.34+deprecated"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde",
|
||||
"unsafe-libyaml",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
@@ -2588,6 +2685,15 @@ dependencies = [
|
||||
"rb_tree",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spin"
|
||||
version = "0.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3"
|
||||
dependencies = [
|
||||
"lock_api",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "stable_deref_trait"
|
||||
version = "1.2.1"
|
||||
@@ -2736,7 +2842,7 @@ dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.3.4",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"rustix 1.1.4",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
@@ -2920,6 +3026,18 @@ version = "0.2.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
||||
|
||||
[[package]]
|
||||
name = "unit-prefix"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3"
|
||||
|
||||
[[package]]
|
||||
name = "unsafe-libyaml"
|
||||
version = "0.2.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861"
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.9.0"
|
||||
@@ -3421,7 +3539,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"rustix",
|
||||
"rustix 1.1.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -8,6 +8,7 @@ mod merge;
|
||||
mod numa;
|
||||
mod rebuild;
|
||||
mod reindex;
|
||||
mod sankoff;
|
||||
mod select;
|
||||
mod siblings;
|
||||
mod stats;
|
||||
@@ -19,4 +20,8 @@ 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::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
||||
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,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
+207
-99
@@ -14,18 +14,25 @@
|
||||
//! mask by callers, not stored (see `FamilyMask` and
|
||||
//! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below).
|
||||
//!
|
||||
//! Per layer, an `obipipeline` `Flat` stage (throttled — see
|
||||
//! `obipipeline::throttle`) generates each k-mer's 3 central variants,
|
||||
//! interleaved across many in-flight k-mers by the scheduler's shared
|
||||
//! worker pool rather than processed on a single thread. The actual
|
||||
//! cross-partition lookup, though, reuses
|
||||
//! `KmerPartition::query_partition_with` — the same partition-batching
|
||||
//! mechanism `obikmer query` already uses (open a partition's files once,
|
||||
//! answer a whole batch of queries against it) — rather than a per-item
|
||||
//! pipeline stage: an earlier per-item design reopened/re-mmap'd every
|
||||
//! target partition's files on every single lookup, which was fine at
|
||||
//! toy scale but manifested as ~90% system time against a real index.
|
||||
//! See `docmd/theory/evolutionary_distances.md`, Step 2b, "Mechanism".
|
||||
//! Per layer, an `obipipeline` batch transform (throttled — see
|
||||
//! `obipipeline::throttle`) generates a whole batch's central variants at
|
||||
//! once (`BATCH_SIZE` source k-mers in, that batch's variants out as one
|
||||
//! pipeline message), interleaved across many in-flight batches by the
|
||||
//! scheduler's shared worker pool rather than processed on a single
|
||||
//! thread. The actual cross-partition lookup reuses a `PartitionCache` of
|
||||
//! every partition's already-open MPHF layers, built once for the whole
|
||||
//! `build_sibling_annex` run, rather than reopening files per lookup or
|
||||
//! per source layer. Two earlier, coarser-grained designs were tried and
|
||||
//! measured (not guessed) to be worse, in order: (1) reopening/re-mmap'ing
|
||||
//! every target partition's files on every single lookup — fine at toy
|
||||
//! scale, ~90% system time against a real index; (2) a `Flat` pipeline
|
||||
//! stage pushing one message per generated *variant* (up to 3 per source
|
||||
//! k-mer) — cheaper than reopening files, but sampling a real run showed
|
||||
//! most wall-clock time going into per-message channel send/notify
|
||||
//! syscalls rather than the lookup itself, because a single k-mer's ≤3
|
||||
//! variants is far too fine a granularity to amortise a pipeline's
|
||||
//! synchronisation cost over. See `docmd/theory/evolutionary_distances.md`,
|
||||
//! Step 2b, "Mechanism".
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicU8, Ordering};
|
||||
@@ -107,34 +114,30 @@ fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
|
||||
|
||||
// ── obipipeline data types ─────────────────────────────────────────────────
|
||||
|
||||
/// One distinct k-mer of the layer currently being processed, at its local
|
||||
/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared,
|
||||
/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this
|
||||
/// one fans out into via the `Flat` stage) that is only released once every
|
||||
/// one of those descendants has been fully processed — see the module docs'
|
||||
/// "Throttling" note for why this is required, not optional, once a `Flat`
|
||||
/// stage is in the pipeline.
|
||||
struct SourceItem {
|
||||
slot: usize,
|
||||
kmer: CanonicalKmer,
|
||||
_permit: Arc<ThrottleGuard>,
|
||||
/// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the
|
||||
/// pipeline's source item — batched, not one k-mer per item, so that
|
||||
/// pipeline messages and their synchronisation cost stay amortised over
|
||||
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on
|
||||
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
|
||||
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
|
||||
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
|
||||
struct SourceBatch {
|
||||
items: Vec<(usize, CanonicalKmer)>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
/// One of a source k-mer's (up to 3) central-substitution variants, already
|
||||
/// routed to its destination partition and carrying its own central base
|
||||
/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of
|
||||
/// the source item's throttle permit.
|
||||
struct VariantQuery {
|
||||
source_slot: usize,
|
||||
dest_partition: usize,
|
||||
variant: CanonicalKmer,
|
||||
base: u8,
|
||||
_permit: Arc<ThrottleGuard>,
|
||||
/// One batch's worth of central-substitution variants (up to 3 per source
|
||||
/// k-mer), each already routed to its destination partition and carrying
|
||||
/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
|
||||
/// hit. `(dest_partition, variant, source_slot, base)` per entry.
|
||||
struct VariantBatch {
|
||||
items: Vec<(usize, CanonicalKmer, usize, u8)>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
enum SibData {
|
||||
Item(SourceItem),
|
||||
Query(VariantQuery),
|
||||
Batch(SourceBatch),
|
||||
Variants(VariantBatch),
|
||||
}
|
||||
|
||||
/// Every partition's already-open MPHF layers, built **once** for the whole
|
||||
@@ -341,57 +344,72 @@ impl KmerIndex {
|
||||
mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// ── obipipeline: Flat stage generates variants only — the actual
|
||||
// cross-partition lookup reuses `KmerPartition::query_partition_with`
|
||||
// (the same batching mechanism `obikmer query` already uses: open a
|
||||
// partition's files once, answer a whole batch of queries against
|
||||
// it) instead of one lookup per pipeline item. A per-item lookup
|
||||
// (tried first) reopened/re-mmap'd every target partition's files on
|
||||
// every single variant — fine at the scale of a handful of test
|
||||
// k-mers, but with billions of lookups against a real index this
|
||||
// manifested as ~90% system time, observed in practice. ───────────
|
||||
// ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
|
||||
// the actual cross-partition lookup reuses
|
||||
// `KmerPartition::query_partition_with` (the same batching mechanism
|
||||
// `obikmer query` already uses: open a partition's files once,
|
||||
// answer a whole batch of queries against it) instead of one lookup
|
||||
// per pipeline item. A per-item lookup (tried first) reopened/
|
||||
// re-mmap'd every target partition's files on every single variant
|
||||
// — fine at toy scale, but ~90% system time against a real index,
|
||||
// observed in practice. A *later* attempt still pushed one pipeline
|
||||
// message per generated variant (a `Flat` stage, `SourceItem` =>
|
||||
// `VariantQuery`, one k-mer in => up to 3 variants out as separate
|
||||
// messages) — cheaper than reopening files, but sampling a real run
|
||||
// showed most wall-clock time going into per-message channel
|
||||
// send/notify syscalls instead of the lookup itself: the pipeline's
|
||||
// whole point is amortising synchronisation over a batch, and a
|
||||
// single k-mer's ≤3 variants is far too fine a granularity for
|
||||
// that. Batching `BATCH_SIZE` source k-mers into one pipeline item
|
||||
// — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of
|
||||
// variants out, one message either way — keeps the per-message
|
||||
// synchronisation cost amortised over thousands of lookups instead
|
||||
// of one to three. ──────────────────────────────────────────────
|
||||
const BATCH_SIZE: usize = 4096;
|
||||
let n_workers = obisys::effective_parallelism();
|
||||
let capacity = 256;
|
||||
|
||||
// Throttling is not optional once a `Flat` stage is in the pipeline
|
||||
// (see `obipipeline::throttle`'s docs): without it, every worker can
|
||||
// become a simultaneous `Flat` producer, saturate the shared output
|
||||
// channel, and deadlock against the scheduler's own dispatch loop —
|
||||
// also observed in practice. The permit acquired here for a source
|
||||
// k-mer is held (via the `Arc`-shared guard carried through
|
||||
// `SourceItem` -> `VariantQuery`) until every one of its (up to 3)
|
||||
// descendants has been read out of the pipeline by the accumulation
|
||||
// loop below, not just until the `Flat` stage itself returns.
|
||||
// Throttling limits how many *batches* are in flight at once — the
|
||||
// permit is acquired per batch (not per k-mer) in the source
|
||||
// thread, and released once its `VariantBatch` has been read out of
|
||||
// the pipeline by the accumulation loop below. See
|
||||
// `obipipeline::throttle`'s docs for why this is required, not
|
||||
// optional, once a `Flat`-style stage sits in the pipeline.
|
||||
let sources: Vec<(usize, CanonicalKmer)> = slot_kmer
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
|
||||
.collect();
|
||||
let throttled = obipipeline::throttle(sources.into_iter(), n_workers).map(|t| SourceItem {
|
||||
slot: t.item.0,
|
||||
kmer: t.item.1,
|
||||
_permit: Arc::new(t.guard),
|
||||
let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources
|
||||
.chunks(BATCH_SIZE)
|
||||
.map(|chunk| chunk.to_vec())
|
||||
.collect();
|
||||
let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
|
||||
items: t.item,
|
||||
_permit: t.guard,
|
||||
});
|
||||
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
SibData : SourceItem => VariantQuery,
|
||||
|| {
|
||||
move |item: SourceItem| -> Vec<VariantQuery> {
|
||||
let kmer = item.kmer;
|
||||
let permit = item._permit;
|
||||
kmer.central_canonical_neighbors()
|
||||
.into_iter()
|
||||
.filter(|variant| *variant != kmer)
|
||||
.map(|variant| VariantQuery {
|
||||
source_slot: item.slot,
|
||||
dest_partition: partition_of(variant, n_parts),
|
||||
variant,
|
||||
base: central_base(variant, k),
|
||||
_permit: Arc::clone(&permit),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
SibData : SourceBatch => VariantBatch,
|
||||
| {
|
||||
move |batch: SourceBatch| -> VariantBatch {
|
||||
let mut items = Vec::with_capacity(batch.items.len() * 3);
|
||||
for (slot, kmer) in batch.items {
|
||||
for variant in kmer.central_canonical_neighbors() {
|
||||
if variant == kmer {
|
||||
continue;
|
||||
}
|
||||
items.push((
|
||||
partition_of(variant, n_parts),
|
||||
variant,
|
||||
slot,
|
||||
central_base(variant, k),
|
||||
));
|
||||
}
|
||||
}
|
||||
VariantBatch { items, _permit: batch._permit }
|
||||
}
|
||||
} : Item => Query,
|
||||
} : Batch => Variants,
|
||||
};
|
||||
|
||||
// ── Group generated variants by destination partition. `cache`
|
||||
@@ -400,11 +418,13 @@ impl KmerIndex {
|
||||
// lookup is not free just because the file isn't reopened. Grouping
|
||||
// keeps one partition's pages hot while its whole batch is resolved,
|
||||
// instead of faulting pages in and out as lookups jump between
|
||||
// partitions in whatever order the `Flat` stage happens to produce
|
||||
// them. The throttle permit drops here, once accumulated. ─────────
|
||||
// partitions in whatever order the pipeline happens to produce
|
||||
// them. Each batch's throttle permit drops here, once accumulated.
|
||||
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
|
||||
for vq in pipe.apply(throttled, n_workers, capacity) {
|
||||
outgoing[vq.dest_partition].push((vq.variant, vq.source_slot, vq.base));
|
||||
for vb in pipe.apply(throttled, n_workers, capacity) {
|
||||
for (dest_partition, variant, source_slot, base) in vb.items {
|
||||
outgoing[dest_partition].push((variant, source_slot, base));
|
||||
}
|
||||
}
|
||||
|
||||
// ── Resolve each partition's batch against the cache in one
|
||||
@@ -635,9 +655,27 @@ pub struct RawSnpDistanceOutput {
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
|
||||
/// (run [`build_sibling_annex`](Self::build_sibling_annex) first).
|
||||
pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
|
||||
/// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance)
|
||||
/// and [`base_pair_tally`](Self::base_pair_tally): for every family
|
||||
/// (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`.
|
||||
fn scan_family_pairs<Acc, F, C>(
|
||||
&self,
|
||||
label: &str,
|
||||
zero: impl Fn() -> Acc + Sync,
|
||||
on_pair: F,
|
||||
combine: C,
|
||||
) -> OKIResult<Acc>
|
||||
where
|
||||
Acc: Send,
|
||||
F: Fn(&mut Acc, usize, usize, u8, u8) + Sync,
|
||||
C: Fn(Acc, Acc) -> Acc,
|
||||
{
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
@@ -673,12 +711,11 @@ impl KmerIndex {
|
||||
}
|
||||
}
|
||||
|
||||
let pb = progress_bar("raw_snp_distance", layer_dirs.len() as u64, "layers");
|
||||
let partials: Vec<(Array2<u64>, Array2<u64>)> = layer_dirs
|
||||
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
|
||||
let partials: Vec<Acc> = layer_dirs
|
||||
.par_iter()
|
||||
.map(|layer_dir| -> OKIResult<(Array2<u64>, Array2<u64>)> {
|
||||
let mut snp = Array2::<u64>::zeros((n_genomes, n_genomes));
|
||||
let mut shared = Array2::<u64>::zeros((n_genomes, n_genomes));
|
||||
.map(|layer_dir| -> OKIResult<Acc> {
|
||||
let mut acc = zero();
|
||||
|
||||
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)?;
|
||||
@@ -754,31 +791,102 @@ impl KmerIndex {
|
||||
continue;
|
||||
}
|
||||
let Some(bj) = single_form[j] else { continue };
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
shared[[j, i]] += 1;
|
||||
} else {
|
||||
snp[[i, j]] += 1;
|
||||
snp[[j, i]] += 1;
|
||||
}
|
||||
on_pair(&mut acc, i, j, bi, bj);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pb.inc(1);
|
||||
Ok((snp, shared))
|
||||
Ok(acc)
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
pb.finish_and_clear();
|
||||
|
||||
let mut snp = Array2::<u64>::zeros((n_genomes, n_genomes));
|
||||
let mut shared = Array2::<u64>::zeros((n_genomes, n_genomes));
|
||||
for (s, sh) in partials {
|
||||
snp += &s;
|
||||
shared += &sh;
|
||||
let mut total = zero();
|
||||
for partial in partials {
|
||||
total = combine(total, partial);
|
||||
}
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
|
||||
/// (run [`build_sibling_annex`](Self::build_sibling_annex) first).
|
||||
pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
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| {
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
shared[[j, i]] += 1;
|
||||
} else {
|
||||
snp[[i, j]] += 1;
|
||||
snp[[j, i]] += 1;
|
||||
}
|
||||
},
|
||||
|(mut snp, mut shared), (s, sh)| {
|
||||
snp += &s;
|
||||
shared += &sh;
|
||||
(snp, shared)
|
||||
},
|
||||
)?;
|
||||
Ok(RawSnpDistanceOutput { snp, shared })
|
||||
}
|
||||
|
||||
/// 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.
|
||||
///
|
||||
/// A second full pass over the annex, sharing
|
||||
/// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by
|
||||
/// it, not a blind re-scan) — needed because `raw_snp_distance` only
|
||||
/// keeps aggregate SNP/shared counts per genome pair, not which bases
|
||||
/// were actually involved at each locus, and the ratio-ceiling filter
|
||||
/// can only be evaluated once the aggregate counts are known.
|
||||
pub fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
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 counts = self.scan_family_pairs(
|
||||
"base_pair_tally",
|
||||
|| [[0u64; 4]; 4],
|
||||
|counts, i, j, bi, bj| {
|
||||
if bi != bj && included[[i, j]] {
|
||||
counts[bi as usize][bj as usize] += 1;
|
||||
counts[bj as usize][bi as usize] += 1;
|
||||
}
|
||||
},
|
||||
|mut total, partial| {
|
||||
for a in 0..4 {
|
||||
for b in 0..4 {
|
||||
total[a][b] += partial[a][b];
|
||||
}
|
||||
}
|
||||
total
|
||||
},
|
||||
)?;
|
||||
Ok(BasePairTally { counts })
|
||||
}
|
||||
}
|
||||
|
||||
/// See [`KmerIndex::base_pair_tally`].
|
||||
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).
|
||||
pub counts: [[u64; 4]; 4],
|
||||
}
|
||||
|
||||
/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
|
||||
|
||||
@@ -22,15 +22,17 @@ obikindex = { path = "../obikindex", default-features = false }
|
||||
obitaxonomy = { path = "../obitaxonomy" }
|
||||
obilayeredmap = { path = "../obilayeredmap" }
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
serde_yaml = "0.9.33"
|
||||
csv = "1"
|
||||
kodama = "0.2"
|
||||
kodama = "0.3.0"
|
||||
speedytree = "0.1"
|
||||
rayon = "1"
|
||||
indicatif = "0.17"
|
||||
indicatif = "0.18"
|
||||
tracing = "0.1.44"
|
||||
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
|
||||
pprof = { version = "0.13", features = ["prost-codec"], optional = true }
|
||||
pprof = { version = "0.15", features = ["prost-codec"], optional = true }
|
||||
|
||||
[features]
|
||||
default = ["numa"]
|
||||
|
||||
+495
-10
@@ -4,7 +4,11 @@ use std::path::PathBuf;
|
||||
use clap::Args;
|
||||
use kodama::{Method, linkage};
|
||||
use obifastwrite::{JsonVal, write_record};
|
||||
use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
||||
use obikindex::{
|
||||
BasePairTally, DistanceMetric, KmerIndex, PHatEstimate, RawSnpDistanceOutput, SankoffWeights,
|
||||
SiblingAnnexStats, SnpAlignment, build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat,
|
||||
mean_substitution_cost, substitution_costs_from_tally,
|
||||
};
|
||||
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
|
||||
use tracing::info;
|
||||
|
||||
@@ -94,9 +98,61 @@ pub struct DistanceArgs {
|
||||
#[arg(long)]
|
||||
pub snp: bool,
|
||||
|
||||
/// Calibrate a 16-state Sankoff cost matrix and its matching
|
||||
/// pseudo-alignment from an already-built sibling annex (run with
|
||||
/// `--sibling-annex` first, in this invocation or an earlier one), for
|
||||
/// use with TNT/PhyG. See `docmd/theory/evolutionary_distances.md`,
|
||||
/// "Sankoff parsimony as the resolution of the 16-state model problem".
|
||||
#[arg(long)]
|
||||
pub sankoff: bool,
|
||||
|
||||
/// Exclude genome pairs whose raw SNP ratio exceeds this value from the
|
||||
/// `p_hat` calibration pooled by `--sankoff` — a pair this close to
|
||||
/// saturation carries no information about `p_hat` and would bias it
|
||||
/// upward if pooled in (unlike a low eligible-loci count, which barely
|
||||
/// moves the pooled estimate either way — see design doc).
|
||||
#[arg(long, default_value = "0.5")]
|
||||
pub sankoff_ratio_ceiling: f64,
|
||||
|
||||
/// Also write <prefix>_sankoff.tnt, a ready-to-run TNT script (`proc
|
||||
/// <file>;`) for the same matrix/alignment `--sankoff` computes —
|
||||
/// recoded to TNT's default xread alphabet (0-9A-F only; TNT rejects
|
||||
/// the wider IUPAC set `--sankoff`'s own output uses unless `nstates
|
||||
/// dna` is set, which imposes TNT's own incompatible DNA encoding
|
||||
/// instead) with integer-scaled costs (TNT's smatrix/cost commands
|
||||
/// reject decimals). Implies `--sankoff`.
|
||||
#[arg(long)]
|
||||
pub tnt: bool,
|
||||
|
||||
/// Also write <prefix>_sankoff.tcm and <prefix>_sankoff.pg, a
|
||||
/// custom-alphabet cost matrix and a ready-to-run PhyG script (`read`/
|
||||
/// `search`/`report`) for the same matrix/alignment `--sankoff`
|
||||
/// computes. Reuses `--sankoff`'s own `_sankoff.fasta` directly — PhyG's
|
||||
/// `tcm:` alphabet is read from the matrix file itself, so the IUPAC+`0`
|
||||
/// alphabet needs no recoding here, unlike `--tnt`. Implies `--sankoff`.
|
||||
#[arg(long)]
|
||||
pub phyg: bool,
|
||||
|
||||
/// Scale factor applied before rounding real-valued costs to the
|
||||
/// integers both `--tnt`'s smatrix/cost commands and `--phyg`'s `tcm:`
|
||||
/// matrix require. Keep this small: the total tree score is this scale
|
||||
/// times the sum of per-character costs across every character (908k+
|
||||
/// for a typical run here), and there are hints in TNT's own manual
|
||||
/// that at least some of its internal accumulators are 32-bit — a large
|
||||
/// scale risks a silent integer overflow (undetectable, not just a
|
||||
/// crash) far more costly than the resolution a bigger factor would
|
||||
/// buy. Shared between `--tnt` and `--phyg` rather than split into two
|
||||
/// flags: both scale the same calibrated matrix for the same reason
|
||||
/// (integer-only cost commands), and no PhyG-specific accumulator-width
|
||||
/// constraint has actually been found to justify a different default.
|
||||
#[arg(long, default_value = "100")]
|
||||
pub sankoff_cost_scale: f64,
|
||||
|
||||
/// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
|
||||
/// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta,
|
||||
/// <prefix>_nj.nwk, <prefix>_upgma.nwk.
|
||||
/// <prefix>_sankoff_matrix.csv, <prefix>_sankoff_params.yaml,
|
||||
/// <prefix>_sankoff.fasta, <prefix>_sankoff.tnt, <prefix>_sankoff.tcm,
|
||||
/// <prefix>_sankoff.pg, <prefix>_nj.nwk, <prefix>_upgma.nwk.
|
||||
/// If omitted, the distance matrix is written to stdout.
|
||||
#[arg(short, long)]
|
||||
pub output: Option<PathBuf>,
|
||||
@@ -117,6 +173,13 @@ pub fn run(args: DistanceArgs) {
|
||||
// deliberately decoupled: the annex is meant to be (re)built routinely,
|
||||
// the distribution only occasionally, on demand.
|
||||
if args.sibling_annex {
|
||||
// Writes into the index directory — hold an exclusive lock for the
|
||||
// duration so a second, concurrent `--sibling-annex` run on the same
|
||||
// index can't corrupt these writes (see obisys::DirLock).
|
||||
let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| {
|
||||
eprintln!("error locking index directory {}: {e}", args.index.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
info!("building sibling-count/minorant annex");
|
||||
idx.build_sibling_annex().unwrap_or_else(|e| {
|
||||
eprintln!("error building sibling annex: {e}");
|
||||
@@ -144,15 +207,56 @@ pub fn run(args: DistanceArgs) {
|
||||
});
|
||||
write_snp_fasta(&alignment, &labels, &args.output);
|
||||
}
|
||||
if args.sankoff || args.tnt || args.phyg {
|
||||
let raw = idx.raw_snp_distance().unwrap_or_else(|e| {
|
||||
eprintln!("error computing raw SNP distance: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
let estimate = calibrate_p_hat(&raw, args.sankoff_ratio_ceiling);
|
||||
let m = (idx.kmer_size() - 1) / 2;
|
||||
|
||||
// `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp` are
|
||||
// their own operation, not a modifier on top of a distance-metric
|
||||
// computation — a metric was never requested by asking for any of them,
|
||||
// so there is nothing for the rest of this function to compute. Not a
|
||||
// historical accident to keep: stop here rather than always also
|
||||
// running a Jaccard (or whichever `--metric` defaults to) pass and
|
||||
// printing an unrequested matrix.
|
||||
if args.sibling_annex || args.sibling_stats || args.raw_snp_distance || args.snp {
|
||||
let 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 alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP pseudo-alignment: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
write_sankoff_alignment_fasta(&alignment, &labels, &args.output);
|
||||
|
||||
if args.tnt {
|
||||
write_sankoff_tnt(&matrix, &alignment, &labels, &args.output, args.sankoff_cost_scale);
|
||||
}
|
||||
if args.phyg {
|
||||
write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale);
|
||||
}
|
||||
}
|
||||
|
||||
// `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp`/
|
||||
// `--sankoff`/`--tnt` are their own operation, not a modifier on top of
|
||||
// a distance-metric computation — a metric was never requested by
|
||||
// asking for any of them, so there is nothing for the rest of this
|
||||
// function to compute. Not a historical accident to keep: stop here
|
||||
// rather than always also running a Jaccard (or whichever `--metric`
|
||||
// defaults to) pass and printing an unrequested matrix.
|
||||
if args.sibling_annex
|
||||
|| args.sibling_stats
|
||||
|| args.raw_snp_distance
|
||||
|| args.snp
|
||||
|| args.sankoff
|
||||
|| args.tnt
|
||||
|| args.phyg
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -364,6 +468,387 @@ fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option<
|
||||
if n_sites == 1 { "" } else { "s" });
|
||||
}
|
||||
|
||||
// ── Sankoff pseudo-alignment → FASTA ────────────────────────────────────────
|
||||
//
|
||||
// Same data as `--snp`'s pseudo-alignment (`SnpAlignment`/
|
||||
// `snp_pseudo_alignment`), re-coded so its symbols match the accompanying
|
||||
// `--sankoff-matrix` output exactly: `0` for the empty/absent state instead
|
||||
// of `-`, which TNT/PhyG would otherwise read as their own gap character
|
||||
// rather than our "family absent" state.
|
||||
|
||||
fn write_sankoff_alignment_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option<PathBuf>) {
|
||||
let path = output.as_ref()
|
||||
.map(|p| format!("{}_sankoff.fasta", p.display()))
|
||||
.unwrap_or_else(|| "sankoff.fasta".into());
|
||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {path}: {e}");
|
||||
std::process::exit(1);
|
||||
}));
|
||||
let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
|
||||
for (label, seq) in labels.iter().zip(alignment.sequences.iter()) {
|
||||
let recoded: Vec<u8> = seq.iter().map(|&b| if b == b'-' { b'0' } else { b }).collect();
|
||||
write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| {
|
||||
eprintln!("error writing {path}: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
}
|
||||
info!("Sankoff pseudo-alignment → {path} ({n_sites} site{})",
|
||||
if n_sites == 1 { "" } else { "s" });
|
||||
}
|
||||
|
||||
// ── Sankoff cost matrix → CSV ────────────────────────────────────────────────
|
||||
//
|
||||
// 16 states indexed by bitmask (bit 0=A, 1=C, 2=G, 3=T; state 0 is `∅`),
|
||||
// matching the convention already used for `--snp`'s IUPAC-coded output and
|
||||
// for the external TNT/PhyG scripts this feeds. Calibration report (p_hat,
|
||||
// its variance, how many pairs/loci went into it, the resulting c_ctx) goes
|
||||
// to the log, not the CSV, since it's a run-level fact, not per-cell data.
|
||||
|
||||
// IUPAC ambiguity code per state (same mapping as `siblings::iupac_code`,
|
||||
// already used for `--snp`'s pseudo-alignment — a biologist reads "R" as
|
||||
// "A or G" without needing this file's convention explained), with `0`
|
||||
// standing in for the empty state (`-` would collide with TNT/PhyG's own
|
||||
// gap/range syntax). Bit order: 0=A, 1=C, 2=G, 3=T. This project's
|
||||
// canonical alphabet — `write_sankoff_tnt` below recodes it to TNT's own
|
||||
// default alphabet at the adapter boundary, rather than using it here.
|
||||
const STATE_SYMBOL: [char; 16] = [
|
||||
'0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N',
|
||||
];
|
||||
|
||||
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());
|
||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {path}: {e}");
|
||||
std::process::exit(1);
|
||||
}));
|
||||
write!(f, "state").unwrap();
|
||||
for sym in STATE_SYMBOL { write!(f, ",{sym}").unwrap(); }
|
||||
writeln!(f).unwrap();
|
||||
for (s, row) in matrix.iter().enumerate() {
|
||||
write!(f, "{}", STATE_SYMBOL[s]).unwrap();
|
||||
for cost in row { write!(f, ",{cost:.4}").unwrap(); }
|
||||
writeln!(f).unwrap();
|
||||
}
|
||||
info!("Sankoff cost matrix → {path}");
|
||||
}
|
||||
|
||||
// ── Sankoff calibration parameters → YAML report ────────────────────────────
|
||||
//
|
||||
// Everything `--sankoff` estimates from real data, in one durable,
|
||||
// machine-readable file: `p_hat` and its variance (with how many pairs/loci
|
||||
// went into it), the derived `c_ctx`, and the base-pair substitution tally
|
||||
// (raw counts, not just the derived costs) — kept for the same reason raw
|
||||
// counts are kept anywhere else in this project: costs are a modelling
|
||||
// choice built *from* the counts, and reproducing/re-deriving them later
|
||||
// needs the counts, not just their current derived value. Structured (YAML,
|
||||
// not an ad hoc key=value text file) so R/Python/etc. can load it directly
|
||||
// rather than re-parsing free text.
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct SankoffSubstitution {
|
||||
pair: String,
|
||||
count: u64,
|
||||
cost: 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>,
|
||||
}
|
||||
|
||||
fn write_sankoff_params(
|
||||
estimate: &PHatEstimate,
|
||||
tally: &BasePairTally,
|
||||
weights: &SankoffWeights,
|
||||
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 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 path = output.as_ref()
|
||||
.map(|p| format!("{}_sankoff_params.yaml", p.display()))
|
||||
.unwrap_or_else(|| "sankoff_params.yaml".into());
|
||||
let f = std::fs::File::create(&path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {path}: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
serde_yaml::to_writer(f, &report).unwrap_or_else(|e| {
|
||||
eprintln!("error writing {path}: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
info!("Sankoff calibration parameters → {path}");
|
||||
}
|
||||
|
||||
// ── Sankoff cost matrix + alignment → ready-to-run TNT script ──────────────
|
||||
//
|
||||
// TNT's *default* xread reader only accepts its own 0-9A-F alphabet (see
|
||||
// its manual: "up to 16 states are allowed by xread, using symbols 0-9 ...
|
||||
// and A-F") — the wider IUPAC set `STATE_SYMBOL` uses is rejected as an
|
||||
// "alien symbol" unless `nstates dna` is set, which imposes TNT's own fixed
|
||||
// DNA encoding instead, incompatible with a custom smatrix. And TNT's
|
||||
// `smatrix`/`cost` commands reject decimal costs ("found symbol . when
|
||||
// reading transformation costs"). So: recode to TNT's alphabet and
|
||||
// integer-scale the costs here, at this adapter's boundary, rather than
|
||||
// degrading the project's own canonical (IUPAC, real-valued) output.
|
||||
|
||||
const TNT_STATE_SYMBOL: [char; 16] = [
|
||||
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F',
|
||||
];
|
||||
|
||||
fn write_sankoff_tnt(
|
||||
matrix: &[[f64; 16]; 16],
|
||||
alignment: &SnpAlignment,
|
||||
labels: &[String],
|
||||
output: &Option<PathBuf>,
|
||||
cost_scale: f64,
|
||||
) {
|
||||
let path = output.as_ref()
|
||||
.map(|p| format!("{}_sankoff.tnt", p.display()))
|
||||
.unwrap_or_else(|| "sankoff.tnt".into());
|
||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {path}: {e}");
|
||||
std::process::exit(1);
|
||||
}));
|
||||
|
||||
// IUPAC-ish symbol -> bitmask, to translate the alignment (which uses
|
||||
// `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment`
|
||||
// callers) into TNT's alphabet without re-deriving state indices.
|
||||
let mut iupac_to_state = [0u8; 128];
|
||||
for (state, &sym) in STATE_SYMBOL.iter().enumerate() {
|
||||
iupac_to_state[sym as usize] = state as u8;
|
||||
}
|
||||
|
||||
let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
|
||||
writeln!(f, "xread").unwrap();
|
||||
writeln!(f, "'obikmer central-position SNP families, calibrated Sankoff 16-state encoding'").unwrap();
|
||||
writeln!(f, "{n_sites} {}", labels.len()).unwrap();
|
||||
for (label, seq) in labels.iter().zip(alignment.sequences.iter()) {
|
||||
write!(f, "{label} ").unwrap();
|
||||
for &b in seq {
|
||||
let b = if b == b'-' { b'0' } else { b };
|
||||
let state = iupac_to_state[b as usize];
|
||||
write!(f, "{}", TNT_STATE_SYMBOL[state as usize]).unwrap();
|
||||
}
|
||||
writeln!(f).unwrap();
|
||||
}
|
||||
writeln!(f, ";\n").unwrap();
|
||||
|
||||
let scaled_matrix = scaled_metric_matrix(matrix, cost_scale);
|
||||
|
||||
writeln!(f, "smatrix =0 (family16)").unwrap();
|
||||
for i in 0..16 {
|
||||
for j in (i + 1)..16 {
|
||||
writeln!(f, "{}/{} {}", TNT_STATE_SYMBOL[i], TNT_STATE_SYMBOL[j], scaled_matrix[i][j]).unwrap();
|
||||
}
|
||||
}
|
||||
writeln!(f, ";\n").unwrap();
|
||||
|
||||
writeln!(f, "ccode ( 0.{} ;", n_sites - 1).unwrap();
|
||||
writeln!(f, "smatrix +0 0.{} ;", n_sites - 1).unwrap();
|
||||
writeln!(f).unwrap();
|
||||
|
||||
// Basename only (not the full `path`/`output` prefix): TNT's natural
|
||||
// workflow is to `cd` into the output directory before `proc`-ing the
|
||||
// script, and an absolute path here would break if that directory is
|
||||
// later moved or copied elsewhere.
|
||||
let tre_name = output.as_ref()
|
||||
.and_then(|p| p.file_name())
|
||||
.map(|n| format!("{}_sankoff.tre", n.to_string_lossy()))
|
||||
.unwrap_or_else(|| "sankoff.tre".into());
|
||||
|
||||
// TNT's plain command parser has no comment syntax of its own — `/* */`
|
||||
// and `[ ]` are only recognised inside the (separately-enabled) macro
|
||||
// scripting language, and fail with "No command!" here otherwise
|
||||
// (verified against this file with the local TNT binary). `quote` is
|
||||
// the closest working equivalent: it prints free text and does not
|
||||
// otherwise affect parsing, so it doubles as an explanation of the
|
||||
// defaults below when the script is run. `;` ends a `quote` block like
|
||||
// any other TNT command, so the text itself must avoid semicolons.
|
||||
writeln!(f, "quote").unwrap();
|
||||
writeln!(f, "Default search below (edit or delete this block to run your own strategy):").unwrap();
|
||||
writeln!(f, " hold N : size of TNT's tree buffer (how many equally-parsimonious").unwrap();
|
||||
writeln!(f, " trees it keeps in memory at once), 20 is a small, fast").unwrap();
|
||||
writeln!(f, " default, raise it if mult reports it had to drop trees.").unwrap();
|
||||
writeln!(f, " mult : traditional search (random addition sequences followed by").unwrap();
|
||||
writeln!(f, " TBR branch-swapping, TNT's own default replication count),").unwrap();
|
||||
writeln!(f, " a reasonable first-pass strategy on this data's memory").unwrap();
|
||||
writeln!(f, " footprint, xmult's ratchet/drift/tree-fusion buffers ran").unwrap();
|
||||
writeln!(f, " this out of RAM at TNT's default mxram on this dataset.").unwrap();
|
||||
writeln!(f, " export - F : write the trees held in the buffer to file F, in").unwrap();
|
||||
writeln!(f, " TNT/Hennig86 format ('-' means trees, as opposed to data).").unwrap();
|
||||
writeln!(f, ";").unwrap();
|
||||
writeln!(f, "hold 20;").unwrap();
|
||||
writeln!(f, "mult;").unwrap();
|
||||
writeln!(f, "export - {tre_name};").unwrap();
|
||||
|
||||
info!(
|
||||
"TNT script → {path} (costs scaled x{cost_scale:.0}, runs a default `hold 20; mult;` \
|
||||
search and writes trees to {tre_name} in TNT's working directory — edit the trailing \
|
||||
comment block in the script to change this)\n\
|
||||
Run it with:\n \
|
||||
printf 'proc {path};\\nquit;\\n' | tnt\n\
|
||||
(or start `tnt` interactively and type `proc {path};`)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Scale `matrix` by `cost_scale` and round to integers (TNT's smatrix/cost
|
||||
/// commands reject decimals), then take the *metric 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 `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
|
||||
/// TNT-version-dependent.
|
||||
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 {
|
||||
for j in 0..16 {
|
||||
m[i][j] = (matrix[i][j] * cost_scale).round() as i64;
|
||||
}
|
||||
}
|
||||
for k in 0..16 {
|
||||
for i in 0..16 {
|
||||
for j in 0..16 {
|
||||
let via = m[i][k] + m[k][j];
|
||||
if via < m[i][j] {
|
||||
m[i][j] = via;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
m
|
||||
}
|
||||
|
||||
// ── Sankoff cost matrix → PhyG custom-alphabet TCM + ready-to-run script ────
|
||||
//
|
||||
// PhyG's `tcm:STRING` format needs no alphabet recoding, unlike `--tnt`:
|
||||
// its parser reads the alphabet straight from the tcm file's own first
|
||||
// line, so `--sankoff`'s own `_sankoff.fasta` (already IUPAC+`0`) is reused
|
||||
// as-is via `prefasta:`. PhyG auto-adds its own indel/gap state as an
|
||||
// (n+1)-th row/column of the tcm — inert here since the alignment already
|
||||
// encodes absence as an ordinary state (`0`), never as `-` (see
|
||||
// `write_sankoff_alignment_fasta`'s own comment on why, and the RAxML-era
|
||||
// bug that motivated it). The gap row/column below reuses `matrix[i][0]`/
|
||||
// `matrix[0][j]` (cost to/from `∅`) as the closest principled value for a
|
||||
// state that, in practice, is never actually triggered.
|
||||
|
||||
fn write_sankoff_phyg(matrix: &[[f64; 16]; 16], output: &Option<PathBuf>, cost_scale: f64) {
|
||||
let scaled_matrix = scaled_metric_matrix(matrix, cost_scale);
|
||||
|
||||
let basename = |suffix: &str| -> String {
|
||||
output.as_ref()
|
||||
.and_then(|p| p.file_name())
|
||||
.map(|n| format!("{}{suffix}", n.to_string_lossy()))
|
||||
.unwrap_or_else(|| format!("sankoff{suffix}"))
|
||||
};
|
||||
let full_path = |suffix: &str| -> String {
|
||||
output.as_ref()
|
||||
.map(|p| format!("{}{suffix}", p.display()))
|
||||
.unwrap_or_else(|| format!("sankoff{suffix}"))
|
||||
};
|
||||
|
||||
let tcm_path = full_path("_sankoff.tcm");
|
||||
let mut f = BufWriter::new(std::fs::File::create(&tcm_path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {tcm_path}: {e}");
|
||||
std::process::exit(1);
|
||||
}));
|
||||
let alphabet_line = STATE_SYMBOL.iter().map(|c| c.to_string()).collect::<Vec<_>>().join(" ");
|
||||
writeln!(f, "{alphabet_line}").unwrap();
|
||||
for i in 0..16 {
|
||||
let mut row: Vec<i64> = (0..16).map(|j| scaled_matrix[i][j]).collect();
|
||||
row.push(scaled_matrix[i][0]); // gap column: same cost as to/from ∅
|
||||
writeln!(f, "{}", row.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(" ")).unwrap();
|
||||
}
|
||||
let mut gap_row: Vec<i64> = (0..16).map(|j| scaled_matrix[0][j]).collect();
|
||||
gap_row.push(0);
|
||||
writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(" ")).unwrap();
|
||||
info!("PhyG TCM → {tcm_path}");
|
||||
|
||||
let pg_path = full_path("_sankoff.pg");
|
||||
let mut f = BufWriter::new(std::fs::File::create(&pg_path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {pg_path}: {e}");
|
||||
std::process::exit(1);
|
||||
}));
|
||||
let fasta_name = basename("_sankoff.fasta");
|
||||
let tcm_name = basename("_sankoff.tcm");
|
||||
let tre_name = basename("_sankoff.tre");
|
||||
writeln!(f, "read(prefasta:\"{fasta_name}\", tcm:\"{tcm_name}\")").unwrap();
|
||||
writeln!(f, "search(seconds:300, instances:4)").unwrap();
|
||||
writeln!(f, "report(\"{tre_name}\", graphs, newick, overwrite)").unwrap();
|
||||
|
||||
let pg_dir = std::path::Path::new(&pg_path).parent()
|
||||
.filter(|d| !d.as_os_str().is_empty())
|
||||
.map(|d| d.display().to_string())
|
||||
.unwrap_or_else(|| ".".into());
|
||||
let pg_name = std::path::Path::new(&pg_path).file_name()
|
||||
.map(|n| n.to_string_lossy().into_owned())
|
||||
.unwrap_or_else(|| pg_path.clone());
|
||||
info!(
|
||||
"PhyG script → {pg_path} (costs scaled x{cost_scale:.0}, runs a default 300s/4-instance \
|
||||
search and writes trees to {tre_name})\n\
|
||||
Run it with:\n \
|
||||
cd {pg_dir} && phyg {pg_name}\n\
|
||||
(`phyg` must run from that directory — `read()`/`report()` in the script use relative \
|
||||
file names)"
|
||||
);
|
||||
}
|
||||
|
||||
// ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────
|
||||
|
||||
fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String {
|
||||
|
||||
@@ -78,6 +78,13 @@ pub fn run(args: FilterCmdArgs) {
|
||||
filters.push(Box::new(MinComplexity { level_max: args.complexity_level_max, theta }));
|
||||
}
|
||||
|
||||
// Source is opened read-only above and needs no lock; only the
|
||||
// destination is written.
|
||||
let _lock = obisys::DirLock::acquire(&args.output).unwrap_or_else(|e| {
|
||||
eprintln!("error locking output directory {}: {e}", args.output.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
let mut rep = Reporter::new();
|
||||
KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep)
|
||||
.unwrap_or_else(|e| {
|
||||
|
||||
@@ -156,6 +156,16 @@ pub fn run(args: IndexArgs) {
|
||||
let output = args.output.clone();
|
||||
let mut rep = Reporter::new();
|
||||
|
||||
// Locked for the whole build (including a possible --force removal +
|
||||
// recreation below): a second `index` run resuming/overwriting the same
|
||||
// output directory concurrently would otherwise corrupt it. Unlinking
|
||||
// the lock file via --force's remove_dir_all is safe — the held file
|
||||
// descriptor keeps the lock regardless of the directory entry.
|
||||
let _lock = obisys::DirLock::acquire(&output).unwrap_or_else(|e| {
|
||||
eprintln!("error locking output directory {}: {e}", output.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
// ── Resolve evidence kind ────────────────────────────────────────────────
|
||||
let (evidence, effective_kmer_size) = if args.approx {
|
||||
let (z, b, fp) = resolve_approx_params(args.findere_z, args.evidence_bits, args.fp);
|
||||
|
||||
@@ -64,6 +64,13 @@ pub fn run(args: MergeArgs) {
|
||||
sources.len(), n_genomes, args.output.display()
|
||||
);
|
||||
|
||||
// Only the destination is written; sources are opened read-only above
|
||||
// and need no lock.
|
||||
let _lock = obisys::DirLock::acquire(&args.output).unwrap_or_else(|e| {
|
||||
eprintln!("error locking output directory {}: {e}", args.output.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
let mut rep = Reporter::new();
|
||||
KmerIndex::merge(&args.output, &source_refs, mode, args.force, args.rename_duplicates, args.budget_fraction, &mut rep).unwrap_or_else(|e| {
|
||||
eprintln!("error merging: {e}");
|
||||
|
||||
@@ -12,6 +12,13 @@ pub struct PackArgs {
|
||||
}
|
||||
|
||||
pub fn run(args: PackArgs) {
|
||||
// Modifies the index in place; acquired before opening so a concurrent
|
||||
// writer can't slip in between the open and the pack below.
|
||||
let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| {
|
||||
eprintln!("error locking index directory {}: {e}", args.index.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
||||
eprintln!("error opening index: {e}");
|
||||
std::process::exit(1);
|
||||
|
||||
@@ -47,6 +47,13 @@ pub fn run(args: ReindexArgs) {
|
||||
IndexMode::Exact
|
||||
};
|
||||
|
||||
// Modifies the index in place; acquired before opening so a concurrent
|
||||
// writer can't slip in between the open and the reindex below.
|
||||
let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| {
|
||||
eprintln!("error locking index directory {}: {e}", args.index.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
let mut idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
||||
eprintln!("error opening index: {e}");
|
||||
std::process::exit(1);
|
||||
|
||||
@@ -215,6 +215,16 @@ pub fn run(args: SelectArgs) {
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
// Lock whichever directory actually gets written: the source itself in
|
||||
// --in-place mode, otherwise the (distinct) --output directory. Acquired
|
||||
// before opening the source so a concurrent writer can't slip in between
|
||||
// the open and the write below.
|
||||
let lock_target = if args.in_place { &args.source } else { args.output.as_ref().unwrap() };
|
||||
let _lock = obisys::DirLock::acquire(lock_target).unwrap_or_else(|e| {
|
||||
eprintln!("error locking {}: {e}", lock_target.display());
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
let mut src = KmerIndex::open(&args.source).unwrap_or_else(|e| {
|
||||
eprintln!("error opening source index: {e}");
|
||||
std::process::exit(1);
|
||||
|
||||
@@ -8,3 +8,4 @@ libc = "0.2"
|
||||
sysinfo = "0.33"
|
||||
indicatif = "0.17"
|
||||
tracing = "0.1"
|
||||
fs4 = "0.9"
|
||||
|
||||
@@ -8,6 +8,49 @@ use tracing::{debug, info, warn};
|
||||
|
||||
const BRAILLE: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||
|
||||
// ── DirLock ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Exclusive advisory lock on an index directory, held for the duration of
|
||||
/// any command that writes into an already-existing index (building the
|
||||
/// sibling annex, merging into a destination, filtering/selecting in place,
|
||||
/// ...). Two such commands racing on the same directory can otherwise
|
||||
/// corrupt each other's writes with no error from either side.
|
||||
///
|
||||
/// Only the directory actually being *written to* needs a lock — a command
|
||||
/// like `merge` that reads several source indexes to build one destination
|
||||
/// only needs to lock the destination.
|
||||
///
|
||||
/// Uses the OS's advisory file lock (`flock` on Unix, `LockFileEx` on
|
||||
/// Windows) via `fs4`, not a hand-rolled PID file: the OS releases it
|
||||
/// automatically on process exit, including a crash — no stale-lock cleanup
|
||||
/// logic needed.
|
||||
pub struct DirLock {
|
||||
_file: std::fs::File,
|
||||
}
|
||||
|
||||
impl DirLock {
|
||||
/// Block until the exclusive lock on `dir` is acquired (creating `dir`
|
||||
/// and the lock file within it if needed). Logs once if the wait is
|
||||
/// non-trivial, so a blocked command doesn't look silently hung.
|
||||
pub fn acquire(dir: &std::path::Path) -> std::io::Result<Self> {
|
||||
use fs4::fs_std::FileExt;
|
||||
|
||||
std::fs::create_dir_all(dir)?;
|
||||
let lock_path = dir.join(".obikmer.lock");
|
||||
let file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.write(true)
|
||||
.open(&lock_path)?;
|
||||
|
||||
if file.try_lock_exclusive().is_err() {
|
||||
info!(dir = %dir.display(), "waiting for another obikmer process to release this index");
|
||||
file.lock_exclusive()?;
|
||||
}
|
||||
Ok(Self { _file: file })
|
||||
}
|
||||
}
|
||||
|
||||
// ── TracedBar ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Wrapper around `ProgressBar` that emits `tracing` events when stderr is not
|
||||
|
||||
Reference in New Issue
Block a user