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///
269/// Marked `#[non_exhaustive]`: alphabets are a judgement call, not a closed set,
270/// so adding one should not break callers. Alphabets are normally passed in
271/// rather than matched on, which makes the cost of the annotation close to zero.
272#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
273#[non_exhaustive]
274pub enum Alphabet {
275    /// `ACGT` plus `N`, either case.
276    Dna,
277    /// `ACGU` plus `N`, either case.
278    Rna,
279    /// The full IUPAC nucleotide set plus `-` and `.` gaps.
280    Iupac,
281    /// The 20 amino acids plus `BXZJUO*-.`, either case.
282    Protein,
283    /// Any printable, non-whitespace byte.
284    Any,
285}
286
287impl Alphabet {
288    /// Whether `byte` is a member of the alphabet.
289    pub fn contains(self, byte: u8) -> bool {
290        let upper = byte.to_ascii_uppercase();
291        match self {
292            Alphabet::Dna => matches!(upper, b'A' | b'C' | b'G' | b'T' | b'N'),
293            Alphabet::Rna => matches!(upper, b'A' | b'C' | b'G' | b'U' | b'N'),
294            Alphabet::Iupac => {
295                matches!(
296                    upper,
297                    b'A' | b'C'
298                        | b'G'
299                        | b'T'
300                        | b'U'
301                        | b'R'
302                        | b'Y'
303                        | b'S'
304                        | b'W'
305                        | b'K'
306                        | b'M'
307                        | b'B'
308                        | b'D'
309                        | b'H'
310                        | b'V'
311                        | b'N'
312                        | b'-'
313                        | b'.'
314                )
315            }
316            Alphabet::Protein => upper.is_ascii_uppercase() || matches!(upper, b'*' | b'-' | b'.'),
317            Alphabet::Any => byte.is_ascii_graphic(),
318        }
319    }
320
321    /// Check every byte of `seq`, reporting the first violation.
322    pub fn validate(self, seq: &[u8]) -> Result<()> {
323        self.validate_named(seq, "")
324    }
325
326    /// Like [`Alphabet::validate`] but attaches a record id to the error.
327    pub fn validate_named(self, seq: &[u8], id: &str) -> Result<()> {
328        match seq.iter().position(|&b| !self.contains(b)) {
329            None => Ok(()),
330            Some(pos) => Err(Error::InvalidByte {
331                id: id.to_string(),
332                pos,
333                byte: seq[pos],
334            }),
335        }
336    }
337}
338
339/// Uppercase a sequence in place (soft-masked genomes use lowercase for repeats).
340pub fn make_uppercase(seq: &mut [u8]) {
341    seq.make_ascii_uppercase();
342}
343
344/// Replace `U`/`u` with `T`/`t`, turning RNA into DNA in place.
345pub fn rna_to_dna(seq: &mut [u8]) {
346    for b in seq.iter_mut() {
347        match *b {
348            b'U' => *b = b'T',
349            b'u' => *b = b't',
350            _ => {}
351        }
352    }
353}
354
355/// The standard genetic code (NCBI translation table 1), in `TCAG` codon order.
356const CODON_TABLE: &[u8; 64] = b"FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG";
357
358/// Sentinel stored in [`CODON_INDEX`] for bytes that are not `T`/`U`, `C`, `A` or `G`.
359const NOT_A_BASE: u8 = u8::MAX;
360
361/// Maps every byte to its position in the `TCAG` codon ordering.
362static CODON_INDEX: [u8; 256] = build_codon_index_table();
363
364const fn build_codon_index_table() -> [u8; 256] {
365    let mut table = [NOT_A_BASE; 256];
366    table[b'T' as usize] = 0;
367    table[b't' as usize] = 0;
368    table[b'U' as usize] = 0;
369    table[b'u' as usize] = 0;
370    table[b'C' as usize] = 1;
371    table[b'c' as usize] = 1;
372    table[b'A' as usize] = 2;
373    table[b'a' as usize] = 2;
374    table[b'G' as usize] = 3;
375    table[b'g' as usize] = 3;
376    table
377}
378
379#[inline]
380fn base_index(base: u8) -> Option<usize> {
381    match CODON_INDEX[base as usize] {
382        NOT_A_BASE => None,
383        index => Some(index as usize),
384    }
385}
386
387/// Translate a single codon using the standard genetic code.
388///
389/// Stop codons become `*` and codons containing an ambiguous base become `X`.
390///
391/// ```
392/// # use fastx::seq::translate_codon;
393/// assert_eq!(translate_codon(b"ATG"), b'M');
394/// assert_eq!(translate_codon(b"tga"), b'*');
395/// assert_eq!(translate_codon(b"ANG"), b'X');
396/// ```
397pub fn translate_codon(codon: &[u8]) -> u8 {
398    if codon.len() < 3 {
399        return b'X';
400    }
401    match (
402        base_index(codon[0]),
403        base_index(codon[1]),
404        base_index(codon[2]),
405    ) {
406        (Some(a), Some(b), Some(c)) => CODON_TABLE[a * 16 + b * 4 + c],
407        _ => b'X',
408    }
409}
410
411/// Translate a nucleotide sequence into a protein sequence.
412///
413/// `frame` is 0, 1 or 2 and skips that many leading bases. A trailing partial
414/// codon is dropped. When `stop_at_stop` is true, translation ends before the
415/// first stop codon.
416///
417/// ```
418/// # use fastx::seq::translate;
419/// assert_eq!(translate(b"ATGGCCTGA", 0, false), b"MA*");
420/// assert_eq!(translate(b"ATGGCCTGA", 0, true), b"MA");
421/// assert_eq!(translate(b"AATGGCC", 1, false), b"MA");
422/// ```
423pub fn translate(seq: &[u8], frame: usize, stop_at_stop: bool) -> Vec<u8> {
424    let seq = if frame < seq.len() {
425        &seq[frame..]
426    } else {
427        &[][..]
428    };
429    let mut out = Vec::with_capacity(seq.len() / 3);
430    for codon in seq.chunks_exact(3) {
431        let aa = translate_codon(codon);
432        if stop_at_stop && aa == b'*' {
433            break;
434        }
435        out.push(aa);
436    }
437    out
438}
439
440/// Iterator over the overlapping k-mers of a sequence.
441///
442/// ```
443/// # use fastx::seq::kmers;
444/// let all: Vec<&[u8]> = kmers(b"ACGTA", 3).collect();
445/// assert_eq!(all, vec![&b"ACG"[..], &b"CGT"[..], &b"GTA"[..]]);
446/// assert_eq!(kmers(b"AC", 3).count(), 0);
447/// ```
448pub fn kmers(seq: &[u8], k: usize) -> impl Iterator<Item = &[u8]> {
449    let n = if k == 0 || seq.len() < k {
450        0
451    } else {
452        seq.len() - k + 1
453    };
454    (0..n).map(move |i| &seq[i..i + k])
455}
456
457/// Canonical k-mer: the lexicographically smaller of a k-mer and its reverse
458/// complement, so that both strands hash to the same value.
459pub fn canonical_kmer(kmer: &[u8]) -> Vec<u8> {
460    let rc = reverse_complement(kmer);
461    if rc.as_slice() < kmer {
462        rc
463    } else {
464        kmer.to_vec()
465    }
466}
467
468/// Hamming distance between two equal-length sequences, case-insensitive.
469///
470/// Returns `None` if the lengths differ.
471pub fn hamming_distance(a: &[u8], b: &[u8]) -> Option<usize> {
472    if a.len() != b.len() {
473        return None;
474    }
475    Some(
476        a.iter()
477            .zip(b)
478            .filter(|(x, y)| !x.eq_ignore_ascii_case(y))
479            .count(),
480    )
481}
482
483/// N50 of a set of lengths: the length `L` such that contigs of at least `L`
484/// cover half of the total assembly length.
485///
486/// ```
487/// # use fastx::seq::n50;
488/// assert_eq!(n50(&mut [2, 3, 4, 5, 6]), Some(5));
489/// assert_eq!(n50(&mut []), None);
490/// ```
491pub fn n50(lengths: &mut [u64]) -> Option<u64> {
492    nx(lengths, 0.5)
493}
494
495/// Generalised N-statistic: `nx(lengths, 0.9)` is the N90.
496///
497/// `fraction` is clamped to `0.0..=1.0`. The slice is sorted as a side effect.
498pub fn nx(lengths: &mut [u64], fraction: f64) -> Option<u64> {
499    if lengths.is_empty() {
500        return None;
501    }
502    let total: u64 = lengths.iter().sum();
503    if total == 0 {
504        return Some(0);
505    }
506    lengths.sort_unstable_by(|a, b| b.cmp(a));
507    let target = total as f64 * fraction.clamp(0.0, 1.0);
508    let mut acc = 0u64;
509    for &len in lengths.iter() {
510        acc += len;
511        if acc as f64 >= target {
512            return Some(len);
513        }
514    }
515    lengths.last().copied()
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521
522    #[test]
523    fn complements_iupac_and_preserves_case() {
524        assert_eq!(reverse_complement(b"ACGTacgt"), b"acgtACGT");
525        assert_eq!(reverse_complement(b"RYKM"), b"KMRY");
526        assert_eq!(reverse_complement(b""), b"");
527        // Non-nucleotide bytes are preserved.
528        assert_eq!(reverse_complement(b"AC-GT"), b"AC-GT");
529    }
530
531    #[test]
532    fn in_place_matches_allocating() {
533        for s in [&b""[..], b"A", b"AC", b"ACG", b"ACGTN", b"acgtRYn"] {
534            let mut owned = s.to_vec();
535            reverse_complement_in_place(&mut owned);
536            assert_eq!(
537                owned,
538                reverse_complement(s),
539                "{:?}",
540                String::from_utf8_lossy(s)
541            );
542        }
543    }
544
545    #[test]
546    fn counts_bases() {
547        let c = BaseCounts::of(b"AACCGGTTNNxx");
548        assert_eq!((c.a, c.c, c.g, c.t, c.n, c.other), (2, 2, 2, 2, 2, 2));
549        assert_eq!(c.total(), 12);
550        assert_eq!(c.acgt(), 8);
551        assert_eq!(c.gc_content(), Some(0.5));
552    }
553
554    #[test]
555    fn validates_alphabets() {
556        assert!(Alphabet::Dna.validate(b"ACGTN").is_ok());
557        assert!(Alphabet::Dna.validate(b"ACGU").is_err());
558        assert!(Alphabet::Rna.validate(b"ACGU").is_ok());
559        assert!(Alphabet::Iupac.validate(b"ACGTRYKM-").is_ok());
560        assert!(Alphabet::Protein.validate(b"MEEPQSDPSV*").is_ok());
561        assert!(Alphabet::Any.validate(b"anything!").is_ok());
562        assert!(Alphabet::Any.validate(b"tab\there").is_err());
563
564        match Alphabet::Dna.validate_named(b"ACG!T", "read1") {
565            Err(Error::InvalidByte { id, pos, byte }) => {
566                assert_eq!((id.as_str(), pos, byte), ("read1", 3, b'!'));
567            }
568            other => panic!("expected InvalidByte, got {other:?}"),
569        }
570    }
571
572    #[test]
573    fn translates_all_codons() {
574        // Spot-check every position of the table via the classic ordering.
575        assert_eq!(translate(b"TTTTTCTTATTG", 0, false), b"FFLL");
576        assert_eq!(translate(b"ATGCATTAA", 0, false), b"MH*");
577        assert_eq!(translate(b"AUGCAU", 0, false), b"MH"); // RNA input
578                                                           // Trailing partial codon dropped.
579        assert_eq!(translate(b"ATGCA", 0, false), b"M");
580        assert_eq!(translate(b"AT", 0, false), b"");
581        assert_eq!(translate(b"ATG", 5, false), b"");
582    }
583
584    #[test]
585    fn kmers_and_canonical() {
586        assert_eq!(kmers(b"AAAA", 4).count(), 1);
587        assert_eq!(kmers(b"AAAA", 0).count(), 0);
588        assert_eq!(canonical_kmer(b"TTT"), b"AAA");
589        assert_eq!(canonical_kmer(b"AAA"), b"AAA");
590    }
591
592    #[test]
593    fn hamming() {
594        assert_eq!(hamming_distance(b"ACGT", b"acgt"), Some(0));
595        assert_eq!(hamming_distance(b"ACGT", b"ACGA"), Some(1));
596        assert_eq!(hamming_distance(b"ACGT", b"ACG"), None);
597    }
598
599    #[test]
600    fn n_statistics() {
601        // Total 100; sorted desc 50,30,15,5 -> 50 reaches 50.
602        assert_eq!(n50(&mut [5, 50, 30, 15]), Some(50));
603        assert_eq!(nx(&mut [5, 50, 30, 15], 0.9), Some(15));
604        assert_eq!(nx(&mut [0, 0], 0.5), Some(0));
605    }
606
607    #[test]
608    fn rna_dna_conversion() {
609        let mut s = b"ACGUacgu".to_vec();
610        rna_to_dna(&mut s);
611        assert_eq!(s, b"ACGTacgt");
612        make_uppercase(&mut s);
613        assert_eq!(s, b"ACGTACGT");
614    }
615}