rename KmerPartition to KmerPartitions and update Mat enum
Rename the KmerPartition type to KmerPartitions across obikindex, obikpartitionner, and obikphylo/siblings to reflect an updated data model. Update the Mat enum in siblings/cache.rs to add a SparsePresence variant and simplify opening logic by delegating sparse versus dense detection to PersistentBitMatrix. Apply consistent code formatting, import reordering, and multi-line refactoring throughout the affected modules.
This commit is contained in:
Eric Coissac
@@ -1,13 +1,13 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::family_scan::{scan_layer_families, Selection};
|
||||
use super::family_scan::{Selection, scan_layer_families};
|
||||
use super::subsample::EntropyBias;
|
||||
|
||||
/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set
|
||||
@@ -70,18 +70,26 @@ pub trait SnpAlignmentExt {
|
||||
/// that sample (or, with `subsample = None`, filter the whole index)
|
||||
/// by a Gaussian kernel on each family's entropy — see
|
||||
/// `DevDocMD/architecture/siblings.md`, "Entropy-biased selection".
|
||||
fn snp_pseudo_alignment(&self, subsample: Option<usize>, entropy_bias: Option<EntropyBias>) -> OKIResult<SnpAlignment>;
|
||||
fn snp_pseudo_alignment(
|
||||
&self,
|
||||
subsample: Option<usize>,
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<SnpAlignment>;
|
||||
}
|
||||
|
||||
impl SnpAlignmentExt for KmerIndex {
|
||||
fn snp_pseudo_alignment(&self, subsample: Option<usize>, entropy_bias: Option<EntropyBias>) -> OKIResult<SnpAlignment> {
|
||||
fn snp_pseudo_alignment(
|
||||
&self,
|
||||
subsample: Option<usize>,
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<SnpAlignment> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta().genomes.len();
|
||||
let with_counts = self.meta().config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
self.root_path(),
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
@@ -90,7 +98,8 @@ impl SnpAlignmentExt for KmerIndex {
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
|
||||
let selections = super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
|
||||
let selections =
|
||||
super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
|
||||
|
||||
let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers");
|
||||
// One layer at a time, not `par_iter()` over layers — same
|
||||
@@ -109,14 +118,23 @@ impl SnpAlignmentExt for KmerIndex {
|
||||
None => Selection::All,
|
||||
Some(set) => Selection::Some(set),
|
||||
};
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, mask, genome_mask| {
|
||||
if mask.family_size() < 2 {
|
||||
return; // monomorphic family — no signal, skip
|
||||
}
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
sequences[g].push(iupac_code(m));
|
||||
}
|
||||
})?;
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&selection,
|
||||
|_family_idx, mask, genome_mask| {
|
||||
if mask.family_size() < 2 {
|
||||
return; // monomorphic family — no signal, skip
|
||||
}
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
sequences[g].push(iupac_code(m));
|
||||
}
|
||||
},
|
||||
)?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
+225
-195
@@ -1,22 +1,22 @@
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obipipeline::ThrottleGuard;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::MphfLayer;
|
||||
use obilayeredmap::meta::IndexMode;
|
||||
use obipipeline::ThrottleGuard;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::helpers::{central_base, is_minorant};
|
||||
use super::{olm_to_ok, FamilyMask, SiblingAnnexBuilder, ANNEX_FILE_NAME};
|
||||
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder, olm_to_ok};
|
||||
|
||||
// ── obipipeline data types ─────────────────────────────────────────────────
|
||||
|
||||
@@ -78,7 +78,7 @@ impl SiblingAnnexBuildExt for KmerIndex {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
self.root_path(),
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
@@ -87,7 +87,11 @@ impl SiblingAnnexBuildExt for KmerIndex {
|
||||
.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 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;
|
||||
@@ -102,12 +106,19 @@ impl SiblingAnnexBuildExt for KmerIndex {
|
||||
let mut part_slots: u64 = 0;
|
||||
for l in 0..meta.n_layers {
|
||||
part_slots += build_layer_sibling_annex(
|
||||
self, &self.partition().layer_dir(part, l), &meta.mode, n_parts, l, &cache,
|
||||
self,
|
||||
&self.partition().layer_dir(part, l),
|
||||
&meta.mode,
|
||||
n_parts,
|
||||
l,
|
||||
&cache,
|
||||
)?;
|
||||
}
|
||||
total_slots += part_slots;
|
||||
pb.inc(1);
|
||||
pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
|
||||
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");
|
||||
@@ -139,213 +150,232 @@ fn build_layer_sibling_annex(
|
||||
// agree; see `PartitionCache::fast_mode`'s docs.
|
||||
let fast_mode = cache.fast_mode();
|
||||
|
||||
// ── The annex file itself is the 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 `DevDocMD/architecture/siblings.md`: the MPHF is not
|
||||
// invertible, and evidence answers membership, not identity). No
|
||||
// separate in-memory accumulator: every concurrent write goes
|
||||
// straight through `SiblingAnnexBuilder::atomic_slot` into the
|
||||
// mmap — a layer can hold billions of k-mers, so a `Vec` shadowing
|
||||
// the whole file in RAM just to copy it out again at the end (the
|
||||
// previous design) doubles memory for no benefit once the builder
|
||||
// itself can be written concurrently.
|
||||
//
|
||||
// `Arc<SiblingAnnexBuilder>`, not a bare `SiblingAnnexBuilder` —
|
||||
// `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. `atomic_slot`,
|
||||
// not `set`, for every concurrent write below: the gather phase
|
||||
// parallelises across destination partitions (independent
|
||||
// `cache.find` calls, safe to run concurrently) and their hits can
|
||||
// land on arbitrary, possibly-shared source entries — a lock-free
|
||||
// `fetch_or` avoids needing any synchronisation beyond that.
|
||||
// Reclaimed as a plain owned value (`Arc::try_unwrap`) once every
|
||||
// concurrent phase below is done, for the sequential finalisation
|
||||
// pass. ─────────────────────────────────────────────────────────
|
||||
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||||
let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?);
|
||||
// ── The annex file itself is the 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 `DevDocMD/architecture/siblings.md`: the MPHF is not
|
||||
// invertible, and evidence answers membership, not identity). No
|
||||
// separate in-memory accumulator: every concurrent write goes
|
||||
// straight through `SiblingAnnexBuilder::atomic_slot` into the
|
||||
// mmap — a layer can hold billions of k-mers, so a `Vec` shadowing
|
||||
// the whole file in RAM just to copy it out again at the end (the
|
||||
// previous design) doubles memory for no benefit once the builder
|
||||
// itself can be written concurrently.
|
||||
//
|
||||
// `Arc<SiblingAnnexBuilder>`, not a bare `SiblingAnnexBuilder` —
|
||||
// `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. `atomic_slot`,
|
||||
// not `set`, for every concurrent write below: the gather phase
|
||||
// parallelises across destination partitions (independent
|
||||
// `cache.find` calls, safe to run concurrently) and their hits can
|
||||
// land on arbitrary, possibly-shared source entries — a lock-free
|
||||
// `fetch_or` avoids needing any synchronisation beyond that.
|
||||
// Reclaimed as a plain owned value (`Arc::try_unwrap`) once every
|
||||
// concurrent phase below is done, for the sequential finalisation
|
||||
// pass. ─────────────────────────────────────────────────────────
|
||||
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||||
let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?);
|
||||
|
||||
// ── 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. `BATCH_SIZE` was originally tuned back when the
|
||||
// per-k-mer cost inside this stage (minimiser via `RollingStat`)
|
||||
// was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`,
|
||||
// a direct O(k) bit-arithmetic replacement) — at the old per-k-mer
|
||||
// cost, dispatch overhead (the scheduler's single-threaded `Select`
|
||||
// loop: one channel round-trip per batch) was negligible next to
|
||||
// the work each batch represented; now that the work itself is
|
||||
// cheaper, that fixed per-batch overhead is proportionally larger,
|
||||
// so a bigger batch amortises it over more k-mers again. ─────────
|
||||
const BATCH_SIZE: usize = 32768;
|
||||
let n_workers = obisys::effective_parallelism();
|
||||
let capacity = 256;
|
||||
// ── 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. `BATCH_SIZE` was originally tuned back when the
|
||||
// per-k-mer cost inside this stage (minimiser via `RollingStat`)
|
||||
// was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`,
|
||||
// a direct O(k) bit-arithmetic replacement) — at the old per-k-mer
|
||||
// cost, dispatch overhead (the scheduler's single-threaded `Select`
|
||||
// loop: one channel round-trip per batch) was negligible next to
|
||||
// the work each batch represented; now that the work itself is
|
||||
// cheaper, that fixed per-batch overhead is proportionally larger,
|
||||
// so a bigger batch amortises it over more k-mers again. ─────────
|
||||
const BATCH_SIZE: usize = 32768;
|
||||
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 straight into the annex in
|
||||
// this same pass, since this is exactly the iteration order the
|
||||
// annex is keyed on — no separate seeding pass needed.
|
||||
let seed_builder = Arc::clone(&builder);
|
||||
let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| {
|
||||
// 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 straight into the annex in
|
||||
// this same pass, since this is exactly the iteration order the
|
||||
// annex is keyed on — no separate seeding pass needed.
|
||||
let seed_builder = Arc::clone(&builder);
|
||||
let batches = mphf
|
||||
.enumerate_kmers_batch(BATCH_SIZE)
|
||||
.map(move |(start, kmers)| {
|
||||
kmers
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, kmer)| {
|
||||
let base = central_base(kmer, k);
|
||||
let bits = if fast_mode { FamilyMask::layer_bits(base, l) } else { FamilyMask::presence_bits(base) };
|
||||
seed_builder.atomic_slot(start + i).fetch_or(bits, Ordering::Relaxed);
|
||||
let bits = if fast_mode {
|
||||
FamilyMask::layer_bits(base, l)
|
||||
} else {
|
||||
FamilyMask::presence_bits(base)
|
||||
};
|
||||
seed_builder
|
||||
.atomic_slot(start + i)
|
||||
.fetch_or(bits, Ordering::Relaxed);
|
||||
(start + i, kmer)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
|
||||
// Diagnostic: count still-empty annex slots after seeding + cross-partition
|
||||
// resolution. A non-zero count here means some k-mer's own base was never
|
||||
// written (should be impossible after the batch_offset fix above).
|
||||
let check_empty = |label: &str| {
|
||||
let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count();
|
||||
if empty > 0 {
|
||||
tracing::warn!("{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})");
|
||||
}
|
||||
};
|
||||
let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch {
|
||||
items: t.item,
|
||||
_permit: t.guard,
|
||||
});
|
||||
// Diagnostic: count still-empty annex slots after seeding + cross-partition
|
||||
// resolution. A non-zero count here means some k-mer's own base was never
|
||||
// written (should be impossible after the batch_offset fix above).
|
||||
let check_empty = |label: &str| {
|
||||
let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count();
|
||||
if empty > 0 {
|
||||
tracing::warn!(
|
||||
"{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})"
|
||||
);
|
||||
}
|
||||
};
|
||||
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),
|
||||
));
|
||||
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,
|
||||
};
|
||||
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));
|
||||
// ── 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));
|
||||
}
|
||||
}
|
||||
|
||||
check_empty("after_seeding");
|
||||
|
||||
// ── Resolve each partition's batch against the cache, parallelised
|
||||
// over roughly equal-sized *chunks*, not over partitions — a plain
|
||||
// `outgoing.par_iter()` over `n_parts` buckets gives one thread the
|
||||
// whole of one partition's bucket, however large, and a lot of
|
||||
// buckets are far from equal: a central-base substitution changes
|
||||
// the minimiser (and thus routes to a different partition) only
|
||||
// when the winning minimiser window overlaps the central position.
|
||||
// For k=31/m=11 that is ~11 of the 21 possible windows — so *most*
|
||||
// of the remaining ~10/21 send the variant right back to the
|
||||
// partition already being built. That self-partition bucket ends up
|
||||
// far larger than any other, so the naive per-partition split
|
||||
// leaves one thread grinding through it alone long after every
|
||||
// other partition's (tiny) bucket is done — confirmed by sampling a
|
||||
// real run: one thread solely in `MphfLayer::find` while the rest
|
||||
// of the pool sits idle. Splitting each non-empty bucket into
|
||||
// `chunk_size`-sized pieces first keeps this pass's per-partition
|
||||
// locality (each chunk is still contiguous within one partition,
|
||||
// still resolved via `cache.find` grouped by that partition) while
|
||||
// letting Rayon spread a single oversized bucket across several
|
||||
// threads instead of pinning it to one. ──────────────────────────
|
||||
let chunk_size = ((outgoing.iter().map(Vec::len).sum::<usize>() / n_workers).max(1)).min(4096);
|
||||
let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, q)| !q.is_empty())
|
||||
.flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c)))
|
||||
.collect();
|
||||
work.par_iter().for_each(|&(dest, chunk)| {
|
||||
for &(variant, source_order, base) in chunk {
|
||||
if let Some(layer) = cache.find(dest, variant) {
|
||||
let bits = if fast_mode {
|
||||
FamilyMask::layer_bits(base, layer)
|
||||
} else {
|
||||
FamilyMask::presence_bits(base)
|
||||
};
|
||||
builder
|
||||
.atomic_slot(source_order)
|
||||
.fetch_or(bits, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
check_empty("after_seeding");
|
||||
check_empty("after_resolution");
|
||||
|
||||
// ── Resolve each partition's batch against the cache, parallelised
|
||||
// over roughly equal-sized *chunks*, not over partitions — a plain
|
||||
// `outgoing.par_iter()` over `n_parts` buckets gives one thread the
|
||||
// whole of one partition's bucket, however large, and a lot of
|
||||
// buckets are far from equal: a central-base substitution changes
|
||||
// the minimiser (and thus routes to a different partition) only
|
||||
// when the winning minimiser window overlaps the central position.
|
||||
// For k=31/m=11 that is ~11 of the 21 possible windows — so *most*
|
||||
// of the remaining ~10/21 send the variant right back to the
|
||||
// partition already being built. That self-partition bucket ends up
|
||||
// far larger than any other, so the naive per-partition split
|
||||
// leaves one thread grinding through it alone long after every
|
||||
// other partition's (tiny) bucket is done — confirmed by sampling a
|
||||
// real run: one thread solely in `MphfLayer::find` while the rest
|
||||
// of the pool sits idle. Splitting each non-empty bucket into
|
||||
// `chunk_size`-sized pieces first keeps this pass's per-partition
|
||||
// locality (each chunk is still contiguous within one partition,
|
||||
// still resolved via `cache.find` grouped by that partition) while
|
||||
// letting Rayon spread a single oversized bucket across several
|
||||
// threads instead of pinning it to one. ──────────────────────────
|
||||
let chunk_size = ((outgoing.iter().map(Vec::len).sum::<usize>() / n_workers).max(1)).min(4096);
|
||||
let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, q)| !q.is_empty())
|
||||
.flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c)))
|
||||
.collect();
|
||||
work.par_iter().for_each(|&(dest, chunk)| {
|
||||
for &(variant, source_order, base) in chunk {
|
||||
if let Some(layer) = cache.find(dest, variant) {
|
||||
let bits = if fast_mode { FamilyMask::layer_bits(base, layer) } else { FamilyMask::presence_bits(base) };
|
||||
builder.atomic_slot(source_order).fetch_or(bits, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
check_empty("after_resolution");
|
||||
|
||||
// ── 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).
|
||||
// Every concurrent phase above is done and its `Arc` clone (the
|
||||
// seeding pass's `seed_builder`) has already been dropped along
|
||||
// with the now-fully-drained `batches` iterator — this is the only
|
||||
// remaining strong reference, so reclaiming plain ownership for the
|
||||
// sequential `set` calls below is guaranteed to succeed.
|
||||
let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| panic!("sibling-annex builder still shared after every concurrent phase completed"));
|
||||
for (order, kmer) in mphf.enumerate_kmers() {
|
||||
let final_mask = builder.get(order).expect("seeded by construction");
|
||||
let minorant = is_minorant(kmer, final_mask, k);
|
||||
builder.set(order, final_mask.with_minorant(minorant));
|
||||
}
|
||||
builder.close()?;
|
||||
// ── 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).
|
||||
// Every concurrent phase above is done and its `Arc` clone (the
|
||||
// seeding pass's `seed_builder`) has already been dropped along
|
||||
// with the now-fully-drained `batches` iterator — this is the only
|
||||
// remaining strong reference, so reclaiming plain ownership for the
|
||||
// sequential `set` calls below is guaranteed to succeed.
|
||||
let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| {
|
||||
panic!("sibling-annex builder still shared after every concurrent phase completed")
|
||||
});
|
||||
for (order, kmer) in mphf.enumerate_kmers() {
|
||||
let final_mask = builder.get(order).expect("seeded by construction");
|
||||
let minorant = is_minorant(kmer, final_mask, k);
|
||||
builder.set(order, final_mask.with_minorant(minorant));
|
||||
}
|
||||
builder.close()?;
|
||||
|
||||
Ok(n as u64)
|
||||
}
|
||||
|
||||
@@ -2,17 +2,17 @@ use rayon::prelude::*;
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, PersistentSparseBitMatrix};
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::{Layer, OLMResult};
|
||||
use obilayeredmap::meta::IndexMode;
|
||||
use obilayeredmap::{Layer, OLMResult};
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::OKIResult;
|
||||
|
||||
use super::iter::SiblingLayerExt;
|
||||
use super::SiblingAnnex;
|
||||
use super::iter::SiblingLayerExt;
|
||||
|
||||
/// Every partition's already-open layers, built **once** for the whole
|
||||
/// `build_sibling_annex` run and shared (read-only) across every lookup, in
|
||||
@@ -41,47 +41,40 @@ use super::SiblingAnnex;
|
||||
/// going back through the low-level pieces `Layer` already assembles.
|
||||
pub(super) enum Mat {
|
||||
Count(Layer<PersistentCompactIntMatrix>),
|
||||
/// Dense *or* sparse (`pack --sparse`d) presence — `PersistentBitMatrix`
|
||||
/// itself is a 4-way enum (`Columnar`/`Packed`/`Sparse`/`Implicit`) that
|
||||
/// already auto-detects sparse storage in its own `open()` (checks
|
||||
/// `presence/sparse_meta.json`) and dispatches every method
|
||||
/// (`row`/`fill_row`/`nonzero_iter`/...) across all four internally —
|
||||
/// see `DevDocMD/implementation/partition_layer_cache.md`. A separate
|
||||
/// `SparsePresence(Layer<PersistentSparseBitMatrix>)` arm used to exist
|
||||
/// here, opened via its own `presence/is_multi.prsb` check; it
|
||||
/// predated `PersistentBitMatrix` growing native sparse support and
|
||||
/// was pure duplication by the time it was removed — every method
|
||||
/// below dispatched it identically to this arm.
|
||||
Presence(Layer<PersistentBitMatrix>),
|
||||
/// Same role as `Presence`, over a layer packed by `pack --sparse`
|
||||
/// (`PersistentSparseBitMatrix`) instead of the dense form — see
|
||||
/// `DevDocMD/architecture/siblings.md`'s sparse-matrix section. Every
|
||||
/// `Mat` method below dispatches to the same `Layer<D>` generic code
|
||||
/// (`D: LayerData`) as `Presence`, so this arm is purely a storage
|
||||
/// choice, not a behavioural difference.
|
||||
SparsePresence(Layer<PersistentSparseBitMatrix>),
|
||||
}
|
||||
|
||||
impl Mat {
|
||||
/// Open one layer's matrix, auto-detecting count vs. dense-presence vs.
|
||||
/// sparse-presence from what is actually on disk — the single source of
|
||||
/// truth for *every* caller that opens a layer's own matrix (this
|
||||
/// module's cross-partition [`PartitionCache::build`] and
|
||||
/// `family_scan::scan_layer_families`'s own-layer lookup alike), so the
|
||||
/// two can never disagree about which format a layer's `presence/` was
|
||||
/// packed to. Before this existed, `scan_layer_families` open-coded its
|
||||
/// own (non-sparse-aware) copy of this decision: on a `pack --sparse`d
|
||||
/// layer, `presence/matrix.pbmx` no longer exists (removed by
|
||||
/// `pack_sparse_bit_matrix`), so a caller that only checks for it and
|
||||
/// otherwise unconditionally opens `PersistentBitMatrix` doesn't error —
|
||||
/// `PersistentBitMatrix::open`'s own auto-detection falls all the way
|
||||
/// through to the `Implicit` (mono-genome, all-present) case, silently
|
||||
/// corrupting every read.
|
||||
/// Open one layer's matrix, auto-detecting count vs. presence from what
|
||||
/// is actually on disk — the single source of truth for *every* caller
|
||||
/// that opens a layer's own matrix (this module's cross-partition
|
||||
/// [`PartitionCache::build`] and `family_scan::scan_layer_families`'s
|
||||
/// own-layer lookup alike), so the two can never disagree. Sparse vs.
|
||||
/// dense presence storage is `PersistentBitMatrix::open`'s own concern
|
||||
/// (see the [`Presence`](Mat::Presence) variant's docs), not decided
|
||||
/// here.
|
||||
pub(super) fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> {
|
||||
if with_counts && layer_dir.join("counts").exists() {
|
||||
return Layer::<PersistentCompactIntMatrix>::open(layer_dir, mode).map(Mat::Count);
|
||||
}
|
||||
if layer_dir.join("presence").join("is_multi.prsb").exists() {
|
||||
Layer::<PersistentSparseBitMatrix>::open(layer_dir, mode).map(Mat::SparsePresence)
|
||||
} else {
|
||||
Layer::<PersistentBitMatrix>::open(layer_dir, mode).map(Mat::Presence)
|
||||
}
|
||||
Layer::<PersistentBitMatrix>::open(layer_dir, mode).map(Mat::Presence)
|
||||
}
|
||||
|
||||
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
|
||||
match self {
|
||||
Mat::Count(l) => l.find_slot(kmer),
|
||||
Mat::Presence(l) => l.find_slot(kmer),
|
||||
Mat::SparsePresence(l) => l.find_slot(kmer),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -95,7 +88,6 @@ impl Mat {
|
||||
match self {
|
||||
Mat::Count(l) => l.index_batch(kmers),
|
||||
Mat::Presence(l) => l.index_batch(kmers),
|
||||
Mat::SparsePresence(l) => l.index_batch(kmers),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,7 +102,6 @@ impl Mat {
|
||||
match self {
|
||||
Mat::Count(l) => l.iter_minorants_batch(annex, batch_size),
|
||||
Mat::Presence(l) => l.iter_minorants_batch(annex, batch_size),
|
||||
Mat::SparsePresence(l) => l.iter_minorants_batch(annex, batch_size),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -118,7 +109,6 @@ impl Mat {
|
||||
match self {
|
||||
Mat::Count(l) => l.n_cols(),
|
||||
Mat::Presence(l) => l.n_cols(),
|
||||
Mat::SparsePresence(l) => l.n_cols(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -138,7 +128,6 @@ impl Mat {
|
||||
pub(super) fn fill_sub_matrix_carries(&self, slots: &[usize], out: &mut [Vec<bool>]) {
|
||||
match self {
|
||||
Mat::Presence(l) => l.fill_sub_matrix(slots, out),
|
||||
Mat::SparsePresence(l) => l.fill_sub_matrix(slots, out),
|
||||
Mat::Count(l) => {
|
||||
let mut counts: Vec<Vec<u32>> = out.iter().map(|_| Vec::new()).collect();
|
||||
l.fill_sub_matrix(slots, &mut counts);
|
||||
@@ -169,7 +158,11 @@ pub(super) struct PartitionCache {
|
||||
}
|
||||
|
||||
impl PartitionCache {
|
||||
pub(super) fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
|
||||
pub(super) fn build(
|
||||
partition: &KmerPartitions,
|
||||
n_parts: usize,
|
||||
with_counts: bool,
|
||||
) -> OKIResult<Self> {
|
||||
let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
|
||||
let built: Vec<(Vec<Mat>, usize)> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
@@ -182,7 +175,10 @@ impl PartitionCache {
|
||||
let meta = partition.partition_meta(part)?;
|
||||
let mut mats = Vec::with_capacity(meta.n_layers);
|
||||
for l in 0..meta.n_layers {
|
||||
let Ok(mat) = Mat::open(&partition.layer_dir(part, l), &meta.mode, with_counts) else { continue };
|
||||
let Ok(mat) = Mat::open(&partition.layer_dir(part, l), &meta.mode, with_counts)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
mats.push(mat);
|
||||
}
|
||||
pb.inc(1);
|
||||
@@ -257,7 +253,9 @@ impl PartitionCache {
|
||||
n_genomes: usize,
|
||||
mut on_hit: impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let Some(mats) = self.mats.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).
|
||||
@@ -297,7 +295,9 @@ impl PartitionCache {
|
||||
n_genomes: usize,
|
||||
mut on_hit: impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let Some(mats) = self.mats.get(dest_partition) else { return };
|
||||
let Some(mats) = self.mats.get(dest_partition) else {
|
||||
return;
|
||||
};
|
||||
|
||||
let mut by_layer: Vec<Vec<(CanonicalKmer, usize, u8)>> = vec![Vec::new(); mats.len()];
|
||||
for &(variant, family_idx, base, layer) in queries {
|
||||
@@ -311,7 +311,8 @@ impl PartitionCache {
|
||||
continue;
|
||||
}
|
||||
let mat = &mats[li];
|
||||
let variants: Vec<CanonicalKmer> = entries.iter().map(|&(variant, _, _)| variant).collect();
|
||||
let variants: Vec<CanonicalKmer> =
|
||||
entries.iter().map(|&(variant, _, _)| variant).collect();
|
||||
let slots = mat.index_batch(&variants);
|
||||
let hits: Vec<(usize, usize, u8)> = slots
|
||||
.into_iter()
|
||||
|
||||
@@ -2,15 +2,15 @@ use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::distance::RawSnpDistanceOutput;
|
||||
use super::family_scan::{scan_layer_families, Selection};
|
||||
use super::family_scan::{Selection, scan_layer_families};
|
||||
|
||||
/// See [`KmerIndex::cardinality_tally`].
|
||||
pub struct CardinalityTally {
|
||||
@@ -52,11 +52,19 @@ pub trait CardinalityExt {
|
||||
/// gets the matching restriction via `scan_family_pairs`'s new
|
||||
/// `variable` flag, rather than a `family_size()` check of its own (it
|
||||
/// doesn't have direct access to the family's mask).
|
||||
fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally>;
|
||||
fn cardinality_tally(
|
||||
&self,
|
||||
raw: &RawSnpDistanceOutput,
|
||||
ratio_ceiling: f64,
|
||||
) -> OKIResult<CardinalityTally>;
|
||||
}
|
||||
|
||||
impl CardinalityExt for KmerIndex {
|
||||
fn cardinality_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<CardinalityTally> {
|
||||
fn cardinality_tally(
|
||||
&self,
|
||||
raw: &RawSnpDistanceOutput,
|
||||
ratio_ceiling: f64,
|
||||
) -> OKIResult<CardinalityTally> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta().genomes.len();
|
||||
let with_counts = self.meta().config.with_counts;
|
||||
@@ -72,7 +80,7 @@ impl CardinalityExt for KmerIndex {
|
||||
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
|
||||
});
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
self.root_path(),
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
@@ -89,33 +97,42 @@ impl CardinalityExt for KmerIndex {
|
||||
// family comes back — no per-layer buffer.
|
||||
let mut total = [[0u64; 5]; 5];
|
||||
for layer_dir in &layer_dirs {
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| {
|
||||
if mask.family_size() < 2 {
|
||||
// Fully invariant family (never varies anywhere in
|
||||
// the index) — genome-wide background, not
|
||||
// SNP-adjacent signal; would otherwise swamp the
|
||||
// diagonal (`c=1/c=1` etc.), which needs to reflect
|
||||
// the same variable-families-only population the
|
||||
// `+ASC`-corrected alignment/likelihood actually
|
||||
// models. See `base_pair_tally`'s `variable` gate
|
||||
// on its own `same` diagonal for the matching fix.
|
||||
return;
|
||||
}
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&Selection::All,
|
||||
|_family_idx, mask, genome_mask| {
|
||||
if mask.family_size() < 2 {
|
||||
// Fully invariant family (never varies anywhere in
|
||||
// the index) — genome-wide background, not
|
||||
// SNP-adjacent signal; would otherwise swamp the
|
||||
// diagonal (`c=1/c=1` etc.), which needs to reflect
|
||||
// the same variable-families-only population the
|
||||
// `+ASC`-corrected alignment/likelihood actually
|
||||
// models. See `base_pair_tally`'s `variable` gate
|
||||
// on its own `same` diagonal for the matching fix.
|
||||
return;
|
||||
}
|
||||
|
||||
for i in 0..n_genomes {
|
||||
let card_i = genome_mask[i].count_ones() as usize;
|
||||
for j in (i + 1)..n_genomes {
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
let card_j = genome_mask[j].count_ones() as usize;
|
||||
total[card_i][card_j] += 1;
|
||||
if card_i != card_j {
|
||||
total[card_j][card_i] += 1;
|
||||
for i in 0..n_genomes {
|
||||
let card_i = genome_mask[i].count_ones() as usize;
|
||||
for j in (i + 1)..n_genomes {
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
let card_j = genome_mask[j].count_ones() as usize;
|
||||
total[card_i][card_j] += 1;
|
||||
if card_i != card_j {
|
||||
total[card_j][card_i] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})?;
|
||||
},
|
||||
)?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
@@ -2,14 +2,14 @@ use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::family_scan::{scan_layer_families, Selection};
|
||||
use super::family_scan::{Selection, scan_layer_families};
|
||||
|
||||
/// Raw p-distance restricted to loci that are single-copy in **both**
|
||||
/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility
|
||||
@@ -72,7 +72,7 @@ where
|
||||
let k = index.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
index.root_path(),
|
||||
index.kmer_size(),
|
||||
index.minimizer_size(),
|
||||
@@ -82,15 +82,23 @@ where
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = super::family_scan::sibling_layer_dirs(index)?;
|
||||
|
||||
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
|
||||
// One layer at a time — see `snp_pseudo_alignment`'s comment for why
|
||||
// `par_iter()` over layers would defeat `scan_layer_families`'s
|
||||
// partition-grouped locality.
|
||||
let mut total = zero();
|
||||
for layer_dir in &layer_dirs {
|
||||
let mut acc = zero();
|
||||
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
|
||||
// One layer at a time — see `snp_pseudo_alignment`'s comment for why
|
||||
// `par_iter()` over layers would defeat `scan_layer_families`'s
|
||||
// partition-grouped locality.
|
||||
let mut total = zero();
|
||||
for layer_dir in &layer_dirs {
|
||||
let mut acc = zero();
|
||||
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| {
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&Selection::All,
|
||||
|_family_idx, mask, genome_mask| {
|
||||
let variable = mask.family_size() >= 2;
|
||||
|
||||
// Per genome: which single form (if exactly one) it
|
||||
@@ -109,15 +117,16 @@ where
|
||||
on_pair(&mut acc, i, j, bi, bj, variable);
|
||||
}
|
||||
}
|
||||
})?;
|
||||
},
|
||||
)?;
|
||||
|
||||
total = combine(total, acc);
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(total)
|
||||
total = combine(total, acc);
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// Adds [`raw_snp_distance`](Self::raw_snp_distance) and
|
||||
/// [`base_pair_tally`](Self::base_pair_tally) to `KmerIndex`.
|
||||
@@ -141,7 +150,11 @@ pub trait DistanceExt {
|
||||
/// keeps aggregate SNP/shared counts per genome pair, not which bases
|
||||
/// were actually involved at each locus, and the ratio-ceiling filter
|
||||
/// can only be evaluated once the aggregate counts are known.
|
||||
fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally>;
|
||||
fn base_pair_tally(
|
||||
&self,
|
||||
raw: &RawSnpDistanceOutput,
|
||||
ratio_ceiling: f64,
|
||||
) -> OKIResult<BasePairTally>;
|
||||
}
|
||||
|
||||
impl DistanceExt for KmerIndex {
|
||||
@@ -150,7 +163,12 @@ impl DistanceExt for KmerIndex {
|
||||
let (snp, shared) = scan_family_pairs(
|
||||
self,
|
||||
"raw_snp_distance",
|
||||
|| (Array2::<u64>::zeros((n_genomes, n_genomes)), Array2::<u64>::zeros((n_genomes, n_genomes))),
|
||||
|| {
|
||||
(
|
||||
Array2::<u64>::zeros((n_genomes, n_genomes)),
|
||||
Array2::<u64>::zeros((n_genomes, n_genomes)),
|
||||
)
|
||||
},
|
||||
|(snp, shared), i, j, bi, bj, _variable| {
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
@@ -169,7 +187,11 @@ impl DistanceExt for KmerIndex {
|
||||
Ok(RawSnpDistanceOutput { snp, shared })
|
||||
}
|
||||
|
||||
fn base_pair_tally(&self, raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> OKIResult<BasePairTally> {
|
||||
fn base_pair_tally(
|
||||
&self,
|
||||
raw: &RawSnpDistanceOutput,
|
||||
ratio_ceiling: f64,
|
||||
) -> OKIResult<BasePairTally> {
|
||||
let n_genomes = self.meta().genomes.len();
|
||||
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
|
||||
if i == j {
|
||||
|
||||
@@ -13,14 +13,14 @@ use std::io::{BufWriter, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::cache::PartitionCache;
|
||||
use super::entropy_annex::{EntropyAnnexBuilder, ENTROPY_ANNEX_FILE_NAME};
|
||||
use super::entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnexBuilder};
|
||||
use super::family_scan::{Selection, scan_layer_families};
|
||||
use super::subsample::EntropyBias;
|
||||
|
||||
@@ -120,18 +120,28 @@ pub trait ShannonEntropyExt {
|
||||
/// `DevDocMD/architecture/siblings.md`, "Entropy-biased selection") —
|
||||
/// useful here specifically to see the biased sample's own entropy
|
||||
/// distribution, not just to speed up a distance computation.
|
||||
fn shannon_entropy_csv(&self, path: &Path, subsample: Option<usize>, entropy_bias: Option<EntropyBias>) -> OKIResult<()>;
|
||||
fn shannon_entropy_csv(
|
||||
&self,
|
||||
path: &Path,
|
||||
subsample: Option<usize>,
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<()>;
|
||||
}
|
||||
|
||||
impl ShannonEntropyExt for KmerIndex {
|
||||
fn shannon_entropy_csv(&self, path: &Path, subsample: Option<usize>, entropy_bias: Option<EntropyBias>) -> OKIResult<()> {
|
||||
fn shannon_entropy_csv(
|
||||
&self,
|
||||
path: &Path,
|
||||
subsample: Option<usize>,
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
) -> OKIResult<()> {
|
||||
let n_parts = self.n_partitions();
|
||||
let n_genomes = self.meta().genomes.len();
|
||||
let with_counts = self.meta().config.with_counts;
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
self.root_path(),
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
@@ -140,31 +150,53 @@ impl ShannonEntropyExt for KmerIndex {
|
||||
.map_err(OKIError::Partition)?;
|
||||
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
|
||||
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
|
||||
let selections = super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
|
||||
let selections =
|
||||
super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
|
||||
|
||||
let mut f = BufWriter::new(std::fs::File::create(path).map_err(OKIError::Io)?);
|
||||
writeln!(f, "layer,family_idx,entropy15,entropy4,family_size,n_genomes_present").map_err(OKIError::Io)?;
|
||||
writeln!(
|
||||
f,
|
||||
"layer,family_idx,entropy15,entropy4,family_size,n_genomes_present"
|
||||
)
|
||||
.map_err(OKIError::Io)?;
|
||||
|
||||
let pb = progress_bar("shannon_entropy", layer_dirs.len() as u64, "layers");
|
||||
for (layer_ord, (layer_dir, layer_selection)) in layer_dirs.iter().zip(selections.iter()).enumerate() {
|
||||
for (layer_ord, (layer_dir, layer_selection)) in
|
||||
layer_dirs.iter().zip(selections.iter()).enumerate()
|
||||
{
|
||||
let selection = match layer_selection {
|
||||
None => Selection::All,
|
||||
Some(set) => Selection::Some(set),
|
||||
};
|
||||
let mut write_err = None;
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |family_idx, mask, genome_mask| {
|
||||
if write_err.is_some() {
|
||||
return;
|
||||
}
|
||||
if mask.family_size() < 2 {
|
||||
return; // monomorphic family — entropy trivially 0, no signal, skip
|
||||
}
|
||||
let Some((h15, n_present)) = family_entropy(genome_mask) else { return };
|
||||
let h4 = family_entropy_4(genome_mask).map_or(0.0, |(h, _)| h);
|
||||
if let Err(e) = writeln!(f, "{layer_ord},{family_idx},{h15:.6},{h4:.6},{},{n_present}", mask.family_size()) {
|
||||
write_err = Some(e);
|
||||
}
|
||||
})?;
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&selection,
|
||||
|family_idx, mask, genome_mask| {
|
||||
if write_err.is_some() {
|
||||
return;
|
||||
}
|
||||
if mask.family_size() < 2 {
|
||||
return; // monomorphic family — entropy trivially 0, no signal, skip
|
||||
}
|
||||
let Some((h15, n_present)) = family_entropy(genome_mask) else {
|
||||
return;
|
||||
};
|
||||
let h4 = family_entropy_4(genome_mask).map_or(0.0, |(h, _)| h);
|
||||
if let Err(e) = writeln!(
|
||||
f,
|
||||
"{layer_ord},{family_idx},{h15:.6},{h4:.6},{},{n_present}",
|
||||
mask.family_size()
|
||||
) {
|
||||
write_err = Some(e);
|
||||
}
|
||||
},
|
||||
)?;
|
||||
if let Some(e) = write_err {
|
||||
return Err(OKIError::Io(e));
|
||||
}
|
||||
@@ -194,7 +226,9 @@ impl ShannonEntropyExt for KmerIndex {
|
||||
/// `scan_layer_families`'s expensive per-genome resolution for the ~98%
|
||||
/// of minorants that are monomorphic only to discard the result.
|
||||
pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf]) -> OKIResult<()> {
|
||||
let missing: Vec<usize> = layer_dirs.iter().enumerate()
|
||||
let missing: Vec<usize> = layer_dirs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, dir)| !dir.join(ENTROPY_ANNEX_FILE_NAME).exists())
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
@@ -208,7 +242,7 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf])
|
||||
let k = index.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
index.root_path(),
|
||||
index.kmer_size(),
|
||||
index.minimizer_size(),
|
||||
@@ -221,15 +255,30 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf])
|
||||
let pb = progress_bar("entropy_annex_build", missing.len() as u64, "layers");
|
||||
for &li in &missing {
|
||||
let layer_dir = &layer_dirs[li];
|
||||
let mut builder = EntropyAnnexBuilder::new(minorant_counts[li] as usize, &layer_dir.join(ENTROPY_ANNEX_FILE_NAME))
|
||||
.map_err(OKIError::Io)?;
|
||||
let mut builder = EntropyAnnexBuilder::new(
|
||||
minorant_counts[li] as usize,
|
||||
&layer_dir.join(ENTROPY_ANNEX_FILE_NAME),
|
||||
)
|
||||
.map_err(OKIError::Io)?;
|
||||
let eligible = super::subsample::non_monomorphic_selection_layer(layer_dir)?;
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::Some(&eligible), |family_idx, mask, genome_mask| {
|
||||
debug_assert!(mask.family_size() >= 2, "Selection::Some(eligible) must only visit non-monomorphic minorants");
|
||||
if let Some((h15, _)) = family_entropy(genome_mask) {
|
||||
builder.set(family_idx, h15 as f32);
|
||||
}
|
||||
})?;
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&Selection::Some(&eligible),
|
||||
|family_idx, mask, genome_mask| {
|
||||
debug_assert!(
|
||||
mask.family_size() >= 2,
|
||||
"Selection::Some(eligible) must only visit non-monomorphic minorants"
|
||||
);
|
||||
if let Some((h15, _)) = family_entropy(genome_mask) {
|
||||
builder.set(family_idx, h15 as f32);
|
||||
}
|
||||
},
|
||||
)?;
|
||||
builder.close().map_err(OKIError::Io)?;
|
||||
pb.inc(1);
|
||||
}
|
||||
|
||||
@@ -29,18 +29,18 @@ use std::sync::Arc;
|
||||
|
||||
use ndarray::Array2;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::alignment::{iupac_code, SnpAlignment};
|
||||
use super::alignment::{SnpAlignment, iupac_code};
|
||||
use super::cache::PartitionCache;
|
||||
use super::cardinality::CardinalityTally;
|
||||
use super::distance::{BasePairTally, RawSnpDistanceOutput};
|
||||
use super::family_scan::{scan_layer_families, sibling_layer_dirs, Selection};
|
||||
use super::subsample::{compute_selections, EntropyBias};
|
||||
use super::family_scan::{Selection, scan_layer_families, sibling_layer_dirs};
|
||||
use super::subsample::{EntropyBias, compute_selections};
|
||||
|
||||
/// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs,
|
||||
/// computed together from one shared, possibly-subsampled/entropy-biased
|
||||
@@ -85,7 +85,7 @@ impl SankoffBundleExt for KmerIndex {
|
||||
let k = self.kmer_size();
|
||||
let n_bits = n_parts.trailing_zeros() as usize;
|
||||
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
self.root_path(),
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
@@ -108,25 +108,34 @@ impl SankoffBundleExt for KmerIndex {
|
||||
None => Selection::All,
|
||||
Some(set) => Selection::Some(set),
|
||||
};
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, _mask, genome_mask| {
|
||||
let single_form = |g: usize| -> Option<u8> {
|
||||
let m = genome_mask[g];
|
||||
(m.count_ones() == 1).then(|| m.trailing_zeros() as u8)
|
||||
};
|
||||
for i in 0..n_genomes {
|
||||
let Some(bi) = single_form(i) else { continue };
|
||||
for j in (i + 1)..n_genomes {
|
||||
let Some(bj) = single_form(j) else { continue };
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
shared[[j, i]] += 1;
|
||||
} else {
|
||||
snp[[i, j]] += 1;
|
||||
snp[[j, i]] += 1;
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&selection,
|
||||
|_family_idx, _mask, genome_mask| {
|
||||
let single_form = |g: usize| -> Option<u8> {
|
||||
let m = genome_mask[g];
|
||||
(m.count_ones() == 1).then(|| m.trailing_zeros() as u8)
|
||||
};
|
||||
for i in 0..n_genomes {
|
||||
let Some(bi) = single_form(i) else { continue };
|
||||
for j in (i + 1)..n_genomes {
|
||||
let Some(bj) = single_form(j) else { continue };
|
||||
if bi == bj {
|
||||
shared[[i, j]] += 1;
|
||||
shared[[j, i]] += 1;
|
||||
} else {
|
||||
snp[[i, j]] += 1;
|
||||
snp[[j, i]] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})?;
|
||||
},
|
||||
)?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
@@ -168,69 +177,83 @@ impl SankoffBundleExt for KmerIndex {
|
||||
None => Selection::All,
|
||||
Some(set) => Selection::Some(set),
|
||||
};
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &selection, |_family_idx, mask, genome_mask| {
|
||||
let variable = mask.family_size() >= 2;
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&selection,
|
||||
|_family_idx, mask, genome_mask| {
|
||||
let variable = mask.family_size() >= 2;
|
||||
|
||||
// base-pair tally — same pairwise single-form resolution as pass A.
|
||||
let single_form = |g: usize| -> Option<u8> {
|
||||
let m = genome_mask[g];
|
||||
(m.count_ones() == 1).then(|| m.trailing_zeros() as u8)
|
||||
};
|
||||
for i in 0..n_genomes {
|
||||
let Some(bi) = single_form(i) else { continue };
|
||||
for j in (i + 1)..n_genomes {
|
||||
let Some(bj) = single_form(j) else { continue };
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
if bi != bj {
|
||||
bp_counts[bi as usize][bj as usize] += 1;
|
||||
bp_counts[bj as usize][bi as usize] += 1;
|
||||
} else if variable {
|
||||
bp_same[bi as usize] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cardinality tally — restricted to variable families, see
|
||||
// `CardinalityExt::cardinality_tally`'s docs for why.
|
||||
if variable {
|
||||
// base-pair tally — same pairwise single-form resolution as pass A.
|
||||
let single_form = |g: usize| -> Option<u8> {
|
||||
let m = genome_mask[g];
|
||||
(m.count_ones() == 1).then(|| m.trailing_zeros() as u8)
|
||||
};
|
||||
for i in 0..n_genomes {
|
||||
let card_i = genome_mask[i].count_ones() as usize;
|
||||
let Some(bi) = single_form(i) else { continue };
|
||||
for j in (i + 1)..n_genomes {
|
||||
let Some(bj) = single_form(j) else { continue };
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
let card_j = genome_mask[j].count_ones() as usize;
|
||||
card_counts[card_i][card_j] += 1;
|
||||
if card_i != card_j {
|
||||
card_counts[card_j][card_i] += 1;
|
||||
if bi != bj {
|
||||
bp_counts[bi as usize][bj as usize] += 1;
|
||||
bp_counts[bj as usize][bi as usize] += 1;
|
||||
} else if variable {
|
||||
bp_same[bi as usize] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pseudo-alignment — same variable-family gate as
|
||||
// `snp_pseudo_alignment`. Every family the shared selection
|
||||
// actually visits already satisfies this when the
|
||||
// selection is `Selection::Some` (only non-monomorphic
|
||||
// minorants are ever selected), but the explicit check
|
||||
// still matters under `Selection::All` (no `--subsample`),
|
||||
// which visits monomorphic minorants too.
|
||||
if variable {
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
sequences[g].push(iupac_code(m));
|
||||
// cardinality tally — restricted to variable families, see
|
||||
// `CardinalityExt::cardinality_tally`'s docs for why.
|
||||
if variable {
|
||||
for i in 0..n_genomes {
|
||||
let card_i = genome_mask[i].count_ones() as usize;
|
||||
for j in (i + 1)..n_genomes {
|
||||
if !included[[i, j]] {
|
||||
continue;
|
||||
}
|
||||
let card_j = genome_mask[j].count_ones() as usize;
|
||||
card_counts[card_i][card_j] += 1;
|
||||
if card_i != card_j {
|
||||
card_counts[card_j][card_i] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})?;
|
||||
|
||||
// pseudo-alignment — same variable-family gate as
|
||||
// `snp_pseudo_alignment`. Every family the shared selection
|
||||
// actually visits already satisfies this when the
|
||||
// selection is `Selection::Some` (only non-monomorphic
|
||||
// minorants are ever selected), but the explicit check
|
||||
// still matters under `Selection::All` (no `--subsample`),
|
||||
// which visits monomorphic minorants too.
|
||||
if variable {
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
sequences[g].push(iupac_code(m));
|
||||
}
|
||||
}
|
||||
},
|
||||
)?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Ok(SankoffBundle {
|
||||
raw,
|
||||
base_pair_tally: BasePairTally { counts: bp_counts, same: bp_same },
|
||||
cardinality_tally: CardinalityTally { counts: card_counts },
|
||||
base_pair_tally: BasePairTally {
|
||||
counts: bp_counts,
|
||||
same: bp_same,
|
||||
},
|
||||
cardinality_tally: CardinalityTally {
|
||||
counts: card_counts,
|
||||
},
|
||||
alignment: SnpAlignment { sequences },
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2,16 +2,16 @@ use std::sync::Arc;
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
use obikpartitionner::KmerPartition;
|
||||
use obikpartitionner::KmerPartitions;
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use super::ANNEX_FILE_NAME;
|
||||
use super::SiblingAnnex;
|
||||
use super::cache::PartitionCache;
|
||||
use super::family_scan::{scan_layer_families, Selection};
|
||||
use super::family_scan::{Selection, scan_layer_families};
|
||||
|
||||
/// Distribution of family sizes (1-4), read back from an already-built
|
||||
/// annex (see [`super::build::SiblingAnnexBuildExt::build_sibling_annex`])
|
||||
@@ -81,7 +81,9 @@ impl SiblingStatsExt for KmerIndex {
|
||||
|mut counts, layer_dir| -> OKIResult<[u64; 4]> {
|
||||
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
|
||||
for slot in 0..annex.len() {
|
||||
let Some(mask) = annex.get(slot) else { continue };
|
||||
let Some(mask) = annex.get(slot) else {
|
||||
continue;
|
||||
};
|
||||
if !mask.is_minorant() {
|
||||
continue; // family tallied once, at its minorant
|
||||
}
|
||||
@@ -112,7 +114,7 @@ impl SiblingStatsExt for KmerIndex {
|
||||
// Same whole-run cache as `build_sibling_annex` — see its docs for
|
||||
// why re-opening per lookup (or per call to a batching helper) is
|
||||
// not good enough on a real index.
|
||||
let partition = KmerPartition::open_with_config(
|
||||
let partition = KmerPartitions::open_with_config(
|
||||
self.root_path(),
|
||||
self.kmer_size(),
|
||||
self.minimizer_size(),
|
||||
@@ -131,18 +133,27 @@ impl SiblingStatsExt for KmerIndex {
|
||||
..Default::default()
|
||||
};
|
||||
for layer_dir in &layer_dirs {
|
||||
scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, &Selection::All, |_family_idx, mask, genome_mask| {
|
||||
// "Genome g represents this family" means g carries
|
||||
// *any* of its members, not just the minorant's own —
|
||||
// `genome_mask[g] != 0` is exactly that.
|
||||
let s = mask.siblings() as usize;
|
||||
stats.counts[s] += 1;
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
if m != 0 {
|
||||
stats.per_genome[g][s] += 1;
|
||||
scan_layer_families(
|
||||
layer_dir,
|
||||
n_parts,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
k,
|
||||
&cache,
|
||||
&Selection::All,
|
||||
|_family_idx, mask, genome_mask| {
|
||||
// "Genome g represents this family" means g carries
|
||||
// *any* of its members, not just the minorant's own —
|
||||
// `genome_mask[g] != 0` is exactly that.
|
||||
let s = mask.siblings() as usize;
|
||||
stats.counts[s] += 1;
|
||||
for (g, &m) in genome_mask.iter().enumerate() {
|
||||
if m != 0 {
|
||||
stats.per_genome[g][s] += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
})?;
|
||||
},
|
||||
)?;
|
||||
pb.inc(1);
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
|
||||
Reference in New Issue
Block a user