Extract phylogenetic sibling logic into new obikphylo crate
Relocate the `siblings` and `cardcomp` modules from `obikindex` to a dedicated `obikphylo` workspace member. Convert inherent methods on `KmerIndex` into extension traits, update import paths across `obikmer`, and add supporting accessor methods to `obikseq` and `obilayeredmap`. This restructuring reduces the public API surface of `obikindex` while organizing phylogenetic iteration, caching, and distance calculation logic under a dedicated crate.
This commit is contained in:
Eric Coissac
@@ -11,8 +11,6 @@ obiskio = { path = "../obiskio" }
|
||||
obisys = { path = "../obisys" }
|
||||
obicompactvec = { path = "../obicompactvec" }
|
||||
obilayeredmap = { path = "../obilayeredmap" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
obipipeline = { path = "../obipipeline" }
|
||||
ndarray = "0.16"
|
||||
rayon = "1"
|
||||
crossbeam-channel = "0.5"
|
||||
|
||||
@@ -1,292 +0,0 @@
|
||||
//! Cardinality × composition decomposition of the 16-state transition cost
|
||||
//! matrix — replaces `sankoff::build_cost_matrix`'s elementary-edge graph
|
||||
//! and its Floyd-Warshall shortest-path closure, which double-counts
|
||||
//! multi-hop transitions once IQ-TREE's own matrix exponential composes
|
||||
//! them again. See `docmd/theory/evolutionary_distances.md`, "`R` via
|
||||
//! `exp(-cost)` is wrong for a CTMC; cardinality/composition decomposition".
|
||||
//!
|
||||
//! Two small, directly-estimated first-order Markov models (diagonals
|
||||
//! included — "stay the same" is a real, calibrated outcome, not implicit):
|
||||
//! a 5-state cardinality model (`0..=4` members of a family) and a 4-state
|
||||
//! base-composition model (`A,C,G,T`, unambiguous/cardinality-1 loci only —
|
||||
//! composition bias is a substitution phenomenon, only meaningful when
|
||||
//! cardinality is conserved). Composed per pair of 16 states via a
|
||||
//! parsimony-style pairing of the non-shared elements (see
|
||||
//! `pairwise_transition_matrix`'s doc), then symmetrised — Sankoff
|
||||
//! parsimony's score is only independent of root placement (as required,
|
||||
//! since TNT/PhyG search *unrooted* trees) if the cost matrix is
|
||||
//! symmetric, the discrete-parsimony analogue of CTMC reversibility.
|
||||
|
||||
use crate::{BasePairTally, CardinalityTally};
|
||||
|
||||
/// Row-stochastic 5×5 cardinality transition probabilities (`0..=4`),
|
||||
/// diagonal included ("stay at the same cardinality"), from
|
||||
/// [`CardinalityTally`]'s pooled, symmetric co-occurrence counts.
|
||||
pub fn cardinality_transition_probs(tally: &CardinalityTally) -> [[f64; 5]; 5] {
|
||||
let mut p = [[0.0f64; 5]; 5];
|
||||
for a in 0..5 {
|
||||
let row_sum: u64 = tally.counts[a].iter().sum();
|
||||
if row_sum == 0 {
|
||||
continue;
|
||||
}
|
||||
for b in 0..5 {
|
||||
p[a][b] = tally.counts[a][b] as f64 / row_sum as f64;
|
||||
}
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
/// Row-stochastic 4×4 base-composition transition probabilities
|
||||
/// (`A,C,G,T`), diagonal included ("stay the same base"), from
|
||||
/// [`BasePairTally`]'s pooled substitution (off-diagonal) and agreement
|
||||
/// (`same`, diagonal) counts — unambiguous, cardinality-1 loci only.
|
||||
pub fn composition_transition_probs(tally: &BasePairTally) -> [[f64; 4]; 4] {
|
||||
let mut p = [[0.0f64; 4]; 4];
|
||||
for a in 0..4 {
|
||||
let row_sum = tally.same[a] + (0..4).map(|b| tally.counts[a][b]).sum::<u64>();
|
||||
if row_sum == 0 {
|
||||
continue;
|
||||
}
|
||||
p[a][a] = tally.same[a] as f64 / row_sum as f64;
|
||||
for b in 0..4 {
|
||||
if a != b {
|
||||
p[a][b] = tally.counts[a][b] as f64 / row_sum as f64;
|
||||
}
|
||||
}
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
/// Enumerate every injective mapping from `small` (indices `0..small.len()`)
|
||||
/// into `large` (indices `0..large.len()`, `large.len() >= small.len()`),
|
||||
/// calling `visit` with each full permutation of chosen `large` indices.
|
||||
/// `small.len() <= 4` always here (at most 4 bases can be lost or gained),
|
||||
/// so a plain recursive enumeration (at most `4! = 24` calls) is simplest —
|
||||
/// not worth a dependency for.
|
||||
fn for_each_injection(small_len: usize, large_len: usize, visit: &mut dyn FnMut(&[usize])) {
|
||||
let mut chosen = vec![0usize; small_len];
|
||||
let mut used = vec![false; large_len];
|
||||
fn go(pos: usize, chosen: &mut [usize], used: &mut [bool], visit: &mut dyn FnMut(&[usize])) {
|
||||
if pos == chosen.len() {
|
||||
visit(chosen);
|
||||
return;
|
||||
}
|
||||
for l in 0..used.len() {
|
||||
if used[l] {
|
||||
continue;
|
||||
}
|
||||
used[l] = true;
|
||||
chosen[pos] = l;
|
||||
go(pos + 1, chosen, used, visit);
|
||||
used[l] = false;
|
||||
}
|
||||
}
|
||||
if small_len == 0 {
|
||||
visit(&[]);
|
||||
return;
|
||||
}
|
||||
go(0, &mut chosen, &mut used, visit);
|
||||
}
|
||||
|
||||
/// Parsimony-style pairing cost for the non-shared elements of a
|
||||
/// transition: pair every element of the smaller of `lost`/`gained` with
|
||||
/// some element of the larger one, choosing the pairing that minimises
|
||||
/// total `-ln(P_composition)` (maximises likelihood) — the discrete
|
||||
/// analogue of "prefer a substitution over an independent loss+gain".
|
||||
/// Leftover elements of the larger side (`|large|-|small|` of them) are
|
||||
/// *not* charged here — they're pure cardinality change, already priced by
|
||||
/// `P_cardinality(|A|→|B|)` in the caller.
|
||||
fn best_pairing_cost(lost: &[u8], gained: &[u8], p_comp: &[[f64; 4]; 4]) -> f64 {
|
||||
let (small, large) = if lost.len() <= gained.len() { (lost, gained) } else { (gained, lost) };
|
||||
if small.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let mut best = f64::INFINITY;
|
||||
for_each_injection(small.len(), large.len(), &mut |perm| {
|
||||
let mut cost = 0.0;
|
||||
for (i, &li) in perm.iter().enumerate() {
|
||||
let (x, y) = (small[i], large[li]);
|
||||
let p = p_comp[x as usize][y as usize];
|
||||
cost += if p > 0.0 { -p.ln() } else { f64::INFINITY };
|
||||
}
|
||||
if cost < best {
|
||||
best = cost;
|
||||
}
|
||||
});
|
||||
best
|
||||
}
|
||||
|
||||
/// The full 16-state transition cost matrix, replacing
|
||||
/// `sankoff::build_cost_matrix`. For every pair of states `A`, `B`
|
||||
/// (bitmasks, bit `0..4` = `A,C,G,T`): `shared = A∩B` contributes
|
||||
/// `P_composition(x→x)` per shared base (diagonal — "stayed"); the
|
||||
/// remaining `lost = A\B`, `gained = B\A` are parsimony-paired into
|
||||
/// substitutions via [`best_pairing_cost`], and whatever's left over after
|
||||
/// pairing (only possible on one side, since one side is fully consumed)
|
||||
/// is *pure* cardinality change, priced once via `P_cardinality(|A|→|B|)`
|
||||
/// — never chained through intermediate states. Row-normalised (`P(A→·)`
|
||||
/// sums to `1` over all `B`), converted to a cost via `-ln`, then
|
||||
/// symmetrised (`(cost(A,B)+cost(B,A))/2` — the row-normalised `P` is not
|
||||
/// symmetric in general, but a Sankoff parsimony cost must be, so the
|
||||
/// score is independent of where an unrooted tree gets rooted).
|
||||
///
|
||||
/// `free_loss`: drop the `P_cardinality(|A|→|B|)` factor entirely (never
|
||||
/// added to `log_p`) — the same low/incomplete-coverage argument that
|
||||
/// justifies recoding whole-family non-detection as `?` (see
|
||||
/// `docmd/theory/evolutionary_distances.md`, "Locus dropout under
|
||||
/// incomplete coverage") applies one level down: whether a genome shows 1
|
||||
/// vs 2 (etc.) detected members of a *present* family is exactly as
|
||||
/// vulnerable to sampling failure as whether the family was detected at
|
||||
/// all. Without this, `∅`-involving transitions are neutralised (via the
|
||||
/// `?` recoding, bypassing this matrix's row/column 0 entirely) but
|
||||
/// cardinality changes *between two otherwise-detected, non-empty* states
|
||||
/// (e.g. `{A} -> {A,C}`) still carried the same calibrated
|
||||
/// `P_cardinality` penalty as any other gain/loss — inconsistent with
|
||||
/// `--free-loss`'s own rationale. With the factor dropped, cost is driven
|
||||
/// only by composition matching (shared-base retention and paired
|
||||
/// substitutions), never by a state pair's cardinality difference alone.
|
||||
pub fn pairwise_cost_matrix(p_card: &[[f64; 5]; 5], p_comp: &[[f64; 4]; 4], free_loss: bool) -> [[f64; 16]; 16] {
|
||||
let mut raw = [[0.0f64; 16]; 16];
|
||||
for a in 0u8..16 {
|
||||
for b in 0u8..16 {
|
||||
let shared = a & b;
|
||||
let lost: Vec<u8> = (0..4).filter(|&i| a & (1 << i) != 0 && b & (1 << i) == 0).collect();
|
||||
let gained: Vec<u8> = (0..4).filter(|&i| b & (1 << i) != 0 && a & (1 << i) == 0).collect();
|
||||
|
||||
let mut log_p = 0.0; // accumulate ln(P), so 0.0 = probability 1
|
||||
if !free_loss {
|
||||
let card_a = a.count_ones() as usize;
|
||||
let card_b = b.count_ones() as usize;
|
||||
let p_c = p_card[card_a][card_b];
|
||||
log_p += if p_c > 0.0 { p_c.ln() } else { f64::NEG_INFINITY };
|
||||
}
|
||||
|
||||
for i in 0..4u8 {
|
||||
if shared & (1 << i) != 0 {
|
||||
let p = p_comp[i as usize][i as usize];
|
||||
log_p += if p > 0.0 { p.ln() } else { f64::NEG_INFINITY };
|
||||
}
|
||||
}
|
||||
log_p -= best_pairing_cost(&lost, &gained, p_comp);
|
||||
|
||||
raw[a as usize][b as usize] = log_p.exp();
|
||||
}
|
||||
}
|
||||
|
||||
// Row-normalise to a proper transition probability matrix.
|
||||
let mut p = [[0.0f64; 16]; 16];
|
||||
for a in 0..16 {
|
||||
let row_sum: f64 = raw[a].iter().sum();
|
||||
if row_sum > 0.0 {
|
||||
for b in 0..16 {
|
||||
p[a][b] = raw[a][b] / row_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cost = -ln(P), then symmetrise (see doc comment: parsimony on an
|
||||
// unrooted tree requires a symmetric cost matrix).
|
||||
let mut cost = [[0.0f64; 16]; 16];
|
||||
for a in 0..16 {
|
||||
for b in 0..16 {
|
||||
cost[a][b] = if p[a][b] > 0.0 { -p[a][b].ln() } else { f64::INFINITY };
|
||||
}
|
||||
}
|
||||
let mut sym = [[0.0f64; 16]; 16];
|
||||
for a in 0..16 {
|
||||
for b in 0..16 {
|
||||
sym[a][b] = if a == b { 0.0 } else { (cost[a][b] + cost[b][a]) / 2.0 };
|
||||
}
|
||||
}
|
||||
sym
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cardinality_probs_row_normalised() {
|
||||
let mut counts = [[0u64; 5]; 5];
|
||||
counts[0][0] = 10;
|
||||
counts[0][1] = 5;
|
||||
counts[1][0] = 5;
|
||||
counts[1][1] = 2;
|
||||
let tally = CardinalityTally { counts };
|
||||
let p = cardinality_transition_probs(&tally);
|
||||
assert!((p[0].iter().sum::<f64>() - 1.0).abs() < 1e-12);
|
||||
assert!((p[1].iter().sum::<f64>() - 1.0).abs() < 1e-12);
|
||||
assert!((p[0][0] - 10.0 / 15.0).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn composition_probs_include_diagonal() {
|
||||
let mut tally = BasePairTally { counts: [[0u64; 4]; 4], same: [0u64; 4] };
|
||||
tally.same[0] = 90; // A stays A 90 times
|
||||
tally.counts[0][1] = 10; // A -> C 10 times
|
||||
tally.counts[1][0] = 10;
|
||||
let p = composition_transition_probs(&tally);
|
||||
assert!((p[0][0] - 0.9).abs() < 1e-12);
|
||||
assert!((p[0][1] - 0.1).abs() < 1e-12);
|
||||
assert!((p[0].iter().sum::<f64>() - 1.0).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pairwise_cost_matrix_is_symmetric_and_zero_diagonal() {
|
||||
// A plausible-looking pair of models (not calibrated from real
|
||||
// data — just needs every state reachable with nonzero probability
|
||||
// for this structural test).
|
||||
let p_card = [
|
||||
[0.5, 0.3, 0.15, 0.04, 0.01],
|
||||
[0.3, 0.4, 0.2, 0.08, 0.02],
|
||||
[0.15, 0.2, 0.4, 0.2, 0.05],
|
||||
[0.04, 0.08, 0.2, 0.5, 0.18],
|
||||
[0.01, 0.02, 0.05, 0.18, 0.74],
|
||||
];
|
||||
let p_comp = [
|
||||
[0.7, 0.1, 0.15, 0.05],
|
||||
[0.1, 0.7, 0.05, 0.15],
|
||||
[0.15, 0.05, 0.7, 0.1],
|
||||
[0.05, 0.15, 0.1, 0.7],
|
||||
];
|
||||
let cost = pairwise_cost_matrix(&p_card, &p_comp, false);
|
||||
for a in 0..16 {
|
||||
assert_eq!(cost[a][a], 0.0);
|
||||
for b in 0..16 {
|
||||
assert!((cost[a][b] - cost[b][a]).abs() < 1e-9, "cost[{a}][{b}]={} cost[{b}][{a}]={}", cost[a][b], cost[b][a]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn free_loss_ignores_cardinality_transition_probs() {
|
||||
// A skewed cardinality model (cardinality change made artificially
|
||||
// expensive) must have zero effect on the cost matrix once
|
||||
// `free_loss` is set — the whole point of the flag.
|
||||
let p_card_uniform = [[0.2f64; 5]; 5];
|
||||
let p_card_skewed = [
|
||||
[0.96, 0.01, 0.01, 0.01, 0.01],
|
||||
[0.01, 0.96, 0.01, 0.01, 0.01],
|
||||
[0.01, 0.01, 0.96, 0.01, 0.01],
|
||||
[0.01, 0.01, 0.01, 0.96, 0.01],
|
||||
[0.01, 0.01, 0.01, 0.01, 0.96],
|
||||
];
|
||||
let p_comp = [
|
||||
[0.7, 0.1, 0.15, 0.05],
|
||||
[0.1, 0.7, 0.05, 0.15],
|
||||
[0.15, 0.05, 0.7, 0.1],
|
||||
[0.05, 0.15, 0.1, 0.7],
|
||||
];
|
||||
let cost_uniform = pairwise_cost_matrix(&p_card_uniform, &p_comp, true);
|
||||
let cost_skewed = pairwise_cost_matrix(&p_card_skewed, &p_comp, true);
|
||||
for a in 0..16 {
|
||||
for b in 0..16 {
|
||||
assert!(
|
||||
(cost_uniform[a][b] - cost_skewed[a][b]).abs() < 1e-9,
|
||||
"cost[{a}][{b}] differs between cardinality models under free_loss: {} vs {}",
|
||||
cost_uniform[a][b], cost_skewed[a][b],
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -132,6 +132,10 @@ impl KmerIndex {
|
||||
IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty)
|
||||
}
|
||||
|
||||
/// The index's root directory — needed by out-of-crate extension code
|
||||
/// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto
|
||||
/// the same on-disk index.
|
||||
pub fn root_path(&self) -> &Path { &self.root_path }
|
||||
pub fn meta(&self) -> &IndexMeta { &self.meta }
|
||||
pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta }
|
||||
pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size }
|
||||
|
||||
@@ -2,7 +2,6 @@ pub mod error;
|
||||
pub mod meta;
|
||||
pub mod predicate;
|
||||
pub mod state;
|
||||
mod cardcomp;
|
||||
mod distance;
|
||||
mod dump;
|
||||
mod index;
|
||||
@@ -11,7 +10,6 @@ mod numa;
|
||||
mod rebuild;
|
||||
mod reindex;
|
||||
mod select;
|
||||
mod siblings;
|
||||
mod stats;
|
||||
|
||||
pub use error::{OKIError, OKIResult};
|
||||
@@ -22,5 +20,3 @@ pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME
|
||||
pub use predicate::{GroupFilterParams, MetaPred};
|
||||
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
|
||||
pub use stats::IndexBitsPerKmer;
|
||||
pub use siblings::{BasePairTally, CardinalityTally, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
||||
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
|
||||
|
||||
@@ -1,103 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::family_scan::scan_layer_families;
|
||||
|
||||
/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
|
||||
/// iff the genome carries the member whose own central base is `b`):
|
||||
/// 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).
|
||||
fn iupac_code(mask: u8) -> u8 {
|
||||
match mask & 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"),
|
||||
}
|
||||
}
|
||||
|
||||
/// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per
|
||||
/// genome, one column per variable family (`family_size() >= 2` — monomorphic
|
||||
/// families carry no signal and are skipped, unlike `raw_snp_distance`'s
|
||||
/// tally which does count them as `shared`). Column order is the same,
|
||||
/// deterministic sweep order as the annex build (partition, then layer, then
|
||||
/// slot) — arbitrary but stable and identical across genomes, which is all a
|
||||
/// pseudo-alignment needs (there is no natural genomic coordinate to sort by
|
||||
/// once flanks are dropped). See `docmd/theory/evolutionary_distances.md`,
|
||||
/// "Multi-genome framing: family as pseudo-alignment column".
|
||||
pub struct SnpAlignment {
|
||||
/// `sequences[g]` = genome `g`'s IUPAC-coded row, same length for every
|
||||
/// genome (`sequences.len()` columns).
|
||||
pub sequences: Vec<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Build the SNP-only pseudo-alignment from an already-built sibling
|
||||
/// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first).
|
||||
pub fn snp_pseudo_alignment(&self) -> OKIResult<SnpAlignment> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
&self.root_path,
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = self.sibling_layer_dirs()?;
|
||||
|
||||
let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
|
||||
// One layer at a time, not `par_iter()` over layers — same
|
||||
// rationale as `build_sibling_annex`: running many layers'
|
||||
// `scan_layer_families` concurrently would each group their own
|
||||
// lookups by partition internally, but interleave those sweeps
|
||||
// across layers at the OS level, scattering page-cache access over
|
||||
// every partition at once again and defeating the whole point of
|
||||
// the grouping. Columns appended straight into `sequences` as each
|
||||
// family comes back from `scan_layer_families` (bounded-batch, not
|
||||
// whole-layer) — no intermediate per-layer buffer, layer order
|
||||
// giving a single deterministic column order across the whole index.
|
||||
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); n_genomes];
|
||||
for layer_dir in &layer_dirs {
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
|
||||
if mask.family_size() < 2 {
|
||||
return; // monomorphic family — no signal, skip
|
||||
}
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
sequences[g].push(iupac_code(m));
|
||||
}
|
||||
})?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(SnpAlignment { sequences })
|
||||
}
|
||||
}
|
||||
@@ -1,277 +0,0 @@
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicU8, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obicompactvec::{FamilyMask, SiblingAnnexBuilder};
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obipipeline::ThrottleGuard;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::helpers::{central_base, is_minorant, partition_of};
|
||||
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||||
|
||||
// ── obipipeline data types ─────────────────────────────────────────────────
|
||||
|
||||
/// A batch of this layer's distinct k-mers (iteration-order index + k-mer),
|
||||
/// the pipeline's source item — batched, not one k-mer per item, so that
|
||||
/// pipeline messages and their synchronisation cost stay amortised over
|
||||
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on
|
||||
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
|
||||
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
|
||||
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
|
||||
///
|
||||
/// The `usize` is this k-mer's position in `iter_kmers()`'s enumeration
|
||||
/// order, **not** an MPHF slot — see `docmd/architecture/siblings.md`. It
|
||||
/// is the same index the annex file is written under.
|
||||
struct SourceBatch {
|
||||
items: Vec<(usize, CanonicalKmer)>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
/// One batch's worth of central-substitution variants (up to 3 per source
|
||||
/// k-mer), each already routed to its destination partition and carrying
|
||||
/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
|
||||
/// hit. `(dest_partition, variant, source_order, base)` per entry, where
|
||||
/// `source_order` is the source k-mer's iteration-order index (see
|
||||
/// `SourceBatch`), not an MPHF slot.
|
||||
struct VariantBatch {
|
||||
items: Vec<(usize, CanonicalKmer, usize, u8)>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
enum SibData {
|
||||
Batch(SourceBatch),
|
||||
Variants(VariantBatch),
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Build the sibling-count/minorant annex for every layer of every
|
||||
/// partition of this (already built) index, writing one annex file per
|
||||
/// layer alongside its existing index files. Safe to call again later
|
||||
/// (e.g. after a fresh `merge`) — each run simply overwrites the annex
|
||||
/// files of the index it is called on.
|
||||
///
|
||||
/// Construction only — no statistics gathered here on purpose: this is
|
||||
/// meant to run routinely (it is the artefact the SNP-family distances
|
||||
/// will consume), while the sibling-count distribution
|
||||
/// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate,
|
||||
/// occasional diagnostic pass over the result, not run every time.
|
||||
///
|
||||
/// Cross-partition/cross-layer lookups are required (a k-mer's siblings
|
||||
/// can live in any partition), but the layer loop itself — and thus the
|
||||
/// annex file this produces — stays local to one layer at a time.
|
||||
pub fn build_sibling_annex(&self) -> OKIResult<()> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
&self.root_path,
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
|
||||
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?);
|
||||
|
||||
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
|
||||
let mut total_slots: u64 = 0;
|
||||
for part in 0..n_parts {
|
||||
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
pb.inc(1);
|
||||
continue;
|
||||
}
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||||
|
||||
let mut part_slots: u64 = 0;
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||||
part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?;
|
||||
}
|
||||
total_slots += part_slots;
|
||||
pb.inc(1);
|
||||
pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the number of distinct k-mers (annex entries) processed, for
|
||||
/// progress reporting.
|
||||
fn build_layer_sibling_annex(
|
||||
&self,
|
||||
layer_dir: &Path,
|
||||
n_parts: usize,
|
||||
cache: &Arc<PartitionCache>,
|
||||
) -> OKIResult<u64> {
|
||||
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
|
||||
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
|
||||
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
|
||||
let k = self.kmer_size();
|
||||
let n = mphf.n();
|
||||
|
||||
// ── Reconciliation state, indexed by this layer's k-mer iteration
|
||||
// order (the physical layout of `unitigs.bin`), never by MPHF slot
|
||||
// — this layer's own k-mers are known members by construction, so
|
||||
// no evidence check, no MPHF slot, and no slot -> k-mer
|
||||
// reconstruction is needed or legitimate here (see
|
||||
// `docmd/architecture/siblings.md`: the MPHF is not invertible, and
|
||||
// evidence answers membership, not identity). The annex is written
|
||||
// under this same iteration order end to end, so a reader can later
|
||||
// zip `iter_kmers()` with the annex file directly, with no
|
||||
// MPHF/slot indirection at read time either.
|
||||
//
|
||||
// `Arc<Vec<AtomicU8>>`, not `Vec<AtomicU8>` — `Pipe::apply` requires
|
||||
// its source iterator to be `Send + 'static` (its items are
|
||||
// dispatched to worker threads that outlive this call), so the
|
||||
// batch-generating closure below needs an owned handle it can move
|
||||
// in, not a borrow of a local. `AtomicU8`, not `FamilyMask`,
|
||||
// because the gather phase below parallelises across destination
|
||||
// partitions (independent `query_partition_with` calls, safe to run
|
||||
// concurrently) and their `Found` hits can land on arbitrary,
|
||||
// possibly-shared source entries — a lock-free `fetch_or` avoids
|
||||
// needing any synchronisation beyond that. ─────────────────────────
|
||||
let mask: Arc<Vec<AtomicU8>> = Arc::new((0..n).map(|_| AtomicU8::new(0)).collect());
|
||||
|
||||
// ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
|
||||
// the actual cross-partition lookup reuses
|
||||
// `KmerPartition::query_partition_with` (the same batching mechanism
|
||||
// `obikmer query` already uses: open a partition's files once,
|
||||
// answer a whole batch of queries against it) instead of one lookup
|
||||
// per pipeline item. A per-item lookup (tried first) reopened/
|
||||
// re-mmap'd every target partition's files on every single variant
|
||||
// — fine at toy scale, but ~90% system time against a real index,
|
||||
// observed in practice. A *later* attempt still pushed one pipeline
|
||||
// message per generated variant (a `Flat` stage, `SourceItem` =>
|
||||
// `VariantQuery`, one k-mer in => up to 3 variants out as separate
|
||||
// messages) — cheaper than reopening files, but sampling a real run
|
||||
// showed most wall-clock time going into per-message channel
|
||||
// send/notify syscalls instead of the lookup itself: the pipeline's
|
||||
// whole point is amortising synchronisation over a batch, and a
|
||||
// single k-mer's ≤3 variants is far too fine a granularity for
|
||||
// that. Batching `BATCH_SIZE` source k-mers into one pipeline item
|
||||
// — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of
|
||||
// variants out, one message either way — keeps the per-message
|
||||
// synchronisation cost amortised over thousands of lookups instead
|
||||
// of one to three. ──────────────────────────────────────────────
|
||||
const BATCH_SIZE: usize = 4096;
|
||||
let n_workers = obisys::effective_parallelism();
|
||||
let capacity = 256;
|
||||
|
||||
// Throttling limits how many *batches* are in flight at once — the
|
||||
// permit is acquired per batch (not per k-mer) in the source
|
||||
// thread, and released once its `VariantBatch` has been read out of
|
||||
// the pipeline by the accumulation loop below. See
|
||||
// `obipipeline::throttle`'s docs for why this is required, not
|
||||
// optional, once a `Flat`-style stage sits in the pipeline.
|
||||
//
|
||||
// Batches stream straight from `unitigs.bin` via
|
||||
// `enumerate_kmers_batch` — never a full-layer `Vec` collect (a
|
||||
// layer can hold billions of k-mers; see the "no full collect"
|
||||
// rule). Each k-mer's own base is seeded into `mask` in this same
|
||||
// pass, since this is exactly the iteration order `mask` is keyed
|
||||
// on — no separate seeding pass needed.
|
||||
let seed_mask = Arc::clone(&mask);
|
||||
let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| {
|
||||
kmers
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, kmer)| {
|
||||
seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
|
||||
(start + i, kmer)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch {
|
||||
items: t.item,
|
||||
_permit: t.guard,
|
||||
});
|
||||
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
SibData : SourceBatch => VariantBatch,
|
||||
| {
|
||||
move |batch: SourceBatch| -> VariantBatch {
|
||||
let mut items = Vec::with_capacity(batch.items.len() * 3);
|
||||
for (order, kmer) in batch.items {
|
||||
for variant in kmer.central_canonical_neighbors() {
|
||||
if variant == kmer {
|
||||
continue;
|
||||
}
|
||||
items.push((
|
||||
partition_of(variant, n_parts),
|
||||
variant,
|
||||
order,
|
||||
central_base(variant, k),
|
||||
));
|
||||
}
|
||||
}
|
||||
VariantBatch { items, _permit: batch._permit }
|
||||
}
|
||||
} : Batch => Variants,
|
||||
};
|
||||
|
||||
// ── Group generated variants by destination partition. `cache`
|
||||
// holds every partition already mmap'd (no more `open()` cost), but
|
||||
// `mmap` pages are still loaded on demand and can be evicted — a
|
||||
// lookup is not free just because the file isn't reopened. Grouping
|
||||
// keeps one partition's pages hot while its whole batch is resolved,
|
||||
// instead of faulting pages in and out as lookups jump between
|
||||
// partitions in whatever order the pipeline happens to produce
|
||||
// them. Each batch's throttle permit drops here, once accumulated.
|
||||
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
|
||||
for vb in pipe.apply(throttled, n_workers, capacity) {
|
||||
for (dest_partition, variant, source_order, base) in vb.items {
|
||||
outgoing[dest_partition].push((variant, source_order, base));
|
||||
}
|
||||
}
|
||||
|
||||
// ── Resolve each partition's batch against the cache in one
|
||||
// contiguous pass; parallelised across partitions (independent,
|
||||
// read-only) so this keeps using multiple cores without giving up
|
||||
// the per-partition locality above. ─────────────────────────────
|
||||
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
||||
for &(variant, source_order, base) in queries {
|
||||
if cache.find(dest, variant) {
|
||||
mask[source_order].fetch_or(1 << base, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// ── Write the layer's annex file, indexed by iteration order — a
|
||||
// second streamed pass over `unitigs.bin` (via `enumerate_kmers`),
|
||||
// now that every entry's mask is final; never a `Vec` hold of the
|
||||
// whole layer. The minorant flag is computed here, not by every
|
||||
// later reader: this is the one place the whole family's *final*
|
||||
// mask and this entry's own k-mer (already in hand, no extra
|
||||
// lookup) are both available together. Every consumer that only
|
||||
// needs to know "is this k-mer the family's minorant" — the common
|
||||
// case, since a family is tallied once, at its minorant — reads
|
||||
// the bit straight back instead of re-deriving it (re-scanning
|
||||
// `unitigs.bin` and re-hashing through the MPHF, the cost
|
||||
// `is_minorant` was cheap to *compute* but expensive to *get the
|
||||
// inputs for* every time).
|
||||
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||||
let mut builder = SiblingAnnexBuilder::new(n, &annex_path)?;
|
||||
for (order, kmer) in mphf.enumerate_kmers() {
|
||||
let final_mask = FamilyMask::from_bits(mask[order].load(Ordering::Relaxed));
|
||||
let minorant = is_minorant(kmer, final_mask, k);
|
||||
builder.set(order, final_mask.with_minorant(minorant));
|
||||
}
|
||||
builder.close()?;
|
||||
|
||||
Ok(n as u64)
|
||||
}
|
||||
}
|
||||
@@ -1,166 +0,0 @@
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::OKIResult;
|
||||
|
||||
use super::{olm_to_ok, INDEX_SUBDIR};
|
||||
|
||||
/// Every partition's already-open MPHF layers, built **once** for the whole
|
||||
/// `build_sibling_annex` run and shared (read-only) across every lookup, in
|
||||
/// every source layer, for the rest of the run — not reopened/re-mmap'd per
|
||||
/// query, nor per source layer.
|
||||
///
|
||||
/// Confirmed necessary by sampling a real run: routing lookups through
|
||||
/// `KmerPartition::query_partition_with` (the same batching `obikmer query`
|
||||
/// uses) still reopens+re-mmaps every target partition's files on every
|
||||
/// call, and it is called once per destination partition **per source
|
||||
/// layer** — for an index with many layers this repeats the same
|
||||
/// `MphfLayer::open`/`Evidence::open`/`PersistentBitMatrix::open` work over
|
||||
/// and over. Parallelising those calls (see the gather step below) spread
|
||||
/// the redundant work across more cores but did not reduce it: sampling
|
||||
/// showed Rayon workers spending their time inside repeated `open()`
|
||||
/// syscalls, not computation. This cache amortises that cost to once per
|
||||
/// partition for the entire run, regardless of how many source layers or
|
||||
/// lookups follow.
|
||||
/// A cached layer's opened presence/count matrix, alongside its `MphfLayer`.
|
||||
pub(super) enum Mat {
|
||||
Count(PersistentCompactIntMatrix),
|
||||
Presence(PersistentBitMatrix),
|
||||
}
|
||||
|
||||
impl Mat {
|
||||
pub(super) fn n_cols(&self) -> usize {
|
||||
match self {
|
||||
Mat::Count(m) => m.n_cols(),
|
||||
Mat::Presence(m) => m.n_cols(),
|
||||
}
|
||||
}
|
||||
pub(super) fn carries(&self, g: usize, slot: usize) -> bool {
|
||||
match self {
|
||||
Mat::Count(m) => m.col_view(g).get(slot) != 0,
|
||||
Mat::Presence(m) => m.get(g, slot) != 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct PartitionCache {
|
||||
/// `layers[partition][layer]` = that partition's opened MPHF layers,
|
||||
/// paired 1:1 with `mats[partition][layer]`; empty if the partition
|
||||
/// directory doesn't exist. Used by both
|
||||
/// [`crate::index::KmerIndex::build_sibling_annex`] (`layers` only) and
|
||||
/// [`crate::index::KmerIndex::sibling_annex_stats`] (both).
|
||||
layers: Vec<Vec<MphfLayer>>,
|
||||
mats: Vec<Vec<Mat>>,
|
||||
}
|
||||
|
||||
impl PartitionCache {
|
||||
pub(super) fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
|
||||
let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
|
||||
let built: Vec<(Vec<MphfLayer>, Vec<Mat>)> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
.map(|part| -> OKIResult<(Vec<MphfLayer>, Vec<Mat>)> {
|
||||
let index_dir = partition.part_dir(part).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
pb.inc(1);
|
||||
return Ok((Vec::new(), Vec::new()));
|
||||
}
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||||
let mut layers = Vec::with_capacity(meta.n_layers);
|
||||
let mut mats = Vec::with_capacity(meta.n_layers);
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||||
let Ok(mphf) = MphfLayer::open(&layer_dir, &meta.mode) else { continue };
|
||||
let use_counts = with_counts && layer_dir.join("counts").exists();
|
||||
let mat = if use_counts {
|
||||
PersistentCompactIntMatrix::open(&layer_dir).ok().map(Mat::Count)
|
||||
} else {
|
||||
PersistentBitMatrix::open(&layer_dir).ok().map(Mat::Presence)
|
||||
};
|
||||
let Some(mat) = mat else { continue };
|
||||
layers.push(mphf);
|
||||
mats.push(mat);
|
||||
}
|
||||
pb.inc(1);
|
||||
Ok((layers, mats))
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
pb.finish_and_clear();
|
||||
let (layers, mats) = built.into_iter().unzip();
|
||||
Ok(Self { layers, mats })
|
||||
}
|
||||
|
||||
/// Existence-only lookup of `variant` in partition `dest_partition`:
|
||||
/// tries each of the partition's already-open layers in turn, stopping
|
||||
/// at the first hit.
|
||||
pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool {
|
||||
self.layers
|
||||
.get(dest_partition)
|
||||
.is_some_and(|layers| layers.iter().any(|mphf| mphf.find(variant).is_some()))
|
||||
}
|
||||
|
||||
/// Resolve many `(variant, family_idx, base)` queries against one
|
||||
/// partition's matrices at once, calling `on_hit(family_idx, base, g)`
|
||||
/// for every genome `g` that carries the resolved variant.
|
||||
///
|
||||
/// Genome-major, not query-major: `PersistentBitMatrix`/
|
||||
/// `PersistentCompactIntMatrix` are stored one contiguous block per
|
||||
/// genome (column), slot as the offset within it (see
|
||||
/// `obicompactvec::bitmatrix::packed::PackedBitMatrix` — each column's
|
||||
/// own `mmap` region). Resolving query-by-query (`for query { for genome
|
||||
/// { mat.carries(genome, slot) } }`, this function's predecessor) visits
|
||||
/// every genome's block once *per query* — for a batch of thousands of
|
||||
/// queries against ~90 genomes, that is thousands of jumps into each of
|
||||
/// ~90 widely separated multi-MB regions, in query order, not genome
|
||||
/// order: the access pattern a column-major layout is least suited to.
|
||||
/// Grouping first (by layer, since each layer's matrix is a separate
|
||||
/// column set) and sorting each group by slot, then visiting genome by
|
||||
/// genome, turns that into ~90 mostly-sequential sweeps through one
|
||||
/// column's own bytes — the layout's fast axis — confirmed by sampling
|
||||
/// a real run: `PersistentBitMatrix::get` dominated wall-clock time,
|
||||
/// mostly blocked on page faults, even after every partition/batch
|
||||
/// locality fix above this in the traversal.
|
||||
pub(super) fn find_presence_batch(
|
||||
&self,
|
||||
dest_partition: usize,
|
||||
queries: &[(CanonicalKmer, usize, u8)],
|
||||
n_genomes: usize,
|
||||
mut on_hit: impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let Some(layers) = self.layers.get(dest_partition) else { return };
|
||||
let Some(mats) = self.mats.get(dest_partition) else { return };
|
||||
|
||||
// First hit wins, same semantics as the old per-query loop (a
|
||||
// variant present in an earlier layer shadows later ones).
|
||||
let mut by_layer: Vec<Vec<(usize, usize, u8)>> = vec![Vec::new(); layers.len()];
|
||||
for &(variant, family_idx, base) in queries {
|
||||
for (li, mphf) in layers.iter().enumerate() {
|
||||
if let Some(slot) = mphf.find(variant) {
|
||||
by_layer[li].push((slot, family_idx, base));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (li, mut hits) in by_layer.into_iter().enumerate() {
|
||||
if hits.is_empty() {
|
||||
continue;
|
||||
}
|
||||
hits.sort_unstable_by_key(|&(slot, _, _)| slot);
|
||||
let mat = &mats[li];
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
for g in 0..n_cols {
|
||||
for &(slot, family_idx, base) in &hits {
|
||||
if mat.carries(g, slot) {
|
||||
on_hit(family_idx, base, g);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,120 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::distance::RawSnpDistanceOutput;
|
||||
use super::family_scan::scan_layer_families;
|
||||
|
||||
/// See [`KmerIndex::cardinality_tally`].
|
||||
pub struct CardinalityTally {
|
||||
/// `counts[a][b] == counts[b][a]` = number of family sites, pooled over
|
||||
/// included genome pairs, where one genome's family cardinality
|
||||
/// (popcount of its presence mask, `0..=4`) is `a` and the other's is
|
||||
/// `b`. Diagonal is real data here (both genomes at the same
|
||||
/// cardinality), unlike [`super::distance::BasePairTally::counts`].
|
||||
pub counts: [[u64; 5]; 5],
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Cardinality co-occurrence, pooled only over genome pairs whose
|
||||
/// overall SNP ratio in `raw` is at or below `ratio_ceiling` — same
|
||||
/// saturation/no-data exclusion discipline as
|
||||
/// [`base_pair_tally`](Self::base_pair_tally). Unlike
|
||||
/// `scan_family_pairs` (which resolves each genome to a single form and
|
||||
/// silently drops any genome carrying more than one member of the
|
||||
/// family), this needs the *full* per-genome presence mask — a family
|
||||
/// member count of 2, 3 or 4 is exactly the signal being tallied, not
|
||||
/// noise to discard — so it re-implements the traversal rather than
|
||||
/// reusing that helper.
|
||||
///
|
||||
/// Restricted to variable families (`family_size() >= 2`), matching
|
||||
/// `snp_pseudo_alignment`'s own scope — briefly removed, then
|
||||
/// reinstated: without it, the diagonal is dominated by genome-wide
|
||||
/// invariant background (family_size()<2 loci vastly outnumber the
|
||||
/// ones that ever vary anywhere), which is inconsistent with the
|
||||
/// `+ASC`-corrected alignment this matrix is ultimately used with —
|
||||
/// `+ASC` exists specifically because the likelihood only ever sees
|
||||
/// variable sites, so a rate model calibrated mostly from invariant
|
||||
/// background sites doesn't describe the population it's applied to.
|
||||
/// Verified empirically: removing the filter measurably worsened a
|
||||
/// real IQ-TREE run (log-likelihood dropped, `NNI search needs
|
||||
/// unusual large number of steps to converge` warnings appeared) — see
|
||||
/// `docmd/theory/evolutionary_distances.md` for the full account.
|
||||
/// [`base_pair_tally`](Self::base_pair_tally)'s own diagonal (`same`)
|
||||
/// gets the matching restriction via `scan_family_pairs`'s new
|
||||
/// `variable` flag, rather than a `family_size()` check of its own (it
|
||||
/// doesn't have direct access to the family's mask).
|
||||
pub fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
|
||||
if i == j {
|
||||
return false;
|
||||
}
|
||||
let snp = raw.snp[[i, j]];
|
||||
let total = snp + raw.shared[[i, j]];
|
||||
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
|
||||
});
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
&self.root_path,
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = self.sibling_layer_dirs()?;
|
||||
|
||||
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
|
||||
// One layer at a time — see `snp_pseudo_alignment`'s comment for why
|
||||
// `par_iter()` over layers would defeat `scan_layer_families`'s
|
||||
// partition-grouped locality. Tallied straight into `total` as each
|
||||
// family comes back — no per-layer buffer.
|
||||
let mut total = [[0u64; 5]; 5];
|
||||
for layer_dir in &layer_dirs {
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
|
||||
if mask.family_size() < 2 {
|
||||
// Fully invariant family (never varies anywhere in
|
||||
// the index) — genome-wide background, not
|
||||
// SNP-adjacent signal; would otherwise swamp the
|
||||
// diagonal (`c=1/c=1` etc.), which needs to reflect
|
||||
// the same variable-families-only population the
|
||||
// `+ASC`-corrected alignment/likelihood actually
|
||||
// models. See `base_pair_tally`'s `variable` gate
|
||||
// on its own `same` diagonal for the matching fix.
|
||||
return;
|
||||
}
|
||||
|
||||
for i in 0..n_genomes {
|
||||
let card_i = genome_mask[i].count_ones() as usize;
|
||||
for j in (i + 1)..n_genomes {
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
let card_j = genome_mask[j].count_ones() as usize;
|
||||
total[card_i][card_j] += 1;
|
||||
if card_i != card_j {
|
||||
total[card_j][card_i] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
})?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(CardinalityTally { counts: total })
|
||||
}
|
||||
}
|
||||
@@ -1,222 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::family_scan::scan_layer_families;
|
||||
|
||||
/// Raw p-distance restricted to loci that are single-copy in **both**
|
||||
/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility
|
||||
/// rule (`docmd/theory/evolutionary_distances.md`, "Locus eligibility"),
|
||||
/// without the JC/K2P/LogDet finalisation math: `snp[i,j]` + `shared[i,j]`
|
||||
/// is the number of eligible loci, `snp[i,j] / (snp[i,j] + shared[i,j])` is
|
||||
/// `p_hat`. A quick, self-contained way to sanity-check the estimator
|
||||
/// against a real index before the full `SnpTally` design is built.
|
||||
///
|
||||
/// A locus (family, tallied once at its minorant) is eligible for pair
|
||||
/// `(i, j)` iff genome `i` carries exactly one of the family's observed
|
||||
/// forms **and** genome `j` carries exactly one (possibly a different one)
|
||||
/// — presence-only: a genome carrying the same form twice (a same-allele
|
||||
/// duplicate) is indistinguishable from carrying it once when only a
|
||||
/// presence matrix is available, so such cases are not excluded here even
|
||||
/// when a count index exists. See "Locus eligibility", stringent rule, for
|
||||
/// why this matters and how a count index would close the gap — left as a
|
||||
/// follow-up, not applied here.
|
||||
pub struct RawSnpDistanceOutput {
|
||||
/// n×n count of eligible loci where the two genomes' single forms differ.
|
||||
pub snp: Array2<u64>,
|
||||
/// n×n count of eligible loci where the two genomes' single forms agree.
|
||||
pub shared: Array2<u64>,
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// Shared traversal behind [`raw_snp_distance`](Self::raw_snp_distance)
|
||||
/// and [`base_pair_tally`](Self::base_pair_tally): for every family
|
||||
/// (tallied once, at its minorant) of every layer of the already-built
|
||||
/// sibling annex, resolves each genome's single observed form (`None`
|
||||
/// if absent or ambiguous/multi-copy), then calls `on_pair(acc, i, j,
|
||||
/// bi, bj, variable)` for every genome pair `(i, j)` where both are
|
||||
/// unambiguous and single-copy (`bi == bj` means shared at that locus,
|
||||
/// `bi != bj` means a SNP). `variable` is the family's own
|
||||
/// `family_size() >= 2` (true if more than one member is observed
|
||||
/// *anywhere* in the family, i.e. it isn't fully invariant across the
|
||||
/// whole index) — `raw_snp_distance` ignores it (a fully-invariant
|
||||
/// family is still legitimately "shared"), but callers whose diagonal
|
||||
/// should only reflect genuine SNP-adjacent agreement, not the
|
||||
/// genome-wide invariant background, need it (see
|
||||
/// [`base_pair_tally`](Self::base_pair_tally)'s `same` field). Layers
|
||||
/// are processed one at a time, not in parallel — see
|
||||
/// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers
|
||||
/// would defeat `scan_layer_families`'s partition-grouped locality;
|
||||
/// each layer gets its own accumulator from `zero()`, combined
|
||||
/// pairwise via `combine`.
|
||||
fn scan_family_pairs<Acc, F, C>(
|
||||
&self,
|
||||
label: &str,
|
||||
zero: impl Fn() -> Acc + Sync,
|
||||
on_pair: F,
|
||||
combine: C,
|
||||
) -> OKIResult<Acc>
|
||||
where
|
||||
Acc: Send,
|
||||
F: Fn(&mut Acc, usize, usize, u8, u8, bool) + Sync,
|
||||
C: Fn(Acc, Acc) -> Acc,
|
||||
{
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
&self.root_path,
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = self.sibling_layer_dirs()?;
|
||||
|
||||
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
|
||||
// One layer at a time — see `snp_pseudo_alignment`'s comment for why
|
||||
// `par_iter()` over layers would defeat `scan_layer_families`'s
|
||||
// partition-grouped locality.
|
||||
let mut total = zero();
|
||||
for layer_dir in &layer_dirs {
|
||||
let mut acc = zero();
|
||||
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
|
||||
let variable = mask.family_size() >= 2;
|
||||
|
||||
// Per genome: which single form (if exactly one) it
|
||||
// carries — `None` (a `popcount != 1` mask) once a
|
||||
// second form is seen, ambiguous/not single-copy,
|
||||
// ineligible for either side of a pair.
|
||||
let single_form = |g: usize| -> Option<u8> {
|
||||
let m = genome_mask[g];
|
||||
(m.count_ones() == 1).then(|| m.trailing_zeros() as u8)
|
||||
};
|
||||
|
||||
for i in 0..n_genomes {
|
||||
let Some(bi) = single_form(i) else { continue };
|
||||
for j in (i + 1)..n_genomes {
|
||||
let Some(bj) = single_form(j) else { continue };
|
||||
on_pair(&mut acc, i, j, bi, bj, variable);
|
||||
}
|
||||
}
|
||||
})?;
|
||||
|
||||
total = combine(total, acc);
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
|
||||
/// (run [`build_sibling_annex`](Self::build_sibling_annex) first).
|
||||
pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let (snp, shared) = self.scan_family_pairs(
|
||||
"raw_snp_distance",
|
||||
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
|
||||
|(snp, shared), i, j, bi, bj, _variable| {
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
shared[[j, i]] += 1;
|
||||
} else {
|
||||
snp[[i, j]] += 1;
|
||||
snp[[j, i]] += 1;
|
||||
}
|
||||
},
|
||||
|(mut snp, mut shared), (s, sh)| {
|
||||
snp += &s;
|
||||
shared += &sh;
|
||||
(snp, shared)
|
||||
},
|
||||
)?;
|
||||
Ok(RawSnpDistanceOutput { snp, shared })
|
||||
}
|
||||
|
||||
/// Symmetric 6-category base-pair substitution tally (AC, AG, AT, CG,
|
||||
/// CT, GT — indexed `0=A,1=C,2=G,3=T`), pooled only over genome pairs
|
||||
/// whose overall SNP ratio in `raw` is at or below `ratio_ceiling` —
|
||||
/// same saturation-exclusion discipline as
|
||||
/// [`cardinality_tally`](Self::cardinality_tally), for the same reason:
|
||||
/// a saturated pair's observed base-pair mix trends toward neutral base
|
||||
/// composition, not the true point-mutation spectrum.
|
||||
///
|
||||
/// A second full pass over the annex, sharing
|
||||
/// [`raw_snp_distance`](Self::raw_snp_distance)'s traversal (guided by
|
||||
/// it, not a blind re-scan) — needed because `raw_snp_distance` only
|
||||
/// keeps aggregate SNP/shared counts per genome pair, not which bases
|
||||
/// were actually involved at each locus, and the ratio-ceiling filter
|
||||
/// can only be evaluated once the aggregate counts are known.
|
||||
pub fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
|
||||
if i == j {
|
||||
return false;
|
||||
}
|
||||
let snp = raw.snp[[i, j]];
|
||||
let total = snp + raw.shared[[i, j]];
|
||||
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
|
||||
});
|
||||
|
||||
let (counts, same) = self.scan_family_pairs(
|
||||
"base_pair_tally",
|
||||
|| ([[0u64; 4]; 4], [0u64; 4]),
|
||||
|(counts, same), i, j, bi, bj, variable| {
|
||||
if !included[[i, j]] {
|
||||
return;
|
||||
}
|
||||
if bi != bj {
|
||||
counts[bi as usize][bj as usize] += 1;
|
||||
counts[bj as usize][bi as usize] += 1;
|
||||
} else if variable {
|
||||
// Only count "stayed the same" from families that vary
|
||||
// *somewhere* in the index — a fully invariant family
|
||||
// (never varies anywhere) isn't a SNP-adjacent
|
||||
// agreement, it's genome-wide background, and would
|
||||
// otherwise swamp the diagonal (see
|
||||
// `docmd/theory/evolutionary_distances.md`, the
|
||||
// ascertainment-bias regression this was reverting).
|
||||
same[bi as usize] += 1;
|
||||
}
|
||||
},
|
||||
|(mut counts, mut same), (partial_counts, partial_same)| {
|
||||
for a in 0..4 {
|
||||
same[a] += partial_same[a];
|
||||
for b in 0..4 {
|
||||
counts[a][b] += partial_counts[a][b];
|
||||
}
|
||||
}
|
||||
(counts, same)
|
||||
},
|
||||
)?;
|
||||
Ok(BasePairTally { counts, same })
|
||||
}
|
||||
}
|
||||
|
||||
/// See [`KmerIndex::base_pair_tally`].
|
||||
pub struct BasePairTally {
|
||||
/// `counts[a][b] == counts[b][a]` = number of eligible loci, pooled
|
||||
/// over included genome pairs, where the two genomes' single forms are
|
||||
/// `a` and `b` (0=A, 1=C, 2=G, 3=T). Diagonal always `0` — an `a == b`
|
||||
/// locus is counted in `same`, not here.
|
||||
pub counts: [[u64; 4]; 4],
|
||||
/// `same[a]` = number of eligible loci, pooled over included genome
|
||||
/// pairs, where both genomes' single forms are `a` — the diagonal
|
||||
/// `counts` omits, needed to build a proper row-stochastic composition
|
||||
/// probability matrix (the "stay the same base" entries), not just the
|
||||
/// substitution-cost off-diagonal.
|
||||
pub same: [u64; 4],
|
||||
}
|
||||
@@ -1,311 +0,0 @@
|
||||
//! Shared per-layer family traversal, used by every sibling-annex consumer
|
||||
//! (`snp_pseudo_alignment`, `cardinality_tally`, `scan_family_pairs`,
|
||||
//! `sibling_annex_stats`) — resolves each minorant family's per-genome
|
||||
//! base-presence, with the same locality discipline as
|
||||
//! `build_sibling_annex` (see `cache.rs`/`build.rs`'s docs): cross-partition
|
||||
//! lookups are grouped by destination partition and resolved in one
|
||||
//! contiguous sweep per partition, instead of one lookup jumping between
|
||||
//! partitions in family order.
|
||||
//!
|
||||
//! Two concerns, kept on two different mechanisms because they need
|
||||
//! opposite things:
|
||||
//!
|
||||
//! - **Generating** each batch's cross-partition queries is pure CPU (bit
|
||||
//! tests, one hash per variant) — cheap, and safe to run for several
|
||||
//! batches at once. It runs on an `obipipeline::throttle` + `make_pipe!`
|
||||
//! pipeline, the same mechanism `build_layer_sibling_annex` uses, for the
|
||||
//! same reason: it overlaps this CPU work with the *previous* batch's
|
||||
//! resolution instead of leaving cores idle between batches.
|
||||
//! - **Resolving** those queries against `PartitionCache` is I/O (mmap page
|
||||
//! faults on a real index far larger than RAM) — one batch at a time,
|
||||
//! `rayon`-parallel *across partitions* (like `build_sibling_annex`'s own
|
||||
//! `outgoing.par_iter()`), never several batches concurrently. An earlier
|
||||
//! version resolved each batch fully inside its own pipeline worker, so
|
||||
//! `n_workers` batches were resolved concurrently — each one spreading its
|
||||
//! queries thin across every partition of the layer at once. Sampling a
|
||||
//! real run showed the fix hadn't helped at all: 16 threads "busy" in
|
||||
//! `find_presence`, but mostly blocked on page faults (~1.7 GB/s pagein),
|
||||
//! because 16 concurrent sweeps across the same partition space is exactly
|
||||
//! the scattering the grouping was meant to prevent — just spread over
|
||||
//! threads instead of over layers this time. Resolving one batch's worth
|
||||
//! at a time, with each of `rayon`'s threads owning one partition
|
||||
//! contiguously until that batch is done, keeps only one partition set
|
||||
//! "hot" at once, restoring the locality the batching is for.
|
||||
//!
|
||||
//! `obipipeline` does not guarantee output order, so generated batches carry
|
||||
//! their own starting family index and are replayed through a small reorder
|
||||
//! buffer before resolution — bounded by the throttle's own concurrency
|
||||
//! (`n_workers` batches), not by the layer's size.
|
||||
//!
|
||||
//! [`FAMILY_BATCH`] bounds a single batch's memory to `FAMILY_BATCH *
|
||||
//! genome_count` bytes, and is chosen large enough that a batch still gives
|
||||
//! each partition a decent number of queries to resolve in one go — too
|
||||
//! small a batch starves that density regardless of how the resolution step
|
||||
//! is parallelised. An earlier version materialised the whole layer's
|
||||
//! families up front before resolving anything: O(layer's family count ×
|
||||
//! genome count) memory, measured pushing a real run into heavy VM
|
||||
//! compression/swap.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU8, Ordering};
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obipipeline::{ThrottleGuard, throttle};
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::{Mat, PartitionCache};
|
||||
use super::helpers::{central_base, partition_of};
|
||||
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||||
|
||||
/// Families per batch — see the module docs for the memory-vs-per-partition-
|
||||
/// density trade-off this picks a point on. At ~90 genomes and a few
|
||||
/// hundred partitions, this keeps a batch's resolution memory in the low
|
||||
/// tens of MB while still giving each partition on the order of a thousand
|
||||
/// queries per batch to amortise against.
|
||||
const FAMILY_BATCH: usize = 65536;
|
||||
|
||||
impl KmerIndex {
|
||||
/// Every (partition, layer) directory carrying a sibling annex, checked
|
||||
/// up front so a missing one is reported before any real work starts.
|
||||
pub(super) fn sibling_layer_dirs(&self) -> OKIResult<Vec<PathBuf>> {
|
||||
let n_parts = self.n_partitions();
|
||||
let mut layer_dirs = Vec::new();
|
||||
for part in 0..n_parts {
|
||||
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
continue;
|
||||
}
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||||
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||||
if !annex_path.exists() {
|
||||
return Err(OKIError::InvalidInput(format!(
|
||||
"no sibling annex at {} — run build_sibling_annex first",
|
||||
annex_path.display()
|
||||
)));
|
||||
}
|
||||
layer_dirs.push(layer_dir);
|
||||
}
|
||||
}
|
||||
Ok(layer_dirs)
|
||||
}
|
||||
}
|
||||
|
||||
/// Read-only state shared (via `Arc`) across every pipeline worker
|
||||
/// generating this layer's batches — opened once, not per batch. No `cache`
|
||||
/// here: generation never touches the cross-partition cache, only this
|
||||
/// layer's own already-open matrix. No `annex` either — the filter step
|
||||
/// that builds `minorant_slots` already carries each slot's mask through,
|
||||
/// so a worker never needs to re-read it.
|
||||
struct LayerCtx {
|
||||
slot_kmer: Vec<Option<CanonicalKmer>>,
|
||||
mat: Mat,
|
||||
n_parts: usize,
|
||||
n_genomes: usize,
|
||||
n_cols: usize,
|
||||
k: usize,
|
||||
}
|
||||
|
||||
struct SourceBatch {
|
||||
start_family_idx: usize,
|
||||
/// `(slot, mask)` — the mask is carried through from the filter step
|
||||
/// below rather than re-read from the annex per batch; it's the same
|
||||
/// byte either way, just already in hand.
|
||||
slots: Vec<(usize, FamilyMask)>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
/// One batch's generated work: local presence already resolved (straight
|
||||
/// from this layer's own matrix, no lookup needed), cross-partition queries
|
||||
/// collected but not yet resolved against the cache.
|
||||
struct GeneratedBatch {
|
||||
start_family_idx: usize,
|
||||
masks: Vec<FamilyMask>,
|
||||
/// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`.
|
||||
genome_mask: Vec<u8>,
|
||||
/// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base)`.
|
||||
outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
enum FamData {
|
||||
Batch(SourceBatch),
|
||||
Generated(GeneratedBatch),
|
||||
}
|
||||
|
||||
/// Visits every minorant family of one layer, in slot order, batched
|
||||
/// [`FAMILY_BATCH`] at a time — see the module docs for why generation and
|
||||
/// resolution use different concurrency. `on_family` is called once per
|
||||
/// family, in slot order, with its own annex mask and its per-genome
|
||||
/// base-presence (`genome_mask[g]`: bit `b` set iff genome `g` carries the
|
||||
/// member whose own canonical central base is `b`), backed by a scratch
|
||||
/// buffer reused across every call — callers that need to keep data past
|
||||
/// the call must copy it themselves.
|
||||
pub(super) fn scan_layer_families(
|
||||
layer_dir: &Path,
|
||||
n_parts: usize,
|
||||
n_genomes: usize,
|
||||
with_counts: bool,
|
||||
k: usize,
|
||||
cache: &Arc<PartitionCache>,
|
||||
mut on_family: impl FnMut(FamilyMask, &[u8]),
|
||||
) -> OKIResult<()> {
|
||||
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
|
||||
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
|
||||
|
||||
// `MphfLayer::kmer_at` — direct slot -> k-mer reconstruction (evidence +
|
||||
// direct-access unitigs, no MPHF hashing, no file scan) instead of the
|
||||
// previous scan-`unitigs.bin`-and-hash-every-k-mer-forward approach,
|
||||
// which redundantly re-read the same data `mphf.find()` itself already
|
||||
// reads through `evidence`/`unitigs` to verify each hit.
|
||||
let slot_kmer: Vec<Option<CanonicalKmer>> = (0..annex.len()).map(|slot| mphf.kmer_at(slot)).collect();
|
||||
|
||||
let use_counts = with_counts && layer_dir.join("counts").exists();
|
||||
let mat = if use_counts {
|
||||
Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?)
|
||||
} else {
|
||||
Mat::Presence(PersistentBitMatrix::open(layer_dir)?)
|
||||
};
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
|
||||
// Minorant slots, in order, mask carried along — the flag is already
|
||||
// stored in the annex (set once, at construction — see
|
||||
// `build_layer_sibling_annex`), no need to re-derive it from the kmer
|
||||
// here, and no need for the pipeline workers below to re-read the same
|
||||
// byte from the annex a second time. `slot_kmer[s]` is still required
|
||||
// to exist (used later for `central_canonical_neighbors()`), but not
|
||||
// for this check.
|
||||
let minorant_slots: Vec<(usize, FamilyMask)> = (0..annex.len())
|
||||
.filter_map(|s| {
|
||||
let mask = annex.get(s)?;
|
||||
(slot_kmer[s].is_some() && mask.is_minorant()).then_some((s, mask))
|
||||
})
|
||||
.collect();
|
||||
let total_families = minorant_slots.len();
|
||||
if total_families == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let ctx = Arc::new(LayerCtx { slot_kmer, mat, n_parts, n_genomes, n_cols, k });
|
||||
|
||||
let mut offset = 0usize;
|
||||
let batches: Vec<(usize, Vec<(usize, FamilyMask)>)> = minorant_slots
|
||||
.chunks(FAMILY_BATCH)
|
||||
.map(|chunk| {
|
||||
let start = offset;
|
||||
offset += chunk.len();
|
||||
(start, chunk.to_vec())
|
||||
})
|
||||
.collect();
|
||||
|
||||
let n_workers = obisys::effective_parallelism();
|
||||
let capacity = 4;
|
||||
let throttled = throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
|
||||
start_family_idx: t.item.0,
|
||||
slots: t.item.1,
|
||||
_permit: t.guard,
|
||||
});
|
||||
|
||||
let worker_ctx = Arc::clone(&ctx);
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
FamData : SourceBatch => GeneratedBatch,
|
||||
| {
|
||||
move |batch: SourceBatch| -> GeneratedBatch {
|
||||
let ctx = &worker_ctx;
|
||||
let n = batch.slots.len();
|
||||
let mut masks = Vec::with_capacity(n);
|
||||
let mut bases = Vec::with_capacity(n);
|
||||
let mut genome_mask = vec![0u8; n * ctx.n_genomes];
|
||||
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..ctx.n_parts).map(|_| Vec::new()).collect();
|
||||
|
||||
// Pass 1: cheap, no matrix access — own base and this
|
||||
// batch's cross-partition queries. Mask already in hand
|
||||
// from the filter step, no second annex read.
|
||||
for (i, &(fslot, mask)) in batch.slots.iter().enumerate() {
|
||||
let kmer = ctx.slot_kmer[fslot].expect("minorant slot has a kmer");
|
||||
masks.push(mask);
|
||||
let base = central_base(kmer, ctx.k);
|
||||
bases.push(base);
|
||||
for other in kmer.central_canonical_neighbors() {
|
||||
if other == kmer {
|
||||
continue; // local — resolved below straight from `mat`, no lookup
|
||||
}
|
||||
let b = central_base(other, ctx.k);
|
||||
if !mask.has(b) {
|
||||
continue;
|
||||
}
|
||||
let dest = partition_of(other, ctx.n_parts);
|
||||
outgoing[dest].push((other, i, b));
|
||||
}
|
||||
}
|
||||
|
||||
// Pass 2: genome-major, not family-major — `mat` is stored
|
||||
// one contiguous block per genome (column), slot as the
|
||||
// offset within it (see `PartitionCache::find_presence_batch`'s
|
||||
// docs for the full rationale). `batch.slots` is already
|
||||
// sorted (a contiguous sub-range of the layer's sorted
|
||||
// minorant-slot list), so this sweeps each column roughly
|
||||
// in slot order instead of jumping between all `n_cols`
|
||||
// columns once per family.
|
||||
for g in 0..ctx.n_cols {
|
||||
for (i, &(fslot, _)) in batch.slots.iter().enumerate() {
|
||||
if ctx.mat.carries(g, fslot) {
|
||||
genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
GeneratedBatch { start_family_idx: batch.start_family_idx, masks, genome_mask, outgoing, _permit: batch._permit }
|
||||
}
|
||||
} : Batch => Generated,
|
||||
};
|
||||
|
||||
// Batches finish generation in whatever order their worker completes
|
||||
// them, not submission order — replay in order through a buffer bounded
|
||||
// by the throttle's own concurrency (`n_workers` batches can be in
|
||||
// flight at once, so at most that many can be waiting here).
|
||||
let mut pending: HashMap<usize, GeneratedBatch> = HashMap::new();
|
||||
let mut next_expected = 0usize;
|
||||
let mut scratch = vec![0u8; n_genomes];
|
||||
for generated in pipe.apply(throttled, n_workers, capacity) {
|
||||
pending.insert(generated.start_family_idx, generated);
|
||||
while let Some(batch) = pending.remove(&next_expected) {
|
||||
let n = batch.masks.len();
|
||||
|
||||
// ── Resolve this one batch's cross-partition queries — the
|
||||
// only place that touches `cache`. Parallel across partitions
|
||||
// (one batch at a time, never several concurrently — see the
|
||||
// module docs), each thread owning one partition's queries
|
||||
// contiguously until this batch is done.
|
||||
let genome_mask: Vec<AtomicU8> = batch.genome_mask.into_iter().map(AtomicU8::new).collect();
|
||||
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
||||
cache.find_presence_batch(dest, queries, n_genomes, |i, base, g| {
|
||||
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
||||
});
|
||||
});
|
||||
|
||||
for i in 0..n {
|
||||
for (g, dst) in scratch.iter_mut().enumerate() {
|
||||
*dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed);
|
||||
}
|
||||
on_family(batch.masks[i], &scratch);
|
||||
}
|
||||
next_expected += n;
|
||||
}
|
||||
}
|
||||
debug_assert_eq!(next_expected, total_families, "every batch must have been replayed");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
use obikseq::CanonicalKmer;
|
||||
use obiskbuilder::rolling_stat::RollingStat;
|
||||
|
||||
use obicompactvec::FamilyMask;
|
||||
|
||||
/// Central-position base of a canonical k-mer, in the fixed 0=A/1=C/2=G/3=T
|
||||
/// encoding — the mask's bit index. `k` must be odd (project invariant).
|
||||
#[inline]
|
||||
pub(super) fn central_base(kmer: CanonicalKmer, k: usize) -> u8 {
|
||||
kmer.nucleotide((k - 1) / 2)
|
||||
}
|
||||
|
||||
/// Is `kmer` the minorant of its family, given the family's presence mask?
|
||||
/// Regenerates the family's 4 canonical forms from `kmer` itself (cheap, no
|
||||
/// lookup — see the design doc's "Definitions" section for why this is
|
||||
/// always safe: the set of 4 forms is invariant regardless of which member
|
||||
/// you start from), and compares the raw encodings of whichever are marked
|
||||
/// present in `mask`.
|
||||
pub(super) fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool {
|
||||
kmer.central_canonical_neighbors().into_iter().all(|other| {
|
||||
other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw()
|
||||
})
|
||||
}
|
||||
|
||||
/// Minimiser of a single, isolated canonical k-mer (not part of a streamed
|
||||
/// sequence). `RollingStat` computes minimisers incrementally along a
|
||||
/// sequence; this feeds one k-mer's bases through a fresh instance to get
|
||||
/// the same selection for a single, disconnected k-mer. Not the leanest
|
||||
/// possible primitive (an O(1)-amortised dedicated scan, as originally
|
||||
/// sketched in the design doc's Step 0, would avoid the ASCII round-trip and
|
||||
/// `RollingStat` allocation) but correct and reuses already-tested logic;
|
||||
/// left as a follow-up optimisation.
|
||||
fn lone_kmer_minimizer(kmer: CanonicalKmer) -> obikseq::Minimizer {
|
||||
let ascii = kmer.to_ascii();
|
||||
let mut rs = RollingStat::new(0);
|
||||
for b in ascii {
|
||||
rs.push(b);
|
||||
}
|
||||
rs.canonical_minimizer()
|
||||
.expect("RollingStat must be ready after k bases of a valid k-mer")
|
||||
}
|
||||
|
||||
/// Destination partition for a (possibly synthetic) canonical k-mer, using
|
||||
/// the same routing rule as the rest of the index (`minimiser.seq_hash() &
|
||||
/// mask`, `n_partitions` is a power of two).
|
||||
pub(super) fn partition_of(kmer: CanonicalKmer, n_partitions: usize) -> usize {
|
||||
let mask = (n_partitions as u64) - 1;
|
||||
(lone_kmer_minimizer(kmer).seq_hash() & mask) as usize
|
||||
}
|
||||
@@ -1,78 +0,0 @@
|
||||
//! Phylo/sibling-domain iteration over a layer — an extension trait, not a
|
||||
//! new field on `MphfLayer`/`Layer<D>`: "family"/"minorant" are phylo
|
||||
//! concepts, `obilayeredmap` stays kmer/slot-mapping only (see
|
||||
//! `docmd/architecture/siblings.md`).
|
||||
//!
|
||||
//! The sibling annex is persisted in the same order as `iter_kmers()`
|
||||
//! (`build_sibling_annex`, see `build.rs`), so pairing them is a plain zip —
|
||||
//! no MPHF, no slot, no `kmer_at`. Both sides are already `Send + 'static`
|
||||
//! (`KmerIter` owns an `Arc<UnitigFileReader>` clone; `SiblingAnnex` is
|
||||
//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter`
|
||||
//! streams straight from disk and can be fed to `obipipeline` batch by
|
||||
//! batch — never collected whole into memory (see the project's "no full
|
||||
//! collect" rule).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use obicompactvec::{FamilyMask, SiblingAnnex};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::{KmerIter, MphfLayer};
|
||||
|
||||
/// One layer entry: a k-mer's position in the layer's iteration order (the
|
||||
/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer
|
||||
/// itself, and its family mask.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct SiblingEntry {
|
||||
pub order: usize,
|
||||
pub kmer: CanonicalKmer,
|
||||
pub mask: FamilyMask,
|
||||
}
|
||||
|
||||
/// Streams `(order, kmer, mask)` triples for one layer, in iteration order.
|
||||
/// Produced by [`SiblingLayerExt::iter_siblings`].
|
||||
pub struct SiblingIter {
|
||||
kmers: KmerIter,
|
||||
annex: Arc<SiblingAnnex>,
|
||||
order: usize,
|
||||
}
|
||||
|
||||
impl Iterator for SiblingIter {
|
||||
type Item = SiblingEntry;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
loop {
|
||||
let kmer = self.kmers.next()?;
|
||||
let order = self.order;
|
||||
self.order += 1;
|
||||
// `None` means "not yet computed" (see `SiblingAnnex` module
|
||||
// docs) — shouldn't happen against a fully-built annex, but
|
||||
// skip rather than misalign the two streams if it does.
|
||||
if let Some(mask) = self.annex.get(order) {
|
||||
return Some(SiblingEntry { order, kmer, mask });
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds phylo/sibling iteration to `MphfLayer`.
|
||||
pub trait SiblingLayerExt {
|
||||
/// Zip this layer's k-mers with their sibling-annex entry, in iteration
|
||||
/// order. `annex` must have been built from this same layer (its length
|
||||
/// must match the layer's k-mer count).
|
||||
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter;
|
||||
|
||||
/// Like [`iter_siblings`](Self::iter_siblings), filtered to the
|
||||
/// minorant of each family — the common case, since a family is
|
||||
/// tallied once, at its minorant.
|
||||
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> impl Iterator<Item = SiblingEntry>;
|
||||
}
|
||||
|
||||
impl SiblingLayerExt for MphfLayer {
|
||||
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter {
|
||||
SiblingIter { kmers: self.iter_kmers(), annex, order: 0 }
|
||||
}
|
||||
|
||||
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> impl Iterator<Item = SiblingEntry> {
|
||||
self.iter_siblings(annex).filter(|e| e.mask.is_minorant())
|
||||
}
|
||||
}
|
||||
@@ -1,76 +0,0 @@
|
||||
//! Family presence-mask annex construction.
|
||||
//!
|
||||
//! See `docmd/theory/evolutionary_distances.md`, "Definitions: family, and
|
||||
//! the canonical form of a family" and "Step 2b", for the full design
|
||||
//! discussion this implements.
|
||||
//!
|
||||
//! For each distinct k-mer of each layer of the (already built/merged)
|
||||
//! index, computes a 4-bit presence mask for its "family" (the up to 4
|
||||
//! k-mers sharing its flanks, differing only at the central base —
|
||||
//! well-defined for odd k): bit `b` set iff the family member whose own
|
||||
//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere
|
||||
//! in the current multi-genome index — a property of the whole index, not
|
||||
//! of any one genome. Sibling count and minorant are *derived* from the
|
||||
//! mask by callers, not stored (see `FamilyMask` and
|
||||
//! [`sibling_annex_stats`](crate::index::KmerIndex::sibling_annex_stats)
|
||||
//! below).
|
||||
//!
|
||||
//! Per layer, an `obipipeline` batch transform (throttled — see
|
||||
//! `obipipeline::throttle`) generates a whole batch's central variants at
|
||||
//! once (`BATCH_SIZE` source k-mers in, that batch's variants out as one
|
||||
//! pipeline message), interleaved across many in-flight batches by the
|
||||
//! scheduler's shared worker pool rather than processed on a single
|
||||
//! thread. The actual cross-partition lookup reuses a `PartitionCache` of
|
||||
//! every partition's already-open MPHF layers, built once for the whole
|
||||
//! `build_sibling_annex` run, rather than reopening files per lookup or
|
||||
//! per source layer. Two earlier, coarser-grained designs were tried and
|
||||
//! measured (not guessed) to be worse, in order: (1) reopening/re-mmap'ing
|
||||
//! every target partition's files on every single lookup — fine at toy
|
||||
//! scale, ~90% system time against a real index; (2) a `Flat` pipeline
|
||||
//! stage pushing one message per generated *variant* (up to 3 per source
|
||||
//! k-mer) — cheaper than reopening files, but sampling a real run showed
|
||||
//! most wall-clock time going into per-message channel send/notify
|
||||
//! syscalls rather than the lookup itself, because a single k-mer's ≤3
|
||||
//! variants is far too fine a granularity to amortise a pipeline's
|
||||
//! synchronisation cost over. See `docmd/theory/evolutionary_distances.md`,
|
||||
//! Step 2b, "Mechanism".
|
||||
//!
|
||||
//! Submodules, in the order data flows through them: [`cache`] (shared
|
||||
//! whole-run partition cache), [`helpers`] (small pure functions used
|
||||
//! throughout), [`build`] (annex construction), [`stats`] (family-size
|
||||
//! diagnostics), [`distance`] (raw SNP distance + base-pair tally),
|
||||
//! [`cardinality`] (cardinality co-occurrence), [`alignment`] (SNP-only
|
||||
//! pseudo-alignment).
|
||||
|
||||
mod alignment;
|
||||
mod build;
|
||||
mod cache;
|
||||
mod cardinality;
|
||||
mod distance;
|
||||
mod family_scan;
|
||||
mod helpers;
|
||||
mod iter;
|
||||
mod stats;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub use alignment::SnpAlignment;
|
||||
pub use cardinality::CardinalityTally;
|
||||
pub use distance::{BasePairTally, RawSnpDistanceOutput};
|
||||
pub use iter::{SiblingEntry, SiblingIter, SiblingLayerExt};
|
||||
pub use stats::SiblingAnnexStats;
|
||||
|
||||
use obilayeredmap::OLMError;
|
||||
|
||||
use crate::error::OKIError;
|
||||
|
||||
pub(super) const INDEX_SUBDIR: &str = "index";
|
||||
pub(super) const ANNEX_FILE_NAME: &str = "siblings.psib";
|
||||
|
||||
pub(super) fn olm_to_ok(e: OLMError) -> OKIError {
|
||||
match e {
|
||||
OLMError::Io(e) => OKIError::Io(e),
|
||||
other => OKIError::InvalidInput(format!("layered-map error: {other}")),
|
||||
}
|
||||
}
|
||||
@@ -1,126 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use obicompactvec::SiblingAnnex;
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::ANNEX_FILE_NAME;
|
||||
use super::cache::PartitionCache;
|
||||
use super::family_scan::scan_layer_families;
|
||||
|
||||
/// Distribution of family sizes (1-4), read back from an already-built
|
||||
/// annex (see [`KmerIndex::build_sibling_annex`]) plus the index's
|
||||
/// presence/count data — a separate, occasional diagnostic pass, not fused
|
||||
/// into construction.
|
||||
///
|
||||
/// Every count here is **per family, not per slot**: a family with `F`
|
||||
/// members occupies `F` annex slots (one per observed member), all sharing
|
||||
/// the same mask. Counting every slot would count each family up to 4
|
||||
/// times over; only the minorant's slot is tallied (minorant is derived on
|
||||
/// the fly — see `is_minorant` — not stored, but cheap: no lookup, pure
|
||||
/// bit arithmetic on already-in-hand data).
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct SiblingAnnexStats {
|
||||
/// `counts[s]` = number of *families* of size `s + 1` (index 0 = size 1,
|
||||
/// i.e. 0 siblings, ... index 3 = size 4, i.e. 3 siblings).
|
||||
pub counts: [u64; 4],
|
||||
/// `per_genome[g][s]` = number of families of size `s + 1` for which
|
||||
/// genome `g` (index into `KmerIndex::meta().genomes`) carries at least
|
||||
/// one member.
|
||||
pub per_genome: Vec<[u64; 4]>,
|
||||
}
|
||||
|
||||
impl KmerIndex {
|
||||
/// The global family-size histogram alone (`SiblingAnnexStats::counts`,
|
||||
/// no `per_genome`) — reads only the already-built annex (`mask.siblings()`
|
||||
/// + `mask.is_minorant()`, 1 byte/slot, mmap'd), nothing else: no
|
||||
/// `unitigs.bin` scan, no MPHF lookup, no `PartitionCache`. Cost is a
|
||||
/// linear scan of one small file per layer, independent of the rest of
|
||||
/// the index's size or how much of it is paged in. An earlier version
|
||||
/// re-derived minorant-ness per slot (re-reading `unitigs.bin`, hashing
|
||||
/// every k-mer through the MPHF again to reconstruct what
|
||||
/// `build_sibling_annex` already knew) — sampling a real run showed
|
||||
/// `MphfLayer::find` alone at 71% of wall-clock time for what should be
|
||||
/// a near-instant four-bucket count. `mask.is_minorant()` is that same
|
||||
/// fact, computed once at construction (see `build_layer_sibling_annex`)
|
||||
/// and stored in the annex's own spare bits — free to read back.
|
||||
/// [`sibling_annex_stats`](Self::sibling_annex_stats) computes the same
|
||||
/// `counts` but pays the full per-genome cross-partition resolution
|
||||
/// cost to get there too; use this when only the global histogram is
|
||||
/// needed. Requires an annex built after the minorant flag was added —
|
||||
/// re-run `build_sibling_annex` if this reads as all-zero on an older one.
|
||||
pub fn sibling_family_size_histogram(&self) -> OKIResult<[u64; 4]> {
|
||||
let layer_dirs = self.sibling_layer_dirs()?;
|
||||
|
||||
let mut counts = [0u64; 4];
|
||||
for layer_dir in &layer_dirs {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
for slot in 0..annex.len() {
|
||||
let Some(mask) = annex.get(slot) else { continue };
|
||||
if !mask.is_minorant() {
|
||||
continue; // family tallied once, at its minorant
|
||||
}
|
||||
counts[mask.siblings() as usize] += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(counts)
|
||||
}
|
||||
|
||||
/// Tally the family-size distribution of an already-built annex
|
||||
/// (globally, and per genome), counting each family once (at its
|
||||
/// minorant slot). Errors if [`build_sibling_annex`] has not been run on
|
||||
/// this index first.
|
||||
///
|
||||
/// [`build_sibling_annex`]: Self::build_sibling_annex
|
||||
pub fn sibling_annex_stats(&self) -> OKIResult<SiblingAnnexStats> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta.genomes.len();
|
||||
let with_counts = self.meta.config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
// Same whole-run cache as `build_sibling_annex` — see its docs for
|
||||
// why re-opening per lookup (or per call to a batching helper) is
|
||||
// not good enough on a real index.
|
||||
let partition = KmerPartition::open_with_config(
|
||||
&self.root_path,
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = self.sibling_layer_dirs()?;
|
||||
|
||||
// One layer at a time, not parallelised across layers — see
|
||||
// `snp_pseudo_alignment`'s comment for why `par_iter()` over layers
|
||||
// would defeat `scan_layer_families`'s partition-grouped locality.
|
||||
let pb = progress_bar("sibling_annex_stats", layer_dirs.len() as u64, "layers");
|
||||
let mut stats = SiblingAnnexStats {
|
||||
per_genome: vec![[0u64; 4]; n_genomes],
|
||||
..Default::default()
|
||||
};
|
||||
for layer_dir in &layer_dirs {
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
|
||||
// "Genome g represents this family" means g carries
|
||||
// *any* of its members, not just the minorant's own —
|
||||
// `genome_mask[g] != 0` is exactly that.
|
||||
let s = mask.siblings() as usize;
|
||||
stats.counts[s] += 1;
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
if m != 0 {
|
||||
stats.per_genome[g][s] += 1;
|
||||
}
|
||||
}
|
||||
})?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
@@ -1,257 +0,0 @@
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{FamilyMask, SiblingAnnex};
|
||||
use obikseq::{CanonicalKmer, Kmer, Sequence};
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obisys::Reporter;
|
||||
use tempfile::tempdir;
|
||||
|
||||
use crate::index::KmerIndex;
|
||||
use crate::meta::{GenomeInfo, IndexConfig};
|
||||
use crate::merge::MergeMode;
|
||||
|
||||
use super::helpers::is_minorant;
|
||||
use super::{ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||||
|
||||
// k must be >= 11 (project constraint, "k ∈ [11,31]"); k=11, level_max=1,
|
||||
// theta=0.0 mirror `obiskbuilder`'s own tests (smaller k/level_max
|
||||
// combinations trip an unrelated pre-existing bug in `obikentropy`'s
|
||||
// sliding-window ring buffer — not this feature's concern).
|
||||
const K: usize = 11;
|
||||
const M: usize = 5;
|
||||
|
||||
/// Build a single-genome index from one in-memory FASTA sequence, driving
|
||||
/// the same primitives `obikmer`'s `scatter` step uses (minus the
|
||||
/// multi-file `obipipeline` wrapper — a single sequence needs none of
|
||||
/// that): normalise -> build superkmers -> route -> write.
|
||||
/// `cargo test` doesn't install a `tracing` subscriber the way `obikmer`'s
|
||||
/// CLI does, so `debug!`/etc. are silent no-ops by default — including the
|
||||
/// `PartitionRunner` instrumentation that would matter most for
|
||||
/// re-diagnosing a hang here. `try_init` is idempotent across concurrently
|
||||
/// running tests (later calls just find a subscriber already installed).
|
||||
fn init_tracing() {
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||||
)
|
||||
.with_writer(std::io::stderr)
|
||||
.try_init();
|
||||
}
|
||||
|
||||
fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
|
||||
init_tracing();
|
||||
let fasta_path = dir.join(format!("{label}.fasta"));
|
||||
let mut f = std::fs::File::create(&fasta_path).unwrap();
|
||||
writeln!(f, ">{label}").unwrap();
|
||||
f.write_all(seq).unwrap();
|
||||
writeln!(f).unwrap();
|
||||
drop(f);
|
||||
|
||||
let index_path = dir.join(format!("{label}.idx"));
|
||||
let config = IndexConfig {
|
||||
kmer_size: K,
|
||||
minimizer_size: M,
|
||||
n_bits: 0, // 1 partition — keeps the test deterministic and simple
|
||||
with_counts: false,
|
||||
evidence: obilayeredmap::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false)
|
||||
.expect("create");
|
||||
|
||||
let mut rep = Reporter::new();
|
||||
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
|
||||
for page in stream {
|
||||
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
|
||||
idx.partition_mut().write_batch(batch).expect("write_batch");
|
||||
}
|
||||
idx.partition_mut().close().expect("close partition writers");
|
||||
idx.mark_scattered().expect("mark_scattered");
|
||||
idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
|
||||
idx.build_layers(1, None, false, &mut rep).expect("build_layers");
|
||||
idx
|
||||
}
|
||||
|
||||
fn canonical(ascii: &[u8]) -> CanonicalKmer {
|
||||
Kmer::from_ascii(ascii).unwrap().canonical()
|
||||
}
|
||||
|
||||
/// Read back the annex entry for a given canonical k-mer from the merged
|
||||
/// index's (single) partition/layer, asserting it was found at all.
|
||||
fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask {
|
||||
let index_dir = idx.partition().part_dir(0).join(INDEX_SUBDIR);
|
||||
let meta = PartitionMeta::load(&index_dir).unwrap();
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||||
let mphf = MphfLayer::open(&layer_dir, &meta.mode).unwrap();
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).unwrap();
|
||||
return annex.get(slot).expect("slot must have a computed annex entry");
|
||||
}
|
||||
}
|
||||
panic!("kmer not found in any layer of partition 0");
|
||||
}
|
||||
|
||||
fn merge_two(dir: &Path, g1: &KmerIndex, g2: &KmerIndex) -> KmerIndex {
|
||||
let mut rep = Reporter::new();
|
||||
KmerIndex::merge(
|
||||
&dir.join("merged.idx"),
|
||||
&[g1, g2],
|
||||
MergeMode::Presence,
|
||||
false,
|
||||
false,
|
||||
1.0,
|
||||
&mut rep,
|
||||
)
|
||||
.expect("merge")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_annex_one_sibling_each() {
|
||||
// k=11, centre = index 5 (0-based). Two genomes, each exactly one
|
||||
// k-mer, sharing every base except the centre:
|
||||
// g1 = "AACCGCTTAAG" (centre 'C', base index 1)
|
||||
// g2 = "AACCGGTTAAG" (centre 'G', base index 2)
|
||||
// Hand-verified: both stay forward-oriented under canonicalisation
|
||||
// (each is lexicographically smaller than its own reverse
|
||||
// complement, since both start with "AA"), and raw(g1) < raw(g2)
|
||||
// (only differing base: C=0b01 < G=0b10 at the centre) — so g1 is
|
||||
// the minorant, g2 is not. The mask is a family-wide value: both
|
||||
// slots must read back the *same* presence bits (1 and 2 set) — but
|
||||
// only g1's slot should carry the stored minorant flag.
|
||||
let dir = tempdir().unwrap();
|
||||
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
|
||||
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
|
||||
let merged = merge_two(dir.path(), &g1, &g2);
|
||||
merged.build_sibling_annex().expect("build_sibling_annex");
|
||||
|
||||
let g1_kmer = canonical(b"AACCGCTTAAG");
|
||||
let g2_kmer = canonical(b"AACCGGTTAAG");
|
||||
let expected_mask = FamilyMask::EMPTY.with(1).with(2);
|
||||
|
||||
let a = annex_info_for(&merged, g1_kmer);
|
||||
assert_eq!(a.bits(), expected_mask.bits(), "AACCGCTTAAG");
|
||||
assert_eq!(a.siblings(), 1);
|
||||
assert!(is_minorant(g1_kmer, a, K), "g1 should be the minorant");
|
||||
assert!(a.is_minorant(), "g1's stored minorant flag should be set at build time");
|
||||
|
||||
let b = annex_info_for(&merged, g2_kmer);
|
||||
assert_eq!(b.bits(), expected_mask.bits(), "AACCGGTTAAG");
|
||||
assert_eq!(b.siblings(), 1);
|
||||
assert!(!is_minorant(g2_kmer, b, K), "g2 should not be the minorant");
|
||||
assert!(!b.is_minorant(), "g2's stored minorant flag should not be set");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_annex_zero_siblings_when_identical_and_no_variant_exists() {
|
||||
// Same k-mer in both genomes, no other genome around to carry a
|
||||
// variant -> 0 siblings, trivially its own minorant.
|
||||
let dir = tempdir().unwrap();
|
||||
let g1 = build_single_genome_index(dir.path(), "g1", b"GATTACAGATC");
|
||||
let g2 = build_single_genome_index(dir.path(), "g2", b"GATTACAGATC");
|
||||
let merged = merge_two(dir.path(), &g1, &g2);
|
||||
merged.build_sibling_annex().expect("build_sibling_annex");
|
||||
|
||||
let kmer = canonical(b"GATTACAGATC");
|
||||
let mask = annex_info_for(&merged, kmer);
|
||||
assert_eq!(mask.siblings(), 0, "GATTACAGATC");
|
||||
assert_eq!(mask.family_size(), 1);
|
||||
assert!(is_minorant(kmer, mask, K));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_annex_stats_counts_each_family_once_and_per_genome() {
|
||||
// Reuses the one-sibling-each fixture: a single family of size 2
|
||||
// (g1's centre-C form + g2's centre-G form), each genome carrying
|
||||
// exactly one of the two members. Stats must report exactly one
|
||||
// family of size 2 (`counts[1] == 1`, since index 1 = size 2), not
|
||||
// two (which naively summing both slots would give), and both
|
||||
// genomes represented at size 2, neither at any other size.
|
||||
let dir = tempdir().unwrap();
|
||||
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
|
||||
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
|
||||
let merged = merge_two(dir.path(), &g1, &g2);
|
||||
merged.build_sibling_annex().expect("build_sibling_annex");
|
||||
|
||||
let stats = merged.sibling_annex_stats().expect("sibling_annex_stats");
|
||||
|
||||
assert_eq!(stats.counts, [0, 1, 0, 0], "one family of size 2, counted once");
|
||||
assert_eq!(stats.per_genome.len(), 2);
|
||||
for g in 0..2 {
|
||||
assert_eq!(
|
||||
stats.per_genome[g], [0, 1, 0, 0],
|
||||
"genome {g} should represent exactly one size-2 family"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_family_size_histogram_matches_full_stats_global_counts() {
|
||||
// Same fixture as `sibling_annex_stats_counts_each_family_once_and_per_genome`
|
||||
// — the cheap, annex-only histogram must agree with the `counts` half of
|
||||
// the full (cross-partition) stats pass, without needing a `PartitionCache`
|
||||
// at all.
|
||||
let dir = tempdir().unwrap();
|
||||
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
|
||||
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
|
||||
let merged = merge_two(dir.path(), &g1, &g2);
|
||||
merged.build_sibling_annex().expect("build_sibling_annex");
|
||||
|
||||
let histogram = merged.sibling_family_size_histogram().expect("sibling_family_size_histogram");
|
||||
assert_eq!(histogram, [0, 1, 0, 0], "one family of size 2, counted once");
|
||||
|
||||
let stats = merged.sibling_annex_stats().expect("sibling_annex_stats");
|
||||
assert_eq!(histogram, stats.counts, "must agree with the full stats pass's global counts");
|
||||
}
|
||||
|
||||
/// Exercises the four sibling-annex consumers that were rewritten to share
|
||||
/// `family_scan::scan_layer_families` (partition-grouped lookups instead of
|
||||
/// one lookup per family) — same one-sibling-each fixture as the tests
|
||||
/// above (g1 carries the family's centre-C member, g2 the centre-G one),
|
||||
/// hand-verified expected output for each.
|
||||
#[test]
|
||||
fn family_scan_consumers_agree_on_one_sibling_each() {
|
||||
let dir = tempdir().unwrap();
|
||||
let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
|
||||
let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
|
||||
let merged = merge_two(dir.path(), &g1, &g2);
|
||||
merged.build_sibling_annex().expect("build_sibling_annex");
|
||||
|
||||
// Merge doesn't promise to preserve source order, so resolve each
|
||||
// genome's index by label rather than assuming g1 -> 0, g2 -> 1.
|
||||
let idx_of = |label: &str| merged.meta.genomes.iter().position(|g| g.label == label).unwrap();
|
||||
let (i1, i2) = (idx_of("g1"), idx_of("g2"));
|
||||
|
||||
// snp_pseudo_alignment: one variable family, one column — g1's row
|
||||
// reads 'C' (its own member), g2's reads 'G'.
|
||||
let alignment = merged.snp_pseudo_alignment().expect("snp_pseudo_alignment");
|
||||
assert_eq!(alignment.sequences[i1], vec![b'C']);
|
||||
assert_eq!(alignment.sequences[i2], vec![b'G']);
|
||||
|
||||
// raw_snp_distance: g1's single form (C) != g2's (G) at the family's
|
||||
// one eligible locus -> a SNP, not a shared site.
|
||||
let raw = merged.raw_snp_distance().expect("raw_snp_distance");
|
||||
assert_eq!(raw.snp[[i1, i2]], 1);
|
||||
assert_eq!(raw.snp[[i2, i1]], 1);
|
||||
assert_eq!(raw.shared[[i1, i2]], 0);
|
||||
assert_eq!(raw.shared[[i2, i1]], 0);
|
||||
|
||||
// cardinality_tally: both genomes carry exactly one member of the
|
||||
// family (cardinality 1 each) -> one co-occurrence at [1][1].
|
||||
let cardinality = merged.cardinality_tally(&raw, 1.0).expect("cardinality_tally");
|
||||
assert_eq!(cardinality.counts[1][1], 1);
|
||||
let total: u64 = cardinality.counts.iter().flatten().sum();
|
||||
assert_eq!(total, 1, "no other cardinality pair should be tallied");
|
||||
|
||||
// base_pair_tally: the one eligible, differing locus is C (base 1) vs
|
||||
// G (base 2).
|
||||
let base_pairs = merged.base_pair_tally(&raw, 1.0).expect("base_pair_tally");
|
||||
assert_eq!(base_pairs.counts[1][2], 1);
|
||||
assert_eq!(base_pairs.counts[2][1], 1);
|
||||
let total: u64 = base_pairs.counts.iter().flatten().sum();
|
||||
assert_eq!(total, 2, "no other base pair should be tallied");
|
||||
assert_eq!(base_pairs.same, [0, 0, 0, 0], "the two genomes never agree at this locus");
|
||||
}
|
||||
Reference in New Issue
Block a user