feat: add phylogenetic sampling, masking, and plain FASTA writer
Introduces configurable phylogenetic sampling and masking controls via new CLI flags (`--subsample`, `--entropy`, `--exclude-genome`, etc.). Adds a complete SNP pseudo-alignment pipeline featuring entropy-biased Gaussian sampling, post-hoc state masking, and proportional per-layer filtering. Extends the FASTA writer with a `write_plain_record` API for bare-header output without JSON annotations.
This commit is contained in:
Eric Coissac
@@ -0,0 +1,63 @@
|
||||
//! SNP-only pseudo-alignment: one row (IUPAC-coded) per genome, one column
|
||||
//! per sampled family — built directly on [`sample_index`], so every family
|
||||
//! this produces has already survived both its Bernoulli draw and post-hoc
|
||||
//! masking ([`super::masking::survives_masking`]).
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
|
||||
use super::masking::{iupac_code, masked_state};
|
||||
use super::subsample::{EntropyBias, sample_index};
|
||||
|
||||
/// `sequences[i]` is `genome_indices[i]`'s IUPAC-coded row — every row the
|
||||
/// same length (one byte per sampled family, in the order [`sample_index`]
|
||||
/// produced them: layer by layer, `family_idx` order within each). An
|
||||
/// excluded genome (`excluded[g]` in
|
||||
/// [`snp_pseudo_alignment`]) has no row at all, not a blanked one — its
|
||||
/// index simply doesn't appear in `genome_indices`. `pub`, not
|
||||
/// `pub(crate)`: part of the public signature of
|
||||
/// [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`].
|
||||
pub struct SnpAlignment {
|
||||
pub sequences: Vec<Vec<u8>>,
|
||||
/// This index's own genome numbering (matching `IndexCache::meta().genomes()`'s
|
||||
/// order) for each row of [`sequences`](Self::sequences), in the same
|
||||
/// order — the caller's authoritative mapping back to genome labels, so
|
||||
/// it never has to re-derive which genomes were excluded itself.
|
||||
pub genome_indices: Vec<usize>,
|
||||
}
|
||||
|
||||
/// See [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`]
|
||||
/// for the public-facing docs — this is its implementation.
|
||||
pub(crate) fn snp_pseudo_alignment(
|
||||
cache: &IndexCache,
|
||||
n: usize,
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<SnpAlignment> {
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let genome_indices: Vec<usize> = (0..n_genomes)
|
||||
.filter(|&g| !excluded.get(g).copied().unwrap_or(false))
|
||||
.collect();
|
||||
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); genome_indices.len()];
|
||||
|
||||
sample_index(
|
||||
cache,
|
||||
n,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, _family_idx, _mask, genome_mask| {
|
||||
for (row, &g) in genome_indices.iter().enumerate() {
|
||||
let ch = masked_state(genome_mask[g], free_loss, no_ambiguity)
|
||||
.map(iupac_code)
|
||||
.unwrap_or(b'?');
|
||||
sequences[row].push(ch);
|
||||
}
|
||||
},
|
||||
)?;
|
||||
|
||||
Ok(SnpAlignment { sequences, genome_indices })
|
||||
}
|
||||
@@ -9,13 +9,14 @@
|
||||
//! this project, not a 4-symbol reduction.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::path::Path;
|
||||
|
||||
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};
|
||||
use super::minorant_selection::{is_non_monomorphic_minorant, non_monomorphic_reindex_layer};
|
||||
use crate::siblings::{ANNEX_FILE_NAME, ENTROPY_ANNEX_FILE_NAME, EntropyAnnex, EntropyAnnexBuilder, SiblingAnnex};
|
||||
|
||||
/// 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
|
||||
@@ -130,3 +131,36 @@ pub(crate) fn ensure_layer_entropy_annex(
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Streams one entropy value per minorant of the layer, **in `family_idx`
|
||||
/// order** — monomorphic minorants included, valued `0.0` by convention
|
||||
/// (trivially uninformative, never computed or stored in the compact
|
||||
/// entropy annex). Exists so a caller can `.zip()` this directly against
|
||||
/// any other `family_idx`-ordered stream (e.g.
|
||||
/// `SiblingLayerExt::iter_minorants`) instead of separately consulting
|
||||
/// [`non_monomorphic_reindex_layer`]'s map at each site.
|
||||
///
|
||||
/// Requires this layer's entropy annex to already exist
|
||||
/// ([`ensure_layer_entropy_annex`] first) — reads it positionally
|
||||
/// (`compact` advances only on non-monomorphic minorants, mirroring
|
||||
/// exactly how the annex was written), never recomputes anything.
|
||||
pub(crate) fn iter_full_entropy(layer_dir: &Path) -> OKIResult<impl Iterator<Item = f32> + use<>> {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).map_err(OKIError::Io)?;
|
||||
let entropy_annex =
|
||||
EntropyAnnex::open(&layer_dir.join(ENTROPY_ANNEX_FILE_NAME)).map_err(OKIError::Io)?;
|
||||
|
||||
let mut compact = 0usize;
|
||||
Ok((0..annex.len()).filter_map(move |slot| {
|
||||
let mask = annex.get(slot)?;
|
||||
if !mask.is_minorant() {
|
||||
return None; // not a minorant at all — no entry in family_idx numbering either
|
||||
}
|
||||
if is_non_monomorphic_minorant(mask) {
|
||||
let value = entropy_annex.get(compact).unwrap_or(0.0);
|
||||
compact += 1;
|
||||
Some(value)
|
||||
} else {
|
||||
Some(0.0) // monomorphic — trivially zero by convention, never stored
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
//! Post-hoc masking applied to a resolved family's per-genome states before
|
||||
//! judging whether it's still a usable alignment site — `--free-loss`
|
||||
//! (absence not treated as a real character), `--no-ambiguity` (a genome
|
||||
//! carrying more than one family member not treated as a real character),
|
||||
//! and `--exclude-genome` (an excluded genome's state never counted at all).
|
||||
//! All three are per-genome, per-family local properties (decidable from
|
||||
//! one family's own `genome_mask` alone), unlike the project's
|
||||
//! rare-character filter, which is relative to the *whole* eventual sample
|
||||
//! and can't be decided per family in isolation — that stays a separate,
|
||||
//! post-hoc pass over the fully assembled alignment, not implemented here.
|
||||
//!
|
||||
//! `--exclude-genome` is checked here, in [`survives_masking`] itself, not
|
||||
//! as a later filter over an already-sampled/already-written alignment: a
|
||||
//! family whose only polymorphism lived in an excluded genome must be
|
||||
//! rejected *during* sampling, the same reasoning `--free-loss`/
|
||||
//! `--no-ambiguity` already established — checking after the fact would
|
||||
//! waste sample quota on sites that die anyway once the excluded genome's
|
||||
//! data is dropped.
|
||||
|
||||
/// This genome's masked character state, or `None` for "?" (missing). `m`
|
||||
/// is one `genome_mask` entry (bit `b` set iff this genome carries the
|
||||
/// family member whose own canonical central base is `b`; `0` = no member
|
||||
/// present in this genome, `∅` — a real state on its own unless
|
||||
/// `free_loss` says otherwise).
|
||||
#[inline]
|
||||
pub(crate) fn masked_state(m: u8, free_loss: bool, no_ambiguity: bool) -> Option<u8> {
|
||||
if free_loss && m == 0 {
|
||||
return None; // absence -> missing, not a real state
|
||||
}
|
||||
if no_ambiguity && m.count_ones() > 1 {
|
||||
return None; // ambiguous (>1 member in this genome) -> missing
|
||||
}
|
||||
Some(m)
|
||||
}
|
||||
|
||||
/// Whether a family's resolved `genome_mask` still carries real
|
||||
/// polymorphism (`>=2` distinct masked states) once [`masked_state`] is
|
||||
/// applied to every *included* genome — i.e. whether it remains a usable
|
||||
/// alignment site rather than one that collapsed to monomorphic (or fully
|
||||
/// missing) once uninformative states, and excluded genomes entirely, are
|
||||
/// dropped from consideration. `excluded[g]` (`false` if `g` is beyond
|
||||
/// `excluded`'s own length — a caller that never excludes anyone can pass
|
||||
/// an empty slice) means genome `g` is skipped outright, not merely masked
|
||||
/// to "?": it never contributes to `seen`, whether or not `free_loss`/
|
||||
/// `no_ambiguity` would also have masked it.
|
||||
pub(crate) fn survives_masking(
|
||||
genome_mask: &[u8],
|
||||
excluded: &[bool],
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
) -> bool {
|
||||
let mut seen: u16 = 0;
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
if excluded.get(g).copied().unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
if let Some(state) = masked_state(m, free_loss, no_ambiguity) {
|
||||
seen |= 1 << state;
|
||||
}
|
||||
}
|
||||
seen.count_ones() >= 2
|
||||
}
|
||||
|
||||
/// IUPAC ambiguity code for a [`masked_state`] value (`0..=15`): single bit
|
||||
/// -> the plain base; 2 or 3 bits -> the matching IUPAC ambiguity code
|
||||
/// (preserves partial information instead of collapsing to `N`, the same
|
||||
/// convention used for diploid heterozygous VCF/FASTA sites); all 4 bits ->
|
||||
/// `N`; no bits (genome carries none of the family's observed members) ->
|
||||
/// `-` (no data at this locus for this genome — distinct from `?`, which
|
||||
/// [`masked_state`] already handles separately for `--free-loss`/
|
||||
/// `--no-ambiguity`-excluded genomes).
|
||||
#[inline]
|
||||
pub(crate) fn iupac_code(state: u8) -> u8 {
|
||||
match state & 0b1111 {
|
||||
0b0000 => b'-',
|
||||
0b0001 => b'A',
|
||||
0b0010 => b'C',
|
||||
0b0100 => b'G',
|
||||
0b1000 => b'T',
|
||||
0b0101 => b'R', // A/G
|
||||
0b1010 => b'Y', // C/T
|
||||
0b0110 => b'S', // C/G
|
||||
0b1001 => b'W', // A/T
|
||||
0b1100 => b'K', // G/T
|
||||
0b0011 => b'M', // A/C
|
||||
0b1110 => b'B', // C/G/T
|
||||
0b1101 => b'D', // A/G/T
|
||||
0b1011 => b'H', // A/C/T
|
||||
0b0111 => b'V', // A/C/G
|
||||
0b1111 => b'N',
|
||||
_ => unreachable!("masked to 4 bits"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn plain_no_masking_counts_absence_as_its_own_state() {
|
||||
// Two genomes: one absent (0), one carrying base A (1) — 2 distinct
|
||||
// states without free-loss, so this survives.
|
||||
assert!(survives_masking(&[0, 1], &[], false, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn free_loss_drops_absence_and_can_collapse_to_monomorphic() {
|
||||
// Same data as above, but with free-loss the absent genome becomes
|
||||
// "?" (excluded) — only one real state (A) remains.
|
||||
assert!(!survives_masking(&[0, 1], &[], true, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_ambiguity_drops_ambiguous_genomes() {
|
||||
// Genome 0 carries both A and C (ambiguous, mask 0b0011); genome 1
|
||||
// carries only A. Without no-ambiguity: 2 distinct states (0b0011,
|
||||
// 0b0001) -> survives. With no-ambiguity: genome 0 becomes "?",
|
||||
// only one real state left -> collapses.
|
||||
assert!(survives_masking(&[0b0011, 0b0001], &[], false, false));
|
||||
assert!(!survives_masking(&[0b0011, 0b0001], &[], false, true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn excluded_genome_never_counted_even_without_free_loss() {
|
||||
// Genome 0 absent (0), genome 1 carries A (1) — without exclusion
|
||||
// this survives (2 distinct states, matching the plain test above).
|
||||
// Excluding genome 1 (the only real signal) leaves just genome 0's
|
||||
// absence — a single state, collapses.
|
||||
assert!(survives_masking(&[0, 1], &[false, false], false, false));
|
||||
assert!(!survives_masking(&[0, 1], &[false, true], false, false));
|
||||
}
|
||||
}
|
||||
@@ -8,11 +8,12 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::TempBitVecBuilder;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex};
|
||||
|
||||
fn is_non_monomorphic_minorant(mask: FamilyMask) -> bool {
|
||||
pub(crate) fn is_non_monomorphic_minorant(mask: FamilyMask) -> bool {
|
||||
mask.is_minorant() && mask.siblings() >= 1
|
||||
}
|
||||
|
||||
@@ -47,3 +48,37 @@ pub(crate) fn non_monomorphic_reindex_layer(layer_dir: &Path) -> OKIResult<HashM
|
||||
}
|
||||
Ok(reindex)
|
||||
}
|
||||
|
||||
/// A bit per minorant of one layer (indexed by `family_idx`, same numbering
|
||||
/// as everywhere else) — set iff non-monomorphic, i.e. eligible for
|
||||
/// `algorithms::subsample`'s circular-walk sampling. Bit-packed
|
||||
/// (`obicompactvec::TempBitVecBuilder`, disk/mmap-backed, 1 bit/entry) —
|
||||
/// not a `HashSet<usize>`/`Vec<bool>`: sampling both reads (`is this family
|
||||
/// still eligible?`) and writes (`mark it taken`) this same structure
|
||||
/// throughout a layer's whole circular walk, so it stays a `Builder`
|
||||
/// (mutable) rather than ever freezing into the read-only `TempBitVec`.
|
||||
///
|
||||
/// Two streamed, annex-only passes (no cross-partition resolution): the
|
||||
/// first counts this layer's total minorants (needed up front to size the
|
||||
/// bitset), the second sets the eligible bits.
|
||||
pub(crate) fn eligible_bitset_layer(layer_dir: &Path) -> OKIResult<TempBitVecBuilder> {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).map_err(OKIError::Io)?;
|
||||
|
||||
let n_minorants = (0..annex.len())
|
||||
.filter(|&slot| annex.get(slot).is_some_and(|m| m.is_minorant()))
|
||||
.count();
|
||||
|
||||
let mut eligible = TempBitVecBuilder::new(n_minorants).map_err(OKIError::Io)?;
|
||||
let mut family_idx = 0usize;
|
||||
for slot in 0..annex.len() {
|
||||
let Some(mask) = annex.get(slot) else { continue };
|
||||
if !mask.is_minorant() {
|
||||
continue;
|
||||
}
|
||||
if is_non_monomorphic_minorant(mask) {
|
||||
eligible.set(family_idx, true);
|
||||
}
|
||||
family_idx += 1;
|
||||
}
|
||||
Ok(eligible)
|
||||
}
|
||||
|
||||
@@ -7,19 +7,30 @@
|
||||
//!
|
||||
//! [`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.
|
||||
//! (`entropy`, [`subsample`], [`alignment`], and future ones — cardinality,
|
||||
//! stats) is built on; [`minorant_selection`] is the cheap annex-only
|
||||
//! groundwork both `entropy` and `subsample` share; [`masking`] is
|
||||
//! `subsample`/`alignment`'s shared post-hoc survival check and character
|
||||
//! encoding.
|
||||
|
||||
mod alignment;
|
||||
mod annex;
|
||||
mod entropy;
|
||||
mod family_scan;
|
||||
mod masking;
|
||||
mod minorant_selection;
|
||||
mod subsample;
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
|
||||
pub use alignment::SnpAlignment;
|
||||
pub use subsample::EntropyBias;
|
||||
|
||||
pub(crate) use alignment::snp_pseudo_alignment;
|
||||
pub(crate) use annex::build_layer_sibling_annex;
|
||||
pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4};
|
||||
pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4, iter_full_entropy};
|
||||
pub(crate) use family_scan::{Selection, scan_layer_families};
|
||||
pub(crate) use subsample::sample_index;
|
||||
|
||||
/// Whether every layer number in `cache` fits in a `FamilyMask` field
|
||||
/// (`< 7`) — a single fact about the whole index, decided once here so
|
||||
|
||||
@@ -0,0 +1,297 @@
|
||||
//! Proportional per-layer sampling of non-monomorphic minorants, bounded to
|
||||
//! `n` sites index-wide — a cross-partition-resolution-cost bound, not a
|
||||
//! full scan (see `DevDocMD/architecture/siblings.md`, "`--subsample`/
|
||||
//! `--shannon`").
|
||||
//!
|
||||
//! Per layer: a quota proportional to that layer's own non-monomorphic
|
||||
//! minorant count (`quota = round(n * count_layer / total_count)`), reached
|
||||
//! by walking [`minorant_selection::eligible_bitset_layer`]'s bit vector
|
||||
//! circularly from a random start, in the same natural (physical, locality-
|
||||
//! friendly) order the rest of this crate already batches through — never a
|
||||
//! full random permutation, which would scatter access across the whole
|
||||
//! layer instead of keeping nearby reads nearby.
|
||||
//!
|
||||
//! For each eligible family visited: a Bernoulli(`p`) draw decides whether
|
||||
//! to *attempt* it this round — `p = p0` (`quota / count_layer`) uniformly,
|
||||
//! or `p = p0 * kernel(entropy)` when [`EntropyBias`] is given (an
|
||||
//! unnormalised Gaussian kernel over that family's already-built entropy
|
||||
//! value — `1` exactly at `entropy == mu`, decaying smoothly away from it,
|
||||
//! no hard cutoff). Entropy is read by zipping the round's own circular
|
||||
//! index sequence against [`super::entropy::iter_full_entropy`]'s streamed
|
||||
//! values (rotated to start at the same point, via
|
||||
//! [`circular_entropy_values`]) — never a `HashMap`/`Vec` materialised
|
||||
//! up front, just two more small mmap opens per round (cheap: page cache,
|
||||
//! not real re-reads) — so biasing costs nothing beyond the plain uniform
|
||||
//! path's own annex-bits work; the uniform path skips this stream
|
||||
//! entirely. Rejected families stay eligible (bit untouched, may be drawn
|
||||
//! again on a later lap); accepted ones are cleared immediately (bit set to
|
||||
//! `false`) regardless of what happens next — once a family has been
|
||||
//! resolved once, it is never revisited, whether it turned out to survive
|
||||
//! post-hoc masking (`algorithms::masking`) or not.
|
||||
//!
|
||||
//! Accepted families are *not* resolved one at a time: `scan_layer_families`
|
||||
//! re-reads a whole layer's k-mer stream on every call, so resolving in
|
||||
//! small per-family increments would repeatedly re-pay that cost. Instead,
|
||||
//! each round gathers a batch of accepted candidates (sized from the
|
||||
//! previous round's observed survival rate, or optimistically `remaining`
|
||||
//! on the first round) and resolves the whole batch in one
|
||||
//! `scan_layer_families` call, bounded to
|
||||
//! [`MAX_TOPUP_ROUNDS`] rounds per layer — a layer that still hasn't
|
||||
//! reached its quota after that many rounds (or has exhausted its whole
|
||||
//! circular walk) simply contributes fewer than its share; `n` sites
|
||||
//! index-wide is a target, never a guarantee (see the project discussion
|
||||
//! this implements).
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::TempBitVecBuilder;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
use rand::RngExt;
|
||||
|
||||
use super::entropy::iter_full_entropy;
|
||||
use super::family_scan::{Selection, scan_layer_families};
|
||||
use super::is_fast_mode;
|
||||
use super::masking::survives_masking;
|
||||
use super::minorant_selection::eligible_bitset_layer;
|
||||
use crate::siblings::FamilyMask;
|
||||
use crate::siblings::extensions::SiblingBuilder;
|
||||
|
||||
/// Resolution rounds a single layer's sampling may take before giving up on
|
||||
/// reaching its own quota — see the module docs on why a round batches many
|
||||
/// accepted families rather than resolving them one at a time.
|
||||
const MAX_TOPUP_ROUNDS: usize = 5;
|
||||
|
||||
/// Gaussian-kernel entropy bias parameters — `mu`/`sigma` are the
|
||||
/// `--entropy`/`--entropy-sd` CLI values (defaulted by the caller, not
|
||||
/// here). `pub`, not `pub(crate)`: part of the public signature of
|
||||
/// [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`].
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct EntropyBias {
|
||||
pub mu: f64,
|
||||
pub sigma: f64,
|
||||
}
|
||||
|
||||
/// [`iter_full_entropy`], rotated to start at `start` and wrap once back to
|
||||
/// it — the same `family_idx` order [`sample_layer`]'s own circular index
|
||||
/// sequence walks, so the two can be `.zip()`-ed directly. Two more streamed
|
||||
/// passes over the layer's (small) annex/entropy files, not a materialised
|
||||
/// `Vec`/`HashMap` — `iter_full_entropy` itself never allocates more than
|
||||
/// its own iterator state.
|
||||
fn circular_entropy_values(
|
||||
layer_dir: &Path,
|
||||
start: usize,
|
||||
) -> OKIResult<impl Iterator<Item = f32> + use<>> {
|
||||
let head = iter_full_entropy(layer_dir)?.skip(start);
|
||||
let tail = iter_full_entropy(layer_dir)?.take(start);
|
||||
Ok(head.chain(tail))
|
||||
}
|
||||
|
||||
/// Samples up to `n` non-monomorphic families index-wide, proportionally
|
||||
/// per layer, calling `on_family(partition, layer, family_idx, mask,
|
||||
/// genome_mask)` once for every family that was both drawn *and* still
|
||||
/// carries real polymorphism after `free_loss`/`no_ambiguity`/`excluded`
|
||||
/// masking (`algorithms::masking::survives_masking`) — i.e. every family
|
||||
/// this produces is already a usable alignment site, no second resolution
|
||||
/// pass needed by the caller. `excluded[g]` (see `survives_masking`'s own
|
||||
/// docs) never counts genome `g` toward that check — `genome_mask` itself
|
||||
/// still carries every genome's own resolved state, excluded or not, so a
|
||||
/// caller building per-genome output (e.g. an alignment row) must apply the
|
||||
/// same exclusion itself when consuming it. Returns the actual number of
|
||||
/// sites kept, which may be less than `n` (see the module docs).
|
||||
pub(crate) fn sample_index(
|
||||
cache: &IndexCache,
|
||||
n: usize,
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
mut on_family: impl FnMut(usize, usize, usize, FamilyMask, &[u8]),
|
||||
) -> OKIResult<usize> {
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let fast_mode = is_fast_mode(cache);
|
||||
|
||||
// Entropy-biased mode reads each candidate's entropy from the entropy
|
||||
// annex (`EntropyWeights::open`, per layer, below) — build any missing
|
||||
// one now, once, rather than mid-walk.
|
||||
if entropy_bias.is_some() {
|
||||
cache.ensure_entropy_annexes()?;
|
||||
}
|
||||
|
||||
// Every cached layer's eligibility bitset, built up front — annex-only,
|
||||
// no cross-partition resolution, so cheap even summed over the whole
|
||||
// index — needed before any layer's quota can be computed (each
|
||||
// layer's share depends on every other layer's own count).
|
||||
let mut layers: Vec<(usize, usize, TempBitVecBuilder, usize)> = Vec::new();
|
||||
for part in cache.partitions() {
|
||||
let n_layer = cache.n_layer(part).unwrap_or(0);
|
||||
for l in 0..n_layer {
|
||||
let layer = cache
|
||||
.get_layer(part, l)
|
||||
.expect("cache-consistent: n_layer(part) already bounds l");
|
||||
let eligible = eligible_bitset_layer(layer.dir())?;
|
||||
let n_minorants = eligible.view().len();
|
||||
layers.push((part, l, eligible, n_minorants));
|
||||
}
|
||||
}
|
||||
|
||||
let total_eligible: u64 = layers.iter().map(|(_, _, e, _)| e.view().count_ones()).sum();
|
||||
if total_eligible == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut total_kept = 0usize;
|
||||
for (part, l, mut eligible, n_minorants) in layers {
|
||||
let count_layer = eligible.view().count_ones();
|
||||
if count_layer == 0 {
|
||||
continue;
|
||||
}
|
||||
let quota = ((n as u128 * count_layer as u128) / total_eligible as u128) as usize;
|
||||
if quota == 0 {
|
||||
continue;
|
||||
}
|
||||
total_kept += sample_layer(
|
||||
cache,
|
||||
part,
|
||||
l,
|
||||
&mut eligible,
|
||||
n_minorants,
|
||||
quota,
|
||||
n_genomes,
|
||||
fast_mode,
|
||||
free_loss,
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|family_idx, mask, genome_mask| on_family(part, l, family_idx, mask, genome_mask),
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(total_kept)
|
||||
}
|
||||
|
||||
/// One layer's share of [`sample_index`] — see the module docs for the
|
||||
/// per-family accept/reject/mark rules and the round-batching scheme.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn sample_layer(
|
||||
cache: &IndexCache,
|
||||
partition: usize,
|
||||
layer_idx: usize,
|
||||
eligible: &mut TempBitVecBuilder,
|
||||
n_minorants: usize,
|
||||
quota: usize,
|
||||
n_genomes: usize,
|
||||
fast_mode: bool,
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
|
||||
) -> OKIResult<usize> {
|
||||
if n_minorants == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
let count_layer = eligible.view().count_ones() as usize;
|
||||
if count_layer == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
let p0 = (quota as f64 / count_layer as f64).min(1.0);
|
||||
let layer_dir = cache
|
||||
.get_layer(partition, layer_idx)
|
||||
.expect("cache-consistent")
|
||||
.dir()
|
||||
.to_path_buf();
|
||||
|
||||
let mut rng = rand::rng();
|
||||
let start = rng.random_range(0..n_minorants);
|
||||
let mut visited = 0usize;
|
||||
let mut kept = 0usize;
|
||||
let mut resolved = 0usize; // total accepted (and resolved) so far, across every round
|
||||
|
||||
for _round in 0..MAX_TOPUP_ROUNDS {
|
||||
// Stop on success (quota reached) or true exhaustion (every
|
||||
// eligible family has been accepted at least once — rejected ones
|
||||
// stay at `1` and keep getting fresh draws lap after lap, so this
|
||||
// can only reach `0` once nothing is left to even attempt). Not
|
||||
// bounded to one lap: `round_start` below already wraps past
|
||||
// `n_minorants` correctly, so a layer with a low masking-survival
|
||||
// rate can and does loop back around to retry earlier rejections,
|
||||
// up to `MAX_TOPUP_ROUNDS` laps.
|
||||
if kept >= quota || eligible.view().count_ones() == 0 {
|
||||
break;
|
||||
}
|
||||
let remaining = quota - kept;
|
||||
// Optimistic on the first round (assume every accepted family will
|
||||
// survive masking); afterwards, size the next batch from the
|
||||
// observed survival rate so a second round is likely to close the
|
||||
// gap rather than needing all `MAX_TOPUP_ROUNDS`.
|
||||
let survival_rate = if resolved == 0 {
|
||||
1.0
|
||||
} else {
|
||||
(kept as f64 / resolved as f64).max(0.05)
|
||||
};
|
||||
let target_this_round = ((remaining as f64 / survival_rate).ceil() as usize).max(1);
|
||||
|
||||
// This round's own circular walk, resuming exactly where the
|
||||
// previous round left off (`start + visited`, wrapped) — zipped
|
||||
// against that same rotation of the entropy stream when biased, so
|
||||
// each candidate's own kernel weight is read in lockstep, never via
|
||||
// a separate lookup structure. Unbiased draws skip the entropy
|
||||
// stream entirely (a flat `1.0` kernel), never touching the entropy
|
||||
// annex.
|
||||
let round_start = (start + visited) % n_minorants;
|
||||
let entropies: Box<dyn Iterator<Item = f32>> = match entropy_bias {
|
||||
Some(_) => Box::new(circular_entropy_values(&layer_dir, round_start)?),
|
||||
None => Box::new(std::iter::repeat(0.0f32)), // never read: kernel forced to 1.0 below
|
||||
};
|
||||
let candidates = (round_start..n_minorants).chain(0..round_start).zip(entropies);
|
||||
|
||||
let mut accepted: HashSet<usize> = HashSet::new();
|
||||
for (family_idx, entropy) in candidates {
|
||||
visited += 1;
|
||||
if !eligible.view().get(family_idx) {
|
||||
continue;
|
||||
}
|
||||
let kernel = match entropy_bias {
|
||||
Some(bias) => {
|
||||
let d = (entropy as f64 - bias.mu) / bias.sigma;
|
||||
(-0.5 * d * d).exp()
|
||||
}
|
||||
None => 1.0,
|
||||
};
|
||||
if rng.random::<f64>() < (p0 * kernel).min(1.0) {
|
||||
eligible.set(family_idx, false);
|
||||
accepted.insert(family_idx);
|
||||
if accepted.len() >= target_this_round {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if accepted.is_empty() {
|
||||
break; // nothing left to try this lap
|
||||
}
|
||||
resolved += accepted.len();
|
||||
|
||||
scan_layer_families(
|
||||
cache,
|
||||
partition,
|
||||
layer_idx,
|
||||
n_genomes,
|
||||
fast_mode,
|
||||
&Selection::Some(&accepted),
|
||||
|family_idx, mask, genome_mask| {
|
||||
debug_assert!(
|
||||
mask.family_size() >= 2,
|
||||
"eligible bitset must only accept non-monomorphic minorants"
|
||||
);
|
||||
if survives_masking(genome_mask, excluded, free_loss, no_ambiguity) {
|
||||
on_family(family_idx, mask, genome_mask);
|
||||
kept += 1;
|
||||
}
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(kept)
|
||||
}
|
||||
@@ -14,8 +14,8 @@ 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,
|
||||
EntropyBias, Selection, SnpAlignment, build_layer_sibling_annex, family_entropy,
|
||||
family_entropy_4, is_fast_mode, scan_layer_families, snp_pseudo_alignment,
|
||||
};
|
||||
use crate::siblings::extensions::SiblingBuilder;
|
||||
use crate::siblings::ENTROPY_ANNEX_FILE_NAME;
|
||||
@@ -54,6 +54,28 @@ pub trait SiblingExt {
|
||||
/// 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<()>;
|
||||
|
||||
/// Build a SNP-only pseudo-alignment (`algorithms::subsample::sample_index`
|
||||
/// — proportional per-layer sampling, `n` sites index-wide, target not
|
||||
/// guarantee) from an already-built sibling annex (run
|
||||
/// [`build_sibling_annex`](Self::build_sibling_annex) first). `free_loss`
|
||||
/// (`--free-loss`), `no_ambiguity` (`--no-ambiguity`) and `excluded`
|
||||
/// (`--exclude-genome`, `excluded[g]` true to drop genome `g`) all feed
|
||||
/// straight into the sampling itself, not a later filter — see
|
||||
/// `algorithms::masking`'s docs: a family whose polymorphism only
|
||||
/// existed via states/genomes these exclude is caught and discarded
|
||||
/// during sampling, not left in the output as a false site. Every
|
||||
/// column of the result has already survived that check, and an
|
||||
/// excluded genome has no row at all in the result (see
|
||||
/// `SnpAlignment::genome_indices`).
|
||||
fn snp_pseudo_alignment(
|
||||
&self,
|
||||
n: usize,
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<SnpAlignment>;
|
||||
}
|
||||
|
||||
impl SiblingExt for IndexCache {
|
||||
@@ -167,4 +189,15 @@ impl SiblingExt for IndexCache {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn snp_pseudo_alignment(
|
||||
&self,
|
||||
n: usize,
|
||||
free_loss: bool,
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<SnpAlignment> {
|
||||
snp_pseudo_alignment(self, n, free_loss, no_ambiguity, excluded, entropy_bias)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@ mod siblingannex;
|
||||
pub(crate) use entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnex, EntropyAnnexBuilder};
|
||||
pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
||||
|
||||
pub use algorithms::{EntropyBias, SnpAlignment};
|
||||
pub use extensions::SiblingExt;
|
||||
|
||||
pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib";
|
||||
|
||||
Reference in New Issue
Block a user