Push zpwxxpnpktps #67
@@ -1624,10 +1624,33 @@ a fixed reference. Storing instead a **4-bit mask** — one bit per base
|
||||
family's fixed canonical form, i.e. the member with `A` at the centre — see
|
||||
above) is observed anywhere in the index — fixes both:
|
||||
- **Sibling count is derived, not stored**: `siblings = popcount(mask) - 1`.
|
||||
- **Minorant is derived, not stored**: regenerate the family's 4 canonical
|
||||
- ~~**Minorant is derived, not stored**: regenerate the family's 4 canonical
|
||||
forms from the slot's own k-mer (cheap, no lookup — see above), compare
|
||||
the raw encodings of whichever bits are set in the mask, take the
|
||||
smallest.
|
||||
smallest.~~ **Erratum (2026-08-14) — this was wrong, kept struck through
|
||||
rather than deleted.** "Cheap, no lookup" only accounts for the bit
|
||||
algebra (regenerate 4 forms, compare raw encodings) — true in isolation,
|
||||
but it silently assumed "the slot's own k-mer" is a free fact. It isn't:
|
||||
getting from an MPHF slot index back to the actual k-mer sequence means
|
||||
reconstructing `slot_kmer` for the whole layer — scan `unitigs.bin`,
|
||||
`mphf.find()` every k-mer to place it — an O(distinct k-mers) pass
|
||||
through the MPHF, not a per-slot O(1) lookup. That reconstruction is free
|
||||
*only* when the caller already needs k-mer identity for something else in
|
||||
the same traversal (e.g. the SNP sweep below, which needs it anyway to
|
||||
generate `central_canonical_neighbors()`). A caller that wants *only* the
|
||||
minorant flag pays the full reconstruction for nothing: measured on a real
|
||||
run, a bare family-size histogram (four buckets, otherwise near-instant)
|
||||
spent 71% of wall-clock in `MphfLayer::find`, all of it solely to answer
|
||||
"is this slot the minorant". **Current design: minorant *is* stored**
|
||||
after all — a 5th mask bit (4 presence bits + 1 minorant bit, still fits
|
||||
one byte alongside the presence mask below), written once in the
|
||||
construction pass where the k-mer is already in hand for other reasons
|
||||
(`obikindex::siblings::build_layer_sibling_annex`), read back for free by
|
||||
every later consumer (`FamilyMask::is_minorant`). This is genuinely a
|
||||
return to the superseded 3-bit design's core idea (store minorant
|
||||
alongside the count) — the "which variant" blindness that motivated
|
||||
moving away from it is fixed by keeping the full 4-bit presence mask
|
||||
too, not by dropping the stored minorant bit again.
|
||||
- **A future consumer knows exactly which variants to (re-)query** —
|
||||
`popcount(mask) - 1` lookups instead of always 3, and it knows *which*
|
||||
3 (or fewer) to issue, not just how many hits to expect.
|
||||
@@ -1644,7 +1667,11 @@ above) is observed anywhere in the index — fixes both:
|
||||
per-slot packed array (the presence mask above), one per partition — same
|
||||
on-disk shape family as `PersistentBitMatrix`'s `Packed` variant, but
|
||||
simpler (no per-genome columns, a single derived read-only value per
|
||||
slot).
|
||||
slot). As actually implemented (`obicompactvec::siblingannex`): **not**
|
||||
truly bit-packed — 1 byte/slot, 4 presence bits + the minorant bit (see
|
||||
the erratum above) in the low 5 bits, 3 unused. Deliberately simpler for
|
||||
a first implementation; packing to 5 bits/slot is a pure storage-density
|
||||
follow-up, not a behavioural change, still not done as of this note.
|
||||
<details><summary>Superseded 3-bit design (historical)</summary>
|
||||
3 bits, storing minorant status alongside sibling count directly, since
|
||||
it came for free from the same lookups (point 3 below) — 5 real states
|
||||
|
||||
@@ -41,11 +41,22 @@ const MAGIC: [u8; 4] = *b"PSIB";
|
||||
// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows.
|
||||
const HEADER_SIZE: usize = 16;
|
||||
|
||||
/// A family presence mask: bit `b` set iff the member whose own canonical
|
||||
/// central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the index.
|
||||
/// A family presence mask: bit `b` (0..3) set iff the member whose own
|
||||
/// canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the
|
||||
/// index. Bit 4 is a second, independent fact piggy-backed onto the same
|
||||
/// byte: whether *this slot's own k-mer* is its family's minorant (the
|
||||
/// smallest raw encoding among the family's observed members) — computed
|
||||
/// once, when the whole family's mask is already final (see
|
||||
/// `obikindex::siblings::build_sibling_annex`), and read back by every
|
||||
/// consumer that would otherwise have to reconstruct this slot's k-mer from
|
||||
/// `unitigs.bin` and hash it through the MPHF again just to ask the same
|
||||
/// question. Bits 5-7 unused.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct FamilyMask(u8);
|
||||
|
||||
const PRESENCE_BITS: u8 = 0b0000_1111;
|
||||
const MINORANT_BIT: u8 = 0b0001_0000;
|
||||
|
||||
impl FamilyMask {
|
||||
/// The empty mask — never a valid *computed* result (a slot's own base
|
||||
/// is always present in its own family) — used only to build up a mask
|
||||
@@ -69,7 +80,7 @@ impl FamilyMask {
|
||||
/// Number of family members observed anywhere in the index (1..=4).
|
||||
#[inline]
|
||||
pub fn family_size(self) -> u32 {
|
||||
self.0.count_ones()
|
||||
(self.0 & PRESENCE_BITS).count_ones()
|
||||
}
|
||||
|
||||
/// Number of *other* members observed (0..=3) — `family_size() - 1`.
|
||||
@@ -78,18 +89,39 @@ impl FamilyMask {
|
||||
self.family_size() - 1
|
||||
}
|
||||
|
||||
/// Raw bitmask (bit `b` = base `b` present) — for callers that build up
|
||||
/// a mask via their own bit operations (e.g. concurrently, via an
|
||||
/// `AtomicU8`) and only need the `FamilyMask` wrapper at the end.
|
||||
/// Set or clear the minorant flag (see the struct docs).
|
||||
#[inline]
|
||||
pub fn bits(self) -> u8 {
|
||||
self.0
|
||||
pub fn with_minorant(self, is_minorant: bool) -> Self {
|
||||
if is_minorant {
|
||||
FamilyMask(self.0 | MINORANT_BIT)
|
||||
} else {
|
||||
FamilyMask(self.0 & !MINORANT_BIT)
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct from a raw bitmask (only the low 4 bits are kept).
|
||||
/// Is this slot's own k-mer its family's minorant? Only meaningful once
|
||||
/// `with_minorant` has been called with the family's *final* mask (i.e.
|
||||
/// after construction) — see the struct docs.
|
||||
#[inline]
|
||||
pub fn is_minorant(self) -> bool {
|
||||
self.0 & MINORANT_BIT != 0
|
||||
}
|
||||
|
||||
/// Raw presence bitmask (bit `b` = base `b` present, low 4 bits only —
|
||||
/// never includes the minorant flag) — for callers that build up a mask
|
||||
/// via their own bit operations (e.g. concurrently, via an `AtomicU8`)
|
||||
/// and only need the `FamilyMask` wrapper at the end.
|
||||
#[inline]
|
||||
pub fn bits(self) -> u8 {
|
||||
self.0 & PRESENCE_BITS
|
||||
}
|
||||
|
||||
/// Construct from a raw presence bitmask (only the low 4 bits are kept —
|
||||
/// the minorant flag is not part of this, use [`with_minorant`](Self::with_minorant)
|
||||
/// separately).
|
||||
#[inline]
|
||||
pub fn from_bits(bits: u8) -> Self {
|
||||
FamilyMask(bits & 0b1111)
|
||||
FamilyMask(bits & PRESENCE_BITS)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
@@ -101,10 +133,13 @@ impl FamilyMask {
|
||||
fn decode(byte: u8) -> Option<Self> {
|
||||
if byte == 0 {
|
||||
// Unreachable for a real result — reserved as the "not yet
|
||||
// computed" sentinel.
|
||||
// computed" sentinel. Still safe as a sentinel with the
|
||||
// minorant bit added: a real entry always has at least one
|
||||
// presence bit set (bits 0-3), so byte == 0 still means
|
||||
// "nothing written yet" and never a genuine minorant-only value.
|
||||
return None;
|
||||
}
|
||||
Some(FamilyMask(byte & 0b1111))
|
||||
Some(FamilyMask(byte))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obisys::progress_bar;
|
||||
|
||||
@@ -68,7 +70,7 @@ impl KmerIndex {
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
|
||||
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");
|
||||
|
||||
@@ -10,14 +10,13 @@ use obipipeline::ThrottleGuard;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obiskio::UnitigFileReader;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::helpers::{central_base, partition_of};
|
||||
use super::helpers::{central_base, is_minorant, partition_of};
|
||||
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||||
|
||||
// ── obipipeline data types ─────────────────────────────────────────────────
|
||||
@@ -117,15 +116,10 @@ impl KmerIndex {
|
||||
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
|
||||
let n_slots = mphf.n();
|
||||
|
||||
// ── Enumerate this layer's distinct k-mers, one per slot ────────────
|
||||
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; n_slots];
|
||||
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
|
||||
.map_err(OKIError::Partition)?;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
slot_kmer[slot] = Some(kmer);
|
||||
}
|
||||
}
|
||||
// ── Enumerate this layer's distinct k-mers, one per slot — direct
|
||||
// slot -> k-mer reconstruction (evidence + direct-access unitigs, no
|
||||
// MPHF hashing, no file scan), not scan-and-hash-every-k-mer-forward.
|
||||
let slot_kmer: Vec<Option<CanonicalKmer>> = (0..n_slots).map(|slot| mphf.kmer_at(slot)).collect();
|
||||
|
||||
let k = self.kmer_size();
|
||||
|
||||
@@ -240,13 +234,22 @@ impl KmerIndex {
|
||||
});
|
||||
|
||||
// ── Write the layer's annex file ─────────────────────────────────────
|
||||
// The minorant flag is computed here, not by every later reader:
|
||||
// this is the one place the whole family's *final* mask and this
|
||||
// slot's own k-mer (already in hand, no extra lookup) are both
|
||||
// available together. Every consumer that only needs to know
|
||||
// "is this slot 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_slots, &annex_path)?;
|
||||
for (slot, m) in mask.iter().enumerate() {
|
||||
if slot_kmer[slot].is_none() {
|
||||
continue; // unused MPHF slot, if any — leave at the sentinel
|
||||
}
|
||||
builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed)));
|
||||
let Some(kmer) = slot_kmer[slot] else { continue }; // unused MPHF slot, if any — leave at the sentinel
|
||||
let final_mask = FamilyMask::from_bits(m.load(Ordering::Relaxed));
|
||||
let minorant = is_minorant(kmer, final_mask, k);
|
||||
builder.set(slot, final_mask.with_minorant(minorant));
|
||||
}
|
||||
builder.close()?;
|
||||
|
||||
|
||||
@@ -104,18 +104,63 @@ impl PartitionCache {
|
||||
.is_some_and(|layers| layers.iter().any(|mphf| mphf.find(variant).is_some()))
|
||||
}
|
||||
|
||||
/// Per-genome presence vector for `variant` in partition `dest_partition`
|
||||
/// (`true` iff that genome carries it), `None` on a miss. Same shape as
|
||||
/// `find`, but also reads the cached matrix instead of just the MPHF.
|
||||
pub(super) fn find_presence(&self, dest_partition: usize, variant: CanonicalKmer, n_genomes: usize) -> Option<Vec<bool>> {
|
||||
let layers = self.layers.get(dest_partition)?;
|
||||
let mats = self.mats.get(dest_partition)?;
|
||||
for (mphf, mat) in layers.iter().zip(mats.iter()) {
|
||||
/// 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);
|
||||
return Some((0..n_cols).map(|g| mat.carries(g, slot)).collect());
|
||||
for g in 0..n_cols {
|
||||
for &(slot, family_idx, base) in &hits {
|
||||
if mat.carries(g, slot) {
|
||||
on_hit(family_idx, base, g);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
@@ -72,7 +74,7 @@ impl KmerIndex {
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
|
||||
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");
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
@@ -79,7 +81,7 @@ impl KmerIndex {
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
|
||||
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");
|
||||
|
||||
@@ -4,25 +4,51 @@
|
||||
//! 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 at a time jumping
|
||||
//! between partitions in family order.
|
||||
//! contiguous sweep per partition, instead of one lookup jumping between
|
||||
//! partitions in family order.
|
||||
//!
|
||||
//! Processed in bounded batches of [`FAMILY_BATCH`] families, not the whole
|
||||
//! layer at once — same rationale and constant as `build_layer_sibling_annex`'s
|
||||
//! `BATCH_SIZE`. An earlier version of this function materialised every
|
||||
//! family of the whole layer up front (`Vec<FamilyRow>`, each with its own
|
||||
//! heap-allocated `Vec<u8>`) before handing anything back to the caller:
|
||||
//! O(layer's family count × genome count) extra memory, on top of a second
|
||||
//! full copy for the owned per-family vectors — measured pushing a real run
|
||||
//! into heavy VM compression/swap even though the per-lookup I/O pattern was
|
||||
//! already fixed. A batch of a few thousand families needs the same
|
||||
//! partition-grouping to keep lookups local; it does not need the entire
|
||||
//! layer resident at once. Callers get each family's `genome_mask` through
|
||||
//! `on_family`, backed by one reused scratch buffer — no per-family
|
||||
//! allocation at all, matching the O(genome count) working set the original,
|
||||
//! pre-locality-fix per-family code had.
|
||||
//! 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::*;
|
||||
@@ -31,21 +57,21 @@ use obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix,
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obiskio::UnitigFileReader;
|
||||
use obipipeline::{ThrottleGuard, throttle};
|
||||
|
||||
use crate::error::{OKIError, OKIResult};
|
||||
use crate::index::KmerIndex;
|
||||
|
||||
use super::cache::{Mat, PartitionCache};
|
||||
use super::helpers::{central_base, is_minorant, partition_of};
|
||||
use super::helpers::{central_base, partition_of};
|
||||
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||||
|
||||
/// Families per batch — bounds `scan_layer_families`'s extra memory to
|
||||
/// `FAMILY_BATCH * n_genomes` bytes regardless of layer size. Same value and
|
||||
/// rationale as `build_layer_sibling_annex`'s `BATCH_SIZE`: large enough to
|
||||
/// amortise the per-batch partition-grouping/sync cost, small enough to keep
|
||||
/// working memory a rounding error next to the index itself.
|
||||
const FAMILY_BATCH: usize = 4096;
|
||||
/// 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
|
||||
@@ -75,9 +101,52 @@ impl KmerIndex {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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. `on_family` is
|
||||
/// called once per family with its own annex mask and its per-genome
|
||||
/// [`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
|
||||
@@ -88,7 +157,7 @@ pub(super) fn scan_layer_families(
|
||||
n_genomes: usize,
|
||||
with_counts: bool,
|
||||
k: usize,
|
||||
cache: &PartitionCache,
|
||||
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");
|
||||
@@ -96,13 +165,12 @@ pub(super) fn scan_layer_families(
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
|
||||
|
||||
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; annex.len()];
|
||||
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")).map_err(OKIError::Partition)?;
|
||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
slot_kmer[slot] = Some(kmer);
|
||||
}
|
||||
}
|
||||
// `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 {
|
||||
@@ -112,77 +180,132 @@ pub(super) fn scan_layer_families(
|
||||
};
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
|
||||
let n_slots = annex.len();
|
||||
let mut slot = 0usize;
|
||||
let mut batch_slots: Vec<usize> = Vec::with_capacity(FAMILY_BATCH);
|
||||
let mut scratch = vec![0u8; n_genomes]; // reused across every `on_family` call — no per-family allocation
|
||||
|
||||
while slot < n_slots {
|
||||
batch_slots.clear();
|
||||
while slot < n_slots && batch_slots.len() < FAMILY_BATCH {
|
||||
let s = slot;
|
||||
slot += 1;
|
||||
let Some(mask) = annex.get(s) else { continue };
|
||||
let Some(kmer) = slot_kmer[s] else { continue };
|
||||
if !is_minorant(kmer, mask, k) {
|
||||
continue; // family tallied once, at its minorant
|
||||
}
|
||||
batch_slots.push(s);
|
||||
}
|
||||
if batch_slots.is_empty() {
|
||||
continue; // ran out of minorant slots before filling a batch — loop condition ends it
|
||||
// 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(());
|
||||
}
|
||||
|
||||
// ── Local only: seed each family's own minorant presence (already
|
||||
// open, no lookup) and collect this batch's cross-partition
|
||||
// queries, grouped by destination partition.
|
||||
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
|
||||
let genome_mask: Vec<AtomicU8> = (0..batch_slots.len() * n_genomes).map(|_| AtomicU8::new(0)).collect();
|
||||
let ctx = Arc::new(LayerCtx { slot_kmer, mat, n_parts, n_genomes, n_cols, k });
|
||||
|
||||
for (family_idx, &fslot) in batch_slots.iter().enumerate() {
|
||||
let kmer = slot_kmer[fslot].expect("batch slot has a kmer");
|
||||
let mask = annex.get(fslot).expect("batch slot has a mask");
|
||||
let base = central_base(kmer, k);
|
||||
for g in 0..n_cols {
|
||||
if mat.carries(g, fslot) {
|
||||
genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
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 above straight from `mat`, no lookup
|
||||
continue; // local — resolved below straight from `mat`, no lookup
|
||||
}
|
||||
let b = central_base(other, k);
|
||||
let b = central_base(other, ctx.k);
|
||||
if !mask.has(b) {
|
||||
continue;
|
||||
}
|
||||
let dest = partition_of(other, n_parts);
|
||||
outgoing[dest].push((other, family_idx, b));
|
||||
let dest = partition_of(other, ctx.n_parts);
|
||||
outgoing[dest].push((other, i, b));
|
||||
}
|
||||
}
|
||||
|
||||
// ── Resolve this batch's queries, one contiguous sweep per
|
||||
// destination partition — same rationale as `build_sibling_annex`'s
|
||||
// `outgoing` grouping.
|
||||
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
||||
for &(variant, family_idx, base) in queries {
|
||||
let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue };
|
||||
for (g, &present) in presence.iter().enumerate() {
|
||||
if present {
|
||||
genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
||||
// 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 (family_idx, &fslot) in batch_slots.iter().enumerate() {
|
||||
let mask = annex.get(fslot).expect("batch slot has a mask");
|
||||
for i in 0..n {
|
||||
for (g, dst) in scratch.iter_mut().enumerate() {
|
||||
*dst = genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed);
|
||||
*dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed);
|
||||
}
|
||||
on_family(mask, &scratch);
|
||||
on_family(batch.masks[i], &scratch);
|
||||
}
|
||||
next_expected += n;
|
||||
}
|
||||
}
|
||||
debug_assert_eq!(next_expected, total_families, "every batch must have been replayed");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
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;
|
||||
|
||||
@@ -30,6 +34,42 @@ pub struct SiblingAnnexStats {
|
||||
}
|
||||
|
||||
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
|
||||
@@ -53,7 +93,7 @@ impl KmerIndex {
|
||||
n_bits,
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = PartitionCache::build(&partition, n_parts, with_counts)?;
|
||||
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
|
||||
|
||||
@@ -120,7 +120,8 @@ fn sibling_annex_one_sibling_each() {
|
||||
// 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* mask (bits 1 and 2 set).
|
||||
// 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");
|
||||
@@ -132,14 +133,16 @@ fn sibling_annex_one_sibling_each() {
|
||||
let expected_mask = FamilyMask::EMPTY.with(1).with(2);
|
||||
|
||||
let a = annex_info_for(&merged, g1_kmer);
|
||||
assert_eq!(a, expected_mask, "AACCGCTTAAG");
|
||||
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, expected_mask, "AACCGGTTAAG");
|
||||
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]
|
||||
@@ -185,6 +188,25 @@ fn sibling_annex_stats_counts_each_family_once_and_per_genome() {
|
||||
}
|
||||
}
|
||||
|
||||
#[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
|
||||
|
||||
@@ -106,6 +106,17 @@ pub struct PhyloArgs {
|
||||
#[arg(long)]
|
||||
pub sibling_stats: bool,
|
||||
|
||||
/// Print just the global family-size histogram (1-4 members) of an
|
||||
/// already-built annex — the `global` row `--sibling-stats` also
|
||||
/// writes, but without the per-genome breakdown, so it skips
|
||||
/// `--sibling-stats`'s cross-partition resolution entirely: reads only
|
||||
/// the already-open annex mask and each layer's own `unitigs.bin`, cost
|
||||
/// independent of the rest of the index. A quick sanity check that the
|
||||
/// annex itself is sound, decoupled from `--sibling-stats`'s much
|
||||
/// heavier per-genome pass.
|
||||
#[arg(long)]
|
||||
pub sibling_hist: bool,
|
||||
|
||||
/// Compute the raw p-distance restricted to loci that are single-copy
|
||||
/// in both genomes of each pair (an already-built sibling annex is
|
||||
/// required — run with `--sibling-annex` first, in this invocation or
|
||||
@@ -231,7 +242,8 @@ pub struct PhyloArgs {
|
||||
pub sankoff_cost_scale: f64,
|
||||
|
||||
/// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
|
||||
/// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_rawsnp_counts.csv,
|
||||
/// <prefix>_siblings.csv, <prefix>_sibling_hist.csv, <prefix>_rawsnp.csv,
|
||||
/// <prefix>_rawsnp_counts.csv,
|
||||
/// <prefix>_snp.fasta, <prefix>_family_overlap.csv,
|
||||
/// <prefix>_sankoff_matrix.csv, <prefix>_sankoff_params.yaml,
|
||||
/// <prefix>_sankoff.fasta, <prefix>_sankoff.tnt, <prefix>_sankoff.tcm,
|
||||
|
||||
@@ -19,7 +19,7 @@ use tracing::info;
|
||||
pub use args::PhyloArgs;
|
||||
use family_overlap::{apply_min_shared_family_exclusion, write_family_overlap_csv};
|
||||
use iqtree::write_iqtree;
|
||||
use outputs::{write_raw_snp_counts_csv, write_raw_snp_distance_csv, write_sibling_stats_csv, write_snp_fasta, upgma_to_newick};
|
||||
use outputs::{write_raw_snp_counts_csv, write_raw_snp_distance_csv, write_sibling_hist_csv, write_sibling_stats_csv, write_snp_fasta, upgma_to_newick};
|
||||
use phyg::write_sankoff_phyg;
|
||||
use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params};
|
||||
use tnt::write_sankoff_tnt;
|
||||
@@ -135,6 +135,13 @@ pub fn run(args: PhyloArgs) {
|
||||
});
|
||||
write_sibling_stats_csv(&stats, &labels, &args.output);
|
||||
}
|
||||
if args.sibling_hist {
|
||||
let counts = idx.sibling_family_size_histogram().unwrap_or_else(|e| {
|
||||
eprintln!("error computing sibling family-size histogram: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
write_sibling_hist_csv(&counts, &args.output);
|
||||
}
|
||||
if args.raw_snp_distance {
|
||||
let mut result = idx.raw_snp_distance().unwrap_or_else(|e| {
|
||||
eprintln!("error computing raw SNP distance: {e}");
|
||||
|
||||
@@ -76,6 +76,25 @@ pub(super) fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels:
|
||||
info!("raw single-copy SNP distance matrix → {path}");
|
||||
}
|
||||
|
||||
// ── Global family-size histogram (annex-only, no per-genome pass) → CSV ────
|
||||
|
||||
pub(super) fn write_sibling_hist_csv(counts: &[u64; 4], output: &Option<PathBuf>) {
|
||||
let path = output.as_ref()
|
||||
.map(|p| format!("{}_sibling_hist.csv", p.display()))
|
||||
.unwrap_or_else(|| "sibling_hist.csv".into());
|
||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
||||
eprintln!("error creating {path}: {e}");
|
||||
std::process::exit(1);
|
||||
}));
|
||||
writeln!(f, "size,count").unwrap();
|
||||
for (size, count) in counts.iter().enumerate() {
|
||||
writeln!(f, "{},{count}", size + 1).unwrap();
|
||||
}
|
||||
let total: u64 = counts.iter().sum();
|
||||
info!("family-size histogram → {path} (total {total} famil{})",
|
||||
if total == 1 { "y" } else { "ies" });
|
||||
}
|
||||
|
||||
// ── Raw single-copy SNP distance → per-pair diagnostic counts ──────────────
|
||||
//
|
||||
// A pair table (one row per unordered genome pair), not a matrix: the ratio
|
||||
|
||||
@@ -134,6 +134,31 @@ impl MphfLayer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Reconstruct the canonical k-mer stored at `slot` — the inverse of
|
||||
/// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot).
|
||||
/// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives
|
||||
/// `(chunk_id, rank)` directly (no MPHF hashing, no file scan), then
|
||||
/// `unitigs.canonical_raw_kmer` is a direct-access read. `None` on
|
||||
/// `Approx` layers — fingerprints alone can't recover a k-mer, and
|
||||
/// falling back to a full sequential scan here would silently make one
|
||||
/// call cost O(n); callers needing this on an `Approx` layer should
|
||||
/// scan `unitigs.bin` themselves and decide how to handle that cost
|
||||
/// explicitly.
|
||||
pub fn kmer_at(&self, slot: usize) -> Option<CanonicalKmer> {
|
||||
if slot >= self.n {
|
||||
return None;
|
||||
}
|
||||
match &self.ev {
|
||||
LayerEvidence::Exact { evidence, unitigs } |
|
||||
LayerEvidence::Hybrid { evidence, unitigs, .. } => {
|
||||
let (chunk_id, rank) = evidence.decode(slot);
|
||||
let raw = unitigs.canonical_raw_kmer(chunk_id as usize, rank as usize);
|
||||
Some(CanonicalKmer::from_raw_unchecked(raw))
|
||||
}
|
||||
LayerEvidence::Approx { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn n(&self) -> usize { self.n }
|
||||
}
|
||||
|
||||
|
||||
@@ -157,6 +157,16 @@ impl UnitigFileReader {
|
||||
canonical_raw(self.raw_kmer(i, j), self.k) == query.raw()
|
||||
}
|
||||
|
||||
/// Canonical raw k-mer at position `j` of chunk `i` — [`raw_kmer`](Self::raw_kmer)
|
||||
/// reduced to its canonical (`min(raw, revcomp)`) form. The reconstruction
|
||||
/// half of [`verify_canonical_kmer`](Self::verify_canonical_kmer): that
|
||||
/// method computes the same value internally just to compare it away: use
|
||||
/// this one when the caller wants the k-mer itself, not a yes/no answer.
|
||||
#[inline]
|
||||
pub fn canonical_raw_kmer(&self, i: usize, j: usize) -> u64 {
|
||||
canonical_raw(self.raw_kmer(i, j), self.k)
|
||||
}
|
||||
|
||||
// ── Sequential iterators (O(n) running-offset cursor) ─────────────────────
|
||||
|
||||
pub(crate) fn iter_chunks_sequential(&self) -> impl Iterator<Item = (usize, Unitig)> + '_ {
|
||||
|
||||
Reference in New Issue
Block a user