Push zunrplorkwkt #70
@@ -4,6 +4,7 @@
|
||||
.kilo/
|
||||
.serena/
|
||||
.zed/
|
||||
.ast-cache/
|
||||
memory/
|
||||
sandbox/
|
||||
src/target
|
||||
|
||||
@@ -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<impl Fn(Progress) + Sync>) -> SKResult<KmerSpectrum>;
|
||||
}
|
||||
```
|
||||
|
||||
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
|
||||
|
||||
Executable
+128
@@ -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"
|
||||
@@ -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<u32, u64>,
|
||||
}
|
||||
|
||||
/// 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<impl Fn(Progress) + Sync>) -> SKResult<KmerSpectrum> {
|
||||
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<SKResult<()>> = (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<u32, u64> = 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::<u32>(), freq.as_u64()) {
|
||||
*counts.entry(c).or_insert(0) += f;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !self.keep_partial {
|
||||
let _ = fs::remove_file(&path);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(KmerSpectrum { f0, f1, counts })
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<u32, u64>,
|
||||
}
|
||||
|
||||
// ── 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<Option<SKFileWriter>>,
|
||||
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<KmerSpectrum> {
|
||||
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<SKResult<()>> = (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<u32, u64> = 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::<u32>(), 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
|
||||
|
||||
@@ -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::<fn(obisys::Progress)>).unwrap();
|
||||
(spectrum.f0, spectrum.f1)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -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| {
|
||||
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}");
|
||||
|
||||
@@ -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::<fn(obisys::Progress)>).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::<fn(obisys::Progress)>).expect("count_kmer");
|
||||
|
||||
idx.mark_scattered().expect("mark_scattered");
|
||||
idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum");
|
||||
|
||||
Reference in New Issue
Block a user