diff --git a/src/Cargo.lock b/src/Cargo.lock index 74a8dabe..3efc8866 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1768,9 +1768,7 @@ dependencies = [ "obikpartitionner", "obikseq", "obilayeredmap", - "obipipeline", "obiread", - "obiskbuilder", "obiskio", "obisys", "obitaxonomy", @@ -1794,6 +1792,7 @@ dependencies = [ "obifastwrite", "obikindex", "obikpartitionner", + "obikphylo", "obikrope", "obikseq", "obilayeredmap", @@ -1843,6 +1842,27 @@ dependencies = [ "tracing", ] +[[package]] +name = "obikphylo" +version = "0.1.0" +dependencies = [ + "ndarray", + "obicompactvec", + "obikindex", + "obikpartitionner", + "obikseq", + "obilayeredmap", + "obipipeline", + "obiread", + "obiskbuilder", + "obiskio", + "obisys", + "rayon", + "tempfile", + "tracing", + "tracing-subscriber", +] + [[package]] name = "obikrope" version = "0.1.0" diff --git a/src/Cargo.toml b/src/Cargo.toml index 5684237b..f34f6918 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -1,5 +1,5 @@ [workspace] resolver = "3" -members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy"] +members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy", "obikphylo"] [profile.release] debug = 1 diff --git a/src/obikindex/Cargo.toml b/src/obikindex/Cargo.toml index ba86dd64..b92aa075 100644 --- a/src/obikindex/Cargo.toml +++ b/src/obikindex/Cargo.toml @@ -11,8 +11,6 @@ obiskio = { path = "../obiskio" } obisys = { path = "../obisys" } obicompactvec = { path = "../obicompactvec" } obilayeredmap = { path = "../obilayeredmap" } -obiskbuilder = { path = "../obiskbuilder" } -obipipeline = { path = "../obipipeline" } ndarray = "0.16" rayon = "1" crossbeam-channel = "0.5" diff --git a/src/obikindex/src/index.rs b/src/obikindex/src/index.rs index 6d2163ed..3aedca98 100644 --- a/src/obikindex/src/index.rs +++ b/src/obikindex/src/index.rs @@ -132,6 +132,10 @@ impl KmerIndex { IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty) } + /// The index's root directory — needed by out-of-crate extension code + /// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto + /// the same on-disk index. + pub fn root_path(&self) -> &Path { &self.root_path } pub fn meta(&self) -> &IndexMeta { &self.meta } pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta } pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size } diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 665e2d86..6ab2e070 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -2,7 +2,6 @@ pub mod error; pub mod meta; pub mod predicate; pub mod state; -mod cardcomp; mod distance; mod dump; mod index; @@ -11,7 +10,6 @@ mod numa; mod rebuild; mod reindex; mod select; -mod siblings; mod stats; pub use error::{OKIError, OKIResult}; @@ -22,5 +20,3 @@ pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME pub use predicate::{GroupFilterParams, MetaPred}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use stats::IndexBitsPerKmer; -pub use siblings::{BasePairTally, CardinalityTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; -pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix}; diff --git a/src/obikindex/src/siblings/cache.rs b/src/obikindex/src/siblings/cache.rs deleted file mode 100644 index 32d49f5d..00000000 --- a/src/obikindex/src/siblings/cache.rs +++ /dev/null @@ -1,166 +0,0 @@ -use rayon::prelude::*; - -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; -use obikpartitionner::KmerPartition; -use obikseq::CanonicalKmer; -use obilayeredmap::MphfLayer; -use obilayeredmap::meta::PartitionMeta; -use obisys::progress_bar; - -use crate::error::OKIResult; - -use super::{olm_to_ok, INDEX_SUBDIR}; - -/// 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`. -pub(super) enum Mat { - Count(PersistentCompactIntMatrix), - Presence(PersistentBitMatrix), -} - -impl Mat { - pub(super) fn n_cols(&self) -> usize { - match self { - Mat::Count(m) => m.n_cols(), - Mat::Presence(m) => m.n_cols(), - } - } - pub(super) 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, - } - } -} - -pub(super) 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 - /// [`crate::index::KmerIndex::build_sibling_annex`] (`layers` only) and - /// [`crate::index::KmerIndex::sibling_annex_stats`] (both). - layers: Vec>, - mats: Vec>, -} - -impl PartitionCache { - pub(super) 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. - pub(super) 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())) - } - - /// 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); - } - } - } - } - } -} diff --git a/src/obikindex/src/siblings/helpers.rs b/src/obikindex/src/siblings/helpers.rs deleted file mode 100644 index f50ee2ce..00000000 --- a/src/obikindex/src/siblings/helpers.rs +++ /dev/null @@ -1,49 +0,0 @@ -use obikseq::CanonicalKmer; -use obiskbuilder::rolling_stat::RollingStat; - -use obicompactvec::FamilyMask; - -/// 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] -pub(super) 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`. -pub(super) 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 -/// the same selection for a single, disconnected k-mer. Not the leanest -/// possible primitive (an O(1)-amortised dedicated scan, as originally -/// sketched in the design doc's Step 0, would avoid the ASCII round-trip and -/// `RollingStat` allocation) but correct and reuses already-tested logic; -/// left as a follow-up optimisation. -fn lone_kmer_minimizer(kmer: CanonicalKmer) -> obikseq::Minimizer { - let ascii = kmer.to_ascii(); - let mut rs = RollingStat::new(0); - for b in ascii { - rs.push(b); - } - rs.canonical_minimizer() - .expect("RollingStat must be ready after k bases of a valid k-mer") -} - -/// Destination partition for a (possibly synthetic) canonical k-mer, using -/// the same routing rule as the rest of the index (`minimiser.seq_hash() & -/// mask`, `n_partitions` is a power of two). -pub(super) fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize { - let mask = (n_partitions as u64) - 1; - (lone_kmer_minimizer(kmer).seq_hash() & mask) as usize -} diff --git a/src/obikindex/src/siblings/iter.rs b/src/obikindex/src/siblings/iter.rs deleted file mode 100644 index 70b26c1c..00000000 --- a/src/obikindex/src/siblings/iter.rs +++ /dev/null @@ -1,78 +0,0 @@ -//! Phylo/sibling-domain iteration over a layer — an extension trait, not a -//! new field on `MphfLayer`/`Layer`: "family"/"minorant" are phylo -//! concepts, `obilayeredmap` stays kmer/slot-mapping only (see -//! `docmd/architecture/siblings.md`). -//! -//! The sibling annex is persisted in the same order as `iter_kmers()` -//! (`build_sibling_annex`, see `build.rs`), so pairing them is a plain zip — -//! no MPHF, no slot, no `kmer_at`. Both sides are already `Send + 'static` -//! (`KmerIter` owns an `Arc` clone; `SiblingAnnex` is -//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter` -//! streams straight from disk and can be fed to `obipipeline` batch by -//! batch — never collected whole into memory (see the project's "no full -//! collect" rule). - -use std::sync::Arc; - -use obicompactvec::{FamilyMask, SiblingAnnex}; -use obikseq::CanonicalKmer; -use obilayeredmap::{KmerIter, MphfLayer}; - -/// One layer entry: a k-mer's position in the layer's iteration order (the -/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer -/// itself, and its family mask. -#[derive(Debug, Clone, Copy)] -pub struct SiblingEntry { - pub order: usize, - pub kmer: CanonicalKmer, - pub mask: FamilyMask, -} - -/// Streams `(order, kmer, mask)` triples for one layer, in iteration order. -/// Produced by [`SiblingLayerExt::iter_siblings`]. -pub struct SiblingIter { - kmers: KmerIter, - annex: Arc, - order: usize, -} - -impl Iterator for SiblingIter { - type Item = SiblingEntry; - - fn next(&mut self) -> Option { - loop { - let kmer = self.kmers.next()?; - let order = self.order; - self.order += 1; - // `None` means "not yet computed" (see `SiblingAnnex` module - // docs) — shouldn't happen against a fully-built annex, but - // skip rather than misalign the two streams if it does. - if let Some(mask) = self.annex.get(order) { - return Some(SiblingEntry { order, kmer, mask }); - } - } - } -} - -/// Adds phylo/sibling iteration to `MphfLayer`. -pub trait SiblingLayerExt { - /// Zip this layer's k-mers with their sibling-annex entry, in iteration - /// order. `annex` must have been built from this same layer (its length - /// must match the layer's k-mer count). - fn iter_siblings(&self, annex: Arc) -> SiblingIter; - - /// Like [`iter_siblings`](Self::iter_siblings), filtered to the - /// minorant of each family — the common case, since a family is - /// tallied once, at its minorant. - fn iter_minorants(&self, annex: Arc) -> impl Iterator; -} - -impl SiblingLayerExt for MphfLayer { - fn iter_siblings(&self, annex: Arc) -> SiblingIter { - SiblingIter { kmers: self.iter_kmers(), annex, order: 0 } - } - - fn iter_minorants(&self, annex: Arc) -> impl Iterator { - self.iter_siblings(annex).filter(|e| e.mask.is_minorant()) - } -} diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 7a16f63c..903c69db 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -19,6 +19,7 @@ obikpartitionner = { path = "../obikpartitionner" } obisys = { path = "../obisys" } obiskio = { path = "../obiskio" } obikindex = { path = "../obikindex", default-features = false } +obikphylo = { path = "../obikphylo" } obitaxonomy = { path = "../obitaxonomy" } obilayeredmap = { path = "../obilayeredmap" } clap = { version = "4", features = ["derive"] } diff --git a/src/obikmer/src/cmd/phylo/family_overlap.rs b/src/obikmer/src/cmd/phylo/family_overlap.rs index fa01c8e2..7a8f9d7b 100644 --- a/src/obikmer/src/cmd/phylo/family_overlap.rs +++ b/src/obikmer/src/cmd/phylo/family_overlap.rs @@ -1,7 +1,8 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; -use obikindex::{KmerIndex, SnpAlignment}; +use obikindex::KmerIndex; +use obikphylo::siblings::{SnpAlignment, SnpAlignmentExt}; use tracing::info; // ── Family overlap: shared-family counts and the `--min-shared-family` / diff --git a/src/obikmer/src/cmd/phylo/iqtree.rs b/src/obikmer/src/cmd/phylo/iqtree.rs index cd906841..d5b17af0 100644 --- a/src/obikmer/src/cmd/phylo/iqtree.rs +++ b/src/obikmer/src/cmd/phylo/iqtree.rs @@ -2,7 +2,7 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; use obifastwrite::{JsonVal, write_record}; -use obikindex::SnpAlignment; +use obikphylo::siblings::SnpAlignment; use tracing::info; use super::sankoff::state_index_table; diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index 9b09a071..b1cb86d3 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -9,9 +9,13 @@ mod tnt; use std::io::{self, BufWriter, Write}; use kodama::{Method, linkage}; -use obikindex::{ - KmerIndex, RawSnpDistanceOutput, SnpAlignment, +use obikindex::KmerIndex; +use obikphylo::{ cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix, + siblings::{ + CardinalityExt, DistanceExt, RawSnpDistanceOutput, SiblingAnnexBuildExt, SiblingStatsExt, + SnpAlignment, SnpAlignmentExt, + }, }; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; diff --git a/src/obikmer/src/cmd/phylo/outputs.rs b/src/obikmer/src/cmd/phylo/outputs.rs index 28b67a66..eef77368 100644 --- a/src/obikmer/src/cmd/phylo/outputs.rs +++ b/src/obikmer/src/cmd/phylo/outputs.rs @@ -2,7 +2,7 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; use obifastwrite::{JsonVal, write_record}; -use obikindex::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; +use obikphylo::siblings::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; use tracing::info; // ── Family-size distribution → CSV ────────────────────────────────────────── diff --git a/src/obikmer/src/cmd/phylo/sankoff.rs b/src/obikmer/src/cmd/phylo/sankoff.rs index ab96e7d3..7c91b0ca 100644 --- a/src/obikmer/src/cmd/phylo/sankoff.rs +++ b/src/obikmer/src/cmd/phylo/sankoff.rs @@ -2,7 +2,7 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; use obifastwrite::{JsonVal, write_record}; -use obikindex::{BasePairTally, CardinalityTally, SnpAlignment}; +use obikphylo::siblings::{BasePairTally, CardinalityTally, SnpAlignment}; use tracing::info; // ── Sankoff pseudo-alignment → FASTA ──────────────────────────────────────── diff --git a/src/obikmer/src/cmd/phylo/tnt.rs b/src/obikmer/src/cmd/phylo/tnt.rs index aa16b419..a4be379f 100644 --- a/src/obikmer/src/cmd/phylo/tnt.rs +++ b/src/obikmer/src/cmd/phylo/tnt.rs @@ -1,7 +1,7 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; -use obikindex::SnpAlignment; +use obikphylo::siblings::SnpAlignment; use tracing::info; use super::sankoff::{scaled_metric_matrix, state_index_table}; diff --git a/src/obikphylo/Cargo.toml b/src/obikphylo/Cargo.toml new file mode 100644 index 00000000..be66a13c --- /dev/null +++ b/src/obikphylo/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "obikphylo" +version = "0.1.0" +edition = "2024" + +[dependencies] +obikindex = { path = "../obikindex" } +obikseq = { path = "../obikseq" } +obikpartitionner = { path = "../obikpartitionner" } +obiskio = { path = "../obiskio" } +obisys = { path = "../obisys" } +obicompactvec = { path = "../obicompactvec" } +obilayeredmap = { path = "../obilayeredmap" } +obiskbuilder = { path = "../obiskbuilder" } +obipipeline = { path = "../obipipeline" } +ndarray = "0.16" +rayon = "1" +tracing = "0.1.44" + +[dev-dependencies] +obiread = { path = "../obiread" } +tempfile = "3" +tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/src/obikindex/src/cardcomp.rs b/src/obikphylo/src/cardcomp.rs similarity index 99% rename from src/obikindex/src/cardcomp.rs rename to src/obikphylo/src/cardcomp.rs index eab104a3..8fdd2fa8 100644 --- a/src/obikindex/src/cardcomp.rs +++ b/src/obikphylo/src/cardcomp.rs @@ -17,7 +17,7 @@ //! since TNT/PhyG search *unrooted* trees) if the cost matrix is //! symmetric, the discrete-parsimony analogue of CTMC reversibility. -use crate::{BasePairTally, CardinalityTally}; +use crate::siblings::{BasePairTally, CardinalityTally}; /// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`), /// diagonal included ("stay at the same cardinality"), from diff --git a/src/obikphylo/src/lib.rs b/src/obikphylo/src/lib.rs new file mode 100644 index 00000000..d745f063 --- /dev/null +++ b/src/obikphylo/src/lib.rs @@ -0,0 +1,14 @@ +//! Library-level phylogenetic functionality for `obikmer`, built as +//! extension traits over `obikindex::KmerIndex` and `obilayeredmap`'s +//! generic layer types — the `phylo` CLI command is a consumer of this +//! crate, not the owner of this logic (see `docmd/architecture/siblings.md`). +//! +//! Starts with [`siblings`] (family presence-mask annex, SNP distance, +//! cardinality, pseudo-alignment); further phylo-domain functionality +//! (currently `obikindex::distance`/`obikindex::cardcomp`) moves here +//! incrementally. + +mod cardcomp; +pub mod siblings; + +pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix}; diff --git a/src/obikindex/src/siblings/alignment.rs b/src/obikphylo/src/siblings/alignment.rs similarity index 84% rename from src/obikindex/src/siblings/alignment.rs rename to src/obikphylo/src/siblings/alignment.rs index e87e1017..123b14d3 100644 --- a/src/obikindex/src/siblings/alignment.rs +++ b/src/obikphylo/src/siblings/alignment.rs @@ -3,8 +3,8 @@ use std::sync::Arc; use obikpartitionner::KmerPartition; use obisys::progress_bar; -use crate::error::{OKIError, OKIResult}; -use crate::index::KmerIndex; +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; use super::cache::PartitionCache; use super::family_scan::scan_layer_families; @@ -53,25 +53,32 @@ pub struct SnpAlignment { pub sequences: Vec>, } -impl KmerIndex { +/// Adds [`snp_pseudo_alignment`](Self::snp_pseudo_alignment) to `KmerIndex` — +/// phylo-domain functionality, kept out of `obikindex` itself (see +/// `docmd/architecture/siblings.md`). +pub trait SnpAlignmentExt { /// Build the SNP-only pseudo-alignment from an already-built sibling - /// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first). - pub fn snp_pseudo_alignment(&self) -> OKIResult { + /// annex (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first). + fn snp_pseudo_alignment(&self) -> OKIResult; +} + +impl SnpAlignmentExt for KmerIndex { + fn snp_pseudo_alignment(&self) -> OKIResult { let n_parts = self.n_partitions(); - let n_genomes = self.meta.genomes.len(); - let with_counts = self.meta.config.with_counts; + let n_genomes = self.meta().genomes.len(); + let with_counts = self.meta().config.with_counts; let k = self.kmer_size(); let n_bits = n_parts.trailing_zeros() as usize; let partition = KmerPartition::open_with_config( - &self.root_path, + self.root_path(), self.kmer_size(), self.minimizer_size(), n_bits, ) .map_err(OKIError::Partition)?; let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); - let layer_dirs = self.sibling_layer_dirs()?; + let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); // One layer at a time, not `par_iter()` over layers — same diff --git a/src/obikindex/src/siblings/build.rs b/src/obikphylo/src/siblings/build.rs similarity index 89% rename from src/obikindex/src/siblings/build.rs rename to src/obikphylo/src/siblings/build.rs index b51a7f32..70955d49 100644 --- a/src/obikindex/src/siblings/build.rs +++ b/src/obikphylo/src/siblings/build.rs @@ -12,11 +12,11 @@ use obilayeredmap::MphfLayer; use obilayeredmap::meta::PartitionMeta; use obisys::progress_bar; -use crate::error::{OKIError, OKIResult}; -use crate::index::KmerIndex; +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; use super::cache::PartitionCache; -use super::helpers::{central_base, is_minorant, partition_of}; +use super::helpers::{central_base, is_minorant}; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; // ── obipipeline data types ───────────────────────────────────────────────── @@ -53,7 +53,8 @@ enum SibData { Variants(VariantBatch), } -impl KmerIndex { +/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `KmerIndex`. +pub trait SiblingAnnexBuildExt { /// Build the sibling-count/minorant annex for every layer of every /// partition of this (already built) index, writing one annex file per /// layer alongside its existing index files. Safe to call again later @@ -63,18 +64,23 @@ impl KmerIndex { /// Construction only — no statistics gathered here on purpose: this is /// meant to run routinely (it is the artefact the SNP-family distances /// will consume), while the sibling-count distribution - /// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate, - /// occasional diagnostic pass over the result, not run every time. + /// ([`sibling_annex_stats`](super::stats::SiblingStatsExt::sibling_annex_stats)) + /// is a separate, occasional diagnostic pass over the result, not run + /// every time. /// /// Cross-partition/cross-layer lookups are required (a k-mer's siblings /// can live in any partition), but the layer loop itself — and thus the /// annex file this produces — stays local to one layer at a time. - pub fn build_sibling_annex(&self) -> OKIResult<()> { + fn build_sibling_annex(&self) -> OKIResult<()>; +} + +impl SiblingAnnexBuildExt for KmerIndex { + fn build_sibling_annex(&self) -> OKIResult<()> { let n_parts = self.n_partitions(); let n_bits = n_parts.trailing_zeros() as usize; let partition = KmerPartition::open_with_config( - &self.root_path, + self.root_path(), self.kmer_size(), self.minimizer_size(), n_bits, @@ -82,7 +88,7 @@ impl KmerIndex { .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 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; @@ -97,7 +103,7 @@ impl KmerIndex { let mut part_slots: u64 = 0; for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); - part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?; + part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, &cache)?; } total_slots += part_slots; pb.inc(1); @@ -108,20 +114,22 @@ impl KmerIndex { Ok(()) } +} - /// Returns the number of distinct k-mers (annex entries) processed, for - /// progress reporting. - fn build_layer_sibling_annex( - &self, - layer_dir: &Path, - n_parts: usize, - 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)?; - let k = self.kmer_size(); - let n = mphf.n(); +/// Returns the number of distinct k-mers (annex entries) processed, for +/// progress reporting. A free function, not a `KmerIndex` method — called +/// only from `build_sibling_annex` above, in the same file. +fn build_layer_sibling_annex( + index: &KmerIndex, + layer_dir: &Path, + n_parts: usize, + 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)?; + let k = index.kmer_size(); + let n = mphf.n(); // ── Reconciliation state, indexed by this layer's k-mer iteration // order (the physical layout of `unitigs.bin`), never by MPHF slot @@ -211,7 +219,7 @@ impl KmerIndex { continue; } items.push(( - partition_of(variant, n_parts), + variant.partition(n_parts), variant, order, central_base(variant, k), @@ -272,6 +280,5 @@ impl KmerIndex { } builder.close()?; - Ok(n as u64) - } + Ok(n as u64) } diff --git a/src/obikphylo/src/siblings/cache.rs b/src/obikphylo/src/siblings/cache.rs new file mode 100644 index 00000000..8d706d27 --- /dev/null +++ b/src/obikphylo/src/siblings/cache.rs @@ -0,0 +1,225 @@ +use rayon::prelude::*; + +use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; +use obikpartitionner::KmerPartition; +use obikseq::CanonicalKmer; +use obilayeredmap::Layer; +use obilayeredmap::meta::PartitionMeta; +use obisys::progress_bar; + +use obikindex::OKIResult; + +use super::iter::SiblingLayerExt; +use super::{olm_to_ok, INDEX_SUBDIR}; + +/// Every partition's already-open 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. +/// +/// One `Layer` per layer (MPHF + matrix bundled), not a separate +/// `MphfLayer` and a separate `PersistentCompactIntMatrix`/ +/// `PersistentBitMatrix` in parallel arrays — `obilayeredmap::Layer` already +/// *is* that bundle, with `find_slot` (MPHF-only, no data read), +/// `n_cols`/`sub_matrix`/`fill_sub_matrix` (batched, sorted-internally +/// column access) on top of it. Reinventing that pairing here would just be +/// going back through the low-level pieces `Layer` already assembles. +pub(super) enum Mat { + Count(Layer), + Presence(Layer), +} + +impl Mat { + fn find_slot(&self, kmer: CanonicalKmer) -> Option { + match self { + Mat::Count(l) => l.find_slot(kmer), + Mat::Presence(l) => l.find_slot(kmer), + } + } + + /// Raw MPHF batch lookup: kmer → slot, no membership check — for + /// callers that already know every kmer is a member of *this* layer + /// (e.g. it came from this layer's own `iter_minorants_batch`), so the + /// evidence check `find_slot`/`find` would perform is redundant work. + /// See `docmd/architecture/siblings.md`: iteration-pipeline kmers use + /// `index`, never `find`. + pub(super) fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec { + match self { + Mat::Count(l) => l.index_batch(kmers), + Mat::Presence(l) => l.index_batch(kmers), + } + } + + /// This layer's own `SiblingLayerExt::iter_minorants_batch` — dispatch + /// only, both arms return the same concrete `MinorantBatchIter` (it + /// doesn't depend on `D`), so no boxing is needed. + pub(super) fn iter_minorants_batch( + &self, + annex: std::sync::Arc, + batch_size: usize, + ) -> super::iter::MinorantBatchIter { + match self { + Mat::Count(l) => l.iter_minorants_batch(annex, batch_size), + Mat::Presence(l) => l.iter_minorants_batch(annex, batch_size), + } + } + + pub(super) fn n_cols(&self) -> usize { + match self { + Mat::Count(l) => l.n_cols(), + Mat::Presence(l) => l.n_cols(), + } + } + + /// Batch, genome-major "carries" for a set of `slots` — `out[g][i]` = + /// whether genome `g` (0..`out.len()`) carries `slots[i]`. `out` must + /// have one entry per genome column, each resized to `slots.len()`. + /// + /// Delegates entirely to `Layer::fill_sub_matrix`, which sorts + /// `slots` once internally for a sequential mmap sweep per column, then + /// restores the original order — the same discipline this call site + /// (and `find_presence_batch`) used to hand-roll with its own sort + + /// genome-major loop. `Count` still needs one intermediate + /// `Vec>` fetch (the underlying store only has an int + /// sub-matrix, not a bool one), converted to presence (`!= 0`) in place + /// — the sort/sequential-access win is unaffected, just one extra + /// allocation pass over already-in-hand data. + pub(super) fn fill_sub_matrix_carries(&self, slots: &[usize], out: &mut [Vec]) { + match self { + Mat::Presence(l) => l.fill_sub_matrix(slots, out), + Mat::Count(l) => { + let mut counts: Vec> = out.iter().map(|_| Vec::new()).collect(); + l.fill_sub_matrix(slots, &mut counts); + for (o, c) in out.iter_mut().zip(counts.iter()) { + o.clear(); + o.extend(c.iter().map(|&v| v != 0)); + } + } + } + } +} + +pub(super) struct PartitionCache { + /// `mats[partition][layer]` = that partition's opened layers. Used by + /// both [`obikindex::KmerIndex::build_sibling_annex`] and + /// [`obikindex::KmerIndex::sibling_annex_stats`]. + mats: Vec>, +} + +impl PartitionCache { + pub(super) 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> = (0..n_parts) + .into_par_iter() + .map(|part| -> OKIResult> { + let index_dir = partition.part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + pb.inc(1); + return Ok(Vec::new()); + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + 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 use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Layer::::open(&layer_dir, &meta.mode).ok().map(Mat::Count) + } else { + Layer::::open(&layer_dir, &meta.mode).ok().map(Mat::Presence) + }; + let Some(mat) = mat else { continue }; + mats.push(mat); + } + pb.inc(1); + Ok(mats) + }) + .collect::>>()?; + pb.finish_and_clear(); + Ok(Self { mats: built }) + } + + /// 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. `find_slot`, not `sub_matrix`/`carries` — no data + /// read needed for a plain existence check. + pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool { + self.mats + .get(dest_partition) + .is_some_and(|mats| mats.iter().any(|mat| mat.find_slot(variant).is_some())) + } + + /// 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 letting `fill_sub_matrix_carries` sort each group by + /// slot internally, 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(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(); mats.len()]; + for &(variant, family_idx, base) in queries { + for (li, mat) in mats.iter().enumerate() { + if let Some(slot) = mat.find_slot(variant) { + by_layer[li].push((slot, family_idx, base)); + break; + } + } + } + + for (li, hits) in by_layer.into_iter().enumerate() { + if hits.is_empty() { + continue; + } + let mat = &mats[li]; + let n_cols = mat.n_cols().min(n_genomes); + let slots: Vec = hits.iter().map(|&(slot, _, _)| slot).collect(); + let mut carries: Vec> = (0..n_cols).map(|_| Vec::new()).collect(); + mat.fill_sub_matrix_carries(&slots, &mut carries); + for (g, col) in carries.iter().enumerate() { + for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) { + if carries_it { + on_hit(family_idx, base, g); + } + } + } + } + } +} diff --git a/src/obikindex/src/siblings/cardinality.rs b/src/obikphylo/src/siblings/cardinality.rs similarity index 89% rename from src/obikindex/src/siblings/cardinality.rs rename to src/obikphylo/src/siblings/cardinality.rs index 5f09415d..c76888fd 100644 --- a/src/obikindex/src/siblings/cardinality.rs +++ b/src/obikphylo/src/siblings/cardinality.rs @@ -5,8 +5,8 @@ use ndarray::Array2; use obikpartitionner::KmerPartition; use obisys::progress_bar; -use crate::error::{OKIError, OKIResult}; -use crate::index::KmerIndex; +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; use super::cache::PartitionCache; use super::distance::RawSnpDistanceOutput; @@ -22,7 +22,8 @@ pub struct CardinalityTally { pub counts: [[u64; 5]; 5], } -impl KmerIndex { +/// Adds [`cardinality_tally`](Self::cardinality_tally) to `KmerIndex`. +pub trait CardinalityExt { /// Cardinality co-occurrence, pooled only over genome pairs whose /// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same /// saturation/no-data exclusion discipline as @@ -51,10 +52,14 @@ impl KmerIndex { /// gets the matching restriction via `scan_family_pairs`'s new /// `variable` flag, rather than a `family_size()` check of its own (it /// doesn't have direct access to the family's mask). - pub fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { + fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult; +} + +impl CardinalityExt for KmerIndex { + fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { let n_parts = self.n_partitions(); - let n_genomes = self.meta.genomes.len(); - let with_counts = self.meta.config.with_counts; + 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; @@ -68,14 +73,14 @@ impl KmerIndex { }); let partition = KmerPartition::open_with_config( - &self.root_path, + self.root_path(), self.kmer_size(), self.minimizer_size(), n_bits, ) .map_err(OKIError::Partition)?; let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); - let layer_dirs = self.sibling_layer_dirs()?; + let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers"); // One layer at a time — see `snp_pseudo_alignment`'s comment for why diff --git a/src/obikindex/src/siblings/distance.rs b/src/obikphylo/src/siblings/distance.rs similarity index 66% rename from src/obikindex/src/siblings/distance.rs rename to src/obikphylo/src/siblings/distance.rs index 3f746303..99347600 100644 --- a/src/obikindex/src/siblings/distance.rs +++ b/src/obikphylo/src/siblings/distance.rs @@ -5,8 +5,8 @@ use ndarray::Array2; use obikpartitionner::KmerPartition; use obisys::progress_bar; -use crate::error::{OKIError, OKIResult}; -use crate::index::KmerIndex; +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; use super::cache::PartitionCache; use super::family_scan::scan_layer_families; @@ -35,54 +35,52 @@ pub struct RawSnpDistanceOutput { pub shared: Array2, } -impl KmerIndex { - /// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance) - /// and [`base_pair_tally`](Self::base_pair_tally): for every family - /// (tallied once, at its minorant) of every layer of the already-built - /// sibling annex, resolves each genome's single observed form (`None` - /// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j, - /// bi, bj, variable)` for every genome pair `(i, j)` where both are - /// unambiguous and single-copy (`bi == bj` means shared at that locus, - /// `bi != bj` means a SNP). `variable` is the family's own - /// `family_size() >= 2` (true if more than one member is observed - /// *anywhere* in the family, i.e. it isn't fully invariant across the - /// whole index) — `raw_snp_distance` ignores it (a fully-invariant - /// family is still legitimately "shared"), but callers whose diagonal - /// should only reflect genuine SNP-adjacent agreement, not the - /// genome-wide invariant background, need it (see - /// [`base_pair_tally`](Self::base_pair_tally)'s `same` field). Layers - /// are processed one at a time, not in parallel — see - /// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers - /// would defeat `scan_layer_families`'s partition-grouped locality; - /// each layer gets its own accumulator from `zero()`, combined - /// pairwise via `combine`. - fn scan_family_pairs( - &self, - label: &str, - zero: impl Fn() -> Acc + Sync, - on_pair: F, - combine: C, - ) -> OKIResult - where - Acc: Send, - F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync, - C: Fn(Acc, Acc) -> Acc, - { - 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; +/// Shared traversal behind [`DistanceExt::raw_snp_distance`] and +/// [`DistanceExt::base_pair_tally`]: for every family (tallied once, at its +/// minorant) of every layer of the already-built sibling annex, resolves +/// each genome's single observed form (`None` if absent or +/// ambiguous/multi-copy), then calls `on_pair(acc, i, j, bi, bj, variable)` +/// for every genome pair `(i, j)` where both are unambiguous and +/// single-copy (`bi == bj` means shared at that locus, `bi != bj` means a +/// SNP). `variable` is the family's own `family_size() >= 2` (true if more +/// than one member is observed *anywhere* in the family, i.e. it isn't +/// fully invariant across the whole index) — `raw_snp_distance` ignores it +/// (a fully-invariant family is still legitimately "shared"), but callers +/// whose diagonal should only reflect genuine SNP-adjacent agreement, not +/// the genome-wide invariant background, need it (see +/// [`DistanceExt::base_pair_tally`]'s `same` field). Layers are processed +/// one at a time, not in parallel — see `snp_pseudo_alignment`'s comment +/// for why `par_iter()` over layers would defeat `scan_layer_families`'s +/// partition-grouped locality; each layer gets its own accumulator from +/// `zero()`, combined pairwise via `combine`. A free function, not a +/// `KmerIndex` method — called only from this file's `DistanceExt` impl. +fn scan_family_pairs( + index: &KmerIndex, + label: &str, + zero: impl Fn() -> Acc + Sync, + on_pair: F, + combine: C, +) -> OKIResult +where + Acc: Send, + F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync, + C: Fn(Acc, Acc) -> Acc, +{ + let n_parts = index.n_partitions(); + let n_genomes = index.meta().genomes.len(); + let with_counts = index.meta().config.with_counts; + let k = index.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; - let partition = KmerPartition::open_with_config( - &self.root_path, - self.kmer_size(), - self.minimizer_size(), - n_bits, - ) - .map_err(OKIError::Partition)?; - let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); - let layer_dirs = self.sibling_layer_dirs()?; + let partition = KmerPartition::open_with_config( + index.root_path(), + index.kmer_size(), + index.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); + let layer_dirs = super::family_scan::sibling_layer_dirs(index)?; let pb = progress_bar(label, layer_dirs.len() as u64, "layers"); // One layer at a time — see `snp_pseudo_alignment`'s comment for why @@ -121,11 +119,36 @@ impl KmerIndex { Ok(total) } +/// Adds [`raw_snp_distance`](Self::raw_snp_distance) and +/// [`base_pair_tally`](Self::base_pair_tally) to `KmerIndex`. +pub trait DistanceExt { /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex - /// (run [`build_sibling_annex`](Self::build_sibling_annex) first). - pub fn raw_snp_distance(&self) -> OKIResult { - let n_genomes = self.meta.genomes.len(); - let (snp, shared) = self.scan_family_pairs( + /// (run [`build_sibling_annex`](super::build::SiblingAnnexBuildExt::build_sibling_annex) first). + fn raw_snp_distance(&self) -> OKIResult; + + /// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG, + /// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs + /// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` — + /// same saturation-exclusion discipline as + /// [`cardinality_tally`](super::cardinality::CardinalityExt::cardinality_tally), + /// for the same reason: a saturated pair's observed base-pair mix + /// trends toward neutral base composition, not the true point-mutation + /// spectrum. + /// + /// A second full pass over the annex, sharing + /// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by + /// it, not a blind re-scan) — needed because `raw_snp_distance` only + /// keeps aggregate SNP/shared counts per genome pair, not which bases + /// were actually involved at each locus, and the ratio-ceiling filter + /// can only be evaluated once the aggregate counts are known. + fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult; +} + +impl DistanceExt for KmerIndex { + fn raw_snp_distance(&self) -> OKIResult { + let n_genomes = self.meta().genomes.len(); + let (snp, shared) = scan_family_pairs( + self, "raw_snp_distance", || (Array2::::zeros((n_genomes, n_genomes)), Array2::::zeros((n_genomes, n_genomes))), |(snp, shared), i, j, bi, bj, _variable| { @@ -146,22 +169,8 @@ impl KmerIndex { Ok(RawSnpDistanceOutput { snp, shared }) } - /// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG, - /// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs - /// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` — - /// same saturation-exclusion discipline as - /// [`cardinality_tally`](Self::cardinality_tally), for the same reason: - /// a saturated pair's observed base-pair mix trends toward neutral base - /// composition, not the true point-mutation spectrum. - /// - /// A second full pass over the annex, sharing - /// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by - /// it, not a blind re-scan) — needed because `raw_snp_distance` only - /// keeps aggregate SNP/shared counts per genome pair, not which bases - /// were actually involved at each locus, and the ratio-ceiling filter - /// can only be evaluated once the aggregate counts are known. - pub fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { - let n_genomes = self.meta.genomes.len(); + fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult { + let n_genomes = self.meta().genomes.len(); let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| { if i == j { return false; @@ -171,7 +180,8 @@ impl KmerIndex { total > 0 && (snp as f64 / total as f64) <= ratio_ceiling }); - let (counts, same) = self.scan_family_pairs( + let (counts, same) = scan_family_pairs( + self, "base_pair_tally", || ([[0u64; 4]; 4], [0u64; 4]), |(counts, same), i, j, bi, bj, variable| { @@ -206,7 +216,7 @@ impl KmerIndex { } } -/// See [`KmerIndex::base_pair_tally`]. +/// See [`DistanceExt::base_pair_tally`]. pub struct BasePairTally { /// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled /// over included genome pairs, where the two genomes' single forms are diff --git a/src/obikindex/src/siblings/family_scan.rs b/src/obikphylo/src/siblings/family_scan.rs similarity index 65% rename from src/obikindex/src/siblings/family_scan.rs rename to src/obikphylo/src/siblings/family_scan.rs index 5aacdd62..56bc5d4c 100644 --- a/src/obikindex/src/siblings/family_scan.rs +++ b/src/obikphylo/src/siblings/family_scan.rs @@ -55,15 +55,16 @@ use rayon::prelude::*; use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex}; use obikseq::CanonicalKmer; -use obilayeredmap::MphfLayer; +use obilayeredmap::Layer; use obilayeredmap::meta::PartitionMeta; use obipipeline::{ThrottleGuard, throttle}; -use crate::error::{OKIError, OKIResult}; -use crate::index::KmerIndex; +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; use super::cache::{Mat, PartitionCache}; -use super::helpers::{central_base, partition_of}; +use super::helpers::central_base; +use super::iter::SiblingEntry; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; /// Families per batch — see the module docs for the memory-vs-per-partition- @@ -73,42 +74,43 @@ use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; /// queries per batch to amortise against. const FAMILY_BATCH: usize = 65536; -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; +/// Every (partition, layer) directory carrying a sibling annex, checked +/// up front so a missing one is reported before any real work starts. +/// Shared by every sibling-annex consumer's extension-trait impl (`alignment`, +/// `build`, `cardinality`, `distance`, `stats`) — a free function, not a +/// `KmerIndex` method, since it is crate-internal only and `KmerIndex` lives +/// in `obikindex`, a foreign crate from here (orphan rule). +pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult> { + let n_parts = index.n_partitions(); + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = index.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); } - let 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); - } + layer_dirs.push(layer_dir); } - Ok(layer_dirs) } + Ok(layer_dirs) } /// 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. +/// layer's own already-open matrix. No separate MPHF/`slot_kmer` either — +/// `mat` (a `Layer`) already bundles the MPHF, and each `SiblingEntry` +/// arrives with its kmer and mask already in hand from `iter_minorants_batch`. struct LayerCtx { - slot_kmer: Vec>, mat: Mat, n_parts: usize, n_genomes: usize, @@ -118,10 +120,11 @@ struct LayerCtx { 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)>, + /// One entry per minorant family in this batch, straight from + /// `iter_minorants_batch` — `order` (iteration-order index, not an MPHF + /// slot; see `docmd/architecture/siblings.md`), `kmer`, and `mask` + /// already carried through, no second annex read. + entries: Vec, _permit: ThrottleGuard, } @@ -143,10 +146,10 @@ enum FamData { Generated(GeneratedBatch), } -/// Visits every minorant family of one layer, in slot order, batched +/// Visits every minorant family of one layer, in iteration order, batched /// [`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 +/// family, in iteration 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 @@ -162,59 +165,38 @@ pub(super) fn scan_layer_families( ) -> 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)?; - - // `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 annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?); let use_counts = with_counts && layer_dir.join("counts").exists(); let mat = if use_counts { - Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + Mat::Count(Layer::::open(layer_dir, &meta.mode).map_err(olm_to_ok)?) } else { - Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + Mat::Presence(Layer::::open(layer_dir, &meta.mode).map_err(olm_to_ok)?) }; let n_cols = mat.n_cols().min(n_genomes); - // 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(()); - } + let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k }); - let ctx = Arc::new(LayerCtx { slot_kmer, mat, n_parts, n_genomes, n_cols, k }); - - 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(); + // Streamed straight from `iter_minorants_batch` (zips this layer's own + // `iter_kmers()` with the annex, both in iteration order — never an + // MPHF slot; see `docmd/architecture/siblings.md`) — never collected + // into a `Vec` first: a layer can hold billions of k-mers, so + // materialising every minorant family up front is exactly the memory + // blowup an earlier version of this traversal was rewritten to avoid + // (see the module docs). `.scan()` computes each batch's starting + // family index lazily, mirroring what the eager `chunks()`+running + // `offset` used to do. + let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| { + let start = *offset; + *offset += batch.len(); + Some((start, batch)) + }); let n_workers = obisys::effective_parallelism(); let capacity = 4; - let throttled = throttle(batches.into_iter(), n_workers).map(|t| SourceBatch { + let throttled = throttle(batches, n_workers).map(|t| SourceBatch { start_family_idx: t.item.0, - slots: t.item.1, + entries: t.item.1, _permit: t.guard, }); @@ -224,17 +206,18 @@ pub(super) fn scan_layer_families( | { move |batch: SourceBatch| -> GeneratedBatch { let ctx = &worker_ctx; - let n = batch.slots.len(); + let n = batch.entries.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"); + // batch's cross-partition queries. Kmer and mask already in + // hand from `iter_minorants_batch`, no second annex read, + // no slot lookup needed for this pass. + for (i, entry) in batch.entries.iter().enumerate() { + let (kmer, mask) = (entry.kmer, entry.mask); masks.push(mask); let base = central_base(kmer, ctx.k); bases.push(base); @@ -246,7 +229,7 @@ pub(super) fn scan_layer_families( if !mask.has(b) { continue; } - let dest = partition_of(other, ctx.n_parts); + let dest = other.partition(ctx.n_parts); outgoing[dest].push((other, i, b)); } } @@ -254,14 +237,21 @@ pub(super) fn scan_layer_families( // 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) { + // docs for the full rationale). `fill_sub_matrix_carries` + // sorts the slots internally for a sequential mmap sweep per + // column, then restores this batch's order — no hand-rolled + // sort/genome-major loop needed here. The presence/count + // matrix is still MPHF-slot-indexed (unlike the annex), so + // each entry's kmer is mapped to its slot via `index_batch` + // — a pure MPHF lookup, no evidence check, since these are + // this layer's own kmers, known members by construction. + let kmers: Vec = batch.entries.iter().map(|e| e.kmer).collect(); + let slots = ctx.mat.index_batch(&kmers); + let mut carries: Vec> = (0..ctx.n_cols).map(|_| Vec::new()).collect(); + ctx.mat.fill_sub_matrix_carries(&slots, &mut carries); + for (g, col) in carries.iter().enumerate() { + for (i, &carries_it) in col.iter().enumerate() { + if carries_it { genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i]; } } @@ -305,7 +295,7 @@ pub(super) fn scan_layer_families( next_expected += n; } } - debug_assert_eq!(next_expected, total_families, "every batch must have been replayed"); + debug_assert!(pending.is_empty(), "every generated batch must have been replayed"); Ok(()) } diff --git a/src/obikphylo/src/siblings/helpers.rs b/src/obikphylo/src/siblings/helpers.rs new file mode 100644 index 00000000..580162b4 --- /dev/null +++ b/src/obikphylo/src/siblings/helpers.rs @@ -0,0 +1,23 @@ +use obikseq::CanonicalKmer; + +use obicompactvec::FamilyMask; + +/// 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] +pub(super) 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`. +pub(super) 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() + }) +} + diff --git a/src/obikphylo/src/siblings/iter.rs b/src/obikphylo/src/siblings/iter.rs new file mode 100644 index 00000000..20bea078 --- /dev/null +++ b/src/obikphylo/src/siblings/iter.rs @@ -0,0 +1,172 @@ +//! Phylo/sibling-domain iteration over a layer — an extension trait, not a +//! new field on `MphfLayer`/`Layer`: "family"/"minorant" are phylo +//! concepts, `obilayeredmap` stays kmer/slot-mapping only (see +//! `docmd/architecture/siblings.md`). +//! +//! The sibling annex is persisted in the same order as `iter_kmers()` +//! (`build_sibling_annex`, see `build.rs`), so pairing them is a plain zip — +//! no MPHF, no slot, no `kmer_at`. Both sides are already `Send + 'static` +//! (`KmerIter` owns an `Arc` clone; `SiblingAnnex` is +//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter` +//! streams straight from disk and can be fed to `obipipeline` batch by +//! batch — never collected whole into memory (see the project's "no full +//! collect" rule). +//! +//! Four iterator types, deliberately mirroring `obilayeredmap`'s own +//! `KmerIter`/`KmerBatchIter` pair (single item vs. `Vec` batch) — plus the +//! minorant-filtered variant of each, since "all siblings" and "one row per +//! family" are both common cases: +//! +//! | | all entries | minorants only | +//! |------------|-------------------|---------------------| +//! | single | [`SiblingIter`] | [`MinorantIter`] | +//! | batch | [`SiblingBatchIter`] | [`MinorantBatchIter`] | +//! +//! No separate "enumerate" variant (unlike `KmerIter`/`enumerate_kmers`): +//! [`SiblingEntry`] already carries `order` for free, since pairing with +//! the annex requires it anyway. + +use std::sync::Arc; + +use obicompactvec::{FamilyMask, SiblingAnnex}; +use obikseq::CanonicalKmer; +use obilayeredmap::{KmerIter, Layer, LayerData}; + +/// One layer entry: a k-mer's position in the layer's iteration order (the +/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer +/// itself, and its family mask. +#[derive(Debug, Clone, Copy)] +pub struct SiblingEntry { + pub order: usize, + pub kmer: CanonicalKmer, + pub mask: FamilyMask, +} + +/// Streams `(order, kmer, mask)` triples for one layer, in iteration order. +/// Produced by [`SiblingLayerExt::iter_siblings`]. +pub struct SiblingIter { + kmers: KmerIter, + annex: Arc, + order: usize, +} + +impl Iterator for SiblingIter { + type Item = SiblingEntry; + + fn next(&mut self) -> Option { + loop { + let kmer = self.kmers.next()?; + let order = self.order; + self.order += 1; + // `None` means "not yet computed" (see `SiblingAnnex` module + // docs) — shouldn't happen against a fully-built annex, but + // skip rather than misalign the two streams if it does. + if let Some(mask) = self.annex.get(order) { + return Some(SiblingEntry { order, kmer, mask }); + } + } + } +} + +/// Batches of [`SiblingIter`]'s entries, `batch_size` at a time — the last +/// batch may be shorter. Produced by [`SiblingLayerExt::iter_siblings_batch`]. +pub struct SiblingBatchIter { + inner: SiblingIter, + batch_size: usize, +} + +impl Iterator for SiblingBatchIter { + type Item = Vec; + + fn next(&mut self) -> Option { + collect_batch(&mut self.inner, self.batch_size) + } +} + +/// Like [`SiblingIter`], filtered to the minorant of each family — the +/// common case, since a family is tallied once, at its minorant. Produced +/// by [`SiblingLayerExt::iter_minorants`]. +pub struct MinorantIter { + inner: SiblingIter, +} + +impl Iterator for MinorantIter { + type Item = SiblingEntry; + + fn next(&mut self) -> Option { + self.inner.by_ref().find(|e| e.mask.is_minorant()) + } +} + +/// Batches of [`MinorantIter`]'s entries, `batch_size` at a time — the last +/// batch may be shorter. Produced by +/// [`SiblingLayerExt::iter_minorants_batch`]. +pub struct MinorantBatchIter { + inner: MinorantIter, + batch_size: usize, +} + +impl Iterator for MinorantBatchIter { + type Item = Vec; + + fn next(&mut self) -> Option { + collect_batch(&mut self.inner, self.batch_size) + } +} + +/// Shared by [`SiblingBatchIter`] and [`MinorantBatchIter`] — pull up to +/// `batch_size` items, `None` once the source is exhausted with nothing left. +fn collect_batch(inner: &mut I, batch_size: usize) -> Option> { + let mut batch = Vec::with_capacity(batch_size); + for _ in 0..batch_size { + match inner.next() { + Some(item) => batch.push(item), + None => break, + } + } + if batch.is_empty() { None } else { Some(batch) } +} + +/// Adds phylo/sibling iteration to any `Layer` — the extension that +/// turns a plain layer into a "sibling layer". Generic over `D` +/// (`LayerData`) rather than implemented once per matrix kind: kmer +/// iteration doesn't depend on the data payload, and `Layer` already +/// delegates `iter_kmers`/`index`/`index_batch` to its inner MPHF for every +/// `D` — reusing that instead of going back through a separate `MphfLayer`. +pub trait SiblingLayerExt { + /// Zip this layer's k-mers with their sibling-annex entry, in iteration + /// order. `annex` must have been built from this same layer (its length + /// must match the layer's k-mer count). + fn iter_siblings(&self, annex: Arc) -> SiblingIter; + + /// Like [`iter_siblings`](Self::iter_siblings), yielding `batch_size` + /// entries at a time. + fn iter_siblings_batch(&self, annex: Arc, batch_size: usize) -> SiblingBatchIter; + + /// Like [`iter_siblings`](Self::iter_siblings), filtered to the + /// minorant of each family — the common case, since a family is + /// tallied once, at its minorant. + fn iter_minorants(&self, annex: Arc) -> MinorantIter; + + /// Like [`iter_minorants`](Self::iter_minorants), yielding `batch_size` + /// minorants at a time. + fn iter_minorants_batch(&self, annex: Arc, batch_size: usize) -> MinorantBatchIter; +} + +impl SiblingLayerExt for Layer { + fn iter_siblings(&self, annex: Arc) -> SiblingIter { + SiblingIter { kmers: self.iter_kmers(), annex, order: 0 } + } + + fn iter_siblings_batch(&self, annex: Arc, batch_size: usize) -> SiblingBatchIter { + SiblingBatchIter { inner: self.iter_siblings(annex), batch_size } + } + + fn iter_minorants(&self, annex: Arc) -> MinorantIter { + MinorantIter { inner: self.iter_siblings(annex) } + } + + fn iter_minorants_batch(&self, annex: Arc, batch_size: usize) -> MinorantBatchIter { + MinorantBatchIter { inner: self.iter_minorants(annex), batch_size } + } +} diff --git a/src/obikindex/src/siblings/mod.rs b/src/obikphylo/src/siblings/mod.rs similarity index 86% rename from src/obikindex/src/siblings/mod.rs rename to src/obikphylo/src/siblings/mod.rs index a91711a1..cd3f5da8 100644 --- a/src/obikindex/src/siblings/mod.rs +++ b/src/obikphylo/src/siblings/mod.rs @@ -12,7 +12,7 @@ //! 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`](crate::index::KmerIndex::sibling_annex_stats) +//! [`sibling_annex_stats`](stats::SiblingStatsExt::sibling_annex_stats) //! below). //! //! Per layer, an `obipipeline` batch transform (throttled — see @@ -55,15 +55,16 @@ mod stats; #[cfg(test)] mod tests; -pub use alignment::SnpAlignment; -pub use cardinality::CardinalityTally; -pub use distance::{BasePairTally, RawSnpDistanceOutput}; -pub use iter::{SiblingEntry, SiblingIter, SiblingLayerExt}; -pub use stats::SiblingAnnexStats; +pub use alignment::{SnpAlignment, SnpAlignmentExt}; +pub use build::SiblingAnnexBuildExt; +pub use cardinality::{CardinalityExt, CardinalityTally}; +pub use distance::{BasePairTally, DistanceExt, RawSnpDistanceOutput}; +pub use iter::{MinorantBatchIter, MinorantIter, SiblingBatchIter, SiblingEntry, SiblingIter, SiblingLayerExt}; +pub use stats::{SiblingAnnexStats, SiblingStatsExt}; use obilayeredmap::OLMError; -use crate::error::OKIError; +use obikindex::OKIError; pub(super) const INDEX_SUBDIR: &str = "index"; pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib"; diff --git a/src/obikindex/src/siblings/stats.rs b/src/obikphylo/src/siblings/stats.rs similarity index 80% rename from src/obikindex/src/siblings/stats.rs rename to src/obikphylo/src/siblings/stats.rs index 0ca0caf9..9cd59da2 100644 --- a/src/obikindex/src/siblings/stats.rs +++ b/src/obikphylo/src/siblings/stats.rs @@ -4,17 +4,17 @@ use obicompactvec::SiblingAnnex; use obikpartitionner::KmerPartition; use obisys::progress_bar; -use crate::error::{OKIError, OKIResult}; -use crate::index::KmerIndex; +use obikindex::{OKIError, OKIResult}; +use obikindex::KmerIndex; use super::ANNEX_FILE_NAME; 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 -/// presence/count data — a separate, occasional diagnostic pass, not fused -/// into construction. +/// annex (see [`super::build::SiblingAnnexBuildExt::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 @@ -33,7 +33,9 @@ pub struct SiblingAnnexStats { pub per_genome: Vec<[u64; 4]>, } -impl KmerIndex { +/// Adds [`sibling_family_size_histogram`](Self::sibling_family_size_histogram) +/// and [`sibling_annex_stats`](Self::sibling_annex_stats) to `KmerIndex`. +pub trait SiblingStatsExt { /// 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 @@ -52,8 +54,19 @@ impl KmerIndex { /// 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()?; + fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]>; + + /// 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`](super::build::SiblingAnnexBuildExt::build_sibling_annex) + /// has not been run on this index first. + fn sibling_annex_stats(&self) -> OKIResult; +} + +impl SiblingStatsExt for KmerIndex { + fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> { + let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; let mut counts = [0u64; 4]; for layer_dir in &layer_dirs { @@ -70,16 +83,10 @@ impl KmerIndex { 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 - /// this index first. - /// - /// [`build_sibling_annex`]: Self::build_sibling_annex - pub fn sibling_annex_stats(&self) -> OKIResult { + 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 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; @@ -87,14 +94,14 @@ impl KmerIndex { // 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.root_path(), self.kmer_size(), self.minimizer_size(), n_bits, ) .map_err(OKIError::Partition)?; let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?); - let layer_dirs = self.sibling_layer_dirs()?; + let layer_dirs = super::family_scan::sibling_layer_dirs(self)?; // One layer at a time, not parallelised across layers — see // `snp_pseudo_alignment`'s comment for why `par_iter()` over layers diff --git a/src/obikindex/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs similarity index 96% rename from src/obikindex/src/siblings/tests.rs rename to src/obikphylo/src/siblings/tests.rs index 866464e5..5aaca00b 100644 --- a/src/obikindex/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -8,11 +8,14 @@ use obilayeredmap::meta::PartitionMeta; use obisys::Reporter; use tempfile::tempdir; -use crate::index::KmerIndex; -use crate::meta::{GenomeInfo, IndexConfig}; -use crate::merge::MergeMode; +use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode}; +use super::alignment::SnpAlignmentExt; +use super::build::SiblingAnnexBuildExt; +use super::cardinality::CardinalityExt; +use super::distance::DistanceExt; use super::helpers::is_minorant; +use super::stats::SiblingStatsExt; use super::{ANNEX_FILE_NAME, INDEX_SUBDIR}; // k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1, @@ -222,7 +225,7 @@ fn family_scan_consumers_agree_on_one_sibling_each() { // 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 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 diff --git a/src/obikseq/src/kmer.rs b/src/obikseq/src/kmer.rs index 365a7851..0fe41101 100644 --- a/src/obikseq/src/kmer.rs +++ b/src/obikseq/src/kmer.rs @@ -367,6 +367,50 @@ impl CanonicalKmerOf { pub fn into_kmer(self) -> KmerOf { KmerOf(self.0, PhantomData) } + + /// This k-mer's own minimiser — as a standalone unit, not part of a + /// streamed sequence. Enumerates all `L::len() - MLen::len() + 1` + /// windows directly on the packed 2-bit representation (no ASCII + /// round-trip, no rolling state, no entropy tracking) and keeps the one + /// with the lowest [`hash_kmer`] — the same selection rule + /// `obiskbuilder::rolling_stat::RollingStat` computes incrementally + /// along a sequence, replicated here in `O(k)` pure bit arithmetic for + /// a single isolated k-mer (the case a streaming rolling scan is + /// unnecessary machinery for). Ties (equal hash) keep the last + /// (highest-position) window, matching `RollingStat`'s monotonic-deque + /// eviction rule (`>=` pops the previous record) — a formal detail, not + /// a practical concern: a `mix64` collision between two distinct + /// m-mers is vanishingly unlikely. + pub fn minimizer(&self) -> Minimizer { + let k = L::len(); + let ml = MLen::len(); + let rc = self.revcomp().raw(); + let mut best_canon: RawKmer = 0; + let mut best_hash = u64::MAX; + for p in 0..=(k - ml) { + let fwd = (self.0 << (2 * p)) >> (KMER_BITS - 2 * ml); + let rev = (rc << (2 * (k - ml - p))) >> (KMER_BITS - 2 * ml); + let canon = fwd.min(rev); + let hash = hash_kmer(canon << (KMER_BITS - 2 * ml)); + if hash < best_hash { + best_hash = hash; + best_canon = canon; + } + } + Minimizer::from_raw_unchecked(best_canon << (KMER_BITS - 2 * ml)) + } + + /// Destination partition for this k-mer, using the project-wide routing + /// rule (`minimizer().seq_hash() & mask`) — the same rule + /// `KmerPartition`/`RoutableSuperKmer` apply to k-mers read from a + /// streamed sequence, here for a standalone k-mer (e.g. a synthetic + /// variant generated outside any sequence). `n_partitions` must be a + /// power of two. + #[inline] + pub fn partition(&self, n_partitions: usize) -> usize { + let mask = (n_partitions as u64) - 1; + (self.minimizer().seq_hash() & mask) as usize + } } impl Sequence for CanonicalKmerOf { diff --git a/src/obilayeredmap/src/layer.rs b/src/obilayeredmap/src/layer.rs index 44797318..c2989a04 100644 --- a/src/obilayeredmap/src/layer.rs +++ b/src/obilayeredmap/src/layer.rs @@ -72,6 +72,15 @@ impl Layer { self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) }) } + /// MPHF + evidence membership check only — no data read. For callers + /// that batch many lookups before touching the matrix at all (e.g. to + /// group hits by column for a later `sub_matrix`/`fill_sub_matrix` + /// sweep), so a plain `query` reading — and discarding — a full row + /// per lookup would be wasted work. + pub fn find_slot(&self, kmer: CanonicalKmer) -> Option { + self.mphf.find(kmer) + } + pub fn n(&self) -> usize { self.mphf.n() } /// Raw MPHF lookup: kmer → slot, no membership check. @@ -187,6 +196,11 @@ impl Layer { .map_err(OLMError::Io) } + /// Number of genome columns in this layer's count matrix. + pub fn n_cols(&self) -> usize { + self.data.n_cols() + } + /// Extract a sub-matrix of counts for the rows at `slots`. /// /// Returns a column-first `Vec>`: one inner `Vec` per genome @@ -218,6 +232,14 @@ impl Layer { .map_err(OLMError::Io) } + /// Number of genome columns in this layer's presence matrix — see + /// `PersistentBitMatrix::n_cols`'s docs for the `Implicit` mono-genome + /// special case (always reports `1`, regardless of the index's real + /// genome count). + pub fn n_cols(&self) -> usize { + self.data.n_cols() + } + /// Extract a sub-matrix of presence/absence for the rows at `slots`. /// /// Returns a column-first `Vec>`: one inner `Vec` per genome