Push zunrplorkwkt #70
@@ -14,7 +14,7 @@ use obikphylo::{
|
||||
cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix,
|
||||
siblings::{
|
||||
DistanceExt, EntropyBias, RawSnpDistanceOutput, SankoffBundleExt, ShannonEntropyExt,
|
||||
SiblingAnnexBuildExt, SiblingStatsExt, SnpAlignment, SnpAlignmentExt,
|
||||
SiblingExt, SiblingStatsExt, SnpAlignment, SnpAlignmentExt,
|
||||
},
|
||||
};
|
||||
use obisys::{Reporter, Stage};
|
||||
|
||||
@@ -5,7 +5,7 @@ use std::sync::Arc;
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::KmerIndex;
|
||||
use obikphylo::siblings::SiblingAnnexBuildExt;
|
||||
use obikphylo::siblings::SiblingExt;
|
||||
use obikphylo::{Metrics, neighbor_joining, upgma};
|
||||
use obisys::{Reporter, Stage};
|
||||
use tracing::info;
|
||||
@@ -33,7 +33,7 @@ pub fn run(args: PhyloArgs) {
|
||||
let mut rep = Reporter::new();
|
||||
|
||||
// Every partition/layer this needs is opened once, up front, and
|
||||
// handed to `Metrics::distance`/`SiblingAnnexBuildExt::build_sibling_annex`
|
||||
// handed to `Metrics::distance`/`SiblingExt::build_sibling_annex`
|
||||
// — see `obikquery`'s own use of `IndexCache` for the same reasoning
|
||||
// (one open, many in-memory reads).
|
||||
let cache = IndexCache::new(Arc::clone(&idx), None);
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
//! `siblings` is being reconnected piece by piece while `obikindex`'s
|
||||
//! stateless audit is caught up with — see the project memory on this. So
|
||||
//! far it only carries the annex-construction path
|
||||
//! (`SiblingAnnexBuildExt::build_sibling_annex`); the rest of the old
|
||||
//! (`SiblingExt::build_sibling_annex`); the rest of the old
|
||||
//! `siblings_old` (SNP distance, cardinality, pseudo-alignment, stats) is
|
||||
//! still disconnected.
|
||||
|
||||
|
||||
@@ -0,0 +1,420 @@
|
||||
//! Shared per-layer family traversal — resolves each minorant family's
|
||||
//! per-genome base-presence, with the same locality discipline as
|
||||
//! [`super::annex::build_layer_sibling_annex`]: 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 `cache` 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: 16
|
||||
//! concurrent sweeps across the same partition space scatters exactly what
|
||||
//! the grouping is meant to prevent, just spread over threads instead of
|
||||
//! over layers. 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.
|
||||
//!
|
||||
//! `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.
|
||||
//!
|
||||
//! The pipeline's own working context (`LayerCtx`) needs an owned,
|
||||
//! `Send + 'static` `KmerLayer` — `IndexCache`'s own copy of this same layer
|
||||
//! is only ever borrowed (`&'_ KmerLayer`, tied to the caller's stack frame),
|
||||
//! so it can't be moved into the worker threads `Pipe::apply` spawns; this
|
||||
//! reopens the layer once, independently, for that purpose. Cross-partition
|
||||
//! resolution, by contrast, runs synchronously on the calling thread (via
|
||||
//! `rayon::scope`-free `par_iter`, no `thread::spawn`), so it can — and
|
||||
//! does — borrow `cache` directly for the whole call.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU8, Ordering};
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::layer::KmerLayer;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obipipeline::ThrottleGuard;
|
||||
|
||||
use crate::siblings::helpers::central_base;
|
||||
use crate::siblings::iter::{SiblingEntry, SiblingLayerExt};
|
||||
use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex};
|
||||
|
||||
/// Families per batch — trades off per-batch memory (`FAMILY_BATCH *
|
||||
/// genome_count` bytes) against giving 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.
|
||||
const FAMILY_BATCH: usize = 65536;
|
||||
|
||||
/// Restricts [`scan_layer_families`] to a subset of a layer's families,
|
||||
/// identified by their iteration-order index. `All` (the only variant any
|
||||
/// current caller uses) keeps a full, unsampled scan; `Some` is exercised
|
||||
/// once sampling (`--subsample`/`--entropy-bias`) is reconnected.
|
||||
pub enum Selection<'a> {
|
||||
All,
|
||||
Some(&'a std::collections::HashSet<usize>),
|
||||
}
|
||||
|
||||
fn is_selected(selected: &Option<std::collections::HashSet<usize>>, family_idx: usize) -> bool {
|
||||
match selected {
|
||||
None => true,
|
||||
Some(set) => set.contains(&family_idx),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
struct LayerCtx {
|
||||
mat: KmerLayer,
|
||||
n_parts: usize,
|
||||
n_genomes: usize,
|
||||
n_cols: usize,
|
||||
k: usize,
|
||||
}
|
||||
|
||||
struct SourceBatch {
|
||||
start_family_idx: usize,
|
||||
/// One entry per minorant family in this batch, straight from
|
||||
/// `iter_minorants_batch` — `order` (iteration-order index, not an MPHF
|
||||
/// slot), `kmer`, and `mask` already carried through, no second annex
|
||||
/// read.
|
||||
entries: Vec<SiblingEntry>,
|
||||
_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,
|
||||
/// layer)` — `layer` is the annex-recorded destination layer
|
||||
/// (`FamilyMask::family_members`'s `layer_value`, `0` if absent),
|
||||
/// trustworthy only when `fast_mode` is true.
|
||||
outgoing: Vec<Vec<(CanonicalKmer, usize, u8, u8)>>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
enum FamData {
|
||||
Batch(SourceBatch),
|
||||
Generated(GeneratedBatch),
|
||||
}
|
||||
|
||||
/// Visits every minorant family of one layer, in iteration 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 iteration order, with that family's iteration-order index
|
||||
/// (the same numbering [`Selection`] indices are drawn from), 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.
|
||||
///
|
||||
/// `source_partition`/`source_layer` identify which of `cache`'s layers to
|
||||
/// scan; `fast_mode` must be the same value `build_layer_sibling_annex`
|
||||
/// used when writing this layer's annex (whether every family member's
|
||||
/// recorded field is a trustworthy layer index — see
|
||||
/// `crate::siblings::FamilyMask`'s docs).
|
||||
pub(crate) fn scan_layer_families(
|
||||
cache: &IndexCache,
|
||||
source_partition: usize,
|
||||
source_layer: usize,
|
||||
n_genomes: usize,
|
||||
fast_mode: bool,
|
||||
selection: &Selection,
|
||||
mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
|
||||
) -> OKIResult<()> {
|
||||
let n_parts = cache.index().n_partitions();
|
||||
let k = cache.index().kmer_size();
|
||||
|
||||
// An owned, independently-opened copy of this same layer — the shared
|
||||
// `cache`'s own copy is only ever borrowed, so it can't be moved into
|
||||
// the `Send + 'static` pipeline workers below (see the module docs).
|
||||
let partition = cache
|
||||
.index()
|
||||
.partition(source_partition)
|
||||
.map_err(|e| OKIError::InvalidInput(format!("partition {source_partition}: {e}")))?;
|
||||
let mat = partition
|
||||
.layer(source_layer)
|
||||
.and_then(KmerLayer::open)
|
||||
.map_err(|e| OKIError::InvalidInput(format!("layer {source_layer}: {e}")))?;
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
let layer_dir: &Path = mat.dir();
|
||||
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).map_err(OKIError::Io)?);
|
||||
|
||||
let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k });
|
||||
|
||||
// Streamed straight from `iter_minorants_batch` (zips this layer's own
|
||||
// `iter_kmers()` with the annex, both in iteration order — never an
|
||||
// MPHF slot) — never collected into a `Vec` first: a layer can hold
|
||||
// billions of k-mers.
|
||||
let batches = ctx
|
||||
.mat
|
||||
.iter_minorants_batch(annex, FAMILY_BATCH)
|
||||
.scan(0usize, |offset, batch| {
|
||||
let start = *offset;
|
||||
*offset += batch.len();
|
||||
Some((start, batch))
|
||||
});
|
||||
|
||||
let n_workers = obisys::effective_parallelism();
|
||||
let capacity = 4;
|
||||
let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch {
|
||||
start_family_idx: t.item.0,
|
||||
entries: t.item.1,
|
||||
_permit: t.guard,
|
||||
});
|
||||
|
||||
// `Selection::Some` borrows a `HashSet` whose lifetime doesn't span the
|
||||
// pipeline's worker threads — cloned once into an `Arc` so every worker
|
||||
// can share it cheaply instead of requiring `selection` itself to be
|
||||
// `'static`. `None` means [`Selection::All`], checked with a plain
|
||||
// `map_or`-style match at each use instead of allocating an
|
||||
// always-true set.
|
||||
let selected: Arc<Option<std::collections::HashSet<usize>>> = Arc::new(match selection {
|
||||
Selection::All => None,
|
||||
Selection::Some(set) => Some((*set).clone()),
|
||||
});
|
||||
|
||||
let worker_ctx = Arc::clone(&ctx);
|
||||
let worker_selected = Arc::clone(&selected);
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
FamData : SourceBatch => GeneratedBatch,
|
||||
| {
|
||||
move |batch: SourceBatch| -> GeneratedBatch {
|
||||
let ctx = &worker_ctx;
|
||||
let n = batch.entries.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, u8)>> =
|
||||
(0..ctx.n_parts).map(|_| Vec::new()).collect();
|
||||
|
||||
// Pass 1: cheap, no matrix access — own base and this
|
||||
// batch's cross-partition queries. Kmer and mask already in
|
||||
// hand from `iter_minorants_batch`, no second annex read,
|
||||
// no slot lookup needed for this pass. `family_members`
|
||||
// (not a hand-rolled `central_canonical_neighbors` +
|
||||
// `mask.has` loop) already filters to present members and
|
||||
// hands back each one's annex-recorded layer alongside.
|
||||
// Families outside `selected` (sampling only) still get a
|
||||
// mask/base entry — pass 2 indexes uniformly by `i` — but
|
||||
// never an outgoing query: that's the expensive part a
|
||||
// sampled run exists to avoid paying for every family.
|
||||
for (i, entry) in batch.entries.iter().enumerate() {
|
||||
let (kmer, mask) = (entry.kmer, entry.mask);
|
||||
masks.push(mask);
|
||||
let base = central_base(kmer, ctx.k);
|
||||
bases.push(base);
|
||||
let family_idx = batch.start_family_idx + i;
|
||||
if !is_selected(&worker_selected, family_idx) {
|
||||
continue;
|
||||
}
|
||||
for (member, layer) in mask.family_members(kmer, ctx.k) {
|
||||
if member == kmer {
|
||||
continue; // local — resolved below straight from `mat`, no lookup
|
||||
}
|
||||
let b = central_base(member, ctx.k);
|
||||
let dest = member.partition(ctx.n_parts);
|
||||
outgoing[dest].push((member, i, b, layer.unwrap_or(0)));
|
||||
}
|
||||
}
|
||||
|
||||
// Pass 2: genome-major, not family-major — the matrix is
|
||||
// stored one contiguous block per genome (column), slot as
|
||||
// the offset within it. `fill_sub_matrix_carries` sorts the
|
||||
// slots internally for a sequential mmap sweep per column,
|
||||
// then restores this batch's order. The presence/count
|
||||
// matrix is still MPHF-slot-indexed (unlike the annex), so
|
||||
// each entry's kmer is mapped to its slot via `hash_batch`
|
||||
// — a pure MPHF lookup, no evidence check, since these are
|
||||
// this layer's own kmers, known members by construction.
|
||||
let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect();
|
||||
let slots = ctx.mat.hash_batch(&kmers);
|
||||
let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect();
|
||||
ctx.mat.fill_sub_matrix_carries(&slots, &mut carries);
|
||||
for (g, col) in carries.iter().enumerate() {
|
||||
for (i, &carries_it) in col.iter().enumerate() {
|
||||
if carries_it {
|
||||
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)| {
|
||||
let on_hit = |i: usize, base: u8, g: usize| {
|
||||
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
||||
};
|
||||
resolve_partition_batch(cache, dest, queries, n_genomes, fast_mode, on_hit);
|
||||
});
|
||||
|
||||
for i in 0..n {
|
||||
let family_idx = batch.start_family_idx + i;
|
||||
if !is_selected(&selected, family_idx) {
|
||||
continue; // not in the sample — never resolved above, nothing to report
|
||||
}
|
||||
for (g, dst) in scratch.iter_mut().enumerate() {
|
||||
*dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed);
|
||||
}
|
||||
on_family(family_idx, batch.masks[i], &scratch);
|
||||
}
|
||||
next_expected += n;
|
||||
}
|
||||
}
|
||||
debug_assert!(
|
||||
pending.is_empty(),
|
||||
"every generated batch must have been replayed"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve one batch's `(variant, family_idx, base, layer_hint)` queries
|
||||
/// against `dest_partition`'s already-cached layers, calling
|
||||
/// `on_hit(family_idx, base, g)` for every genome `g` that carries the
|
||||
/// resolved variant. `fast_mode` selects between trusting `layer_hint`
|
||||
/// outright (a direct `hash_batch`, no evidence check — safe because that
|
||||
/// exact variant's membership in that exact layer was already positively
|
||||
/// established once, at annex-build time) and probing every layer of the
|
||||
/// partition in turn with `find_slot` (first hit wins), for indexes with
|
||||
/// too many layers per partition for `layer_hint` to be trustworthy (see
|
||||
/// `crate::siblings::FamilyMask`'s docs).
|
||||
fn resolve_partition_batch(
|
||||
cache: &IndexCache,
|
||||
dest_partition: usize,
|
||||
queries: &[(CanonicalKmer, usize, u8, u8)],
|
||||
n_genomes: usize,
|
||||
fast_mode: bool,
|
||||
mut on_hit: impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let Some(n_layer) = cache.n_layer(dest_partition) else {
|
||||
return;
|
||||
};
|
||||
|
||||
if fast_mode {
|
||||
let mut by_layer: Vec<Vec<(CanonicalKmer, usize, u8)>> = vec![Vec::new(); n_layer];
|
||||
for &(variant, family_idx, base, layer_hint) in queries {
|
||||
if let Some(bucket) = by_layer.get_mut(layer_hint as usize) {
|
||||
bucket.push((variant, family_idx, base));
|
||||
}
|
||||
}
|
||||
for (li, entries) in by_layer.into_iter().enumerate() {
|
||||
if entries.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let layer = cache
|
||||
.get_layer(dest_partition, li)
|
||||
.expect("cache-consistent: n_layer(dest_partition) already bounds li");
|
||||
let variants: Vec<CanonicalKmer> = entries.iter().map(|&(v, _, _)| v).collect();
|
||||
let slots = layer.hash_batch(&variants);
|
||||
let hits: Vec<(usize, usize, u8)> = slots
|
||||
.into_iter()
|
||||
.zip(entries.iter())
|
||||
.map(|(slot, &(_, family_idx, base))| (slot, family_idx, base))
|
||||
.collect();
|
||||
resolve_layer_hits(layer, &hits, n_genomes, &mut on_hit);
|
||||
}
|
||||
} else {
|
||||
// First hit wins, same semantics as the fast path — a variant
|
||||
// present in an earlier layer shadows later ones.
|
||||
let mut by_layer: Vec<Vec<(usize, usize, u8)>> = vec![Vec::new(); n_layer];
|
||||
for &(variant, family_idx, base, _layer_hint) in queries {
|
||||
for li in 0..n_layer {
|
||||
let layer = cache
|
||||
.get_layer(dest_partition, li)
|
||||
.expect("cache-consistent: n_layer(dest_partition) already bounds li");
|
||||
if let Some(slot) = layer.find_slot(variant) {
|
||||
by_layer[li].push((slot, family_idx, base));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (li, hits) in by_layer.into_iter().enumerate() {
|
||||
if hits.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let layer = cache
|
||||
.get_layer(dest_partition, li)
|
||||
.expect("cache-consistent: n_layer(dest_partition) already bounds li");
|
||||
resolve_layer_hits(layer, &hits, n_genomes, &mut on_hit);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared tail of both branches of [`resolve_partition_batch`]: given one
|
||||
/// layer's already-resolved `(slot, family_idx, base)` hits, sweep
|
||||
/// genome-major (via `fill_sub_matrix_carries`, buffer-reusing, no per-call
|
||||
/// allocation) and call `on_hit` for every carry.
|
||||
fn resolve_layer_hits(
|
||||
layer: &KmerLayer,
|
||||
hits: &[(usize, usize, u8)],
|
||||
n_genomes: usize,
|
||||
on_hit: &mut impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let n_cols = layer.n_cols().min(n_genomes);
|
||||
let slots: Vec<usize> = hits.iter().map(|&(slot, _, _)| slot).collect();
|
||||
let mut carries: Vec<Vec<bool>> = (0..n_cols).map(|_| Vec::new()).collect();
|
||||
layer.fill_sub_matrix_carries(&slots, &mut carries);
|
||||
for (g, col) in carries.iter().enumerate() {
|
||||
for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) {
|
||||
if carries_it {
|
||||
on_hit(family_idx, base, g);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,11 +1,13 @@
|
||||
//! Sibling-annex construction algorithm: the cross-partition sweep that
|
||||
//! fills in one layer's [`crate::siblings::FamilyMask`] slots — driven, one
|
||||
//! layer at a time, by
|
||||
//! [`crate::siblings::extensions::SiblingAnnexBuildExt`]. Kept apart from
|
||||
//! [`crate::siblings::extensions::SiblingExt`]. Kept apart from
|
||||
//! that trait (and from the `siblingannex`/`helpers` data model) the same
|
||||
//! way `obikindexer` splits `algorithms` from `extensions`: this module
|
||||
//! depends on the data model, never the reverse.
|
||||
|
||||
mod annex;
|
||||
mod family_scan;
|
||||
|
||||
pub(crate) use annex::build_layer_sibling_annex;
|
||||
pub(crate) use family_scan::{Selection, scan_layer_families};
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! [`SiblingAnnexBuildExt`] — the public entry point for building the
|
||||
//! [`SiblingExt`] — the public entry point for building the
|
||||
//! sibling-count/minorant annex of an already-cached index. Thin: it only
|
||||
//! walks `cache`'s partitions/layers and delegates each layer's actual
|
||||
//! cross-partition sweep to [`crate::siblings::algorithms::build_layer_sibling_annex`]
|
||||
@@ -13,7 +13,7 @@ use obisys::progress_bar;
|
||||
use crate::siblings::algorithms::build_layer_sibling_annex;
|
||||
|
||||
/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `IndexCache`.
|
||||
pub trait SiblingAnnexBuildExt {
|
||||
pub trait SiblingExt {
|
||||
/// Build the sibling-count/minorant annex for every layer of every
|
||||
/// cached partition, writing one annex file per layer alongside its
|
||||
/// existing index files. Safe to call again later (e.g. after a fresh
|
||||
@@ -31,7 +31,7 @@ pub trait SiblingAnnexBuildExt {
|
||||
fn build_sibling_annex(&self) -> OKIResult<()>;
|
||||
}
|
||||
|
||||
impl SiblingAnnexBuildExt for IndexCache {
|
||||
impl SiblingExt for IndexCache {
|
||||
fn build_sibling_annex(&self) -> OKIResult<()> {
|
||||
let n_parts = self.index().n_partitions();
|
||||
|
||||
|
||||
@@ -6,4 +6,4 @@
|
||||
|
||||
mod annex_build;
|
||||
|
||||
pub use annex_build::SiblingAnnexBuildExt;
|
||||
pub use annex_build::SiblingExt;
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
//! Phylo/sibling-domain iteration over a layer — an extension trait, not a
|
||||
//! new field on `TypedLayer<D>`/`KmerLayer`: "family"/"minorant" are phylo
|
||||
//! concepts, `obikindex::layer` stays kmer/slot-mapping only.
|
||||
//!
|
||||
//! The sibling annex is persisted in the same order as `iter_kmers()` (see
|
||||
//! `crate::siblings::algorithms::build_layer_sibling_annex`), so pairing
|
||||
//! them is a plain zip — no MPHF, no slot. Both sides are already `Send +
|
||||
//! 'static` (`KmerIter` owns an `Arc`-backed reader 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.
|
||||
//!
|
||||
//! Four iterator types, deliberately mirroring `obikindex::layer`'s own
|
||||
//! `KmerIter`/`KmerBatchIter` pair (single item vs. `Vec` batch) — plus the
|
||||
//! minorant-filtered variant of each, since "all siblings" and "one row per
|
||||
//! family" are both common cases:
|
||||
//!
|
||||
//! | | all entries | minorants only |
|
||||
//! |------------|-----------------------|------------------------|
|
||||
//! | single | [`SiblingIter`] | [`MinorantIter`] |
|
||||
//! | batch | [`SiblingBatchIter`] | [`MinorantBatchIter`] |
|
||||
//!
|
||||
//! No separate "enumerate" variant (unlike `KmerIter`/`enumerate_kmers`):
|
||||
//! [`SiblingEntry`] already carries `order` for free, since pairing with
|
||||
//! the annex requires it anyway.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use obikindex::layer::{KmerIter, LayerData, TypedLayer};
|
||||
use obikseq::CanonicalKmer;
|
||||
|
||||
use super::{FamilyMask, SiblingAnnex};
|
||||
|
||||
/// 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(crate) 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(crate) 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 });
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Batches of [`SiblingIter`]'s entries, `batch_size` at a time — the last
|
||||
/// batch may be shorter. Produced by [`SiblingLayerExt::iter_siblings_batch`].
|
||||
pub(crate) struct SiblingBatchIter {
|
||||
inner: SiblingIter,
|
||||
batch_size: usize,
|
||||
}
|
||||
|
||||
impl Iterator for SiblingBatchIter {
|
||||
type Item = Vec<SiblingEntry>;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
collect_batch(&mut self.inner, self.batch_size)
|
||||
}
|
||||
}
|
||||
|
||||
/// Like [`SiblingIter`], filtered to the minorant of each family — the
|
||||
/// common case, since a family is tallied once, at its minorant. Produced
|
||||
/// by [`SiblingLayerExt::iter_minorants`].
|
||||
pub(crate) struct MinorantIter {
|
||||
inner: SiblingIter,
|
||||
}
|
||||
|
||||
impl Iterator for MinorantIter {
|
||||
type Item = SiblingEntry;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
self.inner.by_ref().find(|e| e.mask.is_minorant())
|
||||
}
|
||||
}
|
||||
|
||||
/// Batches of [`MinorantIter`]'s entries, `batch_size` at a time — the last
|
||||
/// batch may be shorter. Produced by
|
||||
/// [`SiblingLayerExt::iter_minorants_batch`].
|
||||
pub(crate) struct MinorantBatchIter {
|
||||
inner: MinorantIter,
|
||||
batch_size: usize,
|
||||
}
|
||||
|
||||
impl Iterator for MinorantBatchIter {
|
||||
type Item = Vec<SiblingEntry>;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
collect_batch(&mut self.inner, self.batch_size)
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared by [`SiblingBatchIter`] and [`MinorantBatchIter`] — pull up to
|
||||
/// `batch_size` items, `None` once the source is exhausted with nothing left.
|
||||
fn collect_batch<I: Iterator>(inner: &mut I, batch_size: usize) -> Option<Vec<I::Item>> {
|
||||
let mut batch = Vec::with_capacity(batch_size);
|
||||
for _ in 0..batch_size {
|
||||
match inner.next() {
|
||||
Some(item) => batch.push(item),
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
if batch.is_empty() { None } else { Some(batch) }
|
||||
}
|
||||
|
||||
/// Adds phylo/sibling iteration to any `TypedLayer<D>` — the extension that
|
||||
/// turns a plain layer into a "sibling layer". Generic over `D`
|
||||
/// (`LayerData`) rather than implemented once per matrix kind: kmer
|
||||
/// iteration doesn't depend on the data payload, and `TypedLayer<D>` already
|
||||
/// carries `iter_kmers`/`hash_batch`/`fill_sub_matrix_carries` for every `D`.
|
||||
pub(crate) 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), yielding `batch_size`
|
||||
/// entries at a time.
|
||||
fn iter_siblings_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> SiblingBatchIter;
|
||||
|
||||
/// 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>) -> MinorantIter;
|
||||
|
||||
/// Like [`iter_minorants`](Self::iter_minorants), yielding `batch_size`
|
||||
/// minorants at a time.
|
||||
fn iter_minorants_batch(
|
||||
&self,
|
||||
annex: Arc<SiblingAnnex>,
|
||||
batch_size: usize,
|
||||
) -> MinorantBatchIter;
|
||||
}
|
||||
|
||||
impl<D: LayerData> SiblingLayerExt for TypedLayer<D> {
|
||||
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter {
|
||||
SiblingIter {
|
||||
kmers: self.iter_kmers(),
|
||||
annex,
|
||||
order: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_siblings_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> SiblingBatchIter {
|
||||
SiblingBatchIter {
|
||||
inner: self.iter_siblings(annex),
|
||||
batch_size,
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> MinorantIter {
|
||||
MinorantIter {
|
||||
inner: self.iter_siblings(annex),
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_minorants_batch(
|
||||
&self,
|
||||
annex: Arc<SiblingAnnex>,
|
||||
batch_size: usize,
|
||||
) -> MinorantBatchIter {
|
||||
MinorantBatchIter {
|
||||
inner: self.iter_minorants(annex),
|
||||
batch_size,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Same extension, over `obikindex::layer::KmerLayer` (the format-erased
|
||||
/// `Count`/`Presence` handle `IndexCache` actually holds) — dispatch only,
|
||||
/// both arms return the same concrete iterator types (they don't depend on
|
||||
/// which `D` is inside), so no boxing is needed. `obikindex::layer` itself
|
||||
/// can't implement this: "family"/"minorant" are phylo concepts, it stays
|
||||
/// kmer/slot-mapping only (see the module docs above).
|
||||
impl SiblingLayerExt for obikindex::layer::KmerLayer {
|
||||
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter {
|
||||
match self {
|
||||
obikindex::layer::KmerLayer::Count { layer, .. } => layer.iter_siblings(annex),
|
||||
obikindex::layer::KmerLayer::Presence { layer, .. } => layer.iter_siblings(annex),
|
||||
obikindex::layer::KmerLayer::Empty { .. } => {
|
||||
panic!("iter_siblings() called on an Empty layer")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_siblings_batch(&self, annex: Arc<SiblingAnnex>, batch_size: usize) -> SiblingBatchIter {
|
||||
match self {
|
||||
obikindex::layer::KmerLayer::Count { layer, .. } => {
|
||||
layer.iter_siblings_batch(annex, batch_size)
|
||||
}
|
||||
obikindex::layer::KmerLayer::Presence { layer, .. } => {
|
||||
layer.iter_siblings_batch(annex, batch_size)
|
||||
}
|
||||
obikindex::layer::KmerLayer::Empty { .. } => {
|
||||
panic!("iter_siblings_batch() called on an Empty layer")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> MinorantIter {
|
||||
match self {
|
||||
obikindex::layer::KmerLayer::Count { layer, .. } => layer.iter_minorants(annex),
|
||||
obikindex::layer::KmerLayer::Presence { layer, .. } => layer.iter_minorants(annex),
|
||||
obikindex::layer::KmerLayer::Empty { .. } => {
|
||||
panic!("iter_minorants() called on an Empty layer")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_minorants_batch(
|
||||
&self,
|
||||
annex: Arc<SiblingAnnex>,
|
||||
batch_size: usize,
|
||||
) -> MinorantBatchIter {
|
||||
match self {
|
||||
obikindex::layer::KmerLayer::Count { layer, .. } => {
|
||||
layer.iter_minorants_batch(annex, batch_size)
|
||||
}
|
||||
obikindex::layer::KmerLayer::Presence { layer, .. } => {
|
||||
layer.iter_minorants_batch(annex, batch_size)
|
||||
}
|
||||
obikindex::layer::KmerLayer::Empty { .. } => {
|
||||
panic!("iter_minorants_batch() called on an Empty layer")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -22,10 +22,11 @@
|
||||
pub mod algorithms;
|
||||
pub mod extensions;
|
||||
mod helpers;
|
||||
mod iter;
|
||||
mod siblingannex;
|
||||
|
||||
pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
||||
|
||||
pub use extensions::SiblingAnnexBuildExt;
|
||||
pub use extensions::SiblingExt;
|
||||
|
||||
pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib";
|
||||
|
||||
Reference in New Issue
Block a user