From ba990a48a0d8559b8ac93466eb7008d3b0de4da9 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 15:27:35 +0200 Subject: [PATCH] feat: add obipipeline for concurrent sibling annex stats Add the `obipipeline` crate and replace sequential scatter/gather logic with a concurrent pipeline using `Flat` and `Transform` stages. Introduce `SiblingAnnexStats` API to compute distributions, and add CLI flags to `distance.rs` for constructing the annex and exporting statistics as CSV. --- src/Cargo.lock | 1 + src/obikindex/Cargo.toml | 1 + src/obikindex/src/lib.rs | 1 + src/obikindex/src/siblings.rs | 291 +++++++++++++++++++++++++------- src/obikmer/src/cmd/distance.rs | 75 +++++++- 5 files changed, 310 insertions(+), 59 deletions(-) diff --git a/src/Cargo.lock b/src/Cargo.lock index 87ae4a4..aa92eee 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1701,6 +1701,7 @@ dependencies = [ "obikpartitionner", "obikseq", "obilayeredmap", + "obipipeline", "obiread", "obiskbuilder", "obiskio", diff --git a/src/obikindex/Cargo.toml b/src/obikindex/Cargo.toml index 05d3d77..20e1dd3 100644 --- a/src/obikindex/Cargo.toml +++ b/src/obikindex/Cargo.toml @@ -11,6 +11,7 @@ obisys = { path = "../obisys" } obicompactvec = { path = "../obicompactvec" } obilayeredmap = { path = "../obilayeredmap" } obiskbuilder = { path = "../obiskbuilder" } +obipipeline = { path = "../obipipeline" } ndarray = "0.16" rayon = "1" crossbeam-channel = "0.5" diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 1f769c6..4d22205 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -19,3 +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; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index 73b09b9..4638bde 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -13,23 +13,22 @@ //! - whether this k-mer is the "minorant" of its family — the smallest //! canonical encoding among the members actually observed. //! -//! Implementation note (deviation from the fully staged `obipipeline` design -//! discussed at length in the doc): this first implementation processes each -//! layer with a straightforward sequential scatter (batch the layer's -//! outgoing variant queries by destination partition) then gather (one -//! `query_partition_with` call per destination partition) — not the -//! multi-stage elementary `obipipeline` pipeline the design settled on. The -//! external semantics (one annex file per layer, sequential outer loop over -//! layers, order-independent reconciliation) match the design exactly; only -//! the internal execution mechanism is simplified, as a scope trade-off. -//! Revisiting this to use `obipipeline` with elementary stages, as designed, -//! is a follow-up, not a behavioural change. +//! Per layer, the computation runs as an `obipipeline` pipeline with several +//! elementary stages (a `Flat` stage generating each k-mer's 3 central +//! variants, a `Transform` stage looking each variant up in its destination +//! partition), so the scheduler's shared worker pool interleaves this work +//! across many in-flight k-mers/variants rather than processing everything +//! on a single thread — see `docmd/theory/evolutionary_distances.md`, Step +//! 2b, "Mechanism", for why elementary stages were chosen deliberately over +//! a few coarse ones. -use std::collections::HashMap; use std::path::Path; +use std::sync::Arc; -use obicompactvec::{SiblingAnnexBuilder, SiblingInfo}; -use obikpartitionner::{KmerDesc, QueryHit}; +use obicompactvec::{ + PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, SiblingInfo, +}; +use obikpartitionner::KmerPartition; use obikseq::{CanonicalKmer, Minimizer}; use obilayeredmap::{MphfLayer, OLMError}; use obilayeredmap::meta::PartitionMeta; @@ -90,6 +89,66 @@ impl Default for RunningState { } } +// ── obipipeline data types ───────────────────────────────────────────────── + +/// One distinct k-mer of the layer currently being processed, at its local +/// MPHF slot — the pipeline's source item. +#[derive(Clone, Copy)] +struct SourceItem { + slot: usize, + kmer: CanonicalKmer, +} + +/// One of a source k-mer's (up to 3) central-substitution variants, already +/// routed to its destination partition and carrying, decided here (both +/// encodings are already in hand — no need to wait for the lookup answer), +/// whether this specific variant would outrank the source as minorant. +#[derive(Clone, Copy)] +struct VariantQuery { + source_slot: usize, + dest_partition: usize, + variant: CanonicalKmer, + smaller: bool, +} + +/// Outcome of looking a [`VariantQuery`] up in its destination partition. +#[derive(Clone, Copy)] +struct AnswerMsg { + source_slot: usize, + hit: bool, + smaller: bool, +} + +#[derive(Clone, Copy)] +enum SibData { + Item(SourceItem), + Query(VariantQuery), + Answer(AnswerMsg), +} + +/// Existence-only lookup of `variant` in partition `dest_partition`: tries +/// each of the partition's layers in turn, stopping at the first hit — the +/// same "try every layer's MPHF" shape as `QueryLayer::find_slot` (private +/// to `obikpartitionner`), reimplemented here directly against the public +/// `MphfLayer::open`/`find` since only existence is needed, not a column +/// fetch. +fn lookup_exists(partition: &KmerPartition, dest_partition: usize, variant: CanonicalKmer) -> bool { + let index_dir = partition.part_dir(dest_partition).join(INDEX_SUBDIR); + if !index_dir.exists() { + return false; + } + let Ok(meta) = PartitionMeta::load(&index_dir) else { return false }; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + if let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) { + if mphf.find(variant).is_some() { + return true; + } + } + } + false +} + impl KmerIndex { /// Build the sibling-count/minorant annex for every layer of every /// partition of this (already built) index, writing one annex file per @@ -97,11 +156,32 @@ impl KmerIndex { /// (e.g. after a fresh `merge`) — each run simply overwrites the annex /// files of the index it is called on. /// + /// Construction only — no statistics gathered here on purpose: this is + /// meant to run routinely (it is the artefact the SNP-family distances + /// will consume), while the sibling-count distribution + /// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate, + /// occasional diagnostic pass over the result, not run every time. + /// /// Cross-partition/cross-layer lookups are required (a k-mer's siblings /// can live in any partition), but the layer loop itself — and thus the /// annex file this produces — stays local to one layer at a time. pub fn build_sibling_annex(&self) -> OKIResult<()> { let n_parts = self.n_partitions(); + let n_bits = n_parts.trailing_zeros() as usize; + + // A fresh, owned `KmerPartition` handle (read-only use only — no + // writers opened), wrapped in `Arc` so pipeline stage closures + // (which must be `'static`, running in spawned threads) can share + // it without borrowing `self`. + let partition = Arc::new( + KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?, + ); for part in 0..n_parts { let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); @@ -112,7 +192,7 @@ impl KmerIndex { for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); - self.build_layer_sibling_annex(&layer_dir, &meta, n_parts)?; + self.build_layer_sibling_annex(&layer_dir, n_parts, &partition)?; } } @@ -122,9 +202,11 @@ impl KmerIndex { fn build_layer_sibling_annex( &self, layer_dir: &Path, - meta: &PartitionMeta, n_parts: usize, + partition: &Arc, ) -> OKIResult<()> { + 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 mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; let n_slots = mphf.n(); @@ -138,53 +220,55 @@ impl KmerIndex { } } - // ── Scatter: bucket outgoing variant queries by destination partition ── - // `KmerDesc.seq_idx` carries the origin (this layer's local slot); - // `KmerDesc.pos` is repurposed as a 0/1 flag: 1 iff this specific - // variant's own encoding is smaller than the source's — decided here, - // at scatter time, since both encodings are already in hand; the - // query response only needs to confirm existence (hit/miss). - let mut outgoing: Vec>> = - (0..n_parts).map(|_| HashMap::new()).collect(); + let sources: Vec = slot_kmer + .iter() + .enumerate() + .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| SourceItem { slot, kmer })) + .collect(); - for (slot, maybe_kmer) in slot_kmer.iter().enumerate() { - let Some(kmer) = maybe_kmer else { continue }; - for variant in kmer.central_canonical_neighbors() { - if variant == *kmer { - continue; // identity substitution — not a real variant + // ── obipipeline: Flat (generate variants) -> Transform (lookup) ───── + let n_workers = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4); + let capacity = 256; + + let partition_for_lookup = Arc::clone(partition); + let pipe = obipipeline::make_pipe! { + SibData : SourceItem => AnswerMsg, + || { + move |item: SourceItem| -> Vec { + let kmer = item.kmer; + 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, + smaller: variant.raw() < kmer.raw(), + }) + .collect::>() } - let smaller = variant.raw() < kmer.raw(); - let dest = partition_of(variant, n_parts); - outgoing[dest].entry(variant).or_default().push(KmerDesc { - seq_idx: slot as u32, - pos: smaller as u32, - }); - } - } + } : Item => Query, + | { + let partition_for_lookup = Arc::clone(&partition_for_lookup); + move |vq: VariantQuery| -> AnswerMsg { + let hit = lookup_exists(&partition_for_lookup, vq.dest_partition, vq.variant); + AnswerMsg { source_slot: vq.source_slot, hit, smaller: vq.smaller } + } + } : Query => Answer, + }; - // ── Gather + reconcile ─────────────────────────────────────────────── + // ── Reconciliation (the pipeline's sink): commutative fold of every + // answer into its origin k-mer's running state, in whatever order + // the pipeline delivers them. ───────────────────────────────────── let mut state = vec![RunningState::default(); n_slots]; - let n_genomes = self.meta.genomes.len(); - let with_counts = self.meta.config.with_counts; - - for (dest, kmers) in outgoing.iter().enumerate() { - if kmers.is_empty() { - continue; + for ans in pipe.apply(sources.into_iter(), n_workers, capacity) { + if ans.hit { + let st = &mut state[ans.source_slot]; + st.siblings = (st.siblings + 1).min(3); + if ans.smaller { + st.minorant = false; + } } - self.partition() - .query_partition_with(dest, kmers, n_genomes, with_counts, |hit| { - if let QueryHit::Found(descs) = hit { - for d in descs { - let slot = d.seq_idx as usize; - let st = &mut state[slot]; - st.siblings = (st.siblings + 1).min(3); - if d.pos == 1 { - st.minorant = false; - } - } - } - }) - .map_err(OKIError::Partition)?; } // ── Write the layer's annex file ───────────────────────────────────── @@ -202,6 +286,97 @@ impl KmerIndex { } } +/// Distribution of sibling counts (0-3), read back from an already-built +/// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's +/// presence/count data — a separate, occasional diagnostic pass, not fused +/// into construction. +#[derive(Debug, Clone, Default)] +pub struct SiblingAnnexStats { + /// `counts[s]` = number of k-mers (slots) with exactly `s` siblings, + /// counted once each regardless of how many genomes carry them. + pub counts: [u64; 4], + /// Of those, how many are minorant. + pub minorant_counts: [u64; 4], + /// `per_genome[g][s]` = number of k-mers with exactly `s` siblings that + /// genome `g` (index into `KmerIndex::meta().genomes`) carries. + pub per_genome: Vec<[u64; 4]>, +} + +impl KmerIndex { + /// Tally the sibling-count distribution of an already-built annex + /// (globally, and per genome). Errors if [`build_sibling_annex`] has not + /// been run on this index first. + /// + /// [`build_sibling_annex`]: Self::build_sibling_annex + pub fn sibling_annex_stats(&self) -> OKIResult { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + + let mut stats = SiblingAnnexStats { + per_genome: vec![[0u64; 4]; n_genomes], + ..Default::default() + }; + + 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() + ))); + } + let annex = SiblingAnnex::open(&annex_path)?; + let use_counts = with_counts && layer_dir.join("counts").exists(); + + // Opened once per layer, outside the slot loop. + enum Mat { + Count(PersistentCompactIntMatrix), + Presence(PersistentBitMatrix), + } + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(&layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(&layer_dir)?) + }; + let n_cols = match &mat { + Mat::Count(m) => m.n_cols(), + Mat::Presence(m) => m.n_cols(), + } + .min(n_genomes); + + for slot in 0..annex.len() { + let Some(info) = annex.get(slot) else { continue }; + let s = info.siblings as usize; + stats.counts[s] += 1; + if info.minorant { + stats.minorant_counts[s] += 1; + } + for g in 0..n_cols { + let carried = match &mat { + Mat::Count(m) => m.col_view(g).get(slot) != 0, + Mat::Presence(m) => m.get(g, slot) != 0, + }; + if carried { + stats.per_genome[g][s] += 1; + } + } + } + } + } + + Ok(stats) + } +} + #[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 999beba..3707387 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}; +use obikindex::{DistanceMetric, KmerIndex, SiblingAnnexStats}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -64,7 +64,21 @@ pub struct DistanceArgs { #[arg(long)] pub upgma: bool, + /// Build the sibling-count/minorant annex on this (multi-genome) index + /// — see `docmd/theory/evolutionary_distances.md`, Step 2b. Construction + /// only; does not by itself compute or write any statistics. + #[arg(long)] + pub sibling_annex: bool, + + /// Tally the sibling-count distribution (CSV) of an already-built annex + /// (run with `--sibling-annex` first, in this invocation or an earlier + /// one). A separate, occasional diagnostic pass — not run every time the + /// annex itself is (re)built. + #[arg(long)] + pub sibling_stats: bool, + /// Output prefix: _dist.csv, _shared.csv, + /// _siblings.csv, _siblings_per_genome.csv, /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] @@ -80,6 +94,26 @@ pub fn run(args: DistanceArgs) { let labels: Vec = idx.meta().genomes.iter().map(|g| g.label.clone()).collect(); let n = labels.len(); + + // ── Sibling-count/minorant annex (independent of the distance metric) ── + // Construction (`--sibling-annex`) and stats (`--sibling-stats`) are + // deliberately decoupled: the annex is meant to be (re)built routinely, + // the distribution only occasionally, on demand. + if args.sibling_annex { + info!("building sibling-count/minorant annex"); + idx.build_sibling_annex().unwrap_or_else(|e| { + eprintln!("error building sibling annex: {e}"); + std::process::exit(1); + }); + } + if args.sibling_stats { + let stats = idx.sibling_annex_stats().unwrap_or_else(|e| { + eprintln!("error computing sibling-annex stats: {e}"); + std::process::exit(1); + }); + write_sibling_stats_csv(&stats, &labels, &args.output); + } + info!( "computing {:?} distances for {} genome(s)", args.metric, n @@ -191,6 +225,45 @@ pub fn run(args: DistanceArgs) { } } +// ── Sibling-count distribution → CSV ──────────────────────────────────────── + +fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: &Option) { + // Global histogram: one row per sibling count (0-3). + let global_path = output.as_ref() + .map(|p| format!("{}_siblings.csv", p.display())) + .unwrap_or_else(|| "siblings.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&global_path).unwrap_or_else(|e| { + eprintln!("error creating {global_path}: {e}"); + std::process::exit(1); + })); + writeln!(f, "siblings,slots,minorant_slots").unwrap(); + for s in 0..4 { + writeln!(f, "{s},{},{}", stats.counts[s], stats.minorant_counts[s]).unwrap(); + } + let total: u64 = stats.counts.iter().sum(); + info!("sibling-count distribution → {global_path} (total {total} slot(s))"); + + // Per-genome breakdown: one row per genome, 4 columns (0-3), + a total row. + let per_genome_path = output.as_ref() + .map(|p| format!("{}_siblings_per_genome.csv", p.display())) + .unwrap_or_else(|| "siblings_per_genome.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&per_genome_path).unwrap_or_else(|e| { + eprintln!("error creating {per_genome_path}: {e}"); + std::process::exit(1); + })); + writeln!(f, "genome,0,1,2,3").unwrap(); + let mut column_totals = [0u64; 4]; + for (label, counts) in labels.iter().zip(stats.per_genome.iter()) { + writeln!(f, "{label},{},{},{},{}", counts[0], counts[1], counts[2], counts[3]).unwrap(); + for s in 0..4 { column_totals[s] += counts[s]; } + } + writeln!( + f, "total,{},{},{},{}", + column_totals[0], column_totals[1], column_totals[2], column_totals[3], + ).unwrap(); + info!("per-genome sibling-count distribution → {per_genome_path}"); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String {