Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
21a20ce7ca |
+7
-14
@@ -1,8 +1,9 @@
|
|||||||
pname: CI
|
name: CI
|
||||||
|
|
||||||
on:
|
on:
|
||||||
pull_request:
|
push:
|
||||||
branches: ['main']
|
branches: ['main']
|
||||||
|
pull_request:
|
||||||
|
|
||||||
jobs:
|
jobs:
|
||||||
build:
|
build:
|
||||||
@@ -25,19 +26,11 @@ jobs:
|
|||||||
~/.cargo/registry
|
~/.cargo/registry
|
||||||
~/.cargo/git
|
~/.cargo/git
|
||||||
src/target
|
src/target
|
||||||
key: ${{ runner.os }}-cargo-v2-${{ hashFiles('src/Cargo.lock') }}
|
key: ${{ runner.os }}-cargo-${{ hashFiles('src/Cargo.lock') }}
|
||||||
restore-keys: ${{ runner.os }}-cargo-v2-
|
restore-keys: ${{ runner.os }}-cargo-
|
||||||
|
|
||||||
# Both `obikmer` and `obikindex` default to the `numa` feature
|
|
||||||
# (hwloc-based topology detection + CPU pinning), which is only useful
|
|
||||||
# on bare-metal multi-socket indexing hosts. Under this runner's
|
|
||||||
# container/cgroup setup it deadlocks at startup — confirmed live
|
|
||||||
# (2026-08-11): the same test binary hangs indefinitely with `numa` on
|
|
||||||
# and passes instantly, repeatedly, with it off, on the same
|
|
||||||
# container. Disable it for CI; it has nothing to do with test
|
|
||||||
# correctness.
|
|
||||||
- name: Build
|
- name: Build
|
||||||
run: cargo build --release --no-default-features
|
run: cargo build --release
|
||||||
|
|
||||||
- name: Test
|
- name: Test
|
||||||
run: cargo test --release --no-default-features
|
run: cargo test --release
|
||||||
|
|||||||
@@ -13,7 +13,7 @@ jobs:
|
|||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
with:
|
with:
|
||||||
fetch-depth: 0
|
fetch-tags: true
|
||||||
|
|
||||||
- name: Create Gitea release
|
- name: Create Gitea release
|
||||||
id: create
|
id: create
|
||||||
@@ -86,11 +86,20 @@ jobs:
|
|||||||
build-macos-arm64:
|
build-macos-arm64:
|
||||||
needs: create-release
|
needs: create-release
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
|
defaults:
|
||||||
|
run:
|
||||||
|
working-directory: src
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
|
|
||||||
- name: Login to registry
|
- name: Install Rust + zigbuild
|
||||||
run: echo "${{ secrets.REGISTRYTOKEN }}" | docker login registry.metabarcoding.org -u ${{ secrets.REGISTRYUSER }} --password-stdin
|
run: |
|
||||||
|
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable
|
||||||
|
echo "$HOME/.cargo/bin" >> $GITHUB_PATH
|
||||||
|
sudo apt-get update -qq && sudo apt-get install -y -qq jq
|
||||||
|
pip install ziglang --quiet --break-system-packages
|
||||||
|
$HOME/.cargo/bin/cargo install cargo-zigbuild
|
||||||
|
$HOME/.cargo/bin/rustup target add aarch64-apple-darwin
|
||||||
|
|
||||||
- name: Cache cargo registry
|
- name: Cache cargo registry
|
||||||
uses: actions/cache@v4
|
uses: actions/cache@v4
|
||||||
@@ -103,24 +112,15 @@ jobs:
|
|||||||
restore-keys: macos-arm64-cargo-
|
restore-keys: macos-arm64-cargo-
|
||||||
|
|
||||||
- name: Build macOS binary
|
- name: Build macOS binary
|
||||||
run: |
|
run: cargo zigbuild --release --target aarch64-apple-darwin --no-default-features
|
||||||
CID=$(docker create \
|
|
||||||
-w /src/src \
|
|
||||||
registry.metabarcoding.org/cibuilder/rustcrossosx:latest \
|
|
||||||
cargo build --release --target aarch64-apple-darwin --no-default-features)
|
|
||||||
docker cp . "$CID:/src"
|
|
||||||
docker start -a "$CID"
|
|
||||||
STATUS=$(docker wait "$CID")
|
|
||||||
mkdir -p /tmp/dist
|
|
||||||
docker cp "$CID:/src/src/target/aarch64-apple-darwin/release/obikmer" /tmp/dist/obikmer-macos-arm64
|
|
||||||
docker rm "$CID" > /dev/null
|
|
||||||
[ "$STATUS" -eq 0 ]
|
|
||||||
|
|
||||||
- name: Prepare and upload artifact
|
- name: Prepare and upload artifact
|
||||||
env:
|
env:
|
||||||
GITEA_TOKEN: ${{ secrets.GITEATOKEN }}
|
GITEA_TOKEN: ${{ secrets.GITEATOKEN }}
|
||||||
RELEASE_ID: ${{ needs.create-release.outputs.release_id }}
|
RELEASE_ID: ${{ needs.create-release.outputs.release_id }}
|
||||||
run: |
|
run: |
|
||||||
|
mkdir -p /tmp/dist
|
||||||
|
cp target/aarch64-apple-darwin/release/obikmer /tmp/dist/obikmer-macos-arm64
|
||||||
curl -s -X POST \
|
curl -s -X POST \
|
||||||
"${{ github.server_url }}/api/v1/repos/${{ github.repository }}/releases/$RELEASE_ID/assets" \
|
"${{ github.server_url }}/api/v1/repos/${{ github.repository }}/releases/$RELEASE_ID/assets" \
|
||||||
-H "Authorization: token $GITEA_TOKEN" \
|
-H "Authorization: token $GITEA_TOKEN" \
|
||||||
|
|||||||
@@ -8,15 +8,12 @@ data-stress
|
|||||||
*.pb
|
*.pb
|
||||||
./**/*.json
|
./**/*.json
|
||||||
*.bin
|
*.bin
|
||||||
*.log
|
|
||||||
*.csv
|
|
||||||
Betula_exilis--IGA-24-33
|
Betula_exilis--IGA-24-33
|
||||||
benchmark/genomes
|
benchmark/genomes
|
||||||
benchmark/simulated_data
|
benchmark/simulated_data
|
||||||
benchmark/specimen_index_presence
|
benchmark/specimen_index_presence
|
||||||
benchmark/specimen_index_count
|
benchmark/specimen_index_count
|
||||||
benchmark/global_index_presence
|
benchmark/global_index_presence
|
||||||
benchmark/all_specific
|
|
||||||
benchmark/global_index_count
|
benchmark/global_index_count
|
||||||
benchmark/stats
|
benchmark/stats
|
||||||
benchmark/reference_index
|
benchmark/reference_index
|
||||||
|
|||||||
@@ -92,7 +92,6 @@ release: bump-version
|
|||||||
git_hash=$$(jj log -r @ --no-graph -T 'commit_id'); \
|
git_hash=$$(jj log -r @ --no-graph -T 'commit_id'); \
|
||||||
commits=$$(jj log -r 'latest(tags())..@' --no-graph -T 'description ++ "\n"' 2>/dev/null || \
|
commits=$$(jj log -r 'latest(tags())..@' --no-graph -T 'description ++ "\n"' 2>/dev/null || \
|
||||||
jj log --no-graph -T 'description ++ "\n"' --limit 30); \
|
jj log --no-graph -T 'description ++ "\n"' --limit 30); \
|
||||||
notes=$$(printf 'Write concise markdown release notes for obikmer (a Rust kmer genomics tool). Be technical and direct. Base them strictly on these commit messages:\n\n%s' "$$commits" | aichat 2>/dev/null); \
|
notes=$$(printf 'Write concise markdown release notes for obikmer (a Rust kmer genomics tool). Be technical and direct. Base them strictly on these commit messages:\n\n%s' "$$commits" | aichat); \
|
||||||
tag_msg="$${notes:-Release v$$new_version}"; \
|
git tag -a "v$$new_version" -m "$$notes" "$$git_hash" && \
|
||||||
git tag -a "v$$new_version" -m "$$tag_msg" "$$git_hash" && \
|
|
||||||
git push origin "v$$new_version"
|
git push origin "v$$new_version"
|
||||||
|
|||||||
@@ -162,158 +162,14 @@ A single `PartitionRunner` instance can be built once per command invocation
|
|||||||
and reused across multiple `run()` calls (e.g. `merge` runs
|
and reused across multiple `run()` calls (e.g. `merge` runs
|
||||||
`merge_partitions` then `pack_matrices`).
|
`merge_partitions` then `pack_matrices`).
|
||||||
|
|
||||||
## Known issue: CPU-only activation signal stalls on I/O-bound stages
|
|
||||||
|
|
||||||
Observed on a real `filter` run (109 genomes, 256 partitions, 8×24-core NUMA):
|
|
||||||
`rebuild` (CPU-bound — k-mer construction) scales cleanly from 9 to 43 active
|
|
||||||
workers as `CpuSample::do_i_activate` (`obisys::lib.rs`) sees efficiency climb.
|
|
||||||
`pack_matrices` (I/O-bound — reopens and recomposes per-genome column files
|
|
||||||
into `.pbmx`/`.pcmx`) activates one extra worker then flatlines at 10/192 for
|
|
||||||
the rest of the stage, even though 256 partitions keep completing over several
|
|
||||||
minutes. This matches the documented intent (§ Adaptive mechanism — "avoids
|
|
||||||
over-provisioning ... I/O-bound ... workloads") but conflates two different
|
|
||||||
things: *"CPU is not the bottleneck"* and *"more workers would not help"*. On
|
|
||||||
storage with real queue depth (NVMe, RAID, parallel FS) the second stage could
|
|
||||||
still benefit from more concurrent workers even with flat CPU usage — a signal
|
|
||||||
the current mechanism cannot see.
|
|
||||||
|
|
||||||
A one-off artefact was also found in the same log: right after a stage
|
|
||||||
transition, `do_i_activate` produced a physically impossible spike (efficiency
|
|
||||||
~94 cores on a 192-core box) because it has no minimum-window guard — unlike
|
|
||||||
its sibling `cpu_efficiency`, which returns `0.0` if `wall < 0.1s`
|
|
||||||
(`obisys::lib.rs:260`). `do_i_activate` unconditionally overwrites
|
|
||||||
`self.wall`/`self.user_secs`/`self.sys_secs` even when the elapsed window is
|
|
||||||
too short to be meaningful, so a burst of rapid completions right after
|
|
||||||
activating a worker can divide a real CPU delta by a near-zero wall delta.
|
|
||||||
|
|
||||||
### Implemented: I/O signal + shared debounce guard
|
|
||||||
|
|
||||||
`IoSample` (`obisys::lib.rs`, alongside `CpuSample`) is fed by
|
|
||||||
`read_bytes`/`write_bytes` from `/proc/self/io` on Linux (actual bytes
|
|
||||||
submitted to the block layer — not `rchar`/`wchar`, which also count
|
|
||||||
page-cache hits, and not `ru_inblock`/`ru_oublock`, unreliable on macOS), with
|
|
||||||
a `proc_pid_rusage(RUSAGE_INFO_V4)` fallback on macOS
|
|
||||||
(`ri_diskio_bytesread`/`ri_diskio_byteswritten`, FFI only via `libc`, no new
|
|
||||||
dependency — same pattern as the existing `getrusage` bindings). Any other
|
|
||||||
target degrades gracefully to a signal that never triggers (falls back to
|
|
||||||
CPU-only activation), same pattern as `cgroup_v2_available`.
|
|
||||||
|
|
||||||
`maybe_activate` (`numa.rs`) activates a worker if *either* signal still shows
|
|
||||||
headroom, making `PartitionRunner` adapt to whichever resource is actually the
|
|
||||||
bottleneck without per-call configuration. Both samplers are called
|
|
||||||
unconditionally — no `||` short-circuit — so neither window starves behind
|
|
||||||
whichever signal fires first:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
let cpu_threshold = CPU_SPAWN_THRESHOLD * activation.last_step() as f64;
|
|
||||||
let cpu_wants_more = cpu_sample.do_i_activate(cpu_threshold);
|
|
||||||
let io_wants_more = io_sample.do_i_activate(IO_SPAWN_THRESHOLD);
|
|
||||||
if cpu_wants_more || io_wants_more {
|
|
||||||
activation.grow(GROWTH_DIVISOR, n_total);
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
The CPU threshold is *not* the flat absolute delta it started as: it scales
|
|
||||||
with `activation.last_step()` — the number of workers activated in the last
|
|
||||||
growth step, tracked by `NodeActivation` (`numa.rs`) and updated every time
|
|
||||||
`grow()` actually grows something. Growing by 8 workers should add ~8 cores of
|
|
||||||
efficiency if the workload is truly CPU-bound; requiring only
|
|
||||||
`CPU_SPAWN_THRESHOLD` (20 %) of that expected gain confirms the growth was
|
|
||||||
useful without demanding perfect linear scaling. Scaling by the *last step's
|
|
||||||
size* rather than the cumulative total keeps the bar equally meaningful
|
|
||||||
whether it's the 2nd growth step or the 20th — a flat absolute threshold
|
|
||||||
(0.2 core) is a strong signal at 8 active workers but pure noise at 150; a
|
|
||||||
threshold scaled by the *cumulative* total instead (considered and rejected)
|
|
||||||
would have made the bar essentially impossible to clear late in the ramp,
|
|
||||||
strangling exactly the CPU-bound saturation the mechanism exists to allow.
|
|
||||||
|
|
||||||
Unlike the CPU signal (an absolute delta in cores — a bounded, portable unit),
|
|
||||||
raw I/O throughput has no natural scale across devices, so `IoSample` uses a
|
|
||||||
**relative** growth threshold instead of an absolute one:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
pub fn do_i_activate(&mut self, threshold: f64) -> bool {
|
|
||||||
let elapsed = self.wall.elapsed().as_secs_f64();
|
|
||||||
if elapsed < 0.1 { return false; } // state untouched — window keeps accumulating
|
|
||||||
|
|
||||||
let n = Self::read_bytes();
|
|
||||||
let rate = n.saturating_sub(self.bytes) as f64 / elapsed;
|
|
||||||
let activate = if self.previous_rate == 0.0 {
|
|
||||||
rate > 0.0 // bootstrap: any measured throughput is signal
|
|
||||||
} else {
|
|
||||||
(rate - self.previous_rate) / self.previous_rate >= threshold
|
|
||||||
};
|
|
||||||
|
|
||||||
self.bytes = n;
|
|
||||||
self.wall = Instant::now(); // reset only on a real sample
|
|
||||||
activate
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
The `elapsed < 0.1s → return false without mutating state` guard was also
|
|
||||||
back-ported into `CpuSample::do_i_activate` (previously missing — source of
|
|
||||||
the ~94-core artefact above) — one fix for both problems, and it removes the
|
|
||||||
need for any arbitrary I/O-rate floor: a short/noisy window is rejected
|
|
||||||
outright rather than papered over with a hardware-dependent constant.
|
|
||||||
|
|
||||||
Both spawn thresholds (`CPU_SPAWN_THRESHOLD`, `IO_SPAWN_THRESHOLD`, module-level
|
|
||||||
`const` in `numa.rs`, both `0.2`) are a starting point, not a derived value:
|
|
||||||
`0.2` (20 % relative growth) for `IoSample` was chosen to match the CPU
|
|
||||||
threshold's *implicit* relative sensitivity (in the observed log, an 8→9
|
|
||||||
worker step raised efficiency by ~12 %) — but I/O throughput is lumpier than
|
|
||||||
CPU time (buffered writes flush in bursts), so it needs empirical validation
|
|
||||||
against a real `pack` run before being considered final.
|
|
||||||
|
|
||||||
## Known issue: ramp-up too slow, and confused with node count
|
|
||||||
|
|
||||||
The original design started `n_nodes` workers (one per node) and grew one
|
|
||||||
worker at a time. On a real `filter` run this took ~10 minutes to climb from
|
|
||||||
9 to ~40 active workers even on the CPU-bound `rebuild` stage — most of a
|
|
||||||
35-minute stage spent under-provisioned while waiting for evidence to
|
|
||||||
accumulate one worker at a time. There is no scale-down mechanism (`n_active`
|
|
||||||
only grows), so the original caution was deliberate — but a quarter of
|
|
||||||
available cores is still far from saturation, and the real risk zone (over-provisioning
|
|
||||||
a memory-bandwidth-bound stage) only shows up much later in the ramp, near
|
|
||||||
full occupancy — not at 25 %.
|
|
||||||
|
|
||||||
The fix decouples ramp speed from node *count*: both the initial size and the
|
|
||||||
growth step are a fraction of `workers_per_node` (node *size*), applied
|
|
||||||
identically on every node. A single-NUMA-node (UMA) machine ramps exactly as
|
|
||||||
fast as an 8-node one — growing by `n_nodes` per step, as first considered,
|
|
||||||
would have degenerated to "grow by 1" on UMA, reproducing the original
|
|
||||||
problem for exactly the machines that need the fix most.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
// NodeActivation::grow — called both at startup (activate_initial) and on
|
|
||||||
// every CPU/IO-triggered growth step, with a different divisor each time.
|
|
||||||
let wanted = (self.caps[idx] / divisor).max(1); // INITIAL_DIVISOR=4 at startup, GROWTH_DIVISOR=8 per step
|
|
||||||
let room = self.caps[idx].saturating_sub(self.active[idx]);
|
|
||||||
let grow = wanted.min(room).min(n_total.saturating_sub(self.total));
|
|
||||||
```
|
|
||||||
|
|
||||||
This also fixed a latent correctness gap: the original single shared
|
|
||||||
`activate_tx`/`activate_rx` pair had *no* per-node addressing — sending one
|
|
||||||
activation signal woke up whichever dormant worker (from any node) happened
|
|
||||||
to win the race on that channel. `crossbeam_channel` gives no fairness
|
|
||||||
guarantee across competing receivers, so "round-robin across nodes" was an
|
|
||||||
assumption the code never actually enforced. `PartitionRunner::run` now opens
|
|
||||||
one activation channel per node (`activate_txs`/`activate_rxs`, one pair per
|
|
||||||
`NodeConfig`); `NodeActivation` (`numa.rs`) tracks how many of each node's
|
|
||||||
dormant workers have been woken and grows every node by the same amount per
|
|
||||||
step, capped by that node's remaining dormant workers and by the run's total
|
|
||||||
budget (`n_total`) — balance across nodes is now guaranteed by construction,
|
|
||||||
not incidental to channel implementation details.
|
|
||||||
|
|
||||||
## Open questions
|
## Open questions
|
||||||
|
|
||||||
- **Error handling**: `run` currently returns the first error; remaining errors
|
- **Error handling**: `run` currently returns the first error; remaining errors
|
||||||
are dropped. A `Vec<E>` return would give complete diagnostics.
|
are dropped. A `Vec<E>` return would give complete diagnostics.
|
||||||
|
|
||||||
- **`INITIAL_DIVISOR` / `GROWTH_DIVISOR` tuning**: currently `4` and `8`
|
- **`workers_per_node` tuning**: currently `(cpus / 8).max(3).min(8)`, calibrated
|
||||||
(start at 1/4 of a node's cores, grow by 1/8 per step), chosen to fix an
|
for merge on BeeGFS. I/O-bound commands (`dump`, `select`) may benefit from
|
||||||
observed too-slow ramp — not yet validated against a real `pack` (I/O-bound)
|
a higher value. A per-call override could be added to the API.
|
||||||
run, where over-provisioning risk is different from the CPU-bound `rebuild`
|
|
||||||
case this was tuned against.
|
|
||||||
|
|
||||||
- **`on_done` ordering**: the runner serialises calls to `on_done` via an
|
- **`on_done` ordering**: the runner serialises calls to `on_done` via an
|
||||||
internal `Arc<Mutex<C>>`. `Send` is required (the Arc clone crosses thread
|
internal `Arc<Mutex<C>>`. `Send` is required (the Arc clone crosses thread
|
||||||
|
|||||||
+38
-255
@@ -16,43 +16,27 @@ Given a set of query sequences, determine for each sequence how many of its k-me
|
|||||||
|
|
||||||
## Algorithm
|
## Algorithm
|
||||||
|
|
||||||
The query follows the same superkmer-based partitioning strategy used at indexing time. Everything below happens inside `process_chunk` (`query.rs`); there is no separate per-stage function, but the internal data flow is staged: k-mer-level dereplication, a two-part MPHF/column-major matrix lookup (`obikpartitionner::query_partition_with`), and a sparse Findere pass, each producing sparse intermediate structures rather than one dense allocation for the whole chunk.
|
The query follows the same superkmer-based partitioning strategy used at indexing time.
|
||||||
|
|
||||||
```
|
```
|
||||||
for each chunk of sequences (parallel workers via obipipeline, one call to process_chunk):
|
for each chunk of sequences (parallel workers via obipipeline):
|
||||||
build QueryBatch (QueryBatch::from_records):
|
build QueryBatch: decompose all sequences into s-mers via superkmers, deduplicate
|
||||||
decompose all sequences into superkmers (SuperKmerIter) — construction only,
|
allocate seq_results[seq_idx][smer_pos] = None ← per-sequence s-mer result vectors
|
||||||
not the dedup key
|
split superkmers by partition via minimiser hash
|
||||||
deduplicate at k-mer granularity, split by partition in the same pass:
|
|
||||||
by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> ← KmerDesc = (seq_idx, pos)
|
|
||||||
allocate SmerIndex (SmerIndex::new): in_index: Vec<bool>, sized total_smers —
|
|
||||||
NOT multiplied by n_genomes
|
|
||||||
allocate by_genome: Vec<Vec<(seq_idx, pos, value)>>, one empty Vec per genome —
|
|
||||||
stays empty (zero cost) for every genome this chunk never matches
|
|
||||||
for each partition p:
|
for each partition p:
|
||||||
query_partition_with(p, kmers_for_p, on_event):
|
query_partition(p, superkmers_routed_to_p)
|
||||||
stage 1 (MPHF-only): for each unique k-mer, try each layer's MphfLayer::find
|
→ load QueryLayer(s) for p
|
||||||
in turn, stop at the first hit; bucket confirmed hits by (layer, slot);
|
→ for each s-mer in each superkmer: MphfLayer::find(smer)
|
||||||
emit QueryHit::Found(descs) once per hit k-mer
|
fill seq_results[seq_idx][kmer_offset + j] from partition results
|
||||||
stage 2 (column-major fetch): for each layer with ≥1 hit, for each genome
|
for each sequence:
|
||||||
column g in 0..layer.n_cols(): scan that layer's bucketed slots, look up
|
apply_findere(seq_results[seq_idx], effective_z) ← per full sequence
|
||||||
col_value(g, slot); emit QueryHit::Value(descs, g, value) on nonzero
|
accumulate confirmed k-mer results into acc and cov
|
||||||
on_event dispatches: Found → SmerIndex::mark_found for every desc;
|
emit annotated sequences
|
||||||
Value → push (seq_idx, pos, value) into by_genome[g]
|
|
||||||
for each genome g with ≥1 hit (sparse_findere_for_genome):
|
|
||||||
sort by_genome[g] by (seq_idx, pos); detect maximal runs of consecutive pos
|
|
||||||
within one seq_idx; monotone-deque window-minimum scoped to each run →
|
|
||||||
confirmed_by_genome[g]: Vec<(seq_idx, pos_out, value)>
|
|
||||||
accumulate genome_totals per sequence from confirmed_by_genome (per genome, direct)
|
|
||||||
accumulate kmer_count / kmer_missing per (sequence, output position), O(1) each,
|
|
||||||
using only the confirmed-any bitmap and SmerIndex — independent of n_genomes
|
|
||||||
if --detail: densify confirmed_by_genome into per-(seq, genome) coverage arrays
|
|
||||||
emit annotated sequences (emit_batch)
|
|
||||||
```
|
```
|
||||||
|
|
||||||
Superkmers that appear more than once in the batch (same sequence or across sequences), or different superkmers that happen to share a k-mer (read overlaps, repeats, a SNP splitting an otherwise-identical run), are deduplicated at k-mer granularity: each unique `CanonicalKmer` triggers at most one MPHF lookup and, on hit, one matrix fetch, broadcast to every `KmerDesc` occurrence referencing it.
|
Superkmers that appear more than once in the batch (same sequence or across sequences) are deduplicated: each unique `RoutableSuperKmer` is queried once per partition, and the result is broadcast to every `SKDesc` entry that references it.
|
||||||
|
|
||||||
**Findere requires full-sequence aggregation.** The sliding window (now per-run, not per-sequence — see [Findere z-window filter](#findere-z-window-filter)) only ever runs after all partitions have contributed their hits to `by_genome`. Applying it per superkmer would produce false negatives at superkmer boundaries, where the z-window spans two superkmers.
|
**Findere requires full-sequence aggregation.** `apply_findere` is applied once per sequence on the complete s-mer result vector, after all partitions have contributed. Applying it per superkmer would produce false negatives at superkmer boundaries, where the z-window spans two superkmers.
|
||||||
|
|
||||||
Batches are processed in parallel via `obipipeline` workers; the `--threads` flag controls the number of worker threads.
|
Batches are processed in parallel via `obipipeline` workers; the `--threads` flag controls the number of worker threads.
|
||||||
|
|
||||||
@@ -60,46 +44,25 @@ Batches are processed in parallel via `obipipeline` workers; the `--threads` fla
|
|||||||
|
|
||||||
## Findere z-window filter
|
## Findere z-window filter
|
||||||
|
|
||||||
For approximate index modes, the index physically stores s-mers of size `s = k_user − z + 1`; `idx.kmer_size()` (bound to `k` in `process_chunk`) is this physically-indexed s-mer size, so decomposing the query at `k` naturally produces s-mer results.
|
For approximate index modes, the index physically stores s-mers of size `s = k_user − z + 1`. At query time, `set_k(s)` is in effect, so queries naturally produce s-mer results. `apply_findere` then aggregates z consecutive s-mer results into one k_user-mer answer:
|
||||||
|
|
||||||
The z-window aggregation is **sparse**, per genome, implemented in `sparse_findere_for_genome` (`query.rs`) — a run-detection pass followed by a monotone-deque sliding-window minimum scoped to each run, not a dense scan over every s-mer position of every sequence:
|
```rust
|
||||||
|
fn apply_findere(
|
||||||
```
|
results: &[Option<Box<[u32]>>], // N s-mer results
|
||||||
sparse_findere_for_genome(hits, z, presence, threshold):
|
z: usize,
|
||||||
// hits: raw (seq_idx, pos_smer, value) triples for this genome, as delivered
|
n_genomes: usize,
|
||||||
// by query_partition_with's QueryHit::Value — only ever nonzero entries;
|
) -> Vec<Option<Box<[u32]>>> // N − z + 1 k_user-mer results
|
||||||
// a position with no hit for this genome simply has no entry at all.
|
|
||||||
sort hits by (seq_idx, pos_smer)
|
|
||||||
|
|
||||||
for each maximal run of consecutive pos_smer values within the same seq_idx:
|
|
||||||
dq: VecDeque<(run-relative index, value)>
|
|
||||||
for k, (_, pos, value) in enumerate(run):
|
|
||||||
maintain dq monotone non-decreasing (pop back while back.value >= value)
|
|
||||||
push (k, value)
|
|
||||||
evict dq entries with run-relative index <= k - z
|
|
||||||
if k + 1 >= z:
|
|
||||||
win_min = dq.front().value
|
|
||||||
if win_min > 0:
|
|
||||||
pos_out = pos + 1 - z
|
|
||||||
confirmed.push((seq_idx, pos_out, adjust(win_min)))
|
|
||||||
return confirmed
|
|
||||||
```
|
```
|
||||||
|
|
||||||
A window can only be confirmed (`win_min > 0`) when all `z` s-mers in it are present *and* nonzero for this genome — which, by construction, can only happen strictly inside one contiguous run of hits (any gap — an absent or zero-valued s-mer — forces `win_min = 0` for every window spanning it, exactly matching the old dense scan's "not in index counts as 0" rule, just never materialising the zero). The deque logic is otherwise identical to the pre-sparsification version; it's scoped to run-relative indices instead of the whole sequence.
|
Input length N (s-mers), output length N − z + 1 (k_user-mers).
|
||||||
|
|
||||||
This runs once per genome that has at least one hit in the chunk (`process_chunk` iterates `by_genome`, one `Vec<(seq_idx, pos_smer, value)>` per genome, built from `QueryHit::Value` during the partition loop — genomes with zero hits in this chunk have an empty `Vec` and cost nothing beyond the iteration itself). Total work is `O(hits log hits)` per genome (the sort) rather than `O(n_smers)` per genome regardless of hit count — a genuine complexity win on top of the memory one, for the common case where most `(chunk, genome)` pairs have no or few hits.
|
For each genome g independently, a sliding window of size z scans the input. Output position i is confirmed for genome g iff all z values `results[i..i+z][g]` are nonzero (`None` counts as zero for all genomes). The scan is O(n) per genome.
|
||||||
|
|
||||||
Output position `pos_out` is confirmed for genome `g` iff its run produced a nonzero `win_min` — equivalent to "all `z` consecutive s-mer values in the window are nonzero for `g`", same semantics as before.
|
Output values come from `results[i]` (leftmost s-mer of each window); genomes not confirmed are zeroed. If all genomes are zero, the position is returned as `None`.
|
||||||
|
|
||||||
**The value reported per confirmed position is the window minimum, not the leftmost s-mer's raw value** — unchanged from the dense version. For presence indexes (0/1 values) this is equivalent to a logical AND either way. For count indexes it is not: the accumulated count for genome `g` at position `pos_out` is the minimum across the window, the weakest link — not the leftmost s-mer's own count. The presence/count adjustment (`u32::from(win_min >= threshold)` vs. raw `win_min`) is applied once, inside `sparse_findere_for_genome`, rather than later during accumulation.
|
**Short sequences**: when the s-mer count is less than z, no complete window can form — `apply_findere` returns an empty vector. K-mers from sequences shorter than k_user are not emitted.
|
||||||
|
|
||||||
**`kmer_missing` bookkeeping is independent of the per-genome sparse structures**, by design (see roadmap point 9): a lightweight dense `SmerIndex` (`in_index: Vec<bool>`, sized `total_smers` — **not** multiplied by `n_genomes`) is populated from `QueryHit::Found` during the partition loop, one entry per hit k-mer regardless of which genome(s) it matched. A position with no genome confirmed counts as `kmer_missing` iff the leftmost s-mer of that window is absent from `SmerIndex` entirely (see [`kmer_missing` semantics](#kmer_missing-semantics)).
|
**Exact indexes**: `z = 1`, `apply_findere` is a passthrough (output length = input length).
|
||||||
|
|
||||||
**Coverage (`--detail`)** is built by re-scanning each genome's confirmed-hit list (already computed, no extra pass over raw data) and densifying into the `[u32; n_kmers_out]` arrays the JSON output format requires — but only when `--detail` is actually requested; the sparse structures cost nothing extra when it isn't.
|
|
||||||
|
|
||||||
**Short sequences**: when a sequence's s-mer count is less than `z`, its run(s) — if any hits exist at all — can never reach length `z`, so no window is ever confirmed for it; no k_user-mer is emitted, same outcome as the dense version's `n_kmers_out == 0` early-skip, reached here as a natural consequence rather than a separate check.
|
|
||||||
|
|
||||||
**Exact indexes**: `z = 1`, every single-hit "run" of length 1 immediately satisfies `k + 1 >= z`, so every hit is its own confirmed window with `win_min` equal to its own value — a passthrough, as before.
|
|
||||||
|
|
||||||
### Effective z at query time
|
### Effective z at query time
|
||||||
|
|
||||||
@@ -122,17 +85,14 @@ The `-z` CLI option overrides the index metadata value. A higher z increases str
|
|||||||
|
|
||||||
### `QueryLayer` variant selection
|
### `QueryLayer` variant selection
|
||||||
|
|
||||||
`QueryLayer::open` (`obikpartitionner/src/query_layer.rs:28-45`) only ever returns two variants — `Presence` or `Count`, checked in this order:
|
`QueryLayer::open` in `query_layer.rs` selects the data matrix to pair with `MphfLayer`:
|
||||||
|
|
||||||
| Order | Condition | Variant | Data returned per k-mer |
|
| Condition | Variant | Data returned per k-mer |
|
||||||
|---|---|---|---|
|
|---|---|---|
|
||||||
| 1 | `with_counts=true` and `counts/` exists | `Count` | raw count per genome |
|
| `with_counts=true` and `counts/` exists | `Count` | raw count per genome |
|
||||||
| 2 | (else) `presence/` exists, or `counts/` doesn't exist at all | `Presence` | see below |
|
| `presence/` exists | `Presence` | 0/1 per genome (bit matrix) |
|
||||||
| 3 | (else — `counts/` exists, `presence/` doesn't, `with_counts=false`) | `Count` | counts used as-is |
|
| only `counts/` exists | `Count` | counts used as-is |
|
||||||
|
| neither exists | `SetOnly` | 1 for every genome |
|
||||||
There is no `QueryLayer::SetOnly` variant. The "no on-disk matrix at all" case is handled one level down: `Presence` wraps `PersistentBitMatrix`, whose own `open()` (`obicompactvec/src/bitmatrix.rs:260-288`) auto-detects among **three** internal representations — `Packed` (`presence/matrix.pbmx`), `Columnar` (`presence/meta.json`), or `Implicit { n_rows, n_cols }` when neither file exists (built from `layer_meta.json`, `fill_row` returning all-`1`s without touching disk). This is where "1 for every genome" actually happens — not at the `QueryLayer` level.
|
|
||||||
|
|
||||||
**Worth double-checking, not confirmed as a bug**: `PersistentBitMatrix::open`'s `Implicit` branch constructs `Implicit { n_rows: meta.n, n_cols: 1 }` — `n_cols` is hardcoded to `1`, not to the layer's actual `n_genomes`. `fill_row` for `Implicit` only writes `buf[..1]`, leaving the rest of a longer `n_genomes`-sized buffer untouched (zeroed by the caller beforehand). If this path is ever reached for a layer covering more than one genome, only genome index 0 would read as present. Whether that's reachable in practice (layers might always be single-genome when they fall back to `Implicit`) wasn't verified here — flagging for follow-up, not fixing.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -163,7 +123,7 @@ Coverage reflects confirmed k_user-mers only. The vectors are emitted in the JSO
|
|||||||
|
|
||||||
## `kmer_missing` semantics
|
## `kmer_missing` semantics
|
||||||
|
|
||||||
`kmer_missing` counts k_user-mer positions where the leftmost s-mer of the window (`smer_index.is_in_index(seq_idx, pos)`, `SmerIndex`) is `false` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero (but the leftmost one is present) are not counted as missing — the leftmost s-mer being present is used as proxy for index membership.
|
`kmer_missing` counts k_user-mer positions where the first s-mer (`seq_results[seq_idx][pos]`) is `None` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero are not counted as missing (the first s-mer being present is used as proxy for index membership).
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -192,8 +152,8 @@ Genome keys follow the iteration order of `meta.genomes`.
|
|||||||
| Key | Type | Condition | Semantics |
|
| Key | Type | Condition | Semantics |
|
||||||
|---|---|---|---|
|
|---|---|---|---|
|
||||||
| `kmer_count` | int | always | k-mers confirmed (post-Findere) with at least one genome match |
|
| `kmer_count` | int | always | k-mers confirmed (post-Findere) with at least one genome match |
|
||||||
| `kmer_missing` | int | `--count-missing` | k-mers absent from the index entirely (leftmost s-mer of the window not found) |
|
| `kmer_missing` | int | `--count-missing` | k-mers absent from the index entirely (pre-Findere None) |
|
||||||
| `kmer_strict_matches` | object | always | per-genome accumulated value, non-zero entries only (label → count or 0/1) |
|
| `kmer_strict_matches` | object | always | per-genome accumulated value (label → count or 0/1) |
|
||||||
| `coverage` | object | `--detail` | per-genome array of per-position contributions (label → [u32]) |
|
| `coverage` | object | `--detail` | per-genome array of per-position contributions (label → [u32]) |
|
||||||
|
|
||||||
`kmer_count + kmer_missing` ≤ total k_user-mers in the sequence. The gap corresponds to k_user-mers whose z-window was not fully confirmed (at least one s-mer absent or zero for all genomes) but whose first s-mer was present in the index.
|
`kmer_count + kmer_missing` ≤ total k_user-mers in the sequence. The gap corresponds to k_user-mers whose z-window was not fully confirmed (at least one s-mer absent or zero for all genomes) but whose first s-mer was present in the index.
|
||||||
@@ -205,7 +165,7 @@ Genome keys follow the iteration order of `meta.genomes`.
|
|||||||
```
|
```
|
||||||
obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
||||||
[--force-presence] [--presence-threshold <n>]
|
[--force-presence] [--presence-threshold <n>]
|
||||||
[-z <z>] [-T <threads>] [--chunk-size <MiB>]
|
[-z <z>] [-T <threads>]
|
||||||
<query.fa> [<query2.fa> ...]
|
<query.fa> [<query2.fa> ...]
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -217,7 +177,6 @@ obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
|||||||
| `--force-presence` | off | Report 0/1 per genome regardless of index counts |
|
| `--force-presence` | off | Report 0/1 per genome regardless of index counts |
|
||||||
| `--presence-threshold` | 1 | Minimum count to declare genome present |
|
| `--presence-threshold` | 1 | Minimum count to declare genome present |
|
||||||
| `-T` / `--threads` | all CPUs | Worker threads |
|
| `-T` / `--threads` | all CPUs | Worker threads |
|
||||||
| `--chunk-size` | auto (from available RAM and thread count) | I/O chunk size in MiB — see [Future work, point 3](#throughput--parallelism--identified-potential-not-yet-implemented) for why the auto-sizing formula currently under-estimates memory on indexes with many genomes |
|
|
||||||
|
|
||||||
`--mismatch` is accepted but currently ignored with a warning on stderr.
|
`--mismatch` is accepted but currently ignored with a warning on stderr.
|
||||||
|
|
||||||
@@ -228,179 +187,3 @@ obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
|||||||
- **`--mismatch`**: 1-mismatch approximate matching — generate `3·k` single-substitution variants per k-mer, look each up independently.
|
- **`--mismatch`**: 1-mismatch approximate matching — generate `3·k` single-substitution variants per k-mer, look each up independently.
|
||||||
- **Read classification** (`--classify`): assign each read to the genome with the highest match score.
|
- **Read classification** (`--classify`): assign each read to the genome with the highest match score.
|
||||||
- **Whitelist / blacklist filtering**: threshold-based accept/reject on per-genome match scores.
|
- **Whitelist / blacklist filtering**: threshold-based accept/reject on per-genome match scores.
|
||||||
|
|
||||||
### Throughput & parallelism — identified potential (not yet implemented)
|
|
||||||
|
|
||||||
Observed on a 192-core (8×24 NUMA) machine: `query` uses ~10 cores or fewer, and the default chunk size gets the process OOM-killed. Root causes and candidate fixes, in dependency order:
|
|
||||||
|
|
||||||
**1. Single-threaded I/O source (main core-utilization bottleneck).**
|
|
||||||
`run()` builds `all_chunks` via `paths.into_iter().flat_map(read_sequence_chunks_sized(...))` and passes it directly as the `input` iterator to `pipe.apply()`. In `obipipeline::Pipe::apply` (`scheduler.rs`), `input.next()` is called exclusively from the dedicated source thread — so file opening, decompression, and FASTA/FASTQ chunk-boundary parsing for *all* input files run serially in one thread, regardless of `--threads`. Compare with `steps::scatter` (used by `index`) and `cmd/superkmer.rs`: there, file opening + streaming is itself a `Flat` pipeline stage (`||?`), executed across the `n_workers` pool, with `obipipeline::throttle(paths, max_open)` bounding concurrently-open files in the source thread. That pattern parallelises I/O across files (and NUMA nodes); `query.rs` cannot.
|
|
||||||
Fix direction: restructure `query`'s pipe with an initial `Flat` stage analogous to `scatter`'s, opening/chunking files across workers instead of in `flat_map`.
|
|
||||||
|
|
||||||
**2. Gzip decompression is inherently single-threaded per file.**
|
|
||||||
`niffler`/`flate2` (used by `xopen`) do standard DEFLATE, which has no parallel-decodable structure for an arbitrary stream. Fix (1) parallelises *across* files but not *within* one large gzip file. Parking a possible fix (`rapidgzip-rs`) is tracked in [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen).
|
|
||||||
|
|
||||||
**3. Chunk-size memory formula ignores `n_genomes`.**
|
|
||||||
`chunk_bytes = available_memory_bytes() / (n_workers * 16)` (`query.rs:407-414`) assumes a fixed ~8–16× overhead per raw input byte. But `KmerResults::new` (`query.rs:165-179`) allocates `data: Vec<u32>` sized `total_kmers_in_chunk × n_genomes` — dense, **for every k-mer position in the chunk, hit or not** — plus `win_min` and (with `--detail`) `cov`, same scaling. Real per-chunk memory is `O(n_genomes)`, not constant; the formula doesn't know `n_genomes` at all. This is the direct cause of the OOM kill on indexes with many reference genomes.
|
|
||||||
|
|
||||||
**4. MPHF lookup and matrix-row fetch are fused, not staged.**
|
|
||||||
`QueryLayer::find_into` (`obikpartitionner/src/query_layer.rs:48-67`) does the MPHF `find` *and* the `fill_row` matrix read in one call per k-mer, inside a single-threaded loop (`query_partition_with`). There is no separation between "is this k-mer indexed" (cheap, `O(1)`, independent of `n_genomes`) and "what are its per-genome values" (the expensive, `n_genomes`-scaling part).
|
|
||||||
|
|
||||||
**5. Dereplication should happen at k-mer granularity, directly — not via an intermediate superkmer-level dedup.**
|
|
||||||
`QueryBatch::from_records` currently dereplicates at the *superkmer* level (`HashMap<RoutableSuperKmer, Vec<SKDesc>>`, `query.rs:112`). This misses redundancy between k-mers shared by *different* superkmers (read overlaps, repeats, a SNP splitting an otherwise-identical run). Superkmer *construction* (`SuperKmerIter`) stays mandatory — it is the mechanism that computes minimizers/partition routing, not an optional dedup layer — but the dedup structure built on top of it should key directly on `CanonicalKmer`, in the same pass: `HashMap<CanonicalKmer, Vec<(seq_idx, pos)>>`. This also means the MPHF `find` itself runs once per **distinct** k-mer instead of once per occurrence — a win independent of the matrix-fetch cost below.
|
|
||||||
|
|
||||||
**6. Stage 1 output: bucket confirmed hits by layer, keyed by MPHF slot.**
|
|
||||||
For each unique canonical k-mer, MPHF lookup across a partition's layers stops at the first match (`query_partition_with:105-111`) — a k-mer belongs to at most one layer. So stage 1's output can be reshaped directly into:
|
|
||||||
```
|
|
||||||
HashMap<layer_idx, HashMap<slot, Vec<(seq_idx, pos)>>>
|
|
||||||
```
|
|
||||||
replacing the `CanonicalKmer` key by the resolved `slot` (compact integer, and exactly what stage 2 needs to address the matrix). K-mers matching no layer simply have no entry here (they still count toward `in_index`/`kmer_missing` bookkeeping, which stays `O(1)` per position, independent of `n_genomes`).
|
|
||||||
|
|
||||||
**7. Partition-level parallelism is currently absent — and a NUMA-aware mechanism for exactly this already exists, unused, in `obikindex`.**
|
|
||||||
`process_chunk`'s partition loop (`query.rs:250-278`, `for (part_idx, part_sks) in by_part.iter().enumerate()`) processes every partition of a chunk sequentially on the single worker thread that owns that chunk. This is a parallelism axis on its own, independent of the column question below.
|
|
||||||
More importantly: `docmd/architecture/numa_partition_runner.md` and `numa_worker_pools.md` document `PartitionRunner` (`obikindex/src/numa.rs`), **already implemented** and already used by `merge.rs`, `index.rs` (`build_layers`), `select.rs`, `reindex.rs`, `rebuild.rs` — one controller thread per NUMA node, a Rayon pool pinned to that node's CPUs (`hwlocality`, `numa` feature, default-on in `obikindex/Cargo.toml`), adaptive worker activation driven by *both* a CPU-efficiency signal and an I/O-throughput signal (`CpuSample`/`IoSample`, `/proc/self/io` on Linux). It exists precisely because a naive `into_par_iter()` on the global Rayon pool measurably degrades ×60 on this codebase's own 192-core/8-NUMA reference machine (`numa_worker_pools.md`, § Problem) once workers contend for cross-socket memory bandwidth on shared mmap'd/hashed structures — exactly the shape of the matrix-column scan in point 8 below.
|
|
||||||
`obikmer` already depends on `obikindex` (`obikmer/Cargo.toml`, for `KmerIndex`), so `PartitionRunner` is directly reachable from `cmd/query.rs` — no new dependency. Both the partition-level loop and (see point 8) the genome-column scan should be driven through it rather than through ad-hoc `rayon::into_par_iter()`, to avoid reproducing the already-measured-and-fixed contention problem. Also relevant: the "CPU-only signal stalls on I/O-bound stages" issue documented for `pack_matrices` (mmap-heavy, page-fault-bound) applies just as much to a column-major mmap scan over persistent matrices — reuse the existing dual CPU/IO activation signal rather than re-deriving one.
|
|
||||||
>
|
|
||||||
> **Correction from implementation (Phase 4 below)**: this turned out not to be viable as described. `PartitionRunner::run()`'s actual body spawns roughly one OS thread per worker slot across every NUMA node **on every call** (confirmed by reading `numa.rs`, not just its doc comments) — fine for the one-call-per-command-invocation batch usage in `merge`/`build_layers`, but `query_partition_with` runs once per `(chunk, partition)`, far too frequently to absorb that spawn cost. Partition-level parallelism via `PartitionRunner` is deferred, not implemented. See Phase 4's "What did not ship, and why" for the detail.
|
|
||||||
|
|
||||||
**8. Stage 2: column-major matrix fetch, parallel across genome columns — via `PartitionRunner`, not naive `rayon`.**
|
|
||||||
Both persistent matrix formats are column-oriented on disk: `ColumnarCompactIntMatrix`/`ColumnarBitMatrix` (`obicompactvec/src/{intmatrix,bitmatrix}.rs`) mmap one file per genome column; `PackedCompactIntMatrix`/`PackedBitMatrix` mmap one region-offset per column in a single file. `fill_row(slot, buf)` as used today (`query.rs:262-272` via `on_hit`) reads **one slot across all `n_genomes` columns** per hit — the worst possible access pattern for this layout (up to `n_genomes` scattered mmap regions touched per single k-mer).
|
|
||||||
Better: for each layer, walk the matrix **column by column** (genome by genome): for each genome, scan the `slot` keys collected in step 6 for that layer and call `col.get(slot)`, keeping only nonzero results, and broadcast to the associated `(seq_idx, pos)` list. Total `get()` calls are unchanged (`n_hits × n_genomes` in the worst case) — the win is locality (sequential access within one mmap'd column at a time, not scattered across all columns per hit), not fewer operations.
|
|
||||||
Columns are independent (read-only, disjoint mmap regions) → embarrassingly parallel across genomes, *but* — per point 7 — `obicompactvec`'s existing `into_par_iter()` over `0..n_cols` (`sum()`, `count_nonzero()`, pairwise distance matrices) is the **naive, unpinned** pattern the rest of the codebase is actively migrating away from, not a model to copy here. Route this through `PartitionRunner` (or the same NUMA-pool machinery) instead. Two things to settle when this is designed: how the partition axis (point 7), the column axis, and the existing chunk-level `n_workers` `obipipeline` pool compose without oversubscribing the machine (three different concurrency mechanisms — raw-thread pipe workers, `PartitionRunner`'s pinned Rayon pools, and whatever drives the column scan — need a single reconciled thread budget, not three independent ones); and the threshold below which per-column dispatch overhead outweighs the gain (small `n_genomes` or small per-layer hit counts) — to be measured, not assumed.
|
|
||||||
>
|
|
||||||
> **Correction from implementation (Phase 4 below)**: column-major fetch is implemented — but as a plain sequential loop, not parallelised via `PartitionRunner`. Same reason as point 7's correction above. The column-major *locality* win (the actual claim of this point) does not depend on adding parallelism on top of it, and is validated independently. Column-level parallelism is deferred pending a mechanism that fits this call frequency (candidates noted in Phase 4).
|
|
||||||
|
|
||||||
**9. Sparse per-genome representation, fed directly to Findere.**
|
|
||||||
Stage 2's output should be `HashMap<genome_idx, Vec<(seq_idx, position, count)>>`, **sorted by `(seq_idx, position)`** once collected, instead of a dense `KmerResults`-style matrix — the key must carry `seq_idx`, not just `genome_idx`, because a chunk batches many sequences and `position` is only meaningful within one; a plain `Vec<(position, count)>` per genome would silently mix positions from different sequences and corrupt the sliding-window scan. This bounds retained memory by actual nonzero hits on both axes (position sparsity from non-matching k-mers, genome sparsity from a matched k-mer typically belonging to only a handful of genomes out of possibly many). The Findere sliding-window (`process_chunk`, the `win_min`/deque loop) would need reworking to run per `(sequence, genome)` over its sparse, sorted `(position, count)` list — detect runs of ≥`z` consecutive positions, window-min within each run — instead of today's dense `O(total_kmers × n_genomes)` scan. This is also a genuine complexity win (`O(hits log hits)` per genome vs. dense scan), not just memory.
|
|
||||||
**Not covered by this sparsification**: `--detail`'s `cov` accumulator (`query.rs:304-308`) has the identical `n_genomes`-dense scaling problem and wasn't folded into points above. It doesn't need to be retained densely throughout processing, though — only the final JSON serialization (`emit_batch`) requires a dense `[u32]` per `(seq, genome)`, and only for the sequences actually being output with `--detail`. Densification can stay a late, output-time-only step, reconstructed from the sparse per-genome lists.
|
|
||||||
|
|
||||||
**Secondary patterns available from `scatter.rs`/`superkmer.rs`, not yet in `query.rs`:**
|
|
||||||
- `throttle()` + `CommonArgs::effective_max_open()` to bound concurrently-open input files (query.rs defines its own `QueryArgs`, doesn't reuse this).
|
|
||||||
- Progress bar with EMA throughput + live active-worker gauges (`obisys::spinner`, `flat_active`/`transform_active` counters) — diagnostic value for locating the bottleneck.
|
|
||||||
- `obisys::Reporter`/`Stage::start`/`stop` timing per phase (used by `index`, `filter`; absent from `query`).
|
|
||||||
|
|
||||||
None of this is implemented yet — parked here as a coherent roadmap while the design is discussed further. Suggested dependency order: (1) I/O parallelism → (3) genome-aware chunk sizing → (4)–(9) staged/k-mer-deduped/NUMA-aware-partition-and-column-major/sparse query engine (larger refactor, biggest structural payoff — reuses `PartitionRunner` rather than inventing a new parallelism mechanism) → (2) parallel gzip (separate, orthogonal, tracked in chunkreader.md) → secondary diagnostics patterns.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Implementation plan
|
|
||||||
|
|
||||||
Concrete, phased translation of the roadmap above. Phases 0–2 are small, independent, low-risk, and each individually testable against current `query` output — land them first, in order, and measure on the reference 192-core/8-NUMA machine before deciding whether phases 3–5 (the staged/sparse engine, the larger structural payoff) are still worth their cost. Phases 3–5 are one coordinated change spanning `obikmer`, `obikpartitionner`, and `obicompactvec` — they should not be split across releases mid-way, because the intermediate state (e.g. k-mer-level dedup feeding the old dense `KmerResults`) has no correctness or performance benefit on its own. Phase 6 is unrelated to phases 0–5 and can happen any time, independently, if `rapidgzip-rs` is validated (see [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen)).
|
|
||||||
|
|
||||||
Instrumentation is deliberately sequenced *before* the I/O fix (reordering the roadmap's own listed order), because every later phase's justification rests on a measurement ("to be measured, not assumed" appears throughout the roadmap above) — without it, phases 3–5 would be undertaken on faith.
|
|
||||||
|
|
||||||
Performance measurement on the reference 192-core/8-NUMA machine is done by the project owner, not from this development environment (macOS, 16 cores — `PartitionRunner`'s NUMA pinning is Linux-only, so even phase 4's mechanism can't be functionally exercised for its actual purpose here). Each phase below is therefore written to be *self-measuring*: the debug-level logging it adds must be enough, on its own, to judge whether that phase's algorithmic choice paid off from a cluster run's logs, without needing to attach a profiler.
|
|
||||||
|
|
||||||
### Conventions applied to every phase below
|
|
||||||
|
|
||||||
**Debug logging.** Every phase that changes an algorithmic choice (not phase 0, which *is* the logging) adds `tracing::debug!`/`trace!` at points that let a cluster run's logs answer "did this help": counts, ratios, and timings that quantify the specific claim that phase makes — e.g. phase 3 must log how many MPHF `find` calls were saved by k-mer-level dedup (the whole justification for that phase), phase 4 must log per-column scan timings, phase 5 must log actual retained-memory / sparsity ratios achieved. Prefer one structured `debug!` per chunk (fields, not prose) over free-text — the cluster logs will be the only evidence available for judging these choices, so they need to be grep/awk-able, not just readable.
|
|
||||||
|
|
||||||
**Unit tests.** This project's convention (`obiread`, `obikseq`, `obidebruinj`, `obicompactvec`, `obilayeredmap`, `obiskio`, `obifastwrite`) is `#[cfg(test)] #[path = "tests/<name>.rs"] mod tests;` at the bottom of the source file, with the actual test code in a sibling `src/tests/<name>.rs`. Neither `obikmer` nor `obikpartitionner` (the two crates phases 3 and 5 touch most) currently have a `src/tests/` directory at all — this needs creating, following the existing pattern exactly, not inventing a new one.
|
|
||||||
|
|
||||||
**Workflow (`jj`).** Work happens in a fresh `jj` commit, easy to abandon. `jj new` between phases is reasonable where it helps isolate a phase for review, but only when the working copy compiles at that point (project convention) — phase 3's internal sub-steps (batch dedup change, then `query_layer.rs` split, then the new return shape) will likely not each compile independently since they're one coupled change, so treat "commit boundary" and "plan phase boundary" as related but not forced to match 1:1; use judgement per phase rather than mechanically splitting on every bullet.
|
|
||||||
|
|
||||||
### Phase 0 — Instrumentation (prerequisite for measuring every later phase)
|
|
||||||
|
|
||||||
**Goal**: make core utilization, throughput, and per-stage timing visible on a real run, so phases 1–5 can be justified with numbers instead of assumption.
|
|
||||||
|
|
||||||
- `obikmer/src/cmd/query.rs`: wrap `run()`'s main loop with `obisys::Reporter`/`Stage::start("query")`/`.stop()`, printed at the end via `rep.print()` — same pattern as `index.rs`/`filter.rs`.
|
|
||||||
- Add an `obisys::spinner("query")` progress bar around the `pipe.apply(...)` loop, with an EMA throughput readout (bases/s or k-mers/s, mirroring `steps::scatter`'s `ema_rate` computation, `scatter.rs:88-118`) and live gauges for "chunks in flight" / "workers busy" — reuse the `AtomicU32` counter pattern from `scatter.rs` (`flat_active`, `transform_active`) rather than inventing a new one.
|
|
||||||
- Add `max_open_files: Option<usize>` to `QueryArgs` and a `effective_max_open()` method mirroring `CommonArgs::effective_max_open()` (`obikmer/src/cli.rs:90-94`) — needed by phase 1's `throttle()` call. (`QueryArgs` can't just embed `CommonArgs` — it doesn't take `kmer_size`/`minimizer_size`/`partitions`/`level_max`/`theta` from the CLI, those come from the index metadata — so this is a small standalone addition, not a flatten.)
|
|
||||||
- Add one structured `debug!` per `process_chunk` call: chunk byte size, sequence count, s-mer count, wall time, and (once later phases exist) the fields they add — this single log line is the baseline every later phase's own logging gets compared against.
|
|
||||||
- **Validation**: none needed beyond "the numbers appear and look sane" — this phase changes no query logic or output.
|
|
||||||
- **Deliverable used by every phase below**: a before/after throughput and core-utilization measurement on the reference machine.
|
|
||||||
|
|
||||||
### Phase 1 — Parallel per-file I/O (fixes root cause of low core utilization)
|
|
||||||
|
|
||||||
**Goal**: file opening, decompression, and chunk-boundary parsing run across the `n_workers` pool instead of serially in the pipe's dedicated source thread.
|
|
||||||
|
|
||||||
- `obikmer/src/cmd/query.rs`:
|
|
||||||
- Replace the `paths.into_iter().flat_map(read_sequence_chunks_sized(...))` construction (current `run()`, building `all_chunks`) with `obipipeline::throttle(paths.into_iter(), args.effective_max_open())`, passed as the pipe's `input`.
|
|
||||||
- Add a new `QueryData::Path(PathBuf)` variant (alongside `Chunk`/`Output`) to carry the throttled path through the pipe's type-erasure mechanism.
|
|
||||||
- Add a new **first** pipe stage, `Flat`/fallible (`||?`), modeled on `scatter.rs:60-86` and `superkmer.rs:54-65`: given a `Throttled<PathBuf>`, call `read_sequence_chunks_sized(path, chunk_bytes)` and yield each `Rope` chunk, keeping `pw.guard` alive until the file's iterator is exhausted (reuse or adapt `scatter.rs`'s `GuardedIter` wrapper — same lifetime problem, same fix).
|
|
||||||
- The existing `process_chunk` transform stage becomes the pipe's **second** stage, unchanged in its own logic — it still receives one `Rope` chunk at a time, just no longer all coming from one serial source.
|
|
||||||
- `make_pipe!` invocation grows from one stage (`Chunk => Output`) to two (`Path => Chunk => Output`).
|
|
||||||
- Log, per file: time spent waiting on the `throttle()` slot (queueing due to `max_open`), and time spent opening/decompressing/producing the first chunk — this is what directly proves (or disproves) that I/O is now spread across workers instead of serialized.
|
|
||||||
- **Validation**: run `query` on a small multi-file input, diff output against the pre-change version — content must be identical; **record order across files is not guaranteed to be preserved** even before this change (chunk-level dispatch across `n_workers` already reorders completions), so the diff must be order-insensitive (sort by read id, or compare as sets) if it wasn't already.
|
|
||||||
- **Measure**: core utilization on the reference machine with several large input files, compare against phase 0's baseline.
|
|
||||||
|
|
||||||
### Phase 2 — Genome-aware chunk-size formula (fixes OOM)
|
|
||||||
|
|
||||||
**Goal**: `chunk_bytes` reflects actual per-chunk memory (`O(n_genomes)`), not a fixed multiplier.
|
|
||||||
|
|
||||||
- `obikmer/src/cmd/query.rs`, `run()`: `n_genomes` and `args.detail` are already computed above the `chunk_bytes` calculation (`n_genomes` at the top of `run()`, before line 407 in the current file) — reorder if needed, then replace:
|
|
||||||
```rust
|
|
||||||
let computed = avail / (n_workers as u64 * 16);
|
|
||||||
```
|
|
||||||
with a formula that scales the divisor by `n_genomes` (and roughly doubles it when `--detail` is set, since `cov` duplicates the per-genome accumulation): e.g. `per_chunk_multiplier = base_overhead + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * if detail { 2 } else { 1 }`, replacing the flat `16`. `BYTES_PER_KMER_PER_GENOME` should be derived from `KmerResults`'s actual layout (`4` bytes per `u32` entry in `data`, plus the `bool` in `in_index`, plus `win_min`'s equal-sized buffer) rather than guessed.
|
|
||||||
- `args.chunk_size` (manual `--chunk-size` override) keeps taking priority, unchanged.
|
|
||||||
- Log the resolved `chunk_bytes`, `n_genomes`, and the estimated peak per-chunk memory (`chunk_bytes` × the same multiplier used to derive it) once at startup — lets a cluster run confirm the estimate was actually respected, not just that the process didn't get OOM-killed (which could also happen to be true for the wrong reason).
|
|
||||||
- **Validation**: build a test index with a large `n_genomes` (e.g. hundreds), run `query` with default chunk sizing under a memory limit (`ulimit -v` or a cgroup), confirm it no longer gets OOM-killed and that memory scales as predicted when `n_genomes` grows.
|
|
||||||
- **Note**: this phase is superseded once phase 5 lands (sparse retained memory no longer scales with `n_genomes × total_kmers` at all) — but it's needed immediately regardless, since phases 3–5 are a bigger, riskier change and users need a working `query` in the meantime.
|
|
||||||
|
|
||||||
### Phase 3 — K-mer-level dereplication, staged MPHF/matrix lookup
|
|
||||||
|
|
||||||
**Goal**: replace superkmer-level dedup with k-mer-level dedup (roadmap point 5), and split the fused MPHF-find/matrix-fetch (point 4) so stage 1's output is bucketed by layer and MPHF slot (point 6).
|
|
||||||
|
|
||||||
- `obikmer/src/cmd/query.rs`:
|
|
||||||
- Replace `QueryBatch::from_records`'s dedup map (`HashMap<RoutableSuperKmer, Vec<SKDesc>>`, current `query.rs:112`) with a per-partition `HashMap<CanonicalKmer, Vec<(seq_idx: u32, pos: u32)>>`, built in the same `SuperKmerIter` pass: superkmer construction and partition routing (`part_idx` from the superkmer's minimizer hash) are unchanged, only the granularity of what gets deduplicated changes — each `CanonicalKmer` within a superkmer is inserted individually instead of the whole superkmer being the dedup key.
|
|
||||||
- **Verified**: `CanonicalKmer` (`obikseq/src/kmer.rs:390`, `pub type CanonicalKmer = CanonicalKmerOf<KLen>`) — the underlying `CanonicalKmerOf<L>` derives `Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash` (`kmer.rs:269`). Usable as a `HashMap`/`HashSet` key as-is, no change needed.
|
|
||||||
- `obikpartitionner/src/query_layer.rs`:
|
|
||||||
- Split `QueryLayer::find_into` (`query_layer.rs:48-67`) into two methods: `find_slot(&self, kmer: CanonicalKmer) -> Option<usize>` (MPHF only, no matrix touch) and keep `fill_row` as-is for phase 4 to call later.
|
|
||||||
- Replace `query_partition_with`'s inner loop (`query_layer.rs:103-113`) with a version that, for each unique `CanonicalKmer`, calls `find_slot` across the partition's layers (stopping at first hit, same as today), and instead of immediately filling a row, records `(layer_idx, slot)`.
|
|
||||||
- New return shape for the partition-level query, replacing today's `on_hit(sk_idx, kmer_idx, row)` callback: `HashMap<layer_idx, HashMap<slot, Vec<(seq_idx, pos)>>>` (roadmap point 6) — built directly from the k-mer dedup map's `Vec<(seq_idx,pos)>` values, keyed by the resolved slot instead of the k-mer.
|
|
||||||
- **This phase alone has no throughput benefit yet** (matrix fetch still happens, just deferred) beyond the k-mer-level dedup itself (fewer MPHF calls when queries have overlapping/repeated k-mers) — its purpose is to produce the input phase 4 needs. Land phase 3+4 together, not phase 3 alone, per the "don't split 3–5 across releases" note above.
|
|
||||||
- Log, per chunk: total k-mer occurrences vs. unique `CanonicalKmer` count (the dedup ratio — the entire justification for this phase) and the resulting MPHF `find` call count. If the dedup ratio is close to `1.0` on real query data (little redundancy), that's the cluster run telling us this phase wasn't worth it — the logging needs to be able to say that, not just confirm the happy path.
|
|
||||||
- **Unit tests**: create `obikmer/src/cmd/tests/query.rs` (new `src/tests/` dir for this crate, following the project's `#[cfg(test)] #[path = "tests/query.rs"] mod tests;` convention) and `obikpartitionner/src/tests/query_layer.rs` (likewise new for this crate). Cover: the k-mer-level dedup map construction on synthetic sequences with known repeated/overlapping k-mers (assert unique-kmer count and occurrence lists); the `find_slot`/bucket-by-layer-and-slot construction against a small hand-built `QueryLayer` fixture, asserting the `(layer_idx, slot, seq_idx, pos)` tuples match what the old per-occurrence loop would have produced.
|
|
||||||
|
|
||||||
### Phase 4 — Column-major matrix fetch (roadmap points 7–8) — implemented, NUMA parallelism deferred
|
|
||||||
|
|
||||||
**Goal (revised during implementation)**: replace `fill_row`-per-hit (row-major, worst-case mmap locality) with a column-major scan. `PartitionRunner` turned out to be the wrong mechanism for this at this call granularity — see below; the column-major fetch itself is implemented and validated, without it.
|
|
||||||
|
|
||||||
**What shipped:**
|
|
||||||
- `obicompactvec`: the per-column accessors this phase needed **already existed** — `PersistentCompactIntMatrix::col_view(c)` and `PersistentBitMatrix::col_view(c)` are public, and `IntSliceView::get(slot)`/`BitSliceView::get(slot)` are public — the original plan underestimated how much of this plumbing the pairwise-distance code (`dump`/`select`/`stats`) had already required. The one real gap: `PersistentBitMatrix::col_view()` panics on the `Implicit` variant (the documented mono-genome fast path, `bitmatrix.rs`). Added `PersistentBitMatrix::get(c, slot) -> u32` (`bitmatrix.rs`), a non-panicking column-major point lookup that returns `1` for `Implicit` regardless of `c` — the smallest surface needed, not a new `col_get` API from scratch.
|
|
||||||
- `obikpartitionner/src/query_layer.rs`: `query_partition_with` is now two explicit stages, matching roadmap points 6–8: **stage 1** (MPHF-only, per unique k-mer, bucket hits by `(layer_idx, slot)`, emits `QueryHit::Found`) then **stage 2** (per layer with ≥1 hit, column-major: for each genome column `g` in `0..layer.n_cols().min(n_genomes)`, scan that layer's bucketed slots and call `col_value(g, slot)`, emitting `QueryHit::Value(descs, g, value)` on nonzero). `QueryHit` is a single enum delivered through one `FnMut(QueryHit)` callback — an earlier two-closure design (`on_found` + `on_value`) didn't borrow-check, since the caller's single mutable accumulator (`KmerResults`) can't be captured by two separate `FnMut` closures passed to the same call.
|
|
||||||
- `obikmer/src/cmd/query.rs`: `KmerResults::set` (row-major, whole-row-at-once) replaced by `mark_found` (stage 1: flag a position as indexed, independent of any genome's value) and `set_one` (stage 2: write one genome's value at one position). `QueryStats` extended with `n_columns_scanned`/`n_col_get_calls`, logged per chunk.
|
|
||||||
- Total `get()`-equivalent calls are unchanged from the row-major version (`n_hits × n_cols` in the worst case, confirmed by `n_col_get_calls` in the debug log) — the win is locality (sequential access within one layer's column at a time, across `mmap`'d regions, instead of jumping across all columns per hit), exactly as predicted.
|
|
||||||
|
|
||||||
**What did not ship, and why — `PartitionRunner` is architecturally the wrong tool here:**
|
|
||||||
Reading `obikindex/src/numa.rs`'s actual `run()` body (not just its doc comments) shows every call spawns a timer thread **plus one OS thread per worker slot on every NUMA node** (`std::thread::scope` + one `s.spawn()` per node per `max_workers`) — on the 192-core/8-NUMA reference machine, that's on the order of 190+ fresh OS threads spawned **per call**. This is fine for its actual, established usage in this codebase (`merge.rs`, `index.rs`'s `build_layers`): one `PartitionRunner::new()` + one `run()` call per command invocation, amortised over a batch of ~256 long-running partitions. It is not fine for `query`'s call pattern: `query_partition_with` runs once per `(chunk, partition)`, potentially thousands of times per second — spawning ~190 OS threads that often to scan a handful of genome columns would very likely cost far more than the row-major approach it's meant to replace. This is exactly the "resolve empirically, don't assume" composition risk the roadmap flagged, just resolved by reading the mechanism's actual cost before wiring it in, rather than by measuring a regression on the cluster after the fact.
|
|
||||||
The column-major loop in stage 2 is therefore a **plain sequential loop** for now — it captures the whole, provable locality win (roadmap point 8's actual claim) without adding any parallelism mechanism. Genome-column-level parallelism (point 8's "bonus" axis) and partition-level parallelism (point 7) are both deferred — not abandoned. Candidates for a follow-up, once there's a concrete profiling need: (a) `rayon`'s already-warm global pool (`into_par_iter()`) for the column axis specifically — cheap to invoke repeatedly since it doesn't spawn threads per call, though it's the same "naive rayon" pattern `numa_worker_pools.md` warns about for a *different* workload (random pointer-chasing over large hash maps); a column scan's access pattern (sequential reads within one `mmap`'d region) has a different contention profile and hasn't been shown to have the same problem — needs its own measurement, not an assumption either way; (b) restructuring so `PartitionRunner` is invoked once per whole `query` run (or per large batch of chunks) rather than per `(chunk, partition)`, amortising its spawn cost the way `merge`/`build_layers` do — a bigger structural change than this phase's scope.
|
|
||||||
- Log (implemented): `QueryStats::n_columns_scanned`/`n_col_get_calls`, folded into the existing per-chunk `debug!("k-mer dedup + column-major fetch", ...)` line (`query.rs`) alongside phase 3's dedup counters.
|
|
||||||
- **Unit tests**: extended `obikpartitionner/src/tests/query_layer.rs` (phase 3's file) — `query_partition_with`'s empty/missing-index paths updated for the new `QueryStats` fields and single-callback signature.
|
|
||||||
- **Validation performed**: full workspace build + `cargo test --workspace`, zero failures. Functional validation against real indexes: (1) a single-genome index — output byte-identical to pre-phase-4 (same `kmer_count`/`kmer_strict_matches` on every record); (2) the existing 20-genome `benchmark/global_index_presence` index — runs correctly, `n_hits=0` for an unrelated query (expected: no shared k-mers between a plant read and a bacterial reference set), no panics, confirming the `Implicit`/multi-column bounds logic doesn't crash on a real multi-genome, mixed-format index; (3) **the critical correctness case**: built two single-sequence-pair test genomes, merged into one 2-genome index, queried with reads from both — reads from `genomeA` matched **only** `genomeA` (`kmer_count` identical to the pre-dedup occurrence count, zero leakage into `genomeB`'s column) and vice versa. This is the test that would have caught a column-index mixup, an off-by-one in `n_cols`, or cross-genome bleed from the stage-1/stage-2 split — it passed cleanly.
|
|
||||||
- **Not yet done**: the microbenchmark comparing column-major vs. the old row-major access pattern's wall time / page-fault counters on a large-`n_genomes` layer — needs a realistically large multi-genome index and, for the page-fault counters specifically, Linux (not available from this development environment). Left for cluster validation alongside phases 1–3's own pending measurements.
|
|
||||||
|
|
||||||
### Phase 5 — Sparse Findere rework (roadmap point 9)
|
|
||||||
|
|
||||||
**Goal**: replace the dense `KmerResults`/`win_min` sliding-window scan with one operating on phase 4's sparse per-genome output.
|
|
||||||
|
|
||||||
- `obikmer/src/cmd/query.rs`, `process_chunk`:
|
|
||||||
- Remove `KmerResults` (`query.rs:157-202`) and the dense `win_min` allocation (`query.rs:290-291`, sized `max_n_kmers × n_genomes`).
|
|
||||||
- Keep a lightweight dense `in_index: Vec<bool>` per chunk (sized `total_kmers`, independent of `n_genomes`) from phase 3's stage 1 — still needed for `kmer_missing` bookkeeping (leftmost-s-mer-of-window membership test), which phase 4's sparse structure doesn't carry (a k-mer with no genome hit has no entry there at all).
|
|
||||||
- New per-`(seq_idx, genome)` scan: for each genome's `Vec<(seq_idx, pos, count)>` (sorted, per phase 4), group by `seq_idx` (contiguous after sort), then within each sequence's positions detect runs of `pos, pos+1, pos+2, ...` of length ≥ `z`; within each run, the existing monotone-deque window-minimum logic (`query.rs`'s current `dq` loop, conceptually unchanged) applies — but the deque now only scans real entries in the run, never zero-filled gaps.
|
|
||||||
- Update `SeqAcc` accumulation and `emit_batch` to consume this per-genome sparse iteration instead of `results.val`/`results.is_in_index`.
|
|
||||||
- `--detail`/`cov`: build sparsely during the same scan (only positions with a confirmed contribution get an entry), densify into the `[u32]` JSON array only in `emit_batch`, only for genomes/sequences actually being serialized (per roadmap point 9's note, `query.rs:304-308`'s current dense allocation goes away).
|
|
||||||
- Log, per chunk: total sparse entries retained vs. what the old dense `KmerResults` would have allocated (`total_smers × n_genomes`) — the sparsity ratio is this phase's entire reason for existing, so it must be directly visible in the logs, not inferred from process RSS. Also log the run-detection stats (number of runs found, average run length) — a low average run length relative to `z` would mean most positions still fail to form a full window, worth knowing.
|
|
||||||
- **Unit tests**: `obikmer/src/cmd/tests/query.rs` (extended from phase 3) — the property test described below is the primary deliverable here, not an afterthought; write it as an actual `#[test]` (or a small internal fuzz/property-style loop over randomized fixtures if a property-testing crate isn't already a dependency — check before adding one, per this project's dependency-approval rule) rather than a one-off manual comparison.
|
|
||||||
- **Validation — this is the correctness-critical phase**: property-test comparing old (dense, pre-phase-3) and new (sparse) implementations on the same randomized input/index fixtures, asserting identical `kmer_count`, `kmer_missing`, `kmer_strict_matches`, and (with `--detail`) `coverage` for every sequence. Keep both implementations compiled side by side (behind a debug-only flag or a temporary parallel code path) only for the duration of this validation; delete the dense path once parity is confirmed — per this project's own convention, superseded code is not kept "just in case."
|
|
||||||
- **Update `docmd/architecture/query.md` itself**: once this phase lands, the "Findere z-window filter" section (which currently — correctly — describes the dense deque-over-`0..n_smers` scan) needs another pass to describe the sparse run-detection algorithm instead, as already flagged when this phase was discussed.
|
|
||||||
|
|
||||||
**Implemented as planned, no deviations discovered this time.** What shipped:
|
|
||||||
- `KmerResults` removed entirely, replaced by `SmerIndex` (`in_index: Vec<bool>` + `offsets`, unchanged size/purpose, renamed since it's no longer "results" — just the O(1)-per-position "was this k-mer found at all" bookkeeping) and `by_genome: Vec<Vec<(seq_idx, pos, value)>>` (one empty `Vec` per genome until a hit arrives — genomes with zero hits in a chunk cost nothing beyond the outer `Vec`'s own allocation).
|
|
||||||
- New `sparse_findere_for_genome(hits, z, presence, threshold) -> (Vec<ConfirmedHit>, n_runs, total_run_len)` (`query.rs`): sorts one genome's raw hits by `(seq_idx, pos)`, detects maximal runs of consecutive `pos` within one sequence, runs the same monotone-deque window-minimum as before but scoped to each run (run-relative indices for eviction, absolute `pos` for computing `pos_out`). Presence/count adjustment (`u32::from(win_min >= threshold)` vs. raw) is applied inside this function, once per confirmed hit, rather than later during accumulation.
|
|
||||||
- `process_chunk` restructured into three passes after the partition loop: (1) run `sparse_findere_for_genome` per genome, collecting `confirmed_by_genome` and run-detection stats; (2) accumulate `genome_totals` directly from `confirmed_by_genome` and mark a `confirmed_any: Vec<bool>` (sized `total_kmers_out`, not `× n_genomes`); (3) a position-only pass (`O(total_kmers_out)`, no genome factor) computing `kmer_count`/`kmer_missing` from `confirmed_any` + `SmerIndex`. `cov` (`--detail`) is populated by re-scanning `confirmed_by_genome` — only when `--detail` is actually set, otherwise skipped entirely.
|
|
||||||
- Debug log added (`"sparse Findere"`): `n_dense_would_be` (`n_occurrences × n_genomes` — what the deleted dense path would have allocated), `n_sparse_entries` (what's actually retained), `n_runs`/`avg_run_len` (per the plan's ask, to see whether hits mostly fail to form complete windows).
|
|
||||||
- **Unit tests**: `sparse_findere_matches_dense_reference_on_random_inputs` (`obikmer/src/cmd/tests/query.rs`) — 200 randomized cases (sequence count/length, `z`, presence/count mode, threshold, hit density from sparse to fully-dense) comparing `sparse_findere_for_genome` against `dense_reference_findere`, a faithful reimplementation of the deleted dense algorithm kept only as a test-local correctness oracle (no property-testing crate added — checked first, none was a workspace dependency; a small `std`-only xorshift64 PRNG stands in for one, deterministic and dependency-free). All 200 cases pass.
|
|
||||||
- **Functional validation performed**: full workspace build + `cargo test --workspace`, zero failures. End-to-end against real indexes: baseline output (no flags) unchanged from pre-phase-5 recorded values on the same fixtures; `--count-missing` correct (`kmer_missing: 0` on a self-match); `--detail` correct — coverage array length matches `kmer_count`, and critically, re-ran the two-genome cross-contamination check from phase 4 with `--detail --count-missing`: `genomeA` reads show coverage sum `106` for `genomeA` and `0` for `genomeB` (and vice versa) — confirms the sparse-to-dense `cov` reconstruction doesn't leak across genomes either, not just the scalar `kmer_strict_matches` path.
|
|
||||||
- This phase's roadmap item ("update the Findere z-window filter section") — done, see above; the "Algorithm" section's pseudocode was also updated, since it still named `KmerResults`/`SKDesc` from before phases 3–4.
|
|
||||||
|
|
||||||
### Phase 6 — Parallel gzip decompression (independent, optional)
|
|
||||||
|
|
||||||
Tracked separately in [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen); parked pending validation of `rapidgzip-rs` on real data. Not a dependency of, or a dependency for, phases 0–5 — `xopen` is shared infrastructure (`obiread`), phase 1 benefits from it but doesn't require it (phase 1 parallelises *across* files; this phase would additionally parallelise *within* one large file).
|
|
||||||
|
|
||||||
### Cross-cutting risks
|
|
||||||
|
|
||||||
- **Thread-budget oversubscription** (phase 4): the single biggest unresolved design question in this whole plan — see phase 4's composition note. Should be settled with real measurements early in phase 4, not assumed from the design alone.
|
|
||||||
- **`obicompactvec` API surface growth** (phase 4): new public per-column accessors are additive (existing `fill_row`/`row` stay for other callers — `dump`, `select`, distance computations) — no breaking change expected, but worth checking `obicompactvec`'s other callers aren't already relying on `fill_row` being the only/cheapest access path in a way that would make maintaining two access patterns (row-major and column-major) a real maintenance cost rather than a one-off addition.
|
|
||||||
- **`PersistentBitMatrix::Implicit`'s hardcoded `n_cols: 1` — resolved, not a bug.** `LayerMeta`'s own doc comment (`obicompactvec/src/layer_meta.rs:1-9`) states it is written "alongside `mphf.bin`" and read by `PersistentBitMatrix::open` "to determine `n_rows` for **the implicit (mono-genome presence/absence) case**" — i.e. `Implicit` is a documented single-genome fast path (no presence matrix needed when there is trivially one genome), not a generic "no matrix built yet" fallback. `n_cols: 1` is correct by design for the case it's meant to handle. Phase 4's column loop is safe as planned — this was worth checking once, doesn't need further action.
|
|
||||||
|
|||||||
@@ -107,19 +107,3 @@ stateDiagram-v2
|
|||||||
`restart` is updated each time a `+` is found. When any state fails its expected input, the scan jumps back to `restart` and continues from there — guaranteeing that a `@` in a quality line cannot be accepted as a record start, because the `\n+\n` structure immediately following it (going backward) will not be found.
|
`restart` is updated each time a `+` is found. When any state fails its expected input, the scan jumps back to `restart` and continues from there — guaranteeing that a `@` in a quality line cannot be accepted as a record start, because the `\n+\n` structure immediately following it (going backward) will not be found.
|
||||||
|
|
||||||
Returns the byte offset of the `@` that starts the last complete record.
|
Returns the byte offset of the `@` that starts the last complete record.
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Future work — parallel gzip decompression in `xopen`
|
|
||||||
|
|
||||||
`obiread::xopen` (`xopen.rs`) decompresses gzip via `niffler` → `flate2`, which is single-threaded (standard DEFLATE has no parallel-decodable structure). For large local gzip inputs this single-threaded decompression can become the throughput bottleneck feeding the `query`/`index`/`superkmer` pipelines, since chunk/page production for a given file is serialized ahead of the worker pool.
|
|
||||||
|
|
||||||
Candidate: special-case local, on-disk, gzip-magic-detected paths in `open_raw`/`xopen` to use [`rapidgzip-rs`](https://github.com/alekseizarubin/rapidgzip-rs) (`ReaderBuilder::new().parallelism(n).open(path)`, implements `Read + Seek`) instead of `niffler`, keeping `niffler` for every other case: `stdin` (`-`), HTTP(S) sources, and all non-gzip formats (bzip2, xz, zstd — less used in practice here).
|
|
||||||
|
|
||||||
Constraints identified so far (not yet validated against real data):
|
|
||||||
- Branch point must move earlier than the current `decompress()` call in `open_raw` — rapidgzip's fast path needs the file **path**, not an already-opened generic `Read`, so the gzip/local-file detection has to happen before the generic `File::open` + `niffler::send::get_reader` path is taken.
|
|
||||||
- `stdin` and HTTP sources are not seekable — they stay on `niffler` regardless; the gain only applies to local on-disk `.gz` files.
|
|
||||||
- `rapidgzip-sys` vendors a native C++ engine: requires CMake ≥ 3.17, a C++17 compiler, and `nasm` on x86 targets — a real build-toolchain addition, not just a pure-Rust crate.
|
|
||||||
- Low maturity of the Rust binding at review time (2 GitHub stars, ~15 commits, April 2026 latest release) — the underlying C++ engine is validated (HPDC 2023 paper), but the binding itself has limited production track record.
|
|
||||||
|
|
||||||
Decision: parked for now. Before adopting, validate on real data: throughput vs. `niffler` on representative large `.gz` inputs, and byte-for-byte correctness of decompressed output.
|
|
||||||
|
|||||||
@@ -92,48 +92,18 @@ For each genome:
|
|||||||
|
|
||||||
| Flag | Applies to | Meaning |
|
| Flag | Applies to | Meaning |
|
||||||
|------|-----------|---------|
|
|------|-----------|---------|
|
||||||
| `--min-count N` | ingroup | k-mer present in at least N ingroup genomes (N may be negative, see below) |
|
| `--min-count N` | ingroup | k-mer present in at least N ingroup genomes |
|
||||||
| `--max-count N` | ingroup | k-mer present in at most N ingroup genomes (N may be negative, see below) |
|
| `--max-count N` | ingroup | k-mer present in at most N ingroup genomes |
|
||||||
| `--min-frac F` | ingroup | k-mer present in at least fraction F of ingroup genomes |
|
| `--min-frac F` | ingroup | k-mer present in at least fraction F of ingroup genomes |
|
||||||
| `--max-frac F` | ingroup | k-mer present in at most fraction F of ingroup genomes |
|
| `--max-frac F` | ingroup | k-mer present in at most fraction F of ingroup genomes |
|
||||||
| `--min-outgroup-count N` | outgroup | k-mer present in at least N outgroup genomes (N may be negative, see below) |
|
| `--min-outgroup-count N` | outgroup | k-mer present in at least N outgroup genomes |
|
||||||
| `--max-outgroup-count N` | outgroup | k-mer present in at most N outgroup genomes (N may be negative, see below) |
|
| `--max-outgroup-count N` | outgroup | k-mer present in at most N outgroup genomes |
|
||||||
| `--min-outgroup-frac F` | outgroup | k-mer present in at least fraction F of outgroup genomes |
|
| `--min-outgroup-frac F` | outgroup | k-mer present in at least fraction F of outgroup genomes |
|
||||||
| `--max-outgroup-frac F` | outgroup | k-mer present in at most fraction F of outgroup genomes |
|
| `--max-outgroup-frac F` | outgroup | k-mer present in at most fraction F of outgroup genomes |
|
||||||
| `--min-total-count N` | all genomes | sum of per-genome counts ≥ N (`filter` only) |
|
| `--min-total-count N` | all genomes | sum of per-genome counts ≥ N (`filter` only) |
|
||||||
| `--max-total-count N` | all genomes | sum of per-genome counts ≤ N (`filter` only) |
|
| `--max-total-count N` | all genomes | sum of per-genome counts ≤ N (`filter` only) |
|
||||||
| `--presence-threshold N` | all | per-genome count > N to be considered "present" (default 0) |
|
| `--presence-threshold N` | all | per-genome count > N to be considered "present" (default 0) |
|
||||||
|
|
||||||
### Negative counts — offset from group size
|
|
||||||
|
|
||||||
The four integer count flags (`--min-count`, `--max-count`, `--min-outgroup-count`,
|
|
||||||
`--max-outgroup-count`) accept **negative** values, interpreted as an offset counted
|
|
||||||
down from the group size `n`, resolved at run time once `n` is known:
|
|
||||||
|
|
||||||
| Value | Effective threshold |
|
|
||||||
|-------|---------------------|
|
|
||||||
| `N ≥ 0` | literal absolute count `N` |
|
|
||||||
| `-x` (x > 0) | `max(1, n − x)` — "all but x" |
|
|
||||||
|
|
||||||
`-1` literally means *all but one*, `-2` *all but two*, and so on. This expresses
|
|
||||||
a quorum relative to the group size that a plain fraction cannot state exactly
|
|
||||||
(e.g. "present in every genome except at most one" is `n−1`, which is `0.9` for
|
|
||||||
`n = 10` but `0.857…` for `n = 7`).
|
|
||||||
|
|
||||||
The threshold is **floored at 1**, never 0: the negative form always keeps
|
|
||||||
constraining the group. Without the floor, `--min-count -1` on a singleton
|
|
||||||
ingroup (`n = 1`) would resolve to `0` ("at least 0") and silently drop the
|
|
||||||
constraint; the floor makes it `1` ("present in that one genome") instead.
|
|
||||||
|
|
||||||
To express a count of `0` (e.g. "absent from the ingroup"), use the literal `0`,
|
|
||||||
not a negative — `0` and `-0` are indistinguishable, so the offset form starts at
|
|
||||||
`-1`.
|
|
||||||
|
|
||||||
> **Edge case** — on an *empty* group (`n = 0`, e.g. a predicate matching no
|
|
||||||
> genome), a negative count still resolves to `1`, an impossible constraint that
|
|
||||||
> rejects every k-mer. This is consistent with an empty group letting nothing
|
|
||||||
> through, but differs from the "no constraint" behaviour of the fraction flags.
|
|
||||||
|
|
||||||
**Conditional defaults** — the defaults for `--min-frac` and `--max-outgroup-count` depend on two conditions:
|
**Conditional defaults** — the defaults for `--min-frac` and `--max-outgroup-count` depend on two conditions:
|
||||||
whether the corresponding group was declared, **and** whether any quorum flag for that group was explicitly set.
|
whether the corresponding group was declared, **and** whether any quorum flag for that group was explicitly set.
|
||||||
|
|
||||||
@@ -245,17 +215,6 @@ obikmer filter src --output dst \
|
|||||||
--max-outgroup-count 0
|
--max-outgroup-count 0
|
||||||
```
|
```
|
||||||
|
|
||||||
Noise-tolerant core — keep k-mers present in *all but one* ingroup genome
|
|
||||||
(`-1` = `n−1`) and absent from *all but one* of the outgroup:
|
|
||||||
|
|
||||||
```sh
|
|
||||||
obikmer filter src --output dst \
|
|
||||||
--ingroup "genus=Betula" \
|
|
||||||
--outgroup "*" \
|
|
||||||
--min-count -1 \
|
|
||||||
--max-outgroup-count -1
|
|
||||||
```
|
|
||||||
|
|
||||||
To dump only k-mers specific to *Betula nana*:
|
To dump only k-mers specific to *Betula nana*:
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
|
|||||||
@@ -347,24 +347,11 @@ Provided finalisations:
|
|||||||
| `relfreq_euclidean_dist_matrix()` | `√partial_relfreq_euclidean[i,j]` |
|
| `relfreq_euclidean_dist_matrix()` | `√partial_relfreq_euclidean[i,j]` |
|
||||||
| `hellinger_dist_matrix()` | `√partial_hellinger[i,j] / √2` |
|
| `hellinger_dist_matrix()` | `√partial_hellinger[i,j] / √2` |
|
||||||
| `hellinger_euclidean_dist_matrix()` | `√partial_hellinger[i,j]` |
|
| `hellinger_euclidean_dist_matrix()` | `√partial_hellinger[i,j]` |
|
||||||
| `threshold_mash_dist_matrix(k, t)` | Mash distance, derived from `threshold_jaccard_dist_matrix(t)` — no separate partial |
|
|
||||||
|
|
||||||
### BitPartials
|
### BitPartials
|
||||||
|
|
||||||
Required: `partial_jaccard() -> (Array2<u64>, Array2<u64>)`, `partial_hamming() -> Array2<u64>`. Both additive across layers and partitions.
|
Required: `partial_jaccard() -> (Array2<u64>, Array2<u64>)`, `partial_hamming() -> Array2<u64>`. Both additive across layers and partitions.
|
||||||
|
|
||||||
Provided finalisations also include `jaccard_dist_matrix()`, `hamming_dist_matrix()`, and `mash_dist_matrix(k)`.
|
|
||||||
|
|
||||||
### Mash distance
|
|
||||||
|
|
||||||
`mash_dist_matrix`/`threshold_mash_dist_matrix` add no new additive primitive: both are a pointwise transform of the existing Jaccard distance matrix, per the Mash mutation-rate estimator [@Mash-distances-doc; @Fan2015-mash-formula]:
|
|
||||||
|
|
||||||
```
|
|
||||||
D = -1/k · ln(2J / (1+J)), J = 1 - d_jaccard
|
|
||||||
```
|
|
||||||
|
|
||||||
`J ≤ 0` (i.e. `d_jaccard ≥ 1`, no shared k-mers) maps to `D = 1` (maximal distance) rather than the `ln` singularity at `J = 0`.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Temp-file-backed types
|
## Temp-file-backed types
|
||||||
|
|||||||
+1
-1
@@ -13,7 +13,7 @@
|
|||||||
| `query` | Query an index with sequences and annotate matches |
|
| `query` | Query an index with sequences and annotate matches |
|
||||||
| `dump` | Dump all indexed k-mers as CSV (kmer + per-genome counts or presence); supports the shared [kmer filtering](implementation/filtering.md) system; `--head N` limits output to the first N k-mers |
|
| `dump` | Dump all indexed k-mers as CSV (kmer + per-genome counts or presence); supports the shared [kmer filtering](implementation/filtering.md) system; `--head N` limits output to the first N k-mers |
|
||||||
| `annotate` | Add or update genome metadata from a CSV file; or dump metadata as CSV |
|
| `annotate` | Add or update genome metadata from a CSV file; or dump metadata as CSV |
|
||||||
| `distance` | Compute pairwise distance matrix between genomes (`--metric jaccard\|mash\|hamming\|bray-curtis\|relfreq-bray-curtis\|euclidean\|relfreq-euclidean\|hellinger\|hellinger-euclidean`); optionally build NJ/UPGMA trees; `--presence-threshold N` sets the minimum count to consider a k-mer present when computing Jaccard/Mash on count indexes (default 1) |
|
| `distance` | Compute pairwise distance matrix between genomes; optionally build NJ/UPGMA trees; `--presence-threshold N` sets the minimum count to consider a k-mer present when computing Jaccard on count indexes (default 1) |
|
||||||
| `unitig` | Build a global de Bruijn graph across all partitions and enumerate its unitigs as FASTA; supports the shared [kmer filtering](implementation/filtering.md) system |
|
| `unitig` | Build a global de Bruijn graph across all partitions and enumerate its unitigs as FASTA; supports the shared [kmer filtering](implementation/filtering.md) system |
|
||||||
| `select` | Project and/or aggregate genome columns into a new or in-place index; the column-axis counterpart of `filter` (see [select](implementation/select.md)) |
|
| `select` | Project and/or aggregate genome columns into a new or in-place index; the column-axis counterpart of `filter` (see [select](implementation/select.md)) |
|
||||||
| `estimate` | Estimate approximate-index parameters (z, evidence bits, FP rates) before indexing |
|
| `estimate` | Estimate approximate-index parameters (z, evidence bits, FP rates) before indexing |
|
||||||
|
|||||||
@@ -241,21 +241,3 @@
|
|||||||
volume = 33,
|
volume = 33,
|
||||||
year = 2017,
|
year = 2017,
|
||||||
bdsk-url-1 = {http://dx.doi.org/10.1093/bioinformatics/btw832}}
|
bdsk-url-1 = {http://dx.doi.org/10.1093/bioinformatics/btw832}}
|
||||||
|
|
||||||
@misc{Mash-distances-doc,
|
|
||||||
author = {{Marbl Lab}},
|
|
||||||
howpublished = {Mash documentation},
|
|
||||||
title = {Mash Distance},
|
|
||||||
url = {https://mash.readthedocs.io/en/latest/distances.html},
|
|
||||||
urldate = {2026-07-09},
|
|
||||||
year = 2026}
|
|
||||||
|
|
||||||
@article{Fan2015-mash-formula,
|
|
||||||
author = {Fan, Huan and Ives, Anthony R and Surget-Groba, Yann and Cannon, Charles H},
|
|
||||||
doi = {10.1186/s12864-015-1647-5},
|
|
||||||
journal = {BMC Genomics},
|
|
||||||
number = 1,
|
|
||||||
title = {An assembly and alignment-free method of phylogeny reconstruction from next-generation sequencing data},
|
|
||||||
url = {https://doi.org/10.1186/s12864-015-1647-5},
|
|
||||||
volume = 16,
|
|
||||||
year = 2015}
|
|
||||||
|
|||||||
+16
-32
@@ -1,6 +1,6 @@
|
|||||||
# Kmer entropy filter
|
# Kmer entropy filter
|
||||||
|
|
||||||
Low-complexity kmers (polyA, polyT, tandem repeats) are detected and excluded during phase 1. The filter computes a **normalized Shannon entropy** over sub-words of multiple sizes, corrected for one source of bias: the small number of observations within a single kmer relative to the number of possible sub-words.
|
Low-complexity kmers (polyA, polyT, tandem repeats) are detected and excluded during phase 1. The filter computes a **normalized Shannon entropy** over sub-words of multiple sizes, corrected for two sources of bias: the small number of observations within a single kmer, and the unequal sizes of circular equivalence classes.
|
||||||
|
|
||||||
## Sub-word frequencies
|
## Sub-word frequencies
|
||||||
|
|
||||||
@@ -8,15 +8,17 @@ For a kmer of length k and a sub-word size ws (1 ≤ ws ≤ ws_max, typically ws
|
|||||||
|
|
||||||
$$w_i = \text{kmer}[i \mathinner{..} i+ws-1], \quad i = 0, \ldots, n_{\text{words}}-1$$
|
$$w_i = \text{kmer}[i \mathinner{..} i+ws-1], \quad i = 0, \ldots, n_{\text{words}}-1$$
|
||||||
|
|
||||||
Each sub-word is tallied under its own raw 2-bit-packed value — **no canonicalization**. Let $f_j$ be the count of raw word $j$ among the $n_{\text{words}}$ sub-words ($\sum_j f_j = n_{\text{words}}$), over the $4^{ws}$ possible raw words.
|
Each sub-word is mapped to its **circular canonical form**: the lexicographic minimum among all cyclic rotations of the word **and all cyclic rotations of its reverse complement**. This extended equivalence relation ensures that entropy(K) = entropy(revcomp(K)) — the filter is strand-symmetric. Let $s_j$ be the size of equivalence class $j$ (number of distinct raw words mapping to canonical form $j$), and $f_j$ the count of canonical form $j$ among the $n_{\text{words}}$ sub-words ($\sum_j f_j = n_{\text{words}}$).
|
||||||
|
|
||||||
An earlier version of this filter first folded each sub-word into a circular+reverse-complement equivalence class, then "unfolded" the observed class frequency back onto its members to correct for unequal class sizes. That machinery bought nothing it was claimed for — see *Why no equivalence classes* below — while measurably weakening detection of the very sequences the filter exists to catch, so it was removed.
|
|
||||||
|
|
||||||
## Corrected Shannon entropy
|
## Corrected Shannon entropy
|
||||||
|
|
||||||
$$H_{\text{corr}} = \log(n_{\text{words}}) - \frac{1}{n_{\text{words}}} \sum_j f_j \log f_j$$
|
The circular equivalence classes have unequal sizes: under a uniform distribution over all $4^{ws}$ raw words, class $j$ is visited with probability $s_j / 4^{ws}$, not $1/n_a$. Computing entropy directly over canonical classes therefore underestimates the entropy of a random sequence.
|
||||||
|
|
||||||
This is a plain Shannon entropy over the observed raw-word frequencies.
|
The correction "unfolds" each canonical class back to its member raw words, redistributing each observation of class $j$ equally among its $s_j$ members:
|
||||||
|
|
||||||
|
$$H_{\text{corr}} = \log(n_{\text{words}}) - \frac{1}{n_{\text{words}}} \sum_j f_j \log f_j + \frac{1}{n_{\text{words}}} \sum_j f_j \log s_j$$
|
||||||
|
|
||||||
|
The last term is the correction for unequal class sizes. For a uniformly random sequence ($f_j \approx n_{\text{words}} \cdot s_j / 4^{ws}$), this gives $H_{\text{corr}} \approx \log(4^{ws}) = 2 \cdot ws \cdot \log 2$, the maximum entropy over raw words.
|
||||||
|
|
||||||
## Maximum entropy correction for small samples
|
## Maximum entropy correction for small samples
|
||||||
|
|
||||||
@@ -40,45 +42,27 @@ $$\text{entropy}(kmer) = \min_{ws=1}^{ws_{\max}} \hat{H}(ws)$$
|
|||||||
|
|
||||||
A value near 0 indicates low complexity (e.g. AAAA…); near 1 indicates high complexity. A kmer is rejected if $\text{entropy}(kmer) < \theta$, where $\theta$ is a collection parameter (default 0.7). The minimum across word sizes ensures that any scale of repetition is detected independently: polyA is caught at ws=1, dinucleotide repeats at ws=2, etc.
|
A value near 0 indicates low complexity (e.g. AAAA…); near 1 indicates high complexity. A kmer is rejected if $\text{entropy}(kmer) < \theta$, where $\theta$ is a collection parameter (default 0.7). The minimum across word sizes ensures that any scale of repetition is detected independently: polyA is caught at ws=1, dinucleotide repeats at ws=2, etc.
|
||||||
|
|
||||||
## Why no equivalence classes
|
|
||||||
|
|
||||||
A prior design folded each sub-word into the canonical form of its circular-rotation + reverse-complement equivalence class before tallying, on the reasoning that (a) it guarantees $\text{entropy}(K) = \text{entropy}(\text{revcomp}(K))$, and (b) collapsing phase-shifted repeats (e.g. `ATG` ≡ `TGA` ≡ `GAT`) into one class better reflects that they are "the same" low-complexity pattern.
|
|
||||||
|
|
||||||
Both properties already hold for the raw, unfolded entropy above, without any class machinery:
|
|
||||||
|
|
||||||
- **Reverse complement**: for any K of length n, window $j$ of $\text{revcomp}(K)$ equals $\text{revcomp}$ of window $(n{-}ws{-}j)$ of K. This is a bijection between the window sets under which each window maps to its own revcomp — and revcomp is itself a bijection (involution) on the space of raw ws-mers. So the multiset of raw-word frequencies for $\text{revcomp}(K)$ is exactly a relabeling of the multiset for K, and Shannon entropy — a function of the frequency multiset alone — is exactly invariant. No folding required, for any K.
|
|
||||||
- **Tandem repeats**: a period-p repeat sampled by a stride-1 sliding window naturally cycles through its own rotations as raw tokens (e.g. `ATGATGATG…` yields the raw words `ATG`, `TGA`, `GAT` in rotation as the window slides). The low diversity this represents (few distinct raw words out of $4^{ws}$ possible) is already visible in the raw frequency distribution — no folding needed to detect it.
|
|
||||||
|
|
||||||
What the fold-then-unfold step actually did was credit each observed class with the frequency of equivalence-class members that were **never observed on the read strand**, inflating $H_{\text{corr}}$ for genuine repeats. Worked example: k=31, ws=3, kmer = `ATG` repeated ($n_{\text{words}}=29$, all 29 windows fall into one class of size 6 under the old scheme — 3 rotations × forward/revcomp):
|
|
||||||
|
|
||||||
| | $H_{\text{corr}}$ | normalized |
|
|
||||||
|---|---|---|
|
|
||||||
| old (folded, class size 6) | $\log 6 \approx 1.79$ | $\approx 0.53$ |
|
|
||||||
| current (raw, unfolded) | $\log 3 \approx 1.10$ | $\approx 0.33$ |
|
|
||||||
|
|
||||||
The gap is not a rounding artifact: per sub-word order, the folded score for this same repeat swings from 0.53 (ws=3, aligned with the period) up to **1.03** (ws=5, misaligned with the period) — i.e. a period-3 repeat could score *above* the theoretical maximum for a random sequence, depending on which ws happens to divide the repeat's period. The raw formula stays flat at ≈0.33–0.40 across ws=2..6 regardless of alignment, which is the robustness the "minimum across ws" design was meant to provide in the first place.
|
|
||||||
|
|
||||||
## Interpretation as an effective number of classes
|
## Interpretation as an effective number of classes
|
||||||
|
|
||||||
$H_{\text{corr}}$ is a standard Shannon entropy over raw words, so the classical perplexity interpretation holds directly: $N_{\text{eff}} = e^{H_{\text{corr}}}$ is the number of equiprobable raw words that would yield the same entropy.
|
$H_{\text{corr}}$ is a standard Shannon entropy over raw words (after unfolding the equivalence classes), so the classical perplexity interpretation holds directly: $N_{\text{eff}} = e^{H_{\text{corr}}}$ is the number of equiprobable classes that would yield the same entropy.
|
||||||
|
|
||||||
For the normalised score $\hat{H}$, dividing by $H_{\max}$ changes the logarithm base:
|
For the normalised score $\hat{H}$, dividing by $H_{\text{max}}$ changes the logarithm base:
|
||||||
|
|
||||||
$$\hat{H} = \frac{\log N_{\text{eff}}}{\log N_{\max}} = \log_{N_{\max}} N_{\text{eff}} \quad \Longleftrightarrow \quad N_{\text{eff}} = N_{\max}^{\,\hat{H}}$$
|
$$\hat{H} = \frac{\log N_{\text{eff}}}{\log N_{\text{max}}} = \log_{N_{\text{max}}} N_{\text{eff}} \quad \Longleftrightarrow \quad N_{\text{eff}} = N_{\text{max}}^{\,\hat{H}}$$
|
||||||
|
|
||||||
The property is preserved: $\hat{H}$ is the logarithm (in base $N_{\max}$) of the effective number of equi-represented raw words.
|
The property is preserved: $\hat{H}$ is the logarithm (in base $N_{\text{max}}$) of the effective number of equi-represented classes.
|
||||||
|
|
||||||
In the large-sample limit ($n_{\text{words}} \gg 4^{ws}$), $N_{\max} \approx 4^{ws}$, giving:
|
In the large-sample limit ($n_{\text{words}} \gg 4^{ws}$), $N_{\text{max}} \approx 4^{ws}$, giving:
|
||||||
|
|
||||||
$$N_{\text{eff}} \approx 4^{ws \cdot \hat{H}}$$
|
$$N_{\text{eff}} \approx 4^{ws \cdot \hat{H}}$$
|
||||||
|
|
||||||
This has a clean interpretation: $ws \cdot \hat{H}$ is the **effective word length** (in bases) of a perfectly uniform distribution that would produce the same entropy. At $\hat{H} = 1$ the full space of $4^{ws}$ words is used; at $\hat{H} = 0.5$ with ws=2, only $4^1 = 4$ effective words out of 16 are occupied.
|
This has a clean interpretation: $ws \cdot \hat{H}$ is the **effective word length** (in bases) of a perfectly uniform distribution that would produce the same entropy. At $\hat{H} = 1$ the full space of $4^{ws}$ words is used; at $\hat{H} = 0.5$ with ws=2, only $4^1 = 4$ effective classes out of 16 are occupied.
|
||||||
|
|
||||||
In our actual regime, $n_{\text{words}}$ is small and $4^{ws}$ can exceed $n_{\text{words}}$, so $H_{\max} < \log(4^{ws})$ due to the small-sample correction. The exact effective count is $N_{\max}^{\hat{H}}$, not $4^{ws \cdot \hat{H}}$.
|
In our actual regime, $n_{\text{words}}$ is small and $4^{ws}$ can exceed $n_{\text{words}}$, so $H_{\text{max}} < \log(4^{ws})$ due to the small-sample correction. The exact effective count is $N_{\text{max}}^{\hat{H}}$, not $4^{ws \cdot \hat{H}}$.
|
||||||
|
|
||||||
## Properties
|
## Properties
|
||||||
|
|
||||||
The entropy score is a function of the kmer sequence alone — it does not depend on the surrounding context or on the position within any genome. Two consequences:
|
The entropy score is a function of the kmer sequence alone — it does not depend on the surrounding context or on the position within any genome. Two consequences:
|
||||||
|
|
||||||
- **Orientation invariance**: $\text{entropy}(K) = \text{entropy}(\text{revcomp}(K))$ — see *Why no equivalence classes* above for why this holds without any explicit strand-folding step.
|
- **Orientation invariance**: $\text{entropy}(K) = \text{entropy}(\text{revcomp}(K))$, guaranteed by the strand-symmetric canonical form.
|
||||||
- **Context independence**: the same kmer is always rejected or always kept, regardless of which genome it occurs in, where in that genome it appears, or which strand is considered. The filter defines a fixed partition of the kmer space into low-complexity and valid kmers.
|
- **Context independence**: the same kmer is always rejected or always kept, regardless of which genome it occurs in, where in that genome it appears, or which strand is considered. The filter defines a fixed partition of the kmer space into low-complexity and valid kmers.
|
||||||
|
|||||||
@@ -3,14 +3,10 @@
|
|||||||
|
|
||||||
## Code couvert
|
## Code couvert
|
||||||
|
|
||||||
- `obikentropy/src/table.rs`, `obikentropy/src/tracker.rs` — formule d'entropie et tables de correction petits effectifs
|
- `obiskbuilder/src/entropy_table.rs` — filtre Shannon sur les kmers à basse complexité
|
||||||
- `obikentropy/src/kmer_entropy.rs` — entropie d'un kmer isolé (`KmerEntropy`)
|
- `obiskbuilder/src/lib.rs` — application du filtre lors du scatter (phase 1)
|
||||||
- `obiskbuilder/src/rolling_stat.rs` — composition de `obikentropy::EntropyTracker` dans le suivi streaming (sélection de minimiseur + entropie)
|
|
||||||
- `obiskbuilder/src/iter.rs`, `obiskbuilder/src/stream_iter.rs` — application du filtre lors du scatter (phase 1)
|
|
||||||
|
|
||||||
## Notes
|
## Notes
|
||||||
|
|
||||||
Le repli en classes d'équivalence circulaires + brin inverse (décrit dans une version antérieure de ce document) a été supprimé : voir la section « Why no equivalence classes » de `entropy.md` pour la justification théorique et numérique.
|
Document théorique stable. Vérifier que les paramètres `theta` et `level_max` dans le CLI
|
||||||
|
|
||||||
Vérifier que les paramètres `theta` et `level_max` dans le CLI
|
|
||||||
(`obikmer/src/cli.rs` → `CommonArgs`) correspondent bien à ce qui est décrit.
|
(`obikmer/src/cli.rs` → `CommonArgs`) correspondent bien à ce qui est décrit.
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -36,7 +36,6 @@ nav:
|
|||||||
- Entropy filter: theory/entropy.md
|
- Entropy filter: theory/entropy.md
|
||||||
- Minimizer selection: theory/minimizer.md
|
- Minimizer selection: theory/minimizer.md
|
||||||
- Partitioning architecture: theory/indexing.md
|
- Partitioning architecture: theory/indexing.md
|
||||||
- Central-position SNP distance (discussion): theory/evolutionary_distances.md
|
|
||||||
- Implementation:
|
- Implementation:
|
||||||
- SuperKmer: implementation/superkmer.md
|
- SuperKmer: implementation/superkmer.md
|
||||||
- Kmer: implementation/kmer.md
|
- Kmer: implementation/kmer.md
|
||||||
|
|||||||
Generated
+1
-15
@@ -1682,13 +1682,6 @@ dependencies = [
|
|||||||
"xxhash-rust",
|
"xxhash-rust",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "obikentropy"
|
|
||||||
version = "0.1.0"
|
|
||||||
dependencies = [
|
|
||||||
"obikseq",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "obikindex"
|
name = "obikindex"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
@@ -1701,21 +1694,17 @@ dependencies = [
|
|||||||
"obikpartitionner",
|
"obikpartitionner",
|
||||||
"obikseq",
|
"obikseq",
|
||||||
"obilayeredmap",
|
"obilayeredmap",
|
||||||
"obipipeline",
|
|
||||||
"obiread",
|
|
||||||
"obiskbuilder",
|
|
||||||
"obiskio",
|
"obiskio",
|
||||||
"obisys",
|
"obisys",
|
||||||
"rayon",
|
"rayon",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"tempfile",
|
|
||||||
"tracing",
|
"tracing",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "obikmer"
|
name = "obikmer"
|
||||||
version = "1.1.43"
|
version = "1.1.17"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"clap",
|
"clap",
|
||||||
"csv",
|
"csv",
|
||||||
@@ -1753,7 +1742,6 @@ dependencies = [
|
|||||||
"niffler 3.0.0",
|
"niffler 3.0.0",
|
||||||
"obicompactvec",
|
"obicompactvec",
|
||||||
"obidebruinj",
|
"obidebruinj",
|
||||||
"obikentropy",
|
|
||||||
"obikrope",
|
"obikrope",
|
||||||
"obikseq",
|
"obikseq",
|
||||||
"obilayeredmap",
|
"obilayeredmap",
|
||||||
@@ -1836,9 +1824,7 @@ dependencies = [
|
|||||||
name = "obiskbuilder"
|
name = "obiskbuilder"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"criterion2",
|
|
||||||
"lazy_static",
|
"lazy_static",
|
||||||
"obikentropy",
|
|
||||||
"obikrope",
|
"obikrope",
|
||||||
"obikseq",
|
"obikseq",
|
||||||
"obiread",
|
"obiread",
|
||||||
|
|||||||
+1
-1
@@ -1,5 +1,5 @@
|
|||||||
[workspace]
|
[workspace]
|
||||||
resolver = "3"
|
resolver = "3"
|
||||||
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy"]
|
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy"]
|
||||||
[profile.release]
|
[profile.release]
|
||||||
debug = 1
|
debug = 1
|
||||||
|
|||||||
Binary file not shown.
@@ -1,5 +1,5 @@
|
|||||||
use std::fs::{self, File};
|
use std::fs::{self, File};
|
||||||
use std::io::{self, BufWriter, Read as _, Write as _};
|
use std::io::{self, BufWriter, Write as _};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use memmap2::Mmap;
|
use memmap2::Mmap;
|
||||||
@@ -171,43 +171,19 @@ impl PackedBitMatrix {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reads just the `n_cols` field from an existing packed matrix's header,
|
|
||||||
/// without mapping the file. Used by `pack_bit_matrix` to tell a genuinely
|
|
||||||
/// complete pack from a stale one that predates a later column-widening.
|
|
||||||
fn packed_bit_matrix_n_cols(path: &Path) -> io::Result<usize> {
|
|
||||||
let mut f = File::open(path)?;
|
|
||||||
let mut header = [0u8; PBMX_HEADER];
|
|
||||||
f.read_exact(&mut header)?;
|
|
||||||
Ok(u64::from_le_bytes(header[16..24].try_into().unwrap()) as usize)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Build `presence/matrix.pbmx` from existing `col_*.pbiv` files.
|
/// Build `presence/matrix.pbmx` from existing `col_*.pbiv` files.
|
||||||
pub fn pack_bit_matrix(dir: &Path) -> io::Result<()> {
|
pub fn pack_bit_matrix(dir: &Path) -> io::Result<()> {
|
||||||
let packed_path = dir.join("matrix.pbmx");
|
let packed_path = dir.join("matrix.pbmx");
|
||||||
|
if packed_path.exists() {
|
||||||
let meta = match MatrixMeta::load(dir) {
|
// Matrix complete; remove any leftover column files from a killed cleanup.
|
||||||
Ok(meta) => meta,
|
if let Ok(meta) = MatrixMeta::load(dir) {
|
||||||
Err(e) => {
|
|
||||||
// No columnar data pending: either this layer was already
|
|
||||||
// packed and cleaned up (matrix.pbmx complete, nothing left to
|
|
||||||
// do), or genuinely nothing was ever written here.
|
|
||||||
return if packed_path.exists() { Ok(()) } else { Err(e) };
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// A `matrix.pbmx` can already exist here even though columnar data is
|
|
||||||
// still pending — e.g. copied verbatim from a merge's base source
|
|
||||||
// before this layer was widened with more genome columns (see
|
|
||||||
// `obikpartitionner::merge_partition`). Only skip (re-)packing if the
|
|
||||||
// existing file already reflects the current column count; otherwise
|
|
||||||
// the columnar files are newer and must be (re-)packed, overwriting the
|
|
||||||
// stale one — never silently discarded as "leftover cleanup".
|
|
||||||
if packed_bit_matrix_n_cols(&packed_path).ok() == Some(meta.n_cols) {
|
|
||||||
for c in 0..meta.n_cols { let _ = fs::remove_file(col_path(dir, c)); }
|
for c in 0..meta.n_cols { let _ = fs::remove_file(col_path(dir, c)); }
|
||||||
let _ = fs::remove_file(dir.join("meta.json"));
|
let _ = fs::remove_file(dir.join("meta.json"));
|
||||||
|
}
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let meta = MatrixMeta::load(dir)?;
|
||||||
let n_cols = meta.n_cols;
|
let n_cols = meta.n_cols;
|
||||||
|
|
||||||
// Compute offsets from file sizes — no column data loaded into RAM.
|
// Compute offsets from file sizes — no column data loaded into RAM.
|
||||||
@@ -318,19 +294,6 @@ impl PersistentBitMatrix {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Column-major point lookup: value at column `c`, slot `slot`, as 0/1.
|
|
||||||
///
|
|
||||||
/// Unlike [`col_view`](Self::col_view), this never panics on `Implicit`
|
|
||||||
/// (every column reads as present, per the mono-genome fast path) — safe
|
|
||||||
/// to call for any `c < self.n_cols()`.
|
|
||||||
pub fn get(&self, c: usize, slot: usize) -> u32 {
|
|
||||||
match self {
|
|
||||||
Self::Columnar(m) => m.col(c).get(slot) as u32,
|
|
||||||
Self::Packed(m) => m.col_slice(c).get(slot) as u32,
|
|
||||||
Self::Implicit { .. } => 1,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
|
pub fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
|
||||||
match self {
|
match self {
|
||||||
Self::Columnar(m) => PersistentBitVecBuilder::build_from(m.col(c), path),
|
Self::Columnar(m) => PersistentBitVecBuilder::build_from(m.col(c), path),
|
||||||
@@ -537,26 +500,17 @@ where T: Clone + Default {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Compute a symmetric `n×n` matrix in parallel by evaluating `f(i,j)` for
|
/// Compute a symmetric `n×n` matrix in parallel by evaluating `f(i,j)` for
|
||||||
/// all upper-triangle pairs, plus `f(i,i)` for the diagonal. `T: Copy` avoids
|
/// all upper-triangle pairs. `T: Copy` avoids the `.clone()` needed for the
|
||||||
/// the `.clone()` needed for the lower-triangle mirror.
|
/// lower-triangle mirror.
|
||||||
///
|
|
||||||
/// The diagonal is *not* generally `T::default()`: for a self-comparison,
|
|
||||||
/// `f(i,i)` is often the column's own weight (e.g. intersection-with-self —
|
|
||||||
/// see `pairwise2_matrix`), not zero. Distance finalisations that need a
|
|
||||||
/// zero diagonal (self-distance) already overwrite it explicitly.
|
|
||||||
pub(crate) fn pairwise_matrix<T>(n: usize, f: impl Fn(usize, usize) -> T + Sync) -> Array2<T>
|
pub(crate) fn pairwise_matrix<T>(n: usize, f: impl Fn(usize, usize) -> T + Sync) -> Array2<T>
|
||||||
where T: Copy + Default + Send {
|
where T: Copy + Default + Send {
|
||||||
let results: Vec<(usize, usize, T)> = upper_pairs(n)
|
let results: Vec<(usize, usize, T)> = upper_pairs(n)
|
||||||
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
|
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
|
||||||
let mut m = fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)));
|
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
|
||||||
for i in 0..n { m[[i, i]] = f(i, i); }
|
|
||||||
m
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Same as `pairwise_matrix` but `f` returns two values that fill two
|
/// Same as `pairwise_matrix` but `f` returns two values that fill two
|
||||||
/// symmetric matrices simultaneously (e.g. intersection + union for Jaccard).
|
/// symmetric matrices simultaneously (e.g. intersection + union for Jaccard).
|
||||||
/// The diagonal is `f(i,i)` (e.g. a genome's kmer count intersected with
|
|
||||||
/// itself), not `T::default()` — see `pairwise_matrix` for why that matters.
|
|
||||||
pub(crate) fn pairwise2_matrix<T>(n: usize, f: impl Fn(usize, usize) -> (T, T) + Sync) -> (Array2<T>, Array2<T>)
|
pub(crate) fn pairwise2_matrix<T>(n: usize, f: impl Fn(usize, usize) -> (T, T) + Sync) -> (Array2<T>, Array2<T>)
|
||||||
where T: Copy + Default + Send {
|
where T: Copy + Default + Send {
|
||||||
let results: Vec<(usize, usize, T, T)> = upper_pairs(n)
|
let results: Vec<(usize, usize, T, T)> = upper_pairs(n)
|
||||||
@@ -569,10 +523,5 @@ where T: Copy + Default + Send {
|
|||||||
m0[[i, j]] = a; m0[[j, i]] = a;
|
m0[[i, j]] = a; m0[[j, i]] = a;
|
||||||
m1[[i, j]] = b; m1[[j, i]] = b;
|
m1[[i, j]] = b; m1[[j, i]] = b;
|
||||||
}
|
}
|
||||||
for i in 0..n {
|
|
||||||
let (a, b) = f(i, i);
|
|
||||||
m0[[i, i]] = a;
|
|
||||||
m1[[i, i]] = b;
|
|
||||||
}
|
|
||||||
(m0, m1)
|
(m0, m1)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
use std::fs::{self, File};
|
use std::fs::{self, File};
|
||||||
use std::io::{self, BufWriter, Read as _, Write as _};
|
use std::io::{self, BufWriter, Write as _};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use memmap2::Mmap;
|
use memmap2::Mmap;
|
||||||
@@ -228,44 +228,17 @@ impl PackedCompactIntMatrix {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reads just the `n_cols` field from an existing packed matrix's header,
|
|
||||||
/// without mapping the file. Used by `pack_compact_int_matrix` to tell a
|
|
||||||
/// genuinely complete pack from a stale one that predates a later
|
|
||||||
/// column-widening.
|
|
||||||
fn packed_int_matrix_n_cols(path: &Path) -> io::Result<usize> {
|
|
||||||
let mut f = File::open(path)?;
|
|
||||||
let mut header = [0u8; PCMX_HEADER];
|
|
||||||
f.read_exact(&mut header)?;
|
|
||||||
Ok(u64::from_le_bytes(header[16..24].try_into().unwrap()) as usize)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Build `counts/matrix.pcmx` from existing `col_*.pciv` files.
|
/// Build `counts/matrix.pcmx` from existing `col_*.pciv` files.
|
||||||
pub fn pack_compact_int_matrix(dir: &Path) -> io::Result<()> {
|
pub fn pack_compact_int_matrix(dir: &Path) -> io::Result<()> {
|
||||||
let packed_path = dir.join("matrix.pcmx");
|
let packed_path = dir.join("matrix.pcmx");
|
||||||
|
if packed_path.exists() {
|
||||||
let meta = match MatrixMeta::load(dir) {
|
if let Ok(meta) = MatrixMeta::load(dir) {
|
||||||
Ok(meta) => meta,
|
|
||||||
Err(e) => {
|
|
||||||
// No columnar data pending: either this layer was already
|
|
||||||
// packed and cleaned up (matrix.pcmx complete, nothing left to
|
|
||||||
// do), or genuinely nothing was ever written here.
|
|
||||||
return if packed_path.exists() { Ok(()) } else { Err(e) };
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// A `matrix.pcmx` can already exist here even though columnar data is
|
|
||||||
// still pending — e.g. copied verbatim from a merge's base source
|
|
||||||
// before this layer was widened with more genome columns (see
|
|
||||||
// `obikpartitionner::merge_partition`). Only skip (re-)packing if the
|
|
||||||
// existing file already reflects the current column count; otherwise
|
|
||||||
// the columnar files are newer and must be (re-)packed, overwriting the
|
|
||||||
// stale one — never silently discarded as "leftover cleanup".
|
|
||||||
if packed_int_matrix_n_cols(&packed_path).ok() == Some(meta.n_cols) {
|
|
||||||
for c in 0..meta.n_cols { let _ = fs::remove_file(col_path(dir, c)); }
|
for c in 0..meta.n_cols { let _ = fs::remove_file(col_path(dir, c)); }
|
||||||
let _ = fs::remove_file(dir.join("meta.json"));
|
let _ = fs::remove_file(dir.join("meta.json"));
|
||||||
|
}
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
let meta = MatrixMeta::load(dir)?;
|
||||||
let n_cols = meta.n_cols;
|
let n_cols = meta.n_cols;
|
||||||
let col_sizes: Vec<u64> = (0..n_cols)
|
let col_sizes: Vec<u64> = (0..n_cols)
|
||||||
.map(|c| fs::metadata(col_path(dir, c)).map(|m| m.len()))
|
.map(|c| fs::metadata(col_path(dir, c)).map(|m| m.len()))
|
||||||
|
|||||||
@@ -7,7 +7,6 @@ mod intmatrix;
|
|||||||
mod layer_meta;
|
mod layer_meta;
|
||||||
mod meta;
|
mod meta;
|
||||||
mod reader;
|
mod reader;
|
||||||
mod siblingannex;
|
|
||||||
mod tempbitvec;
|
mod tempbitvec;
|
||||||
mod tempintvec;
|
mod tempintvec;
|
||||||
mod views;
|
mod views;
|
||||||
@@ -19,7 +18,6 @@ pub use builder::PersistentCompactIntVecBuilder;
|
|||||||
pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask};
|
pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask};
|
||||||
pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix};
|
pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix};
|
||||||
pub use layer_meta::LayerMeta;
|
pub use layer_meta::LayerMeta;
|
||||||
pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
|
||||||
pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter};
|
pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter};
|
||||||
pub use tempbitvec::{TempBitVec, TempBitVecBuilder};
|
pub use tempbitvec::{TempBitVec, TempBitVecBuilder};
|
||||||
pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
|
pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
|
||||||
|
|||||||
@@ -1,245 +0,0 @@
|
|||||||
//! Family presence-mask annex: a compact, read-only-after-build, per-slot
|
|
||||||
//! derived value used by the central-position SNP distance estimator (see
|
|
||||||
//! `docmd/theory/evolutionary_distances.md`, "Step 2b" and "Definitions:
|
|
||||||
//! family, and the canonical form of a family").
|
|
||||||
//!
|
|
||||||
//! One byte is stored per MPHF slot of a partition/layer, its low 4 bits
|
|
||||||
//! encoding a **presence mask** for the slot's k-mer's "family" (the up to 4
|
|
||||||
//! k-mers sharing the same flanks, differing only at the central base):
|
|
||||||
//! bit `b` (`b` = 0..3, in the fixed A/C/G/T = 0/1/2/3 encoding already used
|
|
||||||
//! for a single nucleotide) is set iff the family member whose *own* central
|
|
||||||
//! base — in its own canonical orientation — is `b`, is observed anywhere in
|
|
||||||
//! the current multi-genome index. This is a property of the whole index,
|
|
||||||
//! not of any one genome.
|
|
||||||
//!
|
|
||||||
//! Both facts the earlier (superseded) 3-bit design stored explicitly are
|
|
||||||
//! derived from the mask instead, not stored:
|
|
||||||
//! - sibling count = `popcount(mask) - 1`;
|
|
||||||
//! - minorant = regenerate the family's 4 canonical forms from the slot's
|
|
||||||
//! own k-mer (`CanonicalKmerOf::central_canonical_neighbors`, cheap, no
|
|
||||||
//! lookup), compare the raw encodings of whichever are set in the mask,
|
|
||||||
//! take the smallest — see `obikindex::siblings`.
|
|
||||||
//!
|
|
||||||
//! Mask value 0 is logically unreachable as a real result (a slot's own base
|
|
||||||
//! is always present in its own family) and is reused as the "not yet
|
|
||||||
//! computed" sentinel: annex files are pre-initialised to all-zero, and a
|
|
||||||
//! real value is only ever written once, by the computation pass.
|
|
||||||
//!
|
|
||||||
//! Deliberately simpler than a true 4-bit pack (1 byte/slot instead of 4
|
|
||||||
//! bits/slot): correctness and simplicity first, for a first implementation.
|
|
||||||
//! Packing to 4 bits/slot is a pure storage-density follow-up, not a
|
|
||||||
//! behavioural change, left for later.
|
|
||||||
|
|
||||||
use std::fs::{File, OpenOptions};
|
|
||||||
use std::io;
|
|
||||||
use std::path::{Path, PathBuf};
|
|
||||||
|
|
||||||
use memmap2::{Mmap, MmapMut};
|
|
||||||
|
|
||||||
const MAGIC: [u8; 4] = *b"PSIB";
|
|
||||||
|
|
||||||
// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows.
|
|
||||||
const HEADER_SIZE: usize = 16;
|
|
||||||
|
|
||||||
/// A family presence mask: bit `b` set iff the member whose own canonical
|
|
||||||
/// central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the index.
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
||||||
pub struct FamilyMask(u8);
|
|
||||||
|
|
||||||
impl FamilyMask {
|
|
||||||
/// The empty mask — never a valid *computed* result (a slot's own base
|
|
||||||
/// is always present in its own family) — used only to build up a mask
|
|
||||||
/// via repeated [`with`](Self::with) calls before storing it.
|
|
||||||
pub const EMPTY: FamilyMask = FamilyMask(0);
|
|
||||||
|
|
||||||
/// Set bit `base` (0=A, 1=C, 2=G, 3=T).
|
|
||||||
#[inline]
|
|
||||||
pub fn with(self, base: u8) -> Self {
|
|
||||||
debug_assert!(base < 4, "base out of range: {base}");
|
|
||||||
FamilyMask(self.0 | (1 << base))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Is the member with central base `base` (0..3) present?
|
|
||||||
#[inline]
|
|
||||||
pub fn has(self, base: u8) -> bool {
|
|
||||||
debug_assert!(base < 4, "base out of range: {base}");
|
|
||||||
self.0 & (1 << base) != 0
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Number of family members observed anywhere in the index (1..=4).
|
|
||||||
#[inline]
|
|
||||||
pub fn family_size(self) -> u32 {
|
|
||||||
self.0.count_ones()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Number of *other* members observed (0..=3) — `family_size() - 1`.
|
|
||||||
#[inline]
|
|
||||||
pub fn siblings(self) -> u32 {
|
|
||||||
self.family_size() - 1
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Raw bitmask (bit `b` = base `b` present) — for callers that build up
|
|
||||||
/// a mask via their own bit operations (e.g. concurrently, via an
|
|
||||||
/// `AtomicU8`) and only need the `FamilyMask` wrapper at the end.
|
|
||||||
#[inline]
|
|
||||||
pub fn bits(self) -> u8 {
|
|
||||||
self.0
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Construct from a raw bitmask (only the low 4 bits are kept).
|
|
||||||
#[inline]
|
|
||||||
pub fn from_bits(bits: u8) -> Self {
|
|
||||||
FamilyMask(bits & 0b1111)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn encode(self) -> u8 {
|
|
||||||
self.0
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn decode(byte: u8) -> Option<Self> {
|
|
||||||
if byte == 0 {
|
|
||||||
// Unreachable for a real result — reserved as the "not yet
|
|
||||||
// computed" sentinel.
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
Some(FamilyMask(byte & 0b1111))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── SiblingAnnex (reader) ───────────────────────────────────────────────────
|
|
||||||
|
|
||||||
pub struct SiblingAnnex {
|
|
||||||
mmap: Mmap,
|
|
||||||
n: usize,
|
|
||||||
path: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl SiblingAnnex {
|
|
||||||
pub fn open(path: &Path) -> io::Result<Self> {
|
|
||||||
let mmap = unsafe { Mmap::map(&File::open(path)?)? };
|
|
||||||
if mmap.len() < HEADER_SIZE {
|
|
||||||
return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file too short"));
|
|
||||||
}
|
|
||||||
if mmap[0..4] != MAGIC {
|
|
||||||
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PSIB magic"));
|
|
||||||
}
|
|
||||||
let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
|
|
||||||
if mmap.len() < HEADER_SIZE + n {
|
|
||||||
return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file truncated"));
|
|
||||||
}
|
|
||||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn path(&self) -> &Path { &self.path }
|
|
||||||
pub fn len(&self) -> usize { self.n }
|
|
||||||
pub fn is_empty(&self) -> bool { self.n == 0 }
|
|
||||||
|
|
||||||
/// `None` means the slot has not (yet) been computed — see module docs.
|
|
||||||
pub fn get(&self, slot: usize) -> Option<FamilyMask> {
|
|
||||||
FamilyMask::decode(self.mmap[HEADER_SIZE + slot])
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── SiblingAnnexBuilder (writer) ────────────────────────────────────────────
|
|
||||||
|
|
||||||
pub struct SiblingAnnexBuilder {
|
|
||||||
mmap: MmapMut,
|
|
||||||
n: usize,
|
|
||||||
path: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl SiblingAnnexBuilder {
|
|
||||||
/// Create a new annex of `n` slots at `path`, pre-initialised to the
|
|
||||||
/// "not yet computed" sentinel (all-zero).
|
|
||||||
pub fn new(n: usize, path: &Path) -> io::Result<Self> {
|
|
||||||
let file_size = HEADER_SIZE + n;
|
|
||||||
let file = OpenOptions::new()
|
|
||||||
.read(true).write(true).create(true).truncate(true)
|
|
||||||
.open(path)?;
|
|
||||||
file.set_len(file_size as u64)?;
|
|
||||||
let mut mmap = unsafe { MmapMut::map_mut(&file)? };
|
|
||||||
mmap[0..4].copy_from_slice(&MAGIC);
|
|
||||||
mmap[4..8].copy_from_slice(&[0u8; 4]);
|
|
||||||
mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes());
|
|
||||||
// Data region left at 0 by `set_len`/mmap — the sentinel value.
|
|
||||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn len(&self) -> usize { self.n }
|
|
||||||
pub fn is_empty(&self) -> bool { self.n == 0 }
|
|
||||||
|
|
||||||
pub fn get(&self, slot: usize) -> Option<FamilyMask> {
|
|
||||||
FamilyMask::decode(self.mmap[HEADER_SIZE + slot])
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set(&mut self, slot: usize, mask: FamilyMask) {
|
|
||||||
// Redundant concurrent writes from independent recomputation paths
|
|
||||||
// converge to the same encoded byte for a given slot, so a plain
|
|
||||||
// store here is safe even without external synchronisation, as long
|
|
||||||
// as the byte write itself is atomic (true for a single aligned
|
|
||||||
// byte on every platform this project targets).
|
|
||||||
self.mmap[HEADER_SIZE + slot] = mask.encode();
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn close(self) -> io::Result<()> { self.mmap.flush() }
|
|
||||||
|
|
||||||
pub fn finish(self) -> io::Result<SiblingAnnex> {
|
|
||||||
let path = self.path.clone();
|
|
||||||
self.close()?;
|
|
||||||
SiblingAnnex::open(&path)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use tempfile::tempdir;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sentinel_is_zero_and_unset_slots_read_as_uncomputed() {
|
|
||||||
let dir = tempdir().unwrap();
|
|
||||||
let path = dir.path().join("test.psib");
|
|
||||||
let builder = SiblingAnnexBuilder::new(4, &path).unwrap();
|
|
||||||
for slot in 0..4 {
|
|
||||||
assert_eq!(builder.get(slot), None);
|
|
||||||
}
|
|
||||||
builder.close().unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn roundtrip_all_valid_masks() {
|
|
||||||
let dir = tempdir().unwrap();
|
|
||||||
let path = dir.path().join("test.psib");
|
|
||||||
let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap();
|
|
||||||
|
|
||||||
let masks = [
|
|
||||||
FamilyMask::EMPTY.with(0), // just A: family size 1
|
|
||||||
FamilyMask::EMPTY.with(0).with(3), // A + T: size 2
|
|
||||||
FamilyMask::EMPTY.with(1).with(2).with(3), // C+G+T: size 3
|
|
||||||
FamilyMask::EMPTY.with(0).with(1).with(2).with(3), // all 4
|
|
||||||
];
|
|
||||||
for (slot, mask) in masks.iter().enumerate() {
|
|
||||||
builder.set(slot, *mask);
|
|
||||||
}
|
|
||||||
let annex = builder.finish().unwrap();
|
|
||||||
for (slot, mask) in masks.iter().enumerate() {
|
|
||||||
assert_eq!(annex.get(slot), Some(*mask));
|
|
||||||
}
|
|
||||||
assert_eq!(annex.get(0).unwrap().siblings(), 0);
|
|
||||||
assert_eq!(annex.get(1).unwrap().siblings(), 1);
|
|
||||||
assert_eq!(annex.get(2).unwrap().siblings(), 2);
|
|
||||||
assert_eq!(annex.get(3).unwrap().siblings(), 3);
|
|
||||||
assert_eq!(annex.get(3).unwrap().family_size(), 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn has_reflects_individual_bits() {
|
|
||||||
let mask = FamilyMask::EMPTY.with(0).with(2);
|
|
||||||
assert!(mask.has(0));
|
|
||||||
assert!(!mask.has(1));
|
|
||||||
assert!(mask.has(2));
|
|
||||||
assert!(!mask.has(3));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,16 +1,5 @@
|
|||||||
use ndarray::{Array1, Array2};
|
use ndarray::{Array1, Array2};
|
||||||
|
|
||||||
/// Convert a Jaccard distance matrix (`1 - J`) into a Mash distance matrix, per
|
|
||||||
/// https://mash.readthedocs.io/en/latest/distances.html:
|
|
||||||
/// `D = -1/k * ln(2J / (1+J))`.
|
|
||||||
fn jaccard_to_mash(d_jaccard: &Array2<f64>, k: usize) -> Array2<f64> {
|
|
||||||
d_jaccard.mapv(|d| {
|
|
||||||
let j = 1.0 - d;
|
|
||||||
if j <= 0.0 { 1.0 }
|
|
||||||
else { -1.0 / k as f64 * (2.0 * j / (1.0 + j)).ln() }
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Column-level weight statistic — total count or presence count per column.
|
// ── Column-level weight statistic — total count or presence count per column.
|
||||||
/// Additive across layers and partitions; used as denominator in normalised distances.
|
/// Additive across layers and partitions; used as denominator in normalised distances.
|
||||||
///
|
///
|
||||||
@@ -85,12 +74,6 @@ pub trait CountPartials: ColumnWeights {
|
|||||||
m
|
m
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived
|
|
||||||
/// from the presence-threshold Jaccard distance.
|
|
||||||
fn threshold_mash_dist_matrix(&self, k: usize, threshold: u32) -> Array2<f64> {
|
|
||||||
jaccard_to_mash(&self.threshold_jaccard_dist_matrix(threshold), k)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn relfreq_bray_dist_matrix(&self) -> Array2<f64> {
|
fn relfreq_bray_dist_matrix(&self) -> Array2<f64> {
|
||||||
let global = self.col_weights();
|
let global = self.col_weights();
|
||||||
let mut m = self.partial_relfreq_bray(&global).mapv(|v| 1.0 - v);
|
let mut m = self.partial_relfreq_bray(&global).mapv(|v| 1.0 - v);
|
||||||
@@ -143,12 +126,6 @@ pub trait BitPartials: ColumnWeights {
|
|||||||
m
|
m
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived
|
|
||||||
/// from the Jaccard distance.
|
|
||||||
fn mash_dist_matrix(&self, k: usize) -> Array2<f64> {
|
|
||||||
jaccard_to_mash(&self.jaccard_dist_matrix(), k)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn hamming_dist_matrix(&self) -> Array2<u64> {
|
fn hamming_dist_matrix(&self) -> Array2<u64> {
|
||||||
self.partial_hamming()
|
self.partial_hamming()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,10 +0,0 @@
|
|||||||
[package]
|
|
||||||
name = "obikentropy"
|
|
||||||
version = "0.1.0"
|
|
||||||
edition = "2024"
|
|
||||||
|
|
||||||
[dependencies]
|
|
||||||
obikseq = { path = "../obikseq" }
|
|
||||||
|
|
||||||
[dev-dependencies]
|
|
||||||
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
|
||||||
@@ -1,41 +0,0 @@
|
|||||||
//! Normalized entropy of an isolated, already-built k-mer (e.g. one
|
|
||||||
//! reconstructed from an index's `unitigs.bin`, with no surrounding
|
|
||||||
//! sequence) — drives the window through [`EntropyTracker`] one base at a
|
|
||||||
//! time, exactly like the streaming path, so a `theta` threshold means the
|
|
||||||
//! same thing whether applied during index construction or after the fact
|
|
||||||
//! (e.g. `obikmer filter`).
|
|
||||||
|
|
||||||
use obikseq::CanonicalKmer;
|
|
||||||
|
|
||||||
use crate::tracker::EntropyTracker;
|
|
||||||
|
|
||||||
/// Extension trait: compute the normalized entropy of a single canonical
|
|
||||||
/// k-mer, independent of any surrounding sequence.
|
|
||||||
pub trait KmerEntropy {
|
|
||||||
/// Normalized entropy across sub-word orders `1..=level_max` (the
|
|
||||||
/// minimum is taken across orders). Lower means less complex; `theta`
|
|
||||||
/// in `index`/`filter` rejects k-mers with a score `< theta`.
|
|
||||||
fn entropy(&self, level_max: usize) -> f64;
|
|
||||||
}
|
|
||||||
|
|
||||||
impl KmerEntropy for CanonicalKmer {
|
|
||||||
fn entropy(&self, level_max: usize) -> f64 {
|
|
||||||
let raw = self.raw(); // left-aligned, 2 bits/base, MSB-first
|
|
||||||
let k = obikseq::params::k();
|
|
||||||
let mask = (!0u64) >> (64 - k * 2);
|
|
||||||
|
|
||||||
let mut tracker = EntropyTracker::new(k);
|
|
||||||
let mut rolling: u64 = 0;
|
|
||||||
for i in 0..k {
|
|
||||||
let shift = 64 - 2 * (i + 1);
|
|
||||||
let base = (raw >> shift) & 3;
|
|
||||||
rolling = ((rolling << 2) | base) & mask;
|
|
||||||
tracker.push(i + 1, rolling);
|
|
||||||
}
|
|
||||||
tracker.normalized_entropy(level_max)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
#[path = "tests/kmer_entropy.rs"]
|
|
||||||
mod tests;
|
|
||||||
@@ -1,17 +0,0 @@
|
|||||||
//! Normalized k-mer entropy: formulas, tables, and a streaming tracker.
|
|
||||||
//!
|
|
||||||
//! This crate holds every piece of the entropy computation described in
|
|
||||||
//! `docmd/theory/entropy.md`: the compile-time tables ([`table`], private),
|
|
||||||
//! the incremental accumulator ([`EntropyTracker`]) that callers compose
|
|
||||||
//! into their own streaming state, and the [`KmerEntropy`] convenience trait
|
|
||||||
//! for scoring a single, already-built k-mer.
|
|
||||||
|
|
||||||
#![deny(missing_docs)]
|
|
||||||
|
|
||||||
mod kmer_entropy;
|
|
||||||
mod ring;
|
|
||||||
mod table;
|
|
||||||
mod tracker;
|
|
||||||
|
|
||||||
pub use kmer_entropy::KmerEntropy;
|
|
||||||
pub use tracker::EntropyTracker;
|
|
||||||
@@ -1,40 +0,0 @@
|
|||||||
//! Stack-allocated ring buffer backing the sliding sub-word windows.
|
|
||||||
|
|
||||||
/// Fixed-capacity ring buffer backed by a stack array.
|
|
||||||
/// N must be a power of two; operations are branchless via `% N`.
|
|
||||||
pub(crate) struct Ring<T: Copy + Default, const N: usize> {
|
|
||||||
buf: [T; N],
|
|
||||||
head: usize,
|
|
||||||
len: usize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<T: Copy + Default, const N: usize> Ring<T, N> {
|
|
||||||
#[inline]
|
|
||||||
pub(crate) fn new() -> Self {
|
|
||||||
Self {
|
|
||||||
buf: [T::default(); N],
|
|
||||||
head: 0,
|
|
||||||
len: 0,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
pub(crate) fn clear(&mut self) {
|
|
||||||
self.len = 0;
|
|
||||||
self.head = 0;
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
pub(crate) fn push_back(&mut self, val: T) {
|
|
||||||
self.buf[(self.head + self.len) % N] = val;
|
|
||||||
self.len += 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
pub(crate) fn pop_front(&mut self) -> T {
|
|
||||||
let val = self.buf[self.head];
|
|
||||||
self.head = (self.head + 1) % N;
|
|
||||||
self.len -= 1;
|
|
||||||
val
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,30 +0,0 @@
|
|||||||
//! Compile-time tables backing the normalized k-mer entropy formula: the
|
|
||||||
//! max-entropy correction for small samples. See `docmd/theory/entropy.md`.
|
|
||||||
//!
|
|
||||||
//! Entropy is computed directly on raw (non-canonicalized) sub-words — no
|
|
||||||
//! equivalence-class folding. Empirically (see the discussion that produced
|
|
||||||
//! this crate's history), folding sub-words into circular/revcomp classes
|
|
||||||
//! before unfolding them back buys nothing for the invariances it was meant
|
|
||||||
//! to guarantee (both hold for raw sub-word entropy already, by a direct
|
|
||||||
//! bijection argument for revcomp and by the sliding window's own dynamics
|
|
||||||
//! for tandem repeats), while it measurably *weakens* detection of the
|
|
||||||
//! low-complexity sequences the filter exists to catch.
|
|
||||||
|
|
||||||
include!(concat!(env!("OUT_DIR"), "/entropy_tables.rs"));
|
|
||||||
|
|
||||||
pub(crate) const WS_MAX: usize = 6;
|
|
||||||
|
|
||||||
#[inline(always)]
|
|
||||||
pub(crate) const fn n_log_n(n: usize) -> f64 {
|
|
||||||
N_LOG_N[n]
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline(always)]
|
|
||||||
pub(crate) const fn emax(k: usize, ws: usize) -> f64 {
|
|
||||||
EMAX[k][ws]
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline(always)]
|
|
||||||
pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 {
|
|
||||||
LOG_NWORDS[k][ws]
|
|
||||||
}
|
|
||||||
@@ -1,52 +0,0 @@
|
|||||||
use super::*;
|
|
||||||
use obikseq::Sequence;
|
|
||||||
use obikseq::kmer::Kmer;
|
|
||||||
|
|
||||||
const K: usize = 21;
|
|
||||||
const LEVEL_MAX: usize = 6;
|
|
||||||
|
|
||||||
fn kmer_from_ascii(seq: &[u8]) -> CanonicalKmer {
|
|
||||||
obikseq::set_k(K);
|
|
||||||
Kmer::from_ascii(seq).expect("valid k-mer sequence").canonical()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn homopolymer_scores_lower_than_diverse_sequence() {
|
|
||||||
let homopolymer = kmer_from_ascii(b"AAAAAAAAAAAAAAAAAAAAA"); // 21 bases
|
|
||||||
let diverse = kmer_from_ascii(b"CATTAGCGTACCTGATCAGGT"); // 21 bases, same as used elsewhere in this workspace's tests
|
|
||||||
|
|
||||||
let e_homopolymer = homopolymer.entropy(LEVEL_MAX);
|
|
||||||
let e_diverse = diverse.entropy(LEVEL_MAX);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
e_homopolymer < e_diverse,
|
|
||||||
"homopolymer ({e_homopolymer}) should score lower than a diverse sequence ({e_diverse})"
|
|
||||||
);
|
|
||||||
// A pure homopolymer is the most degenerate case representable — its
|
|
||||||
// score should sit near the bottom of the range, not just "somewhat lower".
|
|
||||||
assert!(e_homopolymer < 0.3, "homopolymer entropy unexpectedly high: {e_homopolymer}");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn entropy_is_deterministic_for_the_same_kmer() {
|
|
||||||
let a = kmer_from_ascii(b"CATTAGCGTACCTGATCAGGT");
|
|
||||||
let b = kmer_from_ascii(b"CATTAGCGTACCTGATCAGGT");
|
|
||||||
assert_eq!(a.entropy(LEVEL_MAX), b.entropy(LEVEL_MAX));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn entropy_is_within_zero_one_range() {
|
|
||||||
let mut repeat = "AT".repeat(K / 2 + 1);
|
|
||||||
repeat.truncate(K);
|
|
||||||
|
|
||||||
for seq in [
|
|
||||||
"AAAAAAAAAAAAAAAAAAAAA".to_string(),
|
|
||||||
repeat,
|
|
||||||
"CATTAGCGTACCTGATCAGGT".to_string(),
|
|
||||||
] {
|
|
||||||
assert_eq!(seq.len(), K, "test sequence must be exactly K bases: {seq:?}");
|
|
||||||
let kmer = kmer_from_ascii(seq.as_bytes());
|
|
||||||
let e = kmer.entropy(LEVEL_MAX);
|
|
||||||
assert!((0.0..=1.0).contains(&e), "entropy {e} out of [0,1] for {seq:?}");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,255 +0,0 @@
|
|||||||
//! Incremental (streaming) normalized k-mer entropy.
|
|
||||||
//!
|
|
||||||
//! [`EntropyTracker`] maintains, over a sliding window of the last `k` bases,
|
|
||||||
//! the per-sub-word-size raw-word frequency statistics needed to evaluate
|
|
||||||
//! the corrected Shannon entropy described in `docmd/theory/entropy.md`,
|
|
||||||
//! updated in O(1) per base rather than recomputed from scratch. No
|
|
||||||
//! canonicalization is applied — each sub-word is tallied under its own raw
|
|
||||||
//! 2-bit-packed value; only the small-sample max-entropy correction departs
|
|
||||||
//! from a textbook Shannon entropy.
|
|
||||||
//!
|
|
||||||
//! It carries no notion of minimizers or superkmer segmentation — callers
|
|
||||||
//! that need both (e.g. `obiskbuilder::RollingStat`) compose an
|
|
||||||
//! `EntropyTracker` as a plain field alongside their own state, so the two
|
|
||||||
//! concerns update in the same streaming pass without being conflated in one
|
|
||||||
//! struct.
|
|
||||||
|
|
||||||
use crate::ring::Ring;
|
|
||||||
use crate::table::{WS_MAX, emax, log_nwords, n_log_n};
|
|
||||||
|
|
||||||
/// Incremental normalized-entropy accumulator over a sliding window of `k`
|
|
||||||
/// bases. Composed as a plain field by callers that also need other
|
|
||||||
/// per-base state (e.g. minimizer selection) in the same streaming pass.
|
|
||||||
pub struct EntropyTracker {
|
|
||||||
k: usize,
|
|
||||||
steady: bool,
|
|
||||||
|
|
||||||
// Sliding-window queues over the last `k` raw sub-words, one per word
|
|
||||||
// size — stack-allocated, capacity ≤ k ≤ 31.
|
|
||||||
k1q: Ring<u64, 32>,
|
|
||||||
k2q: Ring<u64, 32>,
|
|
||||||
k3q: Ring<u64, 32>,
|
|
||||||
k4q: Ring<u64, 32>,
|
|
||||||
k5q: Ring<u64, 32>,
|
|
||||||
k6q: Ring<u64, 32>,
|
|
||||||
|
|
||||||
// Frequency count arrays, indexed by the raw sub-word value (2 bits per
|
|
||||||
// base). Max count per cell ≤ k ≤ 31 → u8 is sufficient.
|
|
||||||
k1c: [u8; 4],
|
|
||||||
k2c: [u8; 16],
|
|
||||||
k3c: [u8; 64],
|
|
||||||
k4c: [u8; 256],
|
|
||||||
k5c: [u8; 1024],
|
|
||||||
k6c: [u8; 4096],
|
|
||||||
|
|
||||||
sum_f_log_f: [f64; WS_MAX + 1],
|
|
||||||
}
|
|
||||||
|
|
||||||
impl EntropyTracker {
|
|
||||||
/// New tracker for a window of `k` bases (1..=31).
|
|
||||||
pub fn new(k: usize) -> Self {
|
|
||||||
Self {
|
|
||||||
k,
|
|
||||||
steady: false,
|
|
||||||
k1q: Ring::new(),
|
|
||||||
k2q: Ring::new(),
|
|
||||||
k3q: Ring::new(),
|
|
||||||
k4q: Ring::new(),
|
|
||||||
k5q: Ring::new(),
|
|
||||||
k6q: Ring::new(),
|
|
||||||
k1c: [0; 4],
|
|
||||||
k2c: [0; 16],
|
|
||||||
k3c: [0; 64],
|
|
||||||
k4c: [0; 256],
|
|
||||||
k5c: [0; 1024],
|
|
||||||
k6c: [0; 4096],
|
|
||||||
sum_f_log_f: [0.0; WS_MAX + 1],
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Clear all accumulated state, ready to track a new window from
|
|
||||||
/// scratch (`k` is unchanged).
|
|
||||||
pub fn reset(&mut self) {
|
|
||||||
self.steady = false;
|
|
||||||
|
|
||||||
self.k1c.fill(0);
|
|
||||||
self.k2c.fill(0);
|
|
||||||
self.k3c.fill(0);
|
|
||||||
self.k4c.fill(0);
|
|
||||||
self.k5c.fill(0);
|
|
||||||
self.k6c.fill(0);
|
|
||||||
|
|
||||||
self.k1q.clear();
|
|
||||||
self.k2q.clear();
|
|
||||||
self.k3q.clear();
|
|
||||||
self.k4q.clear();
|
|
||||||
self.k5q.clear();
|
|
||||||
self.k6q.clear();
|
|
||||||
|
|
||||||
self.sum_f_log_f = [0.0; WS_MAX + 1];
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn update_sums_decrement<const K: usize>(sum_f_log_f: &mut [f64; WS_MAX + 1], f: usize) {
|
|
||||||
sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn update_sums_increment<const K: usize>(sum_f_log_f: &mut [f64; WS_MAX + 1], g: usize) {
|
|
||||||
sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g);
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Advance the window by one base. `received` is the caller's running
|
|
||||||
/// count of bases pushed so far (1-based, i.e. after this base);
|
|
||||||
/// `rolling_kmer` is the current right-aligned, 2-bit-packed k-mer
|
|
||||||
/// window (same convention as `obiskbuilder::RollingStat::rolling_k`).
|
|
||||||
pub fn push(&mut self, received: usize, rolling_kmer: u64) {
|
|
||||||
let raw1 = rolling_kmer & 3;
|
|
||||||
let raw2 = rolling_kmer & 15;
|
|
||||||
let raw3 = rolling_kmer & 63;
|
|
||||||
let raw4 = rolling_kmer & 255;
|
|
||||||
let raw5 = rolling_kmer & 1023;
|
|
||||||
let raw6 = rolling_kmer & 4095;
|
|
||||||
|
|
||||||
if received > self.k {
|
|
||||||
let old1 = self.k1q.pop_front();
|
|
||||||
let f1 = self.k1c[old1 as usize] as usize;
|
|
||||||
Self::update_sums_decrement::<1>(&mut self.sum_f_log_f, f1);
|
|
||||||
self.k1c[old1 as usize] -= 1;
|
|
||||||
|
|
||||||
let old2 = self.k2q.pop_front();
|
|
||||||
let f2 = self.k2c[old2 as usize] as usize;
|
|
||||||
Self::update_sums_decrement::<2>(&mut self.sum_f_log_f, f2);
|
|
||||||
self.k2c[old2 as usize] -= 1;
|
|
||||||
|
|
||||||
let old3 = self.k3q.pop_front();
|
|
||||||
let f3 = self.k3c[old3 as usize] as usize;
|
|
||||||
Self::update_sums_decrement::<3>(&mut self.sum_f_log_f, f3);
|
|
||||||
self.k3c[old3 as usize] -= 1;
|
|
||||||
|
|
||||||
let old4 = self.k4q.pop_front();
|
|
||||||
let f4 = self.k4c[old4 as usize] as usize;
|
|
||||||
Self::update_sums_decrement::<4>(&mut self.sum_f_log_f, f4);
|
|
||||||
self.k4c[old4 as usize] -= 1;
|
|
||||||
|
|
||||||
let old5 = self.k5q.pop_front();
|
|
||||||
let f5 = self.k5c[old5 as usize] as usize;
|
|
||||||
Self::update_sums_decrement::<5>(&mut self.sum_f_log_f, f5);
|
|
||||||
self.k5c[old5 as usize] -= 1;
|
|
||||||
|
|
||||||
let old6 = self.k6q.pop_front();
|
|
||||||
let f6 = self.k6c[old6 as usize] as usize;
|
|
||||||
Self::update_sums_decrement::<6>(&mut self.sum_f_log_f, f6);
|
|
||||||
self.k6c[old6 as usize] -= 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
if self.steady {
|
|
||||||
let g1 = self.k1c[raw1 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<1>(&mut self.sum_f_log_f, g1);
|
|
||||||
self.k1c[raw1 as usize] += 1;
|
|
||||||
self.k1q.push_back(raw1);
|
|
||||||
|
|
||||||
let g2 = self.k2c[raw2 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<2>(&mut self.sum_f_log_f, g2);
|
|
||||||
self.k2c[raw2 as usize] += 1;
|
|
||||||
self.k2q.push_back(raw2);
|
|
||||||
|
|
||||||
let g3 = self.k3c[raw3 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<3>(&mut self.sum_f_log_f, g3);
|
|
||||||
self.k3c[raw3 as usize] += 1;
|
|
||||||
self.k3q.push_back(raw3);
|
|
||||||
|
|
||||||
let g4 = self.k4c[raw4 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<4>(&mut self.sum_f_log_f, g4);
|
|
||||||
self.k4c[raw4 as usize] += 1;
|
|
||||||
self.k4q.push_back(raw4);
|
|
||||||
|
|
||||||
let g5 = self.k5c[raw5 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<5>(&mut self.sum_f_log_f, g5);
|
|
||||||
self.k5c[raw5 as usize] += 1;
|
|
||||||
self.k5q.push_back(raw5);
|
|
||||||
|
|
||||||
let g6 = self.k6c[raw6 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<6>(&mut self.sum_f_log_f, g6);
|
|
||||||
self.k6c[raw6 as usize] += 1;
|
|
||||||
self.k6q.push_back(raw6);
|
|
||||||
} else {
|
|
||||||
self.push_warmup_increments(received, raw1, raw2, raw3, raw4, raw5, raw6);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cold]
|
|
||||||
#[inline(never)]
|
|
||||||
fn push_warmup_increments(
|
|
||||||
&mut self,
|
|
||||||
received: usize,
|
|
||||||
raw1: u64, raw2: u64, raw3: u64,
|
|
||||||
raw4: u64, raw5: u64, raw6: u64,
|
|
||||||
) {
|
|
||||||
let g1 = self.k1c[raw1 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<1>(&mut self.sum_f_log_f, g1);
|
|
||||||
self.k1c[raw1 as usize] += 1;
|
|
||||||
self.k1q.push_back(raw1);
|
|
||||||
|
|
||||||
if received >= 2 {
|
|
||||||
let g2 = self.k2c[raw2 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<2>(&mut self.sum_f_log_f, g2);
|
|
||||||
self.k2c[raw2 as usize] += 1;
|
|
||||||
self.k2q.push_back(raw2);
|
|
||||||
|
|
||||||
if received >= 3 {
|
|
||||||
let g3 = self.k3c[raw3 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<3>(&mut self.sum_f_log_f, g3);
|
|
||||||
self.k3c[raw3 as usize] += 1;
|
|
||||||
self.k3q.push_back(raw3);
|
|
||||||
|
|
||||||
if received >= 4 {
|
|
||||||
let g4 = self.k4c[raw4 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<4>(&mut self.sum_f_log_f, g4);
|
|
||||||
self.k4c[raw4 as usize] += 1;
|
|
||||||
self.k4q.push_back(raw4);
|
|
||||||
|
|
||||||
if received >= 5 {
|
|
||||||
let g5 = self.k5c[raw5 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<5>(&mut self.sum_f_log_f, g5);
|
|
||||||
self.k5c[raw5 as usize] += 1;
|
|
||||||
self.k5q.push_back(raw5);
|
|
||||||
|
|
||||||
if received >= 6 {
|
|
||||||
let g6 = self.k6c[raw6 as usize] as usize;
|
|
||||||
Self::update_sums_increment::<6>(&mut self.sum_f_log_f, g6);
|
|
||||||
self.k6c[raw6 as usize] += 1;
|
|
||||||
self.k6q.push_back(raw6);
|
|
||||||
self.steady = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Normalized entropy at sub-word size `order` (1..=6). The caller is
|
|
||||||
/// responsible for not calling this before the window is full (`k`
|
|
||||||
/// bases pushed) — an empty/partial window yields a meaningless value.
|
|
||||||
pub fn entropy(&self, order: usize) -> f64 {
|
|
||||||
let k = self.k;
|
|
||||||
let em = emax(k, order);
|
|
||||||
if em <= 0.0 {
|
|
||||||
return 1.0;
|
|
||||||
}
|
|
||||||
let nwords = k - order + 1;
|
|
||||||
let log_nw = log_nwords(k, order);
|
|
||||||
let nw_f = nwords as f64;
|
|
||||||
let h_corr = log_nw - self.sum_f_log_f[order] / nw_f;
|
|
||||||
(h_corr / em).max(0.0)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Minimum of [`Self::entropy`] over sub-word sizes `1..=order_max`, same
|
|
||||||
/// caller responsibility re: window readiness as `entropy`.
|
|
||||||
pub fn normalized_entropy(&self, order_max: usize) -> f64 {
|
|
||||||
let min_e = (1..=order_max)
|
|
||||||
.map(|ws| self.entropy(ws))
|
|
||||||
.fold(f64::MAX, f64::min);
|
|
||||||
if min_e == f64::MAX { 1.0 } else { min_e }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -10,8 +10,6 @@ obiskio = { path = "../obiskio" }
|
|||||||
obisys = { path = "../obisys" }
|
obisys = { path = "../obisys" }
|
||||||
obicompactvec = { path = "../obicompactvec" }
|
obicompactvec = { path = "../obicompactvec" }
|
||||||
obilayeredmap = { path = "../obilayeredmap" }
|
obilayeredmap = { path = "../obilayeredmap" }
|
||||||
obiskbuilder = { path = "../obiskbuilder" }
|
|
||||||
obipipeline = { path = "../obipipeline" }
|
|
||||||
ndarray = "0.16"
|
ndarray = "0.16"
|
||||||
rayon = "1"
|
rayon = "1"
|
||||||
crossbeam-channel = "0.5"
|
crossbeam-channel = "0.5"
|
||||||
@@ -21,10 +19,6 @@ indicatif = "0.17"
|
|||||||
tracing = "0.1.44"
|
tracing = "0.1.44"
|
||||||
hwlocality = { version = "1.0.0-alpha.11", features = ["vendored"], optional = true }
|
hwlocality = { version = "1.0.0-alpha.11", features = ["vendored"], optional = true }
|
||||||
|
|
||||||
[dev-dependencies]
|
|
||||||
obiread = { path = "../obiread" }
|
|
||||||
tempfile = "3"
|
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["numa"]
|
default = ["numa"]
|
||||||
numa = ["hwlocality"]
|
numa = ["hwlocality"]
|
||||||
|
|||||||
@@ -14,8 +14,6 @@ pub enum DistanceMetric {
|
|||||||
Jaccard,
|
Jaccard,
|
||||||
/// Hamming distance (number of differing kmer positions) on presence/absence data.
|
/// Hamming distance (number of differing kmer positions) on presence/absence data.
|
||||||
Hamming,
|
Hamming,
|
||||||
/// Mash distance on presence/absence data (Jaccard-derived mutation-rate estimate).
|
|
||||||
Mash,
|
|
||||||
/// Bray-Curtis dissimilarity on raw counts.
|
/// Bray-Curtis dissimilarity on raw counts.
|
||||||
BrayCurtis,
|
BrayCurtis,
|
||||||
/// Bray-Curtis dissimilarity normalised by per-genome total counts.
|
/// Bray-Curtis dissimilarity normalised by per-genome total counts.
|
||||||
@@ -86,7 +84,6 @@ impl KmerIndex {
|
|||||||
DistanceMetric::Hellinger => CountPartials::hellinger_dist_matrix(&global),
|
DistanceMetric::Hellinger => CountPartials::hellinger_dist_matrix(&global),
|
||||||
DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global),
|
DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global),
|
||||||
DistanceMetric::Jaccard => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold),
|
DistanceMetric::Jaccard => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold),
|
||||||
DistanceMetric::Mash => CountPartials::threshold_mash_dist_matrix(&global, self.kmer_size(), presence_threshold),
|
|
||||||
DistanceMetric::Hamming => {
|
DistanceMetric::Hamming => {
|
||||||
return Err(OKIError::InvalidInput(
|
return Err(OKIError::InvalidInput(
|
||||||
"Hamming is only available for presence/absence indexes".into(),
|
"Hamming is only available for presence/absence indexes".into(),
|
||||||
@@ -111,7 +108,6 @@ impl KmerIndex {
|
|||||||
|
|
||||||
let matrix = match metric {
|
let matrix = match metric {
|
||||||
DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global),
|
DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global),
|
||||||
DistanceMetric::Mash => BitPartials::mash_dist_matrix(&global, self.kmer_size()),
|
|
||||||
DistanceMetric::Hamming => {
|
DistanceMetric::Hamming => {
|
||||||
BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64)
|
BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ mod numa;
|
|||||||
mod rebuild;
|
mod rebuild;
|
||||||
mod reindex;
|
mod reindex;
|
||||||
mod select;
|
mod select;
|
||||||
mod siblings;
|
|
||||||
mod stats;
|
mod stats;
|
||||||
|
|
||||||
pub use error::{OKIError, OKIResult};
|
pub use error::{OKIError, OKIResult};
|
||||||
@@ -19,4 +18,3 @@ pub use merge::MergeMode;
|
|||||||
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
|
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
|
||||||
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
|
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
|
||||||
pub use stats::IndexBitsPerKmer;
|
pub use stats::IndexBitsPerKmer;
|
||||||
pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
|
||||||
|
|||||||
+82
-192
@@ -20,7 +20,7 @@ use hwlocality::cpu::binding::CpuBindingFlags;
|
|||||||
use hwlocality::cpu::cpuset::CpuSet;
|
use hwlocality::cpu::cpuset::CpuSet;
|
||||||
#[cfg(feature = "numa")]
|
#[cfg(feature = "numa")]
|
||||||
use hwlocality::object::types::ObjectType;
|
use hwlocality::object::types::ObjectType;
|
||||||
use obisys::{CpuSample, IoSample};
|
use obisys::CpuSample;
|
||||||
use tracing::debug;
|
use tracing::debug;
|
||||||
|
|
||||||
// ── Public interface ──────────────────────────────────────────────────────────
|
// ── Public interface ──────────────────────────────────────────────────────────
|
||||||
@@ -70,16 +70,15 @@ pub fn build() -> NumaSetup {
|
|||||||
nodes.len(),
|
nodes.len(),
|
||||||
nodes.first().map_or(0, |v| v.len()),
|
nodes.first().map_or(0, |v| v.len()),
|
||||||
);
|
);
|
||||||
return NumaSetup {
|
return NumaSetup { pools, cpus_per_node: nodes };
|
||||||
pools,
|
|
||||||
cpus_per_node: nodes,
|
|
||||||
};
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// UMA fallback: single synthetic node, all cores, no pool, no pinning.
|
// UMA fallback: single synthetic node, all cores, no pool, no pinning.
|
||||||
let n_cores = obisys::effective_parallelism();
|
let n_cores = std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(1);
|
||||||
debug!("UMA: single synthetic node, {} core(s)", n_cores);
|
debug!("UMA: single synthetic node, {} core(s)", n_cores);
|
||||||
NumaSetup {
|
NumaSetup {
|
||||||
pools: vec![None],
|
pools: vec![None],
|
||||||
@@ -89,7 +88,9 @@ pub fn build() -> NumaSetup {
|
|||||||
|
|
||||||
#[cfg(not(feature = "numa"))]
|
#[cfg(not(feature = "numa"))]
|
||||||
pub fn build() -> NumaSetup {
|
pub fn build() -> NumaSetup {
|
||||||
let n_cores = obisys::effective_parallelism();
|
let n_cores = std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(1);
|
||||||
debug!("UMA: single synthetic node, {} core(s)", n_cores);
|
debug!("UMA: single synthetic node, {} core(s)", n_cores);
|
||||||
NumaSetup {
|
NumaSetup {
|
||||||
pools: vec![None],
|
pools: vec![None],
|
||||||
@@ -101,9 +102,7 @@ pub fn build() -> NumaSetup {
|
|||||||
/// Silently returns on any error so the thread still runs, just unbound.
|
/// Silently returns on any error so the thread still runs, just unbound.
|
||||||
#[cfg(feature = "numa")]
|
#[cfg(feature = "numa")]
|
||||||
pub fn pin_current_thread(cpu_indices: &[usize]) {
|
pub fn pin_current_thread(cpu_indices: &[usize]) {
|
||||||
let Ok(topology) = Topology::new() else {
|
let Ok(topology) = Topology::new() else { return };
|
||||||
return;
|
|
||||||
};
|
|
||||||
let mut cpuset = CpuSet::new();
|
let mut cpuset = CpuSet::new();
|
||||||
for &idx in cpu_indices {
|
for &idx in cpu_indices {
|
||||||
cpuset.set(idx);
|
cpuset.set(idx);
|
||||||
@@ -133,22 +132,7 @@ fn build_pool(cpus: &[usize]) -> Option<rayon::ThreadPool> {
|
|||||||
.ok()
|
.ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── PartitionRunner ─────────────────────────────────────────────────────────
|
// ── PartitionRunner ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
/// Growth step (fraction of a node's worker capacity added per activation
|
|
||||||
/// event, see [`NodeActivation::grow`]).
|
|
||||||
const GROWTH_DIVISOR: usize = 8;
|
|
||||||
/// Minimum CPU efficiency growth to activate more workers, as a fraction of
|
|
||||||
/// the size of the *last growth step* (e.g. `0.2` after adding 8 workers
|
|
||||||
/// requires the next check to show at least +1.6 cores of growth — 20 % of
|
|
||||||
/// the ~8 cores those 8 workers should contribute if the workload is truly
|
|
||||||
/// CPU-bound). Scaling by the last step's size — not the cumulative total —
|
|
||||||
/// keeps the bar meaningful regardless of how many workers are already
|
|
||||||
/// active, instead of demanding an ever-larger absolute jump as the pool
|
|
||||||
/// grows.
|
|
||||||
const CPU_SPAWN_THRESHOLD: f64 = 0.2;
|
|
||||||
/// Minimum I/O throughput growth (relative) to activate more workers.
|
|
||||||
const IO_SPAWN_THRESHOLD: f64 = 0.2;
|
|
||||||
|
|
||||||
struct NodeConfig {
|
struct NodeConfig {
|
||||||
pool: Option<Arc<rayon::ThreadPool>>,
|
pool: Option<Arc<rayon::ThreadPool>>,
|
||||||
@@ -158,23 +142,19 @@ struct NodeConfig {
|
|||||||
|
|
||||||
/// Generic NUMA-aware runner for partition-level parallel work.
|
/// Generic NUMA-aware runner for partition-level parallel work.
|
||||||
///
|
///
|
||||||
/// Workers are distributed evenly across NUMA nodes and pinned to their
|
/// Workers are distributed round-robin across NUMA nodes and pinned to their
|
||||||
/// node's CPUs. UMA is the degenerate case: one node, no pinning.
|
/// node's CPUs. UMA is the degenerate case: one node, no pinning.
|
||||||
///
|
///
|
||||||
/// Workers are pre-spawned dormant, one activation channel per node so
|
/// Workers are pre-spawned dormant and activated one by one as CPU efficiency
|
||||||
/// growth always targets a specific node rather than whichever dormant
|
/// falls below `SPAWN_THRESHOLD`. This avoids over-provisioning on I/O-bound
|
||||||
/// worker happens to wake up first on a shared channel. Growth (both the
|
/// or memory-bandwidth-bound workloads while saturating CPU-bound ones.
|
||||||
/// initial count and each subsequent step) is expressed as a fraction of
|
|
||||||
/// `workers_per_node`, applied identically to every node, so the pace of
|
|
||||||
/// ramp-up depends on node size rather than node count — a single-NUMA-node
|
|
||||||
/// (UMA) machine ramps just as fast as an 8-node one.
|
|
||||||
///
|
///
|
||||||
/// # Termination
|
/// # Termination
|
||||||
///
|
///
|
||||||
/// ```text
|
/// ```text
|
||||||
/// drop(part_tx) → part_rx drains → workers exit → drop their result_tx
|
/// drop(part_tx) → part_rx drains → workers exit → drop their result_tx
|
||||||
/// drop(result_tx) → result_rx closes → controller loop exits
|
/// drop(result_tx) → result_rx closes → controller loop exits
|
||||||
/// drop(activate_txs) → dormant workers exit cleanly
|
/// drop(activate_tx) → dormant workers exit cleanly
|
||||||
/// ```
|
/// ```
|
||||||
pub struct PartitionRunner {
|
pub struct PartitionRunner {
|
||||||
nodes: Vec<NodeConfig>,
|
nodes: Vec<NodeConfig>,
|
||||||
@@ -195,8 +175,7 @@ impl PartitionRunner {
|
|||||||
ns.pools.len(),
|
ns.pools.len(),
|
||||||
wpn,
|
wpn,
|
||||||
);
|
);
|
||||||
let nodes = ns
|
let nodes = ns.pools
|
||||||
.pools
|
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.zip(ns.cpus_per_node)
|
.zip(ns.cpus_per_node)
|
||||||
.map(|(pool, cpu_ids)| NodeConfig {
|
.map(|(pool, cpu_ids)| NodeConfig {
|
||||||
@@ -210,24 +189,23 @@ impl PartitionRunner {
|
|||||||
|
|
||||||
/// Run `f(i)` for every index in `order`.
|
/// Run `f(i)` for every index in `order`.
|
||||||
///
|
///
|
||||||
/// Workers are pre-spawned dormant and activated adaptively, per node:
|
/// Workers are pre-spawned dormant and activated adaptively. A timer thread
|
||||||
/// `(workers_per_node / INITIAL_DIVISOR).max(1)` are woken immediately on
|
/// fires a CPU-efficiency check every `TIMER_SECS` seconds; each completed
|
||||||
/// every node, then `(workers_per_node / GROWTH_DIVISOR).max(1)` more per
|
/// partition resets that timer (forcing an immediate check) and also
|
||||||
/// node each time the check below fires. A timer thread fires that check
|
/// triggers its own inline check. A new worker is activated whenever
|
||||||
/// every `TIMER_SECS` seconds; each completed partition resets that timer
|
/// efficiency falls below `SPAWN_THRESHOLD`.
|
||||||
/// (forcing an immediate check) and also triggers its own inline check. A
|
|
||||||
/// growth step happens whenever CPU efficiency grows by at least
|
|
||||||
/// `CPU_SPAWN_THRESHOLD` of what the last growth step should have
|
|
||||||
/// contributed, or I/O throughput grows by at least `IO_SPAWN_THRESHOLD`
|
|
||||||
/// (relative) since the last check — whichever resource is the actual
|
|
||||||
/// bottleneck still shows headroom.
|
|
||||||
///
|
///
|
||||||
/// `on_done(i, result, elapsed)` is called from the controller thread as
|
/// `on_done(i, result, elapsed)` is called from the controller thread as
|
||||||
/// each partition completes — suitable for progress bars and result
|
/// each partition completes — suitable for progress bars and result
|
||||||
/// aggregation.
|
/// aggregation.
|
||||||
///
|
///
|
||||||
/// Returns the first error produced by `f`, if any.
|
/// Returns the first error produced by `f`, if any.
|
||||||
pub fn run<F, R, E, C>(&self, order: &[usize], f: F, mut on_done: C) -> Result<(), E>
|
pub fn run<F, R, E, C>(
|
||||||
|
&self,
|
||||||
|
order: &[usize],
|
||||||
|
f: F,
|
||||||
|
mut on_done: C,
|
||||||
|
) -> Result<(), E>
|
||||||
where
|
where
|
||||||
F: Fn(usize) -> Result<R, E> + Send + Sync,
|
F: Fn(usize) -> Result<R, E> + Send + Sync,
|
||||||
R: Send,
|
R: Send,
|
||||||
@@ -239,28 +217,26 @@ impl PartitionRunner {
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const SPAWN_THRESHOLD: f64 = 0.95;
|
||||||
const TIMER_SECS: u64 = 30;
|
const TIMER_SECS: u64 = 30;
|
||||||
const INITIAL_DIVISOR: usize = 4;
|
|
||||||
|
let n_cores = std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(1);
|
||||||
|
|
||||||
// ── Channels ──────────────────────────────────────────────────────────
|
// ── Channels ──────────────────────────────────────────────────────────
|
||||||
let (part_tx, part_rx) = unbounded::<usize>();
|
let (part_tx, part_rx) = unbounded::<usize>();
|
||||||
|
let (activate_tx, activate_rx) = unbounded::<()>();
|
||||||
// reset_tx: controller → timer ("reset the 30 s window")
|
// reset_tx: controller → timer ("reset the 30 s window")
|
||||||
let (reset_tx, reset_rx) = unbounded::<()>();
|
let (reset_tx, reset_rx) = unbounded::<()>();
|
||||||
// event_tx: workers + timer → controller (unified event stream)
|
// event_tx: workers + timer → controller (unified event stream)
|
||||||
let (event_tx, event_rx) = unbounded::<WorkerEvent<R, E>>();
|
let (event_tx, event_rx) = unbounded::<WorkerEvent<R, E>>();
|
||||||
// One activation channel per node: growth always targets a specific
|
|
||||||
// node, rather than whichever dormant worker happens to win the race
|
|
||||||
// on a channel shared across all nodes.
|
|
||||||
let (activate_txs, activate_rxs): (Vec<_>, Vec<_>) =
|
|
||||||
(0..self.nodes.len()).map(|_| unbounded::<()>()).unzip();
|
|
||||||
|
|
||||||
for &i in order {
|
for &i in order { part_tx.send(i).ok(); }
|
||||||
part_tx.send(i).ok();
|
|
||||||
}
|
|
||||||
drop(part_tx);
|
drop(part_tx);
|
||||||
|
|
||||||
let max_workers = self.max_workers();
|
let max_workers = self.max_workers();
|
||||||
let node_caps: Vec<usize> = self.nodes.iter().map(|n| n.max_workers).collect();
|
let n_nodes = self.nodes.len();
|
||||||
let f = &f;
|
let f = &f;
|
||||||
|
|
||||||
let mut first_err: Option<E> = None;
|
let mut first_err: Option<E> = None;
|
||||||
@@ -284,23 +260,18 @@ impl PartitionRunner {
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
// ── Pre-spawn workers dormant, grouped by node ────────────────────
|
// ── Pre-spawn workers dormant, round-robin across NUMA nodes ──────
|
||||||
// Each worker listens on its own node's activation channel only.
|
for w in 0..max_workers {
|
||||||
for (node, arx) in self.nodes.iter().zip(activate_rxs.iter()) {
|
let node = &self.nodes[w % n_nodes];
|
||||||
let cpu_ids = &node.cpu_ids;
|
|
||||||
for _ in 0..node.max_workers {
|
|
||||||
let prx = part_rx.clone();
|
let prx = part_rx.clone();
|
||||||
let etx = event_tx.clone();
|
let etx = event_tx.clone();
|
||||||
let arx = arx.clone();
|
let arx = activate_rx.clone();
|
||||||
let pool = node.pool.clone();
|
let pool = node.pool.clone();
|
||||||
|
let cpu_ids = &node.cpu_ids;
|
||||||
|
|
||||||
s.spawn(move || {
|
s.spawn(move || {
|
||||||
if arx.recv().is_err() {
|
if arx.recv().is_err() { return; }
|
||||||
return;
|
if !cpu_ids.is_empty() { pin_current_thread(cpu_ids); }
|
||||||
}
|
|
||||||
if !cpu_ids.is_empty() {
|
|
||||||
pin_current_thread(cpu_ids);
|
|
||||||
}
|
|
||||||
for i in &prx {
|
for i in &prx {
|
||||||
let t = Instant::now();
|
let t = Instant::now();
|
||||||
let r = match &pool {
|
let r = match &pool {
|
||||||
@@ -311,17 +282,15 @@ impl PartitionRunner {
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
|
||||||
// Drop controller's event_tx: event_rx closes when all workers +
|
// Drop controller's event_tx: event_rx closes when all workers +
|
||||||
// timer have exited.
|
// timer have exited.
|
||||||
drop(event_tx);
|
drop(event_tx);
|
||||||
|
|
||||||
// ── Controller ────────────────────────────────────────────────────
|
// ── Controller ────────────────────────────────────────────────────
|
||||||
let mut activation = NodeActivation::new(&activate_txs, &node_caps, max_workers);
|
activate_tx.send(()).ok();
|
||||||
activation.activate_initial(INITIAL_DIVISOR, n_total);
|
let mut n_active = 1usize;
|
||||||
|
|
||||||
let mut cpu_sample = CpuSample::now();
|
let mut cpu_sample = CpuSample::now();
|
||||||
let mut io_sample = IoSample::now();
|
let mut eff_at_last_spawn = 0.0f64; // 0 = no previous spawn to evaluate
|
||||||
let mut completed = 0usize;
|
let mut completed = 0usize;
|
||||||
|
|
||||||
while completed < n_total {
|
while completed < n_total {
|
||||||
@@ -330,39 +299,30 @@ impl PartitionRunner {
|
|||||||
WorkerEvent::Completed(i, r, dur) => {
|
WorkerEvent::Completed(i, r, dur) => {
|
||||||
match r {
|
match r {
|
||||||
Ok(v) => on_done(i, v, dur),
|
Ok(v) => on_done(i, v, dur),
|
||||||
Err(e) => {
|
Err(e) => { if first_err.is_none() { first_err = Some(e); } }
|
||||||
if first_err.is_none() {
|
|
||||||
first_err = Some(e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
completed += 1;
|
completed += 1;
|
||||||
// Reset the 30 s timer.
|
// Reset the 30 s timer.
|
||||||
reset_tx.send(()).ok();
|
reset_tx.send(()).ok();
|
||||||
// Inline check: same logic as a timer tick.
|
// Inline check: same logic as a timer tick.
|
||||||
maybe_activate(
|
maybe_activate(
|
||||||
&mut activation,
|
&activate_tx, &mut n_active, max_workers,
|
||||||
&mut cpu_sample,
|
&mut cpu_sample, &mut eff_at_last_spawn,
|
||||||
&mut io_sample,
|
n_cores, SPAWN_THRESHOLD, completed, n_total,
|
||||||
completed,
|
|
||||||
n_total,
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
WorkerEvent::TimerTick => {
|
WorkerEvent::TimerTick => {
|
||||||
maybe_activate(
|
maybe_activate(
|
||||||
&mut activation,
|
&activate_tx, &mut n_active, max_workers,
|
||||||
&mut cpu_sample,
|
&mut cpu_sample, &mut eff_at_last_spawn,
|
||||||
&mut io_sample,
|
n_cores, SPAWN_THRESHOLD, completed, n_total,
|
||||||
completed,
|
|
||||||
n_total,
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Dormant workers exit once every sender for their node's channel
|
// Dormant workers exit when activate_tx closes.
|
||||||
// is dropped — `activate_txs` holds the only ones.
|
drop(activate_tx);
|
||||||
drop(activate_txs);
|
|
||||||
// Timer thread exits when reset_tx closes.
|
// Timer thread exits when reset_tx closes.
|
||||||
drop(reset_tx);
|
drop(reset_tx);
|
||||||
});
|
});
|
||||||
@@ -381,113 +341,43 @@ enum WorkerEvent<R, E> {
|
|||||||
TimerTick,
|
TimerTick,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tracks how many of each node's dormant workers have been woken, and
|
|
||||||
/// grows every node by the same amount at each step (capped by that node's
|
|
||||||
/// remaining dormant workers and by the run's total budget) so load stays
|
|
||||||
/// balanced across nodes at every point in time — never just "one more
|
|
||||||
/// worker somewhere". Also remembers the size of the last real growth step
|
|
||||||
/// (`last_step`), used to scale the CPU activation threshold to what that
|
|
||||||
/// step could plausibly have contributed (see `maybe_activate`).
|
|
||||||
struct NodeActivation<'a> {
|
|
||||||
txs: &'a [crossbeam_channel::Sender<()>],
|
|
||||||
caps: &'a [usize],
|
|
||||||
active: Vec<usize>,
|
|
||||||
total: usize,
|
|
||||||
max: usize,
|
|
||||||
last_step: usize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<'a> NodeActivation<'a> {
|
|
||||||
fn new(txs: &'a [crossbeam_channel::Sender<()>], caps: &'a [usize], max: usize) -> Self {
|
|
||||||
Self {
|
|
||||||
txs,
|
|
||||||
caps,
|
|
||||||
active: vec![0; txs.len()],
|
|
||||||
total: 0,
|
|
||||||
max,
|
|
||||||
last_step: 0,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn total(&self) -> usize {
|
|
||||||
self.total
|
|
||||||
}
|
|
||||||
fn last_step(&self) -> usize {
|
|
||||||
self.last_step
|
|
||||||
}
|
|
||||||
fn max(&self) -> usize {
|
|
||||||
self.max
|
|
||||||
}
|
|
||||||
fn is_full(&self) -> bool {
|
|
||||||
self.total >= self.max
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Wake up to `(node_cap / divisor).max(1)` dormant workers on every
|
|
||||||
/// node, capped by `n_total`. Called once at startup, unconditionally.
|
|
||||||
fn activate_initial(&mut self, divisor: usize, n_total: usize) {
|
|
||||||
self.grow(divisor, n_total);
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Same per-node sizing as [`activate_initial`](Self::activate_initial),
|
|
||||||
/// applied as a growth step. Returns the number of workers actually
|
|
||||||
/// activated (may be less than requested once a node or the total
|
|
||||||
/// budget is exhausted). Updates `last_step` when it actually grew.
|
|
||||||
fn grow(&mut self, divisor: usize, n_total: usize) -> usize {
|
|
||||||
let before = self.total;
|
|
||||||
for idx in 0..self.txs.len() {
|
|
||||||
let wanted = (self.caps[idx] / divisor).max(1);
|
|
||||||
let room = self.caps[idx].saturating_sub(self.active[idx]);
|
|
||||||
let grow = wanted.min(room).min(n_total.saturating_sub(self.total));
|
|
||||||
for _ in 0..grow {
|
|
||||||
self.txs[idx].send(()).ok();
|
|
||||||
}
|
|
||||||
self.active[idx] += grow;
|
|
||||||
self.total += grow;
|
|
||||||
}
|
|
||||||
let grew = self.total - before;
|
|
||||||
if grew > 0 {
|
|
||||||
self.last_step = grew;
|
|
||||||
}
|
|
||||||
grew
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn maybe_activate(
|
fn maybe_activate(
|
||||||
activation: &mut NodeActivation,
|
activate_tx: &crossbeam_channel::Sender<()>,
|
||||||
|
n_active: &mut usize,
|
||||||
|
max_workers: usize,
|
||||||
cpu_sample: &mut CpuSample,
|
cpu_sample: &mut CpuSample,
|
||||||
io_sample: &mut IoSample,
|
eff_at_last_spawn: &mut f64,
|
||||||
|
n_cores: usize,
|
||||||
|
threshold: f64,
|
||||||
completed: usize,
|
completed: usize,
|
||||||
n_total: usize,
|
n_total: usize,
|
||||||
) {
|
) {
|
||||||
if activation.is_full() || completed >= n_total {
|
if *n_active >= max_workers || completed >= n_total { return; }
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Expect roughly 1 core of extra efficiency per worker activated in the
|
let eff = cpu_sample.cpu_efficiency(n_cores);
|
||||||
// last growth step (CPU-bound case); require at least CPU_SPAWN_THRESHOLD
|
if eff >= threshold { return; } // CPU already saturated
|
||||||
// (20 %) of that expected gain before growing again. Scaling by the last
|
|
||||||
// step's size — not the cumulative total — keeps the bar meaningful
|
|
||||||
// regardless of how many workers are already active: growing by 8 should
|
|
||||||
// always take ~+1.6 cores to confirm, whether that's the 2nd growth step
|
|
||||||
// or the 20th.
|
|
||||||
let cpu_threshold = CPU_SPAWN_THRESHOLD * activation.last_step() as f64;
|
|
||||||
|
|
||||||
// Call both unconditionally (no `||` short-circuit): each sampler must
|
// Check that the previous activation was beneficial enough.
|
||||||
// advance its own window every tick, regardless of what the other one
|
// Going from k-1 → k workers, the minimum acceptable speedup is (k-1+0.2)/(k-1).
|
||||||
// reports, or it would starve behind whichever signal fires first.
|
// For the very first extra worker (n_active == 1, no previous spawn), skip this
|
||||||
let cpu_wants_more = cpu_sample.do_i_activate(cpu_threshold);
|
// check: eff_at_last_spawn == 0 acts as the sentinel.
|
||||||
let io_wants_more = io_sample.do_i_activate(IO_SPAWN_THRESHOLD * activation.last_step() as f64);
|
let last_spawn_was_beneficial = if *eff_at_last_spawn < 1e-9 {
|
||||||
if !(cpu_wants_more || io_wants_more) {
|
true // first additional worker: no prior data to evaluate
|
||||||
return;
|
} else {
|
||||||
}
|
let k_before = (*n_active - 1) as f64;
|
||||||
|
let min_speedup = (k_before + 0.2) / k_before;
|
||||||
|
let actual_speedup = eff / *eff_at_last_spawn;
|
||||||
|
actual_speedup >= min_speedup
|
||||||
|
};
|
||||||
|
|
||||||
let grew = activation.grow(GROWTH_DIVISOR, n_total);
|
if last_spawn_was_beneficial {
|
||||||
if grew > 0 {
|
activate_tx.send(()).ok();
|
||||||
|
*eff_at_last_spawn = eff;
|
||||||
|
*n_active += 1;
|
||||||
|
*cpu_sample = CpuSample::now();
|
||||||
debug!(
|
debug!(
|
||||||
"activated {} worker(s) — {}/{} active",
|
"activated worker {}/{} — efficiency {:.0}%",
|
||||||
grew,
|
n_active, max_workers, eff * 100.0,
|
||||||
activation.total(),
|
|
||||||
activation.max()
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "obikmer"
|
name = "obikmer"
|
||||||
version = "1.1.43"
|
version = "1.1.17"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
|
||||||
[[bin]]
|
[[bin]]
|
||||||
@@ -18,7 +18,7 @@ obikrope = { path = "../obikrope" }
|
|||||||
obikpartitionner = { path = "../obikpartitionner" }
|
obikpartitionner = { path = "../obikpartitionner" }
|
||||||
obisys = { path = "../obisys" }
|
obisys = { path = "../obisys" }
|
||||||
obiskio = { path = "../obiskio" }
|
obiskio = { path = "../obiskio" }
|
||||||
obikindex = { path = "../obikindex", default-features = false }
|
obikindex = { path = "../obikindex" }
|
||||||
obitaxonomy = { path = "../obitaxonomy" }
|
obitaxonomy = { path = "../obitaxonomy" }
|
||||||
obilayeredmap = { path = "../obilayeredmap" }
|
obilayeredmap = { path = "../obilayeredmap" }
|
||||||
clap = { version = "4", features = ["derive"] }
|
clap = { version = "4", features = ["derive"] }
|
||||||
@@ -33,6 +33,4 @@ tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
|
|||||||
pprof = { version = "0.13", features = ["prost-codec"], optional = true }
|
pprof = { version = "0.13", features = ["prost-codec"], optional = true }
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["numa"]
|
|
||||||
numa = ["obikindex/numa"]
|
|
||||||
profiling = ["dep:pprof"]
|
profiling = ["dep:pprof"]
|
||||||
|
|||||||
@@ -38,7 +38,9 @@ pub struct CommonArgs {
|
|||||||
#[arg(
|
#[arg(
|
||||||
short = 'T',
|
short = 'T',
|
||||||
long,
|
long,
|
||||||
default_value_t = obisys::effective_parallelism()
|
default_value_t = std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(1)
|
||||||
)]
|
)]
|
||||||
pub threads: usize,
|
pub threads: usize,
|
||||||
|
|
||||||
|
|||||||
@@ -3,15 +3,13 @@ use std::path::PathBuf;
|
|||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use kodama::{Method, linkage};
|
use kodama::{Method, linkage};
|
||||||
use obifastwrite::{JsonVal, write_record};
|
use obikindex::{DistanceMetric, KmerIndex};
|
||||||
use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment};
|
|
||||||
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
|
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
#[derive(clap::ValueEnum, Clone, Copy, Debug)]
|
#[derive(clap::ValueEnum, Clone, Copy, Debug)]
|
||||||
pub enum MetricArg {
|
pub enum MetricArg {
|
||||||
Jaccard,
|
Jaccard,
|
||||||
Mash,
|
|
||||||
Hamming,
|
Hamming,
|
||||||
BrayCurtis,
|
BrayCurtis,
|
||||||
#[value(name = "relfreq-bray-curtis")]
|
#[value(name = "relfreq-bray-curtis")]
|
||||||
@@ -28,7 +26,6 @@ impl From<MetricArg> for DistanceMetric {
|
|||||||
fn from(m: MetricArg) -> Self {
|
fn from(m: MetricArg) -> Self {
|
||||||
match m {
|
match m {
|
||||||
MetricArg::Jaccard => DistanceMetric::Jaccard,
|
MetricArg::Jaccard => DistanceMetric::Jaccard,
|
||||||
MetricArg::Mash => DistanceMetric::Mash,
|
|
||||||
MetricArg::Hamming => DistanceMetric::Hamming,
|
MetricArg::Hamming => DistanceMetric::Hamming,
|
||||||
MetricArg::BrayCurtis => DistanceMetric::BrayCurtis,
|
MetricArg::BrayCurtis => DistanceMetric::BrayCurtis,
|
||||||
MetricArg::RelfreqBrayCurtis => DistanceMetric::RelfreqBrayCurtis,
|
MetricArg::RelfreqBrayCurtis => DistanceMetric::RelfreqBrayCurtis,
|
||||||
@@ -65,37 +62,7 @@ pub struct DistanceArgs {
|
|||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub upgma: bool,
|
pub upgma: bool,
|
||||||
|
|
||||||
/// Build the sibling-count/minorant annex on this (multi-genome) index
|
|
||||||
/// — see `docmd/theory/evolutionary_distances.md`, Step 2b. Construction
|
|
||||||
/// only; does not by itself compute or write any statistics.
|
|
||||||
#[arg(long)]
|
|
||||||
pub sibling_annex: bool,
|
|
||||||
|
|
||||||
/// Tally the sibling-count distribution (CSV) of an already-built annex
|
|
||||||
/// (run with `--sibling-annex` first, in this invocation or an earlier
|
|
||||||
/// one). A separate, occasional diagnostic pass — not run every time the
|
|
||||||
/// annex itself is (re)built.
|
|
||||||
#[arg(long)]
|
|
||||||
pub sibling_stats: bool,
|
|
||||||
|
|
||||||
/// Compute the raw p-distance restricted to loci that are single-copy
|
|
||||||
/// in both genomes of each pair (an already-built sibling annex is
|
|
||||||
/// required — run with `--sibling-annex` first, in this invocation or
|
|
||||||
/// an earlier one). A quick way to test the central-position SNP
|
|
||||||
/// estimator against a real index; not the full `SnpTally` design.
|
|
||||||
#[arg(long)]
|
|
||||||
pub raw_snp_distance: bool,
|
|
||||||
|
|
||||||
/// Write a SNP-only pseudo-alignment (FASTA, IUPAC-coded) from an
|
|
||||||
/// already-built sibling annex — one row per genome, one column per
|
|
||||||
/// variable family (monomorphic families skipped), no flanking
|
|
||||||
/// sequence. See `docmd/theory/evolutionary_distances.md`,
|
|
||||||
/// "Multi-genome framing: family as pseudo-alignment column".
|
|
||||||
#[arg(long)]
|
|
||||||
pub snp: bool,
|
|
||||||
|
|
||||||
/// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
|
/// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
|
||||||
/// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta,
|
|
||||||
/// <prefix>_nj.nwk, <prefix>_upgma.nwk.
|
/// <prefix>_nj.nwk, <prefix>_upgma.nwk.
|
||||||
/// If omitted, the distance matrix is written to stdout.
|
/// If omitted, the distance matrix is written to stdout.
|
||||||
#[arg(short, long)]
|
#[arg(short, long)]
|
||||||
@@ -111,51 +78,6 @@ pub fn run(args: DistanceArgs) {
|
|||||||
|
|
||||||
let labels: Vec<String> = idx.meta().genomes.iter().map(|g| g.label.clone()).collect();
|
let labels: Vec<String> = idx.meta().genomes.iter().map(|g| g.label.clone()).collect();
|
||||||
let n = labels.len();
|
let n = labels.len();
|
||||||
|
|
||||||
// ── Sibling-count/minorant annex (independent of the distance metric) ──
|
|
||||||
// Construction (`--sibling-annex`) and stats (`--sibling-stats`) are
|
|
||||||
// deliberately decoupled: the annex is meant to be (re)built routinely,
|
|
||||||
// the distribution only occasionally, on demand.
|
|
||||||
if args.sibling_annex {
|
|
||||||
info!("building sibling-count/minorant annex");
|
|
||||||
idx.build_sibling_annex().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error building sibling annex: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
}
|
|
||||||
if args.sibling_stats {
|
|
||||||
let stats = idx.sibling_annex_stats().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error computing sibling-annex stats: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
write_sibling_stats_csv(&stats, &labels, &args.output);
|
|
||||||
}
|
|
||||||
if args.raw_snp_distance {
|
|
||||||
let result = idx.raw_snp_distance().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error computing raw SNP distance: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
write_raw_snp_distance_csv(&result, &labels, &args.output);
|
|
||||||
}
|
|
||||||
if args.snp {
|
|
||||||
let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error computing SNP pseudo-alignment: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
write_snp_fasta(&alignment, &labels, &args.output);
|
|
||||||
}
|
|
||||||
|
|
||||||
// `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp` are
|
|
||||||
// their own operation, not a modifier on top of a distance-metric
|
|
||||||
// computation — a metric was never requested by asking for any of them,
|
|
||||||
// so there is nothing for the rest of this function to compute. Not a
|
|
||||||
// historical accident to keep: stop here rather than always also
|
|
||||||
// running a Jaccard (or whichever `--metric` defaults to) pass and
|
|
||||||
// printing an unrequested matrix.
|
|
||||||
if args.sibling_annex || args.sibling_stats || args.raw_snp_distance || args.snp {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
info!(
|
info!(
|
||||||
"computing {:?} distances for {} genome(s)",
|
"computing {:?} distances for {} genome(s)",
|
||||||
args.metric, n
|
args.metric, n
|
||||||
@@ -267,103 +189,6 @@ pub fn run(args: DistanceArgs) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Family-size distribution → CSV ──────────────────────────────────────────
|
|
||||||
//
|
|
||||||
// Each row is a family (the up-to-4 k-mers sharing flanks, differing only at
|
|
||||||
// the centre), counted once — at its minorant — regardless of how many of
|
|
||||||
// its members are observed. Family size 1..4 (not "sibling count" 0..3):
|
|
||||||
// see `docmd/theory/evolutionary_distances.md`, "Definitions".
|
|
||||||
|
|
||||||
fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: &Option<PathBuf>) {
|
|
||||||
// One row per genome (4 columns, family size 1-4: number of families of
|
|
||||||
// that size for which the genome carries at least one member), plus a
|
|
||||||
// `global` row — the actual deduplicated family-size histogram
|
|
||||||
// (`stats.counts`), NOT a sum of the per-genome columns (a family shared
|
|
||||||
// by several genomes would otherwise be counted once per genome it
|
|
||||||
// appears in, inflating the total beyond the real family count).
|
|
||||||
let path = output.as_ref()
|
|
||||||
.map(|p| format!("{}_siblings.csv", p.display()))
|
|
||||||
.unwrap_or_else(|| "siblings.csv".into());
|
|
||||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error creating {path}: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
}));
|
|
||||||
writeln!(f, "genome,1,2,3,4").unwrap();
|
|
||||||
for (label, counts) in labels.iter().zip(stats.per_genome.iter()) {
|
|
||||||
writeln!(f, "{label},{},{},{},{}", counts[0], counts[1], counts[2], counts[3]).unwrap();
|
|
||||||
}
|
|
||||||
writeln!(
|
|
||||||
f, "global,{},{},{},{}",
|
|
||||||
stats.counts[0], stats.counts[1], stats.counts[2], stats.counts[3],
|
|
||||||
).unwrap();
|
|
||||||
let total: u64 = stats.counts.iter().sum();
|
|
||||||
info!("family-size distribution → {path} (total {total} famil{})",
|
|
||||||
if total == 1 { "y" } else { "ies" });
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Raw single-copy SNP distance → CSV ──────────────────────────────────────
|
|
||||||
//
|
|
||||||
// p_hat[i,j] = snp[i,j] / (snp[i,j] + shared[i,j]) over loci single-copy in
|
|
||||||
// both i and j — see `RawSnpDistanceOutput` / `KmerIndex::raw_snp_distance`.
|
|
||||||
// A single file: the distance matrix, with an eligible-loci count alongside
|
|
||||||
// each value so a 0/0 pair (no eligible locus at all) is distinguishable
|
|
||||||
// from a genuinely identical pair.
|
|
||||||
|
|
||||||
fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option<PathBuf>) {
|
|
||||||
let path = output.as_ref()
|
|
||||||
.map(|p| format!("{}_rawsnp.csv", p.display()))
|
|
||||||
.unwrap_or_else(|| "rawsnp.csv".into());
|
|
||||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error creating {path}: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
}));
|
|
||||||
let n = labels.len();
|
|
||||||
write!(f, "genome").unwrap();
|
|
||||||
for g in labels { write!(f, ",{g}").unwrap(); }
|
|
||||||
writeln!(f).unwrap();
|
|
||||||
for (i, g) in labels.iter().enumerate() {
|
|
||||||
write!(f, "{g}").unwrap();
|
|
||||||
for j in 0..n {
|
|
||||||
let snp = result.snp[[i, j]];
|
|
||||||
let shared = result.shared[[i, j]];
|
|
||||||
let eligible = snp + shared;
|
|
||||||
if eligible == 0 {
|
|
||||||
write!(f, ",NA").unwrap();
|
|
||||||
} else {
|
|
||||||
write!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
writeln!(f).unwrap();
|
|
||||||
}
|
|
||||||
info!("raw single-copy SNP distance matrix → {path}");
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── SNP-only pseudo-alignment → FASTA ───────────────────────────────────────
|
|
||||||
//
|
|
||||||
// One record per genome, IUPAC-coded, no flanking sequence — see
|
|
||||||
// `SnpAlignment` / `KmerIndex::snp_pseudo_alignment`. Uses the project's
|
|
||||||
// existing FASTA writer (`obifastwrite::write_record`) rather than
|
|
||||||
// hand-rolling one.
|
|
||||||
|
|
||||||
fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option<PathBuf>) {
|
|
||||||
let path = output.as_ref()
|
|
||||||
.map(|p| format!("{}_snp.fasta", p.display()))
|
|
||||||
.unwrap_or_else(|| "snp.fasta".into());
|
|
||||||
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error creating {path}: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
}));
|
|
||||||
let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
|
|
||||||
for (label, seq) in labels.iter().zip(alignment.sequences.iter()) {
|
|
||||||
write_record(seq, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error writing {path}: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
}
|
|
||||||
info!("SNP pseudo-alignment → {path} ({n_sites} site{})",
|
|
||||||
if n_sites == 1 { "" } else { "s" });
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────
|
// ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────
|
||||||
|
|
||||||
fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String {
|
fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String {
|
||||||
|
|||||||
@@ -2,14 +2,14 @@ use std::path::PathBuf;
|
|||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikindex::{KmerIndex, MergeMode};
|
use obikindex::{KmerIndex, MergeMode};
|
||||||
use obikpartitionner::filter::{MaxTotalCount, MinComplexity, MinTotalCount};
|
use obikpartitionner::filter::{MaxTotalCount, MinTotalCount};
|
||||||
use obisys::Reporter;
|
use obisys::Reporter;
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
use super::predicate::FilterArgs as KmerFilterArgs;
|
use super::predicate::FilterArgs as KmerFilterArgs;
|
||||||
|
|
||||||
#[derive(Args)]
|
#[derive(Args)]
|
||||||
pub struct FilterCmdArgs {
|
pub struct FilterArgs {
|
||||||
/// Source index directory
|
/// Source index directory
|
||||||
pub source: PathBuf,
|
pub source: PathBuf,
|
||||||
|
|
||||||
@@ -28,18 +28,6 @@ pub struct FilterCmdArgs {
|
|||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub max_total_count: Option<u32>,
|
pub max_total_count: Option<u32>,
|
||||||
|
|
||||||
/// Minimum normalized entropy (complexity) to keep a k-mer — same metric
|
|
||||||
/// as `obikmer index`'s --theta, applied here to k-mers already committed
|
|
||||||
/// to the source index (reconstructed from unitigs.bin). K-mers scoring
|
|
||||||
/// below this are removed.
|
|
||||||
#[arg(long)]
|
|
||||||
pub min_complexity: Option<f64>,
|
|
||||||
|
|
||||||
/// Maximum sub-word size for the complexity computation (see `obikmer
|
|
||||||
/// index`'s --level-max). Only used when --min-complexity is set.
|
|
||||||
#[arg(long, default_value_t = 6)]
|
|
||||||
pub complexity_level_max: usize,
|
|
||||||
|
|
||||||
/// Output as presence/absence instead of counts
|
/// Output as presence/absence instead of counts
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub presence: bool,
|
pub presence: bool,
|
||||||
@@ -49,7 +37,7 @@ pub struct FilterCmdArgs {
|
|||||||
pub force: bool,
|
pub force: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn run(args: FilterCmdArgs) {
|
pub fn run(args: FilterArgs) {
|
||||||
let src = KmerIndex::open(&args.source).unwrap_or_else(|e| {
|
let src = KmerIndex::open(&args.source).unwrap_or_else(|e| {
|
||||||
eprintln!("error opening source index: {e}");
|
eprintln!("error opening source index: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
@@ -74,9 +62,6 @@ pub fn run(args: FilterCmdArgs) {
|
|||||||
if let Some(v) = args.max_total_count {
|
if let Some(v) = args.max_total_count {
|
||||||
filters.push(Box::new(MaxTotalCount { total: v }));
|
filters.push(Box::new(MaxTotalCount { total: v }));
|
||||||
}
|
}
|
||||||
if let Some(theta) = args.min_complexity {
|
|
||||||
filters.push(Box::new(MinComplexity { level_max: args.complexity_level_max, theta }));
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut rep = Reporter::new();
|
let mut rep = Reporter::new();
|
||||||
KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep)
|
KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep)
|
||||||
|
|||||||
@@ -151,14 +151,12 @@ pub struct FilterArgs {
|
|||||||
pub outgroup: Vec<String>,
|
pub outgroup: Vec<String>,
|
||||||
|
|
||||||
/// Minimum number of ingroup genomes containing the k-mer
|
/// Minimum number of ingroup genomes containing the k-mer
|
||||||
/// (negative: offset from group size, e.g. -1 = all but one)
|
#[arg(long)]
|
||||||
#[arg(long, allow_hyphen_values = true)]
|
pub min_count: Option<usize>,
|
||||||
pub min_count: Option<isize>,
|
|
||||||
|
|
||||||
/// Maximum number of ingroup genomes containing the k-mer
|
/// Maximum number of ingroup genomes containing the k-mer
|
||||||
/// (negative: offset from group size, e.g. -1 = all but one)
|
#[arg(long)]
|
||||||
#[arg(long, allow_hyphen_values = true)]
|
pub max_count: Option<usize>,
|
||||||
pub max_count: Option<isize>,
|
|
||||||
|
|
||||||
/// Minimum fraction of ingroup genomes containing the k-mer [0.0–1.0]
|
/// Minimum fraction of ingroup genomes containing the k-mer [0.0–1.0]
|
||||||
/// (default 1.0 when --ingroup is set, 0.0 otherwise)
|
/// (default 1.0 when --ingroup is set, 0.0 otherwise)
|
||||||
@@ -170,15 +168,13 @@ pub struct FilterArgs {
|
|||||||
pub max_frac: Option<f64>,
|
pub max_frac: Option<f64>,
|
||||||
|
|
||||||
/// Minimum number of outgroup genomes containing the k-mer
|
/// Minimum number of outgroup genomes containing the k-mer
|
||||||
/// (negative: offset from outgroup size, e.g. -1 = all but one)
|
#[arg(long)]
|
||||||
#[arg(long, allow_hyphen_values = true)]
|
pub min_outgroup_count: Option<usize>,
|
||||||
pub min_outgroup_count: Option<isize>,
|
|
||||||
|
|
||||||
/// Maximum number of outgroup genomes containing the k-mer
|
/// Maximum number of outgroup genomes containing the k-mer
|
||||||
/// (default 0 when --outgroup is set, no constraint otherwise;
|
/// (default 0 when --outgroup is set, no constraint otherwise)
|
||||||
/// negative: offset from outgroup size, e.g. -1 = all but one)
|
#[arg(long)]
|
||||||
#[arg(long, allow_hyphen_values = true)]
|
pub max_outgroup_count: Option<usize>,
|
||||||
pub max_outgroup_count: Option<isize>,
|
|
||||||
|
|
||||||
/// Minimum fraction of outgroup genomes containing the k-mer [0.0–1.0]
|
/// Minimum fraction of outgroup genomes containing the k-mer [0.0–1.0]
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
@@ -243,12 +239,12 @@ pub fn matching_genome_indices(pred_str: &str, genomes: &[GenomeInfo]) -> Result
|
|||||||
|
|
||||||
pub struct GroupFilterParams {
|
pub struct GroupFilterParams {
|
||||||
pub threshold: u32,
|
pub threshold: u32,
|
||||||
pub min_count: Option<isize>,
|
pub min_count: Option<usize>,
|
||||||
pub max_count: Option<isize>,
|
pub max_count: Option<usize>,
|
||||||
pub min_frac: Option<f64>,
|
pub min_frac: Option<f64>,
|
||||||
pub max_frac: Option<f64>,
|
pub max_frac: Option<f64>,
|
||||||
pub min_outgroup_count: Option<isize>,
|
pub min_outgroup_count: Option<usize>,
|
||||||
pub max_outgroup_count: Option<isize>,
|
pub max_outgroup_count: Option<usize>,
|
||||||
pub min_outgroup_frac: Option<f64>,
|
pub min_outgroup_frac: Option<f64>,
|
||||||
pub max_outgroup_frac: Option<f64>,
|
pub max_outgroup_frac: Option<f64>,
|
||||||
}
|
}
|
||||||
@@ -283,20 +279,12 @@ pub fn build_group_filter(
|
|||||||
let default_min_frac = if !ingroup_preds.is_empty() && !ingroup_quorum_explicit { 1.0 } else { 0.0 };
|
let default_min_frac = if !ingroup_preds.is_empty() && !ingroup_quorum_explicit { 1.0 } else { 0.0 };
|
||||||
let default_max_outgroup_count = if !outgroup_preds.is_empty() && !outgroup_quorum_explicit { 0 } else { out_size };
|
let default_max_outgroup_count = if !outgroup_preds.is_empty() && !outgroup_quorum_explicit { 0 } else { out_size };
|
||||||
|
|
||||||
// Resolve a signed count: negative means an offset from the group size
|
let min_count = p.min_count.unwrap_or(0);
|
||||||
// (e.g. -1 = all but one), floored at 1 so the negative form always keeps
|
let max_count = p.max_count.unwrap_or(in_size);
|
||||||
// constraining the group — even a singleton group, where n-1 would be 0
|
|
||||||
// and would otherwise drop the constraint entirely.
|
|
||||||
let resolve = |v: isize, size: usize| -> usize {
|
|
||||||
if v < 0 { (size as isize + v).max(1) as usize } else { v as usize }
|
|
||||||
};
|
|
||||||
|
|
||||||
let min_count = p.min_count.map(|v| resolve(v, in_size)).unwrap_or(0);
|
|
||||||
let max_count = p.max_count.map(|v| resolve(v, in_size)).unwrap_or(in_size);
|
|
||||||
let min_frac = p.min_frac.unwrap_or(default_min_frac);
|
let min_frac = p.min_frac.unwrap_or(default_min_frac);
|
||||||
let max_frac = p.max_frac.unwrap_or(1.0);
|
let max_frac = p.max_frac.unwrap_or(1.0);
|
||||||
let min_outgroup_count = p.min_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(0);
|
let min_outgroup_count = p.min_outgroup_count.unwrap_or(0);
|
||||||
let max_outgroup_count = p.max_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(default_max_outgroup_count);
|
let max_outgroup_count = p.max_outgroup_count.unwrap_or(default_max_outgroup_count);
|
||||||
let min_outgroup_frac = p.min_outgroup_frac.unwrap_or(0.0);
|
let min_outgroup_frac = p.min_outgroup_frac.unwrap_or(0.0);
|
||||||
let max_outgroup_frac = p.max_outgroup_frac.unwrap_or(1.0);
|
let max_outgroup_frac = p.max_outgroup_frac.unwrap_or(1.0);
|
||||||
|
|
||||||
|
|||||||
+177
-502
@@ -2,26 +2,21 @@ use std::collections::{HashMap, VecDeque};
|
|||||||
use std::io::{self, BufWriter, Write};
|
use std::io::{self, BufWriter, Write};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
|
||||||
use std::time::Instant;
|
|
||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikindex::KmerIndex;
|
use obikindex::KmerIndex;
|
||||||
use obikpartitionner::{KmerDesc, QueryHit, QueryStats};
|
|
||||||
use obikrope::Rope;
|
use obikrope::Rope;
|
||||||
use obikseq::CanonicalKmer;
|
use obikseq::RoutableSuperKmer;
|
||||||
use obilayeredmap::IndexMode;
|
use obilayeredmap::IndexMode;
|
||||||
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
|
||||||
use obiread::chunk::read_sequence_chunks_sized;
|
use obiread::chunk::read_sequence_chunks_sized;
|
||||||
use obiread::record::{SeqRecord, parse_chunk};
|
use obiread::record::{SeqRecord, parse_chunk};
|
||||||
use obiskbuilder::SuperKmerIter;
|
use obiskbuilder::SuperKmerIter;
|
||||||
use obisys::{Reporter, Stage, available_memory_bytes, spinner};
|
use obisys::available_memory_bytes;
|
||||||
use tracing::{debug, info};
|
use tracing::info;
|
||||||
|
|
||||||
// ── Pipeline data ─────────────────────────────────────────────────────────────
|
// ── Pipeline data ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
enum QueryData {
|
enum QueryData {
|
||||||
Path(Throttled<PathBuf>),
|
|
||||||
Chunk(Rope),
|
Chunk(Rope),
|
||||||
Output(Vec<u8>),
|
Output(Vec<u8>),
|
||||||
}
|
}
|
||||||
@@ -70,41 +65,30 @@ pub struct QueryArgs {
|
|||||||
#[arg(
|
#[arg(
|
||||||
short = 'T',
|
short = 'T',
|
||||||
long,
|
long,
|
||||||
default_value_t = obisys::effective_parallelism()
|
default_value_t = std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(1)
|
||||||
)]
|
)]
|
||||||
pub threads: usize,
|
pub threads: usize,
|
||||||
|
|
||||||
/// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
|
/// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub chunk_size: Option<usize>,
|
pub chunk_size: Option<usize>,
|
||||||
|
|
||||||
/// Maximum number of input files open simultaneously.
|
|
||||||
/// Defaults to threads/4 (minimum 1). Keep below the number of workers
|
|
||||||
/// to ensure CPU workers are always available for the transform stage.
|
|
||||||
#[arg(long)]
|
|
||||||
pub max_open_files: Option<usize>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl QueryArgs {
|
// ── SKDesc — one occurrence of a superkmer in the batch ───────────────────────
|
||||||
pub fn effective_max_open(&self) -> usize {
|
|
||||||
self.max_open_files
|
/// Describes one occurrence of a superkmer in the query batch.
|
||||||
.unwrap_or_else(|| (self.threads / 4).max(1))
|
pub struct SKDesc {
|
||||||
.max(1)
|
/// Index of the source sequence within the batch.
|
||||||
}
|
pub seq_idx: u32,
|
||||||
|
/// Kmer offset of the first kmer of this superkmer within its sequence.
|
||||||
|
pub kmer_offset: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── QueryBatch ────────────────────────────────────────────────────────────────
|
// ── QueryBatch ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
/// A batch of query sequences, with k-mers deduplicated directly (not just at
|
/// A batch of query sequences with their superkmers deduplicated.
|
||||||
/// the superkmer level) and pre-split by partition.
|
|
||||||
///
|
|
||||||
/// Superkmer *construction* (`SuperKmerIter`) is still required — it's the
|
|
||||||
/// mechanism that computes minimizers and partition routing — but the dedup
|
|
||||||
/// key is the canonical k-mer, not the superkmer: two different superkmers
|
|
||||||
/// that happen to share a k-mer (read overlaps, repeats, a SNP splitting an
|
|
||||||
/// otherwise-identical run) are deduplicated too, not just identical whole
|
|
||||||
/// superkmers. This also means each unique k-mer triggers at most one MPHF
|
|
||||||
/// lookup, not one per occurrence.
|
|
||||||
pub struct QueryBatch {
|
pub struct QueryBatch {
|
||||||
/// Sequence ids in batch order.
|
/// Sequence ids in batch order.
|
||||||
pub ids: Vec<String>,
|
pub ids: Vec<String>,
|
||||||
@@ -112,40 +96,30 @@ pub struct QueryBatch {
|
|||||||
pub seqs: Vec<Vec<u8>>,
|
pub seqs: Vec<Vec<u8>>,
|
||||||
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
||||||
pub n_kmers: Vec<u32>,
|
pub n_kmers: Vec<u32>,
|
||||||
/// Deduplicated k-mer occurrences, one map per partition.
|
/// Deduplicated superkmer map.
|
||||||
pub by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
|
pub map: HashMap<RoutableSuperKmer, Vec<SKDesc>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl QueryBatch {
|
impl QueryBatch {
|
||||||
/// Build a batch from a vec of parsed sequence records, deduplicating
|
/// Build a batch from a vec of parsed sequence records.
|
||||||
/// k-mers and routing them to partitions in the same pass.
|
pub fn from_records(records: Vec<SeqRecord>, k: usize, level_max: usize, theta: f64) -> Self {
|
||||||
pub fn from_records(
|
|
||||||
records: Vec<SeqRecord>,
|
|
||||||
k: usize,
|
|
||||||
level_max: usize,
|
|
||||||
theta: f64,
|
|
||||||
n_partitions: usize,
|
|
||||||
) -> Self {
|
|
||||||
let mut ids = Vec::with_capacity(records.len());
|
let mut ids = Vec::with_capacity(records.len());
|
||||||
let mut seqs = Vec::with_capacity(records.len());
|
let mut seqs = Vec::with_capacity(records.len());
|
||||||
let mut n_kmers = Vec::with_capacity(records.len());
|
let mut n_kmers = Vec::with_capacity(records.len());
|
||||||
let mask = (n_partitions as u64) - 1;
|
// Upper-bound estimate: at most one superkmer per k bases.
|
||||||
let mut by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
|
// Avoids repeated rehash on large chunks.
|
||||||
(0..n_partitions).map(|_| HashMap::new()).collect();
|
let cap = records.iter().map(|r| r.normalized.len()).sum::<usize>() / k.max(1);
|
||||||
|
let mut map: HashMap<RoutableSuperKmer, Vec<SKDesc>> = HashMap::with_capacity(cap);
|
||||||
|
|
||||||
for (seq_idx, record) in records.into_iter().enumerate() {
|
for (seq_idx, record) in records.into_iter().enumerate() {
|
||||||
let mut kmer_offset = 0u32;
|
let mut kmer_offset = 0u32;
|
||||||
|
|
||||||
for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) {
|
for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) {
|
||||||
let part_idx = (rsk.minimizer().seq_hash() & mask) as usize;
|
|
||||||
let map = &mut by_partition[part_idx];
|
|
||||||
for (j, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
|
|
||||||
map.entry(kmer).or_default().push(KmerDesc {
|
|
||||||
seq_idx: seq_idx as u32,
|
|
||||||
pos: kmer_offset + j as u32,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
let n = (rsk.seql() - k + 1) as u32;
|
let n = (rsk.seql() - k + 1) as u32;
|
||||||
|
map.entry(rsk).or_default().push(SKDesc {
|
||||||
|
seq_idx: seq_idx as u32,
|
||||||
|
kmer_offset,
|
||||||
|
});
|
||||||
kmer_offset += n;
|
kmer_offset += n;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -158,27 +132,37 @@ impl QueryBatch {
|
|||||||
ids,
|
ids,
|
||||||
seqs,
|
seqs,
|
||||||
n_kmers,
|
n_kmers,
|
||||||
by_partition,
|
map,
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── SmerIndex — sparse "was this k-mer found at all" bookkeeping ─────────────
|
/// Split the superkmer map by partition index.
|
||||||
|
pub fn split_by_partition(&self, n_partitions: usize) -> Vec<Vec<&RoutableSuperKmer>> {
|
||||||
/// Tracks, per (sequence, s-mer position), whether the k-mer was found in the
|
let mask = (n_partitions as u64) - 1;
|
||||||
/// index at all — independent of *which* genome(s) matched. Sized
|
let mut by_part: Vec<Vec<&RoutableSuperKmer>> = vec![Vec::new(); n_partitions];
|
||||||
/// `total_smers` (one `bool` per s-mer occurrence in the chunk), **not**
|
for rsk in self.map.keys() {
|
||||||
/// multiplied by `n_genomes`: this is the O(1)-per-position bookkeeping that
|
let part = (rsk.minimizer().seq_hash() & mask) as usize;
|
||||||
/// `kmer_missing` needs (the leftmost-s-mer-of-window membership test), kept
|
by_part[part].push(rsk);
|
||||||
/// dense because it's already cheap — the `n_genomes`-scaled data lives in
|
}
|
||||||
/// the sparse per-genome hit lists built alongside it (see `process_chunk`).
|
by_part
|
||||||
struct SmerIndex {
|
}
|
||||||
in_index: Vec<bool>, // total_smers
|
|
||||||
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl SmerIndex {
|
// ── KmerResults — allocation-free ragged result matrix ───────────────────────
|
||||||
fn new(n_kmers_per_seq: &[u32]) -> Self {
|
|
||||||
|
/// Flat storage for per-kmer query results across all sequences in a chunk.
|
||||||
|
///
|
||||||
|
/// Replaces `Vec<Vec<Option<Box<[u32]>>>>` — a single allocation for the whole
|
||||||
|
/// chunk instead of one `Box<[u32]>` per found k-mer.
|
||||||
|
struct KmerResults {
|
||||||
|
data: Vec<u32>, // total_kmers × n_genomes, row-major
|
||||||
|
in_index: Vec<bool>, // total_kmers — true if the kmer was found in the index
|
||||||
|
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = kmer range for sequence i
|
||||||
|
n_genomes: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl KmerResults {
|
||||||
|
fn new(n_kmers_per_seq: &[u32], n_genomes: usize) -> Self {
|
||||||
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
||||||
let mut total = 0usize;
|
let mut total = 0usize;
|
||||||
offsets.push(0);
|
offsets.push(0);
|
||||||
@@ -187,96 +171,34 @@ impl SmerIndex {
|
|||||||
offsets.push(total);
|
offsets.push(total);
|
||||||
}
|
}
|
||||||
Self {
|
Self {
|
||||||
|
data: vec![0u32; total * n_genomes],
|
||||||
in_index: vec![false; total],
|
in_index: vec![false; total],
|
||||||
offsets,
|
offsets,
|
||||||
|
n_genomes,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Mark the k-mer at (seq, kmer) as found in the index — independent of
|
fn n_kmers_for(&self, seq: usize) -> usize {
|
||||||
/// any particular genome's value. Called once per hit k-mer (stage 1 of
|
self.offsets[seq + 1] - self.offsets[seq]
|
||||||
/// `query_partition_with`), regardless of how the column-major fetch
|
}
|
||||||
/// (stage 2) later reports per-genome values.
|
|
||||||
fn mark_found(&mut self, seq: usize, kmer: usize) {
|
fn set(&mut self, seq: usize, kmer: usize, row: &[u32]) {
|
||||||
let abs = self.offsets[seq] + kmer;
|
let abs = self.offsets[seq] + kmer;
|
||||||
self.in_index[abs] = true;
|
self.in_index[abs] = true;
|
||||||
|
let base = abs * self.n_genomes;
|
||||||
|
self.data[base..base + self.n_genomes].copy_from_slice(row);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
||||||
self.in_index[self.offsets[seq] + kmer]
|
self.in_index[self.offsets[seq] + kmer]
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────────
|
/// Value for genome `g` at (seq, kmer); meaningful only when `is_in_index`.
|
||||||
|
#[inline]
|
||||||
/// One confirmed z-window: genome `g`'s window ending at k-mer `pos` (the
|
fn val(&self, seq: usize, kmer: usize, g: usize) -> u32 {
|
||||||
/// *leftmost* s-mer of the window, i.e. the k_user-mer's output position) is
|
self.data[(self.offsets[seq] + kmer) * self.n_genomes + g]
|
||||||
/// fully present and nonzero, with window-minimum `value`.
|
|
||||||
type ConfirmedHit = (u32, u32, u32); // (seq_idx, pos_out, value)
|
|
||||||
|
|
||||||
/// Reduce one genome's raw sparse s-mer hits — `(seq_idx, pos_smer, raw_value)`,
|
|
||||||
/// unsorted, exactly as delivered by `QueryHit::Value` — into confirmed
|
|
||||||
/// z-windows, without ever visiting a position that had no hit at all.
|
|
||||||
///
|
|
||||||
/// A z-window is confirmed only when all z s-mers in it are present *and*
|
|
||||||
/// nonzero for this genome (matching the dense sliding-window's semantics,
|
|
||||||
/// where "not in index" or a zero value both contribute 0 to the window
|
|
||||||
/// minimum) — which can only happen inside a maximal run of consecutive
|
|
||||||
/// `pos_smer` values for the same sequence. `hits` is sorted in place by
|
|
||||||
/// `(seq_idx, pos_smer)` to expose those runs; the monotone-deque
|
|
||||||
/// window-minimum then runs per run, on run-relative indices, identical in
|
|
||||||
/// spirit to the dense version's whole-sequence scan.
|
|
||||||
///
|
|
||||||
/// Returns the confirmed hits plus `(n_runs, total_run_len)` for logging —
|
|
||||||
/// a low average run length relative to `z` means most hits fail to form a
|
|
||||||
/// complete window.
|
|
||||||
fn sparse_findere_for_genome(
|
|
||||||
hits: &mut [(u32, u32, u32)],
|
|
||||||
z: usize,
|
|
||||||
presence: bool,
|
|
||||||
threshold: u32,
|
|
||||||
) -> (Vec<ConfirmedHit>, usize, usize) {
|
|
||||||
hits.sort_unstable_by_key(|&(seq, pos, _)| (seq, pos));
|
|
||||||
|
|
||||||
let mut confirmed = Vec::new();
|
|
||||||
let mut n_runs = 0usize;
|
|
||||||
let mut total_run_len = 0usize;
|
|
||||||
let mut dq: VecDeque<(usize, u32)> = VecDeque::new(); // (run-relative index, value)
|
|
||||||
|
|
||||||
let mut i = 0;
|
|
||||||
while i < hits.len() {
|
|
||||||
let seq = hits[i].0;
|
|
||||||
let mut j = i + 1;
|
|
||||||
while j < hits.len() && hits[j].0 == seq && hits[j].1 == hits[j - 1].1 + 1 {
|
|
||||||
j += 1;
|
|
||||||
}
|
}
|
||||||
let run = &hits[i..j];
|
|
||||||
n_runs += 1;
|
|
||||||
total_run_len += run.len();
|
|
||||||
|
|
||||||
dq.clear();
|
|
||||||
for (k, &(_, pos, val)) in run.iter().enumerate() {
|
|
||||||
while dq.back().map_or(false, |&(_, v)| v >= val) {
|
|
||||||
dq.pop_back();
|
|
||||||
}
|
|
||||||
dq.push_back((k, val));
|
|
||||||
while dq.front().map_or(false, |&(fk, _)| fk + z <= k) {
|
|
||||||
dq.pop_front();
|
|
||||||
}
|
|
||||||
if k + 1 >= z {
|
|
||||||
let win_min = dq.front().unwrap().1;
|
|
||||||
if win_min > 0 {
|
|
||||||
let pos_out = pos + 1 - z as u32;
|
|
||||||
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
|
||||||
confirmed.push((seq, pos_out, c));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
i = j;
|
|
||||||
}
|
|
||||||
|
|
||||||
(confirmed, n_runs, total_run_len)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Per-sequence accumulator ──────────────────────────────────────────────────
|
// ── Per-sequence accumulator ──────────────────────────────────────────────────
|
||||||
@@ -312,91 +234,40 @@ fn process_chunk(
|
|||||||
force_presence: bool,
|
force_presence: bool,
|
||||||
presence_threshold: u32,
|
presence_threshold: u32,
|
||||||
) -> Vec<u8> {
|
) -> Vec<u8> {
|
||||||
let chunk_start = Instant::now();
|
|
||||||
let chunk_bytes = rope.len();
|
|
||||||
|
|
||||||
let records = parse_chunk(&rope, k);
|
let records = parse_chunk(&rope, k);
|
||||||
if records.is_empty() {
|
if records.is_empty() {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
let batch = QueryBatch::from_records(records, k, 6, 0.7);
|
||||||
let n_seqs = batch.ids.len();
|
let n_seqs = batch.ids.len();
|
||||||
|
|
||||||
// Estimate QueryBatch::by_partition's actual memory footprint: the
|
// Flat result matrix — one allocation for the whole chunk.
|
||||||
// k-mer-level dedup map (roadmap point 5) — one HashMap<CanonicalKmer,
|
let mut results = KmerResults::new(&batch.n_kmers, n_genomes);
|
||||||
// Vec<KmerDesc>> per partition, sized by *unique* k-mers, not shrunk by
|
|
||||||
// dedup. On real workloads with a low intra-chunk duplication rate this
|
|
||||||
// can dwarf every other per-chunk structure, including the sparse
|
|
||||||
// Findere ones logged further down — unlike those, chunk_bytes's formula
|
|
||||||
// (run()) does not account for this at all today. Measured by allocated
|
|
||||||
// capacity, not logical length, to reflect real memory pressure
|
|
||||||
// (HashMap/Vec growth slack) — `by_partition` is alive for the entire
|
|
||||||
// process_chunk call (never drained, only iterated by reference), so
|
|
||||||
// this is its footprint for the whole chunk lifetime, not a transient.
|
|
||||||
let hashmap_slot_bytes = (std::mem::size_of::<CanonicalKmer>()
|
|
||||||
+ std::mem::size_of::<Vec<KmerDesc>>()
|
|
||||||
+ 1) as u64; // +1 ≈ hashbrown control byte per slot
|
|
||||||
let by_partition_map_bytes: u64 = batch
|
|
||||||
.by_partition
|
|
||||||
.iter()
|
|
||||||
.map(|m| m.capacity() as u64 * hashmap_slot_bytes)
|
|
||||||
.sum();
|
|
||||||
let by_partition_desc_bytes: u64 = batch
|
|
||||||
.by_partition
|
|
||||||
.iter()
|
|
||||||
.flat_map(|m| m.values())
|
|
||||||
.map(|v| v.capacity() as u64 * std::mem::size_of::<KmerDesc>() as u64)
|
|
||||||
.sum();
|
|
||||||
let by_partition_bytes = by_partition_map_bytes + by_partition_desc_bytes;
|
|
||||||
|
|
||||||
debug!(
|
let by_part = batch.split_by_partition(n_partitions);
|
||||||
n_unique_kmers_total = batch.by_partition.iter().map(|m| m.len() as u64).sum::<u64>(),
|
|
||||||
by_partition_map_bytes,
|
|
||||||
by_partition_desc_bytes,
|
|
||||||
by_partition_bytes,
|
|
||||||
chunk_bytes,
|
|
||||||
"by_partition memory retained"
|
|
||||||
);
|
|
||||||
|
|
||||||
// Sparse bookkeeping for the whole chunk:
|
for (part_idx, part_sks) in by_part.iter().enumerate() {
|
||||||
// - smer_index: O(total_smers) — is this s-mer in the index at all.
|
if part_sks.is_empty() {
|
||||||
// - by_genome[g]: raw (seq_idx, pos_smer, value) hits for genome g, only
|
|
||||||
// ever containing nonzero entries (query_partition_with never emits a
|
|
||||||
// QueryHit::Value for a zero value) — empty for every genome this chunk
|
|
||||||
// never matched, which is the common case for unrelated queries.
|
|
||||||
let mut smer_index = SmerIndex::new(&batch.n_kmers);
|
|
||||||
let mut by_genome: Vec<Vec<(u32, u32, u32)>> = (0..n_genomes).map(|_| Vec::new()).collect();
|
|
||||||
|
|
||||||
// Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
|
|
||||||
// before dedup) vs. unique k-mers actually queried (query_stats) — the
|
|
||||||
// entire justification for k-mer-level dereplication (see query.md,
|
|
||||||
// Future work point 5). If this ratio stays close to 1.0 on real data,
|
|
||||||
// dereplication isn't paying for itself and that should show up here.
|
|
||||||
let n_occurrences: u64 = batch.n_kmers.iter().map(|&n| n as u64).sum();
|
|
||||||
let mut query_stats = QueryStats::default();
|
|
||||||
|
|
||||||
for (part_idx, kmers) in batch.by_partition.iter().enumerate() {
|
|
||||||
if kmers.is_empty() {
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let stats = idx.partition()
|
idx.partition()
|
||||||
.query_partition_with(
|
.query_partition_with(
|
||||||
part_idx,
|
part_idx,
|
||||||
kmers,
|
part_sks,
|
||||||
|
k,
|
||||||
n_genomes,
|
n_genomes,
|
||||||
with_counts,
|
with_counts,
|
||||||
|event| match event {
|
|sk_idx, kmer_idx, row| {
|
||||||
QueryHit::Found(descs) => {
|
let rsk = part_sks[sk_idx];
|
||||||
|
let descs = batch.map.get(rsk).expect("rsk must be in map");
|
||||||
for desc in descs {
|
for desc in descs {
|
||||||
smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
results.set(
|
||||||
}
|
desc.seq_idx as usize,
|
||||||
}
|
desc.kmer_offset as usize + kmer_idx,
|
||||||
QueryHit::Value(descs, g, v) => {
|
row,
|
||||||
for desc in descs {
|
);
|
||||||
by_genome[g].push((desc.seq_idx, desc.pos, v));
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
@@ -404,129 +275,96 @@ fn process_chunk(
|
|||||||
eprintln!("query error on partition {part_idx}: {e}");
|
eprintln!("query error on partition {part_idx}: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
});
|
});
|
||||||
query_stats += stats;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
debug!(
|
// Sliding window minimum — one reusable buffer and one deque per batch.
|
||||||
n_occurrences,
|
|
||||||
n_unique_kmers = query_stats.n_unique_kmers,
|
|
||||||
n_mphf_calls = query_stats.n_mphf_calls,
|
|
||||||
n_hits = query_stats.n_hits,
|
|
||||||
n_columns_scanned = query_stats.n_columns_scanned,
|
|
||||||
n_col_get_calls = query_stats.n_col_get_calls,
|
|
||||||
"k-mer dedup + column-major fetch"
|
|
||||||
);
|
|
||||||
|
|
||||||
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────
|
|
||||||
//
|
//
|
||||||
// Confirmed z-windows, per genome, replace the dense win_min matrix:
|
// win_min[pos * n_genomes + g] = min count across the z-window [pos, pos+z)
|
||||||
// total retained memory is O(actual hits), not O(total_smers × n_genomes)
|
// for genome g, where "not in index" counts as 0.
|
||||||
// — the whole point of this pass. See sparse_findere_for_genome's doc for
|
//
|
||||||
// why run detection is equivalent to the dense scan's semantics.
|
// win_min > 0 ↔ all z consecutive kmers are in the index with count > 0
|
||||||
let presence = force_presence || !with_counts;
|
// ↔ Findere confirmation (for z=1 this degenerates to the
|
||||||
let threshold = presence_threshold;
|
// simple case with no overhead).
|
||||||
let z = effective_z;
|
//
|
||||||
|
// Works uniformly for count matrices and presence/absence (0/1) matrices.
|
||||||
|
let max_n_kmers = batch.n_kmers.iter().map(|&n| n as usize).max().unwrap_or(0);
|
||||||
|
let mut win_min = vec![0u32; max_n_kmers * n_genomes];
|
||||||
|
|
||||||
|
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
||||||
|
|
||||||
let n_kmers_out: Vec<usize> = batch
|
let n_kmers_out: Vec<usize> = batch
|
||||||
.n_kmers
|
.n_kmers
|
||||||
.iter()
|
.iter()
|
||||||
.map(|&n| {
|
.map(|&n| {
|
||||||
let n = n as usize;
|
let n = n as usize;
|
||||||
if n >= z { n - z + 1 } else { 0 }
|
if n >= effective_z { n - effective_z + 1 } else { 0 }
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
let mut out_offsets = Vec::with_capacity(n_seqs + 1);
|
|
||||||
{
|
|
||||||
let mut total = 0usize;
|
|
||||||
out_offsets.push(0);
|
|
||||||
for &n in &n_kmers_out {
|
|
||||||
total += n;
|
|
||||||
out_offsets.push(total);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let total_out = *out_offsets.last().unwrap_or(&0);
|
|
||||||
|
|
||||||
let n_dense_would_be = n_occurrences as u64 * n_genomes as u64;
|
|
||||||
let mut n_sparse_entries = 0u64;
|
|
||||||
let mut n_runs_total = 0usize;
|
|
||||||
let mut run_len_total = 0usize;
|
|
||||||
|
|
||||||
let mut confirmed_by_genome: Vec<Vec<ConfirmedHit>> = Vec::with_capacity(n_genomes);
|
|
||||||
for hits in &mut by_genome {
|
|
||||||
n_sparse_entries += hits.len() as u64;
|
|
||||||
let (confirmed, n_runs, run_len) = sparse_findere_for_genome(hits, z, presence, threshold);
|
|
||||||
n_runs_total += n_runs;
|
|
||||||
run_len_total += run_len;
|
|
||||||
confirmed_by_genome.push(confirmed);
|
|
||||||
}
|
|
||||||
|
|
||||||
debug!(
|
|
||||||
n_dense_would_be,
|
|
||||||
n_sparse_entries,
|
|
||||||
n_runs = n_runs_total,
|
|
||||||
avg_run_len = if n_runs_total > 0 { run_len_total as f64 / n_runs_total as f64 } else { 0.0 },
|
|
||||||
z,
|
|
||||||
"sparse Findere"
|
|
||||||
);
|
|
||||||
|
|
||||||
// Actual bytes retained by the sparse hit structures (by_genome +
|
|
||||||
// confirmed_by_genome, both alive simultaneously at this point — see the
|
|
||||||
// chunk-size formula's comment in `run()`), by allocated capacity rather
|
|
||||||
// than logical length so this reflects real memory pressure including
|
|
||||||
// Vec growth slack. `empirical_multiplier` is directly comparable to
|
|
||||||
// BYTES_PER_KMER_PER_GENOME (`run()`) — the ratio a cluster run's logs
|
|
||||||
// need to judge whether that constant is over- or under-conservative for
|
|
||||||
// real data, instead of guessing.
|
|
||||||
const HIT_ENTRY_BYTES: u64 = std::mem::size_of::<(u32, u32, u32)>() as u64;
|
|
||||||
let by_genome_bytes: u64 = by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
|
|
||||||
let confirmed_bytes: u64 = confirmed_by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
|
|
||||||
let retained_bytes = by_genome_bytes + confirmed_bytes;
|
|
||||||
|
|
||||||
debug!(
|
|
||||||
by_genome_bytes,
|
|
||||||
confirmed_bytes,
|
|
||||||
retained_bytes,
|
|
||||||
chunk_bytes,
|
|
||||||
empirical_multiplier = retained_bytes as f64 / chunk_bytes.max(1) as f64,
|
|
||||||
"sparse memory retained"
|
|
||||||
);
|
|
||||||
|
|
||||||
// ── Accumulate: genome totals (per genome, from confirmed hits) ──────────
|
|
||||||
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
|
||||||
let mut confirmed_any = vec![false; total_out];
|
|
||||||
|
|
||||||
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
|
||||||
for &(seq_idx, pos_out, c) in hits {
|
|
||||||
let abs_out = out_offsets[seq_idx as usize] + pos_out as usize;
|
|
||||||
confirmed_any[abs_out] = true;
|
|
||||||
accs[seq_idx as usize].genome_totals[g] += c;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Accumulate: kmer_count / kmer_missing (per position, genome-independent) ─
|
|
||||||
for seq_idx in 0..n_seqs {
|
|
||||||
let out_n = n_kmers_out[seq_idx];
|
|
||||||
let acc = &mut accs[seq_idx];
|
|
||||||
for pos in 0..out_n {
|
|
||||||
let abs_out = out_offsets[seq_idx] + pos;
|
|
||||||
if confirmed_any[abs_out] {
|
|
||||||
acc.kmer_count += 1;
|
|
||||||
} else if !smer_index.is_in_index(seq_idx, pos) {
|
|
||||||
acc.kmer_missing += 1;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Coverage (--detail): densify only when actually requested ────────────
|
|
||||||
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
||||||
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
||||||
} else {
|
} else {
|
||||||
Vec::new()
|
Vec::new()
|
||||||
};
|
};
|
||||||
if detail {
|
|
||||||
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
let presence = force_presence || !with_counts;
|
||||||
for &(seq_idx, pos_out, c) in hits {
|
let threshold = presence_threshold;
|
||||||
cov[seq_idx as usize][g][pos_out as usize] += c;
|
let z = effective_z;
|
||||||
|
|
||||||
|
// Deque reused across all (seq, genome) pairs.
|
||||||
|
let mut dq: VecDeque<(usize, u32)> = VecDeque::with_capacity(z + 1);
|
||||||
|
|
||||||
|
for seq_idx in 0..n_seqs {
|
||||||
|
let n = results.n_kmers_for(seq_idx);
|
||||||
|
let out_n = n_kmers_out[seq_idx];
|
||||||
|
if out_n == 0 { continue; }
|
||||||
|
|
||||||
|
let mins = &mut win_min[..out_n * n_genomes];
|
||||||
|
mins.fill(0);
|
||||||
|
|
||||||
|
// ── Per-genome sliding window minimum ─────────────────────────────────
|
||||||
|
for g in 0..n_genomes {
|
||||||
|
dq.clear();
|
||||||
|
for i in 0..n {
|
||||||
|
let v_i = if results.is_in_index(seq_idx, i) {
|
||||||
|
results.val(seq_idx, i, g)
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
// Evict elements that have left the window.
|
||||||
|
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
||||||
|
dq.pop_front();
|
||||||
|
}
|
||||||
|
// Maintain monotone non-decreasing back→front for minimum at front.
|
||||||
|
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
||||||
|
dq.pop_back();
|
||||||
|
}
|
||||||
|
dq.push_back((i, v_i));
|
||||||
|
// Window [pos, pos+z) is complete when i = pos + z - 1.
|
||||||
|
if i + 1 >= z {
|
||||||
|
let pos = i + 1 - z;
|
||||||
|
mins[pos * n_genomes + g] = dq.front().unwrap().1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Accumulate ────────────────────────────────────────────────────────
|
||||||
|
let acc = &mut accs[seq_idx];
|
||||||
|
for pos in 0..out_n {
|
||||||
|
let any = (0..n_genomes).any(|g| mins[pos * n_genomes + g] > 0);
|
||||||
|
if !any {
|
||||||
|
if !results.is_in_index(seq_idx, pos) {
|
||||||
|
acc.kmer_missing += 1;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
acc.kmer_count += 1;
|
||||||
|
for g in 0..n_genomes {
|
||||||
|
let v = mins[pos * n_genomes + g];
|
||||||
|
if v == 0 { continue; }
|
||||||
|
let c = if presence { u32::from(v >= threshold) } else { v };
|
||||||
|
acc.genome_totals[g] += c;
|
||||||
|
if detail { cov[seq_idx][g][pos] += c; }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -546,42 +384,9 @@ fn process_chunk(
|
|||||||
&cov,
|
&cov,
|
||||||
&mut buf,
|
&mut buf,
|
||||||
);
|
);
|
||||||
|
|
||||||
debug!(
|
|
||||||
chunk_bytes,
|
|
||||||
n_seqs,
|
|
||||||
n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::<u64>(),
|
|
||||||
wall_ms = chunk_start.elapsed().as_millis() as u64,
|
|
||||||
"process_chunk"
|
|
||||||
);
|
|
||||||
|
|
||||||
buf
|
buf
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ──
|
|
||||||
|
|
||||||
/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`,
|
|
||||||
/// so the guard (and the throttle slot it holds) is only released once the
|
|
||||||
/// file has been fully read — never earlier, never held past that point.
|
|
||||||
struct GuardedChunkIter {
|
|
||||||
inner: Box<dyn Iterator<Item = Rope> + Send>,
|
|
||||||
_guard: ThrottleGuard,
|
|
||||||
files_open: Arc<AtomicU32>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Iterator for GuardedChunkIter {
|
|
||||||
type Item = Rope;
|
|
||||||
fn next(&mut self) -> Option<Rope> {
|
|
||||||
self.inner.next()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for GuardedChunkIter {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
self.files_open.fetch_sub(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Entry point ───────────────────────────────────────────────────────────────
|
// ── Entry point ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
pub fn run(args: QueryArgs) {
|
pub fn run(args: QueryArgs) {
|
||||||
@@ -597,67 +402,17 @@ pub fn run(args: QueryArgs) {
|
|||||||
let n_workers = args.threads.max(1);
|
let n_workers = args.threads.max(1);
|
||||||
|
|
||||||
// Chunk size: each chunk stays in memory for its entire processing lifetime.
|
// Chunk size: each chunk stays in memory for its entire processing lifetime.
|
||||||
//
|
// Overhead per raw byte is ~8× (Rope + parsed records + superkmers + results).
|
||||||
// Per-chunk memory is no longer a dense n_genomes-wide buffer (removed in
|
// We target ≤ 50 % of available RAM across all concurrent workers.
|
||||||
// the sparse Findere rework, see process_chunk) — it now scales with
|
|
||||||
// *actual hit count*, not with total_kmers_in_chunk × n_genomes
|
|
||||||
// unconditionally. BYTES_PER_KMER_PER_GENOME below is therefore a
|
|
||||||
// pathological-case bound, not a typical-case estimate: it protects
|
|
||||||
// against a fully-dense hit pattern (every k-mer of the query matching
|
|
||||||
// every genome — a degenerate case, e.g. low-complexity input theta-
|
|
||||||
// filtering should mostly reject, or an index of near-duplicate genomes),
|
|
||||||
// where by_genome and confirmed_by_genome (process_chunk) both end up
|
|
||||||
// holding one (seq_idx, pos, value) entry — 3 × u32 = 12 bytes, vs. 4
|
|
||||||
// bytes for the old dense encoding, where position was implicit in the
|
|
||||||
// array index — per (k-mer, genome) pair, and *coexist simultaneously*
|
|
||||||
// (by_genome isn't freed before confirmed_by_genome is built), for a
|
|
||||||
// worst case of ~24 bytes/pair before Vec growth slack. `cov` remains
|
|
||||||
// fully dense when --detail is set (unaffected by the sparse rework),
|
|
||||||
// still roughly doubling the n_genomes-scaled cost.
|
|
||||||
//
|
|
||||||
// For realistic, sparse hit patterns actual memory is far below this
|
|
||||||
// bound — see the "sparse memory retained" debug log in process_chunk,
|
|
||||||
// which reports the empirical bytes-per-raw-byte multiplier actually
|
|
||||||
// observed per chunk, directly comparable to BYTES_PER_KMER_PER_GENOME
|
|
||||||
// below. Tightening this constant for typical-case throughput (at the
|
|
||||||
// cost of pathological-case safety margin) is a deliberate tuning
|
|
||||||
// decision to make from that data, not something to guess at here.
|
|
||||||
//
|
|
||||||
// BASE_OVERHEAD approximates what scales with chunk_bytes alone,
|
|
||||||
// independent of n_genomes: the Rope itself, parsed SeqRecord sequence +
|
|
||||||
// normalised bytes, the superkmer dedup map, and the JSON output buffer.
|
|
||||||
// Like the n_genomes-scaled term, this is an estimate — validate against
|
|
||||||
// actual peak RSS (Stage::stop's `rss` in the summary table) on real
|
|
||||||
// workloads rather than trusting it blindly.
|
|
||||||
//
|
|
||||||
// We target ≤ 50 % of available RAM across all concurrent workers
|
|
||||||
// (SAFETY_FACTOR).
|
|
||||||
const BASE_OVERHEAD: u64 = 4;
|
|
||||||
const BYTES_PER_KMER_PER_GENOME: u64 = 8; // pathological-case bound — see comment above
|
|
||||||
const SAFETY_FACTOR: u64 = 2;
|
|
||||||
|
|
||||||
let detail_factor: u64 = if args.detail { 2 } else { 1 };
|
|
||||||
let overhead_multiplier =
|
|
||||||
BASE_OVERHEAD + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * detail_factor;
|
|
||||||
|
|
||||||
let chunk_bytes = args
|
let chunk_bytes = args
|
||||||
.chunk_size
|
.chunk_size
|
||||||
.map(|mb| mb * 1024 * 1024)
|
.map(|mb| mb * 1024 * 1024)
|
||||||
.unwrap_or_else(|| {
|
.unwrap_or_else(|| {
|
||||||
let avail = available_memory_bytes();
|
let avail = available_memory_bytes();
|
||||||
let computed = avail / (n_workers as u64 * overhead_multiplier * SAFETY_FACTOR);
|
let computed = avail / (n_workers as u64 * 16);
|
||||||
computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize
|
computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize
|
||||||
});
|
});
|
||||||
|
|
||||||
debug!(
|
|
||||||
chunk_bytes,
|
|
||||||
n_genomes,
|
|
||||||
detail = args.detail,
|
|
||||||
overhead_multiplier,
|
|
||||||
estimated_peak_chunk_bytes = chunk_bytes as u64 * overhead_multiplier,
|
|
||||||
"chunk-size formula resolved"
|
|
||||||
);
|
|
||||||
|
|
||||||
let effective_z: usize = args
|
let effective_z: usize = args
|
||||||
.findere_z
|
.findere_z
|
||||||
.unwrap_or_else(|| match idx.meta().config.evidence {
|
.unwrap_or_else(|| match idx.meta().config.evidence {
|
||||||
@@ -680,122 +435,48 @@ pub fn run(args: QueryArgs) {
|
|||||||
let force_presence = args.force_presence;
|
let force_presence = args.force_presence;
|
||||||
let presence_threshold = args.presence_threshold;
|
let presence_threshold = args.presence_threshold;
|
||||||
|
|
||||||
// Throttled iterator over input file paths: at most `effective_max_open()`
|
// Flat iterator over all Rope chunks from all input files.
|
||||||
// files are open at once. Opening + decompressing + chunking each file is
|
// I/O runs in the source thread; chunk processing is parallelised by the pipe.
|
||||||
// now a Flat pipeline stage, executed across the `n_workers` pool — not
|
info!("query: chunk_size={}MiB", chunk_bytes / (1024 * 1024));
|
||||||
// serialised in the pipe's dedicated source thread (see steps::scatter /
|
|
||||||
// cmd::superkmer for the same pattern applied to indexing).
|
|
||||||
info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open());
|
|
||||||
|
|
||||||
let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
|
let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
|
||||||
let path_source = throttle(paths.into_iter(), args.effective_max_open());
|
let all_chunks = paths.into_iter().flat_map(move |path| {
|
||||||
|
|
||||||
// Instrumentation: total bytes processed (for the EMA throughput readout),
|
|
||||||
// number of files currently open/being chunked, and number of chunks
|
|
||||||
// currently being processed by a worker — all read from the spinner loop
|
|
||||||
// below, updated from inside the pipe closures.
|
|
||||||
let total_bytes = Arc::new(AtomicU64::new(0));
|
|
||||||
let files_open = Arc::new(AtomicU32::new(0));
|
|
||||||
let chunks_active = Arc::new(AtomicU32::new(0));
|
|
||||||
|
|
||||||
let pipe = obipipeline::make_pipe! {
|
|
||||||
QueryData : Throttled<PathBuf> => Vec<u8>,
|
|
||||||
|| {
|
|
||||||
let files_open = Arc::clone(&files_open);
|
|
||||||
move |pw: Throttled<PathBuf>| -> GuardedChunkIter {
|
|
||||||
let path = pw.item;
|
|
||||||
let guard = pw.guard;
|
|
||||||
let path_str = path.to_str().unwrap_or("").to_owned();
|
let path_str = path.to_str().unwrap_or("").to_owned();
|
||||||
files_open.fetch_add(1, Ordering::Relaxed);
|
match read_sequence_chunks_sized(&path_str, chunk_bytes) {
|
||||||
let open_start = Instant::now();
|
Ok(iter) => Box::new(iter.filter_map(|r| match r {
|
||||||
// Hard-exit on file-open failure (mirrors the previous behaviour):
|
|
||||||
// propagating this as a pipeline Err would hit a known scheduler
|
|
||||||
// hang on early stage errors (obipipeline::scheduler::WorkerPool::run
|
|
||||||
// breaks its main loop without unblocking the still-running source
|
|
||||||
// thread, so the final `h.join()` never returns) — worth fixing in
|
|
||||||
// obipipeline itself, but out of scope here; sidestepping it like the
|
|
||||||
// original code already did is the safe choice for this change.
|
|
||||||
let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening {path_str}: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
debug!(
|
|
||||||
path = %path_str,
|
|
||||||
open_ms = open_start.elapsed().as_millis() as u64,
|
|
||||||
"opened query input file"
|
|
||||||
);
|
|
||||||
let err_path = path_str.clone();
|
|
||||||
GuardedChunkIter {
|
|
||||||
inner: Box::new(iter.filter_map(move |r| match r {
|
|
||||||
Ok(rope) => Some(rope),
|
Ok(rope) => Some(rope),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
eprintln!("read error: {err_path}: {e}");
|
eprintln!("read error: {e}");
|
||||||
None
|
None
|
||||||
}
|
}
|
||||||
})),
|
})) as Box<dyn Iterator<Item = Rope> + Send>,
|
||||||
_guard: guard,
|
Err(e) => {
|
||||||
files_open: Arc::clone(&files_open),
|
eprintln!("error opening {path_str}: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} : Path => Chunk,
|
});
|
||||||
|
|
||||||
|
let pipe = obipipeline::make_pipe! {
|
||||||
|
QueryData : Rope => Vec<u8>,
|
||||||
| {
|
| {
|
||||||
let idx = Arc::clone(&idx);
|
let idx = Arc::clone(&idx);
|
||||||
let total_bytes = Arc::clone(&total_bytes);
|
|
||||||
let chunks_active = Arc::clone(&chunks_active);
|
|
||||||
move |rope: Rope| {
|
move |rope: Rope| {
|
||||||
chunks_active.fetch_add(1, Ordering::Relaxed);
|
process_chunk(
|
||||||
let bytes = rope.len() as u64;
|
|
||||||
let out = process_chunk(
|
|
||||||
&idx, rope, k, n_genomes, n_partitions, with_counts,
|
&idx, rope, k, n_genomes, n_partitions, with_counts,
|
||||||
effective_z, detail, count_missing, force_presence, presence_threshold,
|
effective_z, detail, count_missing, force_presence, presence_threshold,
|
||||||
);
|
)
|
||||||
total_bytes.fetch_add(bytes, Ordering::Relaxed);
|
|
||||||
chunks_active.fetch_sub(1, Ordering::Relaxed);
|
|
||||||
out
|
|
||||||
}
|
}
|
||||||
} : Chunk => Output,
|
} : Chunk => Output,
|
||||||
};
|
};
|
||||||
|
|
||||||
let t = Stage::start("query");
|
|
||||||
let pb = spinner("query");
|
|
||||||
|
|
||||||
let mut ema_rate: f64 = 0.0;
|
|
||||||
let mut last_t = Instant::now();
|
|
||||||
let mut last_bytes: u64 = 0;
|
|
||||||
const ALPHA: f64 = 0.15;
|
|
||||||
|
|
||||||
let mut out = BufWriter::new(io::stdout());
|
let mut out = BufWriter::new(io::stdout());
|
||||||
for block in pipe.apply(path_source, n_workers, 2) {
|
for block in pipe.apply(all_chunks, n_workers, 2) {
|
||||||
if !block.is_empty() {
|
if !block.is_empty() {
|
||||||
out.write_all(&block).expect("write error");
|
out.write_all(&block).expect("write error");
|
||||||
}
|
}
|
||||||
|
|
||||||
let now = Instant::now();
|
|
||||||
let dt = now.duration_since(last_t).as_secs_f64();
|
|
||||||
if dt > 0.1 {
|
|
||||||
let total = total_bytes.load(Ordering::Relaxed);
|
|
||||||
let instant = (total - last_bytes) as f64 / dt;
|
|
||||||
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
|
|
||||||
last_t = now;
|
|
||||||
last_bytes = total;
|
|
||||||
let bp = total as f64;
|
|
||||||
let (count_str, rate_str) = if bp >= 1e9 {
|
|
||||||
(format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6))
|
|
||||||
} else {
|
|
||||||
(format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
|
|
||||||
};
|
|
||||||
let active = chunks_active.load(Ordering::Relaxed);
|
|
||||||
let open = files_open.load(Ordering::Relaxed);
|
|
||||||
pb.set_message(format!("{count_str} {rate_str} [files open: {open}, chunks in flight: {active}]"));
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
out.flush().expect("flush error");
|
out.flush().expect("flush error");
|
||||||
|
|
||||||
pb.finish_and_clear();
|
|
||||||
|
|
||||||
let mut rep = Reporter::new();
|
|
||||||
rep.push(t.stop());
|
|
||||||
rep.print();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Output ────────────────────────────────────────────────────────────────────
|
// ── Output ────────────────────────────────────────────────────────────────────
|
||||||
@@ -820,10 +501,8 @@ fn emit_batch(
|
|||||||
|
|
||||||
let mut match_map = serde_json::Map::new();
|
let mut match_map = serde_json::Map::new();
|
||||||
for (g, genome) in meta.genomes.iter().enumerate() {
|
for (g, genome) in meta.genomes.iter().enumerate() {
|
||||||
if acc.genome_totals[g] != 0 {
|
|
||||||
match_map.insert(genome.label.clone(), acc.genome_totals[g].into());
|
match_map.insert(genome.label.clone(), acc.genome_totals[g].into());
|
||||||
}
|
}
|
||||||
}
|
|
||||||
ann.insert("kmer_strict_matches".into(), match_map.into());
|
ann.insert("kmer_strict_matches".into(), match_map.into());
|
||||||
|
|
||||||
if detail && !cov.is_empty() {
|
if detail && !cov.is_empty() {
|
||||||
@@ -845,7 +524,3 @@ fn emit_batch(
|
|||||||
let _ = out.write_all(b"\n");
|
let _ = out.write_all(b"\n");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
#[path = "tests/query.rs"]
|
|
||||||
mod tests;
|
|
||||||
|
|||||||
@@ -1,228 +0,0 @@
|
|||||||
use super::*;
|
|
||||||
|
|
||||||
const K: usize = 11;
|
|
||||||
const M: usize = 5;
|
|
||||||
|
|
||||||
/// Build a `QueryBatch` from raw FASTA text, going through the same
|
|
||||||
/// `Rope` + `parse_chunk` path `process_chunk` uses — avoids hand-building a
|
|
||||||
/// `normalized` `Rope`, which is an implementation detail of `obiread`.
|
|
||||||
///
|
|
||||||
/// `obikseq`'s global K/M params are thread-local under `test-utils` (see
|
|
||||||
/// `obikseq::params`), so setting them here is per-test-thread and does not
|
|
||||||
/// need coordination with other tests.
|
|
||||||
fn batch_from_fasta(fasta: &str, k: usize, n_partitions: usize) -> QueryBatch {
|
|
||||||
obikseq::set_k(k);
|
|
||||||
obikseq::set_m(M);
|
|
||||||
let mut rope = Rope::new(Some("text/fasta"));
|
|
||||||
rope.push(fasta.as_bytes().to_vec());
|
|
||||||
let records = parse_chunk(&rope, k);
|
|
||||||
QueryBatch::from_records(records, k, 6, 0.7, n_partitions)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn total_occurrences(batch: &QueryBatch) -> u64 {
|
|
||||||
batch.n_kmers.iter().map(|&n| n as u64).sum()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn total_unique_kmers(batch: &QueryBatch) -> u64 {
|
|
||||||
batch.by_partition.iter().map(|m| m.len() as u64).sum()
|
|
||||||
}
|
|
||||||
|
|
||||||
// A 60 bp sequence, arbitrary but fixed — no attempt is made to prove it is
|
|
||||||
// free of internal k=11 repeats; the tests below only rely on inequalities
|
|
||||||
// that hold regardless (see each test's comment).
|
|
||||||
const SEQ: &str = "CATTAGCGTACCTGATCAGGTTACAGCTTAGGCATCCAGTTGACCATGACTGGACTTAGC";
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn single_sequence_yields_plausible_kmer_counts() {
|
|
||||||
// A single record can still contain internal repeats (SEQ isn't
|
|
||||||
// guaranteed repeat-free at k=11) — this only checks the batch is
|
|
||||||
// internally consistent, not a specific dedup ratio. The cross-record
|
|
||||||
// tests below make the actual, unconditional dedup claims.
|
|
||||||
let fasta = format!(">r1\n{SEQ}\n");
|
|
||||||
let batch = batch_from_fasta(&fasta, K, 1);
|
|
||||||
|
|
||||||
assert_eq!(batch.ids, vec!["r1".to_string()]);
|
|
||||||
let occurrences = total_occurrences(&batch);
|
|
||||||
let unique = total_unique_kmers(&batch);
|
|
||||||
assert!(occurrences > 0, "sequence should yield at least one k-mer");
|
|
||||||
assert!(unique > 0 && unique <= occurrences);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn duplicated_sequence_across_records_deduplicates() {
|
|
||||||
// Two records with byte-identical sequences: every k-mer in record 1
|
|
||||||
// exactly duplicates one in record 0, so unique kmers <= n_kmers[0],
|
|
||||||
// strictly less than the summed occurrences (2 * n_kmers[0]) as long as
|
|
||||||
// the sequence yields at least one k-mer. This holds regardless of
|
|
||||||
// whether SEQ has internal repeats.
|
|
||||||
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
|
||||||
let batch = batch_from_fasta(&fasta, K, 1);
|
|
||||||
|
|
||||||
assert_eq!(batch.ids.len(), 2);
|
|
||||||
let occurrences = total_occurrences(&batch);
|
|
||||||
let unique = total_unique_kmers(&batch);
|
|
||||||
|
|
||||||
assert!(batch.n_kmers[0] > 0);
|
|
||||||
assert_eq!(occurrences, batch.n_kmers[0] as u64 + batch.n_kmers[1] as u64);
|
|
||||||
assert!(
|
|
||||||
unique <= batch.n_kmers[0] as u64,
|
|
||||||
"identical sequences must not produce more unique k-mers than one copy has"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
unique < occurrences,
|
|
||||||
"k-mer-level dedup must collapse at least the cross-record duplication"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn duplicated_sequence_broadcasts_to_both_seq_indices() {
|
|
||||||
// Stronger than the ratio check above: pick any k-mer that hit in both
|
|
||||||
// records and confirm its occurrence list actually references both
|
|
||||||
// seq_idx 0 and seq_idx 1 — this is the specific new capability (dedup
|
|
||||||
// reaching across records/superkmers), not just a smaller unique count.
|
|
||||||
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
|
||||||
let batch = batch_from_fasta(&fasta, K, 1);
|
|
||||||
|
|
||||||
let shared = batch.by_partition[0]
|
|
||||||
.values()
|
|
||||||
.find(|descs| descs.iter().any(|d| d.seq_idx == 0) && descs.iter().any(|d| d.seq_idx == 1));
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
shared.is_some(),
|
|
||||||
"expected at least one k-mer shared between the two identical records"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn empty_records_yield_empty_batch() {
|
|
||||||
let batch = batch_from_fasta("", K, 1);
|
|
||||||
assert!(batch.ids.is_empty());
|
|
||||||
assert_eq!(total_occurrences(&batch), 0);
|
|
||||||
assert_eq!(total_unique_kmers(&batch), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partition_routing_is_a_pure_function_of_the_kmer() {
|
|
||||||
// With n_partitions=4, every occurrence of a given k-mer must land in
|
|
||||||
// the same partition bucket as every other occurrence of that k-mer
|
|
||||||
// (partition routing is derived from the minimizer, shared by
|
|
||||||
// definition among instances of the same k-mer's containing superkmer
|
|
||||||
// in this test's single-sequence-pair setup).
|
|
||||||
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
|
||||||
let batch = batch_from_fasta(&fasta, K, 4);
|
|
||||||
|
|
||||||
let total_unique: u64 = batch.by_partition.iter().map(|m| m.len() as u64).sum();
|
|
||||||
assert!(total_unique > 0);
|
|
||||||
|
|
||||||
// No k-mer key appears in more than one partition's map.
|
|
||||||
let mut seen: std::collections::HashSet<CanonicalKmer> = std::collections::HashSet::new();
|
|
||||||
for map in &batch.by_partition {
|
|
||||||
for kmer in map.keys() {
|
|
||||||
assert!(seen.insert(*kmer), "k-mer routed to more than one partition");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── sparse_findere_for_genome vs. a dense reference implementation ──────────
|
|
||||||
//
|
|
||||||
// No property-testing crate (proptest/quickcheck) is a workspace dependency
|
|
||||||
// (checked before writing this — not adding one for a single test module,
|
|
||||||
// per this project's dependency-approval rule). A tiny deterministic xorshift
|
|
||||||
// PRNG, std-only, stands in for one.
|
|
||||||
|
|
||||||
/// Faithful reimplementation of the pre-phase-5 dense sliding-window scan —
|
|
||||||
/// the algorithm `sparse_findere_for_genome` replaced — used here only as a
|
|
||||||
/// correctness oracle, not in production code. Operates on one genome's
|
|
||||||
/// hits across possibly many sequences, exactly like the sparse version.
|
|
||||||
fn dense_reference_findere(
|
|
||||||
hits: &[(u32, u32, u32)],
|
|
||||||
seq_lens: &[usize],
|
|
||||||
z: usize,
|
|
||||||
presence: bool,
|
|
||||||
threshold: u32,
|
|
||||||
) -> Vec<(u32, u32, u32)> {
|
|
||||||
let mut by_seq: Vec<Vec<u32>> = seq_lens.iter().map(|&n| vec![0u32; n]).collect();
|
|
||||||
for &(seq, pos, val) in hits {
|
|
||||||
by_seq[seq as usize][pos as usize] = val;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut confirmed = Vec::new();
|
|
||||||
for (seq_idx, values) in by_seq.iter().enumerate() {
|
|
||||||
let n = values.len();
|
|
||||||
let mut dq: std::collections::VecDeque<(usize, u32)> = std::collections::VecDeque::new();
|
|
||||||
for i in 0..n {
|
|
||||||
let v_i = values[i];
|
|
||||||
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
|
||||||
dq.pop_front();
|
|
||||||
}
|
|
||||||
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
|
||||||
dq.pop_back();
|
|
||||||
}
|
|
||||||
dq.push_back((i, v_i));
|
|
||||||
if i + 1 >= z {
|
|
||||||
let win_min = dq.front().unwrap().1;
|
|
||||||
if win_min > 0 {
|
|
||||||
let pos_out = (i + 1 - z) as u32;
|
|
||||||
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
|
||||||
confirmed.push((seq_idx as u32, pos_out, c));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
confirmed
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Minimal std-only xorshift64 PRNG — deterministic, seedable, no dependency.
|
|
||||||
struct Xorshift64(u64);
|
|
||||||
impl Xorshift64 {
|
|
||||||
fn next(&mut self) -> u64 {
|
|
||||||
self.0 ^= self.0 << 13;
|
|
||||||
self.0 ^= self.0 >> 7;
|
|
||||||
self.0 ^= self.0 << 17;
|
|
||||||
self.0
|
|
||||||
}
|
|
||||||
fn range(&mut self, n: u32) -> u32 {
|
|
||||||
(self.next() % n as u64) as u32
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sparse_findere_matches_dense_reference_on_random_inputs() {
|
|
||||||
let mut rng = Xorshift64(0x5eed_5eed_5eed_5eedu64);
|
|
||||||
|
|
||||||
for case in 0..200 {
|
|
||||||
let n_seqs = 1 + rng.range(4) as usize;
|
|
||||||
let seq_lens: Vec<usize> = (0..n_seqs).map(|_| 1 + rng.range(30) as usize).collect();
|
|
||||||
let z = 1 + rng.range(4) as usize;
|
|
||||||
let presence = rng.range(2) == 0;
|
|
||||||
let threshold = 1 + rng.range(3);
|
|
||||||
|
|
||||||
// Sparse density varies across cases, including edge cases (empty,
|
|
||||||
// fully dense) — deliberately not uniform, to stress both few-hits
|
|
||||||
// and many-overlapping-runs scenarios.
|
|
||||||
let density = rng.range(101);
|
|
||||||
let mut hits: Vec<(u32, u32, u32)> = Vec::new();
|
|
||||||
for (seq_idx, &len) in seq_lens.iter().enumerate() {
|
|
||||||
for pos in 0..len {
|
|
||||||
if rng.range(100) < density {
|
|
||||||
let val = 1 + rng.range(5); // never 0 — matches QueryHit::Value's invariant
|
|
||||||
hits.push((seq_idx as u32, pos as u32, val));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut sparse_input = hits.clone();
|
|
||||||
let (mut sparse_result, _, _) =
|
|
||||||
sparse_findere_for_genome(&mut sparse_input, z, presence, threshold);
|
|
||||||
let mut dense_result = dense_reference_findere(&hits, &seq_lens, z, presence, threshold);
|
|
||||||
|
|
||||||
sparse_result.sort_unstable();
|
|
||||||
dense_result.sort_unstable();
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
sparse_result, dense_result,
|
|
||||||
"case {case}: n_seqs={n_seqs} seq_lens={seq_lens:?} z={z} presence={presence} \
|
|
||||||
threshold={threshold} density={density} hits={hits:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -21,7 +21,7 @@ enum Commands {
|
|||||||
/// Merge multiple built indexes into one
|
/// Merge multiple built indexes into one
|
||||||
Merge(cmd::merge::MergeArgs),
|
Merge(cmd::merge::MergeArgs),
|
||||||
/// Apply row-level selection (σ) to an index: retain only k-mers matching the predicates
|
/// Apply row-level selection (σ) to an index: retain only k-mers matching the predicates
|
||||||
Filter(cmd::filter::FilterCmdArgs),
|
Filter(cmd::filter::FilterArgs),
|
||||||
/// Project and/or aggregate genome columns into a new or in-place index
|
/// Project and/or aggregate genome columns into a new or in-place index
|
||||||
Select(cmd::select::SelectArgs),
|
Select(cmd::select::SelectArgs),
|
||||||
/// Query an index with sequences and annotate matches
|
/// Query an index with sequences and annotate matches
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ edition = "2024"
|
|||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tempfile = "3"
|
tempfile = "3"
|
||||||
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
||||||
|
obiskbuilder = { path = "../obiskbuilder" }
|
||||||
obiread = { path = "../obiread" }
|
obiread = { path = "../obiread" }
|
||||||
obikrope = { path = "../obikrope" }
|
obikrope = { path = "../obikrope" }
|
||||||
|
|
||||||
@@ -13,8 +14,6 @@ obikrope = { path = "../obikrope" }
|
|||||||
niffler = "3.0.0"
|
niffler = "3.0.0"
|
||||||
remove_dir_all = "0.8"
|
remove_dir_all = "0.8"
|
||||||
obikseq = { path = "../obikseq" }
|
obikseq = { path = "../obikseq" }
|
||||||
obikentropy = { path = "../obikentropy" }
|
|
||||||
obiskbuilder = { path = "../obiskbuilder" }
|
|
||||||
obiskio = { path = "../obiskio" }
|
obiskio = { path = "../obiskio" }
|
||||||
obidebruinj = { path = "../obidebruinj" }
|
obidebruinj = { path = "../obidebruinj" }
|
||||||
obilayeredmap = { path = "../obilayeredmap" }
|
obilayeredmap = { path = "../obilayeredmap" }
|
||||||
|
|||||||
@@ -62,7 +62,7 @@ impl KmerPartition {
|
|||||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||||
if let Some(slot) = mphf.find(kmer) {
|
if let Some(slot) = mphf.find(kmer) {
|
||||||
let row = mat.row(slot);
|
let row = mat.row(slot);
|
||||||
if passes_all(filters, kmer, &row, n_genomes) {
|
if passes_all(filters, &row, n_genomes) {
|
||||||
cont = cb(kmer, row);
|
cont = cb(kmer, row);
|
||||||
if !cont { break; }
|
if !cont { break; }
|
||||||
}
|
}
|
||||||
@@ -75,7 +75,7 @@ impl KmerPartition {
|
|||||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||||
if let Some(slot) = mphf.find(kmer) {
|
if let Some(slot) = mphf.find(kmer) {
|
||||||
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
|
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
|
||||||
if passes_all(filters, kmer, &row, n_genomes) {
|
if passes_all(filters, &row, n_genomes) {
|
||||||
cont = cb(kmer, row);
|
cont = cb(kmer, row);
|
||||||
if !cont { break; }
|
if !cont { break; }
|
||||||
}
|
}
|
||||||
@@ -83,19 +83,18 @@ impl KmerPartition {
|
|||||||
}
|
}
|
||||||
cont
|
cont
|
||||||
} else {
|
} else {
|
||||||
// No data matrix: implicit presence — all values are 1. `row`
|
// No data matrix: implicit presence — all values are 1.
|
||||||
// is identical for every kmer, but a filter can still depend
|
// The filter result is identical for every kmer, so evaluate once.
|
||||||
// on the kmer's own sequence (e.g. MinComplexity), so this
|
|
||||||
// cannot be evaluated once for the whole layer — filters must
|
|
||||||
// still be tested per kmer.
|
|
||||||
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
|
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
|
||||||
let mut cont = true;
|
let mut cont = true;
|
||||||
|
if passes_all(filters, &all_present, n_genomes) {
|
||||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
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() {
|
||||||
cont = cb(kmer, all_present.clone());
|
cont = cb(kmer, all_present.clone());
|
||||||
if !cont { break; }
|
if !cont { break; }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
cont
|
cont
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -141,7 +140,7 @@ impl KmerPartition {
|
|||||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||||
if let Some(slot) = mphf.find(kmer) {
|
if let Some(slot) = mphf.find(kmer) {
|
||||||
let row = mat.row(slot);
|
let row = mat.row(slot);
|
||||||
if passes_all(filters, kmer, &row, n_genomes) {
|
if passes_all(filters, &row, n_genomes) {
|
||||||
cont = cb(part, layer, kmer, row);
|
cont = cb(part, layer, kmer, row);
|
||||||
if !cont { break; }
|
if !cont { break; }
|
||||||
}
|
}
|
||||||
@@ -154,7 +153,7 @@ impl KmerPartition {
|
|||||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||||||
if let Some(slot) = mphf.find(kmer) {
|
if let Some(slot) = mphf.find(kmer) {
|
||||||
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
|
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
|
||||||
if passes_all(filters, kmer, &row, n_genomes) {
|
if passes_all(filters, &row, n_genomes) {
|
||||||
cont = cb(part, layer, kmer, row);
|
cont = cb(part, layer, kmer, row);
|
||||||
if !cont { break; }
|
if !cont { break; }
|
||||||
}
|
}
|
||||||
@@ -162,17 +161,16 @@ impl KmerPartition {
|
|||||||
}
|
}
|
||||||
cont
|
cont
|
||||||
} else {
|
} else {
|
||||||
// Same as iter_partition_kmers: row is constant but a filter
|
|
||||||
// may still depend on the kmer's own sequence, so this must
|
|
||||||
// be tested per kmer, not once for the whole layer.
|
|
||||||
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
|
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
|
||||||
let mut cont = true;
|
let mut cont = true;
|
||||||
|
if passes_all(filters, &all_present, n_genomes) {
|
||||||
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
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() {
|
||||||
cont = cb(part, layer, kmer, all_present.clone());
|
cont = cb(part, layer, kmer, all_present.clone());
|
||||||
if !cont { break; }
|
if !cont { break; }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
cont
|
cont
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -1,24 +1,17 @@
|
|||||||
use obicompactvec::FilterMask;
|
use obicompactvec::FilterMask;
|
||||||
use obikseq::CanonicalKmer;
|
|
||||||
|
|
||||||
/// Trait for kmer filters.
|
/// Trait for kmer row filters.
|
||||||
///
|
///
|
||||||
/// `kmer` is the k-mer's own canonical sequence, reconstructed from the
|
|
||||||
/// source index's `unitigs.bin` (always present — see `rebuild_layer.rs`);
|
|
||||||
/// `row` contains raw per-genome counts (or 0/1 for presence/absence data).
|
/// `row` contains raw per-genome counts (or 0/1 for presence/absence data).
|
||||||
/// `n_genomes` equals `row.len()`. Most filters only need `row` — `kmer` is
|
/// `n_genomes` equals `row.len()`.
|
||||||
/// there for filters that reason about the k-mer's sequence itself (e.g.
|
|
||||||
/// [`MinComplexity`]).
|
|
||||||
pub trait KmerFilter: Send + Sync {
|
pub trait KmerFilter: Send + Sync {
|
||||||
fn passes(&self, kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool;
|
fn passes(&self, row: &[u32], n_genomes: usize) -> bool;
|
||||||
|
|
||||||
/// Express this filter as a [`FilterMask`] column-operation expression.
|
/// Express this filter as a [`FilterMask`] column-operation expression.
|
||||||
///
|
///
|
||||||
/// Returns `Some(expr)` if the filter can be evaluated solely from matrix
|
/// Returns `Some(expr)` if the filter can be evaluated solely from matrix
|
||||||
/// column aggregates (no per-kmer row scan needed). Returns `None` if the
|
/// column aggregates (no per-kmer row scan needed). Returns `None` if the
|
||||||
/// filter requires row-level inspection — always the case for a filter
|
/// filter requires row-level inspection.
|
||||||
/// that needs the k-mer's sequence, since a `FilterMask` only expresses
|
|
||||||
/// per-genome column aggregates, never per-slot sequence data.
|
|
||||||
///
|
///
|
||||||
/// `threshold` semantics in the returned mask use `>= threshold`, matching
|
/// `threshold` semantics in the returned mask use `>= threshold`, matching
|
||||||
/// [`obicompactvec::MatrixGroupOps`]. Implementations must add 1 to any
|
/// [`obicompactvec::MatrixGroupOps`]. Implementations must add 1 to any
|
||||||
@@ -30,13 +23,8 @@ pub trait KmerFilter: Send + Sync {
|
|||||||
|
|
||||||
/// True when `row` passes every filter in `filters`.
|
/// True when `row` passes every filter in `filters`.
|
||||||
/// Returns `true` if `filters` is empty.
|
/// Returns `true` if `filters` is empty.
|
||||||
pub fn passes_all(
|
pub fn passes_all(filters: &[Box<dyn KmerFilter>], row: &[u32], n_genomes: usize) -> bool {
|
||||||
filters: &[Box<dyn KmerFilter>],
|
filters.iter().all(|f| f.passes(row, n_genomes))
|
||||||
kmer: CanonicalKmer,
|
|
||||||
row: &[u32],
|
|
||||||
n_genomes: usize,
|
|
||||||
) -> bool {
|
|
||||||
filters.iter().all(|f| f.passes(kmer, row, n_genomes))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Quorum filters ─────────────────────────────────────────────────────────────
|
// ── Quorum filters ─────────────────────────────────────────────────────────────
|
||||||
@@ -52,7 +40,7 @@ pub struct MinGenomeFraction {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for MinGenomeFraction {
|
impl KmerFilter for MinGenomeFraction {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], n_genomes: usize) -> bool {
|
||||||
let p = present_count(row, self.threshold);
|
let p = present_count(row, self.threshold);
|
||||||
p as f64 / n_genomes as f64 >= self.frac
|
p as f64 / n_genomes as f64 >= self.frac
|
||||||
}
|
}
|
||||||
@@ -75,7 +63,7 @@ pub struct MaxGenomeFraction {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for MaxGenomeFraction {
|
impl KmerFilter for MaxGenomeFraction {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], n_genomes: usize) -> bool {
|
||||||
let p = present_count(row, self.threshold);
|
let p = present_count(row, self.threshold);
|
||||||
p as f64 / n_genomes as f64 <= self.frac
|
p as f64 / n_genomes as f64 <= self.frac
|
||||||
}
|
}
|
||||||
@@ -98,7 +86,7 @@ pub struct MinGenomeCount {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for MinGenomeCount {
|
impl KmerFilter for MinGenomeCount {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
|
||||||
present_count(row, self.threshold) >= self.count
|
present_count(row, self.threshold) >= self.count
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -119,7 +107,7 @@ pub struct MaxGenomeCount {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for MaxGenomeCount {
|
impl KmerFilter for MaxGenomeCount {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
|
||||||
present_count(row, self.threshold) <= self.count
|
present_count(row, self.threshold) <= self.count
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -141,7 +129,7 @@ pub struct MinTotalCount {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for MinTotalCount {
|
impl KmerFilter for MinTotalCount {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
|
||||||
row.iter().sum::<u32>() >= self.total
|
row.iter().sum::<u32>() >= self.total
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -159,7 +147,7 @@ pub struct MaxTotalCount {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for MaxTotalCount {
|
impl KmerFilter for MaxTotalCount {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
|
||||||
row.iter().sum::<u32>() <= self.total
|
row.iter().sum::<u32>() <= self.total
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -224,7 +212,7 @@ impl GroupQuorumFilter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl KmerFilter for GroupQuorumFilter {
|
impl KmerFilter for GroupQuorumFilter {
|
||||||
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
|
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
|
||||||
if !self.ingroup_idx.is_empty() {
|
if !self.ingroup_idx.is_empty() {
|
||||||
let n = self.ingroup_idx.iter()
|
let n = self.ingroup_idx.iter()
|
||||||
.filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
|
.filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
|
||||||
@@ -272,27 +260,3 @@ impl KmerFilter for GroupQuorumFilter {
|
|||||||
Some(FilterMask::And(parts))
|
Some(FilterMask::And(parts))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Complexity filter (post-hoc, sequence-based) ──────────────────────────────
|
|
||||||
|
|
||||||
/// Reject k-mers with normalized entropy below `theta` — the same complexity
|
|
||||||
/// metric `obikmer index`'s `--theta`/`--level-max` apply *during* superkmer
|
|
||||||
/// construction (see [`obikentropy::KmerEntropy`]), applied here after the
|
|
||||||
/// fact, to k-mers already committed to a built index.
|
|
||||||
///
|
|
||||||
/// Unlike every other filter in this module, this one needs the k-mer's own
|
|
||||||
/// sequence, not its per-genome row — `column_mask_expr` is never overridden
|
|
||||||
/// (stays `None`), so this filter always forces the row-level scan path in
|
|
||||||
/// `rebuild_layer.rs` (which reconstructs the sequence from `unitigs.bin`
|
|
||||||
/// regardless, so no extra I/O beyond what filtering already requires).
|
|
||||||
pub struct MinComplexity {
|
|
||||||
pub level_max: usize,
|
|
||||||
pub theta: f64,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl KmerFilter for MinComplexity {
|
|
||||||
fn passes(&self, kmer: CanonicalKmer, _row: &[u32], _n_genomes: usize) -> bool {
|
|
||||||
use obikentropy::KmerEntropy;
|
|
||||||
kmer.entropy(self.level_max) >= self.theta
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -14,5 +14,4 @@ mod select_layer;
|
|||||||
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
|
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
|
||||||
pub use merge_layer::MergeMode;
|
pub use merge_layer::MergeMode;
|
||||||
pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
|
pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
|
||||||
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
|
|
||||||
pub use select_layer::{AggOp, OutputCol};
|
pub use select_layer::{AggOp, OutputCol};
|
||||||
|
|||||||
@@ -1,8 +1,7 @@
|
|||||||
use std::collections::HashMap;
|
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||||
use obikseq::CanonicalKmer;
|
use obikseq::{CanonicalKmer, RoutableSuperKmer};
|
||||||
use obiskio::{SKError, SKResult};
|
use obiskio::{SKError, SKResult};
|
||||||
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
|
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
|
||||||
use obilayeredmap::meta::PartitionMeta;
|
use obilayeredmap::meta::PartitionMeta;
|
||||||
@@ -45,133 +44,53 @@ impl QueryLayer {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// MPHF lookup only — no matrix access. `Some(slot)` on hit.
|
/// Write per-genome values into `buf` if `kmer` is indexed; returns true on hit.
|
||||||
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
|
fn find_into(&self, kmer: CanonicalKmer, n_genomes: usize, buf: &mut [u32]) -> bool {
|
||||||
match self {
|
match self {
|
||||||
QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
|
QueryLayer::Presence(mphf, mat) => {
|
||||||
|
if let Some(slot) = mphf.find(kmer) {
|
||||||
|
mat.fill_row(slot, &mut buf[..n_genomes]);
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
QueryLayer::Count(mphf, mat) => {
|
||||||
|
if let Some(slot) = mphf.find(kmer) {
|
||||||
|
mat.fill_row(slot, &mut buf[..n_genomes]);
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Number of genome columns this layer's matrix actually has. Bounds
|
// ── KmerPartition::query_partition* ──────────────────────────────────────────
|
||||||
/// column-major iteration — usually equal to the index's `n_genomes`, but
|
|
||||||
/// `PersistentBitMatrix::Implicit` (the documented mono-genome fast path)
|
|
||||||
/// always reports exactly `1`, regardless of the index's real genome
|
|
||||||
/// count, so callers must use this rather than assuming `n_genomes`.
|
|
||||||
fn n_cols(&self) -> usize {
|
|
||||||
match self {
|
|
||||||
QueryLayer::Presence(_, mat) => mat.n_cols(),
|
|
||||||
QueryLayer::Count(_, mat) => mat.n_cols(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Column-major point lookup: value for genome column `g` at `slot`.
|
|
||||||
/// `g` must be `< self.n_cols()`; `slot` must come from [`find_slot`] on
|
|
||||||
/// this same layer.
|
|
||||||
fn col_value(&self, g: usize, slot: usize) -> u32 {
|
|
||||||
match self {
|
|
||||||
QueryLayer::Presence(_, mat) => mat.get(g, slot),
|
|
||||||
QueryLayer::Count(_, mat) => mat.col_view(g).get(slot),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── KmerDesc — one occurrence of a k-mer in the query batch ──────────────────
|
|
||||||
|
|
||||||
/// Describes one occurrence of a (deduplicated) k-mer in the query batch:
|
|
||||||
/// which sequence it came from, and its absolute s-mer position within it.
|
|
||||||
#[derive(Debug, Clone, Copy)]
|
|
||||||
pub struct KmerDesc {
|
|
||||||
pub seq_idx: u32,
|
|
||||||
pub pos: u32,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Aggregate counters for one `query_partition_with` call — feeds the
|
|
||||||
/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences
|
|
||||||
/// vs. unique k-mers is the whole justification for k-mer-level
|
|
||||||
/// dereplication; columns scanned / `get()` calls quantify the column-major
|
|
||||||
/// fetch's locality claim).
|
|
||||||
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
|
|
||||||
pub struct QueryStats {
|
|
||||||
/// Distinct canonical k-mers queried in this partition.
|
|
||||||
pub n_unique_kmers: usize,
|
|
||||||
/// Total `MphfLayer::find` calls issued (a k-mer tried against more than
|
|
||||||
/// one layer before a hit, or against all layers on a miss, counts once
|
|
||||||
/// per layer attempted).
|
|
||||||
pub n_mphf_calls: usize,
|
|
||||||
/// Distinct canonical k-mers that matched some layer.
|
|
||||||
pub n_hits: usize,
|
|
||||||
/// Total genome columns scanned across all hit layers (sum of
|
|
||||||
/// `layer.n_cols()` over layers with at least one hit).
|
|
||||||
pub n_columns_scanned: usize,
|
|
||||||
/// Total `col_value` calls issued during the column-major fetch pass
|
|
||||||
/// (`n_columns_scanned` × hits-per-layer, summed over layers).
|
|
||||||
pub n_col_get_calls: usize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::ops::AddAssign for QueryStats {
|
|
||||||
fn add_assign(&mut self, other: Self) {
|
|
||||||
self.n_unique_kmers += other.n_unique_kmers;
|
|
||||||
self.n_mphf_calls += other.n_mphf_calls;
|
|
||||||
self.n_hits += other.n_hits;
|
|
||||||
self.n_columns_scanned += other.n_columns_scanned;
|
|
||||||
self.n_col_get_calls += other.n_col_get_calls;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── QueryHit — one event delivered to query_partition_with's callback ───────
|
|
||||||
|
|
||||||
/// One event from [`KmerPartition::query_partition_with`]'s two-stage query:
|
|
||||||
/// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as
|
|
||||||
/// indexed regardless of any genome's value), then a `Value` event per
|
|
||||||
/// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2,
|
|
||||||
/// column-major fetch). Carried as one enum, not two separate callbacks, so
|
|
||||||
/// the caller only needs one `FnMut` closure — passing two closures that each
|
|
||||||
/// need to mutably borrow the same accumulator does not borrow-check.
|
|
||||||
pub enum QueryHit<'a> {
|
|
||||||
Found(&'a [KmerDesc]),
|
|
||||||
Value(&'a [KmerDesc], usize, u32),
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── KmerPartition::query_partition_with ──────────────────────────────────────
|
|
||||||
|
|
||||||
impl KmerPartition {
|
impl KmerPartition {
|
||||||
/// Query a single partition for a pre-deduplicated map of canonical
|
/// Query a single partition, calling `on_hit(sk_idx, kmer_idx, row)` for
|
||||||
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
|
/// every found k-mer without allocating intermediate result vectors.
|
||||||
///
|
|
||||||
/// Two stages:
|
|
||||||
/// 1. **MPHF-only pass**: for each unique k-mer, try each layer's MPHF in
|
|
||||||
/// turn (stopping at the first hit) and bucket confirmed hits by
|
|
||||||
/// `(layer, slot)`. Emits one `QueryHit::Found` per hit k-mer. This
|
|
||||||
/// stage's cost is independent of the index's genome count.
|
|
||||||
/// 2. **Column-major fetch**: for each layer with at least one hit, walk
|
|
||||||
/// its matrix **column by column** (genome by genome) — for each
|
|
||||||
/// genome, scan the slots bucketed in stage 1 and look up their value.
|
|
||||||
/// Emits one `QueryHit::Value` per nonzero `(k-mer, genome)` pair.
|
|
||||||
/// Total lookups are the same as a row-major pass (`n_hits × n_cols`
|
|
||||||
/// in the worst case); the win is memory locality — both persistent
|
|
||||||
/// matrix formats are column-oriented on disk (one `mmap`'d region per
|
|
||||||
/// genome), so scanning one column at a time touches far fewer
|
|
||||||
/// distinct mmap regions than fetching one full row per hit.
|
|
||||||
pub fn query_partition_with<F>(
|
pub fn query_partition_with<F>(
|
||||||
&self,
|
&self,
|
||||||
part_idx: usize,
|
part_idx: usize,
|
||||||
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
superkmers: &[&RoutableSuperKmer],
|
||||||
|
_k: usize,
|
||||||
n_genomes: usize,
|
n_genomes: usize,
|
||||||
with_counts: bool,
|
with_counts: bool,
|
||||||
mut on_event: F,
|
mut on_hit: F,
|
||||||
) -> SKResult<QueryStats>
|
) -> SKResult<()>
|
||||||
where
|
where
|
||||||
F: FnMut(QueryHit),
|
F: FnMut(usize, usize, &[u32]),
|
||||||
{
|
{
|
||||||
let mut stats = QueryStats::default();
|
if superkmers.is_empty() {
|
||||||
|
return Ok(());
|
||||||
if kmers.is_empty() {
|
|
||||||
return Ok(stats);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
||||||
if !index_dir.exists() {
|
if !index_dir.exists() {
|
||||||
return Ok(stats);
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
||||||
@@ -179,47 +98,69 @@ impl KmerPartition {
|
|||||||
.map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
|
.map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
|
||||||
.collect::<SKResult<_>>()?;
|
.collect::<SKResult<_>>()?;
|
||||||
|
|
||||||
// ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
|
let mut buf = vec![0u32; n_genomes];
|
||||||
let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
|
|
||||||
(0..layers.len()).map(|_| HashMap::new()).collect();
|
|
||||||
|
|
||||||
for (kmer, descs) in kmers {
|
for (sk_idx, rsk) in superkmers.iter().enumerate() {
|
||||||
stats.n_unique_kmers += 1;
|
for (kmer_idx, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
|
||||||
for (layer_idx, layer) in layers.iter().enumerate() {
|
for layer in &layers {
|
||||||
stats.n_mphf_calls += 1;
|
if layer.find_into(kmer, n_genomes, &mut buf) {
|
||||||
if let Some(slot) = layer.find_slot(*kmer) {
|
on_hit(sk_idx, kmer_idx, &buf);
|
||||||
by_layer[layer_idx].insert(slot, descs);
|
buf.fill(0);
|
||||||
on_event(QueryHit::Found(descs));
|
|
||||||
stats.n_hits += 1;
|
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Stage 2: column-major fetch, per layer ───────────────────────────
|
|
||||||
for (layer_idx, slots) in by_layer.iter().enumerate() {
|
|
||||||
if slots.is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
let layer = &layers[layer_idx];
|
|
||||||
let n_cols = layer.n_cols().min(n_genomes);
|
|
||||||
stats.n_columns_scanned += n_cols;
|
|
||||||
|
|
||||||
for g in 0..n_cols {
|
|
||||||
for (&slot, descs) in slots {
|
|
||||||
stats.n_col_get_calls += 1;
|
|
||||||
let v = layer.col_value(g, slot);
|
|
||||||
if v != 0 {
|
|
||||||
on_event(QueryHit::Value(descs, g, v));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(stats)
|
Ok(())
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
/// Query a single partition for a slice of super-kmers, returning per-kmer rows.
|
||||||
#[path = "tests/query_layer.rs"]
|
/// Prefer [`query_partition_with`] to avoid per-kmer heap allocations.
|
||||||
mod tests;
|
pub fn query_partition(
|
||||||
|
&self,
|
||||||
|
part_idx: usize,
|
||||||
|
superkmers: &[&RoutableSuperKmer],
|
||||||
|
_k: usize,
|
||||||
|
n_genomes: usize,
|
||||||
|
with_counts: bool,
|
||||||
|
) -> SKResult<Vec<Vec<Option<Box<[u32]>>>>> {
|
||||||
|
if superkmers.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
||||||
|
|
||||||
|
if !index_dir.exists() {
|
||||||
|
return Ok(superkmers
|
||||||
|
.iter()
|
||||||
|
.map(|rsk| vec![None; rsk.seql()])
|
||||||
|
.collect());
|
||||||
|
}
|
||||||
|
|
||||||
|
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
||||||
|
let layers: Vec<QueryLayer> = (0..meta.n_layers)
|
||||||
|
.map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
|
||||||
|
.collect::<SKResult<_>>()?;
|
||||||
|
|
||||||
|
let mut buf = vec![0u32; n_genomes];
|
||||||
|
Ok(superkmers
|
||||||
|
.iter()
|
||||||
|
.map(|rsk| {
|
||||||
|
rsk.superkmer()
|
||||||
|
.iter_canonical_kmers()
|
||||||
|
.map(|kmer| {
|
||||||
|
for layer in &layers {
|
||||||
|
if layer.find_into(kmer, n_genomes, &mut buf) {
|
||||||
|
let row: Box<[u32]> = buf[..n_genomes].into();
|
||||||
|
buf.fill(0);
|
||||||
|
return Some(row);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -126,7 +126,7 @@ fn iter_src_kmers_masked(
|
|||||||
Some(m) => m.get(slot),
|
Some(m) => m.get(slot),
|
||||||
None => {
|
None => {
|
||||||
let row = src_data.fill_row_by_slot(slot, n_genomes);
|
let row = src_data.fill_row_by_slot(slot, n_genomes);
|
||||||
filters.iter().all(|f| f.passes(kmer, &row, n_genomes))
|
filters.iter().all(|f| f.passes(&row, n_genomes))
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
if passes { cb(kmer); }
|
if passes { cb(kmer); }
|
||||||
@@ -165,7 +165,7 @@ fn iter_src_layers(
|
|||||||
cb(kmer, row.into_boxed_slice());
|
cb(kmer, row.into_boxed_slice());
|
||||||
} else {
|
} else {
|
||||||
let row = src_data.fill_row_by_slot(slot, n_genomes);
|
let row = src_data.fill_row_by_slot(slot, n_genomes);
|
||||||
if filters.iter().all(|f| f.passes(kmer, &row, n_genomes)) {
|
if filters.iter().all(|f| f.passes(&row, n_genomes)) {
|
||||||
cb(kmer, row.into_boxed_slice());
|
cb(kmer, row.into_boxed_slice());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,79 +0,0 @@
|
|||||||
use super::*;
|
|
||||||
|
|
||||||
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn query_stats_add_assign_sums_fields() {
|
|
||||||
let mut total = QueryStats {
|
|
||||||
n_unique_kmers: 3,
|
|
||||||
n_mphf_calls: 5,
|
|
||||||
n_hits: 2,
|
|
||||||
n_columns_scanned: 1,
|
|
||||||
n_col_get_calls: 7,
|
|
||||||
};
|
|
||||||
total += QueryStats {
|
|
||||||
n_unique_kmers: 1,
|
|
||||||
n_mphf_calls: 4,
|
|
||||||
n_hits: 1,
|
|
||||||
n_columns_scanned: 2,
|
|
||||||
n_col_get_calls: 3,
|
|
||||||
};
|
|
||||||
assert_eq!(total.n_unique_kmers, 4);
|
|
||||||
assert_eq!(total.n_mphf_calls, 9);
|
|
||||||
assert_eq!(total.n_hits, 3);
|
|
||||||
assert_eq!(total.n_columns_scanned, 3);
|
|
||||||
assert_eq!(total.n_col_get_calls, 10);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn query_stats_default_is_zero() {
|
|
||||||
let s = QueryStats::default();
|
|
||||||
assert_eq!(s.n_unique_kmers, 0);
|
|
||||||
assert_eq!(s.n_mphf_calls, 0);
|
|
||||||
assert_eq!(s.n_hits, 0);
|
|
||||||
assert_eq!(s.n_columns_scanned, 0);
|
|
||||||
assert_eq!(s.n_col_get_calls, 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── query_partition_with on a not-yet-indexed partition ─────────────────────
|
|
||||||
|
|
||||||
/// A `KmerPartition` created but never taken through `build_layers` has no
|
|
||||||
/// `index/` subdirectory under any partition — `query_partition_with` must
|
|
||||||
/// recognise this and return default (all-zero) stats rather than erroring,
|
|
||||||
/// exactly like an empty `kmers` map.
|
|
||||||
#[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 mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
|
||||||
// Any well-formed canonical k-mer works here — the call must return
|
|
||||||
// before ever attempting an MPHF lookup, since `index/` doesn't exist.
|
|
||||||
let kmer = CanonicalKmer::from_raw_unchecked(0u64);
|
|
||||||
kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]);
|
|
||||||
|
|
||||||
let stats = partition
|
|
||||||
.query_partition_with(0, &kmers, 1, false, |_event| {
|
|
||||||
panic!("on_event must not be called: no index was built");
|
|
||||||
})
|
|
||||||
.expect("query_partition_with should not error on a missing index dir");
|
|
||||||
|
|
||||||
assert_eq!(stats, QueryStats::default());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[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 kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
|
||||||
let stats = partition
|
|
||||||
.query_partition_with(0, &kmers, 1, false, |_event| {
|
|
||||||
panic!("on_event must not be called on an empty kmer map");
|
|
||||||
})
|
|
||||||
.expect("query_partition_with on an empty map should not error");
|
|
||||||
|
|
||||||
assert_eq!(stats, QueryStats::default());
|
|
||||||
}
|
|
||||||
@@ -341,27 +341,6 @@ impl<L: KmerLength> CanonicalKmerOf<L> {
|
|||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Return the four central canonical neighbours (each already canonical),
|
|
||||||
/// substituting the base at the middle position `m = (L::len()-1)/2`
|
|
||||||
/// (well-defined for odd `L::len()`). Each of the 4 substitutions is
|
|
||||||
/// canonicalised independently — this correctly handles the case where a
|
|
||||||
/// substitution flips the canonical orientation, unlike inferring the
|
|
||||||
/// variant from a fixed-orientation flank key. One of the 4 equals
|
|
||||||
/// `self`'s own canonical form (the identity substitution); callers that
|
|
||||||
/// only want the 3 genuine variants should skip it.
|
|
||||||
pub fn central_canonical_neighbors(&self) -> [CanonicalKmerOf<L>; 4] {
|
|
||||||
let k = L::len();
|
|
||||||
let m = (k - 1) / 2;
|
|
||||||
let shift = KMER_BITS - 2 - 2 * m;
|
|
||||||
let cleared = self.0 & !((0b11 as RawKmer) << shift);
|
|
||||||
[
|
|
||||||
KmerOf::<L>(cleared | ((0 as RawKmer) << shift), PhantomData).canonical(),
|
|
||||||
KmerOf::<L>(cleared | ((1 as RawKmer) << shift), PhantomData).canonical(),
|
|
||||||
KmerOf::<L>(cleared | ((2 as RawKmer) << shift), PhantomData).canonical(),
|
|
||||||
KmerOf::<L>(cleared | ((3 as RawKmer) << shift), PhantomData).canonical(),
|
|
||||||
]
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Return the inner value as a raw [`KmerOf<L>`].
|
/// Return the inner value as a raw [`KmerOf<L>`].
|
||||||
#[inline]
|
#[inline]
|
||||||
pub fn into_kmer(self) -> KmerOf<L> {
|
pub fn into_kmer(self) -> KmerOf<L> {
|
||||||
|
|||||||
@@ -210,46 +210,4 @@ mod tests {
|
|||||||
check!(31);
|
check!(31);
|
||||||
check!(32);
|
check!(32);
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── central_canonical_neighbors ─────────────────────────────────────────
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn central_canonical_neighbors_hand_checked_k3() {
|
|
||||||
// k=3, centre = index 1. For "ACG", every one of the 4 central
|
|
||||||
// substitutions ("AAG","ACG","AGG","ATG") happens to stay in forward
|
|
||||||
// orientation when canonicalised (verified by hand: each is already
|
|
||||||
// lexicographically <= its own reverse complement), so this case
|
|
||||||
// exercises the substitution logic without the RC-flip edge case.
|
|
||||||
let ck = KmerOf::<ConstLen<3>>::from_ascii(b"ACG").unwrap().canonical();
|
|
||||||
let neighbours = ck.central_canonical_neighbors();
|
|
||||||
let ascii: Vec<Vec<u8>> = neighbours.iter().map(|n| n.to_ascii()).collect();
|
|
||||||
assert_eq!(ascii, vec![b"AAG".to_vec(), b"ACG".to_vec(), b"AGG".to_vec(), b"ATG".to_vec()]);
|
|
||||||
// The identity substitution (centre unchanged) must reproduce `ck`.
|
|
||||||
assert!(neighbours.contains(&ck));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn central_canonical_neighbors_identity_present_for_various_k() {
|
|
||||||
macro_rules! check {
|
|
||||||
($n:expr) => {{
|
|
||||||
let ck = KmerOf::<ConstLen<$n>>::from_ascii(&make_seq::<$n>())
|
|
||||||
.unwrap()
|
|
||||||
.canonical();
|
|
||||||
let neighbours = ck.central_canonical_neighbors();
|
|
||||||
assert!(
|
|
||||||
neighbours.contains(&ck),
|
|
||||||
"identity substitution missing from central_canonical_neighbors for k={}",
|
|
||||||
$n
|
|
||||||
);
|
|
||||||
// Every returned neighbour must itself already be canonical.
|
|
||||||
for n in &neighbours {
|
|
||||||
assert_eq!(n.into_kmer().canonical(), *n, "neighbour not canonical for k={}", $n);
|
|
||||||
}
|
|
||||||
}};
|
|
||||||
}
|
|
||||||
check!(1);
|
|
||||||
check!(3);
|
|
||||||
check!(5);
|
|
||||||
check!(31);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -96,5 +96,162 @@ impl<S: BitPartials> BitPartials for LayeredStore<S> {
|
|||||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
#[path = "tests/layered_store.rs"]
|
mod tests {
|
||||||
mod tests;
|
use super::*;
|
||||||
|
use obicompactvec::{
|
||||||
|
PersistentBitMatrix, PersistentBitMatrixBuilder,
|
||||||
|
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
|
||||||
|
};
|
||||||
|
use tempfile::tempdir;
|
||||||
|
|
||||||
|
fn make_int_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
|
||||||
|
let n = cols.first().map_or(0, |c| c.len());
|
||||||
|
let dir = tempdir().unwrap();
|
||||||
|
let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
|
||||||
|
for &col in cols {
|
||||||
|
let mut cb = b.add_col().unwrap();
|
||||||
|
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
|
||||||
|
cb.close().unwrap();
|
||||||
|
}
|
||||||
|
b.close().unwrap();
|
||||||
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
||||||
|
(dir, m)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn make_bit_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
|
||||||
|
let n = cols.first().map_or(0, |c| c.len());
|
||||||
|
let dir = tempdir().unwrap();
|
||||||
|
let mut b = PersistentBitMatrixBuilder::new(n, &dir.path().join("presence")).unwrap();
|
||||||
|
for &col in cols {
|
||||||
|
let mut cb = b.add_col().unwrap();
|
||||||
|
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
|
||||||
|
cb.close().unwrap();
|
||||||
|
}
|
||||||
|
b.close().unwrap();
|
||||||
|
let m = PersistentBitMatrix::open(dir.path()).unwrap();
|
||||||
|
(dir, m)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── ColumnWeights ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn col_weights_sums_across_layers() {
|
||||||
|
// layer 0: col0=[1,2], col1=[3,4] → weights [3, 7]
|
||||||
|
// layer 1: col0=[10,0], col1=[0,10] → weights [10, 10]
|
||||||
|
// combined: [13, 17]
|
||||||
|
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[3, 4]]);
|
||||||
|
let (_d1, m1) = make_int_matrix(&[&[10, 0], &[0, 10]]);
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
let w = store.col_weights();
|
||||||
|
assert_eq!(w[0], 13);
|
||||||
|
assert_eq!(w[1], 17);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn col_weights_bit_sums_across_layers() {
|
||||||
|
// layer 0: col0=[T,F,T], col1=[F,T,T] → counts [2, 2]
|
||||||
|
// layer 1: col0=[F,F,T], col1=[T,T,F] → counts [1, 2]
|
||||||
|
// combined: [3, 4]
|
||||||
|
let (_d0, m0) = make_bit_matrix(&[&[true, false, true], &[false, true, true]]);
|
||||||
|
let (_d1, m1) = make_bit_matrix(&[&[false, false, true], &[true, true, false]]);
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
let w = store.col_weights();
|
||||||
|
assert_eq!(w[0], 3);
|
||||||
|
assert_eq!(w[1], 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── CountPartials — layered (one partition) ───────────────────────────────
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn layered_bray_matches_combined() {
|
||||||
|
// Split [1,2,3,4,5] across two layers; bray dist should equal direct computation
|
||||||
|
// on [1,2,3,4,5] for each column pair.
|
||||||
|
// col0=[1,2,3,4,5], col1=[5,4,3,2,1]
|
||||||
|
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]); // slots 0-1
|
||||||
|
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]); // slots 2-4
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
|
||||||
|
// direct on full data
|
||||||
|
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
||||||
|
let expected = CountPartials::bray_dist_matrix(&mf);
|
||||||
|
let got = CountPartials::bray_dist_matrix(&store);
|
||||||
|
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "bray [0,1]");
|
||||||
|
assert!((got[[1, 0]] - expected[[1, 0]]).abs() < 1e-12, "bray [1,0]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn layered_relfreq_bray_matches_combined() {
|
||||||
|
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]);
|
||||||
|
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
|
||||||
|
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
||||||
|
let expected = CountPartials::relfreq_bray_dist_matrix(&mf);
|
||||||
|
let got = CountPartials::relfreq_bray_dist_matrix(&store);
|
||||||
|
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "relfreq_bray [0,1]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn layered_euclidean_matches_combined() {
|
||||||
|
let (_d0, m0) = make_int_matrix(&[&[3, 0], &[0, 4]]);
|
||||||
|
let (_d1, m1) = make_int_matrix(&[&[1, 1], &[2, 2]]);
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
|
||||||
|
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 4, 2, 2]]);
|
||||||
|
let expected = CountPartials::euclidean_dist_matrix(&mf);
|
||||||
|
let got = CountPartials::euclidean_dist_matrix(&store);
|
||||||
|
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "euclidean [0,1]");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── CountPartials — partitioned (LayeredStore<LayeredStore<_>>) ───────────
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn partitioned_bray_matches_combined() {
|
||||||
|
// partition 0: slots [1,2,3,4,5] col0 vs col1
|
||||||
|
// partition 1: slots [10,20] col0 vs col1
|
||||||
|
let (_d0, p0) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
||||||
|
let (_d1, p1) = make_int_matrix(&[&[10, 20], &[20, 10]]);
|
||||||
|
|
||||||
|
let partitioned = LayeredStore::new(vec![
|
||||||
|
LayeredStore::new(vec![p0]),
|
||||||
|
LayeredStore::new(vec![p1]),
|
||||||
|
]);
|
||||||
|
|
||||||
|
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5, 10, 20], &[5, 4, 3, 2, 1, 20, 10]]);
|
||||||
|
let expected = CountPartials::bray_dist_matrix(&mf);
|
||||||
|
let got = CountPartials::bray_dist_matrix(&partitioned);
|
||||||
|
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "partitioned bray [0,1]");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── BitPartials ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn layered_jaccard_matches_combined() {
|
||||||
|
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
|
||||||
|
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[true, false]]);
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
|
||||||
|
let (_df, mf) = make_bit_matrix(&[
|
||||||
|
&[true, false, true, true],
|
||||||
|
&[false, true, true, false],
|
||||||
|
]);
|
||||||
|
let expected = BitPartials::jaccard_dist_matrix(&mf);
|
||||||
|
let got = BitPartials::jaccard_dist_matrix(&store);
|
||||||
|
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "jaccard [0,1]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn layered_hamming_matches_combined() {
|
||||||
|
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
|
||||||
|
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[false, false]]);
|
||||||
|
let store = LayeredStore::new(vec![m0, m1]);
|
||||||
|
|
||||||
|
let (_df, mf) = make_bit_matrix(&[
|
||||||
|
&[true, false, true, true],
|
||||||
|
&[false, true, false, false],
|
||||||
|
]);
|
||||||
|
let expected = BitPartials::hamming_dist_matrix(&mf);
|
||||||
|
let got = BitPartials::hamming_dist_matrix(&store);
|
||||||
|
assert_eq!(got[[0, 1]], expected[[0, 1]], "hamming [0,1]");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,381 +0,0 @@
|
|||||||
use super::*;
|
|
||||||
use obicompactvec::{
|
|
||||||
PersistentBitMatrix, PersistentBitMatrixBuilder,
|
|
||||||
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
|
|
||||||
};
|
|
||||||
use tempfile::tempdir;
|
|
||||||
|
|
||||||
fn make_int_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
|
|
||||||
let n = cols.first().map_or(0, |c| c.len());
|
|
||||||
let dir = tempdir().unwrap();
|
|
||||||
let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
|
|
||||||
for &col in cols {
|
|
||||||
let mut cb = b.add_col().unwrap();
|
|
||||||
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
|
|
||||||
cb.close().unwrap();
|
|
||||||
}
|
|
||||||
b.close().unwrap();
|
|
||||||
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
||||||
(dir, m)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_bit_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
|
|
||||||
let n = cols.first().map_or(0, |c| c.len());
|
|
||||||
let dir = tempdir().unwrap();
|
|
||||||
let mut b = PersistentBitMatrixBuilder::new(n, &dir.path().join("presence")).unwrap();
|
|
||||||
for &col in cols {
|
|
||||||
let mut cb = b.add_col().unwrap();
|
|
||||||
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
|
|
||||||
cb.close().unwrap();
|
|
||||||
}
|
|
||||||
b.close().unwrap();
|
|
||||||
let m = PersistentBitMatrix::open(dir.path()).unwrap();
|
|
||||||
(dir, m)
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── ColumnWeights ─────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn col_weights_sums_across_layers() {
|
|
||||||
// layer 0: col0=[1,2], col1=[3,4] → weights [3, 7]
|
|
||||||
// layer 1: col0=[10,0], col1=[0,10] → weights [10, 10]
|
|
||||||
// combined: [13, 17]
|
|
||||||
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[3, 4]]);
|
|
||||||
let (_d1, m1) = make_int_matrix(&[&[10, 0], &[0, 10]]);
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
let w = store.col_weights();
|
|
||||||
assert_eq!(w[0], 13);
|
|
||||||
assert_eq!(w[1], 17);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn col_weights_bit_sums_across_layers() {
|
|
||||||
// layer 0: col0=[T,F,T], col1=[F,T,T] → counts [2, 2]
|
|
||||||
// layer 1: col0=[F,F,T], col1=[T,T,F] → counts [1, 2]
|
|
||||||
// combined: [3, 4]
|
|
||||||
let (_d0, m0) = make_bit_matrix(&[&[true, false, true], &[false, true, true]]);
|
|
||||||
let (_d1, m1) = make_bit_matrix(&[&[false, false, true], &[true, true, false]]);
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
let w = store.col_weights();
|
|
||||||
assert_eq!(w[0], 3);
|
|
||||||
assert_eq!(w[1], 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── CountPartials — layered (one partition) ───────────────────────────────
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn layered_bray_matches_combined() {
|
|
||||||
// Split [1,2,3,4,5] across two layers; bray dist should equal direct computation
|
|
||||||
// on [1,2,3,4,5] for each column pair.
|
|
||||||
// col0=[1,2,3,4,5], col1=[5,4,3,2,1]
|
|
||||||
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]); // slots 0-1
|
|
||||||
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]); // slots 2-4
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
|
|
||||||
// direct on full data
|
|
||||||
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
|
||||||
let expected = CountPartials::bray_dist_matrix(&mf);
|
|
||||||
let got = CountPartials::bray_dist_matrix(&store);
|
|
||||||
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "bray [0,1]");
|
|
||||||
assert!((got[[1, 0]] - expected[[1, 0]]).abs() < 1e-12, "bray [1,0]");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn layered_relfreq_bray_matches_combined() {
|
|
||||||
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]);
|
|
||||||
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
|
||||||
let expected = CountPartials::relfreq_bray_dist_matrix(&mf);
|
|
||||||
let got = CountPartials::relfreq_bray_dist_matrix(&store);
|
|
||||||
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "relfreq_bray [0,1]");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn layered_euclidean_matches_combined() {
|
|
||||||
let (_d0, m0) = make_int_matrix(&[&[3, 0], &[0, 4]]);
|
|
||||||
let (_d1, m1) = make_int_matrix(&[&[1, 1], &[2, 2]]);
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 4, 2, 2]]);
|
|
||||||
let expected = CountPartials::euclidean_dist_matrix(&mf);
|
|
||||||
let got = CountPartials::euclidean_dist_matrix(&store);
|
|
||||||
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "euclidean [0,1]");
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── CountPartials — partitioned (LayeredStore<LayeredStore<_>>) ───────────
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partitioned_bray_matches_combined() {
|
|
||||||
// partition 0: slots [1,2,3,4,5] col0 vs col1
|
|
||||||
// partition 1: slots [10,20] col0 vs col1
|
|
||||||
let (_d0, p0) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
|
||||||
let (_d1, p1) = make_int_matrix(&[&[10, 20], &[20, 10]]);
|
|
||||||
|
|
||||||
let partitioned = LayeredStore::new(vec![
|
|
||||||
LayeredStore::new(vec![p0]),
|
|
||||||
LayeredStore::new(vec![p1]),
|
|
||||||
]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5, 10, 20], &[5, 4, 3, 2, 1, 20, 10]]);
|
|
||||||
let expected = CountPartials::bray_dist_matrix(&mf);
|
|
||||||
let got = CountPartials::bray_dist_matrix(&partitioned);
|
|
||||||
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "partitioned bray [0,1]");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partitioned_threshold_jaccard_off_diagonal_is_pairwise() {
|
|
||||||
// 3 genomes, 2 partitions, 1 layer each — mirrors distance.rs's
|
|
||||||
// LayeredStore<LayeredStore<PersistentCompactIntMatrix>> shape.
|
|
||||||
// partition 0: col0=[3,0], col1=[0,3], col2=[3,3]
|
|
||||||
// partition 1: col0=[1,1], col1=[1,0], col2=[0,1]
|
|
||||||
let (_d0, p0) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
|
|
||||||
let (_d1, p1) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
|
|
||||||
|
|
||||||
let partitioned = LayeredStore::new(vec![
|
|
||||||
LayeredStore::new(vec![p0]),
|
|
||||||
LayeredStore::new(vec![p1]),
|
|
||||||
]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 3, 1, 0], &[3, 3, 0, 1]]);
|
|
||||||
let threshold = 1u32;
|
|
||||||
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
|
|
||||||
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
|
|
||||||
|
|
||||||
let n = 3;
|
|
||||||
for i in 0..n {
|
|
||||||
for j in 0..n {
|
|
||||||
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
|
|
||||||
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partitioned_threshold_jaccard_packed_off_diagonal_is_pairwise() {
|
|
||||||
// Same as `partitioned_threshold_jaccard_off_diagonal_is_pairwise` but
|
|
||||||
// each partition matrix is packed into a single .pcmx file first —
|
|
||||||
// the on-disk format actually used in production after `pack_matrices`.
|
|
||||||
use obicompactvec::pack_compact_int_matrix;
|
|
||||||
|
|
||||||
let (d0, _p0) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
|
|
||||||
pack_compact_int_matrix(&d0.path().join("counts")).unwrap();
|
|
||||||
let p0 = PersistentCompactIntMatrix::open(d0.path()).unwrap();
|
|
||||||
|
|
||||||
let (d1, _p1) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
|
|
||||||
pack_compact_int_matrix(&d1.path().join("counts")).unwrap();
|
|
||||||
let p1 = PersistentCompactIntMatrix::open(d1.path()).unwrap();
|
|
||||||
|
|
||||||
let partitioned = LayeredStore::new(vec![
|
|
||||||
LayeredStore::new(vec![p0]),
|
|
||||||
LayeredStore::new(vec![p1]),
|
|
||||||
]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 3, 1, 0], &[3, 3, 0, 1]]);
|
|
||||||
let threshold = 1u32;
|
|
||||||
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
|
|
||||||
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
|
|
||||||
|
|
||||||
let n = 3;
|
|
||||||
for i in 0..n {
|
|
||||||
for j in 0..n {
|
|
||||||
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
|
|
||||||
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partitioned_multilayer_threshold_jaccard_off_diagonal_is_pairwise() {
|
|
||||||
// 2 partitions, 2 layers each — the shape production indexes actually
|
|
||||||
// have (MPHF collision layers within a partition).
|
|
||||||
// partition 0, layer 0: col0=[3,0], col1=[0,3], col2=[3,3]
|
|
||||||
// partition 0, layer 1: col0=[2,0], col1=[0,0], col2=[2,0]
|
|
||||||
// partition 1, layer 0: col0=[1,1], col1=[1,0], col2=[0,1]
|
|
||||||
// partition 1, layer 1: col0=[0,5], col1=[5,5], col2=[0,0]
|
|
||||||
let (_d0a, p0a) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
|
|
||||||
let (_d0b, p0b) = make_int_matrix(&[&[2, 0], &[0, 0], &[2, 0]]);
|
|
||||||
let (_d1a, p1a) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
|
|
||||||
let (_d1b, p1b) = make_int_matrix(&[&[0, 5], &[5, 5], &[0, 0]]);
|
|
||||||
|
|
||||||
let partitioned = LayeredStore::new(vec![
|
|
||||||
LayeredStore::new(vec![p0a, p0b]),
|
|
||||||
LayeredStore::new(vec![p1a, p1b]),
|
|
||||||
]);
|
|
||||||
|
|
||||||
// Flattened equivalent: concatenate every layer's slots into one matrix.
|
|
||||||
let (_df, mf) = make_int_matrix(&[
|
|
||||||
&[3, 0, 2, 0, 1, 1, 0, 5],
|
|
||||||
&[0, 3, 0, 0, 1, 0, 5, 5],
|
|
||||||
&[3, 3, 2, 0, 0, 1, 0, 0],
|
|
||||||
]);
|
|
||||||
let threshold = 1u32;
|
|
||||||
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
|
|
||||||
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
|
|
||||||
|
|
||||||
let n = 3;
|
|
||||||
for i in 0..n {
|
|
||||||
for j in 0..n {
|
|
||||||
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
|
|
||||||
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── BitPartials ───────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn layered_jaccard_matches_combined() {
|
|
||||||
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
|
|
||||||
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[true, false]]);
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_bit_matrix(&[
|
|
||||||
&[true, false, true, true],
|
|
||||||
&[false, true, true, false],
|
|
||||||
]);
|
|
||||||
let expected = BitPartials::jaccard_dist_matrix(&mf);
|
|
||||||
let got = BitPartials::jaccard_dist_matrix(&store);
|
|
||||||
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "jaccard [0,1]");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn layered_hamming_matches_combined() {
|
|
||||||
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
|
|
||||||
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[false, false]]);
|
|
||||||
let store = LayeredStore::new(vec![m0, m1]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_bit_matrix(&[
|
|
||||||
&[true, false, true, true],
|
|
||||||
&[false, true, false, false],
|
|
||||||
]);
|
|
||||||
let expected = BitPartials::hamming_dist_matrix(&mf);
|
|
||||||
let got = BitPartials::hamming_dist_matrix(&store);
|
|
||||||
assert_eq!(got[[0, 1]], expected[[0, 1]], "hamming [0,1]");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partitioned_bit_jaccard_off_diagonal_is_pairwise() {
|
|
||||||
// Same shape as the count-based `partitioned_multilayer_threshold_jaccard_*`
|
|
||||||
// tests, but for the presence/bit path (`with_counts = false` — what
|
|
||||||
// `all_specifics` actually uses in production).
|
|
||||||
// 4 genomes, 3 partitions, 2 layers in the last one.
|
|
||||||
let (_d0, p0) = make_bit_matrix(&[
|
|
||||||
&[true, false, true],
|
|
||||||
&[false, true, true],
|
|
||||||
&[true, true, false],
|
|
||||||
&[false, false, true],
|
|
||||||
]);
|
|
||||||
let (_d1, p1) = make_bit_matrix(&[
|
|
||||||
&[true, true],
|
|
||||||
&[false, true],
|
|
||||||
&[true, false],
|
|
||||||
&[true, true],
|
|
||||||
]);
|
|
||||||
let (_d2a, p2a) = make_bit_matrix(&[
|
|
||||||
&[false, true],
|
|
||||||
&[true, true],
|
|
||||||
&[false, false],
|
|
||||||
&[true, false],
|
|
||||||
]);
|
|
||||||
let (_d2b, p2b) = make_bit_matrix(&[
|
|
||||||
&[true],
|
|
||||||
&[false],
|
|
||||||
&[true],
|
|
||||||
&[true],
|
|
||||||
]);
|
|
||||||
|
|
||||||
let partitioned = LayeredStore::new(vec![
|
|
||||||
LayeredStore::new(vec![p0]),
|
|
||||||
LayeredStore::new(vec![p1]),
|
|
||||||
LayeredStore::new(vec![p2a, p2b]),
|
|
||||||
]);
|
|
||||||
|
|
||||||
// Flattened equivalent: concatenate every partition/layer's slots.
|
|
||||||
let (_df, mf) = make_bit_matrix(&[
|
|
||||||
&[true, false, true, true, true, false, true, true],
|
|
||||||
&[false, true, true, false, true, true, true, false],
|
|
||||||
&[true, true, false, true, false, false, false, true],
|
|
||||||
&[false, false, true, true, true, true, false, true],
|
|
||||||
]);
|
|
||||||
|
|
||||||
let (inter_p, union_p) = BitPartials::partial_jaccard(&partitioned);
|
|
||||||
let (inter_f, union_f) = BitPartials::partial_jaccard(&mf);
|
|
||||||
|
|
||||||
let n = 4;
|
|
||||||
for i in 0..n {
|
|
||||||
for j in 0..n {
|
|
||||||
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
|
|
||||||
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn partitioned_bit_jaccard_packed_off_diagonal_is_pairwise() {
|
|
||||||
// Same as `partitioned_bit_jaccard_off_diagonal_is_pairwise` but every
|
|
||||||
// partition's presence matrix is packed into a single .pbmx file —
|
|
||||||
// the on-disk format actually used in production after `pack_matrices`.
|
|
||||||
use obicompactvec::pack_bit_matrix;
|
|
||||||
|
|
||||||
let (d0, _p0) = make_bit_matrix(&[
|
|
||||||
&[true, false, true],
|
|
||||||
&[false, true, true],
|
|
||||||
&[true, true, false],
|
|
||||||
&[false, false, true],
|
|
||||||
]);
|
|
||||||
pack_bit_matrix(&d0.path().join("presence")).unwrap();
|
|
||||||
let p0 = PersistentBitMatrix::open(d0.path()).unwrap();
|
|
||||||
|
|
||||||
let (d1, _p1) = make_bit_matrix(&[
|
|
||||||
&[true, true],
|
|
||||||
&[false, true],
|
|
||||||
&[true, false],
|
|
||||||
&[true, true],
|
|
||||||
]);
|
|
||||||
pack_bit_matrix(&d1.path().join("presence")).unwrap();
|
|
||||||
let p1 = PersistentBitMatrix::open(d1.path()).unwrap();
|
|
||||||
|
|
||||||
let (d2a, _p2a) = make_bit_matrix(&[
|
|
||||||
&[false, true],
|
|
||||||
&[true, true],
|
|
||||||
&[false, false],
|
|
||||||
&[true, false],
|
|
||||||
]);
|
|
||||||
pack_bit_matrix(&d2a.path().join("presence")).unwrap();
|
|
||||||
let p2a = PersistentBitMatrix::open(d2a.path()).unwrap();
|
|
||||||
|
|
||||||
let (d2b, _p2b) = make_bit_matrix(&[
|
|
||||||
&[true],
|
|
||||||
&[false],
|
|
||||||
&[true],
|
|
||||||
&[true],
|
|
||||||
]);
|
|
||||||
pack_bit_matrix(&d2b.path().join("presence")).unwrap();
|
|
||||||
let p2b = PersistentBitMatrix::open(d2b.path()).unwrap();
|
|
||||||
|
|
||||||
let partitioned = LayeredStore::new(vec![
|
|
||||||
LayeredStore::new(vec![p0]),
|
|
||||||
LayeredStore::new(vec![p1]),
|
|
||||||
LayeredStore::new(vec![p2a, p2b]),
|
|
||||||
]);
|
|
||||||
|
|
||||||
let (_df, mf) = make_bit_matrix(&[
|
|
||||||
&[true, false, true, true, true, false, true, true],
|
|
||||||
&[false, true, true, false, true, true, true, false],
|
|
||||||
&[true, true, false, true, false, false, false, true],
|
|
||||||
&[false, false, true, true, true, true, false, true],
|
|
||||||
]);
|
|
||||||
|
|
||||||
let (inter_p, union_p) = BitPartials::partial_jaccard(&partitioned);
|
|
||||||
let (inter_f, union_f) = BitPartials::partial_jaccard(&mf);
|
|
||||||
|
|
||||||
let n = 4;
|
|
||||||
for i in 0..n {
|
|
||||||
for j in 0..n {
|
|
||||||
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
|
|
||||||
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -7,13 +7,7 @@ edition = "2024"
|
|||||||
obikseq = { path = "../obikseq" }
|
obikseq = { path = "../obikseq" }
|
||||||
obikrope = { path = "../obikrope" }
|
obikrope = { path = "../obikrope" }
|
||||||
obiread = { path = "../obiread" }
|
obiread = { path = "../obiread" }
|
||||||
obikentropy = { path = "../obikentropy" }
|
|
||||||
lazy_static = "1.5.0"
|
lazy_static = "1.5.0"
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
||||||
criterion2 = { version = "3", features = ["cargo_bench_support"] }
|
|
||||||
|
|
||||||
[[bench]]
|
|
||||||
name = "superkmer_stream"
|
|
||||||
harness = false
|
|
||||||
|
|||||||
@@ -1,58 +0,0 @@
|
|||||||
//! Throughput of the streaming superkmer pipeline (`RollingStat`'s hot path:
|
|
||||||
//! minimizer selection + entropy tracking fused in a single pass).
|
|
||||||
//!
|
|
||||||
//! Reference point for the `obikentropy` extraction: the entropy bookkeeping
|
|
||||||
//! that used to live inline in `RollingStat` was pulled out into a composed
|
|
||||||
//! `EntropyTracker`. This benchmark is run before and after that change to
|
|
||||||
//! confirm no regression.
|
|
||||||
|
|
||||||
use criterion::{Criterion, Throughput, criterion_group, criterion_main};
|
|
||||||
use obikrope::Rope;
|
|
||||||
use obiskbuilder::SuperKmerIter;
|
|
||||||
|
|
||||||
const K: usize = 21;
|
|
||||||
const M: usize = 9;
|
|
||||||
const LEVEL_MAX: usize = 6;
|
|
||||||
const THETA: f64 = 0.7;
|
|
||||||
const SEQ_LEN: usize = 200_000;
|
|
||||||
|
|
||||||
/// Deterministic pseudo-random ACGT sequence — high enough complexity that
|
|
||||||
/// the entropy filter rarely rejects, so the bench stays on the steady-state
|
|
||||||
/// path rather than repeatedly resetting.
|
|
||||||
fn make_sequence(len: usize) -> Vec<u8> {
|
|
||||||
let mut state: u64 = 0x9E3779B97F4A7C15;
|
|
||||||
(0..len)
|
|
||||||
.map(|_| {
|
|
||||||
state ^= state << 13;
|
|
||||||
state ^= state >> 7;
|
|
||||||
state ^= state << 17;
|
|
||||||
b"ACGT"[(state % 4) as usize]
|
|
||||||
})
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_rope(seq: &[u8]) -> Rope {
|
|
||||||
let mut rope = Rope::new(None);
|
|
||||||
rope.push(seq.to_vec());
|
|
||||||
rope
|
|
||||||
}
|
|
||||||
|
|
||||||
fn bench_build_superkmers(c: &mut Criterion) {
|
|
||||||
obikseq::set_k(K);
|
|
||||||
obikseq::set_m(M);
|
|
||||||
|
|
||||||
let seq = make_sequence(SEQ_LEN);
|
|
||||||
let rope = make_rope(&seq);
|
|
||||||
|
|
||||||
let mut group = c.benchmark_group("build_superkmers");
|
|
||||||
group.throughput(Throughput::Bytes(SEQ_LEN as u64));
|
|
||||||
group.bench_function("stream", |b| {
|
|
||||||
b.iter(|| {
|
|
||||||
SuperKmerIter::new(std::hint::black_box(&rope), K, LEVEL_MAX, THETA).count()
|
|
||||||
});
|
|
||||||
});
|
|
||||||
group.finish();
|
|
||||||
}
|
|
||||||
|
|
||||||
criterion_group!(benches, bench_build_superkmers);
|
|
||||||
criterion_main!(benches);
|
|
||||||
@@ -4,6 +4,57 @@ use std::path::PathBuf;
|
|||||||
const K_MAX: usize = 32;
|
const K_MAX: usize = 32;
|
||||||
const WS_MAX: usize = 6;
|
const WS_MAX: usize = 6;
|
||||||
|
|
||||||
|
fn normalize_circular(kmer: u64, ws: usize) -> u64 {
|
||||||
|
let mask = (1u64 << (ws * 2)) - 1;
|
||||||
|
let mut canonical = kmer & mask;
|
||||||
|
let mut current = canonical;
|
||||||
|
for _ in 0..ws - 1 {
|
||||||
|
let top = (current >> ((ws - 1) * 2)) & 3;
|
||||||
|
current = ((current << 2) | top) & mask;
|
||||||
|
if current < canonical {
|
||||||
|
canonical = current;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
canonical
|
||||||
|
}
|
||||||
|
|
||||||
|
fn revcomp_raw(x: u64, k: usize) -> u64 {
|
||||||
|
let x = !x;
|
||||||
|
let x = x.swap_bytes();
|
||||||
|
let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
|
||||||
|
let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
|
||||||
|
x << (64 - 2 * k)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_normalized_kmer(k: usize) -> Vec<u64> {
|
||||||
|
let n = 1usize << (k * 2);
|
||||||
|
let shift = 64 - k * 2;
|
||||||
|
let mut result = vec![0u64; n];
|
||||||
|
for i in 0..n {
|
||||||
|
let la = (i as u64) << shift;
|
||||||
|
let ra = i as u64;
|
||||||
|
let rc_ra = revcomp_raw(la, k) >> shift;
|
||||||
|
let circ = normalize_circular(ra, k);
|
||||||
|
let circ_rc = normalize_circular(rc_ra, k);
|
||||||
|
result[i] = if circ < circ_rc { circ } else { circ_rc };
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_ln_class(norm: &[u64]) -> Vec<f64> {
|
||||||
|
let n = norm.len();
|
||||||
|
let mut sizes = vec![0u32; n];
|
||||||
|
for &c in norm {
|
||||||
|
sizes[c as usize] += 1;
|
||||||
|
}
|
||||||
|
norm.iter()
|
||||||
|
.map(|&c| {
|
||||||
|
let s = sizes[c as usize];
|
||||||
|
if s > 0 { (s as f64).ln() } else { 0.0 }
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
fn build_n_log_n() -> [f64; K_MAX + 1] {
|
fn build_n_log_n() -> [f64; K_MAX + 1] {
|
||||||
let mut t = [0.0f64; K_MAX + 1];
|
let mut t = [0.0f64; K_MAX + 1];
|
||||||
for n in 1..=K_MAX {
|
for n in 1..=K_MAX {
|
||||||
@@ -12,9 +63,6 @@ fn build_n_log_n() -> [f64; K_MAX + 1] {
|
|||||||
t
|
t
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Max achievable entropy over `4^ws` raw sub-words given only `nwords`
|
|
||||||
/// observations (most-uniform integer partition), per
|
|
||||||
/// `docmd/theory/entropy.md`.
|
|
||||||
fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
|
fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
|
||||||
let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
|
let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
|
||||||
for k in 2..=K_MAX {
|
for k in 2..=K_MAX {
|
||||||
@@ -77,6 +125,13 @@ fn main() {
|
|||||||
let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
|
let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
|
||||||
let mut out = String::new();
|
let mut out = String::new();
|
||||||
|
|
||||||
|
for k in 1..=6usize {
|
||||||
|
let n = 1usize << (k * 2);
|
||||||
|
let norm = build_normalized_kmer(k);
|
||||||
|
let ln_class = build_ln_class(&norm);
|
||||||
|
emit_f64_1d(&mut out, &format!("LN_CLASS{k}"), n, &ln_class);
|
||||||
|
}
|
||||||
|
|
||||||
let n_log_n = build_n_log_n();
|
let n_log_n = build_n_log_n();
|
||||||
emit_f64_1d(&mut out, "N_LOG_N", K_MAX + 1, &n_log_n);
|
emit_f64_1d(&mut out, "N_LOG_N", K_MAX + 1, &n_log_n);
|
||||||
|
|
||||||
@@ -86,5 +141,5 @@ fn main() {
|
|||||||
let log_nwords = build_log_nwords();
|
let log_nwords = build_log_nwords();
|
||||||
emit_f64_2d(&mut out, "LOG_NWORDS", K_MAX + 1, WS_MAX + 1, &log_nwords);
|
emit_f64_2d(&mut out, "LOG_NWORDS", K_MAX + 1, WS_MAX + 1, &log_nwords);
|
||||||
|
|
||||||
fs::write(out_dir.join("entropy_tables.rs"), out).unwrap();
|
fs::write(out_dir.join("ln_class_tables.rs"), out).unwrap();
|
||||||
}
|
}
|
||||||
@@ -0,0 +1,108 @@
|
|||||||
|
pub(crate) const NORMK1: [u64; 4] = build_normalized_kmer::<4>();
|
||||||
|
pub(crate) const NORMK2: [u64; 16] = build_normalized_kmer::<16>();
|
||||||
|
pub(crate) const NORMK3: [u64; 64] = build_normalized_kmer::<64>();
|
||||||
|
pub(crate) const NORMK4: [u64; 256] = build_normalized_kmer::<256>();
|
||||||
|
pub(crate) const NORMK5: [u64; 1024] = build_normalized_kmer::<1024>();
|
||||||
|
pub(crate) const NORMK6: [u64; 4096] = build_normalized_kmer::<4096>();
|
||||||
|
|
||||||
|
include!(concat!(env!("OUT_DIR"), "/ln_class_tables.rs"));
|
||||||
|
|
||||||
|
const fn normalize_circular(kmer: u64, ws: usize) -> u64 {
|
||||||
|
let mask = (1u64 << (ws * 2)) - 1;
|
||||||
|
let mut canonical = kmer & mask;
|
||||||
|
let mut current = canonical;
|
||||||
|
let mut i = 0;
|
||||||
|
while i < (ws - 1) {
|
||||||
|
let top = (current >> ((ws - 1) * 2)) & 3;
|
||||||
|
current = ((current << 2) | top) & mask;
|
||||||
|
if current < canonical {
|
||||||
|
canonical = current;
|
||||||
|
}
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
canonical
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn build_normalized_kmer<const N: usize>() -> [u64; N] {
|
||||||
|
let mut result = [0u64; N];
|
||||||
|
let k = k_from_n::<N>();
|
||||||
|
let shift = 64 - k * 2;
|
||||||
|
let mut i = 0;
|
||||||
|
while i < N {
|
||||||
|
let la = (i as u64) << shift;
|
||||||
|
let ra = i as u64;
|
||||||
|
let rc_ra = revcomp_raw(la, k) >> shift;
|
||||||
|
let circ = normalize_circular(ra, k);
|
||||||
|
let circ_rc = normalize_circular(rc_ra, k);
|
||||||
|
result[i] = if circ < circ_rc { circ } else { circ_rc };
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn revcomp_raw(x: u64, k: usize) -> u64 {
|
||||||
|
let x = !x;
|
||||||
|
let x = x.swap_bytes();
|
||||||
|
let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
|
||||||
|
let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
|
||||||
|
x << (64 - 2 * k)
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn k_from_n<const N: usize>() -> usize {
|
||||||
|
match N {
|
||||||
|
4 => 1,
|
||||||
|
16 => 2,
|
||||||
|
64 => 3,
|
||||||
|
256 => 4,
|
||||||
|
1024 => 5,
|
||||||
|
4096 => 6,
|
||||||
|
_ => panic!("N must be a power of 4"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) const WS_MAX: usize = 6;
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
pub(crate) const fn n_log_n(n: usize) -> f64 {
|
||||||
|
N_LOG_N[n]
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
pub(crate) const fn emax(k: usize, ws: usize) -> f64 {
|
||||||
|
EMAX[k][ws]
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 {
|
||||||
|
LOG_NWORDS[k][ws]
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
pub(crate) const fn entropy_norm_kmer<const LEFT: bool, const K: usize>(kmer: u64) -> u64 {
|
||||||
|
const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
|
||||||
|
const NORM: [&[u64]; 7] = [&[], &NORMK1, &NORMK2, &NORMK3, &NORMK4, &NORMK5, &NORMK6];
|
||||||
|
|
||||||
|
let shift = SHIFT[K];
|
||||||
|
let ra = if LEFT { kmer >> shift } else { kmer };
|
||||||
|
let canonical_ra = NORM[K][ra as usize];
|
||||||
|
if LEFT {
|
||||||
|
canonical_ra << shift
|
||||||
|
} else {
|
||||||
|
canonical_ra
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
pub(crate) const fn ln_class_size<const LEFT: bool, const K: usize>(kmer: u64) -> f64 {
|
||||||
|
const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
|
||||||
|
let ra = if LEFT { kmer >> SHIFT[K] } else { kmer };
|
||||||
|
match K {
|
||||||
|
1 => LN_CLASS1[ra as usize],
|
||||||
|
2 => LN_CLASS2[ra as usize],
|
||||||
|
3 => LN_CLASS3[ra as usize],
|
||||||
|
4 => LN_CLASS4[ra as usize],
|
||||||
|
5 => LN_CLASS5[ra as usize],
|
||||||
|
6 => LN_CLASS6[ra as usize],
|
||||||
|
_ => panic!("k must be 1..=6"),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -149,5 +149,157 @@ impl Iterator for SuperKmerIter<'_> {
|
|||||||
// ── tests ─────────────────────────────────────────────────────────────────────
|
// ── tests ─────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
#[path = "tests/iter.rs"]
|
mod tests {
|
||||||
mod tests;
|
use super::*;
|
||||||
|
use obikrope::Rope;
|
||||||
|
|
||||||
|
fn setup() {
|
||||||
|
obikseq::params::set_k(K);
|
||||||
|
obikseq::params::set_m(5);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn make_rope(data: &[u8]) -> Rope {
|
||||||
|
let mut r = Rope::new(None);
|
||||||
|
r.push(data.to_vec());
|
||||||
|
r
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run_nofilter(data: &[u8], k: usize) -> Vec<Vec<u8>> {
|
||||||
|
let rope = make_rope(data);
|
||||||
|
SuperKmerIter::new(&rope, k, 1, 0.0)
|
||||||
|
.map(|rsk| rsk.superkmer().to_ascii())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
// k=11, m=5 — valeurs minimales du projet (k ∈ [11,31])
|
||||||
|
const K: usize = 11;
|
||||||
|
|
||||||
|
/// Collect the set of canonical k-mers from a raw ASCII sequence (no NUL).
|
||||||
|
fn direct_canonical_kmers(seq: &[u8]) -> std::collections::HashSet<Vec<u8>> {
|
||||||
|
(0..seq.len().saturating_sub(K - 1))
|
||||||
|
.map(|i| obikseq::SuperKmer::from_ascii(&seq[i..i + K]).to_ascii())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Collect the set of canonical k-mers emitted by SuperKmerIter over a rope.
|
||||||
|
fn iter_canonical_kmers(rope: &Rope) -> std::collections::HashSet<Vec<u8>> {
|
||||||
|
SuperKmerIter::new(rope, K, 1, 0.0)
|
||||||
|
.flat_map(|rsk| {
|
||||||
|
rsk.superkmer()
|
||||||
|
.iter_canonical_kmers()
|
||||||
|
.map(|km| km.to_ascii())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn coverage_single_segment() {
|
||||||
|
setup();
|
||||||
|
let seq = b"ACGTACGTACGTACGTACGT";
|
||||||
|
let rope = make_rope(&[seq.as_ref(), b"\x00"].concat());
|
||||||
|
let direct = direct_canonical_kmers(seq);
|
||||||
|
let from_iter = iter_canonical_kmers(&rope);
|
||||||
|
let missing: Vec<_> = direct.difference(&from_iter).collect();
|
||||||
|
assert!(
|
||||||
|
missing.is_empty(),
|
||||||
|
"k-mers perdus dans segment unique : {missing:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn coverage_two_segments() {
|
||||||
|
setup();
|
||||||
|
let seg1 = b"ACGTACGTACGTACGTACGT";
|
||||||
|
let seg2 = b"TGCATGCATGCATGCATGCA";
|
||||||
|
let rope = make_rope(&[seg1.as_ref(), b"\x00", seg2.as_ref(), b"\x00"].concat());
|
||||||
|
let mut direct = direct_canonical_kmers(seg1);
|
||||||
|
direct.extend(direct_canonical_kmers(seg2));
|
||||||
|
let from_iter = iter_canonical_kmers(&rope);
|
||||||
|
let missing: Vec<_> = direct.difference(&from_iter).collect();
|
||||||
|
assert!(
|
||||||
|
missing.is_empty(),
|
||||||
|
"k-mers perdus dans deux segments : {missing:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn coverage_minimizer_boundary() {
|
||||||
|
setup();
|
||||||
|
// sequence assez longue pour forcer plusieurs changements de minimiseur
|
||||||
|
let seq: Vec<u8> = (0..80).map(|i| b"ACGT"[i % 4]).collect();
|
||||||
|
let rope = make_rope(&[seq.as_slice(), b"\x00"].concat());
|
||||||
|
let direct = direct_canonical_kmers(&seq);
|
||||||
|
let from_iter = iter_canonical_kmers(&rope);
|
||||||
|
let missing: Vec<_> = direct.difference(&from_iter).collect();
|
||||||
|
assert!(
|
||||||
|
missing.is_empty(),
|
||||||
|
"k-mers perdus à la frontière de minimiseur : {missing:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn single_segment_one_superkmer() {
|
||||||
|
setup();
|
||||||
|
let out = run_nofilter(b"ACGTACGTACGTACGTACGT\x00", K);
|
||||||
|
assert!(!out.is_empty());
|
||||||
|
let total: Vec<u8> = out.into_iter().flatten().collect();
|
||||||
|
assert!(total.len() >= K);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn segment_shorter_than_k_emits_nothing() {
|
||||||
|
setup();
|
||||||
|
let out = run_nofilter(b"ACGTACGT\x00", K);
|
||||||
|
assert_eq!(out, Vec::<Vec<u8>>::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn empty_input_emits_nothing() {
|
||||||
|
setup();
|
||||||
|
let out = run_nofilter(b"", K);
|
||||||
|
assert_eq!(out, Vec::<Vec<u8>>::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn two_segments_both_emitted() {
|
||||||
|
setup();
|
||||||
|
let out = run_nofilter(b"ACGTACGTACGTACGT\x00TGCATGCATGCATGCA\x00", K);
|
||||||
|
assert!(!out.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn low_complexity_kmer_is_rejected() {
|
||||||
|
setup();
|
||||||
|
let out_pass = run_nofilter(b"AAAAAAAAAAAACGTACGTACGT\x00", K);
|
||||||
|
assert!(!out_pass.is_empty());
|
||||||
|
|
||||||
|
let rope = make_rope(b"AAAAAAAAAAAAAAAAAAAA\x00");
|
||||||
|
let out_reject: Vec<Vec<u8>> = SuperKmerIter::new(&rope, K, 6, 0.9)
|
||||||
|
.map(|rsk| rsk.superkmer().to_ascii())
|
||||||
|
.collect();
|
||||||
|
assert!(out_reject.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn multi_slice_rope() {
|
||||||
|
setup();
|
||||||
|
let data = b"ACGTACGTACGTACGTACGT\x00";
|
||||||
|
let mid = data.len() / 2;
|
||||||
|
let mut rope = Rope::new(None);
|
||||||
|
rope.push(data[..mid].to_vec());
|
||||||
|
rope.push(data[mid..].to_vec());
|
||||||
|
let out: Vec<Vec<u8>> = SuperKmerIter::new(&rope, K, 1, 0.0)
|
||||||
|
.map(|rsk| rsk.superkmer().to_ascii())
|
||||||
|
.collect();
|
||||||
|
assert!(!out.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn yields_minimizer_value() {
|
||||||
|
setup();
|
||||||
|
let rope = make_rope(b"ACGTACGTACGTACGTACGT\x00");
|
||||||
|
let results: Vec<RoutableSuperKmer> = SuperKmerIter::new(&rope, K, 1, 0.0).collect();
|
||||||
|
assert!(!results.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -10,8 +10,8 @@ pub mod stream_iter;
|
|||||||
mod scratch;
|
mod scratch;
|
||||||
|
|
||||||
pub(crate) mod encoding;
|
pub(crate) mod encoding;
|
||||||
#[allow(missing_docs)]
|
pub(crate) mod entropy_table;
|
||||||
pub mod rolling_stat;
|
pub(crate) mod rolling_stat;
|
||||||
|
|
||||||
pub use iter::SuperKmerIter;
|
pub use iter::SuperKmerIter;
|
||||||
pub use scratch::SuperKmerScratch;
|
pub use scratch::SuperKmerScratch;
|
||||||
|
|||||||
@@ -1,8 +1,8 @@
|
|||||||
use obikentropy::EntropyTracker;
|
|
||||||
use obikseq::kmer::{Minimizer, hash_kmer};
|
use obikseq::kmer::{Minimizer, hash_kmer};
|
||||||
use obikseq::params;
|
use obikseq::params;
|
||||||
|
|
||||||
use crate::encoding::encode_nuc;
|
use crate::encoding::encode_nuc;
|
||||||
|
use crate::entropy_table::{WS_MAX, emax, entropy_norm_kmer, ln_class_size, log_nwords, n_log_n};
|
||||||
|
|
||||||
// ── Stack-allocated ring buffer ───────────────────────────────────────────────
|
// ── Stack-allocated ring buffer ───────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -83,19 +83,33 @@ pub struct RollingStat {
|
|||||||
entropy_max_k: usize,
|
entropy_max_k: usize,
|
||||||
k: usize,
|
k: usize,
|
||||||
m: usize,
|
m: usize,
|
||||||
|
steady: bool,
|
||||||
rolling_k: u64,
|
rolling_k: u64,
|
||||||
rolling_rck: u64,
|
rolling_rck: u64,
|
||||||
k_mask: u64,
|
k_mask: u64,
|
||||||
m_mask: u64,
|
m_mask: u64,
|
||||||
received: usize,
|
received: usize,
|
||||||
|
|
||||||
// Minimizer selection state.
|
// Sliding-window queues — stack-allocated, capacity ≤ k ≤ 31.
|
||||||
|
k1q: Ring<u64, 32>,
|
||||||
|
k2q: Ring<u64, 32>,
|
||||||
|
k3q: Ring<u64, 32>,
|
||||||
|
k4q: Ring<u64, 32>,
|
||||||
|
k5q: Ring<u64, 32>,
|
||||||
|
k6q: Ring<u64, 32>,
|
||||||
minimier: Ring<MmerItem, 32>,
|
minimier: Ring<MmerItem, 32>,
|
||||||
|
|
||||||
// Entropy tracking, composed as a plain inline field so both concerns
|
// Frequency count arrays.
|
||||||
// update in the same streaming pass without being conflated in one
|
// Max count per cell ≤ k ≤ 31 → u8 is sufficient.
|
||||||
// struct — see `obikentropy::EntropyTracker`.
|
k1c: [u8; 4],
|
||||||
entropy: EntropyTracker,
|
k2c: [u8; 16],
|
||||||
|
k3c: [u8; 64],
|
||||||
|
k4c: [u8; 256],
|
||||||
|
k5c: [u8; 1024],
|
||||||
|
k6c: [u8; 4096],
|
||||||
|
|
||||||
|
sum_f_log_f: [f64; WS_MAX + 1],
|
||||||
|
sum_f_log_s: [f64; WS_MAX + 1],
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RollingStat {
|
impl RollingStat {
|
||||||
@@ -106,13 +120,27 @@ impl RollingStat {
|
|||||||
entropy_max_k,
|
entropy_max_k,
|
||||||
k,
|
k,
|
||||||
m,
|
m,
|
||||||
|
steady: false,
|
||||||
rolling_k: 0,
|
rolling_k: 0,
|
||||||
rolling_rck: 0,
|
rolling_rck: 0,
|
||||||
k_mask: (!0u64) >> (64 - k * 2),
|
k_mask: (!0u64) >> (64 - k * 2),
|
||||||
m_mask: (!0u64) >> (64 - m * 2),
|
m_mask: (!0u64) >> (64 - m * 2),
|
||||||
received: 0,
|
received: 0,
|
||||||
|
k1q: Ring::new(),
|
||||||
|
k2q: Ring::new(),
|
||||||
|
k3q: Ring::new(),
|
||||||
|
k4q: Ring::new(),
|
||||||
|
k5q: Ring::new(),
|
||||||
|
k6q: Ring::new(),
|
||||||
minimier: Ring::new(),
|
minimier: Ring::new(),
|
||||||
entropy: EntropyTracker::new(k),
|
k1c: [0; 4],
|
||||||
|
k2c: [0; 16],
|
||||||
|
k3c: [0; 64],
|
||||||
|
k4c: [0; 256],
|
||||||
|
k5c: [0; 1024],
|
||||||
|
k6c: [0; 4096],
|
||||||
|
sum_f_log_f: [0.0; WS_MAX + 1],
|
||||||
|
sum_f_log_s: [0.0; WS_MAX + 1],
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -120,9 +148,54 @@ impl RollingStat {
|
|||||||
self.rolling_k = 0;
|
self.rolling_k = 0;
|
||||||
self.rolling_rck = 0;
|
self.rolling_rck = 0;
|
||||||
self.received = 0;
|
self.received = 0;
|
||||||
|
self.steady = false;
|
||||||
|
|
||||||
|
// for i in self.k1q.iter() { self.k1c[i as usize] = 0; }
|
||||||
|
// for i in self.k2q.iter() { self.k2c[i as usize] = 0; }
|
||||||
|
// for i in self.k3q.iter() { self.k3c[i as usize] = 0; }
|
||||||
|
// for i in self.k4q.iter() { self.k4c[i as usize] = 0; }
|
||||||
|
// for i in self.k5q.iter() { self.k5c[i as usize] = 0; }
|
||||||
|
// for i in self.k6q.iter() { self.k6c[i as usize] = 0; }
|
||||||
|
|
||||||
|
self.k1c.fill(0);
|
||||||
|
self.k2c.fill(0);
|
||||||
|
self.k3c.fill(0);
|
||||||
|
self.k4c.fill(0);
|
||||||
|
self.k5c.fill(0);
|
||||||
|
self.k6c.fill(0);
|
||||||
|
|
||||||
|
self.k1q.clear();
|
||||||
|
self.k2q.clear();
|
||||||
|
self.k3q.clear();
|
||||||
|
self.k4q.clear();
|
||||||
|
self.k5q.clear();
|
||||||
|
self.k6q.clear();
|
||||||
self.minimier.clear();
|
self.minimier.clear();
|
||||||
self.entropy.reset();
|
|
||||||
|
self.sum_f_log_f = [0.0; WS_MAX + 1];
|
||||||
|
self.sum_f_log_s = [0.0; WS_MAX + 1];
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn update_sums_decrement<const K: usize>(
|
||||||
|
sum_f_log_f: &mut [f64; WS_MAX + 1],
|
||||||
|
sum_f_log_s: &mut [f64; WS_MAX + 1],
|
||||||
|
canonical: u64,
|
||||||
|
f: usize,
|
||||||
|
) {
|
||||||
|
sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f);
|
||||||
|
sum_f_log_s[K] -= ln_class_size::<false, K>(canonical);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn update_sums_increment<const K: usize>(
|
||||||
|
sum_f_log_f: &mut [f64; WS_MAX + 1],
|
||||||
|
sum_f_log_s: &mut [f64; WS_MAX + 1],
|
||||||
|
canonical: u64,
|
||||||
|
g: usize,
|
||||||
|
) {
|
||||||
|
sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g);
|
||||||
|
sum_f_log_s[K] += ln_class_size::<false, K>(canonical);
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn push(&mut self, nuc: u8) {
|
pub fn push(&mut self, nuc: u8) {
|
||||||
@@ -136,6 +209,13 @@ impl RollingStat {
|
|||||||
self.rolling_rck =
|
self.rolling_rck =
|
||||||
((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask;
|
((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask;
|
||||||
|
|
||||||
|
let canonical_k1 = entropy_norm_kmer::<false, 1>(self.rolling_k & 3);
|
||||||
|
let canonical_k2 = entropy_norm_kmer::<false, 2>(self.rolling_k & 15);
|
||||||
|
let canonical_k3 = entropy_norm_kmer::<false, 3>(self.rolling_k & 63);
|
||||||
|
let canonical_k4 = entropy_norm_kmer::<false, 4>(self.rolling_k & 255);
|
||||||
|
let canonical_k5 = entropy_norm_kmer::<false, 5>(self.rolling_k & 1023);
|
||||||
|
let canonical_k6 = entropy_norm_kmer::<false, 6>(self.rolling_k & 4095);
|
||||||
|
|
||||||
self.received += 1;
|
self.received += 1;
|
||||||
|
|
||||||
if self.received >= m {
|
if self.received >= m {
|
||||||
@@ -168,7 +248,153 @@ impl RollingStat {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
self.entropy.push(self.received, self.rolling_k);
|
if self.received > k {
|
||||||
|
let old1 = self.k1q.pop_front();
|
||||||
|
let f1 = self.k1c[old1 as usize] as usize;
|
||||||
|
Self::update_sums_decrement::<1>(
|
||||||
|
&mut self.sum_f_log_f,
|
||||||
|
&mut self.sum_f_log_s,
|
||||||
|
old1,
|
||||||
|
f1,
|
||||||
|
);
|
||||||
|
self.k1c[old1 as usize] -= 1;
|
||||||
|
|
||||||
|
let old2 = self.k2q.pop_front();
|
||||||
|
let f2 = self.k2c[old2 as usize] as usize;
|
||||||
|
Self::update_sums_decrement::<2>(
|
||||||
|
&mut self.sum_f_log_f,
|
||||||
|
&mut self.sum_f_log_s,
|
||||||
|
old2,
|
||||||
|
f2,
|
||||||
|
);
|
||||||
|
self.k2c[old2 as usize] -= 1;
|
||||||
|
|
||||||
|
let old3 = self.k3q.pop_front();
|
||||||
|
let f3 = self.k3c[old3 as usize] as usize;
|
||||||
|
Self::update_sums_decrement::<3>(
|
||||||
|
&mut self.sum_f_log_f,
|
||||||
|
&mut self.sum_f_log_s,
|
||||||
|
old3,
|
||||||
|
f3,
|
||||||
|
);
|
||||||
|
self.k3c[old3 as usize] -= 1;
|
||||||
|
|
||||||
|
let old4 = self.k4q.pop_front();
|
||||||
|
let f4 = self.k4c[old4 as usize] as usize;
|
||||||
|
Self::update_sums_decrement::<4>(
|
||||||
|
&mut self.sum_f_log_f,
|
||||||
|
&mut self.sum_f_log_s,
|
||||||
|
old4,
|
||||||
|
f4,
|
||||||
|
);
|
||||||
|
self.k4c[old4 as usize] -= 1;
|
||||||
|
|
||||||
|
let old5 = self.k5q.pop_front();
|
||||||
|
let f5 = self.k5c[old5 as usize] as usize;
|
||||||
|
Self::update_sums_decrement::<5>(
|
||||||
|
&mut self.sum_f_log_f,
|
||||||
|
&mut self.sum_f_log_s,
|
||||||
|
old5,
|
||||||
|
f5,
|
||||||
|
);
|
||||||
|
self.k5c[old5 as usize] -= 1;
|
||||||
|
|
||||||
|
let old6 = self.k6q.pop_front();
|
||||||
|
let f6 = self.k6c[old6 as usize] as usize;
|
||||||
|
Self::update_sums_decrement::<6>(
|
||||||
|
&mut self.sum_f_log_f,
|
||||||
|
&mut self.sum_f_log_s,
|
||||||
|
old6,
|
||||||
|
f6,
|
||||||
|
);
|
||||||
|
self.k6c[old6 as usize] -= 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.steady {
|
||||||
|
let g1 = self.k1c[canonical_k1 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1);
|
||||||
|
self.k1c[canonical_k1 as usize] += 1;
|
||||||
|
self.k1q.push_back(canonical_k1);
|
||||||
|
|
||||||
|
let g2 = self.k2c[canonical_k2 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2);
|
||||||
|
self.k2c[canonical_k2 as usize] += 1;
|
||||||
|
self.k2q.push_back(canonical_k2);
|
||||||
|
|
||||||
|
let g3 = self.k3c[canonical_k3 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3);
|
||||||
|
self.k3c[canonical_k3 as usize] += 1;
|
||||||
|
self.k3q.push_back(canonical_k3);
|
||||||
|
|
||||||
|
let g4 = self.k4c[canonical_k4 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4);
|
||||||
|
self.k4c[canonical_k4 as usize] += 1;
|
||||||
|
self.k4q.push_back(canonical_k4);
|
||||||
|
|
||||||
|
let g5 = self.k5c[canonical_k5 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5);
|
||||||
|
self.k5c[canonical_k5 as usize] += 1;
|
||||||
|
self.k5q.push_back(canonical_k5);
|
||||||
|
|
||||||
|
let g6 = self.k6c[canonical_k6 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6);
|
||||||
|
self.k6c[canonical_k6 as usize] += 1;
|
||||||
|
self.k6q.push_back(canonical_k6);
|
||||||
|
} else {
|
||||||
|
self.push_warmup_increments(
|
||||||
|
canonical_k1, canonical_k2, canonical_k3,
|
||||||
|
canonical_k4, canonical_k5, canonical_k6,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cold]
|
||||||
|
#[inline(never)]
|
||||||
|
fn push_warmup_increments(
|
||||||
|
&mut self,
|
||||||
|
canonical_k1: u64, canonical_k2: u64, canonical_k3: u64,
|
||||||
|
canonical_k4: u64, canonical_k5: u64, canonical_k6: u64,
|
||||||
|
) {
|
||||||
|
let g1 = self.k1c[canonical_k1 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1);
|
||||||
|
self.k1c[canonical_k1 as usize] += 1;
|
||||||
|
self.k1q.push_back(canonical_k1);
|
||||||
|
|
||||||
|
if self.received >= 2 {
|
||||||
|
let g2 = self.k2c[canonical_k2 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2);
|
||||||
|
self.k2c[canonical_k2 as usize] += 1;
|
||||||
|
self.k2q.push_back(canonical_k2);
|
||||||
|
|
||||||
|
if self.received >= 3 {
|
||||||
|
let g3 = self.k3c[canonical_k3 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3);
|
||||||
|
self.k3c[canonical_k3 as usize] += 1;
|
||||||
|
self.k3q.push_back(canonical_k3);
|
||||||
|
|
||||||
|
if self.received >= 4 {
|
||||||
|
let g4 = self.k4c[canonical_k4 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4);
|
||||||
|
self.k4c[canonical_k4 as usize] += 1;
|
||||||
|
self.k4q.push_back(canonical_k4);
|
||||||
|
|
||||||
|
if self.received >= 5 {
|
||||||
|
let g5 = self.k5c[canonical_k5 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5);
|
||||||
|
self.k5c[canonical_k5 as usize] += 1;
|
||||||
|
self.k5q.push_back(canonical_k5);
|
||||||
|
|
||||||
|
if self.received >= 6 {
|
||||||
|
let g6 = self.k6c[canonical_k6 as usize] as usize;
|
||||||
|
Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6);
|
||||||
|
self.k6c[canonical_k6 as usize] += 1;
|
||||||
|
self.k6q.push_back(canonical_k6);
|
||||||
|
self.steady = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn ready(&self) -> bool {
|
pub fn ready(&self) -> bool {
|
||||||
@@ -196,10 +422,29 @@ impl RollingStat {
|
|||||||
.map(|raw| Minimizer::from_raw_unchecked(raw << (64 - self.m * 2)))
|
.map(|raw| Minimizer::from_raw_unchecked(raw << (64 - self.m * 2)))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn entropy(&self, order: usize) -> Option<f64> {
|
||||||
|
if !self.ready() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let k = self.k;
|
||||||
|
let em = emax(k, order);
|
||||||
|
if em <= 0.0 {
|
||||||
|
return Some(1.0);
|
||||||
|
}
|
||||||
|
let nwords = k - order + 1;
|
||||||
|
let log_nw = log_nwords(k, order);
|
||||||
|
let nw_f = nwords as f64;
|
||||||
|
let h_corr = log_nw + (self.sum_f_log_s[order] - self.sum_f_log_f[order]) / nw_f;
|
||||||
|
Some((h_corr / em).max(0.0))
|
||||||
|
}
|
||||||
|
|
||||||
pub fn normalized_entropy(&self) -> Option<f64> {
|
pub fn normalized_entropy(&self) -> Option<f64> {
|
||||||
if !self.ready() {
|
if !self.ready() {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
Some(self.entropy.normalized_entropy(self.entropy_max_k))
|
let min_e = (1..=self.entropy_max_k)
|
||||||
|
.filter_map(|ws| self.entropy(ws))
|
||||||
|
.fold(f64::MAX, f64::min);
|
||||||
|
Some(if min_e == f64::MAX { 1.0 } else { min_e })
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,152 +0,0 @@
|
|||||||
use super::*;
|
|
||||||
use obikrope::Rope;
|
|
||||||
|
|
||||||
fn setup() {
|
|
||||||
obikseq::params::set_k(K);
|
|
||||||
obikseq::params::set_m(5);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_rope(data: &[u8]) -> Rope {
|
|
||||||
let mut r = Rope::new(None);
|
|
||||||
r.push(data.to_vec());
|
|
||||||
r
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_nofilter(data: &[u8], k: usize) -> Vec<Vec<u8>> {
|
|
||||||
let rope = make_rope(data);
|
|
||||||
SuperKmerIter::new(&rope, k, 1, 0.0)
|
|
||||||
.map(|rsk| rsk.superkmer().to_ascii())
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
// k=11, m=5 — valeurs minimales du projet (k ∈ [11,31])
|
|
||||||
const K: usize = 11;
|
|
||||||
|
|
||||||
/// Collect the set of canonical k-mers from a raw ASCII sequence (no NUL).
|
|
||||||
fn direct_canonical_kmers(seq: &[u8]) -> std::collections::HashSet<Vec<u8>> {
|
|
||||||
(0..seq.len().saturating_sub(K - 1))
|
|
||||||
.map(|i| obikseq::SuperKmer::from_ascii(&seq[i..i + K]).to_ascii())
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Collect the set of canonical k-mers emitted by SuperKmerIter over a rope.
|
|
||||||
fn iter_canonical_kmers(rope: &Rope) -> std::collections::HashSet<Vec<u8>> {
|
|
||||||
SuperKmerIter::new(rope, K, 1, 0.0)
|
|
||||||
.flat_map(|rsk| {
|
|
||||||
rsk.superkmer()
|
|
||||||
.iter_canonical_kmers()
|
|
||||||
.map(|km| km.to_ascii())
|
|
||||||
.collect::<Vec<_>>()
|
|
||||||
})
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coverage_single_segment() {
|
|
||||||
setup();
|
|
||||||
let seq = b"ACGTACGTACGTACGTACGT";
|
|
||||||
let rope = make_rope(&[seq.as_ref(), b"\x00"].concat());
|
|
||||||
let direct = direct_canonical_kmers(seq);
|
|
||||||
let from_iter = iter_canonical_kmers(&rope);
|
|
||||||
let missing: Vec<_> = direct.difference(&from_iter).collect();
|
|
||||||
assert!(
|
|
||||||
missing.is_empty(),
|
|
||||||
"k-mers perdus dans segment unique : {missing:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coverage_two_segments() {
|
|
||||||
setup();
|
|
||||||
let seg1 = b"ACGTACGTACGTACGTACGT";
|
|
||||||
let seg2 = b"TGCATGCATGCATGCATGCA";
|
|
||||||
let rope = make_rope(&[seg1.as_ref(), b"\x00", seg2.as_ref(), b"\x00"].concat());
|
|
||||||
let mut direct = direct_canonical_kmers(seg1);
|
|
||||||
direct.extend(direct_canonical_kmers(seg2));
|
|
||||||
let from_iter = iter_canonical_kmers(&rope);
|
|
||||||
let missing: Vec<_> = direct.difference(&from_iter).collect();
|
|
||||||
assert!(
|
|
||||||
missing.is_empty(),
|
|
||||||
"k-mers perdus dans deux segments : {missing:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coverage_minimizer_boundary() {
|
|
||||||
setup();
|
|
||||||
// sequence assez longue pour forcer plusieurs changements de minimiseur
|
|
||||||
let seq: Vec<u8> = (0..80).map(|i| b"ACGT"[i % 4]).collect();
|
|
||||||
let rope = make_rope(&[seq.as_slice(), b"\x00"].concat());
|
|
||||||
let direct = direct_canonical_kmers(&seq);
|
|
||||||
let from_iter = iter_canonical_kmers(&rope);
|
|
||||||
let missing: Vec<_> = direct.difference(&from_iter).collect();
|
|
||||||
assert!(
|
|
||||||
missing.is_empty(),
|
|
||||||
"k-mers perdus à la frontière de minimiseur : {missing:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn single_segment_one_superkmer() {
|
|
||||||
setup();
|
|
||||||
let out = run_nofilter(b"ACGTACGTACGTACGTACGT\x00", K);
|
|
||||||
assert!(!out.is_empty());
|
|
||||||
let total: Vec<u8> = out.into_iter().flatten().collect();
|
|
||||||
assert!(total.len() >= K);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn segment_shorter_than_k_emits_nothing() {
|
|
||||||
setup();
|
|
||||||
let out = run_nofilter(b"ACGTACGT\x00", K);
|
|
||||||
assert_eq!(out, Vec::<Vec<u8>>::new());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn empty_input_emits_nothing() {
|
|
||||||
setup();
|
|
||||||
let out = run_nofilter(b"", K);
|
|
||||||
assert_eq!(out, Vec::<Vec<u8>>::new());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn two_segments_both_emitted() {
|
|
||||||
setup();
|
|
||||||
let out = run_nofilter(b"ACGTACGTACGTACGT\x00TGCATGCATGCATGCA\x00", K);
|
|
||||||
assert!(!out.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn low_complexity_kmer_is_rejected() {
|
|
||||||
setup();
|
|
||||||
let out_pass = run_nofilter(b"AAAAAAAAAAAACGTACGTACGT\x00", K);
|
|
||||||
assert!(!out_pass.is_empty());
|
|
||||||
|
|
||||||
let rope = make_rope(b"AAAAAAAAAAAAAAAAAAAA\x00");
|
|
||||||
let out_reject: Vec<Vec<u8>> = SuperKmerIter::new(&rope, K, 6, 0.9)
|
|
||||||
.map(|rsk| rsk.superkmer().to_ascii())
|
|
||||||
.collect();
|
|
||||||
assert!(out_reject.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn multi_slice_rope() {
|
|
||||||
setup();
|
|
||||||
let data = b"ACGTACGTACGTACGTACGT\x00";
|
|
||||||
let mid = data.len() / 2;
|
|
||||||
let mut rope = Rope::new(None);
|
|
||||||
rope.push(data[..mid].to_vec());
|
|
||||||
rope.push(data[mid..].to_vec());
|
|
||||||
let out: Vec<Vec<u8>> = SuperKmerIter::new(&rope, K, 1, 0.0)
|
|
||||||
.map(|rsk| rsk.superkmer().to_ascii())
|
|
||||||
.collect();
|
|
||||||
assert!(!out.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn yields_minimizer_value() {
|
|
||||||
setup();
|
|
||||||
let rope = make_rope(b"ACGTACGTACGTACGTACGT\x00");
|
|
||||||
let results: Vec<RoutableSuperKmer> = SuperKmerIter::new(&rope, K, 1, 0.0).collect();
|
|
||||||
assert!(!results.is_empty());
|
|
||||||
}
|
|
||||||
+50
-337
@@ -4,7 +4,7 @@ use std::sync::{Condvar, Mutex};
|
|||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use indicatif::{ProgressBar, ProgressStyle};
|
use indicatif::{ProgressBar, ProgressStyle};
|
||||||
use tracing::{debug, info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
const BRAILLE: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
const BRAILLE: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||||
|
|
||||||
@@ -31,8 +31,7 @@ impl TracedBar {
|
|||||||
let pct10 = (pos * 10) / self.total; // 0..=10
|
let pct10 = (pos * 10) / self.total; // 0..=10
|
||||||
let last = self.last_pct.load(Ordering::Relaxed);
|
let last = self.last_pct.load(Ordering::Relaxed);
|
||||||
if pct10 > last
|
if pct10 > last
|
||||||
&& self
|
&& self.last_pct
|
||||||
.last_pct
|
|
||||||
.compare_exchange(last, pct10, Ordering::Relaxed, Ordering::Relaxed)
|
.compare_exchange(last, pct10, Ordering::Relaxed, Ordering::Relaxed)
|
||||||
.is_ok()
|
.is_ok()
|
||||||
{
|
{
|
||||||
@@ -50,14 +49,14 @@ impl TracedBar {
|
|||||||
let msg = msg.into();
|
let msg = msg.into();
|
||||||
if self.pb.is_hidden() {
|
if self.pb.is_hidden() {
|
||||||
if self.total > 0 {
|
if self.total > 0 {
|
||||||
debug!(stage = %self.label, "{msg}");
|
// bounded bar: always log (already rate-limited by 10% threshold in inc)
|
||||||
|
info!(stage = %self.label, "{msg}");
|
||||||
} else {
|
} else {
|
||||||
// spinner: throttle to ~10 s
|
// spinner: throttle to ~10 s
|
||||||
let now_ms = self.start.elapsed().as_millis() as u64;
|
let now_ms = self.start.elapsed().as_millis() as u64;
|
||||||
let last = self.last_log_ms.load(Ordering::Relaxed);
|
let last = self.last_log_ms.load(Ordering::Relaxed);
|
||||||
if now_ms >= last + 10_000
|
if now_ms >= last + 10_000
|
||||||
&& self
|
&& self.last_log_ms
|
||||||
.last_log_ms
|
|
||||||
.compare_exchange(last, now_ms, Ordering::Relaxed, Ordering::Relaxed)
|
.compare_exchange(last, now_ms, Ordering::Relaxed, Ordering::Relaxed)
|
||||||
.is_ok()
|
.is_ok()
|
||||||
{
|
{
|
||||||
@@ -84,13 +83,8 @@ pub fn spinner(label: &str) -> TracedBar {
|
|||||||
);
|
);
|
||||||
pb.enable_steady_tick(Duration::from_millis(100));
|
pb.enable_steady_tick(Duration::from_millis(100));
|
||||||
TracedBar {
|
TracedBar {
|
||||||
pb,
|
pb, label: label.to_string(), unit: String::new(), total: 0,
|
||||||
label: label.to_string(),
|
start: Instant::now(), last_pct: AtomicU64::new(0), last_log_ms: AtomicU64::new(0),
|
||||||
unit: String::new(),
|
|
||||||
total: 0,
|
|
||||||
start: Instant::now(),
|
|
||||||
last_pct: AtomicU64::new(0),
|
|
||||||
last_log_ms: AtomicU64::new(0),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -107,13 +101,8 @@ pub fn progress_bar(label: &str, n: u64, unit: &str) -> TracedBar {
|
|||||||
);
|
);
|
||||||
pb.enable_steady_tick(Duration::from_millis(100));
|
pb.enable_steady_tick(Duration::from_millis(100));
|
||||||
TracedBar {
|
TracedBar {
|
||||||
pb,
|
pb, label: label.to_string(), unit: unit.to_string(), total: n,
|
||||||
label: label.to_string(),
|
start: Instant::now(), last_pct: AtomicU64::new(0), last_log_ms: AtomicU64::new(0),
|
||||||
unit: unit.to_string(),
|
|
||||||
total: n,
|
|
||||||
start: Instant::now(),
|
|
||||||
last_pct: AtomicU64::new(0),
|
|
||||||
last_log_ms: AtomicU64::new(0),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -202,94 +191,6 @@ fn cgroup_v1_available() -> Option<u64> {
|
|||||||
Some(limit.saturating_sub(used))
|
Some(limit.saturating_sub(used))
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── CPU parallelism query ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
/// Returns the number of cores this process can actually use concurrently.
|
|
||||||
///
|
|
||||||
/// `std::thread::available_parallelism()` reads CPU affinity
|
|
||||||
/// (`sched_getaffinity`), not the container's CPU quota — a Docker/cgroup
|
|
||||||
/// container commonly reports the *host's* full core count this way while
|
|
||||||
/// actually being throttled (via `cpu.max`/`cpu.cfs_quota_us`) to a fraction
|
|
||||||
/// of a core. Sizing a thread/worker pool off the unthrottled count causes
|
|
||||||
/// severe oversubscription: dozens of threads contending for a sliver of
|
|
||||||
/// real CPU time, which can look indistinguishable from a hang for minutes
|
|
||||||
/// or hours (observed in CI). On Linux, this reads the cgroup CPU quota
|
|
||||||
/// first and returns `min(cgroup_quota, host_parallelism)` when a finite
|
|
||||||
/// quota is found; falls back to `available_parallelism()` otherwise (same
|
|
||||||
/// convention as [`available_memory_bytes`]).
|
|
||||||
pub fn effective_parallelism() -> usize {
|
|
||||||
let host = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
{
|
|
||||||
if let Some(quota) = cgroup_v2_cpu_quota() {
|
|
||||||
let effective = quota.clamp(1, host);
|
|
||||||
tracing::debug!(host, quota, effective, source = "cgroup v2", "effective_parallelism");
|
|
||||||
return effective;
|
|
||||||
}
|
|
||||||
if let Some(quota) = cgroup_v1_cpu_quota() {
|
|
||||||
let effective = quota.clamp(1, host);
|
|
||||||
tracing::debug!(host, quota, effective, source = "cgroup v1", "effective_parallelism");
|
|
||||||
return effective;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
tracing::debug!(host, effective = host, source = "available_parallelism (no cgroup quota found)", "effective_parallelism");
|
|
||||||
host
|
|
||||||
}
|
|
||||||
|
|
||||||
/// cgroup v2 (unified hierarchy): reads `cpu.max` ("<quota> <period>", or
|
|
||||||
/// "max <period>" when unlimited) for the current process's cgroup, rounded
|
|
||||||
/// up to whole cores. Returns `None` if unlimited or on any parse error.
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
fn cgroup_v2_cpu_quota() -> Option<usize> {
|
|
||||||
let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
|
|
||||||
let rel = cgroup
|
|
||||||
.lines()
|
|
||||||
.find(|l| l.starts_with("0::"))?
|
|
||||||
.strip_prefix("0::")?
|
|
||||||
.trim();
|
|
||||||
let base = format!("/sys/fs/cgroup{rel}");
|
|
||||||
let raw = std::fs::read_to_string(format!("{base}/cpu.max")).ok()?;
|
|
||||||
let mut parts = raw.split_whitespace();
|
|
||||||
let quota_str = parts.next()?;
|
|
||||||
let period: f64 = parts.next()?.parse().ok()?;
|
|
||||||
if quota_str == "max" {
|
|
||||||
return None; // unlimited
|
|
||||||
}
|
|
||||||
let quota: f64 = quota_str.parse().ok()?;
|
|
||||||
Some((quota / period).ceil().max(1.0) as usize)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// cgroup v1 (cpu subsystem): reads `cpu.cfs_quota_us`/`cpu.cfs_period_us`,
|
|
||||||
/// rounded up to whole cores. Returns `None` if unlimited (quota <= 0) or on
|
|
||||||
/// any parse error.
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
fn cgroup_v1_cpu_quota() -> Option<usize> {
|
|
||||||
let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
|
|
||||||
let path = cgroup
|
|
||||||
.lines()
|
|
||||||
.find(|l| l.contains(":cpu:") || l.contains(":cpu,cpuacct:"))?
|
|
||||||
.split(':')
|
|
||||||
.nth(2)?;
|
|
||||||
let base = format!("/sys/fs/cgroup/cpu{path}");
|
|
||||||
let quota: i64 = std::fs::read_to_string(format!("{base}/cpu.cfs_quota_us"))
|
|
||||||
.ok()?
|
|
||||||
.trim()
|
|
||||||
.parse()
|
|
||||||
.ok()?;
|
|
||||||
if quota <= 0 {
|
|
||||||
return None; // unlimited
|
|
||||||
}
|
|
||||||
let period: i64 = std::fs::read_to_string(format!("{base}/cpu.cfs_period_us"))
|
|
||||||
.ok()?
|
|
||||||
.trim()
|
|
||||||
.parse()
|
|
||||||
.ok()?;
|
|
||||||
if period <= 0 {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
Some(((quota as f64) / (period as f64)).ceil().max(1.0) as usize)
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── raw helpers ───────────────────────────────────────────────────────────────
|
// ── raw helpers ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
fn get_rusage() -> rusage {
|
fn get_rusage() -> rusage {
|
||||||
@@ -303,19 +204,13 @@ fn tv_to_secs(tv: timeval) -> f64 {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(target_os = "macos")]
|
#[cfg(target_os = "macos")]
|
||||||
fn rss_to_bytes(ru: &rusage) -> u64 {
|
fn rss_to_bytes(ru: &rusage) -> u64 { ru.ru_maxrss as u64 }
|
||||||
ru.ru_maxrss as u64
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(not(target_os = "macos"))]
|
#[cfg(not(target_os = "macos"))]
|
||||||
fn rss_to_bytes(ru: &rusage) -> u64 {
|
fn rss_to_bytes(ru: &rusage) -> u64 { ru.ru_maxrss as u64 * 1024 }
|
||||||
ru.ru_maxrss as u64 * 1024
|
|
||||||
}
|
|
||||||
|
|
||||||
// Monotonically increasing counters — negative delta would be a kernel bug.
|
// Monotonically increasing counters — negative delta would be a kernel bug.
|
||||||
fn delta(end: i64, start: i64) -> u64 {
|
fn delta(end: i64, start: i64) -> u64 { (end - start).max(0) as u64 }
|
||||||
(end - start).max(0) as u64
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── CpuSample ─────────────────────────────────────────────────────────────────
|
// ── CpuSample ─────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -326,7 +221,6 @@ pub struct CpuSample {
|
|||||||
wall: Instant,
|
wall: Instant,
|
||||||
user_secs: f64,
|
user_secs: f64,
|
||||||
sys_secs: f64,
|
sys_secs: f64,
|
||||||
previous: f64,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl CpuSample {
|
impl CpuSample {
|
||||||
@@ -336,7 +230,6 @@ impl CpuSample {
|
|||||||
wall: Instant::now(),
|
wall: Instant::now(),
|
||||||
user_secs: tv_to_secs(ru.ru_utime),
|
user_secs: tv_to_secs(ru.ru_utime),
|
||||||
sys_secs: tv_to_secs(ru.ru_stime),
|
sys_secs: tv_to_secs(ru.ru_stime),
|
||||||
previous: 0.0,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -345,129 +238,11 @@ impl CpuSample {
|
|||||||
pub fn cpu_efficiency(&self, n_cores: usize) -> f64 {
|
pub fn cpu_efficiency(&self, n_cores: usize) -> f64 {
|
||||||
let ru = get_rusage();
|
let ru = get_rusage();
|
||||||
let wall = self.wall.elapsed().as_secs_f64();
|
let wall = self.wall.elapsed().as_secs_f64();
|
||||||
if wall < 0.1 {
|
if wall < 0.1 { return 0.0; }
|
||||||
return 0.0;
|
let cpu = (tv_to_secs(ru.ru_utime) - self.user_secs)
|
||||||
}
|
+ (tv_to_secs(ru.ru_stime) - self.sys_secs);
|
||||||
let cpu =
|
|
||||||
(tv_to_secs(ru.ru_utime) - self.user_secs) + (tv_to_secs(ru.ru_stime) - self.sys_secs);
|
|
||||||
cpu / (wall * n_cores as f64)
|
cpu / (wall * n_cores as f64)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn do_i_activate(&mut self, threshold: f64) -> bool {
|
|
||||||
let delta_wall = self.wall.elapsed().as_secs_f64();
|
|
||||||
if delta_wall < 0.1 {
|
|
||||||
// Window too short to be meaningful — leave state untouched so it
|
|
||||||
// keeps accumulating until a real sample can be taken.
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
let n = CpuSample::now();
|
|
||||||
let delta_ru = (n.user_secs - self.user_secs) + (n.sys_secs - self.sys_secs);
|
|
||||||
|
|
||||||
let efficiency = delta_ru / delta_wall;
|
|
||||||
let activate = 0f64.max(efficiency - self.previous) >= threshold;
|
|
||||||
|
|
||||||
debug!(
|
|
||||||
"Do I activate : {} -> {} = {} Activate: {}",
|
|
||||||
self.previous,
|
|
||||||
efficiency,
|
|
||||||
0f64.max(efficiency - self.previous),
|
|
||||||
activate
|
|
||||||
);
|
|
||||||
self.previous = efficiency;
|
|
||||||
self.user_secs = n.user_secs;
|
|
||||||
self.sys_secs = n.sys_secs;
|
|
||||||
self.wall = n.wall;
|
|
||||||
|
|
||||||
activate
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── IoSample ──────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
/// Snapshot of process-wide block I/O (bytes read + written) + wall clock.
|
|
||||||
///
|
|
||||||
/// Same activation protocol as [`CpuSample`], but the growth check in
|
|
||||||
/// [`do_i_activate`](Self::do_i_activate) is *relative* rather than absolute:
|
|
||||||
/// raw I/O throughput has no portable scale across storage devices, unlike a
|
|
||||||
/// core count.
|
|
||||||
pub struct IoSample {
|
|
||||||
wall: Instant,
|
|
||||||
bytes: u64,
|
|
||||||
previous_rate: f64,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl IoSample {
|
|
||||||
pub fn now() -> Self {
|
|
||||||
Self {
|
|
||||||
wall: Instant::now(),
|
|
||||||
bytes: Self::read_bytes(),
|
|
||||||
previous_rate: 0.0,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Bytes actually submitted to the block layer (read + write), summed
|
|
||||||
/// process-wide. Returns 0 if unavailable — degrades gracefully to a
|
|
||||||
/// signal that never triggers activation (CPU-only heuristic).
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
fn read_bytes() -> u64 {
|
|
||||||
let Ok(io) = std::fs::read_to_string("/proc/self/io") else {
|
|
||||||
return 0;
|
|
||||||
};
|
|
||||||
io.lines()
|
|
||||||
.filter_map(|l| {
|
|
||||||
l.strip_prefix("read_bytes: ")
|
|
||||||
.or_else(|| l.strip_prefix("write_bytes: "))
|
|
||||||
})
|
|
||||||
.filter_map(|v| v.trim().parse::<u64>().ok())
|
|
||||||
.sum()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(target_os = "macos")]
|
|
||||||
fn read_bytes() -> u64 {
|
|
||||||
use libc::{RUSAGE_INFO_V4, getpid, proc_pid_rusage, rusage_info_v4};
|
|
||||||
let mut info: rusage_info_v4 = unsafe { std::mem::zeroed() };
|
|
||||||
let ret =
|
|
||||||
unsafe { proc_pid_rusage(getpid(), RUSAGE_INFO_V4, &mut info as *mut _ as *mut _) };
|
|
||||||
if ret != 0 {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
info.ri_diskio_bytesread + info.ri_diskio_byteswritten
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
|
||||||
fn read_bytes() -> u64 {
|
|
||||||
0
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Same protocol as [`CpuSample::do_i_activate`] (0.1 s minimum window,
|
|
||||||
/// state untouched on early return), but growth is measured relative to
|
|
||||||
/// the previous rate. `threshold` is a fraction, e.g. `0.2` for a 20 %
|
|
||||||
/// increase in throughput since the last real sample.
|
|
||||||
pub fn do_i_activate(&mut self, threshold: f64) -> bool {
|
|
||||||
let elapsed = self.wall.elapsed().as_secs_f64();
|
|
||||||
if elapsed < 0.1 {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
let n = Self::read_bytes();
|
|
||||||
let rate = n.saturating_sub(self.bytes) as f64 / elapsed;
|
|
||||||
let activate = if self.previous_rate == 0.0 {
|
|
||||||
rate > 0.0 // bootstrap: any measured throughput is signal enough
|
|
||||||
} else {
|
|
||||||
(rate - self.previous_rate) / self.previous_rate >= threshold
|
|
||||||
};
|
|
||||||
|
|
||||||
debug!(
|
|
||||||
"Do I activate (I/O) : {} -> {} Activate: {}",
|
|
||||||
self.previous_rate, rate, activate
|
|
||||||
);
|
|
||||||
self.previous_rate = rate;
|
|
||||||
self.bytes = n;
|
|
||||||
self.wall = Instant::now();
|
|
||||||
|
|
||||||
activate
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── public API ────────────────────────────────────────────────────────────────
|
// ── public API ────────────────────────────────────────────────────────────────
|
||||||
@@ -484,11 +259,7 @@ impl Stage {
|
|||||||
pub fn start(label: impl Into<String>) -> Self {
|
pub fn start(label: impl Into<String>) -> Self {
|
||||||
let label = label.into();
|
let label = label.into();
|
||||||
info!(stage = %label, "started");
|
info!(stage = %label, "started");
|
||||||
Self {
|
Self { label, wall: Instant::now(), ru: get_rusage() }
|
||||||
label,
|
|
||||||
wall: Instant::now(),
|
|
||||||
ru: get_rusage(),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn stop(self) -> StageStats {
|
pub fn stop(self) -> StageStats {
|
||||||
@@ -547,11 +318,8 @@ pub struct StageStats {
|
|||||||
impl StageStats {
|
impl StageStats {
|
||||||
/// (user + sys) / wall — effective thread count utilisation.
|
/// (user + sys) / wall — effective thread count utilisation.
|
||||||
pub fn parallelism(&self) -> f64 {
|
pub fn parallelism(&self) -> f64 {
|
||||||
if self.wall_secs > 1e-9 {
|
if self.wall_secs > 1e-9 { (self.user_secs + self.sys_secs) / self.wall_secs }
|
||||||
(self.user_secs + self.sys_secs) / self.wall_secs
|
else { 0.0 }
|
||||||
} else {
|
|
||||||
0.0
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// parallelism / n_cores — fraction of available CPU power used (0..1+).
|
/// parallelism / n_cores — fraction of available CPU power used (0..1+).
|
||||||
@@ -567,19 +335,11 @@ pub struct Reporter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl Reporter {
|
impl Reporter {
|
||||||
pub fn new() -> Self {
|
pub fn new() -> Self { Self::default() }
|
||||||
Self::default()
|
pub fn push(&mut self, stats: StageStats) { self.stages.push(stats); }
|
||||||
}
|
pub fn stages(&self) -> &[StageStats] { &self.stages }
|
||||||
pub fn push(&mut self, stats: StageStats) {
|
|
||||||
self.stages.push(stats);
|
|
||||||
}
|
|
||||||
pub fn stages(&self) -> &[StageStats] {
|
|
||||||
&self.stages
|
|
||||||
}
|
|
||||||
/// Print the summary to stderr.
|
/// Print the summary to stderr.
|
||||||
pub fn print(&self) {
|
pub fn print(&self) { eprint!("{self}"); }
|
||||||
eprint!("{self}");
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── diagnosis ─────────────────────────────────────────────────────────────────
|
// ── diagnosis ─────────────────────────────────────────────────────────────────
|
||||||
@@ -627,43 +387,26 @@ fn diagnose(s: &StageStats, n_cores: usize) -> Diagnosis {
|
|||||||
)),
|
)),
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
Diagnosis {
|
Diagnosis { tag: "—", detail: None }
|
||||||
tag: "—",
|
|
||||||
detail: None,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── display helpers ───────────────────────────────────────────────────────────
|
// ── display helpers ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
fn fmt_secs(s: f64) -> String {
|
fn fmt_secs(s: f64) -> String {
|
||||||
if s >= 100.0 {
|
if s >= 100.0 { format!("{:.0}s", s) }
|
||||||
format!("{:.0}s", s)
|
else if s >= 10.0 { format!("{:.1}s", s) }
|
||||||
} else if s >= 10.0 {
|
else if s >= 1.0 { format!("{:.2}s", s) }
|
||||||
format!("{:.1}s", s)
|
else { format!("{:.0}ms", s * 1000.0) }
|
||||||
} else if s >= 1.0 {
|
|
||||||
format!("{:.2}s", s)
|
|
||||||
} else {
|
|
||||||
format!("{:.0}ms", s * 1000.0)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn fmt_bytes(b: u64) -> String {
|
fn fmt_bytes(b: u64) -> String {
|
||||||
if b >= 1 << 30 {
|
if b >= 1 << 30 { format!("{:.1} GB", b as f64 / (1u64 << 30) as f64) }
|
||||||
format!("{:.1} GB", b as f64 / (1u64 << 30) as f64)
|
else if b >= 1 << 20 { format!("{:.0} MB", b as f64 / (1u64 << 20) as f64) }
|
||||||
} else if b >= 1 << 20 {
|
else { format!("{:.0} KB", b as f64 / 1024.0) }
|
||||||
format!("{:.0} MB", b as f64 / (1u64 << 20) as f64)
|
|
||||||
} else {
|
|
||||||
format!("{:.0} KB", b as f64 / 1024.0)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn fmt_efficiency(par: f64, n_cores: usize) -> String {
|
fn fmt_efficiency(par: f64, n_cores: usize) -> String {
|
||||||
format!(
|
format!("{:.1}×/{} ({:.0}%)", par, n_cores, par / n_cores as f64 * 100.0)
|
||||||
"{:.1}×/{} ({:.0}%)",
|
|
||||||
par,
|
|
||||||
n_cores,
|
|
||||||
par / n_cores as f64 * 100.0
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Display ───────────────────────────────────────────────────────────────────
|
// ── Display ───────────────────────────────────────────────────────────────────
|
||||||
@@ -691,11 +434,7 @@ impl MemoryBudget {
|
|||||||
pub fn new(total: u64) -> Self {
|
pub fn new(total: u64) -> Self {
|
||||||
Self {
|
Self {
|
||||||
total,
|
total,
|
||||||
inner: Mutex::new(BudgetInner {
|
inner: Mutex::new(BudgetInner { remaining: total, active: 0, peak_active: 0 }),
|
||||||
remaining: total,
|
|
||||||
active: 0,
|
|
||||||
peak_active: 0,
|
|
||||||
}),
|
|
||||||
condvar: Condvar::new(),
|
condvar: Condvar::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -720,38 +459,24 @@ impl MemoryBudget {
|
|||||||
self.condvar.notify_all();
|
self.condvar.notify_all();
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn total(&self) -> u64 {
|
pub fn total(&self) -> u64 { self.total }
|
||||||
self.total
|
pub fn active(&self) -> usize { self.inner.lock().unwrap().active }
|
||||||
}
|
pub fn remaining(&self) -> u64 { self.inner.lock().unwrap().remaining }
|
||||||
pub fn active(&self) -> usize {
|
pub fn peak_active(&self) -> usize { self.inner.lock().unwrap().peak_active }
|
||||||
self.inner.lock().unwrap().active
|
|
||||||
}
|
|
||||||
pub fn remaining(&self) -> u64 {
|
|
||||||
self.inner.lock().unwrap().remaining
|
|
||||||
}
|
|
||||||
pub fn peak_active(&self) -> usize {
|
|
||||||
self.inner.lock().unwrap().peak_active
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Display ───────────────────────────────────────────────────────────────────
|
// ── Display ───────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
impl fmt::Display for Reporter {
|
impl fmt::Display for Reporter {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
if self.stages.is_empty() {
|
if self.stages.is_empty() { return Ok(()); }
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
let n_cores = effective_parallelism();
|
let n_cores = std::thread::available_parallelism()
|
||||||
|
.map(|n| n.get())
|
||||||
|
.unwrap_or(1);
|
||||||
|
|
||||||
// column widths
|
// column widths
|
||||||
let nw = self
|
let nw = self.stages.iter().map(|s| s.label.len()).max().unwrap_or(5).max(5);
|
||||||
.stages
|
|
||||||
.iter()
|
|
||||||
.map(|s| s.label.len())
|
|
||||||
.max()
|
|
||||||
.unwrap_or(5)
|
|
||||||
.max(5);
|
|
||||||
// efficiency col: worst-case width for this run's n_cores value
|
// efficiency col: worst-case width for this run's n_cores value
|
||||||
let ew = format!("{:.1}×/{} (100%)", 99.9f64, n_cores).len();
|
let ew = format!("{:.1}×/{} (100%)", 99.9f64, n_cores).len();
|
||||||
|
|
||||||
@@ -759,21 +484,18 @@ impl fmt::Display for Reporter {
|
|||||||
let sep = "─".repeat(sep_w);
|
let sep = "─".repeat(sep_w);
|
||||||
|
|
||||||
// header
|
// header
|
||||||
writeln!(
|
writeln!(f, "{:<nw$} {:>7} {:>ew$} {:>8} status",
|
||||||
f,
|
"stage", "wall", "efficiency", "peak RSS")?;
|
||||||
"{:<nw$} {:>7} {:>ew$} {:>8} status",
|
|
||||||
"stage", "wall", "efficiency", "peak RSS"
|
|
||||||
)?;
|
|
||||||
writeln!(f, "{sep}")?;
|
writeln!(f, "{sep}")?;
|
||||||
|
|
||||||
// compute all diagnoses up front (needed for both table and footnotes)
|
// compute all diagnoses up front (needed for both table and footnotes)
|
||||||
let diagnoses: Vec<Diagnosis> = self.stages.iter().map(|s| diagnose(s, n_cores)).collect();
|
let diagnoses: Vec<Diagnosis> = self.stages.iter()
|
||||||
|
.map(|s| diagnose(s, n_cores))
|
||||||
|
.collect();
|
||||||
|
|
||||||
// per-stage rows
|
// per-stage rows
|
||||||
for (s, d) in self.stages.iter().zip(diagnoses.iter()) {
|
for (s, d) in self.stages.iter().zip(diagnoses.iter()) {
|
||||||
writeln!(
|
writeln!(f, "{:<nw$} {:>7} {:>ew$} {:>8} {}",
|
||||||
f,
|
|
||||||
"{:<nw$} {:>7} {:>ew$} {:>8} {}",
|
|
||||||
s.label,
|
s.label,
|
||||||
fmt_secs(s.wall_secs),
|
fmt_secs(s.wall_secs),
|
||||||
fmt_efficiency(s.parallelism(), n_cores),
|
fmt_efficiency(s.parallelism(), n_cores),
|
||||||
@@ -786,18 +508,11 @@ impl fmt::Display for Reporter {
|
|||||||
let tw = self.stages.iter().map(|s| s.wall_secs).sum::<f64>();
|
let tw = self.stages.iter().map(|s| s.wall_secs).sum::<f64>();
|
||||||
let tu = self.stages.iter().map(|s| s.user_secs).sum::<f64>();
|
let tu = self.stages.iter().map(|s| s.user_secs).sum::<f64>();
|
||||||
let ts = self.stages.iter().map(|s| s.sys_secs).sum::<f64>();
|
let ts = self.stages.iter().map(|s| s.sys_secs).sum::<f64>();
|
||||||
let trss = self
|
let trss = self.stages.iter().map(|s| s.max_rss_bytes).max().unwrap_or(0);
|
||||||
.stages
|
|
||||||
.iter()
|
|
||||||
.map(|s| s.max_rss_bytes)
|
|
||||||
.max()
|
|
||||||
.unwrap_or(0);
|
|
||||||
let tpar = if tw > 1e-9 { (tu + ts) / tw } else { 0.0 };
|
let tpar = if tw > 1e-9 { (tu + ts) / tw } else { 0.0 };
|
||||||
|
|
||||||
writeln!(f, "{sep}")?;
|
writeln!(f, "{sep}")?;
|
||||||
writeln!(
|
writeln!(f, "{:<nw$} {:>7} {:>ew$} {:>8}",
|
||||||
f,
|
|
||||||
"{:<nw$} {:>7} {:>ew$} {:>8}",
|
|
||||||
"TOTAL",
|
"TOTAL",
|
||||||
fmt_secs(tw),
|
fmt_secs(tw),
|
||||||
fmt_efficiency(tpar, n_cores),
|
fmt_efficiency(tpar, n_cores),
|
||||||
@@ -805,9 +520,7 @@ impl fmt::Display for Reporter {
|
|||||||
)?;
|
)?;
|
||||||
|
|
||||||
// bottleneck footnotes (only if at least one anomaly detected)
|
// bottleneck footnotes (only if at least one anomaly detected)
|
||||||
let bottlenecks: Vec<(&str, &str)> = self
|
let bottlenecks: Vec<(&str, &str)> = self.stages.iter()
|
||||||
.stages
|
|
||||||
.iter()
|
|
||||||
.zip(diagnoses.iter())
|
.zip(diagnoses.iter())
|
||||||
.filter_map(|(s, d)| d.detail.as_deref().map(|det| (s.label.as_str(), det)))
|
.filter_map(|(s, d)| d.detail.as_deref().map(|det| (s.label.as_str(), det)))
|
||||||
.collect();
|
.collect();
|
||||||
|
|||||||
Reference in New Issue
Block a user