Add per-family Shannon entropy calculation and reporting
Introduces an entropy.pent persisted format and algorithms to compute per-family Shannon entropy across genomic partitions. Refactors IndexCache to use Arc-based layer sharing for efficient iteration and restructures the siblings module with extension traits for annex generation. Adds a --shannon CLI flag to export partition-level entropy metrics to CSV.
This commit is contained in:
Eric Coissac
@@ -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<usize> = 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(())
|
||||
}
|
||||
@@ -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<KmerLayer>` 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<std::collections::HashSet<usize>>, 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<KmerLayer>,
|
||||
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<bool>` 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<Vec<u32>>`).
|
||||
// A real fix means `fill_sub_matrix_carries`'s signature changing to a
|
||||
// bit-packed output (e.g. `bitvec`/a raw `Vec<u64>` 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<usize> = hits.iter().map(|&(slot, _, _)| slot).collect();
|
||||
let mut carries: Vec<Vec<bool>> = (0..n_cols).map(|_| Vec::new()).collect();
|
||||
|
||||
@@ -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<HashMap<usize, usize>> {
|
||||
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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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<Self> {
|
||||
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<f32> {
|
||||
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<Self> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user