From 878b63566fae5c0742ad800f921eb4080eb6a418 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 21 Aug 2026 05:33:23 +0200 Subject: [PATCH] Extract k-mer counting logic into a dedicated counter module Decouple k-mer counting from the partitioner by introducing a new `Counter` struct. The module exposes a fluent builder API with optional partial file retention, executes partition processing in parallel via Rayon with memory-aware chunk sizing, and integrates thread-safe progress callbacks. Update all callers to use the new counter, simplify test pipelines by removing serialization overhead, and clarify algorithm separation in module documentation. --- .gitignore | 1 + .../implementation/partition_layer_cache.md | 81 ++++++++++- scripts/smoke_test_index.sh | 128 +++++++++++++++++ .../{partitionner => counter}/count.rs | 0 .../{partitionner => counter}/kmer_sort.rs | 0 src/obikindexer/src/algorithms/counter/mod.rs | 133 ++++++++++++++++++ src/obikindexer/src/algorithms/mod.rs | 6 +- .../src/algorithms/partitionner/mod.rs | 18 +-- .../src/algorithms/partitionner/router.rs | 101 +------------ .../src/algorithms/partitionner/tests.rs | 15 +- src/obikmer/src/cmd/index/mod.rs | 17 ++- src/obikphylo/src/siblings/tests.rs | 5 +- 12 files changed, 376 insertions(+), 129 deletions(-) create mode 100755 scripts/smoke_test_index.sh rename src/obikindexer/src/algorithms/{partitionner => counter}/count.rs (100%) rename src/obikindexer/src/algorithms/{partitionner => counter}/kmer_sort.rs (100%) create mode 100644 src/obikindexer/src/algorithms/counter/mod.rs diff --git a/.gitignore b/.gitignore index f349fa3a..8834f9ce 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ .kilo/ .serena/ .zed/ +.ast-cache/ memory/ sandbox/ src/target diff --git a/DevDocMD/implementation/partition_layer_cache.md b/DevDocMD/implementation/partition_layer_cache.md index 17f35a3d..74a1286a 100644 --- a/DevDocMD/implementation/partition_layer_cache.md +++ b/DevDocMD/implementation/partition_layer_cache.md @@ -51,7 +51,10 @@ turned out to be bigger than `KmerPartition` alone: `Layer`'s own constructors don't self-name either. Full redesign of both, agreed in detail, session ended (budget) before implementation — see "(5) design agreed" below; **read it before touching `KmerPartition`/`Layer` -signatures**, the shape is fully specified. Earlier mix-up, for +signatures**, the shape is fully specified. (6) done — `Counter`, a third +algorithm, extracted from `PartitionRouter` the same way `Dereplicator` +was in (4); (5) itself still not implemented, still first on the "order of +remaining work" list — see "(6) done" below. Earlier mix-up, for context: an earlier version of this doc used the name `KmerPartition` (singular) for what was actually the *collection* type (later renamed `KmerPartitions`, later @@ -772,6 +775,82 @@ itself or a new crate. examples were judged not enough to be sure of the shape (`Fn+Sync` vs `FnMut` callback bound already diverged between the two that exist). +## (6) done (2026-08-21): `Counter` — third algorithm, extracted the same way as `Dereplicator` + +Between (5) and this, the user did a session of their own crate +restructuring (see the two "Superseded" notes at the top of this file): +`obikpartition`/`obilayeredmap` folded into `obikindex` as submodules +(`obikindex::partition`, `obikindex::layer`), and `obikpartitionner`/ +`obikderep` merged into one sibling crate, `obikindexer`, holding +`obikindexer::algorithms::{partitionner, dereplicator}`. (5)'s design +(`Layer`/`KmerPartition` self-naming by number, `PartitionRouter`'s +`&mut` → `&` fix) was **not** part of that — pure crate/module packaging, +confirmed by reading the actual code (`Layer::open`/`create` still take an +external `dir: &Path`, `KmerPartition` still eagerly opens all layers, +`PartitionRouter` still holds `&mut KmerIndex`). (5) remains exactly as +specified, not yet implemented. + +This step: `count_kmer` (still living on `PartitionRouter`, per (4)'s own +"still not done" note) extracted into `obikindexer::algorithms::counter:: +Counter`, mirroring `Dereplicator` exactly — third data point for the +eventual `obikalgorithm` trait, still not extracted (still only 3 examples +with 2 different callback bounds; holding off per (5)'s "order of +remaining work"). + +```rust +pub struct Counter<'a> { + index: &'a KmerIndex, + n_partitions: usize, + keep_partial: bool, +} + +impl<'a> Counter<'a> { + pub fn new(index: &'a KmerIndex) -> Self; + pub fn keep_partial(mut self, v: bool) -> Self; // setter, mirrors PartitionRouter's style; defaults to false + pub fn run(&self, on_progress: Option) -> SKResult; +} +``` + +Same shape as `Dereplicator` throughout: `Fn(Progress) + Sync` (not +`FnMut`) since counting is also a parallel `par_iter` over partitions, an +`AtomicU64` position counter incremented from inside the parallel closure +so progress reports arrive in real time rather than bursting at the end +once `.collect()` finishes, `total: Some(n_partitions)` (known up front). +`KmerSpectrum` (the `{f0, f1, counts}` aggregate) moved from +`partitionner::router` to `counter`, since it's `Counter::run`'s return +value now, not `PartitionRouter`'s. `count.rs`/`kmer_sort.rs` moved +verbatim from `partitionner/` to `counter/` (unchanged bodies — only +`count_kmer` itself, `KmerSpectrum`, and the imports they pulled in were +removed from `router.rs`). + +One divergence from `Dereplicator`: a `keep_partial` setter exists (no +equivalent on `Dereplicator`, which has no setters at all) — a real, +already-present parameter (`keep_intermediate` at the CLI), not a +speculative addition. + +`count_kmer`'s three former callers (`obikmer::cmd::index`, `obikphylo`'s +test harness, `obikindexer::algorithms::partitionner`'s own +`pipeline_counts` test helper) all updated to `Counter::new(&idx). +run(...)` — the last one simplified further: it used to read back +`kmer_spectrum_raw.json` from disk after calling `count_partition` +directly (white-box), now it just uses the `KmerSpectrum` `Counter::run` +already returns. + +Full workspace suite green (`cargo check --workspace --all-targets` + +`cargo test --workspace`, exit code 0), plus an end-to-end CLI smoke test +against real FASTA data (scatter → dereplicate → count → index-build → +query) — required every time per (3)'s lesson, and it earned its keep +again: the very first smoke-test query returned zero matches, which +looked like a regression until traced to the query sequence itself being +low-complexity ("GGCCCCCCACG", six same-base runs) and rejected by +*query's own* default entropy threshold — nothing to do with this change. +Re-tested with a different substring, confirmed working (kmer found, +count matched the index). + +Still not done: (5) (`Layer`/`KmerPartition` redesign, `PartitionRouter`'s +`&mut`→`&`), the future cache crate, `build_layers` (still a `KmerIndex` +inherent method, not an algorithm), and `obikalgorithm` itself. + ## The problem Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps diff --git a/scripts/smoke_test_index.sh b/scripts/smoke_test_index.sh new file mode 100755 index 00000000..9eb47ca3 --- /dev/null +++ b/scripts/smoke_test_index.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +# smoke_test_index.sh — end-to-end smoke test of `obikmer index` + `obikmer query` +# +# The unit/integration test suite never exercises obikmer's real, +# file-driven CLI path (every test builds indexes by calling library APIs +# directly — see DevDocMD/implementation/partition_layer_cache.md, "(3) +# done"/"(6) done": a green `cargo test --workspace` has twice missed a +# real bug in this exact path). This script is the fast, repeatable +# substitute for hand-rolling that check each time. +# +# What it does: +# 1. builds `obikmer` (debug, via `cargo run`) +# 2. generates a small deterministic random FASTA +# 3. runs `obikmer index` on it +# 4. picks a real k-mer from the source sequence, avoiding low-complexity +# substrings (a homopolymer run tripped up an earlier manual run of +# this check — it gets rejected by query's own entropy filter, which +# looks like a bug but isn't one) +# 5. runs `obikmer query` and checks the k-mer round-trips +# 6. prints total-kmers-indexed and a clear PASS/FAIL, exit code matches +# +# Usage: +# scripts/smoke_test_index.sh [-k KMER_SIZE] [-m MINIMIZER_SIZE] [-p PARTITIONS] [--keep] +# +# --keep leaves the temp directory in place (path printed) instead of +# deleting it on exit, for manual inspection of a failure. + +set -euo pipefail + +K=11 +M=5 +PARTITIONS=4 +KEEP=0 + +while [ $# -gt 0 ]; do + case "$1" in + -k) K="$2"; shift 2 ;; + -m) M="$2"; shift 2 ;; + -p) PARTITIONS="$2"; shift 2 ;; + --keep) KEEP=1; shift ;; + *) echo "unknown argument: $1" >&2; exit 2 ;; + esac +done + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +WORK="$(mktemp -d -t obikmer_smoke.XXXXXX)" + +cleanup() { + if [ "$KEEP" -eq 1 ]; then + echo "kept: $WORK" + else + rm -rf "$WORK" + fi +} +trap cleanup EXIT + +fail() { + echo "FAIL: $1" >&2 + exit 1 +} + +# ── 1. generate a small deterministic random FASTA ───────────────────────── +python3 - "$WORK/test.fasta" "$K" <<'EOF' +import random, sys +path, k = sys.argv[1], int(sys.argv[2]) +random.seed(1234) +bases = "ACGT" +with open(path, "w") as f: + for i in range(3): + seq = "".join(random.choice(bases) for _ in range(300)) + f.write(f">seq{i}\n{seq}\n") +EOF + +# ── 2. build + run index ──────────────────────────────────────────────────── +cd "$REPO_ROOT/src" +INDEX_LOG="$WORK/index.log" +if ! cargo run -q -p obikmer --bin obikmer -- \ + index -k "$K" -m "$M" --theta 0 -p "$PARTITIONS" \ + -o "$WORK/out.idx" "$WORK/test.fasta" > "$INDEX_LOG" 2>&1 +then + cat "$INDEX_LOG" >&2 + fail "obikmer index exited non-zero" +fi + +N_KMERS="$(grep -o '[0-9]* total kmers indexed' "$INDEX_LOG" | grep -o '^[0-9]*' || true)" +[ -n "$N_KMERS" ] || { cat "$INDEX_LOG" >&2; fail "could not find 'N total kmers indexed' in index log"; } +[ "$N_KMERS" -gt 0 ] || fail "index reports 0 kmers indexed" + +# ── 3+4. try candidate k-mers spread across the source sequence until one +# round-trips. `query` applies its own entropy filter to the query +# sequence (not exposed as a CLI flag, undocumented threshold) — a +# window that fails it reports kmer_count:0 even though the k-mer +# is genuinely in the index (not a bug, just a bad test fixture). +# Rather than reverse-engineer the filter's formula, just ask the +# real binary and move to the next candidate on a miss. +readarray -t CANDIDATES < <(python3 - "$WORK/test.fasta" "$K" <<'EOF' +import sys +path, k = sys.argv[1], int(sys.argv[2]) +with open(path) as f: + seq = "".join(l.strip() for l in f if not l.startswith(">")) +for start in range(0, len(seq) - k, 17): + print(seq[start:start+k]) +EOF +) +[ "${#CANDIDATES[@]}" -gt 0 ] || fail "could not extract any candidate k-mer from the source FASTA" + +QUERY_LOG="$WORK/query.log" +FOUND=0 +for QUERY_KMER in "${CANDIDATES[@]}"; do + printf ">q1\n%s\n" "$QUERY_KMER" > "$WORK/query.fasta" + if ! cargo run -q -p obikmer --bin obikmer -- \ + query "$WORK/out.idx" "$WORK/query.fasta" > "$QUERY_LOG" 2>&1 + then + cat "$QUERY_LOG" >&2 + fail "obikmer query exited non-zero" + fi + if grep -q '"kmer_count":1' "$QUERY_LOG"; then + FOUND=1 + break + fi +done + +if [ "$FOUND" -ne 1 ]; then + cat "$QUERY_LOG" >&2 + fail "no candidate k-mer round-tripped (tried ${#CANDIDATES[@]}) — likely a real regression, not a low-complexity fixture" +fi + +echo "PASS: index+query round-trip OK — $N_KMERS kmers indexed, query k-mer '$QUERY_KMER' found" diff --git a/src/obikindexer/src/algorithms/partitionner/count.rs b/src/obikindexer/src/algorithms/counter/count.rs similarity index 100% rename from src/obikindexer/src/algorithms/partitionner/count.rs rename to src/obikindexer/src/algorithms/counter/count.rs diff --git a/src/obikindexer/src/algorithms/partitionner/kmer_sort.rs b/src/obikindexer/src/algorithms/counter/kmer_sort.rs similarity index 100% rename from src/obikindexer/src/algorithms/partitionner/kmer_sort.rs rename to src/obikindexer/src/algorithms/counter/kmer_sort.rs diff --git a/src/obikindexer/src/algorithms/counter/mod.rs b/src/obikindexer/src/algorithms/counter/mod.rs new file mode 100644 index 00000000..7319dc20 --- /dev/null +++ b/src/obikindexer/src/algorithms/counter/mod.rs @@ -0,0 +1,133 @@ +//! Kmer counting — the sibling algorithm to +//! [`crate::algorithms::dereplicator`]: runs after `dereplicator:: +//! Dereplicator::run` has deduplicated each partition's raw superkmer +//! file, the last step before `obikindex::KmerIndex::build_layers` turns +//! the provisional per-partition MPHF/counts this produces into the real +//! layer 0 — see `DevDocMD/implementation/partition_layer_cache.md`. + +mod count; +mod kmer_sort; + +use std::collections::BTreeMap; +use std::fs; +use std::io; +use std::sync::atomic::{AtomicU64, Ordering}; + +use obikindex::KmerIndex; +use obiskio::SKResult; +use obisys::Progress; +use rayon::prelude::*; +use sysinfo::System; + +use count::count_partition; +use kmer_sort::chunk_size_from_ram; + +pub struct KmerSpectrum { + pub f0: u64, + pub f1: u64, + pub counts: BTreeMap, +} + +/// For each partition with a dereplicated superkmer file: +/// 1. Enumerates all unique canonical kmers (two passes over the file). +/// 2. Builds a provisional MPHF (FMPHGO) over those kmers. +/// 3. Writes a flat binary count file (`counts1.bin`, one `u32` per slot, +/// memory-mapped) accumulating kmer abundances from the superkmer counts. +/// 4. Persists the MPHF to `mphf1.bin` for downstream use. +/// +/// Two-phase construction, same shape as `PartitionRouter`/`Dereplicator`: +/// `new` (no disk access), optional `keep_partial` setter, then `run`, +/// which returns the aggregated [`KmerSpectrum`]. +pub struct Counter<'a> { + index: &'a KmerIndex, + n_partitions: usize, + keep_partial: bool, +} + +impl<'a> Counter<'a> { + pub fn new(index: &'a KmerIndex) -> Self { + Self { + index, + n_partitions: index.n_partitions(), + keep_partial: false, + } + } + + /// Keep each partition's own `kmer_spectrum_raw.json` after aggregation + /// instead of deleting it (default: `false`). + pub fn keep_partial(mut self, v: bool) -> Self { + self.keep_partial = v; + self + } + + /// Count every partition in parallel, then aggregate their spectra. + /// + /// `on_progress`, when set, is called once per completed partition, from + /// whichever rayon worker thread finished it — `Fn(...) + Sync`, not + /// `FnMut`, same reason as `Dereplicator::run`: the counting pass itself + /// is a parallel `par_iter`, so the callback must tolerate concurrent + /// calls. `total: Some(n_partitions)` — known up front. This algorithm + /// never renders anything itself. + pub fn run(&self, on_progress: Option) -> SKResult { + let sys = System::new_all(); + // available_memory() can return 0 on macOS when the compressor page count exceeds + // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total. + let available = match sys.available_memory() { + 0 => sys.total_memory() / 2, + n => n, + }; + let n_threads = rayon::current_num_threads().max(1) as u64; + let chunk_kmers = chunk_size_from_ram(available / n_threads); + let done = AtomicU64::new(0); + + let results: Vec> = (0..self.n_partitions) + .into_par_iter() + .map(|i| { + let dir = self.index.layer_dir(i, 0); + let dedup_path = obikindex::layer::dereplicated_superkmers_path(&dir); + let result = if dedup_path.exists() { + count_partition(&dir, &dedup_path, chunk_kmers) + } else { + Ok(()) + }; + if let Some(cb) = &on_progress { + let pos = done.fetch_add(1, Ordering::Relaxed) + 1; + cb(Progress { position: pos, total: Some(self.n_partitions as u64) }); + } + result + }) + .collect(); + + for r in results { + r?; + } + + // Aggregate per-partition spectra. + let mut counts: BTreeMap = BTreeMap::new(); + let mut f0: u64 = 0; + let mut f1: u64 = 0; + + for i in 0..self.n_partitions { + let path = self.index.layer_dir(i, 0).join("kmer_spectrum_raw.json"); + if !path.exists() { + continue; + } + let v: serde_json::Value = + serde_json::from_str(&fs::read_to_string(&path)?).map_err(io::Error::other)?; + f0 += v["f0"].as_u64().unwrap_or(0); + f1 += v["f1"].as_u64().unwrap_or(0); + if let Some(obj) = v["spectrum"].as_object() { + for (c_str, freq) in obj { + if let (Ok(c), Some(f)) = (c_str.parse::(), freq.as_u64()) { + *counts.entry(c).or_insert(0) += f; + } + } + } + if !self.keep_partial { + let _ = fs::remove_file(&path); + } + } + + Ok(KmerSpectrum { f0, f1, counts }) + } +} diff --git a/src/obikindexer/src/algorithms/mod.rs b/src/obikindexer/src/algorithms/mod.rs index c25d8f4d..b20e5765 100644 --- a/src/obikindexer/src/algorithms/mod.rs +++ b/src/obikindexer/src/algorithms/mod.rs @@ -2,8 +2,10 @@ //! [`obikindex::KmerIndex`] to build or transform its content — `obikindex` //! itself is the `index`/`partition`/`layer` data model, not this. Each //! algorithm is its own submodule: [`partitionner`] (routing raw super-kmers -//! into partitions, then counting), [`dereplicator`] (deduplicating a -//! partition's raw super-kmers before counting). +//! into partitions), [`dereplicator`] (deduplicating a partition's raw +//! super-kmers), [`counter`] (counting unique canonical kmers from a +//! partition's dereplicated super-kmers) — the three run in that order. +pub mod counter; pub mod dereplicator; pub mod partitionner; diff --git a/src/obikindexer/src/algorithms/partitionner/mod.rs b/src/obikindexer/src/algorithms/partitionner/mod.rs index e712bf12..2752a752 100644 --- a/src/obikindexer/src/algorithms/partitionner/mod.rs +++ b/src/obikindexer/src/algorithms/partitionner/mod.rs @@ -1,20 +1,16 @@ //! K-mer partitioning: routing super-kmers into per-partition, layer-0 -//! files, and counting unique canonical k-mers. Dereplication of the raw -//! super-kmers in between is [`crate::algorithms::dereplicator`], a -//! sibling algorithm, not a step of this one — see +//! files. Dereplication and counting of the raw super-kmers are separate +//! sibling algorithms ([`crate::algorithms::dereplicator`], +//! [`crate::algorithms::counter`]), not steps of this one — see //! `DevDocMD/implementation/partition_layer_cache.md`. //! -//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the -//! routing/counting lifecycle API — partition/layer path naming itself -//! lives on `obikindex::KmerIndex`, not here), [`count`] (unique-kmer -//! enumeration, MPHF, abundance counting), `kmer_sort` (external sort -//! support for `count`). +//! Submodule: [`router`] (`PartitionRouter`, the routing lifecycle API — +//! partition/layer path naming itself lives on `obikindex::KmerIndex`, +//! not here). -mod count; -mod kmer_sort; mod router; #[cfg(test)] mod tests; -pub use router::{KmerSpectrum, PartitionRouter}; +pub use router::PartitionRouter; diff --git a/src/obikindexer/src/algorithms/partitionner/router.rs b/src/obikindexer/src/algorithms/partitionner/router.rs index f7804673..4b052a04 100644 --- a/src/obikindexer/src/algorithms/partitionner/router.rs +++ b/src/obikindexer/src/algorithms/partitionner/router.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; -use std::fs; use std::io; use std::path::PathBuf; use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; @@ -10,9 +8,7 @@ use obikindex::KmerIndex; use obikseq::RoutableSuperKmer; use obikindex::layer::Layer; use obiskio::SKResult; -use obisys::{progress_bar, Progress}; -use rayon::prelude::*; -use sysinfo::System; +use obisys::Progress; use tracing::info; use niffler::Level; @@ -22,16 +18,6 @@ use obiskio::SKFileWriter; use obipipeline::{throttle, ThrottleGuard, Throttled}; use obiread::NucPage; -use super::kmer_sort::chunk_size_from_ram; - -use super::count::count_partition; - -pub struct KmerSpectrum { - pub f0: u64, - pub f1: u64, - pub counts: BTreeMap, -} - // ── Pipeline plumbing, private to `run` ───────────────────────────────────── /// Carrier enum for `obipipeline::make_pipe!`'s two-stage transform — local @@ -72,17 +58,17 @@ impl Drop for GuardedIter { // ── PartitionRouter ────────────────────────────────────────────────────────── -/// Routes raw super-kmers into per-partition, layer-0 files, then -/// dereplicates (via the sibling [`crate::algorithms::dereplicator`] -/// algorithm) and counts them — the entry point of the indexing pipeline, -/// operating on layer-0 content that doesn't exist yet at this stage (see +/// Routes raw super-kmers into per-partition, layer-0 files — the entry +/// point of the indexing pipeline. Dereplication +/// ([`crate::algorithms::dereplicator`]) and counting +/// ([`crate::algorithms::counter`]) are separate sibling algorithms that +/// run after this one, on the layer-0 content it produces (see /// `DevDocMD/implementation/partition_layer_cache.md`). /// /// Holds `&mut KmerIndex` — this is an algorithm operating on an index, not /// a data structure of its own; it owns no path-naming knowledge (every /// path comes from `index.index_dir`/`obikindex::layer::layer_dir`/ -/// `Layer::create`), only the transient routing/dereplication/counting -/// state a run needs. +/// `Layer::create`), only the transient routing state a run needs. /// /// Two-phase construction: `new` + optional setters (`level_max`/`theta`/ /// `workers`/`max_open`) configure the run, `run` executes it. `run` takes @@ -92,7 +78,6 @@ impl Drop for GuardedIter { /// the caller's decision, not this crate's. pub struct PartitionRouter<'a> { index: &'a mut KmerIndex, - n_partitions: usize, partitions_mask: u64, writers: Vec>, level: Level, @@ -113,7 +98,6 @@ impl<'a> PartitionRouter<'a> { let workers = obisys::effective_parallelism(); Self { index, - n_partitions, partitions_mask: (1u64 << n_bits) - 1, writers: (0..n_partitions).map(|_| None).collect(), level: Level::One, @@ -267,77 +251,6 @@ impl<'a> PartitionRouter<'a> { self.close() } - /// For each partition that has a `dereplicated.{ext}` file: - /// 1. Enumerates all unique canonical kmers (two passes over the file). - /// 2. Builds a provisional MPHF (FMPHGO) over those kmers. - /// 3. Writes a flat binary count file (`counts1.bin`, one `u32` per slot, - /// memory-mapped) accumulating kmer abundances from the superkmer counts. - /// 4. Persists the MPHF to `mphf1.bin` for downstream use. - /// - /// Returns the aggregated `KmerSpectrum`. Per-partition spectrum files are - /// deleted after aggregation unless `keep_partial` is true. - pub fn count_kmer(&self, keep_partial: bool) -> SKResult { - let sys = System::new_all(); - let available = match sys.available_memory() { - 0 => sys.total_memory() / 2, - n => n, - }; - let n_threads = rayon::current_num_threads().max(1) as u64; - let chunk_kmers = chunk_size_from_ram(available / n_threads); - - let pb = progress_bar("counting", self.n_partitions as u64, "partitions"); - - let results: Vec> = (0..self.n_partitions) - .into_par_iter() - .map(|i| { - let dir = self.layer0_dir(i); - let dedup_path = obikindex::layer::dereplicated_superkmers_path(&dir); - if !dedup_path.exists() { - pb.inc(1); - return Ok(()); - } - let t = Instant::now(); - let result = count_partition(&dir, &dedup_path, chunk_kmers); - pb.set_message(format!("last {:.0}ms", t.elapsed().as_millis())); - pb.inc(1); - result - }) - .collect(); - - pb.finish_and_clear(); - for r in results { - r?; - } - - // Aggregate per-partition spectra. - let mut counts: BTreeMap = BTreeMap::new(); - let mut f0: u64 = 0; - let mut f1: u64 = 0; - - for i in 0..self.n_partitions { - let path = self.layer0_dir(i).join("kmer_spectrum_raw.json"); - if !path.exists() { - continue; - } - let v: serde_json::Value = - serde_json::from_str(&fs::read_to_string(&path)?).map_err(io::Error::other)?; - f0 += v["f0"].as_u64().unwrap_or(0); - f1 += v["f1"].as_u64().unwrap_or(0); - if let Some(obj) = v["spectrum"].as_object() { - for (c_str, freq) in obj { - if let (Ok(c), Some(f)) = (c_str.parse::(), freq.as_u64()) { - *counts.entry(c).or_insert(0) += f; - } - } - } - if !keep_partial { - let _ = fs::remove_file(&path); - } - } - - Ok(KmerSpectrum { f0, f1, counts }) - } - // ── private ─────────────────────────────────────────────────────────────── /// Directory of partition `i`'s layer 0 — every raw/dereplicated diff --git a/src/obikindexer/src/algorithms/partitionner/tests.rs b/src/obikindexer/src/algorithms/partitionner/tests.rs index b0ab085b..85713b0a 100644 --- a/src/obikindexer/src/algorithms/partitionner/tests.rs +++ b/src/obikindexer/src/algorithms/partitionner/tests.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use std::fs; +use crate::algorithms::counter::Counter; use crate::algorithms::dereplicator::Dereplicator; use obikindex::{IndexConfig, KmerIndex}; use obikindex::layer::IndexMode; @@ -8,7 +8,6 @@ use obikrope::Rope; use obikseq::SuperKmer; use obiskbuilder::build_superkmers; -use super::count::count_partition; use super::PartitionRouter; const K: usize = 11; @@ -46,7 +45,7 @@ fn direct_counts(seqs: &[&[u8]]) -> (u64, u64) { } /// Run the full pipeline on a list of sequences and return (f0, f1) from -/// the `kmer_spectrum_raw.json` produced by `count_partition`. +/// the aggregated `KmerSpectrum` produced by `Counter::run`. fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) { setup(); @@ -73,14 +72,8 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) { if !dedup_path.exists() { return (0, 0); } - 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 f0 = spec["f0"].as_u64().unwrap_or(0); - let f1 = spec["f1"].as_u64().unwrap_or(0); - (f0, f1) + let spectrum = Counter::new(&index).run(None::).unwrap(); + (spectrum.f0, spectrum.f1) } #[test] diff --git a/src/obikmer/src/cmd/index/mod.rs b/src/obikmer/src/cmd/index/mod.rs index 8372c4ab..7e299b31 100644 --- a/src/obikmer/src/cmd/index/mod.rs +++ b/src/obikmer/src/cmd/index/mod.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use std::time::Instant; use clap::Args; +use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::dereplicator::Dereplicator; use obikindexer::algorithms::partitionner::PartitionRouter; use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex}; @@ -301,15 +302,17 @@ pub fn run(args: IndexArgs) { pb.finish_and_clear(); rep.push(t.stop()); - let router = PartitionRouter::new(&mut idx); - let t = Stage::start("count_kmer"); - let spectrum = router.count_kmer(args.keep_intermediate).unwrap_or_else(|e| { - eprintln!("error: {e}"); - std::process::exit(1); - }); + let pb = progress_bar("counting", idx.n_partitions() as u64, "partitions"); + let spectrum = Counter::new(&idx) + .keep_partial(args.keep_intermediate) + .run(Some(|_: Progress| pb.inc(1))) + .unwrap_or_else(|e| { + eprintln!("error: {e}"); + std::process::exit(1); + }); + pb.finish_and_clear(); rep.push(t.stop()); - drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).unwrap_or_else(|e| { eprintln!("error: {e}"); diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index 90fa56fa..cfe7d71d 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -6,6 +6,7 @@ use obikindex::layer::MphfLayer; use obisys::Reporter; use tempfile::tempdir; +use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::dereplicator::Dereplicator; use obikindexer::algorithms::partitionner::PartitionRouter; use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode}; @@ -80,9 +81,7 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope Dereplicator::new(&idx).run(None::).expect("dereplicate"); - let router = PartitionRouter::new(&mut idx); - let spectrum = router.count_kmer(false).expect("count_kmer"); - drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope + let spectrum = Counter::new(&idx).run(None::).expect("count_kmer"); idx.mark_scattered().expect("mark_scattered"); idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum");