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.
This commit is contained in:
Eric Coissac
2026-09-12 07:45:32 +02:00
parent be838da256
commit a8bcf5ffac
6 changed files with 196 additions and 33 deletions
+74
View File
@@ -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<SankoffBundle>` 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
+2
View File
@@ -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
+49 -29
View File
@@ -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 <N>");
@@ -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)
}
@@ -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)
@@ -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<Array2<f64>> {
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<u8>> = 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,
})
}
@@ -241,6 +241,10 @@ pub(crate) fn snp_distance(
entropy_bias: Option<EntropyBias>,
gamma_shape: GammaShape,
) -> OKIResult<Array2<f64>> {
// 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<Array2<f64>> {
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) }
}))
}