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

refactor: implement RoutableSuperKmer and update k-mer indexing pipeline

Browse Source 
Replace raw SuperkMer routing with a new RoutableSuperKimer type that embeds canonical sequences and precomputed minimizers, enabling direct partition routing via hash. Update the build pipeline to yield RoutableSuperKmers throughout (builder, scatterer), refactor FASTA/unitig export commands to use the new type and compressed outputs (.fasta.gz, .unitigs.fasta.zst), revise SuperKmer header to store n_kmers instead of seql (avoiding 256-byte wrap), and update documentation to reflect minimizer-based theory, two evidence-encoding strategies for unitig-MPHF indexing (global offset vs. ID+rank), and the new obipipeline library architecture with parallel workers, biased scheduling, and error handling.
This commit is contained in:
Eric Coissac
2026-04-29 22:52:42 +02:00
parent 4e26e3bd40
commit 27f5e88a7b
72 changed files with 10093 additions and 1626 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docmd/implementation/superkmer.md
+20 -10
View File
@@ -2,26 +2,34 @@
## Memory layout
A super-kmer is stored as a **32-bit header** followed by a **byte-aligned nucleotide sequence** (2 bits/base, nucleotide 0 at the MSB of the first byte, max 256 nt):
A super-kmer is stored as a **32-bit header** followed by a **byte-aligned nucleotide sequence** (2 bits/base, nucleotide 0 at the MSB of the first byte):
| Field | Bits | Role |
|-------|------|------|
| COUNT | 24 | Occurrence count (≤ 16 M) |
| SEQL | 8 | Sequence length in nucleotides (1–256) |
| NKMERS | 8 | Number of kmers (= seq_length − k + 1, range 1–255) |
Bit layout (MSB to LSB): `[31:8] COUNT [7:0] SEQL`
Bit layout (MSB to LSB): `[31:8] COUNT [7:0] NKMERS`
SEQL is stored as a raw `u8`: values 1–255 represent lengths 1–255; **0 represents 256** (wrapping convention). The public accessor returns a `usize` and performs the conversion:
NKMERS is stored as a raw `u8` in **kmer units**, not nucleotides. The nucleotide length is recovered as `NKMERS + k − 1`. This avoids the awkward wrapping convention (`0 = 256`) that would be needed if nucleotide length were stored directly, and gains k−1 = 30 units of headroom:
| unit | u8 covers | max nucleotides |
|---|---|---|
| nucleotides | 255 nt | 225 kmers |
| **kmers** | **255 kmers** | **285 nt** |
The public accessors:
```rust
fn seql(&self) -> usize { if s == 0 { 256 } else { s as usize } }
fn n_kmers(&self) -> usize { (self.0 & 0xFF) as usize }
fn seql(&self) -> usize { self.n_kmers() + K - 1 }
fn count(&self) -> u32 { self.0 >> 8 }
fn increment(&mut self) { self.0 += 1 << 8; }
fn add(&mut self, n: u32) { self.0 += n << 8; }
fn set_count(&mut self, n: u32) { self.0 = (self.0 & 0xFF) | (n << 8); }
```
The SEQL field is 8 bits, capping the stored sequence at 256 nt. Given the expected length of ~40 nt, this cap is almost never reached; when it is, the super-kmer is split at 256 nt with a k−1 overlap, preserving all kmers without duplication.
In practice, observed super-kmer lengths on metagenomic data (k=31) are below 55 nucleotides (≤ 25 kmers) — far from the 255-kmer cap. If a super-kmer ever exceeds 255 kmers, it is split with a k−1 nucleotide overlap, preserving all kmers without duplication (identical mechanism to partition-boundary splits).
The sequence is always stored in canonical form (lexicographic minimum of forward and reverse complement), with nucleotide 0 at the MSB of the first byte. The byte array can be hashed directly without any adjustment.
@@ -61,10 +69,11 @@ const fn revcomp4(x: u8) -> u8 {
`REVCOMP4` is 256 bytes (fits in L1 cache), computed at compile time. No endianness dependency — all operations are pure arithmetic on byte values.
**Step 2 — realignment.** After step 1, `padding = n × 8 − SEQL × 2` spurious bits (complements of the original padding A's) appear at the start of the array. They are flushed left using `BitSlice<u8, Msb0>::rotate_left(padding)` from the `bitvec` crate, which is SIMD-accelerated. The trailing `padding` bits are then zeroed:
**Step 2 — realignment.** After step 1, `padding = n × 8 − seql × 2` spurious bits (complements of the original padding A's) appear at the start of the array. They are flushed left using `BitSlice<u8, Msb0>::rotate_left(padding)` from the `bitvec` crate, which is SIMD-accelerated. The trailing `padding` bits are then zeroed:
```rust
shift = n * 8 - SEQL * 2 // number of padding bits
let seql = self.n_kmers() + k - 1;
shift = n * 8 - seql * 2 // number of padding bits
bits.rotate_left(shift)
bits[len - shift..].fill(false)
```
@@ -141,8 +150,9 @@ The bit slice `seq[i*2 .. (i+k)*2]` (Msb0 order) is loaded as a big-endian `u64`
!!! abstract "Algorithm — Super-kmer reverse complement"
```text
procedure SuperKmerRevcomp(seq, SEQL):
n ← ⌈SEQL / 4⌉ -- number of bytes
shift ← n × 8 − SEQL × 2 -- padding bits to flush
seql ← NKMERS + k − 1 -- nucleotide length
n ← ⌈seql / 4⌉ -- number of bytes
shift ← n × 8 − seql × 2 -- padding bits to flush
-- step 1: swap bytes outside-in, applying REVCOMP4 to each (256-byte L1 table)
lo ← 0 ; hi ← n − 1
docmd/implementation/unitig_evidence.md
+232
View File
@@ -0,0 +1,232 @@
# Unitig-based MPHF evidence encoding
## Role of unitigs in the index
The MPHF maps each canonical kmer to an integer slot, but provides no way to reconstruct the kmer from its slot. A downstream operation (query, set operation) that receives a slot index and needs the kmer sequence must be able to retrieve it. The **evidence file** serves this purpose: it stores the kmer sequences in compact form and provides, for each MPHF slot, a pointer to where the corresponding kmer can be decoded.
Unitigs are the natural compact representation: a run of L nucleotides encodes L − k + 1 consecutive canonical kmers. The entire kmer set of a partition can be reconstructed from its unitig FASTA file.
---
## Two encoding strategies
### Strategy A — global nucleotide offset
Each MPHF slot stores a single integer: the byte offset of the kmer's first nucleotide within a packed 2-bit nucleotide array that concatenates all unitigs.
```
evidence[slot] = global_offset (bits: ⌈log₂ N_nuc⌉)
```
where `N_nuc` is the total number of nucleotides across all unitigs in the partition.
Decoding: read k nucleotides starting at `global_offset`.
### Strategy B — (unitig_id, rank within unitig)
Each MPHF slot stores a pair:
```
evidence[slot] = (unitig_id, rank)
```
- `unitig_id` : index of the unitig in the partition (0-based)
- `rank` : kmer index within the unitig (0 ≤ rank < n_kmers); kmer i starts at nucleotide i, so the nucleotide offset is identical numerically but the kmer-unit interpretation is the natural one
Decoding: look up the unitig at `unitig_id`, then read k nucleotides starting at `rank`.
---
## Bit-cost analysis
Define for a partition of P kmers with average kmers-per-unitig m:
- total nucleotides: $N_{nuc} = P \cdot \left(1 + \dfrac{k-1}{m}\right)$
- number of unitigs: $U = P / m$
**Strategy A**
$$
b_A = \left\lceil \log_2 N_{nuc} \right\rceil = \left\lceil \log_2 P + \log_2\!\left(1 + \frac{k-1}{m}\right) \right\rceil
$$
**Strategy B**
$$
b_B = \left\lceil \log_2 U \right\rceil + \left\lceil \log_2 L_{max} \right\rceil
$$
where $L_{max}$ is the maximum unitig length (in nucleotides). In practice $L_{max} \ll P$, so the rank field is much cheaper than the full global offset. If unitig lengths are bounded (e.g. by partition structure), the rank field width is a small constant independent of P.
### Empirical bound on unitig length
Lengths and ranks are expressed in **kmer units** (not nucleotides): the nucleotide length is `n_kmers + k − 1`, so storing `n_kmers` instead of `seq_length` saves k−1 = 30 units of headroom in the same field width.
Consequence for `u8` capacity:
| unit | max representable | max nucleotides |
|---|---|---|
| nucleotides | 255 nuc | 225 kmers |
| **kmers** | **255 kmers** | **285 nuc** |
On *Betula nana* (k=31, 256 partitions), m_u ≈ 37.9 kmers/unitig on average; no unitig length distribution data measured yet. The `rank` field (kmer index within the unitig) fits in a `u8` as long as no unitig exceeds 255 kmers — guaranteed by the split strategy below.
### Split strategy for long unitigs
For the rare cases where a unitig exceeds 255 kmers, the unitig is split into chunks of at most 255 kmers, with a **k−1 nucleotide overlap** at each junction — identical to the way super-kmers are delimited at partition boundaries. Each chunk is self-contained and independently decodable.
```
original unitig: kmer_0 … kmer_254 | kmer_255 … kmer_N
↑ cut here
chunk 1: nucleotides 0 … 284 (255 kmers)
chunk 2: nucleotides 255 … N+k-1 (N-255+1 kmers)
shared: nucleotides 255 … 284 (k-1 = 30 nucleotides, stored in both)
```
Cost of one split: k−1 = 30 redundant nucleotides = 60 bits. This event is rare in practice (m_u ≈ 38 for *B. nana*, well below the 255-kmer cap). No kmer is lost: kmer i is in chunk 1 if i < 255, in chunk 2 (at rank i−255) otherwise.
### Savings from u8 length fields
Because all chunks are guaranteed ≤ 255 kmers, the per-chunk length array in the binary index is a flat `u8` array — 1 byte per chunk instead of 8 bytes (usize) or 4 bytes (u32). For a partition with 4 M unitigs:
| length type | bytes/chunk | total (4 M chunks) |
|---|---|---|
| usize (u64) | 8 | 32 MB |
| u32 | 4 | 16 MB |
| **u8** | **1** | **4 MB** |
Random access to chunk i is recovered at load time by a single prefix-sum pass over the u8 array, computing a u32/u64 offset array in O(n_chunks) time and O(n_chunks × 4) bytes — paid once at open time, cached for the lifetime of the partition handle.
Bit costs for *Betula nana* (k=31, 256 partitions, P ≈ 10.4 M, U ≈ 275 k, m_u ≈ 37.9):
| field | strategy A | strategy B |
|---|---|---|
| offset / id | $\lceil\log_2(P \cdot (1 + 30/m_u))\rceil = 25$ bits | $\lceil\log_2(U)\rceil = 19$ bits |
| rank | — | 8 bits (u8, fixed) |
| **total** | **25 bits** | **27 bits** |
Strategy A is 2 bits cheaper. Strategy B's main advantage is **locality**: decoding a kmer touches one unitig's cache lines rather than an arbitrary offset in a large flat array, and the `rank` field doubles as a direct index into the packed nucleotide sequence without pointer arithmetic.
---
## Partition-size tradeoff
The total bits/kmer for the index (sequence + evidence + MPHF) as a function of partition size is:
$$
\text{total} = \underbrace{2\!\left(1 + \frac{k-1}{m}\right)}_{\text{sequence}} + \underbrace{\log_2 P + \log_2\!\left(1+\frac{k-1}{m}\right)}_{\text{evidence}} + \underbrace{c_{MPHF}}_{\approx 2\text{–}4}
$$
### Empirical observation: m_u is set by De Bruijn graph topology, not partition count
Measured on *Betula nana* (k=31, m=11), summing n_kmers and sequence counts across all partition files:
| N partitions | m_sk | m_u | factor m_u/m_sk | nuc ratio (u/sk) |
|---|---|---|---|---|
| 1 | 12.13 | **41.89** | 3.45× | 0.273 |
| 16 | 12.13 | **38.19** | 3.15× | 0.376 |
| 256 | 12.13 | **37.90** | 3.12× | 0.388 |
| 1 024 | 12.13 | **37.89** | 3.12× | 0.389 |
- `m_sk` = avg kmers/super-kmer (invariant — same dataset regardless of partition scheme)
- `m_u` = avg kmers/unitig = total_n_kmers / total_unitigs, summed across all partitions
- `nuc ratio` = (u_symbols + 30·u_reads) / (sk_symbols + 30·sk_reads)
X-axis in both charts: partition bits (0 = 1 partition, 10 = 1024 partitions) — each step doubles the partition count.
```mermaid
xychart-beta
title "m_u (avg kmers/unitig) vs partition bits — B. nana k=31"
x-axis "partition bits" 0 --> 10
y-axis "m_u" 37 --> 43
line [41.89, 40.78, 39.22, 38.52, 38.19, 38.03, 37.96, 37.92, 37.90, 37.89, 37.89]
```
```mermaid
xychart-beta
title "Nucleotide storage: unitigs / super-kmers (%) vs partition bits — B. nana k=31"
x-axis "partition bits" 0 --> 10
y-axis "%" 25 --> 42
line [27.3, 29.7, 33.9, 36.3, 37.6, 38.3, 38.6, 38.7, 38.8, 38.9, 38.9]
```
Key observations:
1. **Partition boundaries have a small but non-zero effect on m_u.** Going from 1 to 1024 partitions reduces m_u by 10% (41.9 → 37.9). Within the practical range 16–1024, the variation is under 1% — m_u is effectively constant.
2. **m_u is a property of the De Bruijn graph, not the partition scheme.** The dominant factor is graph branching (heterozygosity, repeats, sequencing errors).
3. **Unitigs provide substantial compaction over super-kmers.** At 256 partitions, unitigs cover the same unique kmers using 39% of the raw nucleotide content of super-kmers (3.1× compaction factor).
#### Per-partition compaction ratio (sk_symbols / u_symbols)
The ratio measures how much super-kmer kmer-slots are "shared" across different super-kmer records: a ratio of 1.35 means each unique kmer (counted once in unitigs) appears in 1.35 super-kmer kmer-slots on average.
| bits | N partitions | median ratio | min ratio | min partition | min u_reads |
|---|---|---|---|---|---|
| 6 | 64 | 1.355 | 1.073 | — | 4.5 M |
| 7 | 128 | 1.352 | 1.037 | — | 4.1 M |
| 8 | 256 | **1.350** | **1.012** | **145** | **3.8 M** |
| 9 | 512 | 1.350 | 0.998 | 145 | 3.6 M |
| 10 | 1024 | 1.351 | 0.992 | 145 | 3.6 M |
The median stabilises at **1.35** from 64 partitions onward (stdev = 0.027 at 256 partitions). There is one persistent outlier: **partition 145** (at 256-partition resolution) is consistently anomalous across all partition depths — it contains 10–14× more super-kmers and unitigs than the average partition, with a ratio near 1.0, meaning the unitig representation provides almost no kmer deduplication. This is consistent with a highly repetitive or organellar region where the dominant minimiser belongs to a sequence that appears in many reads without forming long overlapping paths in the De Bruijn graph.
Per-partition parameters at 256 partitions (*B. nana*):
| quantity | value |
|---|---|
| P (unique kmers/partition, avg) | ≈ 10.4 M |
| U (unitigs/partition, avg) | ≈ 275 k |
| m_u | ≈ 37.9 |
| Strategy A bits/kmer | ⌈log₂(P·(1+30/m_u))⌉ = 25 |
| Strategy B bits/kmer | ⌈log₂(U)⌉ + 8 = 27 |
Consequence: **the partition count should be as large as memory and parallelism allow.** Each doubling saves 1 bit/kmer in evidence (log₂ P decreases by 1). The sequence term 2·(1 + 30/m_u) ≈ 3.6 bits/kmer is approximately constant.
Strategy B partially decouples evidence cost from P: `log₂(U) = log₂(P/m_u)` grows more slowly than `log₂(P)` by a fixed log₂(m_u) ≈ 5 bits. Strategy B's main benefit remains locality and bounded rank width, not asymptotic compression.
---
## Implementation notes
### Evidence file layout (strategy B)
```
evidence.bin
├── header : k (u8), n_kmers (u64), n_unitigs (u64)
├── id_array : n_kmers × ⌈log₂ n_unitigs⌉ bits — MPHF slot → unitig_id
└── rank_array: n_kmers × 8 bits (u8[n_kmers]) — MPHF slot → rank within unitig
```
`id_array` is a compact bit-packed vector (width = ⌈log₂ n_unitigs⌉; 19 bits for *B. nana* at 256 partitions). `rank_array` is a plain `u8` array — no bit-packing needed. Access is O(1) with a single multiplication and mask for `id_array`, and a direct byte index for `rank_array`.
### Unitig file layout
FASTA with JSON annotation header (xxHash-64 ID, seq_length, kmer_size, n_kmers). The nucleotide sequence is stored in ASCII uppercase; a 2-bit packed version is derived at query time or stored as a parallel `.2bit` file for speed.
```
>c4a1e7f2 {"seq_length":87,"kmer_size":31,"n_kmers":57}
ACGTGGCTA...
```
### Decoding a kmer from slot s
```
unitig_id = id_array[s]
rank = rank_array[s]
kmer = nucleotides(unitig_id)[rank .. rank + k] // 2-bit packed slice
```
One array lookup per field, then a packed slice extraction. The canonical kmer is the stored sequence (by construction — only canonical kmers are inserted into the graph).
### Forward vs reverse complement
The De Bruijn graph stores only canonical kmers. The evidence encodes the canonical orientation. Callers that need the strand of the original kmer must compare the retrieved kmer with its revcomp at query time; this is a single 64-bit comparison.
---
## Open questions
- **Rank field width**: u8 covers 255 kmers; storing lengths and ranks in kmer units (not nucleotides) buys k−1 extra units of headroom at no cost. On *B. nana* (k=31), m_u ≈ 38 — well within u8 range on average, but the maximum unitig length has not been measured yet. For genomes with very long unitigs, u16 may be needed; the header could record the actual width if portability is required.
- **Packed nucleotide cache**: storing a 2-bit packed nucleotide array alongside the FASTA avoids re-encoding at query time; negligible space overhead ($N_{nuc} / 4$ bytes per partition).
- **Cross-partition evidence**: for set operations spanning multiple partitions, strategy B allows unitig-level operations (e.g. mark entire unitigs as present/absent) rather than kmer-level, potentially reducing the operation cost by a factor of m.