From 49f329edd5db9bcac27e4626e1fb9ddd6df69dd8 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 17:31:48 +0200 Subject: [PATCH] feat: add raw SNP distance calculation and CLI flag Exposes RawSnpDistanceOutput and implements KmerIndex::raw_snp_distance() to compute pairwise single-copy locus counts under a paralogy-aware rule. The implementation leverages ndarray for parallel matrix aggregation, producing raw p-distance matrices for sanity-checking. A --raw-snp-distance CLI flag is added to export results as CSV, mapping zero-eligible pairs to NA. --- src/obikindex/src/lib.rs | 2 +- src/obikindex/src/siblings.rs | 172 ++++++++++++++++++++++++++++++++ src/obikmer/src/cmd/distance.rs | 71 +++++++++++-- 3 files changed, 235 insertions(+), 10 deletions(-) diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 4d22205..4b09bc3 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -19,4 +19,4 @@ 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::SiblingAnnexStats; +pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats}; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index e80bd67..6ab50a8 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -31,6 +31,7 @@ use std::path::Path; use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; +use ndarray::Array2; use rayon::prelude::*; use obicompactvec::{ @@ -609,6 +610,177 @@ impl KmerIndex { } } +/// Raw p-distance restricted to loci that are single-copy in **both** +/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility +/// rule (`docmd/theory/evolutionary_distances.md`, "Locus eligibility"), +/// without the JC/K2P/LogDet finalisation math: `snp[i,j]` + `shared[i,j]` +/// is the number of eligible loci, `snp[i,j] / (snp[i,j] + shared[i,j])` is +/// `p_hat`. A quick, self-contained way to sanity-check the estimator +/// against a real index before the full `SnpTally` design is built. +/// +/// A locus (family, tallied once at its minorant) is eligible for pair +/// `(i, j)` iff genome `i` carries exactly one of the family's observed +/// forms **and** genome `j` carries exactly one (possibly a different one) +/// — presence-only: a genome carrying the same form twice (a same-allele +/// duplicate) is indistinguishable from carrying it once when only a +/// presence matrix is available, so such cases are not excluded here even +/// when a count index exists. See "Locus eligibility", stringent rule, for +/// why this matters and how a count index would close the gap — left as a +/// follow-up, not applied here. +pub struct RawSnpDistanceOutput { + /// n×n count of eligible loci where the two genomes' single forms differ. + pub snp: Array2, + /// n×n count of eligible loci where the two genomes' single forms agree. + pub shared: Array2, +} + +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 { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; + + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + layer_dirs.push(layer_dir); + } + } + + let pb = progress_bar("raw_snp_distance", layer_dirs.len() as u64, "layers"); + let partials: Vec<(Array2, Array2)> = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult<(Array2, Array2)> { + let mut snp = Array2::::zeros((n_genomes, n_genomes)); + let mut shared = Array2::::zeros((n_genomes, n_genomes)); + + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); + + // Per family: which single form (if exactly one) each genome + // carries — `None` once a second form is seen (ambiguous, + // not single-copy, ineligible for either side of a pair). + let mut single_form: Vec> = Vec::with_capacity(n_cols); + let mut ambiguous: Vec = Vec::with_capacity(n_cols); + + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant + } + + single_form.clear(); + single_form.resize(n_cols, None); + ambiguous.clear(); + ambiguous.resize(n_cols, false); + + for other in kmer.central_canonical_neighbors() { + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let presence: Option> = if other == kmer { + Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) + } else { + let dest = partition_of(other, n_parts); + cache.find_presence(dest, other, n_genomes) + }; + let Some(presence) = presence else { continue }; + for (g, &present) in presence.iter().enumerate() { + if !present { + continue; + } + if single_form[g].is_some() { + ambiguous[g] = true; + } else { + single_form[g] = Some(base); + } + } + } + + for i in 0..n_cols { + if ambiguous[i] { + continue; + } + let Some(bi) = single_form[i] else { continue }; + for j in (i + 1)..n_cols { + if ambiguous[j] { + 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; + } + } + } + } + + pb.inc(1); + Ok((snp, shared)) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut snp = Array2::::zeros((n_genomes, n_genomes)); + let mut shared = Array2::::zeros((n_genomes, n_genomes)); + for (s, sh) in partials { + snp += &s; + shared += &sh; + } + Ok(RawSnpDistanceOutput { snp, shared }) + } +} + #[cfg(test)] mod tests { use std::io::Write; diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 7a5ad4e..9e02552 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; -use obikindex::{DistanceMetric, KmerIndex, SiblingAnnexStats}; +use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -77,8 +77,17 @@ pub struct DistanceArgs { #[arg(long)] pub sibling_stats: bool, + /// Compute the raw p-distance restricted to loci that are single-copy + /// in both genomes of each pair (an already-built sibling annex is + /// required — run with `--sibling-annex` first, in this invocation or + /// an earlier one). A quick way to test the central-position SNP + /// estimator against a real index; not the full `SnpTally` design. + #[arg(long)] + pub raw_snp_distance: bool, + /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _nj.nwk, _upgma.nwk. + /// _siblings.csv, _rawsnp.csv, _nj.nwk, + /// _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -112,15 +121,22 @@ pub fn run(args: DistanceArgs) { }); write_sibling_stats_csv(&stats, &labels, &args.output); } + if args.raw_snp_distance { + let result = idx.raw_snp_distance().unwrap_or_else(|e| { + eprintln!("error computing raw SNP distance: {e}"); + std::process::exit(1); + }); + write_raw_snp_distance_csv(&result, &labels, &args.output); + } - // `--sibling-annex`/`--sibling-stats` are their own operation, not a - // modifier on top of a distance-metric computation — a metric was - // never requested by asking for either 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 + // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance` 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 { + if args.sibling_annex || args.sibling_stats || args.raw_snp_distance { return; } @@ -269,6 +285,43 @@ fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: if total == 1 { "y" } else { "ies" }); } +// ── Raw single-copy SNP distance → CSV ────────────────────────────────────── +// +// p_hat[i,j] = snp[i,j] / (snp[i,j] + shared[i,j]) over loci single-copy in +// both i and j — see `RawSnpDistanceOutput` / `KmerIndex::raw_snp_distance`. +// A single file: the distance matrix, with an eligible-loci count alongside +// each value so a 0/0 pair (no eligible locus at all) is distinguishable +// from a genuinely identical pair. + +fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_rawsnp.csv", p.display())) + .unwrap_or_else(|| "rawsnp.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); + })); + let n = labels.len(); + write!(f, "genome").unwrap(); + for g in labels { write!(f, ",{g}").unwrap(); } + writeln!(f).unwrap(); + for (i, g) in labels.iter().enumerate() { + write!(f, "{g}").unwrap(); + for j in 0..n { + let snp = result.snp[[i, j]]; + let shared = result.shared[[i, j]]; + let eligible = snp + shared; + if eligible == 0 { + write!(f, ",NA").unwrap(); + } else { + write!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap(); + } + } + writeln!(f).unwrap(); + } + info!("raw single-copy SNP distance matrix → {path}"); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String {