Skip to main content

poa_consensus/
orient.rs

1use crate::types::Strand;
2use std::borrow::Cow;
3
4pub fn reverse_complement(read: &[u8]) -> Vec<u8> {
5    read.iter().rev().map(|&b| complement(b)).collect()
6}
7
8#[inline]
9fn complement(b: u8) -> u8 {
10    match b {
11        b'A' | b'a' => b'T',
12        b'T' | b't' => b'A',
13        b'C' | b'c' => b'G',
14        b'G' | b'g' => b'C',
15        other => other,
16    }
17}
18
19/// Determine the strand of `read` relative to `seed` using k-mer matching.
20///
21/// Counts shared k-mers between `read` and `seed` in the forward orientation,
22/// then between `reverse_complement(read)` and `seed`. Returns `Forward` if the
23/// forward count is greater or equal, `Reverse` otherwise.
24///
25/// Complexity: O(n) in the length of both sequences using a hash map.
26pub fn orient_to_seed(read: &[u8], seed: &[u8], k: usize) -> Strand {
27    if read.len() < k || seed.len() < k {
28        return Strand::Forward;
29    }
30
31    use std::collections::HashMap;
32
33    let mut seed_kmers: HashMap<&[u8], u32> = HashMap::new();
34    for w in seed.windows(k) {
35        *seed_kmers.entry(w).or_insert(0) += 1;
36    }
37
38    let fwd_count: u32 = read
39        .windows(k)
40        .map(|w| seed_kmers.get(w).copied().unwrap_or(0))
41        .sum();
42
43    let rc = reverse_complement(read);
44    let rev_count: u32 = rc
45        .windows(k)
46        .map(|w| seed_kmers.get(w).copied().unwrap_or(0))
47        .sum();
48
49    if fwd_count >= rev_count {
50        Strand::Forward
51    } else {
52        Strand::Reverse
53    }
54}
55
56/// Orient all reads to match the strand of `reads[seed_idx]`.
57///
58/// Returns a `Vec<Cow<[u8]>>` where each element borrows the original read
59/// if it is already in the forward orientation, or owns the reverse complement
60/// if it was flipped.
61pub fn auto_orient<'a>(reads: &'a [Vec<u8>], seed_idx: usize) -> Vec<Cow<'a, [u8]>> {
62    let seed = &reads[seed_idx];
63    let k = 8.min(seed.len().saturating_sub(1).max(1));
64
65    reads
66        .iter()
67        .map(|read| match orient_to_seed(read, seed, k) {
68            Strand::Forward => Cow::Borrowed(read.as_slice()),
69            Strand::Reverse => Cow::Owned(reverse_complement(read)),
70        })
71        .collect()
72}