From 0de078fdf13bb15dd54cbf00a9c103556cab3c5e Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 14 Aug 2026 10:06:44 +0200 Subject: [PATCH] Store minorant flag in family mask to avoid costly k-mer reconstruction Transition the minorant flag from a derived value to a stored field within the family mask, resolving a performance regression where on-the-fly reconstruction consumed significant query time. This change introduces O(1) k-mer reconstruction APIs, shifts minorant computation to the index build phase, and enables direct annex-based statistics. Supporting updates include adopting shared ownership for partition caches and refactoring batch processing pipelines. --- docmd/theory/evolutionary_distances.md | 33 ++- src/obicompactvec/src/siblingannex.rs | 59 +++- src/obikindex/src/siblings/alignment.rs | 4 +- src/obikindex/src/siblings/build.rs | 33 +-- src/obikindex/src/siblings/cache.rs | 67 ++++- src/obikindex/src/siblings/cardinality.rs | 4 +- src/obikindex/src/siblings/distance.rs | 4 +- src/obikindex/src/siblings/family_scan.rs | 313 +++++++++++++++------- src/obikindex/src/siblings/stats.rs | 42 ++- src/obikindex/src/siblings/tests.rs | 28 +- src/obikmer/src/cmd/phylo/args.rs | 14 +- src/obikmer/src/cmd/phylo/mod.rs | 9 +- src/obikmer/src/cmd/phylo/outputs.rs | 19 ++ src/obilayeredmap/src/mphf_layer.rs | 25 ++ src/obiskio/src/unitig_index/reader.rs | 10 + 15 files changed, 519 insertions(+), 145 deletions(-) diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 137ee58d..9405adfc 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -1624,10 +1624,33 @@ a fixed reference. Storing instead a **4-bit mask** — one bit per base family's fixed canonical form, i.e. the member with `A` at the centre — see above) is observed anywhere in the index — fixes both: - **Sibling count is derived, not stored**: `siblings = popcount(mask) - 1`. -- **Minorant is derived, not stored**: regenerate the family's 4 canonical +- ~~**Minorant is derived, not stored**: regenerate the family's 4 canonical forms from the slot's own k-mer (cheap, no lookup — see above), compare the raw encodings of whichever bits are set in the mask, take the - smallest. + smallest.~~ **Erratum (2026-08-14) — this was wrong, kept struck through + rather than deleted.** "Cheap, no lookup" only accounts for the bit + algebra (regenerate 4 forms, compare raw encodings) — true in isolation, + but it silently assumed "the slot's own k-mer" is a free fact. It isn't: + getting from an MPHF slot index back to the actual k-mer sequence means + reconstructing `slot_kmer` for the whole layer — scan `unitigs.bin`, + `mphf.find()` every k-mer to place it — an O(distinct k-mers) pass + through the MPHF, not a per-slot O(1) lookup. That reconstruction is free + *only* when the caller already needs k-mer identity for something else in + the same traversal (e.g. the SNP sweep below, which needs it anyway to + generate `central_canonical_neighbors()`). A caller that wants *only* the + minorant flag pays the full reconstruction for nothing: measured on a real + run, a bare family-size histogram (four buckets, otherwise near-instant) + spent 71% of wall-clock in `MphfLayer::find`, all of it solely to answer + "is this slot the minorant". **Current design: minorant *is* stored** + after all — a 5th mask bit (4 presence bits + 1 minorant bit, still fits + one byte alongside the presence mask below), written once in the + construction pass where the k-mer is already in hand for other reasons + (`obikindex::siblings::build_layer_sibling_annex`), read back for free by + every later consumer (`FamilyMask::is_minorant`). This is genuinely a + return to the superseded 3-bit design's core idea (store minorant + alongside the count) — the "which variant" blindness that motivated + moving away from it is fixed by keeping the full 4-bit presence mask + too, not by dropping the stored minorant bit again. - **A future consumer knows exactly which variants to (re-)query** — `popcount(mask) - 1` lookups instead of always 3, and it knows *which* 3 (or fewer) to issue, not just how many hits to expect. @@ -1644,7 +1667,11 @@ above) is observed anywhere in the index — fixes both: per-slot packed array (the presence mask above), one per partition — same on-disk shape family as `PersistentBitMatrix`'s `Packed` variant, but simpler (no per-genome columns, a single derived read-only value per - slot). + slot). As actually implemented (`obicompactvec::siblingannex`): **not** + truly bit-packed — 1 byte/slot, 4 presence bits + the minorant bit (see + the erratum above) in the low 5 bits, 3 unused. Deliberately simpler for + a first implementation; packing to 5 bits/slot is a pure storage-density + follow-up, not a behavioural change, still not done as of this note.
Superseded 3-bit design (historical) 3 bits, storing minorant status alongside sibling count directly, since it came for free from the same lookups (point 3 below) — 5 real states diff --git a/src/obicompactvec/src/siblingannex.rs b/src/obicompactvec/src/siblingannex.rs index 0de9733e..80e72d56 100644 --- a/src/obicompactvec/src/siblingannex.rs +++ b/src/obicompactvec/src/siblingannex.rs @@ -41,11 +41,22 @@ const MAGIC: [u8; 4] = *b"PSIB"; // Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows. const HEADER_SIZE: usize = 16; -/// A family presence mask: bit `b` set iff the member whose own canonical -/// central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the index. +/// A family presence mask: bit `b` (0..3) set iff the member whose own +/// canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the +/// index. Bit 4 is a second, independent fact piggy-backed onto the same +/// byte: whether *this slot's own k-mer* is its family's minorant (the +/// smallest raw encoding among the family's observed members) — computed +/// once, when the whole family's mask is already final (see +/// `obikindex::siblings::build_sibling_annex`), and read back by every +/// consumer that would otherwise have to reconstruct this slot's k-mer from +/// `unitigs.bin` and hash it through the MPHF again just to ask the same +/// question. Bits 5-7 unused. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct FamilyMask(u8); +const PRESENCE_BITS: u8 = 0b0000_1111; +const MINORANT_BIT: u8 = 0b0001_0000; + impl FamilyMask { /// The empty mask — never a valid *computed* result (a slot's own base /// is always present in its own family) — used only to build up a mask @@ -69,7 +80,7 @@ impl FamilyMask { /// Number of family members observed anywhere in the index (1..=4). #[inline] pub fn family_size(self) -> u32 { - self.0.count_ones() + (self.0 & PRESENCE_BITS).count_ones() } /// Number of *other* members observed (0..=3) — `family_size() - 1`. @@ -78,18 +89,39 @@ impl FamilyMask { self.family_size() - 1 } - /// Raw bitmask (bit `b` = base `b` present) — for callers that build up - /// a mask via their own bit operations (e.g. concurrently, via an - /// `AtomicU8`) and only need the `FamilyMask` wrapper at the end. + /// Set or clear the minorant flag (see the struct docs). #[inline] - pub fn bits(self) -> u8 { - self.0 + pub fn with_minorant(self, is_minorant: bool) -> Self { + if is_minorant { + FamilyMask(self.0 | MINORANT_BIT) + } else { + FamilyMask(self.0 & !MINORANT_BIT) + } } - /// Construct from a raw bitmask (only the low 4 bits are kept). + /// Is this slot's own k-mer its family's minorant? Only meaningful once + /// `with_minorant` has been called with the family's *final* mask (i.e. + /// after construction) — see the struct docs. + #[inline] + pub fn is_minorant(self) -> bool { + self.0 & MINORANT_BIT != 0 + } + + /// Raw presence bitmask (bit `b` = base `b` present, low 4 bits only — + /// never includes the minorant flag) — for callers that build up a mask + /// via their own bit operations (e.g. concurrently, via an `AtomicU8`) + /// and only need the `FamilyMask` wrapper at the end. + #[inline] + pub fn bits(self) -> u8 { + self.0 & PRESENCE_BITS + } + + /// Construct from a raw presence bitmask (only the low 4 bits are kept — + /// the minorant flag is not part of this, use [`with_minorant`](Self::with_minorant) + /// separately). #[inline] pub fn from_bits(bits: u8) -> Self { - FamilyMask(bits & 0b1111) + FamilyMask(bits & PRESENCE_BITS) } #[inline] @@ -101,10 +133,13 @@ impl FamilyMask { fn decode(byte: u8) -> Option { if byte == 0 { // Unreachable for a real result — reserved as the "not yet - // computed" sentinel. + // computed" sentinel. Still safe as a sentinel with the + // minorant bit added: a real entry always has at least one + // presence bit set (bits 0-3), so byte == 0 still means + // "nothing written yet" and never a genuine minorant-only value. return None; } - Some(FamilyMask(byte & 0b1111)) + Some(FamilyMask(byte)) } } diff --git a/src/obikindex/src/siblings/alignment.rs b/src/obikindex/src/siblings/alignment.rs index 93f478d3..e87e1017 100644 --- a/src/obikindex/src/siblings/alignment.rs +++ b/src/obikindex/src/siblings/alignment.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use obikpartitionner::KmerPartition; use obisys::progress_bar; @@ -68,7 +70,7 @@ impl KmerIndex { n_bits, ) .map_err(OKIError::Partition)?; - let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = self.sibling_layer_dirs()?; let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); diff --git a/src/obikindex/src/siblings/build.rs b/src/obikindex/src/siblings/build.rs index a163b84c..a23896ee 100644 --- a/src/obikindex/src/siblings/build.rs +++ b/src/obikindex/src/siblings/build.rs @@ -10,14 +10,13 @@ use obipipeline::ThrottleGuard; use obikseq::CanonicalKmer; use obilayeredmap::MphfLayer; use obilayeredmap::meta::PartitionMeta; -use obiskio::UnitigFileReader; use obisys::progress_bar; use crate::error::{OKIError, OKIResult}; use crate::index::KmerIndex; use super::cache::PartitionCache; -use super::helpers::{central_base, partition_of}; +use super::helpers::{central_base, is_minorant, partition_of}; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; // ── obipipeline data types ───────────────────────────────────────────────── @@ -117,15 +116,10 @@ impl KmerIndex { let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; let n_slots = mphf.n(); - // ── Enumerate this layer's distinct k-mers, one per slot ──────────── - let mut slot_kmer: Vec> = vec![None; n_slots]; - 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); - } - } + // ── Enumerate this layer's distinct k-mers, one per slot — direct + // slot -> k-mer reconstruction (evidence + direct-access unitigs, no + // MPHF hashing, no file scan), not scan-and-hash-every-k-mer-forward. + let slot_kmer: Vec> = (0..n_slots).map(|slot| mphf.kmer_at(slot)).collect(); let k = self.kmer_size(); @@ -240,13 +234,22 @@ impl KmerIndex { }); // ── Write the layer's annex file ───────────────────────────────────── + // The minorant flag is computed here, not by every later reader: + // this is the one place the whole family's *final* mask and this + // slot's own k-mer (already in hand, no extra lookup) are both + // available together. Every consumer that only needs to know + // "is this slot the family's minorant" — the common case, since a + // family is tallied once, at its minorant — reads the bit straight + // back instead of re-deriving it (re-scanning `unitigs.bin` and + // re-hashing through the MPHF, the cost `is_minorant` was cheap to + // *compute* but expensive to *get the inputs for* every time). let annex_path = layer_dir.join(ANNEX_FILE_NAME); let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?; for (slot, m) in mask.iter().enumerate() { - if slot_kmer[slot].is_none() { - continue; // unused MPHF slot, if any — leave at the sentinel - } - builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed))); + let Some(kmer) = slot_kmer[slot] else { continue }; // unused MPHF slot, if any — leave at the sentinel + let final_mask = FamilyMask::from_bits(m.load(Ordering::Relaxed)); + let minorant = is_minorant(kmer, final_mask, k); + builder.set(slot, final_mask.with_minorant(minorant)); } builder.close()?; diff --git a/src/obikindex/src/siblings/cache.rs b/src/obikindex/src/siblings/cache.rs index a2de07e5..32d49f5d 100644 --- a/src/obikindex/src/siblings/cache.rs +++ b/src/obikindex/src/siblings/cache.rs @@ -104,18 +104,63 @@ impl PartitionCache { .is_some_and(|layers| layers.iter().any(|mphf| mphf.find(variant).is_some())) } - /// Per-genome presence vector for `variant` in partition `dest_partition` - /// (`true` iff that genome carries it), `None` on a miss. Same shape as - /// `find`, but also reads the cached matrix instead of just the MPHF. - pub(super) fn find_presence(&self, dest_partition: usize, variant: CanonicalKmer, n_genomes: usize) -> Option> { - let layers = self.layers.get(dest_partition)?; - let mats = self.mats.get(dest_partition)?; - for (mphf, mat) in layers.iter().zip(mats.iter()) { - if let Some(slot) = mphf.find(variant) { - let n_cols = mat.n_cols().min(n_genomes); - return Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()); + /// Resolve many `(variant, family_idx, base)` queries against one + /// partition's matrices at once, calling `on_hit(family_idx, base, g)` + /// for every genome `g` that carries the resolved variant. + /// + /// Genome-major, not query-major: `PersistentBitMatrix`/ + /// `PersistentCompactIntMatrix` are stored one contiguous block per + /// genome (column), slot as the offset within it (see + /// `obicompactvec::bitmatrix::packed::PackedBitMatrix` — each column's + /// own `mmap` region). Resolving query-by-query (`for query { for genome + /// { mat.carries(genome, slot) } }`, this function's predecessor) visits + /// every genome's block once *per query* — for a batch of thousands of + /// queries against ~90 genomes, that is thousands of jumps into each of + /// ~90 widely separated multi-MB regions, in query order, not genome + /// order: the access pattern a column-major layout is least suited to. + /// Grouping first (by layer, since each layer's matrix is a separate + /// column set) and sorting each group by slot, then visiting genome by + /// genome, turns that into ~90 mostly-sequential sweeps through one + /// column's own bytes — the layout's fast axis — confirmed by sampling + /// a real run: `PersistentBitMatrix::get` dominated wall-clock time, + /// mostly blocked on page faults, even after every partition/batch + /// locality fix above this in the traversal. + pub(super) fn find_presence_batch( + &self, + dest_partition: usize, + queries: &[(CanonicalKmer, usize, u8)], + n_genomes: usize, + mut on_hit: impl FnMut(usize, u8, usize), + ) { + let Some(layers) = self.layers.get(dest_partition) else { return }; + let Some(mats) = self.mats.get(dest_partition) else { return }; + + // First hit wins, same semantics as the old per-query loop (a + // variant present in an earlier layer shadows later ones). + let mut by_layer: Vec> = vec![Vec::new(); layers.len()]; + for &(variant, family_idx, base) in queries { + for (li, mphf) in layers.iter().enumerate() { + if let Some(slot) = mphf.find(variant) { + by_layer[li].push((slot, family_idx, base)); + break; + } + } + } + + for (li, mut hits) in by_layer.into_iter().enumerate() { + if hits.is_empty() { + continue; + } + hits.sort_unstable_by_key(|&(slot, _, _)| slot); + let mat = &mats[li]; + let n_cols = mat.n_cols().min(n_genomes); + for g in 0..n_cols { + for &(slot, family_idx, base) in &hits { + if mat.carries(g, slot) { + on_hit(family_idx, base, g); + } + } } } - None } } diff --git a/src/obikindex/src/siblings/cardinality.rs b/src/obikindex/src/siblings/cardinality.rs index 5921be7e..5f09415d 100644 --- a/src/obikindex/src/siblings/cardinality.rs +++ b/src/obikindex/src/siblings/cardinality.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use ndarray::Array2; use obikpartitionner::KmerPartition; @@ -72,7 +74,7 @@ impl KmerIndex { n_bits, ) .map_err(OKIError::Partition)?; - let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = self.sibling_layer_dirs()?; let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers"); diff --git a/src/obikindex/src/siblings/distance.rs b/src/obikindex/src/siblings/distance.rs index 00965acd..3f746303 100644 --- a/src/obikindex/src/siblings/distance.rs +++ b/src/obikindex/src/siblings/distance.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use ndarray::Array2; use obikpartitionner::KmerPartition; @@ -79,7 +81,7 @@ impl KmerIndex { n_bits, ) .map_err(OKIError::Partition)?; - let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = self.sibling_layer_dirs()?; let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); diff --git a/src/obikindex/src/siblings/family_scan.rs b/src/obikindex/src/siblings/family_scan.rs index 2d11f18c..5aacdd62 100644 --- a/src/obikindex/src/siblings/family_scan.rs +++ b/src/obikindex/src/siblings/family_scan.rs @@ -4,25 +4,51 @@ //! base-presence, with the same locality discipline as //! `build_sibling_annex` (see `cache.rs`/`build.rs`'s docs): cross-partition //! lookups are grouped by destination partition and resolved in one -//! contiguous sweep per partition, instead of one lookup at a time jumping -//! between partitions in family order. +//! contiguous sweep per partition, instead of one lookup jumping between +//! partitions in family order. //! -//! Processed in bounded batches of [`FAMILY_BATCH`] families, not the whole -//! layer at once — same rationale and constant as `build_layer_sibling_annex`'s -//! `BATCH_SIZE`. An earlier version of this function materialised every -//! family of the whole layer up front (`Vec`, each with its own -//! heap-allocated `Vec`) before handing anything back to the caller: -//! O(layer's family count × genome count) extra memory, on top of a second -//! full copy for the owned per-family vectors — measured pushing a real run -//! into heavy VM compression/swap even though the per-lookup I/O pattern was -//! already fixed. A batch of a few thousand families needs the same -//! partition-grouping to keep lookups local; it does not need the entire -//! layer resident at once. Callers get each family's `genome_mask` through -//! `on_family`, backed by one reused scratch buffer — no per-family -//! allocation at all, matching the O(genome count) working set the original, -//! pre-locality-fix per-family code had. +//! Two concerns, kept on two different mechanisms because they need +//! opposite things: +//! +//! - **Generating** each batch's cross-partition queries is pure CPU (bit +//! tests, one hash per variant) — cheap, and safe to run for several +//! batches at once. It runs on an `obipipeline::throttle` + `make_pipe!` +//! pipeline, the same mechanism `build_layer_sibling_annex` uses, for the +//! same reason: it overlaps this CPU work with the *previous* batch's +//! resolution instead of leaving cores idle between batches. +//! - **Resolving** those queries against `PartitionCache` is I/O (mmap page +//! faults on a real index far larger than RAM) — one batch at a time, +//! `rayon`-parallel *across partitions* (like `build_sibling_annex`'s own +//! `outgoing.par_iter()`), never several batches concurrently. An earlier +//! version resolved each batch fully inside its own pipeline worker, so +//! `n_workers` batches were resolved concurrently — each one spreading its +//! queries thin across every partition of the layer at once. Sampling a +//! real run showed the fix hadn't helped at all: 16 threads "busy" in +//! `find_presence`, but mostly blocked on page faults (~1.7 GB/s pagein), +//! because 16 concurrent sweeps across the same partition space is exactly +//! the scattering the grouping was meant to prevent — just spread over +//! threads instead of over layers this time. Resolving one batch's worth +//! at a time, with each of `rayon`'s threads owning one partition +//! contiguously until that batch is done, keeps only one partition set +//! "hot" at once, restoring the locality the batching is for. +//! +//! `obipipeline` does not guarantee output order, so generated batches carry +//! their own starting family index and are replayed through a small reorder +//! buffer before resolution — bounded by the throttle's own concurrency +//! (`n_workers` batches), not by the layer's size. +//! +//! [`FAMILY_BATCH`] bounds a single batch's memory to `FAMILY_BATCH * +//! genome_count` bytes, and is chosen large enough that a batch still gives +//! each partition a decent number of queries to resolve in one go — too +//! small a batch starves that density regardless of how the resolution step +//! is parallelised. An earlier version materialised the whole layer's +//! families up front before resolving anything: O(layer's family count × +//! genome count) memory, measured pushing a real run into heavy VM +//! compression/swap. +use std::collections::HashMap; use std::path::{Path, PathBuf}; +use std::sync::Arc; use std::sync::atomic::{AtomicU8, Ordering}; use rayon::prelude::*; @@ -31,21 +57,21 @@ use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, use obikseq::CanonicalKmer; use obilayeredmap::MphfLayer; use obilayeredmap::meta::PartitionMeta; -use obiskio::UnitigFileReader; +use obipipeline::{ThrottleGuard, throttle}; use crate::error::{OKIError, OKIResult}; use crate::index::KmerIndex; use super::cache::{Mat, PartitionCache}; -use super::helpers::{central_base, is_minorant, partition_of}; +use super::helpers::{central_base, partition_of}; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; -/// Families per batch — bounds `scan_layer_families`'s extra memory to -/// `FAMILY_BATCH * n_genomes` bytes regardless of layer size. Same value and -/// rationale as `build_layer_sibling_annex`'s `BATCH_SIZE`: large enough to -/// amortise the per-batch partition-grouping/sync cost, small enough to keep -/// working memory a rounding error next to the index itself. -const FAMILY_BATCH: usize = 4096; +/// Families per batch — see the module docs for the memory-vs-per-partition- +/// density trade-off this picks a point on. At ~90 genomes and a few +/// hundred partitions, this keeps a batch's resolution memory in the low +/// tens of MB while still giving each partition on the order of a thousand +/// queries per batch to amortise against. +const FAMILY_BATCH: usize = 65536; impl KmerIndex { /// Every (partition, layer) directory carrying a sibling annex, checked @@ -75,9 +101,52 @@ impl KmerIndex { } } +/// Read-only state shared (via `Arc`) across every pipeline worker +/// generating this layer's batches — opened once, not per batch. No `cache` +/// here: generation never touches the cross-partition cache, only this +/// layer's own already-open matrix. No `annex` either — the filter step +/// that builds `minorant_slots` already carries each slot's mask through, +/// so a worker never needs to re-read it. +struct LayerCtx { + slot_kmer: Vec>, + mat: Mat, + n_parts: usize, + n_genomes: usize, + n_cols: usize, + k: usize, +} + +struct SourceBatch { + start_family_idx: usize, + /// `(slot, mask)` — the mask is carried through from the filter step + /// below rather than re-read from the annex per batch; it's the same + /// byte either way, just already in hand. + slots: Vec<(usize, FamilyMask)>, + _permit: ThrottleGuard, +} + +/// One batch's generated work: local presence already resolved (straight +/// from this layer's own matrix, no lookup needed), cross-partition queries +/// collected but not yet resolved against the cache. +struct GeneratedBatch { + start_family_idx: usize, + masks: Vec, + /// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`. + genome_mask: Vec, + /// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base)`. + outgoing: Vec>, + _permit: ThrottleGuard, +} + +enum FamData { + Batch(SourceBatch), + Generated(GeneratedBatch), +} + /// Visits every minorant family of one layer, in slot order, batched -/// [`FAMILY_BATCH`] at a time — see the module docs for why. `on_family` is -/// called once per family with its own annex mask and its per-genome +/// [`FAMILY_BATCH`] at a time — see the module docs for why generation and +/// resolution use different concurrency. `on_family` is called once per +/// family, in slot order, with its own annex mask and its per-genome /// base-presence (`genome_mask[g]`: bit `b` set iff genome `g` carries the /// member whose own canonical central base is `b`), backed by a scratch /// buffer reused across every call — callers that need to keep data past @@ -88,7 +157,7 @@ pub(super) fn scan_layer_families( n_genomes: usize, with_counts: bool, k: usize, - cache: &PartitionCache, + cache: &Arc, mut on_family: impl FnMut(FamilyMask, &[u8]), ) -> OKIResult<()> { let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); @@ -96,13 +165,12 @@ pub(super) fn scan_layer_families( 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); - } - } + // `MphfLayer::kmer_at` — direct slot -> k-mer reconstruction (evidence + + // direct-access unitigs, no MPHF hashing, no file scan) instead of the + // previous scan-`unitigs.bin`-and-hash-every-k-mer-forward approach, + // which redundantly re-read the same data `mphf.find()` itself already + // reads through `evidence`/`unitigs` to verify each hit. + let slot_kmer: Vec> = (0..annex.len()).map(|slot| mphf.kmer_at(slot)).collect(); let use_counts = with_counts && layer_dir.join("counts").exists(); let mat = if use_counts { @@ -112,77 +180,132 @@ pub(super) fn scan_layer_families( }; let n_cols = mat.n_cols().min(n_genomes); - let n_slots = annex.len(); - let mut slot = 0usize; - let mut batch_slots: Vec = Vec::with_capacity(FAMILY_BATCH); - let mut scratch = vec![0u8; n_genomes]; // reused across every `on_family` call — no per-family allocation + // Minorant slots, in order, mask carried along — the flag is already + // stored in the annex (set once, at construction — see + // `build_layer_sibling_annex`), no need to re-derive it from the kmer + // here, and no need for the pipeline workers below to re-read the same + // byte from the annex a second time. `slot_kmer[s]` is still required + // to exist (used later for `central_canonical_neighbors()`), but not + // for this check. + let minorant_slots: Vec<(usize, FamilyMask)> = (0..annex.len()) + .filter_map(|s| { + let mask = annex.get(s)?; + (slot_kmer[s].is_some() && mask.is_minorant()).then_some((s, mask)) + }) + .collect(); + let total_families = minorant_slots.len(); + if total_families == 0 { + return Ok(()); + } - while slot < n_slots { - batch_slots.clear(); - while slot < n_slots && batch_slots.len() < FAMILY_BATCH { - let s = slot; - slot += 1; - let Some(mask) = annex.get(s) else { continue }; - let Some(kmer) = slot_kmer[s] else { continue }; - if !is_minorant(kmer, mask, k) { - continue; // family tallied once, at its minorant - } - batch_slots.push(s); - } - if batch_slots.is_empty() { - continue; // ran out of minorant slots before filling a batch — loop condition ends it - } + let ctx = Arc::new(LayerCtx { slot_kmer, mat, n_parts, n_genomes, n_cols, k }); - // ── Local only: seed each family's own minorant presence (already - // open, no lookup) and collect this batch's cross-partition - // queries, grouped by destination partition. - let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); - let genome_mask: Vec = (0..batch_slots.len() * n_genomes).map(|_| AtomicU8::new(0)).collect(); + let mut offset = 0usize; + let batches: Vec<(usize, Vec<(usize, FamilyMask)>)> = minorant_slots + .chunks(FAMILY_BATCH) + .map(|chunk| { + let start = offset; + offset += chunk.len(); + (start, chunk.to_vec()) + }) + .collect(); - for (family_idx, &fslot) in batch_slots.iter().enumerate() { - let kmer = slot_kmer[fslot].expect("batch slot has a kmer"); - let mask = annex.get(fslot).expect("batch slot has a mask"); - let base = central_base(kmer, k); - for g in 0..n_cols { - if mat.carries(g, fslot) { - genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); - } - } - for other in kmer.central_canonical_neighbors() { - if other == kmer { - continue; // local — resolved above straight from `mat`, no lookup - } - let b = central_base(other, k); - if !mask.has(b) { - continue; - } - let dest = partition_of(other, n_parts); - outgoing[dest].push((other, family_idx, b)); - } - } + let n_workers = obisys::effective_parallelism(); + let capacity = 4; + let throttled = throttle(batches.into_iter(), n_workers).map(|t| SourceBatch { + start_family_idx: t.item.0, + slots: t.item.1, + _permit: t.guard, + }); - // ── Resolve this batch's queries, one contiguous sweep per - // destination partition — same rationale as `build_sibling_annex`'s - // `outgoing` grouping. - outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { - for &(variant, family_idx, base) in queries { - let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue }; - for (g, &present) in presence.iter().enumerate() { - if present { - genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); + let worker_ctx = Arc::clone(&ctx); + let pipe = obipipeline::make_pipe! { + FamData : SourceBatch => GeneratedBatch, + | { + move |batch: SourceBatch| -> GeneratedBatch { + let ctx = &worker_ctx; + let n = batch.slots.len(); + let mut masks = Vec::with_capacity(n); + let mut bases = Vec::with_capacity(n); + let mut genome_mask = vec![0u8; n * ctx.n_genomes]; + let mut outgoing: Vec> = (0..ctx.n_parts).map(|_| Vec::new()).collect(); + + // Pass 1: cheap, no matrix access — own base and this + // batch's cross-partition queries. Mask already in hand + // from the filter step, no second annex read. + for (i, &(fslot, mask)) in batch.slots.iter().enumerate() { + let kmer = ctx.slot_kmer[fslot].expect("minorant slot has a kmer"); + masks.push(mask); + let base = central_base(kmer, ctx.k); + bases.push(base); + for other in kmer.central_canonical_neighbors() { + if other == kmer { + continue; // local — resolved below straight from `mat`, no lookup + } + let b = central_base(other, ctx.k); + if !mask.has(b) { + continue; + } + let dest = partition_of(other, ctx.n_parts); + outgoing[dest].push((other, i, b)); } } - } - }); - for (family_idx, &fslot) in batch_slots.iter().enumerate() { - let mask = annex.get(fslot).expect("batch slot has a mask"); - for (g, dst) in scratch.iter_mut().enumerate() { - *dst = genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed); + // Pass 2: genome-major, not family-major — `mat` is stored + // one contiguous block per genome (column), slot as the + // offset within it (see `PartitionCache::find_presence_batch`'s + // docs for the full rationale). `batch.slots` is already + // sorted (a contiguous sub-range of the layer's sorted + // minorant-slot list), so this sweeps each column roughly + // in slot order instead of jumping between all `n_cols` + // columns once per family. + for g in 0..ctx.n_cols { + for (i, &(fslot, _)) in batch.slots.iter().enumerate() { + if ctx.mat.carries(g, fslot) { + genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i]; + } + } + } + + GeneratedBatch { start_family_idx: batch.start_family_idx, masks, genome_mask, outgoing, _permit: batch._permit } } - on_family(mask, &scratch); + } : Batch => Generated, + }; + + // Batches finish generation in whatever order their worker completes + // them, not submission order — replay in order through a buffer bounded + // by the throttle's own concurrency (`n_workers` batches can be in + // flight at once, so at most that many can be waiting here). + let mut pending: HashMap = HashMap::new(); + let mut next_expected = 0usize; + let mut scratch = vec![0u8; n_genomes]; + for generated in pipe.apply(throttled, n_workers, capacity) { + pending.insert(generated.start_family_idx, generated); + while let Some(batch) = pending.remove(&next_expected) { + let n = batch.masks.len(); + + // ── Resolve this one batch's cross-partition queries — the + // only place that touches `cache`. Parallel across partitions + // (one batch at a time, never several concurrently — see the + // module docs), each thread owning one partition's queries + // contiguously until this batch is done. + let genome_mask: Vec = batch.genome_mask.into_iter().map(AtomicU8::new).collect(); + batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { + cache.find_presence_batch(dest, queries, n_genomes, |i, base, g| { + genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); + }); + }); + + for i in 0..n { + for (g, dst) in scratch.iter_mut().enumerate() { + *dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed); + } + on_family(batch.masks[i], &scratch); + } + next_expected += n; } } + debug_assert_eq!(next_expected, total_families, "every batch must have been replayed"); Ok(()) } diff --git a/src/obikindex/src/siblings/stats.rs b/src/obikindex/src/siblings/stats.rs index 2e6009e0..0ca0caf9 100644 --- a/src/obikindex/src/siblings/stats.rs +++ b/src/obikindex/src/siblings/stats.rs @@ -1,9 +1,13 @@ +use std::sync::Arc; + +use obicompactvec::SiblingAnnex; use obikpartitionner::KmerPartition; use obisys::progress_bar; use crate::error::{OKIError, OKIResult}; use crate::index::KmerIndex; +use super::ANNEX_FILE_NAME; use super::cache::PartitionCache; use super::family_scan::scan_layer_families; @@ -30,6 +34,42 @@ pub struct SiblingAnnexStats { } impl KmerIndex { + /// The global family-size histogram alone (`SiblingAnnexStats::counts`, + /// no `per_genome`) — reads only the already-built annex (`mask.siblings()` + /// + `mask.is_minorant()`, 1 byte/slot, mmap'd), nothing else: no + /// `unitigs.bin` scan, no MPHF lookup, no `PartitionCache`. Cost is a + /// linear scan of one small file per layer, independent of the rest of + /// the index's size or how much of it is paged in. An earlier version + /// re-derived minorant-ness per slot (re-reading `unitigs.bin`, hashing + /// every k-mer through the MPHF again to reconstruct what + /// `build_sibling_annex` already knew) — sampling a real run showed + /// `MphfLayer::find` alone at 71% of wall-clock time for what should be + /// a near-instant four-bucket count. `mask.is_minorant()` is that same + /// fact, computed once at construction (see `build_layer_sibling_annex`) + /// and stored in the annex's own spare bits — free to read back. + /// [`sibling_annex_stats`](Self::sibling_annex_stats) computes the same + /// `counts` but pays the full per-genome cross-partition resolution + /// cost to get there too; use this when only the global histogram is + /// needed. Requires an annex built after the minorant flag was added — + /// re-run `build_sibling_annex` if this reads as all-zero on an older one. + pub fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> { + let layer_dirs = self.sibling_layer_dirs()?; + + let mut counts = [0u64; 4]; + for layer_dir in &layer_dirs { + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + if !mask.is_minorant() { + continue; // family tallied once, at its minorant + } + counts[mask.siblings() as usize] += 1; + } + } + + Ok(counts) + } + /// Tally the family-size distribution of an already-built annex /// (globally, and per genome), counting each family once (at its /// minorant slot). Errors if [`build_sibling_annex`] has not been run on @@ -53,7 +93,7 @@ impl KmerIndex { n_bits, ) .map_err(OKIError::Partition)?; - let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); let layer_dirs = self.sibling_layer_dirs()?; // One layer at a time, not parallelised across layers — see diff --git a/src/obikindex/src/siblings/tests.rs b/src/obikindex/src/siblings/tests.rs index 09a1c597..866464e5 100644 --- a/src/obikindex/src/siblings/tests.rs +++ b/src/obikindex/src/siblings/tests.rs @@ -120,7 +120,8 @@ fn sibling_annex_one_sibling_each() { // complement, since both start with "AA"), and raw(g1) < raw(g2) // (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is // the minorant, g2 is not. The mask is a family-wide value: both - // slots must read back the *same* mask (bits 1 and 2 set). + // slots must read back the *same* presence bits (1 and 2 set) — but + // only g1's slot should carry the stored minorant flag. let dir = tempdir().unwrap(); let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); @@ -132,14 +133,16 @@ fn sibling_annex_one_sibling_each() { let expected_mask = FamilyMask::EMPTY.with(1).with(2); let a = annex_info_for(&merged, g1_kmer); - assert_eq!(a, expected_mask, "AACCGCTTAAG"); + assert_eq!(a.bits(), expected_mask.bits(), "AACCGCTTAAG"); assert_eq!(a.siblings(), 1); assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant"); + assert!(a.is_minorant(), "g1's stored minorant flag should be set at build time"); let b = annex_info_for(&merged, g2_kmer); - assert_eq!(b, expected_mask, "AACCGGTTAAG"); + assert_eq!(b.bits(), expected_mask.bits(), "AACCGGTTAAG"); assert_eq!(b.siblings(), 1); assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant"); + assert!(!b.is_minorant(), "g2's stored minorant flag should not be set"); } #[test] @@ -185,6 +188,25 @@ fn sibling_annex_stats_counts_each_family_once_and_per_genome() { } } +#[test] +fn sibling_family_size_histogram_matches_full_stats_global_counts() { + // Same fixture as `sibling_annex_stats_counts_each_family_once_and_per_genome` + // — the cheap, annex-only histogram must agree with the `counts` half of + // the full (cross-partition) stats pass, without needing a `PartitionCache` + // at all. + let dir = tempdir().unwrap(); + let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); + let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); + let merged = merge_two(dir.path(), &g1, &g2); + merged.build_sibling_annex().expect("build_sibling_annex"); + + let histogram = merged.sibling_family_size_histogram().expect("sibling_family_size_histogram"); + assert_eq!(histogram, [0, 1, 0, 0], "one family of size 2, counted once"); + + let stats = merged.sibling_annex_stats().expect("sibling_annex_stats"); + assert_eq!(histogram, stats.counts, "must agree with the full stats pass's global counts"); +} + /// Exercises the four sibling-annex consumers that were rewritten to share /// `family_scan::scan_layer_families` (partition-grouped lookups instead of /// one lookup per family) — same one-sibling-each fixture as the tests diff --git a/src/obikmer/src/cmd/phylo/args.rs b/src/obikmer/src/cmd/phylo/args.rs index deedb5bd..aa294e6e 100644 --- a/src/obikmer/src/cmd/phylo/args.rs +++ b/src/obikmer/src/cmd/phylo/args.rs @@ -106,6 +106,17 @@ pub struct PhyloArgs { #[arg(long)] pub sibling_stats: bool, + /// Print just the global family-size histogram (1-4 members) of an + /// already-built annex — the `global` row `--sibling-stats` also + /// writes, but without the per-genome breakdown, so it skips + /// `--sibling-stats`'s cross-partition resolution entirely: reads only + /// the already-open annex mask and each layer's own `unitigs.bin`, cost + /// independent of the rest of the index. A quick sanity check that the + /// annex itself is sound, decoupled from `--sibling-stats`'s much + /// heavier per-genome pass. + #[arg(long)] + pub sibling_hist: 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 @@ -231,7 +242,8 @@ pub struct PhyloArgs { pub sankoff_cost_scale: f64, /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _rawsnp.csv, _rawsnp_counts.csv, + /// _siblings.csv, _sibling_hist.csv, _rawsnp.csv, + /// _rawsnp_counts.csv, /// _snp.fasta, _family_overlap.csv, /// _sankoff_matrix.csv, _sankoff_params.yaml, /// _sankoff.fasta, _sankoff.tnt, _sankoff.tcm, diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index d510822c..9b09a071 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -19,7 +19,7 @@ use tracing::info; pub use args::PhyloArgs; use family_overlap::{apply_min_shared_family_exclusion, write_family_overlap_csv}; use iqtree::write_iqtree; -use outputs::{write_raw_snp_counts_csv, write_raw_snp_distance_csv, write_sibling_stats_csv, write_snp_fasta, upgma_to_newick}; +use outputs::{write_raw_snp_counts_csv, write_raw_snp_distance_csv, write_sibling_hist_csv, write_sibling_stats_csv, write_snp_fasta, upgma_to_newick}; use phyg::write_sankoff_phyg; use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params}; use tnt::write_sankoff_tnt; @@ -135,6 +135,13 @@ pub fn run(args: PhyloArgs) { }); write_sibling_stats_csv(&stats, &labels, &args.output); } + if args.sibling_hist { + let counts = idx.sibling_family_size_histogram().unwrap_or_else(|e| { + eprintln!("error computing sibling family-size histogram: {e}"); + std::process::exit(1); + }); + write_sibling_hist_csv(&counts, &args.output); + } if args.raw_snp_distance { let mut result = idx.raw_snp_distance().unwrap_or_else(|e| { eprintln!("error computing raw SNP distance: {e}"); diff --git a/src/obikmer/src/cmd/phylo/outputs.rs b/src/obikmer/src/cmd/phylo/outputs.rs index 97b548fd..28b67a66 100644 --- a/src/obikmer/src/cmd/phylo/outputs.rs +++ b/src/obikmer/src/cmd/phylo/outputs.rs @@ -76,6 +76,25 @@ pub(super) fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: info!("raw single-copy SNP distance matrix → {path}"); } +// ── Global family-size histogram (annex-only, no per-genome pass) → CSV ──── + +pub(super) fn write_sibling_hist_csv(counts: &[u64; 4], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_sibling_hist.csv", p.display())) + .unwrap_or_else(|| "sibling_hist.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); + })); + writeln!(f, "size,count").unwrap(); + for (size, count) in counts.iter().enumerate() { + writeln!(f, "{},{count}", size + 1).unwrap(); + } + let total: u64 = counts.iter().sum(); + info!("family-size histogram → {path} (total {total} famil{})", + if total == 1 { "y" } else { "ies" }); +} + // ── Raw single-copy SNP distance → per-pair diagnostic counts ────────────── // // A pair table (one row per unordered genome pair), not a matrix: the ratio diff --git a/src/obilayeredmap/src/mphf_layer.rs b/src/obilayeredmap/src/mphf_layer.rs index 26611230..b9d19901 100644 --- a/src/obilayeredmap/src/mphf_layer.rs +++ b/src/obilayeredmap/src/mphf_layer.rs @@ -134,6 +134,31 @@ impl MphfLayer { } } + /// Reconstruct the canonical k-mer stored at `slot` — the inverse of + /// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot). + /// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives + /// `(chunk_id, rank)` directly (no MPHF hashing, no file scan), then + /// `unitigs.canonical_raw_kmer` is a direct-access read. `None` on + /// `Approx` layers — fingerprints alone can't recover a k-mer, and + /// falling back to a full sequential scan here would silently make one + /// call cost O(n); callers needing this on an `Approx` layer should + /// scan `unitigs.bin` themselves and decide how to handle that cost + /// explicitly. + pub fn kmer_at(&self, slot: usize) -> Option { + if slot >= self.n { + return None; + } + match &self.ev { + LayerEvidence::Exact { evidence, unitigs } | + LayerEvidence::Hybrid { evidence, unitigs, .. } => { + let (chunk_id, rank) = evidence.decode(slot); + let raw = unitigs.canonical_raw_kmer(chunk_id as usize, rank as usize); + Some(CanonicalKmer::from_raw_unchecked(raw)) + } + LayerEvidence::Approx { .. } => None, + } + } + pub fn n(&self) -> usize { self.n } } diff --git a/src/obiskio/src/unitig_index/reader.rs b/src/obiskio/src/unitig_index/reader.rs index 8e1b49c3..e5f0d978 100644 --- a/src/obiskio/src/unitig_index/reader.rs +++ b/src/obiskio/src/unitig_index/reader.rs @@ -157,6 +157,16 @@ impl UnitigFileReader { canonical_raw(self.raw_kmer(i, j), self.k) == query.raw() } + /// Canonical raw k-mer at position `j` of chunk `i` — [`raw_kmer`](Self::raw_kmer) + /// reduced to its canonical (`min(raw, revcomp)`) form. The reconstruction + /// half of [`verify_canonical_kmer`](Self::verify_canonical_kmer): that + /// method computes the same value internally just to compare it away: use + /// this one when the caller wants the k-mer itself, not a yes/no answer. + #[inline] + pub fn canonical_raw_kmer(&self, i: usize, j: usize) -> u64 { + canonical_raw(self.raw_kmer(i, j), self.k) + } + // ── Sequential iterators (O(n) running-offset cursor) ───────────────────── pub(crate) fn iter_chunks_sequential(&self) -> impl Iterator + '_ {