diff --git a/src/obikfilter/src/partition_iter.rs b/src/obikfilter/src/partition_iter.rs index 0d69f570..3793964e 100644 --- a/src/obikfilter/src/partition_iter.rs +++ b/src/obikfilter/src/partition_iter.rs @@ -63,7 +63,7 @@ impl FilteredPartitionIter for IndexCache { ) -> OKIResult { for l in 0..self.n_layer(part).unwrap_or(0) { let layer = self.get_layer(part, l).expect("layer within n_layer(part)"); - if !iter_layer_kmers(layer, use_counts, n_genomes, filters, &mut |kmer, row| cb(kmer, row))? { + if !iter_layer_kmers(&layer, use_counts, n_genomes, filters, &mut |kmer, row| cb(kmer, row))? { return Ok(false); } } @@ -80,7 +80,7 @@ impl FilteredPartitionIter for IndexCache { ) -> OKIResult { for l in 0..self.n_layer(part).unwrap_or(0) { let layer = self.get_layer(part, l).expect("layer within n_layer(part)"); - if !iter_layer_kmers(layer, use_counts, n_genomes, filters, &mut |kmer, row| cb(part, l, kmer, row))? { + if !iter_layer_kmers(&layer, use_counts, n_genomes, filters, &mut |kmer, row| cb(part, l, kmer, row))? { return Ok(false); } } diff --git a/src/obikidxcache/src/index_cache.rs b/src/obikidxcache/src/index_cache.rs index 2f72e141..8dcbfb71 100644 --- a/src/obikidxcache/src/index_cache.rs +++ b/src/obikidxcache/src/index_cache.rs @@ -28,7 +28,28 @@ pub struct IndexCache { /// `0`. A `HashMap` rather than a flat `Vec` because the set of cached /// partitions doesn't have to be `0..n_partitions`, or even contiguous /// — see [`new`](Self::new). - layer_cache: HashMap>, + /// + /// `Arc` per layer, and the whole map itself behind one more + /// `Arc` — every accessor below hands out owned `Arc` + /// clones (never a `&KmerLayer` tied to `&self`'s lifetime), because + /// that's the only thing that can satisfy the `Send + 'static` bound + /// `obipipeline::Pipe::apply` (real `thread::spawn` workers) requires + /// of anything moved into them. A caller needing to hand one specific + /// cached layer to such a pipeline would otherwise have no way to do so + /// short of reopening it independently — a redundant mmap/MPHF open of + /// a layer this cache already holds (confirmed by + /// `obikphylo::siblings::algorithms::family_scan` hitting exactly this + /// before `Arc` was introduced here). + /// + /// The outer `Arc>` is what lets [`iter`](Self::iter)/ + /// [`iter_partition`](Self::iter_partition)/[`iter_indexed`](Self::iter_indexed) + /// return a genuine `Iterator` object (`LayerIter`, below) instead of + /// eagerly materialising every layer's `Arc` into a throwaway `Vec` + /// first: cloning the outer `Arc` is one refcount bump regardless of + /// how many layers this cache holds, so `LayerIter::next()` can look + /// up and clone one `Arc` at a time, on demand, straight out + /// of this same shared map — no separate payload copy ever exists. + layer_cache: Arc>>>, } impl IndexCache { @@ -64,7 +85,7 @@ impl IndexCache { "IndexCache::new: failed to open layer (partition {p}, layer {l}): {e} — partition is not fully built" ) }); - layers.push(layer); + layers.push(Arc::new(layer)); } layer_cache.insert(p, layers); } @@ -73,7 +94,7 @@ impl IndexCache { IndexCache { raw_index: index, meta, - layer_cache, + layer_cache: Arc::new(layer_cache), } } @@ -109,22 +130,52 @@ impl IndexCache { self.layer_cache.keys().copied() } + /// A cheap `Arc` clone (refcount bump, no reopening) of one cached + /// layer — owned, not tied to `&self`'s lifetime, so it can be moved + /// into a context that must own it, e.g. a `Send + 'static` pipeline + /// worker. See the [`layer_cache`](Self::layer_cache) field docs for + /// why every accessor here hands out `Arc` rather than + /// `&KmerLayer`. #[inline] - pub fn get_layer(&self, partition: usize, layer: usize) -> Option<&KmerLayer> { - self.layer_cache.get(&partition)?.get(layer) + pub fn get_layer(&self, partition: usize, layer: usize) -> Option> { + self.layer_cache.get(&partition)?.get(layer).cloned() } - pub fn iter(&self) -> impl Iterator { - self.layer_cache.values().flatten() + /// Every cached layer, across every cached partition, as owned `Arc` + /// clones, in `(partition, layer)` order — genuinely `'static`: see + /// [`iter_indexed`](Self::iter_indexed)/[`LayerIter`]'s own docs. + pub fn iter(&self) -> impl Iterator> + 'static { + self.iter_indexed().map(|(_, _, layer)| layer) } - pub fn iter_indexed(&self) -> impl Iterator { - self.layer_cache.iter().flat_map(|(&p, layers)| { - layers - .iter() - .enumerate() - .map(move |(l, layer)| (p, l, layer)) - }) + /// Every cached layer of *one* partition, in layer order — never the + /// other cached partitions. For a caller that already knows which + /// partition a k-mer routes to (`CanonicalKmer::partition`, the same + /// rule [`find`](Self::find) uses) and only needs to scan that one + /// partition's own layers, not every partition this cache holds. Empty + /// (not an error) if `partition` isn't cached — same convention as + /// [`get_layer`](Self::get_layer). `'static` for the same reason + /// [`iter_indexed`](Self::iter_indexed) is. + pub fn iter_partition(&self, partition: usize) -> impl Iterator> + 'static { + LayerIter { + cache: Arc::clone(&self.layer_cache), + partitions: vec![partition].into_iter(), + current: None, + } + .map(|(_, _, layer)| layer) + } + + /// Every cached layer, tagged with its own `(partition, layer)` + /// position — see [`LayerIter`]'s own docs for why this is a real + /// `next()`-driven iterator, one `Arc` clone at a time out + /// of the shared cache map, rather than a `Vec>` + /// materialised up front. + pub fn iter_indexed(&self) -> impl Iterator)> + 'static { + LayerIter { + cache: Arc::clone(&self.layer_cache), + partitions: self.layer_cache.keys().copied().collect::>().into_iter(), + current: None, + } } /// Iterate over the unitigs of every cached layer, chained — the @@ -133,7 +184,7 @@ impl IndexCache { /// (one partition, several, or the whole index — see [`new`](Self::new)). /// No extra opening: every layer here is already open, this only /// streams what [`iter`](Self::iter) already gives out. - pub fn iter_unitigs(&self) -> impl Iterator + '_ { + pub fn iter_unitigs(&self) -> impl Iterator + 'static { self.iter().flat_map(|l| l.iter_unitigs()) } @@ -224,7 +275,7 @@ impl IndexCache { .flat_map(|(_, layers)| layers.par_iter()) .map(|layer| { let _guard = throttle.acquire_guard(); - map(layer) + map(layer.as_ref()) }) .reduce(&identity, &reduce) } @@ -244,7 +295,7 @@ impl IndexCache { .flat_map(|(_, layers)| layers.par_iter()) .for_each(|layer| { let _guard = throttle.acquire_guard(); - f(layer); + f(layer.as_ref()); }); } @@ -268,8 +319,43 @@ impl IndexCache { .par_iter() .flat_map(|(_, layers)| layers.par_iter()) .map(move |layer| obipipeline::Throttled { - item: layer, + item: layer.as_ref(), guard: throttle.acquire_guard(), }) } } + +/// A real `Iterator` (implements `next()`), not a `Vec` pre-filled and then +/// walked — `next()` looks up and clones exactly one `Arc` at a +/// time, straight out of `cache` (a clone of `IndexCache`'s own shared map, +/// one refcount bump to construct regardless of how many layers are cached), +/// so no separate copy of the layers being walked ever exists. `partitions` +/// is a small, already-owned `Vec` of which partition(s) to walk, in +/// order — holding a `Vec` of plain partition numbers costs nothing close +/// to what holding a `Vec` of the cached layers themselves would. +pub struct LayerIter { + cache: Arc>>>, + partitions: std::vec::IntoIter, + /// `(partition, next layer index)` of the partition currently being + /// walked — `None` once that partition is exhausted or before the + /// first `next()` call, at which point [`partitions`](Self::partitions) + /// is advanced to pick up the next one. + current: Option<(usize, usize)>, +} + +impl Iterator for LayerIter { + type Item = (usize, usize, Arc); + + fn next(&mut self) -> Option { + loop { + if let Some((partition, layer_idx)) = self.current { + if let Some(layer) = self.cache.get(&partition).and_then(|v| v.get(layer_idx)) { + self.current = Some((partition, layer_idx + 1)); + return Some((partition, layer_idx, Arc::clone(layer))); + } + self.current = None; // this partition exhausted + } + self.current = Some((self.partitions.next()?, 0)); + } + } +} diff --git a/src/obikmer2/src/cmd/phylo/args.rs b/src/obikmer2/src/cmd/phylo/args.rs index df31e7e0..b2f6c2a6 100644 --- a/src/obikmer2/src/cmd/phylo/args.rs +++ b/src/obikmer2/src/cmd/phylo/args.rs @@ -51,6 +51,13 @@ pub struct PhyloArgs { #[arg(long)] pub sibling_annex: bool, + /// Write a per-family Shannon entropy report (CSV) — requires an + /// already-built sibling annex (`--sibling-annex` first, in this + /// invocation or an earlier one). Full, unsampled scan of every family; + /// no `--subsample`/`--entropy-bias` yet. + #[arg(long)] + pub shannon: bool, + /// Distance metric to compute #[arg(long, value_enum, default_value = "jaccard")] pub metric: MetricArg, diff --git a/src/obikmer2/src/cmd/phylo/mod.rs b/src/obikmer2/src/cmd/phylo/mod.rs index e7d13270..d6c9ea96 100644 --- a/src/obikmer2/src/cmd/phylo/mod.rs +++ b/src/obikmer2/src/cmd/phylo/mod.rs @@ -58,6 +58,21 @@ pub fn run(args: PhyloArgs) { rep.push(t.stop()); } + // ── Shannon entropy report (`--shannon`) ──────────────────────────────────── + if args.shannon { + let path = args.output.as_ref() + .map(|p| format!("{}_entropy.csv", p.display())) + .unwrap_or_else(|| "entropy.csv".into()); + info!("computing per-family Shannon entropy"); + let t = Stage::start("shannon_entropy"); + cache.shannon_entropy_csv(std::path::Path::new(&path)).unwrap_or_else(|e| { + eprintln!("error computing Shannon entropy: {e}"); + std::process::exit(1); + }); + rep.push(t.stop()); + info!("entropy report → {path}"); + } + info!("computing {:?} distances for {} genome(s)", args.metric, n); let need_shared = args.shared_kmers || args.nj || args.upgma; diff --git a/src/obikphylo/src/siblings/algorithms/entropy.rs b/src/obikphylo/src/siblings/algorithms/entropy.rs new file mode 100644 index 00000000..a96be69d --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/entropy.rs @@ -0,0 +1,132 @@ +//! Shannon entropy (bits) of a family's per-genome states, over the 15 +//! non-empty subsets of `{A,C,G,T}` — see +//! `DevDocMD/theory/evolutionary_distances.md`, "Entropy definition", for +//! the design discussion this implements: the project's calibrated 16-state +//! Sankoff cost matrix treats every observed combination of bases as its +//! own first-class state, so entropy over that same 15-symbol space (state +//! `0`/`∅` excluded — a family is only informative among the genomes where +//! it is actually observed) is the consistent informativeness proxy for +//! this project, not a 4-symbol reduction. + +use std::collections::HashSet; + +use obikidxcache::index_cache::IndexCache; +use obikindex::{OKIError, OKIResult}; + +use super::family_scan::{Selection, scan_layer_families}; +use super::minorant_selection::non_monomorphic_reindex_layer; +use crate::siblings::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnexBuilder}; + +/// Shannon entropy (bits) of one family's per-genome states, over the 15 +/// non-empty subsets of `{A,C,G,T}` actually observed among the genomes +/// carrying it (`genome_mask[g] != 0`) — genomes where the family is absent +/// are excluded from the denominator, not scored as a 16th state. `None` +/// only if the family is absent from every genome (shouldn't happen for a +/// real minorant, guarded against rather than dividing by zero). Second +/// value: the denominator (`n_genomes_present`) used, for traceability. +pub(crate) fn family_entropy(genome_mask: &[u8]) -> Option<(f64, usize)> { + let mut freq = [0u32; 16]; // index 0 (∅) never incremented, kept for direct `m as usize` indexing + let mut present = 0u32; + for &m in genome_mask { + if m != 0 { + freq[m as usize] += 1; + present += 1; + } + } + if present == 0 { + return None; + } + let mut h = 0.0f64; + for &c in &freq[1..] { + if c == 0 { + continue; + } + let p = f64::from(c) / f64::from(present); + h -= p * p.log2(); + } + Some((h, present as usize)) +} + +/// Shannon entropy (bits) of one family's per-genome states, reduced to the +/// 4 plain nucleotide symbols instead of the 15-state space above — kept +/// alongside it (not instead of it) purely to measure how much the two +/// diverge on real data. A genome carrying several bases at once +/// (`genome_mask[g]` with more than one bit set) contributes to *each* of +/// those bases' counts — counted once per base present, not split +/// fractionally, and not folded into a combined state the way +/// [`family_entropy`] does. The denominator is therefore the total number +/// of base-occurrences summed over every present genome, not the genome +/// count. `None` under the same condition as [`family_entropy`]. +pub(crate) fn family_entropy_4(genome_mask: &[u8]) -> Option<(f64, usize)> { + let mut freq = [0u32; 4]; + let mut total = 0u32; + for &m in genome_mask { + for (base, count) in freq.iter_mut().enumerate() { + if m & (1 << base) != 0 { + *count += 1; + total += 1; + } + } + } + if total == 0 { + return None; + } + let mut h = 0.0f64; + for &c in &freq { + if c == 0 { + continue; + } + let p = f64::from(c) / f64::from(total); + h -= p * p.log2(); + } + Some((h, total as usize)) +} + +/// Builds one layer's entropy annex if it doesn't already exist — a no-op, +/// cheap existence check, once it does. Bounded to non-monomorphic +/// minorants only ([`non_monomorphic_reindex_layer`], a cheap annex-only +/// pass — no cross-partition resolution): monomorphism is already known +/// from the annex bits alone, so there is no reason to pay +/// `scan_layer_families`'s expensive per-genome resolution for the ~98% of +/// minorants that are monomorphic only to discard the result. +pub(crate) fn ensure_layer_entropy_annex( + cache: &IndexCache, + partition: usize, + layer_idx: usize, + n_genomes: usize, + fast_mode: bool, +) -> OKIResult<()> { + let layer = cache + .get_layer(partition, layer_idx) + .expect("cache-consistent: caller already bounds layer_idx by n_layer(partition)"); + let path = layer.dir().join(ENTROPY_ANNEX_FILE_NAME); + if path.exists() { + return Ok(()); + } + + let reindex = non_monomorphic_reindex_layer(layer.dir())?; + let mut builder = EntropyAnnexBuilder::new(reindex.len(), &path).map_err(OKIError::Io)?; + let eligible: HashSet = reindex.keys().copied().collect(); + + scan_layer_families( + cache, + partition, + layer_idx, + n_genomes, + fast_mode, + &Selection::Some(&eligible), + |family_idx, mask, genome_mask| { + debug_assert!( + mask.family_size() >= 2, + "Selection::Some(eligible) must only visit non-monomorphic minorants" + ); + if let Some((h15, _)) = family_entropy(genome_mask) { + let compact = reindex[&family_idx]; + builder.set(compact, h15 as f32); + } + }, + )?; + builder.close().map_err(OKIError::Io)?; + + Ok(()) +} diff --git a/src/obikphylo/src/siblings/algorithms/family_scan.rs b/src/obikphylo/src/siblings/algorithms/family_scan.rs index 53f426d0..692b7272 100644 --- a/src/obikphylo/src/siblings/algorithms/family_scan.rs +++ b/src/obikphylo/src/siblings/algorithms/family_scan.rs @@ -30,12 +30,13 @@ //! (`n_workers` batches), not by the layer's size. //! //! The pipeline's own working context (`LayerCtx`) needs an owned, -//! `Send + 'static` `KmerLayer` — `IndexCache`'s own copy of this same layer -//! is only ever borrowed (`&'_ KmerLayer`, tied to the caller's stack frame), -//! so it can't be moved into the worker threads `Pipe::apply` spawns; this -//! reopens the layer once, independently, for that purpose. Cross-partition -//! resolution, by contrast, runs synchronously on the calling thread (via -//! `rayon::scope`-free `par_iter`, no `thread::spawn`), so it can — and +//! `Send + 'static` handle on the source layer — a plain `&KmerLayer` +//! borrowed from `cache` is tied to the caller's stack frame, so it can't be +//! moved into the worker threads `Pipe::apply` spawns. `IndexCache::get_layer` +//! hands out a cheap `Arc` clone (every `IndexCache` accessor +//! does, see its own docs), so no reopening/re-mmap'ing is needed. +//! Cross-partition resolution, by contrast, runs synchronously on the +//! calling thread (`rayon::par_iter`, no `thread::spawn`), so it can — and //! does — borrow `cache` directly for the whole call. use std::collections::HashMap; @@ -55,11 +56,26 @@ use crate::siblings::helpers::central_base; use crate::siblings::iter::{SiblingEntry, SiblingLayerExt}; use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex}; -/// Families per batch — trades off per-batch memory (`FAMILY_BATCH * -/// genome_count` bytes) against giving each partition a decent number of -/// queries to resolve in one go; too small a batch starves that density -/// regardless of how the resolution step is parallelised. -const FAMILY_BATCH: usize = 65536; +/// Target memory budget for one batch's `genome_mask` allocation +/// (`family_batch_size(n_genomes) * n_genomes` bytes, see below) — bounding +/// *that*, not the family count directly, is what actually matters: a fixed +/// family-count batch (this module's original `FAMILY_BATCH = 65536`) costs +/// ~64KB/batch at a handful of genomes but ~655MB/batch at the 10,000-skim- +/// genome scale this project targets, with `n_workers` such batches +/// potentially in flight at once (see [`scan_layer_families`]'s own +/// `throttle` call) — the same fixed constant silently stops being a small +/// number once genome count grows, for no reason the batch size itself +/// expresses. +const FAMILY_BATCH_BYTES: usize = 16 * 1024 * 1024; + +/// Families per batch, derived from [`FAMILY_BATCH_BYTES`] and this run's +/// own genome count — clamped so a very small genome count doesn't produce +/// an absurdly large batch (upper bound: the original fixed default, still +/// a reasonable ceiling once genome count is negligible) nor a very large +/// one starve per-partition query density down to nothing (lower bound). +fn family_batch_size(n_genomes: usize) -> usize { + (FAMILY_BATCH_BYTES / n_genomes.max(1)).clamp(1024, 65536) +} /// Restricts [`scan_layer_families`] to a subset of a layer's families, /// identified by their iteration-order index. `All` (the only variant any @@ -82,7 +98,7 @@ fn is_selected(selected: &Option>, family_idx: /// `cache` here: generation never touches the cross-partition cache, only /// this layer's own already-open matrix. struct LayerCtx { - mat: KmerLayer, + mat: Arc, n_parts: usize, n_genomes: usize, n_cols: usize, @@ -121,8 +137,10 @@ enum FamData { } /// 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_batch_size`] families at a time — see the module docs for why +/// that's derived from `n_genomes` rather than fixed, and for why +/// generation and resolution use different concurrency. `on_family` is +/// called once per /// family, in iteration order, with that family's iteration-order index /// (the same numbering [`Selection`] indices are drawn from), its own annex /// mask, and its per-genome base-presence (`genome_mask[g]`: bit `b` set @@ -147,17 +165,15 @@ pub(crate) fn scan_layer_families( let n_parts = cache.index().n_partitions(); let k = cache.index().kmer_size(); - // An owned, independently-opened copy of this same layer — the shared - // `cache`'s own copy is only ever borrowed, so it can't be moved into - // the `Send + 'static` pipeline workers below (see the module docs). - let partition = cache - .index() - .partition(source_partition) - .map_err(|e| OKIError::InvalidInput(format!("partition {source_partition}: {e}")))?; - let mat = partition - .layer(source_layer) - .and_then(KmerLayer::open) - .map_err(|e| OKIError::InvalidInput(format!("layer {source_layer}: {e}")))?; + // A cheap `Arc` clone of the cache's own already-open layer — `Send + + // 'static`, so it can move into the pipeline workers below without + // reopening/re-mmap'ing anything (see the module docs and + // `IndexCache::get_layer_arc`'s own). + let mat = cache.get_layer(source_partition, source_layer).ok_or_else(|| { + OKIError::InvalidInput(format!( + "layer {source_layer} of partition {source_partition} not in this cache" + )) + })?; let n_cols = mat.n_cols().min(n_genomes); let layer_dir: &Path = mat.dir(); let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).map_err(OKIError::Io)?); @@ -170,7 +186,7 @@ pub(crate) fn scan_layer_families( // billions of k-mers. let batches = ctx .mat - .iter_minorants_batch(annex, FAMILY_BATCH) + .iter_minorants_batch(annex, family_batch_size(n_genomes)) .scan(0usize, |offset, batch| { let start = *offset; *offset += batch.len(); @@ -367,7 +383,7 @@ fn resolve_partition_batch( .zip(entries.iter()) .map(|(slot, &(_, family_idx, base))| (slot, family_idx, base)) .collect(); - resolve_layer_hits(layer, &hits, n_genomes, &mut on_hit); + resolve_layer_hits(&layer, &hits, n_genomes, &mut on_hit); } } else { // First hit wins, same semantics as the fast path — a variant @@ -391,7 +407,7 @@ fn resolve_partition_batch( let layer = cache .get_layer(dest_partition, li) .expect("cache-consistent: n_layer(dest_partition) already bounds li"); - resolve_layer_hits(layer, &hits, n_genomes, &mut on_hit); + resolve_layer_hits(&layer, &hits, n_genomes, &mut on_hit); } } } @@ -406,6 +422,18 @@ fn resolve_layer_hits( n_genomes: usize, on_hit: &mut impl FnMut(usize, u8, usize), ) { + // TODO(memory): `carries` is `n_cols * hits.len()` `bool`s — Rust's + // `Vec` is 1 byte/entry, not bit-packed, so this spends a whole + // byte per bit of real information. Worst case (all of one batch's + // variants landing on the same destination layer) at 10,000 genomes: + // tens to low hundreds of MB per call (see the conversation that + // measured this — `KmerLayer::Count`'s own `fill_sub_matrix_carries` + // transiently doubles that again via an intermediate `Vec>`). + // A real fix means `fill_sub_matrix_carries`'s signature changing to a + // bit-packed output (e.g. `bitvec`/a raw `Vec` bitmask) — that's + // an `obikindex` API change (`content_layer.rs`), not something this + // call site can fix alone; not done yet, flagged for when genome-count + // scale makes it worth the churn. let n_cols = layer.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(); diff --git a/src/obikphylo/src/siblings/algorithms/minorant_selection.rs b/src/obikphylo/src/siblings/algorithms/minorant_selection.rs new file mode 100644 index 00000000..7fd62c09 --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/minorant_selection.rs @@ -0,0 +1,49 @@ +//! Cheap, purely structural (annex-bits-only, no cross-partition +//! resolution) passes over an already-built sibling annex — monomorphism +//! (`FamilyMask::family_size() < 2`) is knowable straight from the annex's +//! own bits, so a downstream consumer that only cares about non-monomorphic +//! families never needs to pay [`super::family_scan::scan_layer_families`]'s +//! expensive cross-partition sweep just to find that out. + +use std::collections::HashMap; +use std::path::Path; + +use obikindex::{OKIError, OKIResult}; + +use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex}; + +fn is_non_monomorphic_minorant(mask: FamilyMask) -> bool { + mask.is_minorant() && mask.siblings() >= 1 +} + +/// `family_idx` (`scan_layer_families`'s own numbering — every minorant, +/// monomorphic included) -> a compact ordinal among *non-monomorphic* +/// minorants only, for one layer. One streamed pass over the layer's +/// already-built sibling annex; no `IndexCache`, no cross-partition +/// resolution. +/// +/// Exists so a downstream per-family array (the entropy annex) can be sized +/// to the non-monomorphic count alone (~2% of a layer's minorants, measured +/// on real data) instead of wasting an entry per monomorphic minorant, which +/// never carries a meaningful value anyway (entropy is trivially undefined +/// for a family with no sibling). The map's own key set doubles as exactly +/// the `Selection` a caller needs to bound `scan_layer_families` to just +/// these families. +pub(crate) fn non_monomorphic_reindex_layer(layer_dir: &Path) -> OKIResult> { + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).map_err(OKIError::Io)?; + let mut reindex = HashMap::new(); + let mut family_idx = 0usize; + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + if !mask.is_minorant() { + continue; // not a minorant at all — doesn't advance `family_idx` either + } + let this_family_idx = family_idx; + family_idx += 1; + if is_non_monomorphic_minorant(mask) { + let compact = reindex.len(); + reindex.insert(this_family_idx, compact); + } + } + Ok(reindex) +} diff --git a/src/obikphylo/src/siblings/algorithms/mod.rs b/src/obikphylo/src/siblings/algorithms/mod.rs index 37595e25..91daf397 100644 --- a/src/obikphylo/src/siblings/algorithms/mod.rs +++ b/src/obikphylo/src/siblings/algorithms/mod.rs @@ -1,13 +1,34 @@ -//! Sibling-annex construction algorithm: the cross-partition sweep that -//! fills in one layer's [`crate::siblings::FamilyMask`] slots — driven, one -//! layer at a time, by -//! [`crate::siblings::extensions::SiblingExt`]. Kept apart from -//! that trait (and from the `siblingannex`/`helpers` data model) the same -//! way `obikindexer` splits `algorithms` from `extensions`: this module -//! depends on the data model, never the reverse. +//! Sibling-annex construction algorithms — driven, one layer (or one +//! partition/layer sweep) at a time, by the traits in +//! [`crate::siblings::extensions`]. Kept apart from those traits (and from +//! the `siblingannex`/`entropy_annex`/`helpers` data model) the same way +//! `obikindexer` splits `algorithms` from `extensions`: this module depends +//! on the data model, never the reverse. +//! +//! [`annex`] builds the sibling-count/minorant annex itself; [`family_scan`] +//! is the shared cross-partition traversal every other consumer +//! (`entropy`, and future ones — cardinality, alignment, stats) is built +//! on; [`minorant_selection`] and [`entropy`] are the first such consumer. mod annex; +mod entropy; mod family_scan; +mod minorant_selection; + +use obikidxcache::index_cache::IndexCache; pub(crate) use annex::build_layer_sibling_annex; +pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4}; pub(crate) use family_scan::{Selection, scan_layer_families}; + +/// Whether every layer number in `cache` fits in a `FamilyMask` field +/// (`< 7`) — a single fact about the whole index, decided once here so +/// every writer/reader of a per-layer `FamilyMask` field agrees. Layer +/// count is guaranteed identical across every partition (a merge adds one +/// layer to all of them at once), so the max over cached partitions speaks +/// for the whole index. Shared by every `extensions` impl that writes or +/// interprets that field (`SiblingExt::build_sibling_annex`, +/// `SiblingBuilder::ensure_entropy_annexes`), so they can never disagree. +pub(crate) fn is_fast_mode(cache: &IndexCache) -> bool { + cache.partitions().filter_map(|p| cache.n_layer(p)).max().unwrap_or(0) <= 7 +} diff --git a/src/obikphylo/src/siblings/entropy_annex.rs b/src/obikphylo/src/siblings/entropy_annex.rs new file mode 100644 index 00000000..3e89fa48 --- /dev/null +++ b/src/obikphylo/src/siblings/entropy_annex.rs @@ -0,0 +1,153 @@ +//! Persisted, per-layer Shannon entropy (`entropy15`, see +//! `algorithms::entropy`), one value per **non-monomorphic minorant +//! family** — compact: indexed by the dedicated ordinal +//! `algorithms::non_monomorphic_reindex_layer` computes (skips monomorphic +//! minorants entirely, ~98% of minorants on real data), never by +//! `scan_layer_families`'s own `family_idx` (which counts every minorant, +//! monomorphic included). +//! +//! Exists so repeated reads of the same layer's entropy values (e.g. +//! entropy-biased family selection, once reconnected) are a plain +//! positional mmap lookup instead of re-paying `scan_layer_families`'s +//! cross-partition resolution — building this file costs exactly that same +//! full scan the first time, it does not make any single computation of +//! entropy itself faster. +//! +//! Mirrors `SiblingAnnex`'s on-disk convention (magic + header + flat +//! per-entry data, read-only after build), just with `f32` payloads and a +//! different (compact) numbering — see the module docs above. + +use std::fs::{File, OpenOptions}; +use std::io; +use std::path::Path; + +use memmap2::{Mmap, MmapMut}; + +const MAGIC: [u8; 4] = *b"PENT"; + +// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (4 bytes/entry) follows. +const HEADER_SIZE: usize = 16; + +/// Sentinel: not yet computed. Entropy is always `>= 0.0`, so any negative +/// value is unambiguous — chosen over `NaN` to avoid `NaN`-comparison +/// footguns at every read site. +const SENTINEL: f32 = -1.0; + +pub(crate) const ENTROPY_ANNEX_FILE_NAME: &str = "entropy.pent"; + +pub(crate) struct EntropyAnnex { + mmap: Mmap, + n: usize, +} + +impl EntropyAnnex { + pub(crate) fn open(path: &Path) -> io::Result { + let mmap = unsafe { Mmap::map(&File::open(path)?)? }; + if mmap.len() < HEADER_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "PENT file too short")); + } + if mmap[0..4] != MAGIC { + return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PENT magic")); + } + let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize; + if mmap.len() < HEADER_SIZE + n * 4 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "PENT file truncated")); + } + Ok(Self { mmap, n }) + } + + pub(crate) fn len(&self) -> usize { + self.n + } + + /// `None` for the sentinel (not yet computed — shouldn't happen against + /// a fully-built annex). + pub(crate) fn get(&self, idx: usize) -> Option { + let off = HEADER_SIZE + idx * 4; + let v = f32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap()); + (v >= 0.0).then_some(v) + } +} + +pub(crate) struct EntropyAnnexBuilder { + mmap: MmapMut, +} + +impl EntropyAnnexBuilder { + /// Create a new annex of `n` entries (a layer's non-monomorphic + /// minorant count) at `path`, pre-initialised to the sentinel. + pub(crate) fn new(n: usize, path: &Path) -> io::Result { + let file_size = HEADER_SIZE + n * 4; + let file = OpenOptions::new() + .read(true).write(true).create(true).truncate(true) + .open(path)?; + file.set_len(file_size as u64)?; + let mut mmap = unsafe { MmapMut::map_mut(&file)? }; + mmap[0..4].copy_from_slice(&MAGIC); + mmap[4..8].copy_from_slice(&[0u8; 4]); + mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes()); + for i in 0..n { + let off = HEADER_SIZE + i * 4; + mmap[off..off + 4].copy_from_slice(&SENTINEL.to_le_bytes()); + } + Ok(Self { mmap }) + } + + pub(crate) fn set(&mut self, idx: usize, value: f32) { + debug_assert!(value >= 0.0, "entropy must be non-negative, got {value}"); + let off = HEADER_SIZE + idx * 4; + self.mmap[off..off + 4].copy_from_slice(&value.to_le_bytes()); + } + + pub(crate) fn close(self) -> io::Result<()> { + self.mmap.flush() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn sentinel_is_negative_and_unset_entries_read_as_none() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.pent"); + let builder = EntropyAnnexBuilder::new(4, &path).unwrap(); + builder.close().unwrap(); + let annex = EntropyAnnex::open(&path).unwrap(); + for i in 0..4 { + assert_eq!(annex.get(i), None); + } + } + + #[test] + fn roundtrip_including_zero_entropy() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.pent"); + let mut builder = EntropyAnnexBuilder::new(4, &path).unwrap(); + let values = [0.0f32, 1.5, 3.9, 0.0001]; + for (i, &v) in values.iter().enumerate() { + builder.set(i, v); + } + builder.close().unwrap(); + let annex = EntropyAnnex::open(&path).unwrap(); + for (i, &v) in values.iter().enumerate() { + assert_eq!(annex.get(i), Some(v), "entry {i}"); + } + } + + #[test] + fn mixed_sentinel_and_real_values() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.pent"); + let mut builder = EntropyAnnexBuilder::new(3, &path).unwrap(); + builder.set(1, 2.0); + builder.close().unwrap(); + let annex = EntropyAnnex::open(&path).unwrap(); + assert_eq!(annex.get(0), None); + assert_eq!(annex.get(1), Some(2.0)); + assert_eq!(annex.get(2), None); + assert_eq!(annex.len(), 3); + } +} diff --git a/src/obikphylo/src/siblings/extensions/annex_build.rs b/src/obikphylo/src/siblings/extensions/annex_build.rs deleted file mode 100644 index acd2519b..00000000 --- a/src/obikphylo/src/siblings/extensions/annex_build.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! [`SiblingExt`] — the public entry point for building the -//! sibling-count/minorant annex of an already-cached index. Thin: it only -//! walks `cache`'s partitions/layers and delegates each layer's actual -//! cross-partition sweep to [`crate::siblings::algorithms::build_layer_sibling_annex`] -//! — same split as `obikphylo::distance::Metrics`, whose `IndexCache` -//! extension impl is itself thin dispatch onto `obicompactvec`'s -//! finalisation methods. - -use obikidxcache::index_cache::IndexCache; -use obikindex::OKIResult; -use obisys::progress_bar; - -use crate::siblings::algorithms::build_layer_sibling_annex; - -/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `IndexCache`. -pub trait SiblingExt { - /// Build the sibling-count/minorant annex for every layer of every - /// cached partition, writing one annex file per layer alongside its - /// existing index files. Safe to call again later (e.g. after a fresh - /// merge, on a freshly-rebuilt `IndexCache`) — each run simply - /// overwrites the annex files of the layers it covers. - /// - /// Construction only — no statistics gathered here on purpose: this is - /// meant to run routinely (it is the artefact the SNP-family distances - /// consume), while the sibling-count distribution is a separate, - /// occasional diagnostic pass over the result. - /// - /// 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. - fn build_sibling_annex(&self) -> OKIResult<()>; -} - -impl SiblingExt for IndexCache { - fn build_sibling_annex(&self) -> OKIResult<()> { - let n_parts = self.index().n_partitions(); - - // Whether every layer number fits in a `FamilyMask` field (`< 7`) — - // a single fact about the whole index, decided once here so the - // writer (this pass) and every future reader agree. Layer count is - // guaranteed identical across every partition (a merge adds one - // layer to all of them at once), so the max over cached partitions - // speaks for the whole index. - let fast_mode = self - .partitions() - .filter_map(|p| self.n_layer(p)) - .max() - .unwrap_or(0) - <= 7; - if !fast_mode { - tracing::warn!( - "index has more than 7 layers per partition — sibling-annex fast layer lookup \ - disabled; consider compacting this index" - ); - } - - let pb = progress_bar("sibling_annex", self.n_partition() as u64, "partitions"); - let mut total_slots: u64 = 0; - for part in self.partitions() { - let n_layer = self.n_layer(part).unwrap_or(0); - let mut part_slots: u64 = 0; - for l in 0..n_layer { - let layer = self - .get_layer(part, l) - .expect("cache-consistent: n_layer(part) already bounds l"); - part_slots += build_layer_sibling_annex(self, layer, l, n_parts, fast_mode)?; - } - total_slots += part_slots; - pb.inc(1); - pb.set_message(format!( - "partition {part}: {part_slots} kmers ({total_slots} total)" - )); - } - pb.finish_and_clear(); - tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions"); - - Ok(()) - } -} diff --git a/src/obikphylo/src/siblings/extensions/mod.rs b/src/obikphylo/src/siblings/extensions/mod.rs index d5be0833..74a0aa94 100644 --- a/src/obikphylo/src/siblings/extensions/mod.rs +++ b/src/obikphylo/src/siblings/extensions/mod.rs @@ -4,6 +4,8 @@ //! other `obikindex`/`obikidxcache` extension trait in this codebase (see //! `obikphylo::distance::Metrics`, `obikindexer::extensions`). -mod annex_build; +mod sibling_builder; +mod sibling_ext; -pub use annex_build::SiblingExt; +pub use sibling_ext::SiblingExt; +pub(crate) use sibling_builder::SiblingBuilder; diff --git a/src/obikphylo/src/siblings/extensions/sibling_builder.rs b/src/obikphylo/src/siblings/extensions/sibling_builder.rs new file mode 100644 index 00000000..b8915941 --- /dev/null +++ b/src/obikphylo/src/siblings/extensions/sibling_builder.rs @@ -0,0 +1,40 @@ +//! [`SiblingBuilder`] — construction triggered *internally*, as a side +//! effect of a [`crate::siblings::extensions::SiblingExt`] method (e.g. +//! `shannon_entropy_csv` ensuring its layers' entropy annexes exist before +//! reading them), never called directly from outside this crate. Unlike +//! `SiblingExt` (explicit, CLI-facing actions), nothing here is meant to be +//! its own CLI flag — same `pub`/`pub(crate)` split `obikindexer` uses for +//! its own `IndexBuilder`/`PrivateBuilder`. + +use obikidxcache::index_cache::IndexCache; +use obikindex::OKIResult; +use obisys::progress_bar; + +use crate::siblings::algorithms::{ensure_layer_entropy_annex, is_fast_mode}; + +pub(crate) trait SiblingBuilder { + /// Build every cached layer's entropy annex that doesn't already exist + /// — a no-op, cheap existence check per layer, once every layer's + /// annex file exists. See `ensure_layer_entropy_annex`'s own docs for + /// why this is bounded to non-monomorphic minorants only. + fn ensure_entropy_annexes(&self) -> OKIResult<()>; +} + +impl SiblingBuilder for IndexCache { + fn ensure_entropy_annexes(&self) -> OKIResult<()> { + let n_genomes = self.meta().genomes().len(); + let fast_mode = is_fast_mode(self); + + let pb = progress_bar("entropy_annex", self.n_partition() as u64, "partitions"); + for part in self.partitions() { + let n_layer = self.n_layer(part).unwrap_or(0); + for l in 0..n_layer { + ensure_layer_entropy_annex(self, part, l, n_genomes, fast_mode)?; + } + pb.inc(1); + } + pb.finish_and_clear(); + + Ok(()) + } +} diff --git a/src/obikphylo/src/siblings/extensions/sibling_ext.rs b/src/obikphylo/src/siblings/extensions/sibling_ext.rs new file mode 100644 index 00000000..9bf2a98a --- /dev/null +++ b/src/obikphylo/src/siblings/extensions/sibling_ext.rs @@ -0,0 +1,170 @@ +//! [`SiblingExt`] — public, explicit, CLI-facing operations on an +//! already-cached index: today, building the sibling-count/minorant annex +//! and computing per-family Shannon entropy. Thin: each method only walks +//! `cache`'s partitions/layers and delegates the actual per-layer work to +//! [`crate::siblings::algorithms`] — same split as +//! `obikphylo::distance::Metrics`, whose `IndexCache` extension impl is +//! itself thin dispatch onto `obicompactvec`'s finalisation methods. + +use std::io::{BufWriter, Write}; +use std::path::Path; + +use obikidxcache::index_cache::IndexCache; +use obikindex::{OKIError, OKIResult}; +use obisys::progress_bar; + +use crate::siblings::algorithms::{ + build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode, scan_layer_families, + Selection, +}; +use crate::siblings::extensions::SiblingBuilder; +use crate::siblings::ENTROPY_ANNEX_FILE_NAME; + +pub trait SiblingExt { + /// Build the sibling-count/minorant annex for every layer of every + /// cached partition, writing one annex file per layer alongside its + /// existing index files, then the entropy annex on top of it + /// ([`SiblingBuilder::ensure_entropy_annexes`] — precomputed here rather + /// than left to build lazily on first `--shannon`/`--entropy-bias` use). + /// Safe to call again later (e.g. after a fresh merge, on a + /// freshly-rebuilt `IndexCache`) — each run simply overwrites the + /// sibling-annex files of the layers it covers, and rebuilds any + /// entropy annex that becomes stale as a result (a fresh sibling annex + /// makes the previous entropy annex's `family_idx` numbering + /// untrustworthy, so it must be rebuilt too, not merely left in place). + /// + /// No statistics gathered here on purpose: this is meant to run + /// routinely (both annexes are inputs the SNP-family distances and + /// entropy report consume), while the sibling-count distribution is a + /// separate, occasional diagnostic pass over the result. + /// + /// 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. + fn build_sibling_annex(&self) -> OKIResult<()>; + + /// Write a CSV (`partition,layer,family_idx,entropy15,entropy4, + /// family_size,n_genomes_present`) of per-family Shannon entropy — both + /// [`family_entropy`] (15 non-empty states, the project's settled + /// definition) and [`family_entropy_4`] (plain 4-symbol reduction, kept + /// alongside for comparison) — one row per non-monomorphic minorant of + /// every cached layer, from an already-built sibling annex (run + /// [`build_sibling_annex`](Self::build_sibling_annex) first). + /// + /// Full, unsampled scan of every cached layer (`Selection::All`) — no + /// `--subsample`/`--entropy-bias` yet (see the project memory on this). + fn shannon_entropy_csv(&self, path: &Path) -> OKIResult<()>; +} + +impl SiblingExt for IndexCache { + fn build_sibling_annex(&self) -> OKIResult<()> { + let n_parts = self.index().n_partitions(); + let fast_mode = is_fast_mode(self); + if !fast_mode { + tracing::warn!( + "index has more than 7 layers per partition — sibling-annex fast layer lookup \ + disabled; consider compacting this index" + ); + } + + let pb = progress_bar("sibling_annex", self.n_partition() as u64, "partitions"); + let mut total_slots: u64 = 0; + for part in self.partitions() { + let n_layer = self.n_layer(part).unwrap_or(0); + let mut part_slots: u64 = 0; + for l in 0..n_layer { + let layer = self + .get_layer(part, l) + .expect("cache-consistent: n_layer(part) already bounds l"); + part_slots += build_layer_sibling_annex(self, &layer, l, n_parts, fast_mode)?; + + // Invalidate this layer's entropy annex, if any — it was + // built against the sibling annex's *previous* + // `family_idx` numbering (`non_monomorphic_reindex_layer` + // reads the sibling annex), now stale. `ensure_entropy_annexes` + // below only builds what's *missing*, so the old file must + // be removed here for it to be rebuilt rather than silently + // left stale. + let entropy_path = layer.dir().join(ENTROPY_ANNEX_FILE_NAME); + if entropy_path.exists() { + std::fs::remove_file(&entropy_path).map_err(OKIError::Io)?; + } + } + total_slots += part_slots; + pb.inc(1); + pb.set_message(format!( + "partition {part}: {part_slots} kmers ({total_slots} total)" + )); + } + pb.finish_and_clear(); + tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions"); + + // Precompute the entropy annex too, now that every layer's fresh + // sibling annex is on disk (`ensure_entropy_annexes` reads it via + // `non_monomorphic_reindex_layer`) — same routine-rebuild footing + // as the sibling annex itself, rather than leaving it to be built + // lazily on first `--shannon`/`--entropy-bias` use. + tracing::info!("building entropy annex"); + self.ensure_entropy_annexes()?; + + Ok(()) + } + + fn shannon_entropy_csv(&self, path: &Path) -> OKIResult<()> { + // Entropy annexes aren't this method's own output, but building + // them is exactly the same full scan this method is about to run + // anyway — see `SiblingBuilder::ensure_entropy_annexes`'s own docs + // for why that's a side effect worth taking here rather than a + // separate, explicit step. + self.ensure_entropy_annexes()?; + + let n_genomes = self.meta().genomes().len(); + let fast_mode = is_fast_mode(self); + + let mut f = BufWriter::new(std::fs::File::create(path).map_err(OKIError::Io)?); + writeln!(f, "partition,layer,family_idx,entropy15,entropy4,family_size,n_genomes_present") + .map_err(OKIError::Io)?; + + let pb = progress_bar("shannon_entropy", self.n_partition() as u64, "partitions"); + for part in self.partitions() { + let n_layer = self.n_layer(part).unwrap_or(0); + for l in 0..n_layer { + let mut write_err = None; + scan_layer_families( + self, + part, + l, + n_genomes, + fast_mode, + &Selection::All, + |family_idx, mask, genome_mask| { + if write_err.is_some() { + return; + } + if mask.family_size() < 2 { + return; // monomorphic family — entropy trivially 0, no signal, skip + } + let Some((h15, n_present)) = family_entropy(genome_mask) else { + return; + }; + let h4 = family_entropy_4(genome_mask).map_or(0.0, |(h, _)| h); + if let Err(e) = writeln!( + f, + "{part},{l},{family_idx},{h15:.6},{h4:.6},{},{n_present}", + mask.family_size() + ) { + write_err = Some(e); + } + }, + )?; + if let Some(e) = write_err { + return Err(OKIError::Io(e)); + } + } + pb.inc(1); + } + pb.finish_and_clear(); + + Ok(()) + } +} diff --git a/src/obikphylo/src/siblings/mod.rs b/src/obikphylo/src/siblings/mod.rs index 5d3377a1..8b37c931 100644 --- a/src/obikphylo/src/siblings/mod.rs +++ b/src/obikphylo/src/siblings/mod.rs @@ -21,10 +21,12 @@ pub mod algorithms; pub mod extensions; +mod entropy_annex; mod helpers; mod iter; mod siblingannex; +pub(crate) use entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnex, EntropyAnnexBuilder}; pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; pub use extensions::SiblingExt;