From beb2951c20214a928760810110b7469953545a56 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 27 Aug 2026 15:26:53 +0200 Subject: [PATCH] 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. --- src/obifastwrite/src/fasta.rs | 20 +- src/obifastwrite/src/lib.rs | 7 +- src/obikmer2/src/cmd/phylo/args.rs | 62 +++- src/obikmer2/src/cmd/phylo/mod.rs | 84 ++++- .../src/siblings/algorithms/alignment.rs | 63 ++++ .../src/siblings/algorithms/entropy.rs | 38 ++- .../src/siblings/algorithms/masking.rs | 132 ++++++++ .../siblings/algorithms/minorant_selection.rs | 37 ++- src/obikphylo/src/siblings/algorithms/mod.rs | 17 +- .../src/siblings/algorithms/subsample.rs | 297 ++++++++++++++++++ .../src/siblings/extensions/sibling_ext.rs | 37 ++- src/obikphylo/src/siblings/mod.rs | 1 + 12 files changed, 768 insertions(+), 27 deletions(-) create mode 100644 src/obikphylo/src/siblings/algorithms/alignment.rs create mode 100644 src/obikphylo/src/siblings/algorithms/masking.rs create mode 100644 src/obikphylo/src/siblings/algorithms/subsample.rs diff --git a/src/obifastwrite/src/fasta.rs b/src/obifastwrite/src/fasta.rs index 3374bbd5..e27ea414 100644 --- a/src/obifastwrite/src/fasta.rs +++ b/src/obifastwrite/src/fasta.rs @@ -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(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( 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(seq: &[u8], id: &str, out: &mut W) -> io::Result<()> { + writeln!(out, ">{id}")?; + write_sequence(out, seq) +} diff --git a/src/obifastwrite/src/lib.rs b/src/obifastwrite/src/lib.rs index 3e066cd1..3df19b58 100644 --- a/src/obifastwrite/src/lib.rs +++ b/src/obifastwrite/src/lib.rs @@ -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 ──────────────────────────────────────────────────────────────── diff --git a/src/obikmer2/src/cmd/phylo/args.rs b/src/obikmer2/src/cmd/phylo/args.rs index b2f6c2a6..dcf9a7d5 100644 --- a/src/obikmer2/src/cmd/phylo/args.rs +++ b/src/obikmer2/src/cmd/phylo/args.rs @@ -36,15 +36,28 @@ impl From 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, + /// 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 ` 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, + + /// 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, + + /// Entropy-biased sampling kernel width (Gaussian standard deviation) + /// for `--pseudo-alignment` — see `--entropy`. + #[arg(long)] + pub entropy_sd: Option, + /// Distance metric to compute #[arg(long, value_enum, default_value = "jaccard")] pub metric: MetricArg, diff --git a/src/obikmer2/src/cmd/phylo/mod.rs b/src/obikmer2/src/cmd/phylo/mod.rs index d6c9ea96..9e0958fb 100644 --- a/src/obikmer2/src/cmd/phylo/mod.rs +++ b/src/obikmer2/src/cmd/phylo/mod.rs @@ -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 = { + 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 "); + 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 = (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}"); diff --git a/src/obikphylo/src/siblings/algorithms/alignment.rs b/src/obikphylo/src/siblings/algorithms/alignment.rs new file mode 100644 index 00000000..71094742 --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/alignment.rs @@ -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>, + /// 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, +} + +/// 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, +) -> OKIResult { + let n_genomes = cache.meta().genomes().len(); + let genome_indices: Vec = (0..n_genomes) + .filter(|&g| !excluded.get(g).copied().unwrap_or(false)) + .collect(); + let mut sequences: Vec> = 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 }) +} diff --git a/src/obikphylo/src/siblings/algorithms/entropy.rs b/src/obikphylo/src/siblings/algorithms/entropy.rs index a96be69d..c1baad05 100644 --- a/src/obikphylo/src/siblings/algorithms/entropy.rs +++ b/src/obikphylo/src/siblings/algorithms/entropy.rs @@ -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 + 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 + } + })) +} diff --git a/src/obikphylo/src/siblings/algorithms/masking.rs b/src/obikphylo/src/siblings/algorithms/masking.rs new file mode 100644 index 00000000..9ae56b34 --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/masking.rs @@ -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 { + 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)); + } +} diff --git a/src/obikphylo/src/siblings/algorithms/minorant_selection.rs b/src/obikphylo/src/siblings/algorithms/minorant_selection.rs index 7fd62c09..8da4a33e 100644 --- a/src/obikphylo/src/siblings/algorithms/minorant_selection.rs +++ b/src/obikphylo/src/siblings/algorithms/minorant_selection.rs @@ -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`/`Vec`: 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 { + 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) +} diff --git a/src/obikphylo/src/siblings/algorithms/mod.rs b/src/obikphylo/src/siblings/algorithms/mod.rs index 91daf397..416a00c2 100644 --- a/src/obikphylo/src/siblings/algorithms/mod.rs +++ b/src/obikphylo/src/siblings/algorithms/mod.rs @@ -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 diff --git a/src/obikphylo/src/siblings/algorithms/subsample.rs b/src/obikphylo/src/siblings/algorithms/subsample.rs new file mode 100644 index 00000000..1767707c --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/subsample.rs @@ -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 + 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, + mut on_family: impl FnMut(usize, usize, usize, FamilyMask, &[u8]), +) -> OKIResult { + 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, + mut on_family: impl FnMut(usize, FamilyMask, &[u8]), +) -> OKIResult { + 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> = 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 = 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::() < (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) +} diff --git a/src/obikphylo/src/siblings/extensions/sibling_ext.rs b/src/obikphylo/src/siblings/extensions/sibling_ext.rs index 9bf2a98a..bf335a99 100644 --- a/src/obikphylo/src/siblings/extensions/sibling_ext.rs +++ b/src/obikphylo/src/siblings/extensions/sibling_ext.rs @@ -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, + ) -> OKIResult; } 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, + ) -> OKIResult { + snp_pseudo_alignment(self, n, free_loss, no_ambiguity, excluded, entropy_bias) + } } diff --git a/src/obikphylo/src/siblings/mod.rs b/src/obikphylo/src/siblings/mod.rs index 8b37c931..3dcdb2c3 100644 --- a/src/obikphylo/src/siblings/mod.rs +++ b/src/obikphylo/src/siblings/mod.rs @@ -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";