Skip to main content

fastx/
seq.rs

1//! Sequence utilities that operate on raw `&[u8]` slices.
2//!
3//! Everything here is case-preserving and allocation-free unless the return type
4//! says otherwise, so these helpers can be used inside hot loops.
5
6use crate::error::{Error, Result};
7
8/// IUPAC pairs shared by the DNA and RNA complement tables.
9const AMBIGUITY_PAIRS: &[(u8, u8)] = &[
10    (b'Y', b'R'),
11    (b'R', b'Y'),
12    (b'S', b'S'),
13    (b'W', b'W'),
14    (b'K', b'M'),
15    (b'M', b'K'),
16    (b'B', b'V'),
17    (b'V', b'B'),
18    (b'D', b'H'),
19    (b'H', b'D'),
20    (b'N', b'N'),
21];
22
23/// Complement table for DNA output: `A` becomes `T`, and `U` becomes `A`.
24static COMPLEMENT: [u8; 256] = build_complement_table(&[
25    (b'A', b'T'),
26    (b'T', b'A'),
27    (b'U', b'A'),
28    (b'G', b'C'),
29    (b'C', b'G'),
30]);
31
32/// Complement table for RNA output: `A` becomes `U`, and both `T` and `U`
33/// become `A`.
34static COMPLEMENT_RNA: [u8; 256] = build_complement_table(&[
35    (b'A', b'U'),
36    (b'U', b'A'),
37    (b'T', b'A'),
38    (b'G', b'C'),
39    (b'C', b'G'),
40]);
41
42const fn build_complement_table(bases: &[(u8, u8)]) -> [u8; 256] {
43    // Bytes that are not nucleotide codes map to themselves.
44    let mut table = [0u8; 256];
45    let mut i = 0;
46    while i < 256 {
47        table[i] = i as u8;
48        i += 1;
49    }
50    let mut p = 0;
51    while p < bases.len() {
52        let (from, to) = bases[p];
53        table[from as usize] = to;
54        table[(from + 32) as usize] = to + 32; // lowercase
55        p += 1;
56    }
57    let mut p = 0;
58    while p < AMBIGUITY_PAIRS.len() {
59        let (from, to) = AMBIGUITY_PAIRS[p];
60        table[from as usize] = to;
61        table[(from + 32) as usize] = to + 32;
62        p += 1;
63    }
64    table
65}
66
67/// Complement a single nucleotide as DNA, preserving case.
68///
69/// Bytes that are not nucleotide codes pass through unchanged, so gaps and
70/// padding survive.
71///
72/// `U` complements to `A`, and `A` complements to `T` — so complementing twice
73/// turns RNA into DNA rather than returning the original byte. That is the only
74/// sensible reading of a DNA complement, but it does mean the operation is not
75/// an involution on RNA. Use [`complement_rna`] to keep uracil.
76///
77/// ```
78/// # use fastx::seq::complement;
79/// assert_eq!(complement(b'A'), b'T');
80/// assert_eq!(complement(b'g'), b'c');
81/// assert_eq!(complement(b'-'), b'-');
82/// assert_eq!(complement(b'U'), b'A');
83/// assert_eq!(complement(complement(b'U')), b'T'); // not 'U'
84/// ```
85#[inline]
86pub fn complement(base: u8) -> u8 {
87    COMPLEMENT[base as usize]
88}
89
90/// Complement a single nucleotide as RNA: `A` becomes `U`, not `T`.
91///
92/// ```
93/// # use fastx::seq::complement_rna;
94/// assert_eq!(complement_rna(b'A'), b'U');
95/// assert_eq!(complement_rna(b'U'), b'A');
96/// assert_eq!(complement_rna(b'T'), b'A');
97/// assert_eq!(complement_rna(complement_rna(b'u')), b'u');
98/// ```
99#[inline]
100pub fn complement_rna(base: u8) -> u8 {
101    COMPLEMENT_RNA[base as usize]
102}
103
104/// Reverse complement of a nucleotide sequence, as DNA.
105///
106/// ```
107/// # use fastx::seq::reverse_complement;
108/// assert_eq!(reverse_complement(b"ACGTn"), b"nACGT");
109/// assert_eq!(reverse_complement(b"ACGU"), b"ACGT"); // uracil becomes thymine
110/// ```
111pub fn reverse_complement(seq: &[u8]) -> Vec<u8> {
112    seq.iter().rev().map(|&b| complement(b)).collect()
113}
114
115/// Reverse complement of a nucleotide sequence, as RNA.
116///
117/// ```
118/// # use fastx::seq::reverse_complement_rna;
119/// assert_eq!(reverse_complement_rna(b"ACGU"), b"ACGU");
120/// assert_eq!(reverse_complement_rna(b"AAGG"), b"CCUU");
121/// ```
122pub fn reverse_complement_rna(seq: &[u8]) -> Vec<u8> {
123    seq.iter().rev().map(|&b| complement_rna(b)).collect()
124}
125
126/// Reverse complement into an existing buffer, which is cleared first.
127///
128/// Use this in loops to avoid one allocation per record.
129pub fn reverse_complement_into(seq: &[u8], out: &mut Vec<u8>) {
130    out.clear();
131    out.reserve(seq.len());
132    out.extend(seq.iter().rev().map(|&b| complement(b)));
133}
134
135/// Reverse complement a sequence in place.
136pub fn reverse_complement_in_place(seq: &mut [u8]) {
137    let n = seq.len();
138    for i in 0..n / 2 {
139        let j = n - 1 - i;
140        let a = complement(seq[i]);
141        let b = complement(seq[j]);
142        seq[i] = b;
143        seq[j] = a;
144    }
145    if n % 2 == 1 {
146        seq[n / 2] = complement(seq[n / 2]);
147    }
148}
149
150/// Per-base counts of a nucleotide sequence, case-insensitive.
151#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
152pub struct BaseCounts {
153    /// Adenine.
154    pub a: u64,
155    /// Cytosine.
156    pub c: u64,
157    /// Guanine.
158    pub g: u64,
159    /// Thymine, or uracil in RNA.
160    pub t: u64,
161    /// `N` and other IUPAC ambiguity codes.
162    pub n: u64,
163    /// Anything that is not a nucleotide code (gaps, `*`, ...).
164    pub other: u64,
165}
166
167/// Class assigned to each byte by [`BASE_CLASS`].
168const CLASS_A: u8 = 0;
169const CLASS_C: u8 = 1;
170const CLASS_G: u8 = 2;
171const CLASS_T: u8 = 3;
172const CLASS_AMBIGUOUS: u8 = 4;
173const CLASS_OTHER: u8 = 5;
174
175/// Maps every byte to a base class, so that counting is one table lookup and one
176/// increment per base rather than a chain of comparisons.
177static BASE_CLASS: [u8; 256] = build_base_class_table();
178
179const fn build_base_class_table() -> [u8; 256] {
180    let mut table = [CLASS_OTHER; 256];
181    table[b'A' as usize] = CLASS_A;
182    table[b'a' as usize] = CLASS_A;
183    table[b'C' as usize] = CLASS_C;
184    table[b'c' as usize] = CLASS_C;
185    table[b'G' as usize] = CLASS_G;
186    table[b'g' as usize] = CLASS_G;
187    table[b'T' as usize] = CLASS_T;
188    table[b't' as usize] = CLASS_T;
189    table[b'U' as usize] = CLASS_T;
190    table[b'u' as usize] = CLASS_T;
191
192    let ambiguous = b"NRYSWKMBDHV";
193    let mut i = 0;
194    while i < ambiguous.len() {
195        table[ambiguous[i] as usize] = CLASS_AMBIGUOUS;
196        table[(ambiguous[i] + 32) as usize] = CLASS_AMBIGUOUS;
197        i += 1;
198    }
199    table
200}
201
202impl BaseCounts {
203    /// Count the bases of `seq`.
204    pub fn of(seq: &[u8]) -> BaseCounts {
205        // Eight slots rather than six so that the index is provably in range and
206        // the bounds check disappears.
207        let mut counts = [0u64; 8];
208        for &b in seq {
209            counts[(BASE_CLASS[b as usize] & 7) as usize] += 1;
210        }
211        BaseCounts {
212            a: counts[CLASS_A as usize],
213            c: counts[CLASS_C as usize],
214            g: counts[CLASS_G as usize],
215            t: counts[CLASS_T as usize],
216            n: counts[CLASS_AMBIGUOUS as usize],
217            other: counts[CLASS_OTHER as usize],
218        }
219    }
220
221    /// Total number of counted bytes.
222    pub fn total(&self) -> u64 {
223        self.a + self.c + self.g + self.t + self.n + self.other
224    }
225
226    /// Unambiguous A/C/G/T bases.
227    pub fn acgt(&self) -> u64 {
228        self.a + self.c + self.g + self.t
229    }
230
231    /// GC fraction over unambiguous bases, or `None` when there are none.
232    pub fn gc_content(&self) -> Option<f64> {
233        let acgt = self.acgt();
234        if acgt == 0 {
235            None
236        } else {
237            Some((self.g + self.c) as f64 / acgt as f64)
238        }
239    }
240
241    /// Add another set of counts.
242    pub fn merge(&mut self, other: &BaseCounts) {
243        self.a += other.a;
244        self.c += other.c;
245        self.g += other.g;
246        self.t += other.t;
247        self.n += other.n;
248        self.other += other.other;
249    }
250}
251
252/// GC fraction of a sequence, ignoring ambiguity codes.
253///
254/// Returns `None` when the sequence has no unambiguous bases.
255///
256/// ```
257/// # use fastx::seq::gc_content;
258/// assert_eq!(gc_content(b"GGCC"), Some(1.0));
259/// assert_eq!(gc_content(b"ATAT"), Some(0.0));
260/// assert_eq!(gc_content(b"GCATNNNN"), Some(0.5));
261/// assert_eq!(gc_content(b"NNNN"), None);
262/// ```
263pub fn gc_content(seq: &[u8]) -> Option<f64> {
264    BaseCounts::of(seq).gc_content()
265}
266
267/// A residue alphabet used for validation.
268#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
269pub enum Alphabet {
270    /// `ACGT` plus `N`, either case.
271    Dna,
272    /// `ACGU` plus `N`, either case.
273    Rna,
274    /// The full IUPAC nucleotide set plus `-` and `.` gaps.
275    Iupac,
276    /// The 20 amino acids plus `BXZJUO*-.`, either case.
277    Protein,
278    /// Any printable, non-whitespace byte.
279    Any,
280}
281
282impl Alphabet {
283    /// Whether `byte` is a member of the alphabet.
284    pub fn contains(self, byte: u8) -> bool {
285        let upper = byte.to_ascii_uppercase();
286        match self {
287            Alphabet::Dna => matches!(upper, b'A' | b'C' | b'G' | b'T' | b'N'),
288            Alphabet::Rna => matches!(upper, b'A' | b'C' | b'G' | b'U' | b'N'),
289            Alphabet::Iupac => {
290                matches!(
291                    upper,
292                    b'A' | b'C'
293                        | b'G'
294                        | b'T'
295                        | b'U'
296                        | b'R'
297                        | b'Y'
298                        | b'S'
299                        | b'W'
300                        | b'K'
301                        | b'M'
302                        | b'B'
303                        | b'D'
304                        | b'H'
305                        | b'V'
306                        | b'N'
307                        | b'-'
308                        | b'.'
309                )
310            }
311            Alphabet::Protein => upper.is_ascii_uppercase() || matches!(upper, b'*' | b'-' | b'.'),
312            Alphabet::Any => byte.is_ascii_graphic(),
313        }
314    }
315
316    /// Check every byte of `seq`, reporting the first violation.
317    pub fn validate(self, seq: &[u8]) -> Result<()> {
318        self.validate_named(seq, "")
319    }
320
321    /// Like [`Alphabet::validate`] but attaches a record id to the error.
322    pub fn validate_named(self, seq: &[u8], id: &str) -> Result<()> {
323        match seq.iter().position(|&b| !self.contains(b)) {
324            None => Ok(()),
325            Some(pos) => Err(Error::InvalidByte {
326                id: id.to_string(),
327                pos,
328                byte: seq[pos],
329            }),
330        }
331    }
332}
333
334/// Uppercase a sequence in place (soft-masked genomes use lowercase for repeats).
335pub fn make_uppercase(seq: &mut [u8]) {
336    seq.make_ascii_uppercase();
337}
338
339/// Replace `U`/`u` with `T`/`t`, turning RNA into DNA in place.
340pub fn rna_to_dna(seq: &mut [u8]) {
341    for b in seq.iter_mut() {
342        match *b {
343            b'U' => *b = b'T',
344            b'u' => *b = b't',
345            _ => {}
346        }
347    }
348}
349
350/// The standard genetic code (NCBI translation table 1), in `TCAG` codon order.
351const CODON_TABLE: &[u8; 64] = b"FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG";
352
353/// Sentinel stored in [`CODON_INDEX`] for bytes that are not `T`/`U`, `C`, `A` or `G`.
354const NOT_A_BASE: u8 = u8::MAX;
355
356/// Maps every byte to its position in the `TCAG` codon ordering.
357static CODON_INDEX: [u8; 256] = build_codon_index_table();
358
359const fn build_codon_index_table() -> [u8; 256] {
360    let mut table = [NOT_A_BASE; 256];
361    table[b'T' as usize] = 0;
362    table[b't' as usize] = 0;
363    table[b'U' as usize] = 0;
364    table[b'u' as usize] = 0;
365    table[b'C' as usize] = 1;
366    table[b'c' as usize] = 1;
367    table[b'A' as usize] = 2;
368    table[b'a' as usize] = 2;
369    table[b'G' as usize] = 3;
370    table[b'g' as usize] = 3;
371    table
372}
373
374#[inline]
375fn base_index(base: u8) -> Option<usize> {
376    match CODON_INDEX[base as usize] {
377        NOT_A_BASE => None,
378        index => Some(index as usize),
379    }
380}
381
382/// Translate a single codon using the standard genetic code.
383///
384/// Stop codons become `*` and codons containing an ambiguous base become `X`.
385///
386/// ```
387/// # use fastx::seq::translate_codon;
388/// assert_eq!(translate_codon(b"ATG"), b'M');
389/// assert_eq!(translate_codon(b"tga"), b'*');
390/// assert_eq!(translate_codon(b"ANG"), b'X');
391/// ```
392pub fn translate_codon(codon: &[u8]) -> u8 {
393    if codon.len() < 3 {
394        return b'X';
395    }
396    match (
397        base_index(codon[0]),
398        base_index(codon[1]),
399        base_index(codon[2]),
400    ) {
401        (Some(a), Some(b), Some(c)) => CODON_TABLE[a * 16 + b * 4 + c],
402        _ => b'X',
403    }
404}
405
406/// Translate a nucleotide sequence into a protein sequence.
407///
408/// `frame` is 0, 1 or 2 and skips that many leading bases. A trailing partial
409/// codon is dropped. When `stop_at_stop` is true, translation ends before the
410/// first stop codon.
411///
412/// ```
413/// # use fastx::seq::translate;
414/// assert_eq!(translate(b"ATGGCCTGA", 0, false), b"MA*");
415/// assert_eq!(translate(b"ATGGCCTGA", 0, true), b"MA");
416/// assert_eq!(translate(b"AATGGCC", 1, false), b"MA");
417/// ```
418pub fn translate(seq: &[u8], frame: usize, stop_at_stop: bool) -> Vec<u8> {
419    let seq = if frame < seq.len() {
420        &seq[frame..]
421    } else {
422        &[][..]
423    };
424    let mut out = Vec::with_capacity(seq.len() / 3);
425    for codon in seq.chunks_exact(3) {
426        let aa = translate_codon(codon);
427        if stop_at_stop && aa == b'*' {
428            break;
429        }
430        out.push(aa);
431    }
432    out
433}
434
435/// Iterator over the overlapping k-mers of a sequence.
436///
437/// ```
438/// # use fastx::seq::kmers;
439/// let all: Vec<&[u8]> = kmers(b"ACGTA", 3).collect();
440/// assert_eq!(all, vec![&b"ACG"[..], &b"CGT"[..], &b"GTA"[..]]);
441/// assert_eq!(kmers(b"AC", 3).count(), 0);
442/// ```
443pub fn kmers(seq: &[u8], k: usize) -> impl Iterator<Item = &[u8]> {
444    let n = if k == 0 || seq.len() < k {
445        0
446    } else {
447        seq.len() - k + 1
448    };
449    (0..n).map(move |i| &seq[i..i + k])
450}
451
452/// Canonical k-mer: the lexicographically smaller of a k-mer and its reverse
453/// complement, so that both strands hash to the same value.
454pub fn canonical_kmer(kmer: &[u8]) -> Vec<u8> {
455    let rc = reverse_complement(kmer);
456    if rc.as_slice() < kmer {
457        rc
458    } else {
459        kmer.to_vec()
460    }
461}
462
463/// Hamming distance between two equal-length sequences, case-insensitive.
464///
465/// Returns `None` if the lengths differ.
466pub fn hamming_distance(a: &[u8], b: &[u8]) -> Option<usize> {
467    if a.len() != b.len() {
468        return None;
469    }
470    Some(
471        a.iter()
472            .zip(b)
473            .filter(|(x, y)| !x.eq_ignore_ascii_case(y))
474            .count(),
475    )
476}
477
478/// N50 of a set of lengths: the length `L` such that contigs of at least `L`
479/// cover half of the total assembly length.
480///
481/// ```
482/// # use fastx::seq::n50;
483/// assert_eq!(n50(&mut [2, 3, 4, 5, 6]), Some(5));
484/// assert_eq!(n50(&mut []), None);
485/// ```
486pub fn n50(lengths: &mut [u64]) -> Option<u64> {
487    nx(lengths, 0.5)
488}
489
490/// Generalised N-statistic: `nx(lengths, 0.9)` is the N90.
491///
492/// `fraction` is clamped to `0.0..=1.0`. The slice is sorted as a side effect.
493pub fn nx(lengths: &mut [u64], fraction: f64) -> Option<u64> {
494    if lengths.is_empty() {
495        return None;
496    }
497    let total: u64 = lengths.iter().sum();
498    if total == 0 {
499        return Some(0);
500    }
501    lengths.sort_unstable_by(|a, b| b.cmp(a));
502    let target = total as f64 * fraction.clamp(0.0, 1.0);
503    let mut acc = 0u64;
504    for &len in lengths.iter() {
505        acc += len;
506        if acc as f64 >= target {
507            return Some(len);
508        }
509    }
510    lengths.last().copied()
511}
512
513#[cfg(test)]
514mod tests {
515    use super::*;
516
517    #[test]
518    fn complements_iupac_and_preserves_case() {
519        assert_eq!(reverse_complement(b"ACGTacgt"), b"acgtACGT");
520        assert_eq!(reverse_complement(b"RYKM"), b"KMRY");
521        assert_eq!(reverse_complement(b""), b"");
522        // Non-nucleotide bytes are preserved.
523        assert_eq!(reverse_complement(b"AC-GT"), b"AC-GT");
524    }
525
526    #[test]
527    fn in_place_matches_allocating() {
528        for s in [&b""[..], b"A", b"AC", b"ACG", b"ACGTN", b"acgtRYn"] {
529            let mut owned = s.to_vec();
530            reverse_complement_in_place(&mut owned);
531            assert_eq!(
532                owned,
533                reverse_complement(s),
534                "{:?}",
535                String::from_utf8_lossy(s)
536            );
537        }
538    }
539
540    #[test]
541    fn counts_bases() {
542        let c = BaseCounts::of(b"AACCGGTTNNxx");
543        assert_eq!((c.a, c.c, c.g, c.t, c.n, c.other), (2, 2, 2, 2, 2, 2));
544        assert_eq!(c.total(), 12);
545        assert_eq!(c.acgt(), 8);
546        assert_eq!(c.gc_content(), Some(0.5));
547    }
548
549    #[test]
550    fn validates_alphabets() {
551        assert!(Alphabet::Dna.validate(b"ACGTN").is_ok());
552        assert!(Alphabet::Dna.validate(b"ACGU").is_err());
553        assert!(Alphabet::Rna.validate(b"ACGU").is_ok());
554        assert!(Alphabet::Iupac.validate(b"ACGTRYKM-").is_ok());
555        assert!(Alphabet::Protein.validate(b"MEEPQSDPSV*").is_ok());
556        assert!(Alphabet::Any.validate(b"anything!").is_ok());
557        assert!(Alphabet::Any.validate(b"tab\there").is_err());
558
559        match Alphabet::Dna.validate_named(b"ACG!T", "read1") {
560            Err(Error::InvalidByte { id, pos, byte }) => {
561                assert_eq!((id.as_str(), pos, byte), ("read1", 3, b'!'));
562            }
563            other => panic!("expected InvalidByte, got {other:?}"),
564        }
565    }
566
567    #[test]
568    fn translates_all_codons() {
569        // Spot-check every position of the table via the classic ordering.
570        assert_eq!(translate(b"TTTTTCTTATTG", 0, false), b"FFLL");
571        assert_eq!(translate(b"ATGCATTAA", 0, false), b"MH*");
572        assert_eq!(translate(b"AUGCAU", 0, false), b"MH"); // RNA input
573                                                           // Trailing partial codon dropped.
574        assert_eq!(translate(b"ATGCA", 0, false), b"M");
575        assert_eq!(translate(b"AT", 0, false), b"");
576        assert_eq!(translate(b"ATG", 5, false), b"");
577    }
578
579    #[test]
580    fn kmers_and_canonical() {
581        assert_eq!(kmers(b"AAAA", 4).count(), 1);
582        assert_eq!(kmers(b"AAAA", 0).count(), 0);
583        assert_eq!(canonical_kmer(b"TTT"), b"AAA");
584        assert_eq!(canonical_kmer(b"AAA"), b"AAA");
585    }
586
587    #[test]
588    fn hamming() {
589        assert_eq!(hamming_distance(b"ACGT", b"acgt"), Some(0));
590        assert_eq!(hamming_distance(b"ACGT", b"ACGA"), Some(1));
591        assert_eq!(hamming_distance(b"ACGT", b"ACG"), None);
592    }
593
594    #[test]
595    fn n_statistics() {
596        // Total 100; sorted desc 50,30,15,5 -> 50 reaches 50.
597        assert_eq!(n50(&mut [5, 50, 30, 15]), Some(50));
598        assert_eq!(nx(&mut [5, 50, 30, 15], 0.9), Some(15));
599        assert_eq!(nx(&mut [0, 0], 0.5), Some(0));
600    }
601
602    #[test]
603    fn rna_dna_conversion() {
604        let mut s = b"ACGUacgu".to_vec();
605        rna_to_dna(&mut s);
606        assert_eq!(s, b"ACGTacgt");
607        make_uppercase(&mut s);
608        assert_eq!(s, b"ACGTACGT");
609    }
610}