Add skip_ambiguous_windows variants

Author: RagnarGrootKoerkamp

CHANGELOG.md

@@ -1,14 +1,13 @@
 # Changelog
 
 ## 2.0
+- **Breaking**: Migrate to `packed-seq` `4.0` with `PaddedIt`.
+- **Breaking**: Improve simd versions for short sequences by reusing allocated buffers.
+- **Feature**: Use `seq-hash` crate to cleanly support multiple kmer hashers; add `*_with_hasher` function variants.
+- **Feature**: Add support for `_skip_ambiguous_windows` variants that omit minimizers for windows containing non-ACTG (such as NYR).
 - Move `simd-minimizers` crate from subdir into the repo root.
-- Rename `simd-minimizers-bench` to just `bench`.
 - Cleanups (use Rust 2024 edition; drop some public-but-unused functions).
 - Improve simd versions for short sequences by preventing over-allocating the output vector.
-- **Breaking**: Migrate to `packed-seq` `4.0` with `PaddedIt`.
-- **Breaking**: Improve simd versions for short sequences by reusing allocated buffers.
-- **Feature**: Use `seq-hash` crate to cleanly support multiple kmer hashers;
-  add `*_with_hasher` function variants.
 
 ## 1.4
 - Make `NtHash` and `MulHash` seedable.

src/lib.rs

@@ -178,8 +178,11 @@ pub mod private {
 use collect::CollectAndDedup;
 use collect::collect_and_dedup_into_scalar;
 use collect::collect_and_dedup_with_index_into_scalar;
+use minimizers::canonical_minimizers_skip_ambiguous_windows;
 /// Re-export of the `packed-seq` crate.
 pub use packed_seq;
+use packed_seq::BitSeq;
+use packed_seq::PackedSeq;
 /// Re-export of the `seq-hash` crate.
 pub use seq_hash;
 
@@ -316,6 +319,35 @@ impl<'h, const CANONICAL: bool, H: KmerHasher> Builder<'h, CANONICAL, H, ()> {
     }
 }
 
+impl<'h, H: KmerHasher> Builder<'h, true, H, ()> {
+    pub fn run_skip_ambiguous_windows_once<'s, 'o>(&self, seq: PackedSeq<'s>, ambi: BitSeq<'s>) -> Vec<u32> {
+        let mut min_pos = vec![];
+        self.run_skip_ambiguous_windows(seq, ambi, &mut min_pos);
+        min_pos
+    }
+    pub fn run_skip_ambiguous_windows<'s, 'o>(
+        &self,
+        seq: PackedSeq<'s>,
+        ambi: BitSeq<'s>,
+        min_pos: &'o mut Vec<u32>,
+    ) -> Output<'o, true, PackedSeq<'s>> {
+        let default_hasher = self.hasher.is_none().then(|| H::new(self.k));
+        let hasher = self
+            .hasher
+            .unwrap_or_else(|| default_hasher.as_ref().unwrap());
+
+        CACHE.with_borrow_mut(|cache| {
+            canonical_minimizers_skip_ambiguous_windows(seq, ambi, hasher, self.w, cache)
+                .collect_and_dedup_into::<true>(min_pos)
+        });
+        Output {
+            k: self.k,
+            seq,
+            min_pos,
+        }
+    }
+}
+
 /// With-superkmer version
 impl<'h, 'o2, const CANONICAL: bool, H: KmerHasher> Builder<'h, CANONICAL, H, &'o2 mut Vec<u32>> {
     pub fn run_scalar_once<'s, SEQ: Seq<'s>>(self, seq: SEQ) -> Vec<u32> {

src/minimizers.rs

@@ -165,3 +165,43 @@ pub fn canonical_minimizers_seq_simd<'s>(
         canonical.blend(lmin, rmin)
     })
 }
+
+/// Use canonical NtHash, and keep both leftmost and rightmost minima.
+#[inline(always)]
+pub fn canonical_minimizers_skip_ambiguous_windows<'s>(
+    seq: packed_seq::PackedSeq<'s>,
+    bitseq: packed_seq::BitSeq<'s>,
+    hasher: &impl KmerHasher,
+    w: usize,
+    cache: &mut Cache,
+) -> PaddedIt<impl ChunkIt<u32x8>> {
+    assert!(hasher.is_canonical());
+
+    let k = hasher.k();
+    let l = k + w - 1;
+    let mut hash_mapper = hasher.in_out_mapper_simd(seq);
+    let (c_delay, mut canonical_mapper) = canonical_mapper_simd(l);
+
+    let mut padded_it = seq
+        .par_iter_bp_delayed_2_with_factor(l, hasher.delay(), c_delay, 2)
+        .zip(bitseq.par_iter_kmer_ambiguity(l, l));
+
+    // Process first k-1 characters separately, to initialize hash values.
+    padded_it.advance_with(k - 1, |((a, rh, rc), _ambi)| {
+        hash_mapper((a, rh));
+        canonical_mapper((a, rc));
+    });
+    let mut sliding_min_mapper = sliding_lr_min_mapper_simd(w, padded_it.it.len(), cache);
+    padded_it.advance_with(w - 1, |((a, rh, rc), _ambi)| {
+        let hash = hash_mapper((a, rh));
+        canonical_mapper((a, rc));
+        sliding_min_mapper(hash);
+    });
+
+    padded_it.map(move |((a, rh, rc), ambi)| {
+        let hash = hash_mapper((a, rh));
+        let canonical = canonical_mapper((a, rc));
+        let (lmin, rmin) = sliding_min_mapper(hash);
+        ambi.blend(u32x8::splat(u32::MAX - 1), canonical.blend(lmin, rmin))
+    })
+}

src/test.rs

@@ -1,8 +1,8 @@
 use super::*;
 use crate::minimizers::*;
 use itertools::Itertools;
-use packed_seq::{AsciiSeq, AsciiSeqVec, PackedSeq, PackedSeqVec, SeqVec};
-use rand::Rng;
+use packed_seq::{AsciiSeq, AsciiSeqVec, BitSeqVec, PackedSeq, PackedSeqVec, PaddedIt, SeqVec};
+use rand::{Rng, random_range};
 use seq_hash::{AntiLexHasher, MulHasher, NtHasher};
 use std::sync::LazyLock;
 
@@ -409,3 +409,64 @@ fn readme_example() {
         .values_u64()
         .collect();
 }
+
+#[test]
+fn skip_ambiguous() {
+    let len = 100;
+    let mut ascii = AsciiSeqVec::random(len);
+    // set 1% to N
+    for _ in 0..len / 100 {
+        ascii.seq[random_range(0..len)] = b'N';
+    }
+    let packed_seq = PackedSeqVec::from_ascii(&ascii.seq);
+    let ambiguous = BitSeqVec::from_ascii(&ascii.seq);
+
+    for k in 1..=64 {
+        for w in 1..64 {
+            if (k + w - 1) % 2 == 0 {
+                continue;
+            }
+            if k + w - 1 > 64 {
+                continue;
+            }
+            let hasher = &<NtHasher>::new(k);
+
+            let mut poss0 = vec![];
+            canonical_minimizers_skip_ambiguous_windows(
+                packed_seq.as_slice(),
+                ambiguous.as_slice(),
+                hasher,
+                w,
+                &mut Default::default(),
+            )
+            .collect_into(&mut poss0);
+            let mut poss1 = vec![];
+            canonical_minimizers_skip_ambiguous_windows(
+                packed_seq.as_slice(),
+                ambiguous.as_slice(),
+                hasher,
+                w,
+                &mut Default::default(),
+            )
+            .collect_and_dedup_into::<false>(&mut poss1);
+            let mut poss2 = vec![];
+            canonical_minimizers_skip_ambiguous_windows(
+                packed_seq.as_slice(),
+                ambiguous.as_slice(),
+                hasher,
+                w,
+                &mut Default::default(),
+            )
+            .collect_and_dedup_into::<true>(&mut poss2);
+
+            let poss = canonical_minimizers(k, w)
+                .run_skip_ambiguous_windows_once(packed_seq.as_slice(), ambiguous.as_slice());
+            for pos in poss {
+                // these should be filtered out
+                assert_ne!(pos, u32::MAX - 1);
+                // check that kmer at pos does not have ambiguous bases
+                assert_eq!(ambiguous.read_kmer_u128(k, pos as usize), 0);
+            }
+        }
+    }
+}