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:
@@ -39,9 +39,12 @@ pub fn seq_id(ascii: &[u8]) -> String {
|
|||||||
format!("{:016X}", xxh64(ascii, 0))
|
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<()> {
|
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) })
|
writeln!(writer, "{}", unsafe { std::str::from_utf8_unchecked(seq) })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -59,3 +62,16 @@ pub fn write_record<W: Write>(
|
|||||||
writeln!(out)?;
|
writeln!(out)?;
|
||||||
write_sequence(out, seq)
|
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)
|
||||||
|
}
|
||||||
|
|||||||
@@ -14,7 +14,10 @@
|
|||||||
//! ```
|
//! ```
|
||||||
//!
|
//!
|
||||||
//! The lower-level primitive [`write_record`] and the [`JsonVal`] type are also
|
//! 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)]
|
#![deny(missing_docs)]
|
||||||
|
|
||||||
@@ -24,7 +27,7 @@ use std::io::{self, Write};
|
|||||||
|
|
||||||
use obikseq::{Minimizer, SuperKmer, Unitig};
|
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 ────────────────────────────────────────────────────────────────
|
// ── public API ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|||||||
@@ -36,15 +36,28 @@ impl From<MetricArg> for DistanceMetric {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric
|
/// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric
|
||||||
/// path (`--metric`/NJ/UPGMA) plus annex construction (`--sibling-annex`) —
|
/// path (`--metric`/NJ/UPGMA), annex construction (`--sibling-annex`),
|
||||||
/// everything else sibling-annex-based (`--snp`, `--sankoff`,
|
/// entropy reporting (`--shannon`) and SNP pseudo-alignment sampling
|
||||||
/// `--tnt`/`--phyg`/`--iqtree`, stats, ...) stays in `obikmer` until the rest
|
/// (`--pseudo-alignment`, `--subsample`, `--free-loss`, `--no-ambiguity`,
|
||||||
/// of `obikphylo::siblings` is reconnected (see the project memory on this).
|
/// `--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)]
|
#[derive(Args)]
|
||||||
pub struct PhyloArgs {
|
pub struct PhyloArgs {
|
||||||
/// Index directory
|
/// Index directory
|
||||||
pub index: PathBuf,
|
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
|
/// Build (or rebuild) the sibling-count/minorant annex — independent of
|
||||||
/// the distance metric below, meant to be run routinely, ahead of any
|
/// the distance metric below, meant to be run routinely, ahead of any
|
||||||
/// SNP-family distance computation that will later consume it.
|
/// 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
|
/// Write a per-family Shannon entropy report (CSV) — requires an
|
||||||
/// already-built sibling annex (`--sibling-annex` first, in this
|
/// already-built sibling annex (`--sibling-annex` first, in this
|
||||||
/// invocation or an earlier one). Full, unsampled scan of every family;
|
/// invocation or an earlier one). Always a full, unsampled scan of
|
||||||
/// no `--subsample`/`--entropy-bias` yet.
|
/// every family (`--subsample`/`--entropy`/`--entropy-sd` below only
|
||||||
|
/// apply to `--pseudo-alignment`, not this).
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub shannon: bool,
|
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
|
/// Distance metric to compute
|
||||||
#[arg(long, value_enum, default_value = "jaccard")]
|
#[arg(long, value_enum, default_value = "jaccard")]
|
||||||
pub metric: MetricArg,
|
pub metric: MetricArg,
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ use std::sync::Arc;
|
|||||||
|
|
||||||
use obikidxcache::index_cache::IndexCache;
|
use obikidxcache::index_cache::IndexCache;
|
||||||
use obikindex::KmerIndex;
|
use obikindex::KmerIndex;
|
||||||
use obikphylo::siblings::SiblingExt;
|
use obikphylo::siblings::{EntropyBias, SiblingExt};
|
||||||
use obikphylo::{Metrics, neighbor_joining, upgma};
|
use obikphylo::{Metrics, neighbor_joining, upgma};
|
||||||
use obisys::{Reporter, Stage};
|
use obisys::{Reporter, Stage};
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
@@ -30,6 +30,25 @@ pub fn run(args: PhyloArgs) {
|
|||||||
.collect();
|
.collect();
|
||||||
let n = labels.len();
|
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();
|
let mut rep = Reporter::new();
|
||||||
|
|
||||||
// Every partition/layer this needs is opened once, up front, and
|
// Every partition/layer this needs is opened once, up front, and
|
||||||
@@ -73,6 +92,48 @@ pub fn run(args: PhyloArgs) {
|
|||||||
info!("entropy report → {path}");
|
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);
|
info!("computing {:?} distances for {} genome(s)", args.metric, n);
|
||||||
|
|
||||||
let need_shared = args.shared_kmers || args.nj || args.upgma;
|
let need_shared = args.shared_kmers || args.nj || args.upgma;
|
||||||
@@ -85,14 +146,19 @@ pub fn run(args: PhyloArgs) {
|
|||||||
});
|
});
|
||||||
rep.push(t.stop());
|
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 ─────────────────────────────────────────────────
|
// ── Distance matrix → CSV ─────────────────────────────────────────────────
|
||||||
let write_dist_csv = |w: &mut dyn Write| {
|
let write_dist_csv = |w: &mut dyn Write| {
|
||||||
write!(w, "genome").unwrap();
|
write!(w, "genome").unwrap();
|
||||||
for g in &labels { write!(w, ",{g}").unwrap(); }
|
for &j in &kept { write!(w, ",{}", labels[j]).unwrap(); }
|
||||||
writeln!(w).unwrap();
|
writeln!(w).unwrap();
|
||||||
for (i, g) in labels.iter().enumerate() {
|
for &i in &kept {
|
||||||
write!(w, "{g}").unwrap();
|
write!(w, "{}", labels[i]).unwrap();
|
||||||
for j in 0..n {
|
for &j in &kept {
|
||||||
write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap();
|
write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap();
|
||||||
}
|
}
|
||||||
writeln!(w).unwrap();
|
writeln!(w).unwrap();
|
||||||
@@ -127,11 +193,11 @@ pub fn run(args: PhyloArgs) {
|
|||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
}));
|
}));
|
||||||
write!(f, "genome").unwrap();
|
write!(f, "genome").unwrap();
|
||||||
for g in &labels { write!(f, ",{g}").unwrap(); }
|
for &j in &kept { write!(f, ",{}", labels[j]).unwrap(); }
|
||||||
writeln!(f).unwrap();
|
writeln!(f).unwrap();
|
||||||
for (i, g) in labels.iter().enumerate() {
|
for &i in &kept {
|
||||||
write!(f, "{g}").unwrap();
|
write!(f, "{}", labels[i]).unwrap();
|
||||||
for j in 0..n { write!(f, ",{}", shared[[i, j]]).unwrap(); }
|
for &j in &kept { write!(f, ",{}", shared[[i, j]]).unwrap(); }
|
||||||
writeln!(f).unwrap();
|
writeln!(f).unwrap();
|
||||||
}
|
}
|
||||||
info!("shared-kmer matrix → {path}");
|
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.
|
//! this project, not a 4-symbol reduction.
|
||||||
|
|
||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
use obikidxcache::index_cache::IndexCache;
|
use obikidxcache::index_cache::IndexCache;
|
||||||
use obikindex::{OKIError, OKIResult};
|
use obikindex::{OKIError, OKIResult};
|
||||||
|
|
||||||
use super::family_scan::{Selection, scan_layer_families};
|
use super::family_scan::{Selection, scan_layer_families};
|
||||||
use super::minorant_selection::non_monomorphic_reindex_layer;
|
use super::minorant_selection::{is_non_monomorphic_minorant, non_monomorphic_reindex_layer};
|
||||||
use crate::siblings::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnexBuilder};
|
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
|
/// 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
|
/// 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(())
|
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::collections::HashMap;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
|
use obicompactvec::TempBitVecBuilder;
|
||||||
use obikindex::{OKIError, OKIResult};
|
use obikindex::{OKIError, OKIResult};
|
||||||
|
|
||||||
use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex};
|
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
|
mask.is_minorant() && mask.siblings() >= 1
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -47,3 +48,37 @@ pub(crate) fn non_monomorphic_reindex_layer(layer_dir: &Path) -> OKIResult<HashM
|
|||||||
}
|
}
|
||||||
Ok(reindex)
|
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`]
|
//! [`annex`] builds the sibling-count/minorant annex itself; [`family_scan`]
|
||||||
//! is the shared cross-partition traversal every other consumer
|
//! is the shared cross-partition traversal every other consumer
|
||||||
//! (`entropy`, and future ones — cardinality, alignment, stats) is built
|
//! (`entropy`, [`subsample`], [`alignment`], and future ones — cardinality,
|
||||||
//! on; [`minorant_selection`] and [`entropy`] are the first such consumer.
|
//! 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 annex;
|
||||||
mod entropy;
|
mod entropy;
|
||||||
mod family_scan;
|
mod family_scan;
|
||||||
|
mod masking;
|
||||||
mod minorant_selection;
|
mod minorant_selection;
|
||||||
|
mod subsample;
|
||||||
|
|
||||||
use obikidxcache::index_cache::IndexCache;
|
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 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 family_scan::{Selection, scan_layer_families};
|
||||||
|
pub(crate) use subsample::sample_index;
|
||||||
|
|
||||||
/// Whether every layer number in `cache` fits in a `FamilyMask` field
|
/// Whether every layer number in `cache` fits in a `FamilyMask` field
|
||||||
/// (`< 7`) — a single fact about the whole index, decided once here so
|
/// (`< 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 obisys::progress_bar;
|
||||||
|
|
||||||
use crate::siblings::algorithms::{
|
use crate::siblings::algorithms::{
|
||||||
build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode, scan_layer_families,
|
EntropyBias, Selection, SnpAlignment, build_layer_sibling_annex, family_entropy,
|
||||||
Selection,
|
family_entropy_4, is_fast_mode, scan_layer_families, snp_pseudo_alignment,
|
||||||
};
|
};
|
||||||
use crate::siblings::extensions::SiblingBuilder;
|
use crate::siblings::extensions::SiblingBuilder;
|
||||||
use crate::siblings::ENTROPY_ANNEX_FILE_NAME;
|
use crate::siblings::ENTROPY_ANNEX_FILE_NAME;
|
||||||
@@ -54,6 +54,28 @@ pub trait SiblingExt {
|
|||||||
/// Full, unsampled scan of every cached layer (`Selection::All`) — no
|
/// Full, unsampled scan of every cached layer (`Selection::All`) — no
|
||||||
/// `--subsample`/`--entropy-bias` yet (see the project memory on this).
|
/// `--subsample`/`--entropy-bias` yet (see the project memory on this).
|
||||||
fn shannon_entropy_csv(&self, path: &Path) -> OKIResult<()>;
|
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 {
|
impl SiblingExt for IndexCache {
|
||||||
@@ -167,4 +189,15 @@ impl SiblingExt for IndexCache {
|
|||||||
|
|
||||||
Ok(())
|
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 entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnex, EntropyAnnexBuilder};
|
||||||
pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
||||||
|
|
||||||
|
pub use algorithms::{EntropyBias, SnpAlignment};
|
||||||
pub use extensions::SiblingExt;
|
pub use extensions::SiblingExt;
|
||||||
|
|
||||||
pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib";
|
pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib";
|
||||||
|
|||||||
Reference in New Issue
Block a user