OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

Push zunrplorkwkt #70

Merged
coissac merged 93 commits from push-zunrplorkwkt into main 2026-08-28 23:15:38 +00:00
Conversation 0 Commits 93 Files Changed 428
+1290 -733
26 changed files with 1290 additions and 733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit f9ef6b8391 - Show all commits
DevDocMD/implementation/partition_layer_cache.md
+76 -7
View File
@@ -5,6 +5,70 @@ themselves not started; ownership split (below) still open. Preparatory
encapsulation work (partition/layer path and metadata accessors on
`KmerPartition`) landed the same day — see "Preparatory work done" below.
## Type-to-concept mapping: Index / Partition / Layer
The conceptual nesting `Index { Partition { Layer { MPHF, Evidence, Matrix
} } } }` does **not** have one Rust type per level — worth stating
explicitly, since two of the names below are misleading.
- **Index** = `obikindex::KmerIndex` — `{ root_path, meta: IndexMeta,
partition: KmerPartition }`. The field is named `partition` (singular),
but it holds the *whole* multi-partition structure below — the name
suggests "one partition", the value is all of them.
- **Partition, the collection (not one partition)** =
`obikpartitionner::KmerPartition`. Despite the singular name, this owns
*every* partition of the index: `root_path`, `n_partitions`, and
per-partition accessors that all take an explicit index `i`
(`partition_dir(i)`, `index_dir(i)`, `layer_dir(i, l)`,
`partition_meta(i)`, `n_layers(i)`, `index_mode(i)` — see the
`part_dir`/`layer_dir` and `PartitionMeta`-encapsulation work above).
It never holds one partition's layers open in memory — no `Vec<Layer<D>>`
here, only path arithmetic and metadata reads. A more honest name would
be `KmerPartitionSet` or `KmerPartitioner`; renaming is out of scope for
now, just worth knowing it's not "a partition."
- **Partition, one of them** = **no dedicated type exists today**. The
structural equivalent of "one partition, its layers held open" is
`obilayeredmap::LayeredMap<D>` — `{ root, meta: PartitionMeta, layers:
Vec<Layer<D>> }` — but `LayeredMap` itself doesn't know it's playing that
role: it operates purely on whatever `root` path it's given and has no
notion of `KmerPartition`, `n_partitions`, or a partition index `i` (see
"Layering: who owns what" below — established independently while
fixing `open_data`/`layer_dir`). Every caller that wants "one partition,
opened" builds the path itself
(`kmer_partition.index_dir(i)`/`.layer_dir(i, l)`) and either passes it
to `LayeredMap::open` or — as `obikphylo::siblings::cache` does —
bypasses `LayeredMap` entirely and opens each `Layer<D>` directly.
- **Layer** = `obilayeredmap::Layer<D>` — `{ mphf: MphfLayer, data: D }`.
- **MPHF** = `MphfLayer.mphf: MemCase<MphfEps>` — kmer → slot.
- **Evidence** = `MphfLayer.ev: LayerEvidence` (`Exact`/`Approx`/
`Hybrid` — `evidence.bin`/`fingerprint.bin`; see `EvidenceKind`).
- **Matrix** = `Layer<D>.data: D` — `PersistentBitMatrix` /
`PersistentCompactIntMatrix` / `PersistentSparseBitMatrix` / `()`.
Today's actual nesting, in Rust terms:
```
KmerIndex
└─ partition: KmerPartition (all N partitions — paths + metadata only)
└─ (opened on demand, per i) LayeredMap<D> ← "one partition", unnamed as such
└─ layers: Vec<Layer<D>>
└─ Layer<D> { mphf: MphfLayer, data: D }
├─ mphf.mphf → MPHF
├─ mphf.ev → Evidence
└─ data → Matrix
```
`obikphylo::siblings::cache::PartitionCache` — the thing (2) is meant to
generalise — skips the middle of this nesting entirely: it doesn't build a
`Vec<LayeredMap<D>>`, it builds `mats: Vec<Vec<Mat>>` directly — outer
index = partition `i` (via `KmerPartition`), inner index = layer `l` (via
`Layer<D>`/`Mat`) — reconstructing "one partition, its layers open" as a
bare nested `Vec` because no owning type for that concept exists to reuse.
That gap is exactly what (2) needs to fill, and (1) is exactly what would
let a per-partition type hold layers of mixed `D` (today `LayeredMap<D>`
can't, being monomorphic — see "The gap in `obilayeredmap`'s existing
cache" below).
## The problem
Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps
@@ -217,13 +281,18 @@ storage decisions, now a *third* copy of the same logic to keep in sync)
about the code. Lesson: for a `pub enum` whose variant list matters,
read the definition unfiltered, don't grep for the variant names you
expect to find.
- One real consequence of that same correction: `obikphylo::siblings::
cache::Mat::SparsePresence(Layer<PersistentSparseBitMatrix>)` looks
redundant now — `Mat::Presence(Layer<PersistentBitMatrix>)` alone
would already handle sparse layers transparently, since
`PersistentBitMatrix` absorbs `Sparse` internally. Likely `Mat` predates
`PersistentBitMatrix` growing native sparse support. Signalled, not
removed — no mandate to touch `obikphylo` for this.
- One real consequence of that same correction, **fixed (2026-08-20)**:
`obikphylo::siblings::cache::Mat::SparsePresence(Layer<
PersistentSparseBitMatrix>)` was redundant — `Mat::Presence(Layer<
PersistentBitMatrix>)` alone already handles sparse layers
transparently, since `PersistentBitMatrix` absorbs `Sparse` internally.
Removed the variant, the `presence/is_multi.prsb` probe in `Mat::open`
(now just opens `Layer::<PersistentBitMatrix>` unconditionally for the
non-count case — sparse-vs-dense is `PersistentBitMatrix::open`'s own
concern), and every now-single-armed match in `find_slot`/`index_batch`/
`iter_minorants_batch`/`n_cols`/`fill_sub_matrix_carries`. Full
workspace test suite green after, including the 27 `obikphylo::siblings`
tests that exercise `pack_matrices(true)`/sparse through `Mat`.
## Remaining instance of the PartitionMeta-encapsulation problem
src/obikindex/src/index.rs
+102 -40
View File
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obikpartitionner::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
use obikpartitionner::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR};
use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*;
use tracing::info;
@@ -16,7 +16,7 @@ use crate::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCAT
pub struct KmerIndex {
pub(crate) root_path: PathBuf,
pub(crate) meta: IndexMeta,
pub(crate) partition: KmerPartition,
pub(crate) partition: KmerPartitions,
}
impl KmerIndex {
@@ -31,7 +31,7 @@ impl KmerIndex {
force: bool,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
let partition = KmerPartition::create(
let partition = KmerPartitions::create(
&root_path,
config.n_bits,
config.kmer_size,
@@ -45,7 +45,11 @@ impl KmerIndex {
meta.genomes.push(info);
}
meta.write(&root_path)?;
Ok(Self { root_path, meta, partition })
Ok(Self {
root_path,
meta,
partition,
})
}
pub fn open<P: AsRef<Path>>(path: P) -> OKIResult<Self> {
@@ -53,13 +57,17 @@ impl KmerIndex {
let meta = IndexMeta::read(&root_path).map_err(OKIError::Io)?;
set_k(meta.config.kmer_size);
set_m(meta.config.minimizer_size);
let partition = KmerPartition::open_with_config(
let partition = KmerPartitions::open_with_config(
&root_path,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?;
Ok(Self { root_path, meta, partition })
Ok(Self {
root_path,
meta,
partition,
})
}
/// Return `true` if `path` contains an `index.meta` file.
@@ -96,12 +104,12 @@ impl KmerIndex {
pub(crate) fn create_skeleton<P: AsRef<Path>>(
output: P,
meta: &IndexMeta,
) -> OKIResult<KmerPartition> {
) -> OKIResult<KmerPartitions> {
let output = output.as_ref();
fs::create_dir_all(output).map_err(OKIError::Io)?;
meta.write(output).map_err(OKIError::Io)?;
fs::create_dir_all(output.join(PARTITIONS_SUBDIR)).map_err(OKIError::Io)?;
Ok(KmerPartition::open_with_config(
Ok(KmerPartitions::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
@@ -134,12 +142,24 @@ impl KmerIndex {
/// The index's root directory — needed by out-of-crate extension code
/// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto
/// the same on-disk index.
pub fn root_path(&self) -> &Path { &self.root_path }
pub fn meta(&self) -> &IndexMeta { &self.meta }
pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta }
pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size }
pub fn minimizer_size(&self) -> usize { self.meta.config.minimizer_size }
pub fn n_partitions(&self) -> usize { self.partition.n_partitions() }
pub fn root_path(&self) -> &Path {
&self.root_path
}
pub fn meta(&self) -> &IndexMeta {
&self.meta
}
pub fn meta_mut(&mut self) -> &mut IndexMeta {
&mut self.meta
}
pub fn kmer_size(&self) -> usize {
self.meta.config.kmer_size
}
pub fn minimizer_size(&self) -> usize {
self.meta.config.minimizer_size
}
pub fn n_partitions(&self) -> usize {
self.partition.n_partitions()
}
/// Number of layers per partition.
///
@@ -152,7 +172,7 @@ impl KmerIndex {
/// Expose the inner partition so the caller can run scatter into it.
/// Call `mark_scattered` once scatter is complete.
pub fn partition_mut(&mut self) -> &mut KmerPartition {
pub fn partition_mut(&mut self) -> &mut KmerPartitions {
&mut self.partition
}
@@ -174,7 +194,11 @@ impl KmerIndex {
///
/// Writes `spectrums/{label}.json` and touches `count.done` upon completion.
/// Per-partition spectrum files are removed unless `keep_intermediate` is true.
pub fn dereplicate_and_count(&self, keep_intermediate: bool, rep: &mut Reporter) -> OKIResult<()> {
pub fn dereplicate_and_count(
&self,
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let t = Stage::start("dereplicate");
self.partition.dereplicate()?;
rep.push(t.stop());
@@ -189,11 +213,17 @@ impl KmerIndex {
}
fn write_spectrum(&self, sp: &KmerSpectrum) -> OKIResult<()> {
let label = self.meta.genomes.first().map(|g| g.label.as_str()).unwrap_or("unknown");
let label = self
.meta
.genomes
.first()
.map(|g| g.label.as_str())
.unwrap_or("unknown");
let spectrums_dir = self.root_path.join("spectrums");
fs::create_dir_all(&spectrums_dir)?;
let path = spectrums_dir.join(format!("{label}.json"));
let spectrum_map: BTreeMap<String, u64> = sp.counts
let spectrum_map: BTreeMap<String, u64> = sp
.counts
.iter()
.map(|(&c, &f)| (format!("{c:010}"), f))
.collect();
@@ -226,17 +256,28 @@ impl KmerIndex {
let order: Vec<usize> = (0..n).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| self.partition.build_index_layer(i, min_ab, max_ab, with_counts, &evidence, block_bits),
|i, n_kmers, _| {
if n_kmers > 0 {
total_kmers += n_kmers;
pb.inc(1);
pb.set_message(format!("{i}: {n_kmers} kmers"));
}
},
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
self.partition.build_index_layer(
i,
min_ab,
max_ab,
with_counts,
&evidence,
block_bits,
)
},
|i, n_kmers, _| {
if n_kmers > 0 {
total_kmers += n_kmers;
pb.inc(1);
pb.set_message(format!("{i}: {n_kmers} kmers"));
}
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
info!("done — {} total kmers indexed", total_kmers);
@@ -253,7 +294,7 @@ impl KmerIndex {
}
/// Borrow the inner partition for direct superkmer-level queries.
pub fn partition(&self) -> &KmerPartition {
pub fn partition(&self) -> &KmerPartitions {
&self.partition
}
@@ -282,12 +323,14 @@ impl KmerIndex {
&order,
|i| -> OKIResult<()> {
let index_dir = self.partition.index_dir(i);
if !index_dir.exists() { return Ok(()); }
if !index_dir.exists() {
return Ok(());
}
let n_layers = self.partition.n_layers(i)?;
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let presence_dir = layer_dir.join("presence");
let counts_dir = layer_dir.join("counts");
let counts_dir = layer_dir.join("counts");
if presence_dir.exists() {
if sparse {
pack_sparse_bit_matrix(&presence_dir).map_err(OKIError::Io)?;
@@ -295,11 +338,15 @@ impl KmerIndex {
pack_bit_matrix(&presence_dir).map_err(OKIError::Io)?;
}
}
if counts_dir.exists() { pack_compact_int_matrix(&counts_dir).map_err(OKIError::Io)?; }
if counts_dir.exists() {
pack_compact_int_matrix(&counts_dir).map_err(OKIError::Io)?;
}
}
Ok(())
},
|_, _, _| { pb.inc(1); },
|_, _, _| {
pb.inc(1);
},
)?;
pb.finish_and_clear();
Ok(())
@@ -318,18 +365,27 @@ impl KmerIndex {
.into_par_iter()
.filter_map(|i| {
let index_dir = self.partition.index_dir(i);
if !index_dir.exists() { return None; }
if !index_dir.exists() {
return None;
}
let n_layers = match self.partition.n_layers(i) {
Ok(n) => n,
Err(e) => return Some(OKIError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))),
Err(e) => {
return Some(OKIError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
e.to_string(),
)));
}
};
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let meta_path = layer_dir.join(LayerMeta::FILENAME);
if meta_path.exists() { continue; }
if meta_path.exists() {
continue;
}
let unitigs_path = layer_dir.join("unitigs.bin");
let n_kmers = match UnitigFileReader::open_sequential(&unitigs_path) {
Ok(r) => r.n_kmers(),
Ok(r) => r.n_kmers(),
Err(e) => return Some(OKIError::Partition(e)),
};
if let Err(e) = LayerMeta::save(&layer_dir, n_kmers) {
@@ -340,7 +396,9 @@ impl KmerIndex {
})
.collect();
if let Some(e) = errors.into_iter().next() { return Err(e); }
if let Some(e) = errors.into_iter().next() {
return Err(e);
}
Ok(())
}
}
@@ -355,7 +413,11 @@ fn label_from_path(path: &Path) -> String {
while let Some(pos) = s.rfind('.') {
s.truncate(pos);
}
if s.is_empty() { "unknown".to_string() } else { s }
if s.is_empty() {
"unknown".to_string()
} else {
s
}
}
fn touch(path: &Path) -> Result<(), std::io::Error> {
src/obikindex/src/merge.rs
+29 -14
View File
@@ -193,7 +193,7 @@ impl KmerIndex {
let block_bits = dst.meta.config.block_bits;
// Pre-build source list once (avoid rebuilding per partition)
let srcs: Vec<(&obikpartitionner::KmerPartition, usize)> = remaining_sources
let srcs: Vec<(&obikpartitionner::KmerPartitions, usize)> = remaining_sources
.iter()
.map(|s| (&s.partition, s.meta.genomes.len()))
.collect();
@@ -221,19 +221,34 @@ impl KmerIndex {
let runner = crate::numa::PartitionRunner::new();
let mut part_stats: Vec<PartStat> = Vec::with_capacity(n_partitions);
runner.run(
&order,
|i| dst_partition.merge_partition(i, srcs, mode, n_dst_genomes, block_bits, evidence),
|i, g_len, dur| {
pb.inc(1);
debug!(
"partition {i}: done in {:.1}s — {} new kmers",
dur.as_secs_f64(),
g_len,
);
part_stats.push(PartStat { id: i, unitig_bytes: partition_sizes[i], g_len });
},
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
dst_partition.merge_partition(
i,
srcs,
mode,
n_dst_genomes,
block_bits,
evidence,
)
},
|i, g_len, dur| {
pb.inc(1);
debug!(
"partition {i}: done in {:.1}s — {} new kmers",
dur.as_secs_f64(),
g_len,
);
part_stats.push(PartStat {
id: i,
unitig_bytes: partition_sizes[i],
g_len,
});
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
src/obikindex/src/reindex.rs
+21 -10
View File
@@ -1,17 +1,17 @@
use std::fs;
use std::path::Path;
use obikpartitionner::KmerPartition;
use obikpartitionner::KmerPartitions;
use obilayeredmap::{IndexMode, layer::Layer};
use obisys::{Reporter, Stage, progress_bar};
use std::fs;
use std::path::Path;
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::state::IndexState;
const EVIDENCE_FILE: &str = "evidence.bin";
const EVIDENCE_FILE: &str = "evidence.bin";
const FINGERPRINT_FILE: &str = "fingerprint.bin";
const UNITIG_IDX_FILE: &str = "unitigs.bin.idx";
const UNITIG_IDX_FILE: &str = "unitigs.bin.idx";
fn olm_to_oki(e: obilayeredmap::OLMError) -> OKIError {
OKIError::InvalidInput(e.to_string())
@@ -48,9 +48,13 @@ impl KmerIndex {
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| reindex_partition(&self.partition, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}"))),
|_, _, _| { pb.inc(1); },
|i| {
reindex_partition(&self.partition, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}")))
},
|_, _, _| {
pb.inc(1);
},
)?;
pb.finish_and_clear();
@@ -66,11 +70,18 @@ impl KmerIndex {
}
/// Process all layers of one partition's index directory.
fn reindex_partition(partition: &KmerPartition, i: usize, target: &IndexMode, block_bits: u8) -> OKIResult<()> {
fn reindex_partition(
partition: &KmerPartitions
, i: usize
, target: &IndexMode
, block_bits: u,
8) -> OKIResult<()> {
if !partition.index_dir(i).exists() {
return Ok(());
}
let n_layers = partition.n_layers(i).map_err(|e| OKIError::InvalidInput(e.to_string()))?;
let n_layers = partition
.n_layers(i)
.map_err(|e| OKIError::InvalidInput(e.to_string()))?;
for layer_idx in 0..n_layers {
reindex_layer(&partition.layer_dir(i, layer_idx), target, block_bits)?;
}
src/obikindex/src/select.rs
+56 -22
View File
@@ -1,6 +1,6 @@
use std::path::Path;
use obikpartitionner::{KmerPartition, OutputCol};
use obikpartitionner::{KmerPartitions, OutputCol};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
@@ -35,31 +35,48 @@ impl KmerIndex {
let mut meta = IndexMeta::new(src.meta.config.clone());
meta.config.with_counts = !output_presence;
meta.genomes = specs.iter()
meta.genomes = specs
.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
info!(
"select: {} partition(s), {} source genome(s) → {} output column(s)",
n_partitions, n_src_genomes, specs.len(),
n_partitions,
n_src_genomes,
specs.len(),
);
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let src_partition = &src.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| dst_partition.select_partition(src_partition, i, specs, n_src_genomes, threshold, output_presence, false),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
dst_partition.select_partition(
src_partition,
i,
specs,
n_src_genomes,
threshold,
output_presence,
false,
)
},
|_, _, _| {
pb.inc(1);
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
@@ -82,9 +99,9 @@ impl KmerIndex {
}
let n_src_genomes = self.meta.genomes.len();
let n_partitions = self.partition.n_partitions();
let n_partitions = self.partition.n_partitions();
let src_partition = KmerPartition::open_with_config(
let src_partition = KmerPartitions::open_with_config(
&self.root_path,
self.meta.config.kmer_size,
self.meta.config.minimizer_size,
@@ -93,26 +110,43 @@ impl KmerIndex {
info!(
"select (in-place): {} partition(s), {} source genome(s) → {} output column(s)",
n_partitions, n_src_genomes, specs.len(),
n_partitions,
n_src_genomes,
specs.len(),
);
let t = Stage::start("select");
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let partition = &self.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| partition.select_partition(&src_partition, i, specs, n_src_genomes, threshold, output_presence, true),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
partition.select_partition(
&src_partition,
i,
specs,
n_src_genomes,
threshold,
output_presence,
true,
)
},
|_, _, _| {
pb.inc(1);
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
self.meta.config.with_counts = !output_presence;
self.meta.genomes = specs.iter()
self.meta.genomes = specs
.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
self.meta.write(&self.root_path)?;
src/obikmer/src/steps/scatter.rs
+22 -13
View File
@@ -1,12 +1,12 @@
use std::path::PathBuf;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Instant;
use obisys::spinner;
use obiread::NucPage;
use obikpartitionner::KmerPartition;
use obikpartitionner::KmerPartitions;
use obipipeline::{ThrottleGuard, Throttled, throttle};
use obiread::NucPage;
use obisys::spinner;
use obisys::{Reporter, Stage};
use tracing::info;
@@ -15,8 +15,8 @@ use crate::cli::PipelineData;
// ── Iterator that keeps the slot guard alive until the file is exhausted ──────
struct GuardedIter {
inner: Box<dyn Iterator<Item = NucPage> + Send>,
_guard: ThrottleGuard,
inner: Box<dyn Iterator<Item = NucPage> + Send>,
_guard: ThrottleGuard,
flat_active: Arc<AtomicU32>,
}
@@ -38,7 +38,7 @@ impl Drop for GuardedIter {
/// Run scatter: normalise → build superkmers → route to partition → close.
/// Reports the "scatter" stage to `rep`.
pub fn scatter(
kp: &mut KmerPartition,
kp: &mut KmerPartitions,
path_source: impl Iterator<Item = PathBuf> + Send + 'static,
k: usize,
level_max: usize,
@@ -52,8 +52,8 @@ pub fn scatter(
// Throttle in the source thread — never in a worker — to prevent deadlock.
let throttled = throttle(path_source, max_open);
let file_count = Arc::new(AtomicU64::new(0));
let flat_active = Arc::new(AtomicU32::new(0));
let file_count = Arc::new(AtomicU64::new(0));
let flat_active = Arc::new(AtomicU32::new(0));
let transform_active = Arc::new(AtomicU32::new(0));
let t = Stage::start("scatter");
@@ -95,7 +95,8 @@ pub fn scatter(
const ALPHA: f64 = 0.15;
for batch in pipe.apply(throttled, n_workers, 1) {
total_bases += batch.iter()
total_bases += batch
.iter()
.map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap))
.sum::<u64>();
let now = Instant::now();
@@ -107,14 +108,22 @@ pub fn scatter(
last_bases = total_bases;
let bp = total_bases as f64;
let (count_str, rate_str) = if bp >= 1e9 {
(format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6))
(
format!("{:.2} Gbp", bp / 1e9),
format!("{:.0} Mbp/s", ema_rate / 1e6),
)
} else {
(format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6))
(
format!("{:.0} Mbp", bp / 1e6),
format!("{:.0} Mbp/s", ema_rate / 1e6),
)
};
let n_files = file_count.load(Ordering::Relaxed);
let r = flat_active.load(Ordering::Relaxed);
let c = transform_active.load(Ordering::Relaxed);
pb.set_message(format!("{count_str} {rate_str} {n_files} files [R:{r} C:{c}]"));
pb.set_message(format!(
"{count_str} {rate_str} {n_files} files [R:{r} C:{c}]"
));
}
kp.write_batch(batch).unwrap_or_else(|e| {
eprintln!("error: {e}");
src/obikpartitionner/src/distance.rs
+9 -5
View File
@@ -1,11 +1,11 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obilayeredmap::{open_data, LayeredStore};
use obilayeredmap::{LayeredStore, open_data};
use obiskio::SKResult;
use crate::common::{load_meta, olm_to_sk};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
impl KmerPartition {
impl KmerPartitions {
/// Open all count matrices for partition `part`, one per layer.
/// Layers without a `counts/` directory are skipped.
pub fn count_store(&self, part: usize) -> SKResult<LayeredStore<PersistentCompactIntMatrix>> {
@@ -16,7 +16,9 @@ impl KmerPartition {
let n_layers = load_meta(&index_dir, "distance")?.n_layers;
let matrices = (0..n_layers)
.filter_map(|l| {
self.layer_dir(part, l).join("counts").exists()
self.layer_dir(part, l)
.join("counts")
.exists()
.then(|| open_data(&index_dir, l).map_err(|e| olm_to_sk(e, "distance")))
})
.collect::<SKResult<Vec<_>>>()?;
@@ -33,7 +35,9 @@ impl KmerPartition {
let n_layers = load_meta(&index_dir, "distance")?.n_layers;
let matrices = (0..n_layers)
.filter_map(|l| {
self.layer_dir(part, l).join("presence").exists()
self.layer_dir(part, l)
.join("presence")
.exists()
.then(|| open_data(&index_dir, l).map_err(|e| olm_to_sk(e, "distance")))
})
.collect::<SKResult<Vec<_>>>()?;
src/obikpartitionner/src/dump_layer.rs
+47 -20
View File
@@ -1,19 +1,22 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obiskio::{SKError, SKResult, UnitigFileReader};
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::filter::{KmerFilter, passes_all};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
OLMError::Io(e) => SKError::Io(e),
other => SKError::InvalidData { context: "dump", detail: other.to_string() },
other => SKError::InvalidData {
context: "dump",
detail: other.to_string(),
},
}
}
impl KmerPartition {
impl KmerPartitions {
/// Iterate all indexed kmers in partition `part`, calling `cb(kmer, row)` for each
/// kmer that passes every filter in `filters`.
///
@@ -43,12 +46,14 @@ impl KmerPartition {
let mut l = 0;
loop {
let layer_dir = self.layer_dir(part, l);
if !layer_dir.exists() { break; }
if !layer_dir.exists() {
break;
}
l += 1;
let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;
let counts_dir = layer_dir.join("counts");
let counts_dir = layer_dir.join("counts");
let presence_dir = layer_dir.join("presence");
let cont = if use_counts && counts_dir.exists() {
@@ -59,7 +64,9 @@ impl KmerPartition {
let row = mat.row(slot);
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(kmer, row);
if !cont { break; }
if !cont {
break;
}
}
}
}
@@ -72,7 +79,9 @@ impl KmerPartition {
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(kmer, row);
if !cont { break; }
if !cont {
break;
}
}
}
}
@@ -86,15 +95,21 @@ impl KmerPartition {
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if mphf.find(kmer).is_some() && passes_all(filters, kmer, &all_present, n_genomes) {
if mphf.find(kmer).is_some()
&& passes_all(filters, kmer, &all_present, n_genomes)
{
cont = cb(kmer, all_present.clone());
if !cont { break; }
if !cont {
break;
}
}
}
cont
};
if !cont { return Ok(false); }
if !cont {
return Ok(false);
}
}
Ok(true)
@@ -121,11 +136,13 @@ impl KmerPartition {
let mut layer = 0;
loop {
let layer_dir = self.layer_dir(part, layer);
if !layer_dir.exists() { break; }
let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
if !layer_dir.exists() {
break;
}
let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;
let counts_dir = layer_dir.join("counts");
let counts_dir = layer_dir.join("counts");
let presence_dir = layer_dir.join("presence");
let cont = if use_counts && counts_dir.exists() {
@@ -136,7 +153,9 @@ impl KmerPartition {
let row = mat.row(slot);
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(part, layer, kmer, row);
if !cont { break; }
if !cont {
break;
}
}
}
}
@@ -149,7 +168,9 @@ impl KmerPartition {
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(part, layer, kmer, row);
if !cont { break; }
if !cont {
break;
}
}
}
}
@@ -161,15 +182,21 @@ impl KmerPartition {
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if mphf.find(kmer).is_some() && passes_all(filters, kmer, &all_present, n_genomes) {
if mphf.find(kmer).is_some()
&& passes_all(filters, kmer, &all_present, n_genomes)
{
cont = cb(part, layer, kmer, all_present.clone());
if !cont { break; }
if !cont {
break;
}
}
}
cont
};
if !cont { return Ok(false); }
if !cont {
return Ok(false);
}
layer += 1;
}
src/obikpartitionner/src/index_layer.rs
+16 -18
View File
@@ -12,7 +12,7 @@ use ptr_hash::{PtrHash, bucket_fn::CubicEps, hash::Xx64};
use crate::common::olm_to_sk;
use crate::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
@@ -24,7 +24,7 @@ fn remove_if_exists(path: &std::path::Path) {
}
}
impl KmerPartition {
impl KmerPartitions {
/// Build the layered MPHF index for partition `i`.
///
/// Returns the number of canonical k-mers indexed, or 0 if the partition
@@ -93,22 +93,20 @@ impl KmerPartition {
}
}
let n_kmers = if with_counts {
let n = write_graph_as_unitigs(g, &layer_dir)?;
Layer::<PersistentCompactIntMatrix>::build(
&layer_dir,
block_bits,
mode,
|kmer| match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
_ => 1,
},
)
.map_err(|e| olm_to_sk(e, "layer build"))?;
n
} else {
materialize_layer(g, &layer_dir, block_bits, mode)?
};
let n_kmers =
if with_counts {
let n = write_graph_as_unitigs(g, &layer_dir)?;
Layer::<PersistentCompactIntMatrix>::build(&layer_dir, block_bits, mode, |kmer| {
match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
_ => 1,
}
})
.map_err(|e| olm_to_sk(e, "layer build"))?;
n
} else {
materialize_layer(g, &layer_dir, block_bits, mode)?
};
let index_dir = layer_dir.parent().expect("layer_dir has a parent");
PartitionMeta {
src/obikpartitionner/src/lib.rs
+2 -2
View File
@@ -1,7 +1,7 @@
pub mod filter;
mod common;
mod distance;
mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod index_layer;
mod kmer_sort;
@@ -13,6 +13,6 @@ mod select_layer;
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use merge_layer::MergeMode;
pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
pub use partition::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR};
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
pub use select_layer::{AggOp, OutputCol};
src/obikpartitionner/src/merge_layer/mod.rs
+129 -68
View File
@@ -12,21 +12,20 @@ use std::io;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tracing::debug;
use obipipeline::{
Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, WorkerPool,
ThrottleGuard, throttle,
make_sink, make_source, make_transform,
Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, ThrottleGuard, WorkerPool,
make_sink, make_source, make_transform, throttle,
};
use tracing::debug;
use obicompactvec::{PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder};
use obikseq::CanonicalKmer;
use obilayeredmap::{layer_dir, IndexMode, Layer, LayeredMap, MphfOnly};
use obilayeredmap::{IndexMode, Layer, LayeredMap, MphfOnly, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{ColBuilder, load_meta, olm_to_sk};
use crate::graph_pipeline::{build_graph, materialize_layer};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
mod src_layer;
@@ -58,14 +57,14 @@ impl MatrixBuilder {
fn new(mode: MergeMode, n: usize, dir: &Path) -> io::Result<Self> {
Ok(match mode {
MergeMode::Presence => MatrixBuilder::Bit(PersistentBitMatrixBuilder::new(n, dir)?),
MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::new(n, dir)?),
MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::new(n, dir)?),
})
}
fn resume(mode: MergeMode, dir: &Path) -> io::Result<Self> {
Ok(match mode {
MergeMode::Presence => MatrixBuilder::Bit(PersistentBitMatrixBuilder::resume(dir)?),
MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::resume(dir)?),
MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::resume(dir)?),
})
}
@@ -117,7 +116,11 @@ mod matrix_builder_tests {
// One source column, filled like pass 2 would.
let mut col = mb.add_col().unwrap();
match &mut col {
ColBuilder::Bit(b) => { b.set(0, true); b.set(1, false); b.set(2, true); }
ColBuilder::Bit(b) => {
b.set(0, true);
b.set(1, false);
b.set(2, true);
}
ColBuilder::Int(_) => unreachable!(),
}
col.close().unwrap();
@@ -140,7 +143,10 @@ mod matrix_builder_tests {
let mut col = mb.add_col().unwrap();
match &mut col {
ColBuilder::Int(b) => { b.set(0, 7); b.set(1, 42); }
ColBuilder::Int(b) => {
b.set(0, 7);
b.set(1, 42);
}
ColBuilder::Bit(_) => unreachable!(),
}
col.close().unwrap();
@@ -171,7 +177,11 @@ mod matrix_builder_tests {
for vals in [[true, false, true], [false, true, false]] {
let mut col = mb.add_col().unwrap();
match &mut col {
ColBuilder::Bit(b) => for (slot, v) in vals.into_iter().enumerate() { b.set(slot, v); },
ColBuilder::Bit(b) => {
for (slot, v) in vals.into_iter().enumerate() {
b.set(slot, v);
}
}
ColBuilder::Int(_) => unreachable!(),
}
col.close().unwrap();
@@ -188,7 +198,7 @@ mod matrix_builder_tests {
// ── KmerPartition::merge_partition ────────────────────────────────────────────
impl KmerPartition {
impl KmerPartitions {
/// Merge `sources` into destination partition `i`.
///
/// Each entry in `sources` is `(partition, n_genomes)` where `n_genomes` is
@@ -201,7 +211,7 @@ impl KmerPartition {
pub fn merge_partition(
&self,
i: usize,
sources: &[(&KmerPartition, usize)],
sources: &[(&KmerPartitions, usize)],
mode: MergeMode,
n_dst_genomes: usize,
block_bits: u8,
@@ -213,7 +223,8 @@ impl KmerPartition {
}
load_meta(&dst_index_dir, "merge")?; // ensure meta.json exists before LayeredMap::open
let dst_map = Arc::new(LayeredMap::<()>::open(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?);
let dst_map =
Arc::new(LayeredMap::<()>::open(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?);
let n_dst_layers = dst_map.n_layers();
let n_src_total: usize = sources.iter().map(|(_, n)| *n).sum();
@@ -221,8 +232,11 @@ impl KmerPartition {
// (all slots true — every kmer in those layers belongs to genome_0).
if n_dst_genomes == 1 && mode == MergeMode::Presence {
for l in 0..n_dst_layers {
Layer::<()>::init_presence_matrix(&layer_dir(&dst_index_dir, l), dst_map.layer(l).n())
.map_err(|e| olm_to_sk(e, "merge"))?;
Layer::<()>::init_presence_matrix(
&layer_dir(&dst_index_dir, l),
dst_map.layer(l).n(),
)
.map_err(|e| olm_to_sk(e, "merge"))?;
}
}
@@ -283,7 +297,10 @@ impl KmerPartition {
)?;
let any_new = g.len() > 0;
debug!("partition {i}: de Bruijn graph done — {} new kmers", g.len());
debug!(
"partition {i}: de Bruijn graph done — {} new kmers",
g.len()
);
// Build new layer from de Bruijn graph if there are new kmers.
let new_layer_idx = n_dst_layers;
@@ -301,11 +318,16 @@ impl KmerPartition {
let t_open = std::time::Instant::now();
let new_mphf: Option<Arc<MphfOnly>> = if any_new {
Some(Arc::new(MphfOnly::open(&new_layer_dir).map_err(|e| olm_to_sk(e, "merge"))?))
Some(Arc::new(
MphfOnly::open(&new_layer_dir).map_err(|e| olm_to_sk(e, "merge"))?,
))
} else {
None
};
debug!("partition {i}: MPHF open in {:.3}s", t_open.elapsed().as_secs_f64());
debug!(
"partition {i}: MPHF open in {:.3}s",
t_open.elapsed().as_secs_f64()
);
// ── Prepare matrix directories for the new layer ──────────────────────
// Absent columns (dst genomes) get an all-zero/false column. Source-genome
@@ -351,7 +373,10 @@ impl KmerPartition {
exist_builders.push(cols);
}
debug!("partition {i}: builders ready in {:.3}s", t_builders.elapsed().as_secs_f64());
debug!(
"partition {i}: builders ready in {:.3}s",
t_builders.elapsed().as_secs_f64()
);
// ── Pass 2: fill builders (pipeline) ─────────────────────────────────
let t_pass2 = std::time::Instant::now();
@@ -391,35 +416,41 @@ impl KmerPartition {
.into_iter()
.map(|b| Arc::new(Mutex::new(b)))
.collect();
let exist_sink: Vec<Vec<Arc<Mutex<ColBuilder>>>> = exist_locked.iter()
let exist_sink: Vec<Vec<Arc<Mutex<ColBuilder>>>> = exist_locked
.iter()
.map(|layer| layer.iter().map(Arc::clone).collect())
.collect();
let new_sink: Vec<Arc<Mutex<ColBuilder>>> = new_locked.iter().map(Arc::clone).collect();
let dst_map_t2 = Arc::clone(&dst_map);
let new_mphf_t2 = new_mphf.clone();
let dst_map_t2 = Arc::clone(&dst_map);
let new_mphf_t2 = new_mphf.clone();
let pass2_err: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
let err_cap2 = Arc::clone(&pass2_err);
let err_cap2 = Arc::clone(&pass2_err);
let capacity = 2;
let throttled_pass2 = throttle(pass2_items.into_iter(), max_open);
let pipeline2 = Pipeline::new(
make_source!(Pass2Data, throttled_pass2.map(|t| {
let (col_offset, src_n, src_layer_dir) = t.item;
(col_offset, src_n, src_layer_dir, t.guard)
}), SrcLayer),
make_source!(
Pass2Data,
throttled_pass2.map(|t| {
let (col_offset, src_n, src_layer_dir) = t.item;
(col_offset, src_n, src_layer_dir, t.guard)
}),
SrcLayer
),
vec![
Stage::Flat(Arc::new(
move |data: Pass2Data,
push: &PipelineSender<Result<Pass2Data, PipelineError>>,
delta: &PipelineSender<isize>|
{
if let Pass2Data::SrcLayer((col_offset, src_n, src_layer_dir, _guard)) = data {
push: &PipelineSender<Result<Pass2Data, PipelineError>>,
delta: &PipelineSender<isize>| {
if let Pass2Data::SrcLayer((col_offset, src_n, src_layer_dir, _guard)) =
data
{
// _guard dropped at end of block, releasing the slot.
let reader = match UnitigFileReader::open_sequential(
&src_layer_dir.join("unitigs.bin"),
) {
Ok(r) => r,
Ok(r) => r,
Err(e) => {
*err_cap2.lock().unwrap() = Some(e.to_string());
delta.send(-1).ok();
@@ -427,7 +458,7 @@ impl KmerPartition {
}
};
let src_data = match SrcLayerData::open(&src_layer_dir, mode) {
Ok(d) => Arc::new(d),
Ok(d) => Arc::new(d),
Err(e) => {
*err_cap2.lock().unwrap() = Some(e.to_string());
delta.send(-1).ok();
@@ -440,66 +471,91 @@ impl KmerPartition {
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
batch.push(kmer);
if batch.len() == BATCH {
let b = std::mem::replace(&mut batch, Vec::with_capacity(BATCH));
let b =
std::mem::replace(&mut batch, Vec::with_capacity(BATCH));
push.send(Ok(Pass2Data::RawBatch((
col_offset, src_n, Arc::clone(&src_data), b,
)))).ok();
col_offset,
src_n,
Arc::clone(&src_data),
b,
))))
.ok();
count += 1;
}
}
if !batch.is_empty() {
push.send(Ok(Pass2Data::RawBatch((
col_offset, src_n, src_data, batch,
)))).ok();
))))
.ok();
count += 1;
}
delta.send(count - 1).ok();
}
}
},
) as SharedFlatFn<Pass2Data>),
make_transform!(Pass2Data, {
move |(col_offset, src_n, src_data, kmers): (usize, usize, Arc<SrcLayerData>, Vec<CanonicalKmer>)|
-> Vec<(Option<usize>, usize, usize, u32)>
make_transform!(
Pass2Data,
{
let mut ops: Vec<(Option<usize>, usize, usize, u32)> = Vec::new();
for kmer in kmers {
let values = src_data.lookup(kmer, src_n);
if let Some((dst_layer, hit)) = dst_map_t2.query(kmer) {
for (g, val) in values.into_iter().enumerate() {
ops.push((Some(dst_layer), col_offset + g, hit.slot, val));
}
} else if let Some(ref mphf) = new_mphf_t2 {
let slot = mphf.index(kmer);
for (g, val) in values.into_iter().enumerate() {
ops.push((None, col_offset + g, slot, val));
move |(col_offset, src_n, src_data, kmers): (
usize,
usize,
Arc<SrcLayerData>,
Vec<CanonicalKmer>,
)|
-> Vec<(Option<usize>, usize, usize, u32)> {
let mut ops: Vec<(Option<usize>, usize, usize, u32)> = Vec::new();
for kmer in kmers {
let values = src_data.lookup(kmer, src_n);
if let Some((dst_layer, hit)) = dst_map_t2.query(kmer) {
for (g, val) in values.into_iter().enumerate() {
ops.push((Some(dst_layer), col_offset + g, hit.slot, val));
}
} else if let Some(ref mphf) = new_mphf_t2 {
let slot = mphf.index(kmer);
for (g, val) in values.into_iter().enumerate() {
ops.push((None, col_offset + g, slot, val));
}
}
}
ops
}
ops
}
}, RawBatch, WriteBatch),
},
RawBatch,
WriteBatch
),
],
make_sink!(Pass2Data, {
move |ops: Vec<(Option<usize>, usize, usize, u32)>| {
for (layer_opt, col, slot, val) in ops {
match layer_opt {
Some(l) => exist_sink[l][col].lock().unwrap().set_val(slot, val),
None => new_sink[col].lock().unwrap().set_val(slot, val),
make_sink!(
Pass2Data,
{
move |ops: Vec<(Option<usize>, usize, usize, u32)>| {
for (layer_opt, col, slot, val) in ops {
match layer_opt {
Some(l) => exist_sink[l][col].lock().unwrap().set_val(slot, val),
None => new_sink[col].lock().unwrap().set_val(slot, val),
}
}
}
}
}, WriteBatch),
},
WriteBatch
),
);
WorkerPool::new(pipeline2, n_workers, capacity).run();
debug!("partition {i}: pass2 pipeline done in {:.3}s", t_pass2.elapsed().as_secs_f64());
debug!(
"partition {i}: pass2 pipeline done in {:.3}s",
t_pass2.elapsed().as_secs_f64()
);
if let Some(msg) = Arc::try_unwrap(pass2_err)
.unwrap_or_else(|_| panic!("pass2: pass2_err not uniquely owned"))
.into_inner()
.unwrap_or_else(|e| e.into_inner())
{
return Err(SKError::InvalidData { context: "merge pass2", detail: msg });
return Err(SKError::InvalidData {
context: "merge pass2",
detail: msg,
});
}
let t_close = std::time::Instant::now();
@@ -527,10 +583,15 @@ impl KmerPartition {
let mut part_meta = self.partition_meta(i)?;
part_meta.n_layers = new_layer_idx + 1;
part_meta.save(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?;
part_meta
.save(&dst_index_dir)
.map_err(|e| olm_to_sk(e, "merge"))?;
}
debug!("partition {i}: builders closed in {:.3}s", t_close.elapsed().as_secs_f64());
debug!(
"partition {i}: builders closed in {:.3}s",
t_close.elapsed().as_secs_f64()
);
Ok(n_new)
}
src/obikpartitionner/src/partition/kmer_partition.rs
+9 -7
View File
@@ -7,9 +7,9 @@ use std::time::Instant;
use obisys::progress_bar;
use obikseq::RoutableSuperKmer;
use obiskio::SKResult;
use obilayeredmap::IndexMode;
use obilayeredmap::meta::PartitionMeta;
use obiskio::SKResult;
use rayon::prelude::*;
use remove_dir_all::remove_dir_all;
use sysinfo::System;
@@ -33,12 +33,12 @@ use super::{PARTITIONS_SUBDIR, SK_EXT};
const INDEX_SUBDIR: &str = "index";
pub struct KmerSpectrum {
pub f0: u64,
pub f1: u64,
pub f0: u64,
pub f1: u64,
pub counts: BTreeMap<u32, u64>,
}
pub struct KmerPartition {
pub struct KmerPartitions {
root_path: PathBuf,
n_partitions: usize,
partitions_mask: u64,
@@ -49,7 +49,7 @@ pub struct KmerPartition {
closed: bool,
}
impl KmerPartition {
impl KmerPartitions {
pub fn create<P: AsRef<Path>>(
path: P,
n_bits: usize,
@@ -179,7 +179,9 @@ impl KmerPartition {
/// Path of partition `i` directory.
pub fn partition_dir(&self, i: usize) -> PathBuf {
self.root_path.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
self.root_path
.join(PARTITIONS_SUBDIR)
.join(format!("part_{i:05}"))
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
@@ -380,7 +382,7 @@ impl KmerPartition {
}
}
impl Drop for KmerPartition {
impl Drop for KmerPartitions {
fn drop(&mut self) {
let _ = self.close();
}
src/obikpartitionner/src/partition/mod.rs
+1 -1
View File
@@ -13,7 +13,7 @@ mod kmer_partition;
#[cfg(test)]
mod tests;
pub use kmer_partition::{KmerPartition, KmerSpectrum};
pub use kmer_partition::{KmerPartitions, KmerSpectrum};
const SK_EXT: &str = "skmer.zst";
pub const PARTITIONS_SUBDIR: &str = "partitions";
src/obikpartitionner/src/partition/tests.rs
+11 -10
View File
@@ -6,7 +6,7 @@ use obikseq::SuperKmer;
use obiskbuilder::build_superkmers;
use super::count::count_partition;
use super::{KmerPartition, PARTITIONS_SUBDIR};
use super::{KmerPartitions, PARTITIONS_SUBDIR};
const K: usize = 11;
const M: usize = 5;
@@ -46,7 +46,7 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);
let dir = tempfile::tempdir().unwrap();
let mut kp = KmerPartition::create(dir.path(), 0, K, M, true).unwrap();
let mut kp = KmerPartitions::create(dir.path(), 0, K, M, true).unwrap();
kp.write_batch(superkmers).unwrap();
kp.close().unwrap();
kp.dereplicate().unwrap();
@@ -58,9 +58,9 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
}
count_partition(&part_dir, &dedup_path, 1 << 20).unwrap();
let spec: serde_json::Value = serde_json::from_reader(
fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap(),
).unwrap();
let spec: serde_json::Value =
serde_json::from_reader(fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap())
.unwrap();
let f0 = spec["f0"].as_u64().unwrap_or(0);
let f1 = spec["f1"].as_u64().unwrap_or(0);
(f0, f1)
@@ -77,10 +77,7 @@ fn single_sequence_f0_f1_match() {
#[test]
fn two_sequences_f0_f1_match() {
let seqs: &[&[u8]] = &[
b"ACGTACGTACGTACGTACGT",
b"TGCATGCATGCATGCATGCA",
];
let seqs: &[&[u8]] = &[b"ACGTACGTACGTACGTACGT", b"TGCATGCATGCATGCATGCA"];
let (ef0, ef1) = direct_counts(seqs);
let (gf0, gf1) = pipeline_counts(seqs);
assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
@@ -103,7 +100,11 @@ fn many_sequences_f0_f1_match() {
// multiple minimizer boundaries per sequence.
let bases = b"ACGT";
let seqs: Vec<Vec<u8>> = (0..20u32)
.map(|i| (0..40).map(|j| bases[((i * 7 + j * 3) % 4) as usize]).collect())
.map(|i| {
(0..40)
.map(|j| bases[((i * 7 + j * 3) % 4) as usize])
.collect()
})
.collect();
let seq_refs: Vec<&[u8]> = seqs.iter().map(|v| v.as_slice()).collect();
let (ef0, ef1) = direct_counts(&seq_refs);
src/obikpartitionner/src/query_layer.rs
+13 -7
View File
@@ -3,15 +3,18 @@ use std::path::Path;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obiskio::{SKError, SKResult};
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
OLMError::Io(io_err) => SKError::Io(io_err),
other => SKError::InvalidData { context: "query", detail: other.to_string() },
other => SKError::InvalidData {
context: "query",
detail: other.to_string(),
},
}
}
@@ -24,8 +27,8 @@ enum QueryLayer {
impl QueryLayer {
fn open(layer_dir: &Path, with_counts: bool, mode: &IndexMode) -> SKResult<Self> {
let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_sk)?;
let counts_dir = layer_dir.join("counts");
let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_sk)?;
let counts_dir = layer_dir.join("counts");
let presence_dir = layer_dir.join("presence");
if with_counts && counts_dir.exists() {
@@ -70,7 +73,10 @@ impl QueryLayer {
/// slot)` `col_value` point lookup, which this layer's `Sparse`
/// presence matrices paid for badly: each such lookup rebuilt the
/// entire row just to return one cell.
fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
fn nonzero_iter<'a>(
&'a self,
slots: &'a [usize],
) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
match self {
QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots),
QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)),
@@ -137,7 +143,7 @@ pub enum QueryHit<'a> {
// ── KmerPartition::query_partition_with ──────────────────────────────────────
impl KmerPartition {
impl KmerPartitions {
/// Query a single partition for a pre-deduplicated map of canonical
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
///
src/obikpartitionner/src/rebuild_layer.rs
+38 -19
View File
@@ -1,27 +1,29 @@
use std::path::Path;
use obicompactvec::{
FilterMask, eval_filter_mask,
PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder,
FilterMask, PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder, eval_filter_mask,
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{layer_dir, IndexMode, MphfLayer};
use obilayeredmap::{IndexMode, MphfLayer, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{load_meta, olm_to_sk};
use crate::filter::KmerFilter;
use crate::graph_pipeline::materialize_layer;
use crate::merge_layer::{MergeMode, SrcLayerData};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
// ── Builders — pair matrix builder + column builders for one mode ─────────────
enum Builders {
Presence(PersistentBitMatrixBuilder, Vec<PersistentBitVecBuilder>),
Count(PersistentCompactIntMatrixBuilder, Vec<PersistentCompactIntVecBuilder>),
Count(
PersistentCompactIntMatrixBuilder,
Vec<PersistentCompactIntVecBuilder>,
),
}
impl Builders {
@@ -30,13 +32,18 @@ impl Builders {
MergeMode::Presence => {
let mut mat = PersistentBitMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut cols = Vec::with_capacity(n_genomes);
for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); }
for _ in 0..n_genomes {
cols.push(mat.add_col().map_err(SKError::Io)?);
}
Ok(Builders::Presence(mat, cols))
}
MergeMode::Count => {
let mut mat = PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut mat =
PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut cols = Vec::with_capacity(n_genomes);
for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); }
for _ in 0..n_genomes {
cols.push(mat.add_col().map_err(SKError::Io)?);
}
Ok(Builders::Count(mat, cols))
}
}
@@ -45,18 +52,22 @@ impl Builders {
fn set_val(&mut self, col: usize, slot: usize, value: u32) {
match self {
Builders::Presence(_, cols) => cols[col].set(slot, value > 0),
Builders::Count(_, cols) => cols[col].set(slot, value),
Builders::Count(_, cols) => cols[col].set(slot, value),
}
}
fn close(self) -> SKResult<()> {
match self {
Builders::Presence(mat, cols) => {
for b in cols { b.close().map_err(SKError::Io)?; }
for b in cols {
b.close().map_err(SKError::Io)?;
}
mat.close().map_err(SKError::Io)
}
Builders::Count(mat, cols) => {
for b in cols { b.close().map_err(SKError::Io)?; }
for b in cols {
b.close().map_err(SKError::Io)?;
}
mat.close().map_err(SKError::Io)
}
}
@@ -112,7 +123,9 @@ fn iter_src_kmers_masked(
for l in 0..src_meta.n_layers {
let src_layer_dir = layer_dir(src_index_dir, l);
let unitigs_path = src_layer_dir.join("unitigs.bin");
if !unitigs_path.exists() { continue; }
if !unitigs_path.exists() {
continue;
}
let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
@@ -127,7 +140,9 @@ fn iter_src_kmers_masked(
filters.iter().all(|f| f.passes(kmer, &row, n_genomes))
}
};
if passes { cb(kmer); }
if passes {
cb(kmer);
}
}
}
Ok(())
@@ -149,7 +164,9 @@ fn iter_src_layers(
for l in 0..src_meta.n_layers {
let src_layer_dir = layer_dir(src_index_dir, l);
let unitigs_path = src_layer_dir.join("unitigs.bin");
if !unitigs_path.exists() { continue; }
if !unitigs_path.exists() {
continue;
}
let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
@@ -158,7 +175,9 @@ fn iter_src_layers(
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
let slot = src_data.slot(kmer);
if let Some(ref m) = mask {
if !m.get(slot) { continue; }
if !m.get(slot) {
continue;
}
let row = src_data.fill_row_by_slot(slot, n_genomes);
cb(kmer, row.into_boxed_slice());
} else {
@@ -174,7 +193,7 @@ fn iter_src_layers(
// ── KmerPartition::rebuild_partition ─────────────────────────────────────────
impl KmerPartition {
impl KmerPartitions {
/// Rebuild partition `i` from `src` into `self` (an empty destination partition).
///
/// Only k-mers whose per-genome row passes all `filters` are written.
@@ -184,7 +203,7 @@ impl KmerPartition {
/// `n_genomes` is the number of genome columns in the source (and destination).
pub fn rebuild_partition(
&self,
src: &KmerPartition,
src: &KmerPartitions,
i: usize,
filters: &[Box<dyn KmerFilter>],
mode: MergeMode,
@@ -222,7 +241,7 @@ impl KmerPartition {
// ── Prepare matrix builders (one column per genome) ───────────────────
let data_dir = match mode {
MergeMode::Presence => dst_layer_dir.join("presence"),
MergeMode::Count => dst_layer_dir.join("counts"),
MergeMode::Count => dst_layer_dir.join("counts"),
};
std::fs::create_dir_all(&data_dir)?;
let mut builders = Builders::new(mode, n_new, &data_dir, n_genomes)?;
src/obikpartitionner/src/select_layer.rs
+115 -47
View File
@@ -3,14 +3,13 @@ use std::io;
use std::path::{Path, PathBuf};
use obicompactvec::{
ColGroup, MatrixGroupOps,
PersistentBitMatrix, PersistentBitMatrixBuilder,
ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use obilayeredmap::OLMError;
use obiskio::{SKError, SKResult};
use crate::partition::KmerPartition;
use crate::partition::KmerPartitions;
// ── AggOp ─────────────────────────────────────────────────────────────────────
@@ -33,9 +32,9 @@ impl AggOp {
// ── OutputCol ─────────────────────────────────────────────────────────────────
pub struct OutputCol {
pub label: String,
pub label: String,
pub indices: Vec<usize>,
pub op: AggOp,
pub op: AggOp,
}
// ── Helpers ───────────────────────────────────────────────────────────────────
@@ -43,7 +42,10 @@ pub struct OutputCol {
fn olm_to_sk(e: OLMError) -> SKError {
match e {
OLMError::Io(e) => SKError::Io(e),
other => SKError::InvalidData { context: "select", detail: other.to_string() },
other => SKError::InvalidData {
context: "select",
detail: other.to_string(),
},
}
}
@@ -77,23 +79,43 @@ fn fill_builders(
if output_presence {
let b = dst_bit.as_deref_mut().unwrap();
match spec.op {
AggOp::Any => b.add_col_from (&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?),
AggOp::All => b.add_col_from (&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?),
AggOp::None => b.add_col_from (&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?),
AggOp::Sum => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
AggOp::Any => {
b.add_col_from(&mat.partial_group_any(&g, threshold).map_err(SKError::Io)?)
}
AggOp::All => {
b.add_col_from(&mat.partial_group_all(&g, threshold).map_err(SKError::Io)?)
}
AggOp::None => {
b.add_col_from(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?)
}
AggOp::Sum => {
b.add_col_from_int(&mat.partial_group_sum(&g).map_err(SKError::Io)?)
}
AggOp::Min => {
b.add_col_from_int(&mat.partial_group_min(&g).map_err(SKError::Io)?)
}
AggOp::Max => {
b.add_col_from_int(&mat.partial_group_max(&g).map_err(SKError::Io)?)
}
}
.map_err(SKError::Io)?;
} else {
let b = dst_int.as_deref_mut().unwrap();
match spec.op {
AggOp::Sum => b.add_col_from (&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from (&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from (&mat.partial_group_max (&g).map_err(SKError::Io)?),
AggOp::Any => b.add_col_from_bit(&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?),
AggOp::All => b.add_col_from_bit(&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?),
AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(SKError::Io)?),
AggOp::Any => b.add_col_from_bit(
&mat.partial_group_any(&g, threshold).map_err(SKError::Io)?,
),
AggOp::All => b.add_col_from_bit(
&mat.partial_group_all(&g, threshold).map_err(SKError::Io)?,
),
AggOp::None => b.add_col_from_bit(
&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?,
),
}
.map_err(SKError::Io)?;
}
}
} else {
@@ -103,23 +125,43 @@ fn fill_builders(
if output_presence {
let b = dst_bit.as_deref_mut().unwrap();
match spec.op {
AggOp::Any => b.add_col_from (&mat.partial_group_any (&g, 1).map_err(SKError::Io)?),
AggOp::All => b.add_col_from (&mat.partial_group_all (&g, 1).map_err(SKError::Io)?),
AggOp::None => b.add_col_from (&mat.partial_group_none(&g, 1).map_err(SKError::Io)?),
AggOp::Sum => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
AggOp::Any => {
b.add_col_from(&mat.partial_group_any(&g, 1).map_err(SKError::Io)?)
}
AggOp::All => {
b.add_col_from(&mat.partial_group_all(&g, 1).map_err(SKError::Io)?)
}
AggOp::None => {
b.add_col_from(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?)
}
AggOp::Sum => {
b.add_col_from_int(&mat.partial_group_sum(&g).map_err(SKError::Io)?)
}
AggOp::Min => {
b.add_col_from_int(&mat.partial_group_min(&g).map_err(SKError::Io)?)
}
AggOp::Max => {
b.add_col_from_int(&mat.partial_group_max(&g).map_err(SKError::Io)?)
}
}
.map_err(SKError::Io)?;
} else {
let b = dst_int.as_deref_mut().unwrap();
match spec.op {
AggOp::Sum => b.add_col_from (&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from (&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from (&mat.partial_group_max (&g).map_err(SKError::Io)?),
AggOp::Any => b.add_col_from_bit(&mat.partial_group_any (&g, 1).map_err(SKError::Io)?),
AggOp::All => b.add_col_from_bit(&mat.partial_group_all (&g, 1).map_err(SKError::Io)?),
AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(SKError::Io)?),
AggOp::Any => {
b.add_col_from_bit(&mat.partial_group_any(&g, 1).map_err(SKError::Io)?)
}
AggOp::All => {
b.add_col_from_bit(&mat.partial_group_all(&g, 1).map_err(SKError::Io)?)
}
AggOp::None => {
b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?)
}
}
.map_err(SKError::Io)?;
}
}
}
@@ -128,7 +170,7 @@ fn fill_builders(
// ── KmerPartition::select_partition ──────────────────────────────────────────
impl KmerPartition {
impl KmerPartitions {
/// Rewrite the data matrices of partition `i` in `src` into `self`.
///
/// `specs` defines the output columns (projection/aggregation).
@@ -136,7 +178,7 @@ impl KmerPartition {
/// `in_place` — `self` and `src` share the same root; write to temp dirs then swap.
pub fn select_partition(
&self,
src: &KmerPartition,
src: &KmerPartitions,
i: usize,
specs: &[OutputCol],
_n_src_genomes: usize,
@@ -159,23 +201,33 @@ impl KmerPartition {
fs::create_dir_all(&dst_index_dir)?;
}
let data_subdir = if output_presence { "presence" } else { "counts" };
let data_subdir = if output_presence {
"presence"
} else {
"counts"
};
for l in 0..n_src_layers {
let src_layer_dir = src.layer_dir(i, l);
if !src_layer_dir.exists() { continue; }
if !src_layer_dir.exists() {
continue;
}
let dst_layer_dir = self.layer_dir(i, l);
let counts_dir = src_layer_dir.join("counts");
let counts_dir = src_layer_dir.join("counts");
let presence_dir = src_layer_dir.join("presence");
let src_is_count = counts_dir.exists() && !presence_dir.exists();
// Determine number of slots and detect implicit layers.
let n = if counts_dir.exists() {
PersistentCompactIntMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n()
PersistentCompactIntMatrix::open(&src_layer_dir)
.map_err(SKError::Io)?
.n()
} else if presence_dir.exists() {
PersistentBitMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n()
PersistentBitMatrix::open(&src_layer_dir)
.map_err(SKError::Io)?
.n()
} else {
// Implicit single-genome layer: no data matrix needed in output either.
if !in_place {
@@ -187,7 +239,7 @@ impl KmerPartition {
// Choose the output data directory (temp name for in-place).
let (dst_data_dir, final_data_dir): (PathBuf, PathBuf) = if in_place {
let tmp = dst_layer_dir.join(format!("{data_subdir}_new"));
let tmp = dst_layer_dir.join(format!("{data_subdir}_new"));
let perm = dst_layer_dir.join(data_subdir);
(tmp, perm)
} else {
@@ -202,14 +254,28 @@ impl KmerPartition {
fs::create_dir_all(&dst_data_dir)?;
let (mut dst_bit, mut dst_int) = if output_presence {
(Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?), None)
(
Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?),
None,
)
} else {
(None, Some(PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?))
(
None,
Some(
PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir)
.map_err(SKError::Io)?,
),
)
};
fill_builders(
specs, &src_layer_dir, src_is_count, threshold, output_presence,
dst_bit.as_mut(), dst_int.as_mut(),
specs,
&src_layer_dir,
src_is_count,
threshold,
output_presence,
dst_bit.as_mut(),
dst_int.as_mut(),
)?;
if output_presence {
@@ -233,7 +299,9 @@ impl KmerPartition {
}
if !in_place {
src.partition_meta(i)?.save(&dst_index_dir).map_err(olm_to_sk)?;
src.partition_meta(i)?
.save(&dst_index_dir)
.map_err(olm_to_sk)?;
}
Ok(())
src/obikpartitionner/src/tests/query_layer.rs
+4 -4
View File
@@ -44,8 +44,8 @@ fn query_stats_default_is_zero() {
#[test]
fn query_partition_with_missing_index_dir_returns_default_stats() {
let tmp = tempfile::tempdir().expect("tempdir");
let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false)
.expect("create partition");
let partition =
KmerPartitions::create(tmp.path().join("idx"), 2, 21, 9, false.expect("create partition");
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
// Any well-formed canonical k-mer works here — the call must return
@@ -65,8 +65,8 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
#[test]
fn query_partition_with_empty_kmers_is_a_noop() {
let tmp = tempfile::tempdir().expect("tempdir");
let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false)
.expect("create partition");
let partition =
KmerPartitions::create(tmp.path().join("idx"), 2, 21, 9, false.expect("create partition");
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
let stats = partition
src/obikphylo/src/siblings/alignment.rs
+33 -15
View File
@@ -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();
src/obikphylo/src/siblings/build.rs
+225 -195
View File
@@ -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)
}
src/obikphylo/src/siblings/cache.rs
+41 -40
View File
@@ -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()
src/obikphylo/src/siblings/cardinality.rs
+47 -30
View File
@@ -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();
src/obikphylo/src/siblings/distance.rs
+44 -22
View File
@@ -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 {
src/obikphylo/src/siblings/entropy.rs
+81 -32
View File
@@ -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);
}
src/obikphylo/src/siblings/sankoff_bundle.rs
+92 -69
View File
@@ -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 },
})
}
src/obikphylo/src/siblings/stats.rs
+27 -16
View File
@@ -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();