From a8bcf5ffac06500b0a1df3dcc978df6e77c32b12 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 11 Sep 2026 17:31:45 +0200 Subject: [PATCH] Refactor distance matrix calculation logic for consistency Ensures that combined distance matrix calculations use the exact same site selection by retaining sampled data within Sankoff bundles. This involves refactoring post-sampling logic into shared functions, implementing `SankoffBundle` to reuse internal tally data, and adding validation to guarantee consistent site selection across all distance calculation paths. --- DevDocMD/theory/evolutionary_distances.md | 74 ++++++++++++++++++ UserDocMD/usage/phylo.md | 2 + src/obikmer/src/cmd/phylo/mod.rs | 78 ++++++++++++------- .../src/siblings/algorithms/pairwise.rs | 4 + .../src/siblings/algorithms/sankoff.rs | 36 ++++++++- .../src/siblings/algorithms/snp_distance.rs | 35 ++++++++- 6 files changed, 196 insertions(+), 33 deletions(-) diff --git a/DevDocMD/theory/evolutionary_distances.md b/DevDocMD/theory/evolutionary_distances.md index 2f403e76..8036cb8d 100644 --- a/DevDocMD/theory/evolutionary_distances.md +++ b/DevDocMD/theory/evolutionary_distances.md @@ -2674,6 +2674,80 @@ written in-house if the Guindon-Gascuel integration above is ever pursued their simulations, so plain NJ may be an acceptable substitute if BIONJ proper is judged not worth the effort). +### Sharing one sample across algorithms (`--sankoff`/`--tnt`/`--phyg`/`--iqtree` + `snp-*` `--distance`) — implemented 2026-09-11 + +**Correctness bug, not just a performance one.** `--sankoff`/`--tnt`/`--phyg`/ +`--iqtree` (via `SiblingExt::sankoff_bundle`) and a `snp-*` `--distance` +(via `SiblingExt::snp_distance`) always consume the *exact same* selection +parameters when requested together in one `obikmer phylo` invocation +(`n`/`free_loss`/`no_ambiguity`/excluded-set/`entropy_bias` — the CLI has +no way to give them different values in one run, see `cmd/phylo/args.rs`). +Before this fix, each independently called `sample_index`, and +`sample_layer`'s `rand::rng()` (`subsample.rs:235`) is a **thread-local +generator that advances across calls, not reseeded each time** — so the +second call silently drew a *different* random sample of sites than the +first, even with identical parameters. Verified empirically (throwaway +8-genome index): running `snp-jc` alone vs. combined with `--sankoff` +produced visibly different distance matrices from the same index/params. +This defeated the actual point of combining these flags — running several +algorithms (Sankoff/TNT/PhyG/IQ-TREE calibration, a `snp-*` distance) on +one *identical* site selection for direct comparison. + +**Fix**: `SankoffBundle` (`sankoff.rs`) now retains its `PairwiseTally`/ +`PartitionDispersion` internally (previously consumed into `raw`/ +`base_pair_tally`/`cardinality_tally` and dropped) and exposes +`SankoffBundle::snp_distance(kind, gamma_shape)`, which computes the +matrix from that *same* already-sampled tally — no second `sample_index` +call. The shared post-sampling logic (the `--gamma-shape` support check, +`alpha` resolution, final matrix build — previously all inline in the +`snp_distance()` free function) was factored into +`pub(crate) fn distance_matrix(tally, dispersion, kind, gamma_shape)` in +`snp_distance.rs`, called by both the standalone `snp_distance()` (after +its own fresh sampling) and `SankoffBundle::snp_distance` (reusing the +bundle's). `cmd/phylo/mod.rs` keeps the `Option` alive past +the Sankoff-family `if` block and, when a `snp-*` `--distance` is also +requested, calls `bundle.snp_distance(...)` instead of +`cache.snp_distance(...)` whenever a bundle was built — since both branches +are driven by the same `args.*` fields, the parameters trivially always +match when both fire; no runtime "do the params match" check needed. +Verified end-to-end: the reused path logs `(reusing the Sankoff bundle's +sample)` and the `snp_distance` stage timer reads `0ms` (`formula(kind)` +is `O(n²)` post-processing, no I/O), vs. ~150-230ms for a fresh sample on +the same tiny test index. + +#### Future direction, not implemented: explicit `--session` + +Raised in discussion, not started. The fix above only covers reuse +*within one process*. A further idea: a `--session DIR` flag naming a +directory (outside the index) that persists the CLI selection parameters +plus every intermediate artifact `sample_index` would otherwise +recompute — the site set itself, `PairwiseTally`/`PartitionDispersion`, +the pseudo-alignment, etc. — across *separate* `obikmer phylo` +invocations, e.g. running `--sankoff` today and `--distance snp-k2p` +tomorrow against the identical sample. Without `--session`, an implicit +*temporary* session would still be created (scoped to the index directory, +matching the tmp-cache design discussed earlier in this file's α-estimation +section) — same mechanism, just not named/kept by the user. + +Two open design points if this is picked up: +- **Cache-key validity**: as established for the tally cache above, the + content depends on `n`/`free_loss`/`no_ambiguity`/excluded-set/ + `entropy_bias` — a session is only reusable for the exact tuple it was + built under. `--subsample` in particular is the parameter most likely to + change between exploratory runs, so an *exhaustive* (`n = None`) session + that later runs subsample *from*, rather than one session per exact `n`, + would likely see far more real reuse. +- **Concurrency-safe cleanup**: `obisys::DirLock` (`lock.rs`) is the + existing pattern to follow — an OS advisory lock (`flock`/`LockFileEx`), + auto-released on process exit *including a crash*, no stale-lock cleanup + logic needed. Applied per cache/session entry (not just once for the + whole index as `DirLock` does today for annex writes): a process + deciding whether to reclaim an old temporary session first tries to + acquire that entry's lock — success means nobody's using it, safe to + delete; failure means another process holds it, leave it alone. Avoids + the hazard of one process deleting another concurrently-running + process's temp session, which a naive "wipe at startup" would risk. + ### Output format: PHYLIP-relaxed by default for the distance matrix **Implemented.** The primary distance-matrix output diff --git a/UserDocMD/usage/phylo.md b/UserDocMD/usage/phylo.md index dd2ed570..916362ca 100644 --- a/UserDocMD/usage/phylo.md +++ b/UserDocMD/usage/phylo.md @@ -202,6 +202,8 @@ A family is eligible for a genome pair $(i,j)$ only if both genomes carry exactl `--subsample`, `--free-loss`, `--no-ambiguity`, `--entropy`/`--entropy-sd` are shared by `--pseudo-alignment`, `--sankoff` (and everything it implies: `--tnt`/`--phyg`/`--iqtree`), and a `snp-*` `--distance` value — one draw feeds all of them in a single invocation. `--subsample` is mandatory for `--pseudo-alignment`/`--sankoff`; for a `snp-*` `--distance` value it is optional (omitted means every non-monomorphic family in the index, not an approximation). +Combining `--sankoff` (or `--tnt`/`--phyg`/`--iqtree`) with a `snp-*` `--distance` value in the same command reuses that one draw for both — the distance and the Sankoff calibration/alignment are guaranteed to be computed from the *identical* set of sampled sites, never two independent samples, so the two outputs are directly comparable. This only holds within a single command; running them as two separate `obikmer phylo` invocations draws two independent samples even with the same flags. + Without `--subsample`, every variable family (family size ≥ 2) is used. With `--subsample N`, roughly `N` families are kept instead, drawn in proportion to how many candidate families each part of the index actually holds, so the sample stays representative of the whole index. If the index has fewer than `N` candidate families, `--subsample` has no effect. ### `--shannon`: measuring how informative a family is diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index a6e657d5..790ea7ba 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -268,6 +268,14 @@ pub fn run(args: PhyloArgs) { } // ── Sankoff cost-matrix calibration (`--sankoff`, `--tnt`, `--phyg`, `--iqtree`) ── + // `bundle` is kept around (not dropped at the end of this block) so the + // `snp-*` `--distance` branch below can reuse its already-sampled + // tally instead of resampling — see `SankoffBundle::snp_distance`'s own + // docs for why that's required for correctness, not just speed, when + // both are requested in the same invocation (they always share the + // same `n`/`free_loss`/`no_ambiguity`/excluded-set/`entropy_bias`: the + // CLI has no way to give them different values in one run). + let mut bundle = None; if args.sankoff || args.tnt || args.phyg || args.iqtree { let Some(subsample_n) = args.subsample else { eprintln!("error: --sankoff requires --subsample "); @@ -276,7 +284,7 @@ pub fn run(args: PhyloArgs) { info!("sampling Sankoff calibration bundle (target {subsample_n} site(s))"); let t = Stage::start("sankoff_bundle"); - let bundle = cache + let b = cache .sankoff_bundle( subsample_n, args.free_loss, @@ -291,25 +299,25 @@ pub fn run(args: PhyloArgs) { }); rep.push(t.stop()); - let p_card = cardinality_transition_probs(&bundle.cardinality_tally); - let p_comp = composition_transition_probs(&bundle.base_pair_tally); + let p_card = cardinality_transition_probs(&b.cardinality_tally); + let p_comp = composition_transition_probs(&b.base_pair_tally); let matrix = pairwise_cost_matrix(&p_card, &p_comp, args.free_loss); write_sankoff_matrix_csv(&matrix, &args.output); write_sankoff_params( - &bundle.cardinality_tally, + &b.cardinality_tally, &p_card, - &bundle.base_pair_tally, + &b.base_pair_tally, &p_comp, args.sankoff_ratio_ceiling, &args.output, ); - write_sankoff_alignment_fasta(&bundle.alignment, &labels, &args.output, args.free_loss); + write_sankoff_alignment_fasta(&b.alignment, &labels, &args.output, args.free_loss); if args.tnt { write_sankoff_tnt( &matrix, - &bundle.alignment, + &b.alignment, &labels, &args.output, args.sankoff_cost_scale, @@ -322,13 +330,14 @@ pub fn run(args: PhyloArgs) { if args.iqtree { write_iqtree( &matrix, - &bundle.alignment, + &b.alignment, &labels, &args.output, args.free_loss, args.iqtree_min_freq, ); } + bundle = Some(b); } // ── Distance computation: classic whole-index metric vs. `snp-*` ─────────── @@ -354,28 +363,39 @@ pub fn run(args: PhyloArgs) { std::process::exit(1); } let kind = args.distance.as_snp().expect("DistanceArg is always classic or snp"); - info!( - "computing {kind:?} SNP distance for {n} genome(s){}", - match args.subsample { - Some(n) => format!(" (subsampled, target {n} site(s))"), - None => " (exhaustive)".into(), - } - ); + let gamma_shape = args.gamma_shape.map_or(GammaShape::Disabled, Into::into); let t = Stage::start("snp_distance"); - let matrix = cache - .snp_distance( - kind, - args.subsample, - args.free_loss, - args.no_ambiguity, - &snp_exclude_mask, - entropy_bias, - args.gamma_shape.map_or(GammaShape::Disabled, Into::into), - ) - .unwrap_or_else(|e| { - eprintln!("error computing SNP distance: {e}"); - std::process::exit(1); - }); + let matrix = match &bundle { + Some(b) => { + info!( + "computing {kind:?} SNP distance for {n} genome(s) \ + (reusing the Sankoff bundle's sample)" + ); + b.snp_distance(kind, gamma_shape) + } + None => { + info!( + "computing {kind:?} SNP distance for {n} genome(s){}", + match args.subsample { + Some(n) => format!(" (subsampled, target {n} site(s))"), + None => " (exhaustive)".into(), + } + ); + cache.snp_distance( + kind, + args.subsample, + args.free_loss, + args.no_ambiguity, + &snp_exclude_mask, + entropy_bias, + gamma_shape, + ) + } + } + .unwrap_or_else(|e| { + eprintln!("error computing SNP distance: {e}"); + std::process::exit(1); + }); rep.push(t.stop()); (matrix, None) } diff --git a/src/obikphylo/src/siblings/algorithms/pairwise.rs b/src/obikphylo/src/siblings/algorithms/pairwise.rs index edfcbb8b..1691522a 100644 --- a/src/obikphylo/src/siblings/algorithms/pairwise.rs +++ b/src/obikphylo/src/siblings/algorithms/pairwise.rs @@ -52,6 +52,10 @@ impl PairwiseTally { Self { n_genomes, pairs: vec![PairStats::default(); n_pairs] } } + pub(crate) fn n_genomes(&self) -> usize { + self.n_genomes + } + /// See `crate::siblings::helpers::triangle_index`. fn flat_index(&self, i: usize, j: usize) -> usize { crate::siblings::helpers::triangle_index(self.n_genomes, i, j) diff --git a/src/obikphylo/src/siblings/algorithms/sankoff.rs b/src/obikphylo/src/siblings/algorithms/sankoff.rs index b41b3081..14d1a05a 100644 --- a/src/obikphylo/src/siblings/algorithms/sankoff.rs +++ b/src/obikphylo/src/siblings/algorithms/sankoff.rs @@ -10,6 +10,7 @@ //! that removes the old two-pass `raw_snp_distance`-before-`base_pair_tally` //! dependency entirely). +use ndarray::Array2; use obikidxcache::index_cache::IndexCache; use obikindex::OKIResult; @@ -18,17 +19,47 @@ use super::pairwise::{ BasePairTally, CardinalityTally, PairwiseTally, PartitionDispersion, RawSnpDistanceOutput, reduce_pairwise, }; +use super::snp_distance::{GammaShape, SnpDistanceKind, distance_matrix}; use super::subsample::{EntropyBias, sample_index}; /// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs, /// computed together from one shared, possibly-subsampled/entropy-biased /// selection — see the module docs. `pub`: part of the public signature of /// [`crate::siblings::extensions::SiblingExt::sankoff_bundle`]. +/// +/// `tally`/`dispersion` are kept (not just consumed into `raw`/ +/// `base_pair_tally`/`cardinality_tally`) so [`Self::snp_distance`] can +/// compute a `snp-*` `--distance` matrix from this *exact* sample — see +/// that method's own docs for why reusing it, rather than a fresh +/// `SiblingExt::snp_distance` call, is a correctness requirement here, not +/// just a performance nicety. pub struct SankoffBundle { pub alignment: SnpAlignment, pub raw: RawSnpDistanceOutput, pub base_pair_tally: BasePairTally, pub cardinality_tally: CardinalityTally, + tally: PairwiseTally, + dispersion: PartitionDispersion, +} + +impl SankoffBundle { + /// A `snp-*` `--distance` matrix computed from this bundle's own + /// already-sampled tally — no new [`sample_index`] call. + /// + /// Exists specifically for the case where `--sankoff`/`--tnt`/`--phyg`/ + /// `--iqtree` and a `snp-*` `--distance` are requested in the same CLI + /// invocation: both consume exactly the same selection parameters + /// (`n`/`free_loss`/`no_ambiguity`/excluded set/`entropy_bias` — the + /// CLI has no way to give them different values in one run), so a + /// second independent `SiblingExt::snp_distance` call wouldn't just + /// waste the sampling work, it would silently sample a **different** + /// set of sites (`sample_layer`'s `rand::rng()` is a thread-local + /// generator that advances across calls, not reseeded each time) — + /// defeating the actual point of running several algorithms together, + /// which is comparing them on one identical site selection. + pub fn snp_distance(&self, kind: SnpDistanceKind, gamma_shape: GammaShape) -> OKIResult> { + distance_matrix(&self.tally, &self.dispersion, kind, gamma_shape) + } } /// See [`crate::siblings::extensions::SiblingExt::sankoff_bundle`] for the @@ -51,9 +82,6 @@ pub(crate) fn sankoff_bundle( .collect(); let mut sequences: Vec> = vec![Vec::new(); genome_indices.len()]; let mut tally = PairwiseTally::new(n_genomes); - // Unused past this call: the Sankoff/TNT/PhyG/IQ-TREE pipeline has no - // `--gamma-shape`, but `reduce_pairwise` folds it in the same pass as - // `tally` regardless — cheaper than a second scan just to skip it. let mut dispersion = PartitionDispersion::default(); sample_index( @@ -79,5 +107,7 @@ pub(crate) fn sankoff_bundle( raw, base_pair_tally, cardinality_tally, + tally, + dispersion, }) } diff --git a/src/obikphylo/src/siblings/algorithms/snp_distance.rs b/src/obikphylo/src/siblings/algorithms/snp_distance.rs index df3a71c2..41d06920 100644 --- a/src/obikphylo/src/siblings/algorithms/snp_distance.rs +++ b/src/obikphylo/src/siblings/algorithms/snp_distance.rs @@ -241,6 +241,10 @@ pub(crate) fn snp_distance( entropy_bias: Option, gamma_shape: GammaShape, ) -> OKIResult> { + // Fail fast, before paying for `sample_index`, not just inside + // `distance_matrix` (which runs after sampling either way — the right + // place for the check when called from `SankoffBundle::snp_distance`, + // where sampling already happened as part of `--sankoff` itself). if gamma_shape != GammaShape::Disabled && !kind.supports_gamma() { return Err(obikindex::OKIError::InvalidInput( "--gamma-shape has no effect on this --distance value".into(), @@ -274,6 +278,34 @@ pub(crate) fn snp_distance( )?; } + distance_matrix(&tally, &dispersion, kind, gamma_shape) +} + +/// The post-sampling half of [`snp_distance`] — everything downstream of +/// an already-built [`PairwiseTally`]/[`PartitionDispersion`]: the +/// `--gamma-shape` support check, `alpha` resolution (fixed/auto/disabled), +/// and the final `n×n` matrix build. Factored out so +/// [`super::sankoff::SankoffBundle::snp_distance`] can reuse a tally it +/// already has (from `--sankoff`/`--tnt`/`--phyg`/`--iqtree` sharing the +/// exact same selection parameters as a `snp-*` `--distance` in the same +/// CLI invocation) instead of re-running [`sample_index`] — which, beyond +/// the wasted work, would also draw a **different** random sample the +/// second time (`rand::rng()` is a thread-local generator that advances +/// across calls, never reseeded per call), silently breaking the +/// "same site selection for every algorithm in this run" guarantee that's +/// the actual point of combining these flags in one invocation. +pub(crate) fn distance_matrix( + tally: &PairwiseTally, + dispersion: &PartitionDispersion, + kind: SnpDistanceKind, + gamma_shape: GammaShape, +) -> OKIResult> { + if gamma_shape != GammaShape::Disabled && !kind.supports_gamma() { + return Err(obikindex::OKIError::InvalidInput( + "--gamma-shape has no effect on this --distance value".into(), + )); + } + let alpha = match gamma_shape { GammaShape::Disabled => None, GammaShape::Fixed(a) => Some(a), @@ -295,9 +327,10 @@ pub(crate) fn snp_distance( }, }; + let n_genomes = tally.n_genomes(); let f = formula(kind); Ok(Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| { - if i == j { 0.0 } else { f(&tally, i, j, alpha) } + if i == j { 0.0 } else { f(tally, i, j, alpha) } })) }