diff --git a/.gitignore b/.gitignore index 793b0a9..d2d43e5 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ data-stress ./**/*.json *.bin *.log +*.csv Betula_exilis--IGA-24-33 benchmark/genomes benchmark/simulated_data diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 6ff266b..2c21a3c 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -65,6 +65,50 @@ reverse-complement (`m -> k-1-m = m`, base complemented). A transition maps to a transition, a transversion to a transversion — the transition/transversion split is well-defined in canonical space. +### Definitions: family, and the canonical form of a family + +**Family.** The family of a k-mer `x` is the set of (up to) 4 k-mers sharing +`x`'s `2m` flanking bases, differing only at the central base `m`. Membership +is a property of the flank pattern, not of `x` itself: any of the 4 possible +central substitutions belongs to the same family. + +**`central_canonical_neighbors()`** (`obikseq`, `CanonicalKmerOf::central_canonical_neighbors`) +generates all 4 members from any one of them (observed or not), each +independently canonicalised (`.canonical()`, i.e. `min(kmer, revcomp(kmer))`). +This independent canonicalisation is necessary because a central substitution +can flip which orientation is lexicographically smaller — two members of the +same family can end up canonicalised in *different* orientations. Despite +that, the **set** of 4 resulting canonical k-mers is invariant: calling +`central_canonical_neighbors()` on any member of a family — present in the +index or not — yields the same 4 values. This is relied upon throughout the +rest of this document. + +**Canonical form of a family.** Because orientation can differ member to +member, "which of the 4 is the reference" cannot be defined relative to +*whichever member happened to be visited first*, nor relative to the +minorant (see below) — both are data-dependent (they depend on what is +actually observed), so using either as the reference would make the +reference itself vary depending on what happens to be present in a given +index. Instead: **the canonical form of a family is, by definition, the +member whose own central base — read in its own already-canonical +orientation — is `A`.** This is well-defined for every family, computed +purely from the flank pattern, whether or not that specific member (or any +member at all) is actually observed anywhere in the index. Concretely: call +`central_canonical_neighbors()` on any member (observed or not) to get the +family's 4 canonical forms; the one among them whose own centre nucleotide is +`A` is the family's canonical form. The other 3 (`C`, `G`, `T`) are labelled +relative to *that* fixed reference, not relative to the calling member's own +orientation. + +**Consequence for the minorant.** With this fixed A-referenced labelling, +`minorant` (the smallest raw encoding among the family's *observed* members, +introduced further below) becomes directly computable rather than needing to +be tracked as extra state: regenerate the family's 4 canonical forms from +any member's own k-mer (cheap, no lookup), compare the raw encodings of +whichever are marked present, and take the smallest. No separate stored bit +is required — see Step 2b below, where this replaces the earlier +minorant-bit design. + ## Locus eligibility: raw definition vs. paralogy filter For each k-mer `x` observed in genome A (source, one MPHF slot; the 3 @@ -482,31 +526,58 @@ here. ### Step 2b — sibling-count / minorant annex (consolidated plan) -Scope: only the precursor annex (sibling count 0-3 per slot, minorant -decided on demand) — not the SNP tally itself, whose Step 2 sweep remains -unresolved above. This piece is simpler than the sweep, because it writes to -an independent per-slot value, not a shared cross-k-mer accumulator, so it -needs no dedup/ownership logic at all at this stage. +Scope: only the precursor annex — not the SNP tally itself, whose Step 2 +sweep remains unresolved above. This piece is simpler than the sweep, +because it writes to an independent per-slot value, not a shared +cross-k-mer accumulator, so it needs no dedup/ownership logic at all at this +stage. + +**Revised annex encoding — 4-bit presence mask, not 3-bit (minorant + +count).** Superseded after settling the "canonical form of a family" +definition above. The 3-bit design (1 minorant bit + 2-bit sibling count, +§ below, kept for the historical record) had two problems: it discards +*which* variants are present (only how many), so any future consumer +(the SNP sweep, or a stats pass — see below) that needs to know which bases +exist still has to regenerate and blindly re-query all 3 candidates; and +the minorant bit's meaning was tied to whichever member was visited, not to +a fixed reference. Storing instead a **4-bit mask** — one bit per base +(A/C/G/T), set iff that member of the family (labelled relative to the +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 + 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. +- **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. +- The all-zero value (no base present at all) is still logically + unreachable as a real result — the slot's *own* base is always present in + its own family — so it remains available as a free "not yet computed" + sentinel, exactly as before. 1. **Primitive.** Reuse `central_canonical_neighbors()` from Step 0 - unchanged — the 3 canonicalised central-substitution variants of a k-mer. -2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 3-bit- - per-slot packed array, 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). 3 bits, not 2: - revised to also store minorant status alongside sibling count, since it - comes for free from the same lookups (point 3 below) — 5 real states + unchanged — the 3 canonicalised central-substitution variants of a k-mer + (plus the identity, i.e. all 4 members of the family — see "Definitions" + above). +2. **New annex type** (`obicompactvec`, alongside `bitmatrix.rs`): a 4-bit- + 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). +
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 (not-minorant; minorant with 0/1/2/3 siblings) fit in 3 bits (8 states, - 3 unused). This lets the future SNP sweep discard a non-minorant slot - **instantly**, with no lookup at all, instead of having to regenerate and - look up its siblings just to rediscover it isn't the designated writer — - moving that cost into this one-time, cached pass instead of repeating it - on every future sweep. The otherwise-unreachable combination - "not-minorant + 0 siblings" (impossible: 0 siblings always implies - minorant, see below) doubles as a free **"not yet computed" sentinel** — - annex files for all partitions/layers can be pre-initialised to this - value before the computation pass runs, distinguishing genuinely-computed - 0-sibling slots from not-yet-processed ones with no extra storage. + 3 unused). This let the future SNP sweep discard a non-minorant slot + instantly, with no lookup at all. The otherwise-unreachable combination + "not-minorant + 0 siblings" (0 siblings always implies minorant) doubled + as the "not yet computed" sentinel. Replaced by the 4-bit mask above, + which subsumes this benefit (minorant still derivable, now for free at + read time rather than stored) while also fixing the "which variant" + blindness. +
3. **Computation pass** (`obikindex`, new `siblings.rs`): **one `obipipeline` run per layer, iterated sequentially over the index's layers** — settled after two false starts, worth recording both. diff --git a/src/obicompactvec/src/lib.rs b/src/obicompactvec/src/lib.rs index 9a136bb..1f85a49 100644 --- a/src/obicompactvec/src/lib.rs +++ b/src/obicompactvec/src/lib.rs @@ -19,7 +19,7 @@ pub use builder::PersistentCompactIntVecBuilder; pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask}; pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix}; pub use layer_meta::LayerMeta; -pub use siblingannex::{SiblingAnnex, SiblingAnnexBuilder, SiblingInfo}; +pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter}; pub use tempbitvec::{TempBitVec, TempBitVecBuilder}; pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; diff --git a/src/obicompactvec/src/siblingannex.rs b/src/obicompactvec/src/siblingannex.rs index 94cb5cc..0de9733 100644 --- a/src/obicompactvec/src/siblingannex.rs +++ b/src/obicompactvec/src/siblingannex.rs @@ -1,27 +1,33 @@ -//! Sibling-count / minorant annex: a compact, read-only-after-build, per-slot +//! Family presence-mask annex: a compact, read-only-after-build, per-slot //! derived value used by the central-position SNP distance estimator (see -//! `docmd/theory/evolutionary_distances.md`, "Step 2b"). +//! `docmd/theory/evolutionary_distances.md`, "Step 2b" and "Definitions: +//! family, and the canonical form of a family"). //! -//! One byte is stored per MPHF slot of a partition/layer, encoding two -//! independent facts about the slot's k-mer's "family" (the up-to-4 k-mers -//! sharing the same flanks, differing only at the central base), both -//! properties of the whole current multi-genome index, not of any one -//! genome: +//! One byte is stored per MPHF slot of a partition/layer, its low 4 bits +//! encoding a **presence mask** for the slot's k-mer's "family" (the up to 4 +//! k-mers sharing the same flanks, differing only at the central base): +//! bit `b` (`b` = 0..3, in the fixed A/C/G/T = 0/1/2/3 encoding already used +//! for a single nucleotide) is set iff the family member whose *own* central +//! base — in its own canonical orientation — is `b`, is observed anywhere in +//! the current multi-genome index. This is a property of the whole index, +//! not of any one genome. //! -//! - bit 0: `minorant` — is this k-mer's canonical encoding the smallest -//! among the family members actually observed in the index? -//! - bits 1-2: `siblings` — how many *other* family members (0-3) are -//! observed anywhere in the index. +//! Both facts the earlier (superseded) 3-bit design stored explicitly are +//! derived from the mask instead, not stored: +//! - sibling count = `popcount(mask) - 1`; +//! - minorant = regenerate the family's 4 canonical forms from the slot's +//! own k-mer (`CanonicalKmerOf::central_canonical_neighbors`, cheap, no +//! lookup), compare the raw encodings of whichever are set in the mask, +//! take the smallest — see `obikindex::siblings`. //! -//! Byte value 0 (`minorant = false`, `siblings = 0`) is logically -//! unreachable as a real result (0 siblings always implies minorant — see -//! the design doc) and is reused as the "not yet computed" sentinel: annex -//! files are pre-initialised to all-zero, and a real value is only ever -//! written once, by the computation pass. +//! Mask value 0 is logically unreachable as a real result (a slot's own base +//! is always present in its own family) and is reused as the "not yet +//! computed" sentinel: annex files are pre-initialised to all-zero, and a +//! real value is only ever written once, by the computation pass. //! -//! Deliberately simpler than a true 3-bit pack (1 byte/slot instead of 3 -//! bits/slot): correctness and simplicity first: for a first implementation. -//! Packing to 3 bits/slot is a pure storage-density follow-up, not a +//! Deliberately simpler than a true 4-bit pack (1 byte/slot instead of 4 +//! bits/slot): correctness and simplicity first, for a first implementation. +//! Packing to 4 bits/slot is a pure storage-density follow-up, not a //! behavioural change, left for later. use std::fs::{File, OpenOptions}; @@ -35,30 +41,70 @@ 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; -/// Decoded value of one slot's annex entry. +/// 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. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct SiblingInfo { - pub minorant: bool, - pub siblings: u8, // 0..=3 -} +pub struct FamilyMask(u8); + +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 + /// via repeated [`with`](Self::with) calls before storing it. + pub const EMPTY: FamilyMask = FamilyMask(0); + + /// Set bit `base` (0=A, 1=C, 2=G, 3=T). + #[inline] + pub fn with(self, base: u8) -> Self { + debug_assert!(base < 4, "base out of range: {base}"); + FamilyMask(self.0 | (1 << base)) + } + + /// Is the member with central base `base` (0..3) present? + #[inline] + pub fn has(self, base: u8) -> bool { + debug_assert!(base < 4, "base out of range: {base}"); + self.0 & (1 << base) != 0 + } + + /// Number of family members observed anywhere in the index (1..=4). + #[inline] + pub fn family_size(self) -> u32 { + self.0.count_ones() + } + + /// Number of *other* members observed (0..=3) — `family_size() - 1`. + #[inline] + pub fn siblings(self) -> u32 { + 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. + #[inline] + pub fn bits(self) -> u8 { + self.0 + } + + /// Construct from a raw bitmask (only the low 4 bits are kept). + #[inline] + pub fn from_bits(bits: u8) -> Self { + FamilyMask(bits & 0b1111) + } -impl SiblingInfo { #[inline] fn encode(self) -> u8 { - (self.siblings << 1) | (self.minorant as u8) + self.0 } #[inline] fn decode(byte: u8) -> Option { if byte == 0 { - // The unreachable "not minorant + 0 siblings" combination — - // reserved as the "not yet computed" sentinel. + // Unreachable for a real result — reserved as the "not yet + // computed" sentinel. return None; } - Some(SiblingInfo { - minorant: byte & 1 != 0, - siblings: (byte >> 1) & 0b11, - }) + Some(FamilyMask(byte & 0b1111)) } } @@ -91,8 +137,8 @@ impl SiblingAnnex { pub fn is_empty(&self) -> bool { self.n == 0 } /// `None` means the slot has not (yet) been computed — see module docs. - pub fn get(&self, slot: usize) -> Option { - SiblingInfo::decode(self.mmap[HEADER_SIZE + slot]) + pub fn get(&self, slot: usize) -> Option { + FamilyMask::decode(self.mmap[HEADER_SIZE + slot]) } } @@ -124,17 +170,17 @@ impl SiblingAnnexBuilder { pub fn len(&self) -> usize { self.n } pub fn is_empty(&self) -> bool { self.n == 0 } - pub fn get(&self, slot: usize) -> Option { - SiblingInfo::decode(self.mmap[HEADER_SIZE + slot]) + pub fn get(&self, slot: usize) -> Option { + FamilyMask::decode(self.mmap[HEADER_SIZE + slot]) } - pub fn set(&mut self, slot: usize, info: SiblingInfo) { + pub fn set(&mut self, slot: usize, mask: FamilyMask) { // Redundant concurrent writes from independent recomputation paths // converge to the same encoded byte for a given slot, so a plain // store here is safe even without external synchronisation, as long // as the byte write itself is atomic (true for a single aligned // byte on every platform this project targets). - self.mmap[HEADER_SIZE + slot] = info.encode(); + self.mmap[HEADER_SIZE + slot] = mask.encode(); } pub fn close(self) -> io::Result<()> { self.mmap.flush() } @@ -163,36 +209,37 @@ mod tests { } #[test] - fn roundtrip_all_valid_states() { + fn roundtrip_all_valid_masks() { let dir = tempdir().unwrap(); let path = dir.path().join("test.psib"); - let mut builder = SiblingAnnexBuilder::new(5, &path).unwrap(); + let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap(); - let cases = [ - SiblingInfo { minorant: true, siblings: 0 }, - SiblingInfo { minorant: true, siblings: 1 }, - SiblingInfo { minorant: true, siblings: 2 }, - SiblingInfo { minorant: true, siblings: 3 }, - SiblingInfo { minorant: false, siblings: 2 }, + let masks = [ + FamilyMask::EMPTY.with(0), // just A: family size 1 + FamilyMask::EMPTY.with(0).with(3), // A + T: size 2 + FamilyMask::EMPTY.with(1).with(2).with(3), // C+G+T: size 3 + FamilyMask::EMPTY.with(0).with(1).with(2).with(3), // all 4 ]; - for (slot, info) in cases.iter().enumerate() { - builder.set(slot, *info); + for (slot, mask) in masks.iter().enumerate() { + builder.set(slot, *mask); } let annex = builder.finish().unwrap(); - for (slot, info) in cases.iter().enumerate() { - assert_eq!(annex.get(slot), Some(*info)); + for (slot, mask) in masks.iter().enumerate() { + assert_eq!(annex.get(slot), Some(*mask)); } + assert_eq!(annex.get(0).unwrap().siblings(), 0); + assert_eq!(annex.get(1).unwrap().siblings(), 1); + assert_eq!(annex.get(2).unwrap().siblings(), 2); + assert_eq!(annex.get(3).unwrap().siblings(), 3); + assert_eq!(annex.get(3).unwrap().family_size(), 4); } #[test] - fn not_minorant_zero_siblings_is_unreachable_via_set_and_decodes_as_sentinel() { - // Documented invariant, not enforced by the type: callers must never - // construct this combination. If they do, it is indistinguishable - // from "not computed" — exercised here to pin the behaviour down. - let dir = tempdir().unwrap(); - let path = dir.path().join("test.psib"); - let mut builder = SiblingAnnexBuilder::new(1, &path).unwrap(); - builder.set(0, SiblingInfo { minorant: false, siblings: 0 }); - assert_eq!(builder.get(0), None); + fn has_reflects_individual_bits() { + let mask = FamilyMask::EMPTY.with(0).with(2); + assert!(mask.has(0)); + assert!(!mask.has(1)); + assert!(mask.has(2)); + assert!(!mask.has(3)); } } diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index 4638bde..e80bd67 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -1,39 +1,49 @@ -//! Sibling-count / minorant annex construction. +//! Family presence-mask annex construction. //! -//! See `docmd/theory/evolutionary_distances.md`, "Step 2b — sibling-count / -//! minorant annex", for the full design discussion this implements. +//! See `docmd/theory/evolutionary_distances.md`, "Definitions: family, and +//! the canonical form of a family" and "Step 2b", for the full design +//! discussion this implements. //! //! For each distinct k-mer of each layer of the (already built/merged) -//! index, computes two facts about its "family" (the up to 4 k-mers sharing -//! its flanks, differing only at the central base — well-defined for odd -//! k), both properties of the whole current multi-genome index rather than -//! of any one genome: -//! - how many *other* family members (0-3) are observed anywhere in the -//! index; -//! - whether this k-mer is the "minorant" of its family — the smallest -//! canonical encoding among the members actually observed. +//! index, computes a 4-bit presence mask for its "family" (the up to 4 +//! k-mers sharing its flanks, differing only at the central base — +//! well-defined for odd k): bit `b` set iff the family member whose own +//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere +//! in the current multi-genome index — a property of the whole index, not +//! of any one genome. Sibling count and minorant are *derived* from the +//! mask by callers, not stored (see `FamilyMask` and +//! [`sibling_annex_stats`](KmerIndex::sibling_annex_stats) below). //! -//! 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. +//! Per layer, an `obipipeline` `Flat` stage (throttled — see +//! `obipipeline::throttle`) generates each k-mer's 3 central variants, +//! interleaved across many in-flight k-mers by the scheduler's shared +//! worker pool rather than processed on a single thread. The actual +//! cross-partition lookup, though, reuses +//! `KmerPartition::query_partition_with` — the same partition-batching +//! mechanism `obikmer query` already uses (open a partition's files once, +//! answer a whole batch of queries against it) — rather than a per-item +//! pipeline stage: an earlier per-item design reopened/re-mmap'd every +//! target partition's files on every single lookup, which was fine at +//! toy scale but manifested as ~90% system time against a real index. +//! See `docmd/theory/evolutionary_distances.md`, Step 2b, "Mechanism". use std::path::Path; +use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; +use rayon::prelude::*; + use obicompactvec::{ - PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, SiblingInfo, + FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex, SiblingAnnexBuilder, }; use obikpartitionner::KmerPartition; +use obipipeline::ThrottleGuard; use obikseq::{CanonicalKmer, Minimizer}; use obilayeredmap::{MphfLayer, OLMError}; use obilayeredmap::meta::PartitionMeta; use obiskbuilder::rolling_stat::RollingStat; use obiskio::UnitigFileReader; +use obisys::progress_bar; use crate::error::{OKIError, OKIResult}; use crate::index::KmerIndex; @@ -48,6 +58,26 @@ fn olm_to_ok(e: OLMError) -> OKIError { } } +/// Central-position base of a canonical k-mer, in the fixed 0=A/1=C/2=G/3=T +/// encoding — the mask's bit index. `k` must be odd (project invariant). +#[inline] +fn central_base(kmer: CanonicalKmer, k: usize) -> u8 { + kmer.nucleotide((k - 1) / 2) +} + +/// Is `kmer` the minorant of its family, given the family's presence mask? +/// Regenerates the family's 4 canonical forms from `kmer` itself (cheap, no +/// lookup — see the design doc's "Definitions" section for why this is +/// always safe: the set of 4 forms is invariant regardless of which member +/// you start from), and compares the raw encodings of whichever are marked +/// present in `mask`. +fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool { + kmer.central_canonical_neighbors().into_iter().all(|other| { + other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw() + }) +} + + /// Minimiser of a single, isolated canonical k-mer (not part of a streamed /// sequence). `RollingStat` computes minimisers incrementally along a /// sequence; this feeds one k-mer's bases through a fresh instance to get @@ -74,79 +104,144 @@ fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize { (lone_kmer_minimizer(kmer).seq_hash() & mask) as usize } -/// Running reconciliation state for one source k-mer, initialised to the -/// trivial "no siblings observed yet" state and folded incrementally (in any -/// order — commutative) as query answers come back. -#[derive(Clone, Copy)] -struct RunningState { - minorant: bool, - siblings: u8, -} - -impl Default for RunningState { - fn default() -> Self { - RunningState { minorant: true, siblings: 0 } - } -} - // ── 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)] +/// MPHF slot — the pipeline's source item. Carries a throttle slot (shared, +/// `Arc`-wrapped so it can be cloned into each of the (up to 3) items this +/// one fans out into via the `Flat` stage) that is only released once every +/// one of those descendants has been fully processed — see the module docs' +/// "Throttling" note for why this is required, not optional, once a `Flat` +/// stage is in the pipeline. struct SourceItem { slot: usize, kmer: CanonicalKmer, + _permit: Arc, } /// 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)] +/// routed to its destination partition and carrying its own central base +/// (0=A/1=C/2=G/3=T) — the mask bit it will set on a hit. Carries a clone of +/// the source item's throttle permit. struct VariantQuery { source_slot: usize, dest_partition: usize, variant: CanonicalKmer, - smaller: bool, + base: u8, + _permit: Arc, } -/// 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; - } +/// Every partition's already-open MPHF layers, built **once** for the whole +/// `build_sibling_annex` run and shared (read-only) across every lookup, in +/// every source layer, for the rest of the run — not reopened/re-mmap'd per +/// query, nor per source layer. +/// +/// Confirmed necessary by sampling a real run: routing lookups through +/// `KmerPartition::query_partition_with` (the same batching `obikmer query` +/// uses) still reopens+re-mmaps every target partition's files on every +/// call, and it is called once per destination partition **per source +/// layer** — for an index with many layers this repeats the same +/// `MphfLayer::open`/`Evidence::open`/`PersistentBitMatrix::open` work over +/// and over. Parallelising those calls (see the gather step below) spread +/// the redundant work across more cores but did not reduce it: sampling +/// showed Rayon workers spending their time inside repeated `open()` +/// syscalls, not computation. This cache amortises that cost to once per +/// partition for the entire run, regardless of how many source layers or +/// lookups follow. +/// A cached layer's opened presence/count matrix, alongside its `MphfLayer`. +enum Mat { + Count(PersistentCompactIntMatrix), + Presence(PersistentBitMatrix), +} + +impl Mat { + fn n_cols(&self) -> usize { + match self { + Mat::Count(m) => m.n_cols(), + Mat::Presence(m) => m.n_cols(), } } - false + fn carries(&self, g: usize, slot: usize) -> bool { + match self { + Mat::Count(m) => m.col_view(g).get(slot) != 0, + Mat::Presence(m) => m.get(g, slot) != 0, + } + } +} + +struct PartitionCache { + /// `layers[partition][layer]` = that partition's opened MPHF layers, + /// paired 1:1 with `mats[partition][layer]`; empty if the partition + /// directory doesn't exist. Used by both [`KmerIndex::build_sibling_annex`] + /// (`layers` only) and [`KmerIndex::sibling_annex_stats`] (both). + layers: Vec>, + mats: Vec>, +} + +impl PartitionCache { + fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult { + let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); + let built: Vec<(Vec, Vec)> = (0..n_parts) + .into_par_iter() + .map(|part| -> OKIResult<(Vec, Vec)> { + let index_dir = partition.part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + pb.inc(1); + return Ok((Vec::new(), Vec::new())); + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + let mut layers = Vec::with_capacity(meta.n_layers); + let mut mats = Vec::with_capacity(meta.n_layers); + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) else { continue }; + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + PersistentCompactIntMatrix::open(&layer_dir).ok().map(Mat::Count) + } else { + PersistentBitMatrix::open(&layer_dir).ok().map(Mat::Presence) + }; + let Some(mat) = mat else { continue }; + layers.push(mphf); + mats.push(mat); + } + pb.inc(1); + Ok((layers, mats)) + }) + .collect::>>()?; + pb.finish_and_clear(); + let (layers, mats) = built.into_iter().unzip(); + Ok(Self { layers, mats }) + } + + /// Existence-only lookup of `variant` in partition `dest_partition`: + /// tries each of the partition's already-open layers in turn, stopping + /// at the first hit. + fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool { + self.layers + .get(dest_partition) + .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. + 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()); + } + } + None + } } impl KmerIndex { @@ -169,42 +264,50 @@ impl KmerIndex { 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)?, - ); + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep"); + let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?); + + let pb = progress_bar("sibling_annex", n_parts as u64, "partitions"); + let mut total_slots: u64 = 0; for part in 0..n_parts { let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); if !index_dir.exists() { + pb.inc(1); continue; } let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + let mut part_slots: u64 = 0; for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); - self.build_layer_sibling_annex(&layer_dir, n_parts, &partition)?; + part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?; } + total_slots += part_slots; + pb.inc(1); + pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)")); } + pb.finish_and_clear(); + tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions"); Ok(()) } + /// Returns the number of distinct k-mers (annex slots) processed, for + /// progress reporting. fn build_layer_sibling_annex( &self, layer_dir: &Path, n_parts: usize, - partition: &Arc, - ) -> OKIResult<()> { + cache: &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)?; @@ -220,22 +323,61 @@ impl KmerIndex { } } - let sources: Vec = slot_kmer - .iter() - .enumerate() - .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| SourceItem { slot, kmer })) - .collect(); + let k = self.kmer_size(); - // ── obipipeline: Flat (generate variants) -> Transform (lookup) ───── + // ── Reconciliation state, initialised with each slot's own base — + // that member is trivially present, no lookup needed. Built before + // the pipeline runs, from the same enumeration, since `sources` + // below is consumed as a throttled iterator, not collected. + // `AtomicU8`, not `FamilyMask`, because the gather phase below + // parallelises across destination partitions (independent + // `query_partition_with` calls, safe to run concurrently) and their + // `Found` hits can land on arbitrary, possibly-shared slots — a + // lock-free `fetch_or` avoids needing any synchronisation beyond + // that. ───────────────────────────────────────────────────────── + let mask: Vec = (0..n_slots).map(|_| AtomicU8::new(0)).collect(); + for (slot, kmer) in slot_kmer.iter().enumerate().filter_map(|(s, k)| k.map(|k| (s, k))) { + mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed); + } + + // ── obipipeline: Flat stage generates variants only — the actual + // cross-partition lookup reuses `KmerPartition::query_partition_with` + // (the same batching mechanism `obikmer query` already uses: open a + // partition's files once, answer a whole batch of queries against + // it) instead of one lookup per pipeline item. A per-item lookup + // (tried first) reopened/re-mmap'd every target partition's files on + // every single variant — fine at the scale of a handful of test + // k-mers, but with billions of lookups against a real index this + // manifested as ~90% system time, observed in practice. ─────────── let n_workers = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4); let capacity = 256; - let partition_for_lookup = Arc::clone(partition); + // Throttling is not optional once a `Flat` stage is in the pipeline + // (see `obipipeline::throttle`'s docs): without it, every worker can + // become a simultaneous `Flat` producer, saturate the shared output + // channel, and deadlock against the scheduler's own dispatch loop — + // also observed in practice. The permit acquired here for a source + // k-mer is held (via the `Arc`-shared guard carried through + // `SourceItem` -> `VariantQuery`) until every one of its (up to 3) + // descendants has been read out of the pipeline by the accumulation + // loop below, not just until the `Flat` stage itself returns. + let sources: Vec<(usize, CanonicalKmer)> = slot_kmer + .iter() + .enumerate() + .filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer))) + .collect(); + let throttled = obipipeline::throttle(sources.into_iter(), n_workers).map(|t| SourceItem { + slot: t.item.0, + kmer: t.item.1, + _permit: Arc::new(t.guard), + }); + let pipe = obipipeline::make_pipe! { - SibData : SourceItem => AnswerMsg, + SibData : SourceItem => VariantQuery, || { move |item: SourceItem| -> Vec { let kmer = item.kmer; + let permit = item._permit; kmer.central_canonical_neighbors() .into_iter() .filter(|variant| *variant != kmer) @@ -243,88 +385,112 @@ impl KmerIndex { source_slot: item.slot, dest_partition: partition_of(variant, n_parts), variant, - smaller: variant.raw() < kmer.raw(), + base: central_base(variant, k), + _permit: Arc::clone(&permit), }) .collect::>() } } : 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, }; - // ── 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]; - 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; + // ── Group generated variants by destination partition. `cache` + // holds every partition already mmap'd (no more `open()` cost), but + // `mmap` pages are still loaded on demand and can be evicted — a + // lookup is not free just because the file isn't reopened. Grouping + // keeps one partition's pages hot while its whole batch is resolved, + // instead of faulting pages in and out as lookups jump between + // partitions in whatever order the `Flat` stage happens to produce + // them. The throttle permit drops here, once accumulated. ───────── + let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); + for vq in pipe.apply(throttled, n_workers, capacity) { + outgoing[vq.dest_partition].push((vq.variant, vq.source_slot, vq.base)); + } + + // ── Resolve each partition's batch against the cache in one + // contiguous pass; parallelised across partitions (independent, + // read-only) so this keeps using multiple cores without giving up + // the per-partition locality above. ───────────────────────────── + outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { + for &(variant, source_slot, base) in queries { + if cache.find(dest, variant) { + mask[source_slot].fetch_or(1 << base, Ordering::Relaxed); } } - } + }); // ── Write the layer's annex file ───────────────────────────────────── let annex_path = layer_dir.join(ANNEX_FILE_NAME); let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?; - for (slot, st) in state.iter().enumerate() { + 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, SiblingInfo { minorant: st.minorant, siblings: st.siblings }); + builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed))); } builder.close()?; - Ok(()) + Ok(n_slots as u64) } } -/// Distribution of sibling counts (0-3), read back from an already-built +/// Distribution of family sizes (1-4), 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. +/// +/// Every count here is **per family, not per slot**: a family with `F` +/// members occupies `F` annex slots (one per observed member), all sharing +/// the same mask. Counting every slot would count each family up to 4 +/// times over; only the minorant's slot is tallied (minorant is derived on +/// the fly — see `is_minorant` — not stored, but cheap: no lookup, pure +/// bit arithmetic on already-in-hand data). #[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. + /// `counts[s]` = number of *families* of size `s + 1` (index 0 = size 1, + /// i.e. 0 siblings, ... index 3 = size 4, i.e. 3 siblings). 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. + /// `per_genome[g][s]` = number of families of size `s + 1` for which + /// genome `g` (index into `KmerIndex::meta().genomes`) carries at least + /// one member. 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. + /// 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 + /// 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 k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; - let mut stats = SiblingAnnexStats { - per_genome: vec![[0u64; 4]; n_genomes], - ..Default::default() - }; + // Same whole-run cache as `build_sibling_annex` — see its docs for + // why re-opening per lookup (or per call to a batching helper) is + // not good enough on a real index. + 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)?; + // Gather the (partition, layer) pairs to process — cheap metadata + // reads only, checking every annex file exists up front so a + // missing one is reported before any real work starts. + 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); @@ -334,45 +500,111 @@ impl KmerIndex { annex_path.display() ))); } - let annex = SiblingAnnex::open(&annex_path)?; - let use_counts = with_counts && layer_dir.join("counts").exists(); + layer_dirs.push(layer_dir); + } + } - // 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)?) + // One layer's worth of work, parallelised across layers with Rayon + // — independent, read-only, each producing its own partial tally + // merged at the end. + let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers"); + let partials: Vec = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult { + let mut stats = SiblingAnnexStats { + per_genome: vec![[0u64; 4]; n_genomes], + ..Default::default() }; - let n_cols = match &mat { - Mat::Count(m) => m.n_cols(), - Mat::Presence(m) => m.n_cols(), + + 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))?; + + // Need each slot's own k-mer to derive minorant — same + // enumeration as construction. + 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); + } } - .min(n_genomes); + + 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); 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; + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // this family is tallied at its minorant's slot only } + let s = mask.siblings() as usize; + stats.counts[s] += 1; + + // "Genome g represents this family" means g carries + // *any* of its members, not just the minorant's own — + // start from the minorant's own presence (already + // open, no lookup) and OR in every other present + // member's presence vector, resolved against the + // whole-run cache (no I/O) — exactly `mask.siblings()` + // of them, the mask tells us precisely which to fetch. + let mut carries = vec![false; n_cols]; 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, - }; + carries[g] = mat.carries(g, slot); + } + for other in kmer.central_canonical_neighbors() { + if other == kmer { + continue; + } + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let dest = partition_of(other, n_parts); + if let Some(other_presence) = cache.find_presence(dest, other, n_genomes) { + for (g, &present) in other_presence.iter().enumerate() { + if present { + carries[g] = true; + } + } + } + } + + for (g, &carried) in carries.iter().enumerate() { if carried { stats.per_genome[g][s] += 1; } } } + + pb.inc(1); + Ok(stats) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut stats = SiblingAnnexStats { + per_genome: vec![[0u64; 4]; n_genomes], + ..Default::default() + }; + for part in partials { + for s in 0..4 { + stats.counts[s] += part.counts[s]; + } + for g in 0..n_genomes { + for s in 0..4 { + stats.per_genome[g][s] += part.per_genome[g][s]; + } } } - Ok(stats) } } @@ -442,7 +674,7 @@ mod tests { /// Read back the annex entry for a given canonical k-mer from the merged /// index's (single) partition/layer, asserting it was found at all. - fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> SiblingInfo { + fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask { let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR); let meta = PartitionMeta::load(&index_dir).unwrap(); for l in 0..meta.n_layers { @@ -474,24 +706,33 @@ mod tests { fn sibling_annex_one_sibling_each() { // k=11, centre = index 5 (0-based). Two genomes, each exactly one // k-mer, sharing every base except the centre: - // g1 = "AACCGCTTAAG" (centre 'C') - // g2 = "AACCGGTTAAG" (centre 'G') + // g1 = "AACCGCTTAAG" (centre 'C', base index 1) + // g2 = "AACCGGTTAAG" (centre 'G', base index 2) // Hand-verified: both stay forward-oriented under canonicalisation // (each is lexicographically smaller than its own reverse // 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, and each is the other's one sibling. + // 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). 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 a = annex_info_for(&merged, canonical(b"AACCGCTTAAG")); - assert_eq!(a, SiblingInfo { minorant: true, siblings: 1 }, "AACCGCTTAAG"); + let g1_kmer = canonical(b"AACCGCTTAAG"); + let g2_kmer = canonical(b"AACCGGTTAAG"); + let expected_mask = FamilyMask::EMPTY.with(1).with(2); - let b = annex_info_for(&merged, canonical(b"AACCGGTTAAG")); - assert_eq!(b, SiblingInfo { minorant: false, siblings: 1 }, "AACCGGTTAAG"); + let a = annex_info_for(&merged, g1_kmer); + assert_eq!(a, expected_mask, "AACCGCTTAAG"); + assert_eq!(a.siblings(), 1); + assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant"); + + let b = annex_info_for(&merged, g2_kmer); + assert_eq!(b, expected_mask, "AACCGGTTAAG"); + assert_eq!(b.siblings(), 1); + assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant"); } #[test] @@ -504,7 +745,36 @@ mod tests { let merged = merge_two(dir.path(), &g1, &g2); merged.build_sibling_annex().expect("build_sibling_annex"); - let info = annex_info_for(&merged, canonical(b"GATTACAGATC")); - assert_eq!(info, SiblingInfo { minorant: true, siblings: 0 }, "GATTACAGATC"); + let kmer = canonical(b"GATTACAGATC"); + let mask = annex_info_for(&merged, kmer); + assert_eq!(mask.siblings(), 0, "GATTACAGATC"); + assert_eq!(mask.family_size(), 1); + assert!(is_minorant(kmer, mask, K)); + } + + #[test] + fn sibling_annex_stats_counts_each_family_once_and_per_genome() { + // Reuses the one-sibling-each fixture: a single family of size 2 + // (g1's centre-C form + g2's centre-G form), each genome carrying + // exactly one of the two members. Stats must report exactly one + // family of size 2 (`counts[1] == 1`, since index 1 = size 2), not + // two (which naively summing both slots would give), and both + // genomes represented at size 2, neither at any other size. + 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 stats = merged.sibling_annex_stats().expect("sibling_annex_stats"); + + assert_eq!(stats.counts, [0, 1, 0, 0], "one family of size 2, counted once"); + assert_eq!(stats.per_genome.len(), 2); + for g in 0..2 { + assert_eq!( + stats.per_genome[g], [0, 1, 0, 0], + "genome {g} should represent exactly one size-2 family" + ); + } } } diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 3707387..7a5ad4e 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -78,8 +78,7 @@ pub struct DistanceArgs { pub sibling_stats: bool, /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _siblings_per_genome.csv, - /// _nj.nwk, _upgma.nwk. + /// _siblings.csv, _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -114,6 +113,17 @@ pub fn run(args: DistanceArgs) { write_sibling_stats_csv(&stats, &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 + // matrix. + if args.sibling_annex || args.sibling_stats { + return; + } + info!( "computing {:?} distances for {} genome(s)", args.metric, n @@ -225,43 +235,38 @@ pub fn run(args: DistanceArgs) { } } -// ── Sibling-count distribution → CSV ──────────────────────────────────────── +// ── Family-size distribution → CSV ────────────────────────────────────────── +// +// Each row is a family (the up-to-4 k-mers sharing flanks, differing only at +// the centre), counted once — at its minorant — regardless of how many of +// its members are observed. Family size 1..4 (not "sibling count" 0..3): +// see `docmd/theory/evolutionary_distances.md`, "Definitions". 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() + // One row per genome (4 columns, family size 1-4: number of families of + // that size for which the genome carries at least one member), plus a + // `global` row — the actual deduplicated family-size histogram + // (`stats.counts`), NOT a sum of the per-genome columns (a family shared + // by several genomes would otherwise be counted once per genome it + // appears in, inflating the total beyond the real family count). + let 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}"); + 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, "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]; + writeln!(f, "genome,1,2,3,4").unwrap(); 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], + f, "global,{},{},{},{}", + stats.counts[0], stats.counts[1], stats.counts[2], stats.counts[3], ).unwrap(); - info!("per-genome sibling-count distribution → {per_genome_path}"); + let total: u64 = stats.counts.iter().sum(); + info!("family-size distribution → {path} (total {total} famil{})", + if total == 1 { "y" } else { "ies" }); } // ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────