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
@@ -0,0 +1,284 @@
|
||||
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 obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::helpers::{central_base, is_minorant};
|
||||
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),
|
||||
}
|
||||
|
||||
/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `KmerIndex`.
|
||||
pub trait SiblingAnnexBuildExt {
|
||||
/// 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`](super::stats::SiblingStatsExt::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.
|
||||
fn build_sibling_annex(&self) -> OKIResult<()>;
|
||||
}
|
||||
|
||||
impl SiblingAnnexBuildExt for KmerIndex {
|
||||
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 += build_layer_sibling_annex(self, &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. A free function, not a `KmerIndex` method — called
|
||||
/// only from `build_sibling_annex` above, in the same file.
|
||||
fn build_layer_sibling_annex(
|
||||
index: &KmerIndex,
|
||||
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 = index.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((
|
||||
variant.partition(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)
|
||||
}
|
||||
Reference in New Issue
Block a user