diff --git a/src/obikindex/src/siblings/alignment.rs b/src/obikindex/src/siblings/alignment.rs index db514a92..f321d671 100644 --- a/src/obikindex/src/siblings/alignment.rs +++ b/src/obikindex/src/siblings/alignment.rs @@ -1,19 +1,13 @@ use rayon::prelude::*; -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; use obikpartitionner::KmerPartition; -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::{Mat, PartitionCache}; -use super::helpers::{central_base, is_minorant, partition_of}; -use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; +use super::cache::PartitionCache; +use super::family_scan::scan_layer_families; /// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set /// iff the genome carries the member whose own central base is `b`): @@ -77,26 +71,7 @@ impl KmerIndex { ) .map_err(OKIError::Partition)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?; - - let mut layer_dirs = Vec::new(); - for part in 0..n_parts { - let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); - if !index_dir.exists() { - continue; - } - let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; - for l in 0..meta.n_layers { - let layer_dir = index_dir.join(format!("layer_{l}")); - let annex_path = layer_dir.join(ANNEX_FILE_NAME); - if !annex_path.exists() { - return Err(OKIError::InvalidInput(format!( - "no sibling annex at {} — run build_sibling_annex first", - annex_path.display() - ))); - } - layer_dirs.push(layer_dir); - } - } + let layer_dirs = self.sibling_layer_dirs()?; let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); // `Vec>` per layer, one entry (column) per variable family; @@ -106,65 +81,12 @@ impl KmerIndex { let partials: Vec>> = layer_dirs .par_iter() .map(|layer_dir| -> 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 annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; - let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; - - let mut slot_kmer: Vec> = vec![None; annex.len()]; - let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) - .map_err(OKIError::Partition)?; - for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { - if let Some(slot) = mphf.find(kmer) { - slot_kmer[slot] = Some(kmer); - } - } - - let use_counts = with_counts && layer_dir.join("counts").exists(); - let mat = if use_counts { - Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) - } else { - Mat::Presence(PersistentBitMatrix::open(layer_dir)?) - }; - let n_cols = mat.n_cols().min(n_genomes); - - let mut columns: Vec> = Vec::new(); - let mut genome_mask: Vec = Vec::with_capacity(n_genomes); - - for slot in 0..annex.len() { - let Some(mask) = annex.get(slot) else { continue }; - let Some(kmer) = slot_kmer[slot] else { continue }; - if !is_minorant(kmer, mask, k) { - continue; // family tallied once, at its minorant - } - if mask.family_size() < 2 { - continue; // monomorphic family — no signal, skip - } - - genome_mask.clear(); - genome_mask.resize(n_genomes, 0); - - for other in kmer.central_canonical_neighbors() { - let base = central_base(other, k); - if !mask.has(base) { - continue; - } - let presence: Option> = if other == kmer { - Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) - } else { - let dest = partition_of(other, n_parts); - cache.find_presence(dest, other, n_genomes) - }; - let Some(presence) = presence else { continue }; - for (g, &present) in presence.iter().enumerate() { - if present { - genome_mask[g] |= 1 << base; - } - } - } - - columns.push(genome_mask.iter().map(|&m| iupac_code(m)).collect()); - } + let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; + let columns = families + .into_iter() + .filter(|f| f.mask.family_size() >= 2) // monomorphic family — no signal, skip + .map(|f| f.genome_mask.iter().map(|&m| iupac_code(m)).collect()) + .collect(); pb.inc(1); Ok(columns) diff --git a/src/obikindex/src/siblings/cardinality.rs b/src/obikindex/src/siblings/cardinality.rs index b9a023fa..84bd5b38 100644 --- a/src/obikindex/src/siblings/cardinality.rs +++ b/src/obikindex/src/siblings/cardinality.rs @@ -1,21 +1,15 @@ use ndarray::Array2; use rayon::prelude::*; -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; use obikpartitionner::KmerPartition; -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::{Mat, PartitionCache}; +use super::cache::PartitionCache; use super::distance::RawSnpDistanceOutput; -use super::helpers::{central_base, is_minorant, partition_of}; -use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; +use super::family_scan::scan_layer_families; /// See [`KmerIndex::cardinality_tally`]. pub struct CardinalityTally { @@ -80,63 +74,17 @@ impl KmerIndex { ) .map_err(OKIError::Partition)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?; - - let mut layer_dirs = Vec::new(); - for part in 0..n_parts { - let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); - if !index_dir.exists() { - continue; - } - let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; - for l in 0..meta.n_layers { - let layer_dir = index_dir.join(format!("layer_{l}")); - let annex_path = layer_dir.join(ANNEX_FILE_NAME); - if !annex_path.exists() { - return Err(OKIError::InvalidInput(format!( - "no sibling annex at {} — run build_sibling_annex first", - annex_path.display() - ))); - } - layer_dirs.push(layer_dir); - } - } + let layer_dirs = self.sibling_layer_dirs()?; let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers"); let partials: Vec<[[u64; 5]; 5]> = layer_dirs .par_iter() .map(|layer_dir| -> OKIResult<[[u64; 5]; 5]> { - let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); - let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; - let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; - let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; - - let mut slot_kmer: Vec> = vec![None; annex.len()]; - let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) - .map_err(OKIError::Partition)?; - for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { - if let Some(slot) = mphf.find(kmer) { - slot_kmer[slot] = Some(kmer); - } - } - - let use_counts = with_counts && layer_dir.join("counts").exists(); - let mat = if use_counts { - Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) - } else { - Mat::Presence(PersistentBitMatrix::open(layer_dir)?) - }; - let n_cols = mat.n_cols().min(n_genomes); - + let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; let mut counts = [[0u64; 5]; 5]; - let mut genome_mask: Vec = Vec::with_capacity(n_genomes); - for slot in 0..annex.len() { - let Some(mask) = annex.get(slot) else { continue }; - let Some(kmer) = slot_kmer[slot] else { continue }; - if !is_minorant(kmer, mask, k) { - continue; // family tallied once, at its minorant - } - if mask.family_size() < 2 { + for family in &families { + if family.mask.family_size() < 2 { // Fully invariant family (never varies anywhere in // the index) — genome-wide background, not // SNP-adjacent signal; would otherwise swamp the @@ -148,28 +96,7 @@ impl KmerIndex { continue; } - genome_mask.clear(); - genome_mask.resize(n_genomes, 0); - - for other in kmer.central_canonical_neighbors() { - let base = central_base(other, k); - if !mask.has(base) { - continue; - } - let presence: Option> = if other == kmer { - Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) - } else { - let dest = partition_of(other, n_parts); - cache.find_presence(dest, other, n_genomes) - }; - let Some(presence) = presence else { continue }; - for (g, &present) in presence.iter().enumerate() { - if present { - genome_mask[g] |= 1 << base; - } - } - } - + let genome_mask = &family.genome_mask; for i in 0..n_genomes { let card_i = genome_mask[i].count_ones() as usize; for j in (i + 1)..n_genomes { diff --git a/src/obikindex/src/siblings/distance.rs b/src/obikindex/src/siblings/distance.rs index 8dcf40fc..6b5d37e0 100644 --- a/src/obikindex/src/siblings/distance.rs +++ b/src/obikindex/src/siblings/distance.rs @@ -1,20 +1,14 @@ use ndarray::Array2; use rayon::prelude::*; -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; use obikpartitionner::KmerPartition; -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::{Mat, PartitionCache}; -use super::helpers::{central_base, is_minorant, partition_of}; -use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; +use super::cache::PartitionCache; +use super::family_scan::scan_layer_families; /// Raw p-distance restricted to loci that are single-copy in **both** /// genomes of a pair — the "stringent / paralogy-aware" locus eligibility @@ -84,108 +78,32 @@ impl KmerIndex { ) .map_err(OKIError::Partition)?; let cache = PartitionCache::build(&partition, n_parts, with_counts)?; - - let mut layer_dirs = Vec::new(); - for part in 0..n_parts { - let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); - if !index_dir.exists() { - continue; - } - let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; - for l in 0..meta.n_layers { - let layer_dir = index_dir.join(format!("layer_{l}")); - let annex_path = layer_dir.join(ANNEX_FILE_NAME); - if !annex_path.exists() { - return Err(OKIError::InvalidInput(format!( - "no sibling annex at {} — run build_sibling_annex first", - annex_path.display() - ))); - } - layer_dirs.push(layer_dir); - } - } + let layer_dirs = self.sibling_layer_dirs()?; let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); let partials: Vec = layer_dirs .par_iter() .map(|layer_dir| -> OKIResult { let mut acc = zero(); + let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; - let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); - let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; - let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; - let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + for family in &families { + let variable = family.mask.family_size() >= 2; + let genome_mask = &family.genome_mask; - 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); - } - } + // Per genome: which single form (if exactly one) it + // carries — `None` (a `popcount != 1` mask) once a + // second form is seen, ambiguous/not single-copy, + // ineligible for either side of a pair. + let single_form = |g: usize| -> Option { + let m = genome_mask[g]; + (m.count_ones() == 1).then(|| m.trailing_zeros() as u8) + }; - let use_counts = with_counts && layer_dir.join("counts").exists(); - let mat = if use_counts { - Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) - } else { - Mat::Presence(PersistentBitMatrix::open(layer_dir)?) - }; - let n_cols = mat.n_cols().min(n_genomes); - - // Per family: which single form (if exactly one) each genome - // carries — `None` once a second form is seen (ambiguous, - // not single-copy, ineligible for either side of a pair). - let mut single_form: Vec> = Vec::with_capacity(n_cols); - let mut ambiguous: Vec = Vec::with_capacity(n_cols); - - for slot in 0..annex.len() { - let Some(mask) = annex.get(slot) else { continue }; - let Some(kmer) = slot_kmer[slot] else { continue }; - if !is_minorant(kmer, mask, k) { - continue; // family tallied once, at its minorant - } - let variable = mask.family_size() >= 2; - - single_form.clear(); - single_form.resize(n_cols, None); - ambiguous.clear(); - ambiguous.resize(n_cols, false); - - for other in kmer.central_canonical_neighbors() { - let base = central_base(other, k); - if !mask.has(base) { - continue; - } - let presence: Option> = if other == kmer { - Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) - } else { - let dest = partition_of(other, n_parts); - cache.find_presence(dest, other, n_genomes) - }; - let Some(presence) = presence else { continue }; - for (g, &present) in presence.iter().enumerate() { - if !present { - continue; - } - if single_form[g].is_some() { - ambiguous[g] = true; - } else { - single_form[g] = Some(base); - } - } - } - - for i in 0..n_cols { - if ambiguous[i] { - continue; - } - let Some(bi) = single_form[i] else { continue }; - for j in (i + 1)..n_cols { - if ambiguous[j] { - continue; - } - let Some(bj) = single_form[j] else { continue }; + for i in 0..n_genomes { + let Some(bi) = single_form(i) else { continue }; + for j in (i + 1)..n_genomes { + let Some(bj) = single_form(j) else { continue }; on_pair(&mut acc, i, j, bi, bj, variable); } } diff --git a/src/obikindex/src/siblings/family_scan.rs b/src/obikindex/src/siblings/family_scan.rs new file mode 100644 index 00000000..4ad7da1f --- /dev/null +++ b/src/obikindex/src/siblings/family_scan.rs @@ -0,0 +1,178 @@ +//! Shared per-layer family traversal, used by every sibling-annex consumer +//! (`snp_pseudo_alignment`, `cardinality_tally`, `scan_family_pairs`, +//! `sibling_annex_stats`) — resolves each minorant family's per-genome +//! 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. The naive, per-family inline lookup +//! was the shape all four consumers used to have independently; sampling a +//! real ~170 GB index against 137 GB of RAM showed it thrashing the page +//! cache (near-continuous ~2 GB/s pagein), the same failure mode +//! `build_sibling_annex` had already been fixed for. + +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU8, Ordering}; + +use rayon::prelude::*; + +use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; +use obikseq::CanonicalKmer; +use obilayeredmap::MphfLayer; +use obilayeredmap::meta::PartitionMeta; +use obiskio::UnitigFileReader; + +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::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; + +/// One layer's minorant family, with its own annex mask and, per genome, +/// the 4-bit base-presence mask resolved across the whole index (bit `b` +/// set iff that genome carries the family member whose own canonical +/// central base is `b`). +pub(super) struct FamilyRow { + pub mask: FamilyMask, + pub genome_mask: Vec, +} + +impl KmerIndex { + /// Every (partition, layer) directory carrying a sibling annex, checked + /// up front so a missing one is reported before any real work starts. + pub(super) fn sibling_layer_dirs(&self) -> OKIResult> { + let n_parts = self.n_partitions(); + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + layer_dirs.push(layer_dir); + } + } + Ok(layer_dirs) + } +} + +/// Every minorant family of one layer, each with its per-genome +/// base-presence mask resolved against the whole (already-open) partition +/// cache. +/// +/// Two passes, not one inline pass: the first is purely local (already-open +/// annex/mphf/matrix, no lookups) and both enumerates the layer's families +/// and collects every cross-partition query they need, grouped by +/// destination partition; the second resolves each partition's whole batch +/// in one contiguous sweep, keeping that partition's mmap'd pages hot for +/// its entire batch instead of faulting them in and out as lookups jump +/// between partitions in family order. +pub(super) fn scan_layer_families( + layer_dir: &Path, + n_parts: usize, + n_genomes: usize, + with_counts: bool, + k: usize, + cache: &PartitionCache, +) -> 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 annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")).map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); + + // ── Pass 1: local only — enumerate minorant families and collect every + // cross-partition query they need, grouped by destination partition. + struct Family { + slot: usize, + mask: FamilyMask, + } + let mut families: Vec = Vec::new(); + let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); + + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant + } + let family_idx = families.len(); + families.push(Family { slot, mask }); + + for other in kmer.central_canonical_neighbors() { + if other == kmer { + continue; // local — resolved below straight from `mat`, no lookup + } + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let dest = partition_of(other, n_parts); + outgoing[dest].push((other, family_idx, base)); + } + } + + let genome_mask: Vec = (0..families.len() * n_genomes).map(|_| AtomicU8::new(0)).collect(); + + // Local presence (the family's own minorant slot) — already open, no + // lookup, so no locality concern either way. + for (family_idx, family) in families.iter().enumerate() { + let kmer = slot_kmer[family.slot].expect("family slot has a kmer"); + let base = central_base(kmer, k); + for g in 0..n_cols { + if mat.carries(g, family.slot) { + genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); + } + } + } + + // ── Pass 2: resolve each destination partition's whole batch in one + // contiguous sweep — 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); + } + } + } + }); + + Ok(families + .into_iter() + .enumerate() + .map(|(family_idx, family)| FamilyRow { + mask: family.mask, + genome_mask: (0..n_genomes) + .map(|g| genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed)) + .collect(), + }) + .collect()) +} diff --git a/src/obikindex/src/siblings/mod.rs b/src/obikindex/src/siblings/mod.rs index edc344fb..2ff427e9 100644 --- a/src/obikindex/src/siblings/mod.rs +++ b/src/obikindex/src/siblings/mod.rs @@ -47,6 +47,7 @@ mod build; mod cache; mod cardinality; mod distance; +mod family_scan; mod helpers; mod stats; diff --git a/src/obikindex/src/siblings/stats.rs b/src/obikindex/src/siblings/stats.rs index 66db7958..ac659773 100644 --- a/src/obikindex/src/siblings/stats.rs +++ b/src/obikindex/src/siblings/stats.rs @@ -1,19 +1,13 @@ use rayon::prelude::*; -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; use obikpartitionner::KmerPartition; -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::{Mat, PartitionCache}; -use super::helpers::{central_base, is_minorant, partition_of}; -use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; +use super::cache::PartitionCache; +use super::family_scan::scan_layer_families; /// Distribution of family sizes (1-4), read back from an already-built /// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's @@ -62,29 +56,7 @@ impl KmerIndex { ) .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); - if !annex_path.exists() { - return Err(OKIError::InvalidInput(format!( - "no sibling annex at {} — run build_sibling_annex first", - annex_path.display() - ))); - } - layer_dirs.push(layer_dir); - } - } + let layer_dirs = self.sibling_layer_dirs()?; // One layer's worth of work, parallelised across layers with Rayon // — independent, read-only, each producing its own partial tally @@ -98,70 +70,15 @@ impl KmerIndex { ..Default::default() }; - 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); - } - } - - 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(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; - + let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; + for family in &families { // "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 { - 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 { + // `genome_mask[g] != 0` is exactly that. + let s = family.mask.siblings() as usize; + stats.counts[s] += 1; + for (g, &m) in family.genome_mask.iter().enumerate() { + if m != 0 { stats.per_genome[g][s] += 1; } } diff --git a/src/obikindex/src/siblings/tests.rs b/src/obikindex/src/siblings/tests.rs index 0aa833b6..09a1c597 100644 --- a/src/obikindex/src/siblings/tests.rs +++ b/src/obikindex/src/siblings/tests.rs @@ -184,3 +184,52 @@ fn sibling_annex_stats_counts_each_family_once_and_per_genome() { ); } } + +/// 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 +/// above (g1 carries the family's centre-C member, g2 the centre-G one), +/// hand-verified expected output for each. +#[test] +fn family_scan_consumers_agree_on_one_sibling_each() { + 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"); + + // Merge doesn't promise to preserve source order, so resolve each + // genome's index by label rather than assuming g1 -> 0, g2 -> 1. + let idx_of = |label: &str| merged.meta.genomes.iter().position(|g| g.label == label).unwrap(); + let (i1, i2) = (idx_of("g1"), idx_of("g2")); + + // snp_pseudo_alignment: one variable family, one column — g1's row + // reads 'C' (its own member), g2's reads 'G'. + let alignment = merged.snp_pseudo_alignment().expect("snp_pseudo_alignment"); + assert_eq!(alignment.sequences[i1], vec![b'C']); + assert_eq!(alignment.sequences[i2], vec![b'G']); + + // raw_snp_distance: g1's single form (C) != g2's (G) at the family's + // one eligible locus -> a SNP, not a shared site. + let raw = merged.raw_snp_distance().expect("raw_snp_distance"); + assert_eq!(raw.snp[[i1, i2]], 1); + assert_eq!(raw.snp[[i2, i1]], 1); + assert_eq!(raw.shared[[i1, i2]], 0); + assert_eq!(raw.shared[[i2, i1]], 0); + + // cardinality_tally: both genomes carry exactly one member of the + // family (cardinality 1 each) -> one co-occurrence at [1][1]. + let cardinality = merged.cardinality_tally(&raw, 1.0).expect("cardinality_tally"); + assert_eq!(cardinality.counts[1][1], 1); + let total: u64 = cardinality.counts.iter().flatten().sum(); + assert_eq!(total, 1, "no other cardinality pair should be tallied"); + + // base_pair_tally: the one eligible, differing locus is C (base 1) vs + // G (base 2). + let base_pairs = merged.base_pair_tally(&raw, 1.0).expect("base_pair_tally"); + assert_eq!(base_pairs.counts[1][2], 1); + assert_eq!(base_pairs.counts[2][1], 1); + let total: u64 = base_pairs.counts.iter().flatten().sum(); + assert_eq!(total, 2, "no other base pair should be tallied"); + assert_eq!(base_pairs.same, [0, 0, 0, 0], "the two genomes never agree at this locus"); +}