perf: optimize entropy computation by pre-filtering monomorphic minorants

Shift monomorphism filtering from the entropy scan layer to a lightweight, annex-only pre-pass. By replacing `Selection::All` with a pre-filtered subset, expensive per-genome resolution is strictly limited to non-monomorphic families. This avoids processing ~98% of minorants that are known to be monomorphic, while preserving positional read speedups for subsequent runs. The change is an internal performance refinement with no public API modifications.
This commit is contained in:
Eric Coissac
2026-08-17 09:32:45 +02:00
parent 7bac0f3850
commit c6cfdac043
3 changed files with 62 additions and 10 deletions
+19 -3
View File
@@ -504,14 +504,30 @@ but indexed by `family_idx` (every minorant of the layer, monomorphic
included — the same numbering `Selection`/`scan_layer_families` already
use), one `f32` entropy15 value per entry, `-1.0` sentinel for monomorphic/
not-yet-computed. First use of `--entropy`/`--entropy-sd` on an index
pays a one-time cost (`ensure_entropy_annexes` in `entropy.rs`: a full,
unsampled `Selection::All` scan, resolving every non-monomorphic
minorant's `genome_mask` once to compute and persist its entropy) — every
pays a one-time cost (`ensure_entropy_annexes` in `entropy.rs`) — every
later run (any `μ`/`σ`, any command) reads the file positionally, no
re-scan, restoring the usual `Selection::Some` "skip resolving excluded
families" speedup that a naive "weigh during the resolving scan" design
would have permanently forfeited.
**Bug found and fixed (2026-08-15): `ensure_entropy_annexes` scanned with
`Selection::All` instead of bounding to non-monomorphic minorants.**
Monomorphism (`family_size() < 2`) is knowable directly from the annex
bits alone, no per-genome resolution needed — but the original
implementation called `scan_layer_families` with `Selection::All`
anyway, so `fill_sub_matrix_carries` (the expensive per-genome
resolution) ran for *every* minorant, ~98% of which are monomorphic
(measured elsewhere in this doc) and had their `genome_mask` immediately
discarded once the callback checked `family_size() < 2`. Fixed by adding
[`subsample::non_monomorphic_selection_layer`] — a cheap, annex-only,
non-sampling pass (same shape as `reservoir_sample_layer`, but keeping
every non-monomorphic minorant's `family_idx` instead of a bounded
reservoir) — and passing `Selection::Some(&eligible)` instead of
`Selection::All`, so the expensive resolution now runs only for the ~2%
of minorants that can actually produce a real entropy value. A
`debug_assert!(mask.family_size() >= 2, ...)` inside the
`scan_layer_families` callback guards the invariant.
**Resolved**: the existing hard "non-monomorphic minorant" eligibility
filter stays a hard gate upstream of the Gaussian weighting — only
qualifying families ever get a stored entropy value or a weighted draw.
+14 -7
View File
@@ -182,9 +182,17 @@ impl ShannonEntropyExt for KmerIndex {
/// [`super::subsample::compute_selections`] whenever entropy-biased
/// selection is requested; see `DevDocMD/architecture/siblings.md`,
/// "Entropy-biased selection", for why this is the one place that pays a
/// full, unsampled scan (every non-monomorphic minorant's `genome_mask`
/// must be resolved to know its entropy) — every later read of the
/// resulting file is a plain positional mmap lookup, not a re-scan.
/// scan (every non-monomorphic minorant's `genome_mask` must be resolved
/// to know its entropy) — every later read of the resulting file is a
/// plain positional mmap lookup, not a re-scan.
///
/// Bounded to non-monomorphic minorants only
/// ([`super::subsample::non_monomorphic_selection_layer`], a cheap
/// annex-only pass — no cross-partition resolution), not
/// `Selection::All`: monomorphism (`family_size() < 2`) 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(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf]) -> OKIResult<()> {
let missing: Vec<usize> = layer_dirs.iter().enumerate()
.filter(|(_, dir)| !dir.join(ENTROPY_ANNEX_FILE_NAME).exists())
@@ -215,10 +223,9 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf])
let layer_dir = &layer_dirs[li];
let mut builder = EntropyAnnexBuilder::new(minorant_counts[li] as usize, &layer_dir.join(ENTROPY_ANNEX_FILE_NAME))
.map_err(OKIError::Io)?;
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |family_idx, mask, genome_mask| {
if mask.family_size() < 2 {
return; // monomorphic minorant — leave the sentinel, never sampled
}
let eligible = super::subsample::non_monomorphic_selection_layer(layer_dir)?;
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &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) {
builder.set(family_idx, h15 as f32);
}
+29
View File
@@ -121,6 +121,35 @@ fn reservoir_sample_layer(layer_dir: &Path, n_layer: usize) -> OKIResult<HashSet
Ok(reservoir.into_iter().collect())
}
/// Every non-monomorphic minorant's `family_idx`, for one layer — not a
/// sample, the *complete* set. Structural, annex-only, no cross-partition
/// resolution (same discipline as [`reservoir_sample_layer`], just keeping
/// everything instead of a bounded reservoir). Used to bound
/// [`super::entropy::ensure_entropy_annexes`]'s `scan_layer_families` call
/// to a `Selection::Some` of exactly the families that need their entropy
/// computed, instead of `Selection::All` — monomorphic minorants (~98% of
/// them, measured) are known to be ineligible straight from the annex bits
/// already, with no need to pay `scan_layer_families`'s expensive
/// per-genome resolution (`fill_sub_matrix_carries`) only to discard the
/// result once `family_size() < 2` is checked *after* resolving.
pub(super) fn non_monomorphic_selection_layer(layer_dir: &Path) -> OKIResult<HashSet<usize>> {
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
let mut selected = HashSet::new();
let mut family_idx: usize = 0;
for slot in 0..annex.len() {
let Some(mask) = annex.get(slot) else { continue };
if !mask.is_minorant() {
continue; // doesn't advance `family_idx` either
}
let this_family_idx = family_idx;
family_idx += 1;
if is_non_monomorphic_minorant(mask) {
selected.insert(this_family_idx);
}
}
Ok(selected)
}
/// Total minorant count per layer (`is_minorant()`, monomorphic included)
/// — sizes the entropy annex (`entropy_annex.rs`), which needs exactly one
/// entry per minorant to keep its `family_idx` numbering aligned with