feat: add entropy-driven k-mer complexity filtering
Introduces a MinComplexity filter driven by new CLI arguments, enabling sequence-aware threshold checks during index reconstruction and partitioning. Adds the kmer_entropy module for normalized complexity scoring, updates the KmerFilter trait to evaluate per-kmer context, and refactors test modules for better organization.
This commit is contained in:
Eric Coissac
@@ -149,157 +149,5 @@ impl Iterator for SuperKmerIter<'_> {
|
||||
// ── tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
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());
|
||||
}
|
||||
}
|
||||
#[path = "tests/iter.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
//! Normalized entropy of an isolated, already-built k-mer.
|
||||
//!
|
||||
//! [`SuperKmerIter`](crate::SuperKmerIter) uses [`RollingStat`] to reject
|
||||
//! low-complexity k-mers *during* superkmer construction, incrementally, over
|
||||
//! a streaming window. [`KmerEntropy`] exposes the same metric for a single
|
||||
//! k-mer taken in isolation (e.g. one already reconstructed from an index's
|
||||
//! `unitigs.bin`, with no surrounding sequence) — built on the identical
|
||||
//! `RollingStat` code path, not a re-derived formula, so a `theta` threshold
|
||||
//! chosen for index-build-time filtering means the same thing when applied
|
||||
//! after the fact (e.g. `obikmer filter`).
|
||||
|
||||
use obikseq::CanonicalKmer;
|
||||
|
||||
use crate::rolling_stat::RollingStat;
|
||||
|
||||
/// 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, same as
|
||||
/// [`RollingStat::normalized_entropy`]). Lower means less complex;
|
||||
/// `theta` in `index`/`filter` rejects k-mers with a score `< theta`.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Requires both `obikseq::params::k()` *and* `params::m()` to already be
|
||||
/// set (via `set_k`/`set_m`), even though the minimizer length plays no
|
||||
/// conceptual role in an entropy score: `RollingStat` is a shared,
|
||||
/// general-purpose struct (also used for minimizer selection during
|
||||
/// superkmer decomposition) and unconditionally sizes an `m`-dependent
|
||||
/// mask in its constructor. Callers must call `set_m` with a valid value
|
||||
/// (e.g. the index's own stored `minimizer_size`) even when only the
|
||||
/// entropy score is needed.
|
||||
fn entropy(&self, level_max: usize) -> f64;
|
||||
}
|
||||
|
||||
impl KmerEntropy for CanonicalKmer {
|
||||
fn entropy(&self, level_max: usize) -> f64 {
|
||||
let mut stat = RollingStat::new(level_max);
|
||||
for &base in self.to_ascii().iter() {
|
||||
stat.push(base);
|
||||
}
|
||||
stat.normalized_entropy().unwrap_or(1.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/kmer_entropy.rs"]
|
||||
mod tests;
|
||||
@@ -6,6 +6,7 @@
|
||||
#![deny(missing_docs)]
|
||||
|
||||
pub mod iter;
|
||||
pub mod kmer_entropy;
|
||||
pub mod stream_iter;
|
||||
mod scratch;
|
||||
|
||||
@@ -14,6 +15,7 @@ pub(crate) mod entropy_table;
|
||||
pub(crate) mod rolling_stat;
|
||||
|
||||
pub use iter::SuperKmerIter;
|
||||
pub use kmer_entropy::KmerEntropy;
|
||||
pub use scratch::SuperKmerScratch;
|
||||
pub use stream_iter::SuperKmerStreamIter;
|
||||
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
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());
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
use super::*;
|
||||
use obikseq::Sequence;
|
||||
use obikseq::kmer::Kmer;
|
||||
|
||||
const K: usize = 21;
|
||||
const M: usize = 9; // RollingStat also tracks the minimizer window internally
|
||||
const LEVEL_MAX: usize = 6;
|
||||
|
||||
fn kmer_from_ascii(seq: &[u8]) -> CanonicalKmer {
|
||||
obikseq::set_k(K);
|
||||
obikseq::set_m(M);
|
||||
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:?}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user