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.

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 <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)
         }

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)

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<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,
     })
 }

src/obikphylo/src/siblings/algorithms/snp_distance.rs

@@ -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) }
     }))
 }