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
2026-08-28 20:54:51 +02:00
parent 5a0b71d105
commit beb2951c20
12 changed files with 768 additions and 27 deletions
+18 -2
View File
@@ -39,9 +39,12 @@ pub fn seq_id(ascii: &[u8]) -> String {
format!("{:016X}", xxh64(ascii, 0))
}
/// Write `seq` as one line of ASCII DNA, followed by a newline.
/// Write `seq` as one line of ASCII text, followed by a newline — plain DNA
/// (A/C/G/T) from every caller in this crate today, but any single-byte
/// ASCII alphabet is sound here (e.g. IUPAC ambiguity codes, `-`/`?`
/// alignment gap/missing markers).
pub fn write_sequence<W: Write>(writer: &mut W, seq: &[u8]) -> io::Result<()> {
// SAFETY: seq is valid ASCII DNA (A/C/G/T).
// SAFETY: every caller in this crate only ever passes valid ASCII bytes.
writeln!(writer, "{}", unsafe { std::str::from_utf8_unchecked(seq) })
}
@@ -59,3 +62,16 @@ pub fn write_record<W: Write>(
writeln!(out)?;
write_sequence(out, seq)
}
/// Plain FASTA record writer — no OBITools JSON annotation block.
///
/// Writes `>{id}\n{sequence}\n`. For callers with no per-record metadata to
/// carry (e.g. a phylogenetic pseudo-alignment's per-genome row, consumed
/// by external tools that don't expect a trailing `{...}` on the header
/// line) — [`write_record`] with an empty `fields` slice would still emit
/// `>{id} {}` (a real, if empty, JSON object plus its leading space), not
/// a bare header.
pub fn write_plain_record<W: Write>(seq: &[u8], id: &str, out: &mut W) -> io::Result<()> {
writeln!(out, ">{id}")?;
write_sequence(out, seq)
}
+5 -2
View File
@@ -14,7 +14,10 @@
//! ```
//!
//! The lower-level primitive [`write_record`] and the [`JsonVal`] type are also
//! public for callers that need custom annotations.
//! public for callers that need custom annotations — or, for callers with no
//! per-record metadata at all (e.g. a phylogenetic pseudo-alignment's rows),
//! [`write_plain_record`] writes a bare `>{id}\n{sequence}\n`, no annotation
//! block.
#![deny(missing_docs)]
@@ -24,7 +27,7 @@ use std::io::{self, Write};
use obikseq::{Minimizer, SuperKmer, Unitig};
pub use fasta::{JsonVal, annotation, seq_id, write_record};
pub use fasta::{JsonVal, annotation, seq_id, write_plain_record, write_record};
// ── public API ────────────────────────────────────────────────────────────────
+56 -6
View File
@@ -36,15 +36,28 @@ impl From<MetricArg> for DistanceMetric {
}
/// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric
/// path (`--metric`/NJ/UPGMA) plus annex construction (`--sibling-annex`)
/// everything else sibling-annex-based (`--snp`, `--sankoff`,
/// `--tnt`/`--phyg`/`--iqtree`, stats, ...) stays in `obikmer` until the rest
/// of `obikphylo::siblings` is reconnected (see the project memory on this).
/// path (`--metric`/NJ/UPGMA), annex construction (`--sibling-annex`),
/// entropy reporting (`--shannon`) and SNP pseudo-alignment sampling
/// (`--pseudo-alignment`, `--subsample`, `--free-loss`, `--no-ambiguity`,
/// `--entropy`/`--entropy-sd`) — everything else sibling-annex-based
/// (`--sankoff`, `--tnt`/`--phyg`/`--iqtree`, stats, ...) stays in `obikmer`
/// until the rest of `obikphylo::siblings` is reconnected (see the project
/// memory on this).
#[derive(Args)]
pub struct PhyloArgs {
/// Index directory
pub index: PathBuf,
/// Exclude a genome (by its exact label) — from `--pseudo-alignment`'s
/// sampling (a family whose only polymorphism lived in an excluded
/// genome is discarded during sampling, not filtered afterward — see
/// `obikphylo::siblings::extensions::SiblingExt::snp_pseudo_alignment`'s
/// own docs) and from the distance matrix / shared-kmer matrix CSV
/// output (row and column both dropped; the underlying computation
/// itself is unaffected). Repeatable.
#[arg(long = "exclude-genome", value_name = "LABEL")]
pub exclude_genome: Vec<String>,
/// Build (or rebuild) the sibling-count/minorant annex — independent of
/// the distance metric below, meant to be run routinely, ahead of any
/// SNP-family distance computation that will later consume it.
@@ -53,11 +66,48 @@ pub struct PhyloArgs {
/// 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.
/// invocation or an earlier one). Always a full, unsampled scan of
/// every family (`--subsample`/`--entropy`/`--entropy-sd` below only
/// apply to `--pseudo-alignment`, not this).
#[arg(long)]
pub shannon: bool,
/// Write a SNP-only pseudo-alignment (FASTA) — requires
/// `--subsample <N>` and an already-built sibling annex
/// (`--sibling-annex` first, in this invocation or an earlier one).
#[arg(long)]
pub pseudo_alignment: bool,
/// Target number of variable sites to sample index-wide for
/// `--pseudo-alignment` — a target, not a guarantee (proportional
/// per-layer sampling; see `obikphylo::siblings::SiblingExt::snp_pseudo_alignment`'s
/// own docs).
#[arg(long)]
pub subsample: Option<usize>,
/// In `--pseudo-alignment`, treat a genome carrying none of a family's
/// observed members (`∅`) as missing data (`?`) rather than a real
/// character state.
#[arg(long)]
pub free_loss: bool,
/// In `--pseudo-alignment`, treat a genome carrying more than one
/// member of a family (ambiguous) as missing data (`?`) rather than an
/// IUPAC ambiguity code.
#[arg(long)]
pub no_ambiguity: bool,
/// Entropy-biased sampling target (Gaussian kernel mean) for
/// `--pseudo-alignment` — activates biasing as soon as this or
/// `--entropy-sd` is given; the other defaults to 1.0/0.5.
#[arg(long)]
pub entropy: Option<f64>,
/// Entropy-biased sampling kernel width (Gaussian standard deviation)
/// for `--pseudo-alignment` — see `--entropy`.
#[arg(long)]
pub entropy_sd: Option<f64>,
/// Distance metric to compute
#[arg(long, value_enum, default_value = "jaccard")]
pub metric: MetricArg,
+75 -9
View File
@@ -5,7 +5,7 @@ use std::sync::Arc;
use obikidxcache::index_cache::IndexCache;
use obikindex::KmerIndex;
use obikphylo::siblings::SiblingExt;
use obikphylo::siblings::{EntropyBias, SiblingExt};
use obikphylo::{Metrics, neighbor_joining, upgma};
use obisys::{Reporter, Stage};
use tracing::info;
@@ -30,6 +30,25 @@ pub fn run(args: PhyloArgs) {
.collect();
let n = labels.len();
// ── Genome exclusion (`--exclude-genome`) ───────────────────────────────────
// Resolved once, up front: `snp_pseudo_alignment` needs it baked into
// sampling itself (see its own docs), and the distance/shared-kmer CSV
// writers below just skip these rows/columns at write time — the
// underlying `cache.distance(...)` computation is unaffected either way.
let exclude_mask: Vec<bool> = {
let mut mask = vec![false; n];
for label in &args.exclude_genome {
match labels.iter().position(|l| l == label) {
Some(i) => mask[i] = true,
None => {
eprintln!("error: --exclude-genome {label:?} does not match any genome in this index");
std::process::exit(1);
}
}
}
mask
};
let mut rep = Reporter::new();
// Every partition/layer this needs is opened once, up front, and
@@ -73,6 +92,48 @@ pub fn run(args: PhyloArgs) {
info!("entropy report → {path}");
}
// ── SNP pseudo-alignment (`--pseudo-alignment`) ─────────────────────────────
if args.pseudo_alignment {
let Some(subsample_n) = args.subsample else {
eprintln!("error: --pseudo-alignment requires --subsample <N>");
std::process::exit(1);
};
// Same activation rule as `--shannon`'s own entropy-biased path
// would use: either flag given activates biasing, the other
// defaults to 1.0/0.5.
let entropy_bias = if args.entropy.is_some() || args.entropy_sd.is_some() {
Some(EntropyBias {
mu: args.entropy.unwrap_or(1.0),
sigma: args.entropy_sd.unwrap_or(0.5),
})
} else {
None
};
info!("sampling SNP pseudo-alignment (target {subsample_n} site(s))");
let t = Stage::start("pseudo_alignment");
let alignment = cache
.snp_pseudo_alignment(subsample_n, args.free_loss, args.no_ambiguity, &exclude_mask, entropy_bias)
.unwrap_or_else(|e| {
eprintln!("error building pseudo-alignment: {e}");
std::process::exit(1);
});
rep.push(t.stop());
let path = args.output.as_ref()
.map(|p| format!("{}_alignment.fasta", p.display()))
.unwrap_or_else(|| "alignment.fasta".into());
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
let n_sites = alignment.sequences.first().map_or(0, Vec::len);
for (&g, seq) in alignment.genome_indices.iter().zip(alignment.sequences.iter()) {
obifastwrite::write_plain_record(seq, &labels[g], &mut f).unwrap();
}
info!("pseudo-alignment ({n_sites} site(s), {} genome(s)) → {path}", alignment.genome_indices.len());
}
info!("computing {:?} distances for {} genome(s)", args.metric, n);
let need_shared = args.shared_kmers || args.nj || args.upgma;
@@ -85,14 +146,19 @@ pub fn run(args: PhyloArgs) {
});
rep.push(t.stop());
// Rows/columns kept in every matrix CSV below — `cache.distance(...)`
// above is computed over every genome regardless; only the writers skip
// excluded ones.
let kept: Vec<usize> = (0..n).filter(|&i| !exclude_mask[i]).collect();
// ── Distance matrix → CSV ─────────────────────────────────────────────────
let write_dist_csv = |w: &mut dyn Write| {
write!(w, "genome").unwrap();
for g in &labels { write!(w, ",{g}").unwrap(); }
for &j in &kept { write!(w, ",{}", labels[j]).unwrap(); }
writeln!(w).unwrap();
for (i, g) in labels.iter().enumerate() {
write!(w, "{g}").unwrap();
for j in 0..n {
for &i in &kept {
write!(w, "{}", labels[i]).unwrap();
for &j in &kept {
write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap();
}
writeln!(w).unwrap();
@@ -127,11 +193,11 @@ pub fn run(args: PhyloArgs) {
std::process::exit(1);
}));
write!(f, "genome").unwrap();
for g in &labels { write!(f, ",{g}").unwrap(); }
for &j in &kept { write!(f, ",{}", labels[j]).unwrap(); }
writeln!(f).unwrap();
for (i, g) in labels.iter().enumerate() {
write!(f, "{g}").unwrap();
for j in 0..n { write!(f, ",{}", shared[[i, j]]).unwrap(); }
for &i in &kept {
write!(f, "{}", labels[i]).unwrap();
for &j in &kept { write!(f, ",{}", shared[[i, j]]).unwrap(); }
writeln!(f).unwrap();
}
info!("shared-kmer matrix → {path}");
@@ -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)
}
+14 -3
View File
@@ -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)
}
}
+1
View File
@@ -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";