Skip to main content

seqtk_rs/
seq.rs

1use crate::bed::{self, BedMap, BedPos};
2use crate::dna;
3use crate::io_utils::{FaReader, FqReader, FxWriter};
4use crate::record::RecordType;
5use crate::sub_cli::SeqArgs;
6
7struct FilterParas {
8    mini_seq_length: usize,
9    drop_ambigous_seq: bool,
10    output_odd_reads: bool,
11    output_even_reads: bool,
12}
13impl FilterParas {
14    fn from(seq: &SeqArgs) -> Self {
15        FilterParas {
16            mini_seq_length: seq.mini_seq_length.unwrap_or(0),
17            drop_ambigous_seq: seq.drop_ambigous_seq,
18            output_odd_reads: seq.output_even,
19            output_even_reads: seq.output_odd,
20        }
21    }
22}
23struct MaskParas {
24    mask_char: Option<char>,
25    uppercases: bool,
26    lowercases_to_char: bool,
27    q_low: u8,
28    q_high: u8,
29    mask_regions: Option<String>,
30    mask_complement_region: bool,
31}
32impl MaskParas {
33    fn from(seq: &SeqArgs) -> Self {
34        let ascii_bases = seq.ascii_bases.unwrap_or(33);
35        MaskParas {
36            mask_char: seq.mask_char,
37            uppercases: seq.uppercases,
38            lowercases_to_char: seq.lowercases_to_char,
39            q_low: (seq.q_low.unwrap_or(0) + ascii_bases),
40            q_high: (seq.q_high.unwrap_or(255 - ascii_bases) + ascii_bases),
41            mask_regions: seq.mask_regions.clone(),
42            mask_complement_region: seq.mask_complement_region,
43        }
44    }
45}
46struct OutArgs {
47    output_qual_shift: u8,
48    fake_fastq_quality: Option<char>,
49    output_fasta: bool,
50    reverse_complement: bool,
51    both_complement: bool,
52    trim_header: bool,
53    line_len: Option<usize>,
54}
55impl OutArgs {
56    fn from(seq: &SeqArgs) -> Self {
57        let ascii_bases = seq.ascii_bases.unwrap_or(33);
58        let out_qual_shift = if seq.output_qual_33 {
59            ascii_bases - 33
60        } else {
61            0
62        };
63        let output_fasta =
64            if !seq.output_fasta && seq.in_fa.is_some() && seq.fake_fastq_quality.is_none() {
65                true
66            } else {
67                seq.output_fasta
68            };
69        OutArgs {
70            output_qual_shift: out_qual_shift,
71            fake_fastq_quality: seq.fake_fastq_quality,
72            output_fasta,
73            reverse_complement: seq.reverse_complement,
74            both_complement: seq.both_complement,
75            trim_header: seq.trim_header,
76            line_len: seq.line_len,
77        }
78    }
79}
80
81/// Parses FASTA file and transforms the sequences according to the arguments.
82/// Outputs the results to [`std::io::stdout()`] in FASTA/Q format.
83///
84/// # Arguments
85///
86/// Check the arguments by `--help`
87///
88/// # Errors
89///
90/// Return an error if the operation cannot be completed.
91pub fn parse_fasta(path: &str, seq: &SeqArgs) -> Result<(), std::io::Error> {
92    let fparas = FilterParas::from(seq);
93    let mparas = MaskParas::from(seq);
94    let oparas = OutArgs::from(seq);
95    let bed_map = match &mparas.mask_regions {
96        Some(bed_path) => BedMap::from(bed_path)?,
97        None => BedMap::new(),
98    };
99    let fa_iter = FaReader::new(path)?;
100    let mut fx_writer = FxWriter::new(oparas.output_fasta);
101    for (i, record) in fa_iter.records().enumerate() {
102        match record {
103            Ok(read) => {
104                let is_pass = is_pass(i + 1, &read, &fparas);
105                if is_pass {
106                    modify_and_print_read(&mut fx_writer, &read, &mparas, &oparas, &bed_map, true)?;
107                }
108            }
109            Err(e) => eprintln!("Error read fASTA: {}", e),
110        }
111    }
112    Ok(())
113}
114/// Parses FASTQ file and transforms the sequences according to the arguments.
115/// Outputs the results to [`std::io::stdout()`] in FASTA/Q format.
116///
117/// # Arguments
118///
119/// Check the arguments by `--help`
120///
121/// # Errors
122///
123/// Return an error if the operation cannot be completed.
124pub fn parse_fastq(path: &str, seq: &SeqArgs) -> Result<(), std::io::Error> {
125    let fparas = FilterParas::from(seq);
126    let mparas = MaskParas::from(seq);
127    let oparas = OutArgs::from(seq);
128    let bed_map = match &mparas.mask_regions {
129        Some(bed_path) => BedMap::from(bed_path)?,
130        None => BedMap::new(),
131    };
132
133    let fq_iter = FqReader::new(path)?;
134    let mut fx_writer = FxWriter::new(oparas.output_fasta);
135    for (i, record) in fq_iter.records().enumerate() {
136        match record {
137            Ok(read) => {
138                let is_pass = is_pass(i + 1, &read, &fparas);
139                if is_pass {
140                    modify_and_print_read(
141                        &mut fx_writer,
142                        &read,
143                        &mparas,
144                        &oparas,
145                        &bed_map,
146                        false,
147                    )?;
148                }
149            }
150            Err(e) => eprintln!("Error read fASTQ: {}", e),
151        }
152    }
153
154    Ok(())
155}
156fn modify_and_print_read(
157    fx_writer: &mut FxWriter,
158    read: &dyn RecordType,
159    mask_paras: &MaskParas,
160    out_paras: &OutArgs,
161    bed_map: &BedMap,
162    is_fasta: bool,
163) -> Result<(), std::io::Error> {
164    let mut seq = modify_seq(read, mask_paras, bed_map, is_fasta);
165    let desc = if out_paras.trim_header {
166        None
167    } else {
168        read.desc()
169    };
170    let mut qual = modify_qual(read, out_paras, is_fasta);
171
172    if let Some(line_len) = out_paras.line_len {
173        add_newlines(&mut seq, line_len);
174        add_newlines(&mut qual, line_len);
175    }
176
177    if out_paras.reverse_complement {
178        revcomp(&mut seq, &mut qual);
179        fx_writer.write(read.id(), &seq, desc, &qual)?;
180    } else if out_paras.both_complement {
181        fx_writer.write(read.id(), &seq, desc, &qual)?;
182        revcomp(&mut seq, &mut qual);
183        fx_writer.write(read.id(), &seq, desc, &qual)?;
184    } else {
185        fx_writer.write(read.id(), &seq, desc, &qual)?;
186    }
187    Ok(())
188}
189fn add_newlines(data: &mut Vec<u8>, line_len: usize) {
190    let mut i = line_len;
191    while i < data.len() {
192        data.insert(i, b'\n'); // Insert '\n' after i
193        i += line_len + 1;
194    }
195}
196fn revcomp(seq: &mut [u8], qual: &mut [u8]) {
197    dna::revcomp(seq);
198    qual.reverse();
199}
200fn modify_qual(read: &dyn RecordType, oparas: &OutArgs, is_fasta: bool) -> Vec<u8> {
201    match oparas.fake_fastq_quality {
202        Some(fake_qual) => {
203            vec![fake_qual as u8; read.seq().len()]
204        }
205        None => {
206            if is_fasta {
207                Vec::new()
208            } else if oparas.output_qual_shift == 0 {
209                read.qual().to_vec()
210            } else {
211                read.qual()
212                    .iter()
213                    .map(|q| q - oparas.output_qual_shift)
214                    .collect()
215            }
216        }
217    }
218}
219fn modify_seq(
220    read: &dyn RecordType,
221    mask_paras: &MaskParas,
222    bed_map: &BedMap,
223    is_fasta: bool,
224) -> Vec<u8> {
225    let default_bed_pos = vec![BedPos(usize::MAX, 0)];
226    let bed_pos = bed_map.get(read.id()).unwrap_or(&default_bed_pos);
227    let mut seq = if mask_paras.uppercases {
228        read.seq().to_ascii_uppercase() // convert all to uppercases
229    } else {
230        read.seq().to_vec()
231    };
232
233    match mask_paras.mask_char {
234        Some(c) => {
235            let c_u8 = c as u8;
236            let mut cur_idx: usize = 0;
237            seq.iter_mut().enumerate().for_each(|(i, ch)| {
238                if mask_paras.lowercases_to_char && ch.is_ascii_lowercase() {
239                    *ch = c_u8;
240                } else {
241                    let (is_overlap, c_idx) = bed::is_overlapping(i, &bed_pos[cur_idx..]);
242                    cur_idx += c_idx;
243                    if (is_overlap && !mask_paras.mask_complement_region)
244                        || (!is_overlap && mask_paras.mask_complement_region)
245                    {
246                        *ch = c_u8;
247                    }
248                }
249            })
250        }
251        None => {
252            let mut cur_idx: usize = 0;
253            seq.iter_mut().enumerate().for_each(|(i, ch)| {
254                let (is_overlap, c_idx) = bed::is_overlapping(i, &bed_pos[cur_idx..]);
255                cur_idx += c_idx;
256                if (is_overlap && !mask_paras.mask_complement_region)
257                    || (!is_overlap && mask_paras.mask_complement_region)
258                {
259                    *ch = ch.to_ascii_lowercase();
260                }
261            })
262        }
263    }
264    if !is_fasta {
265        match mask_paras.mask_char {
266            Some(c) => {
267                let c_u8 = c as u8;
268                read.qual()
269                    .iter()
270                    .zip(seq.iter_mut())
271                    .for_each(|(&qual, ch)| {
272                        if qual < mask_paras.q_low || mask_paras.q_high < qual {
273                            *ch = c_u8; // mask bases by q_low and q_high
274                        }
275                    })
276            }
277            None => {
278                read.qual()
279                    .iter()
280                    .zip(seq.iter_mut())
281                    .for_each(|(&qual, ch)| {
282                        if qual < mask_paras.q_low || mask_paras.q_high < qual {
283                            *ch = ch.to_ascii_lowercase(); // mask bases by q_low and q_high
284                        }
285                    })
286            }
287        }
288    }
289    seq
290}
291fn is_pass(i: usize, read: &dyn RecordType, fparas: &FilterParas) -> bool {
292    if fparas.output_even_reads && i % 2 == 0 {
293        return false;
294    }
295    if fparas.output_odd_reads && i % 2 == 1 {
296        return false;
297    }
298    if fparas.drop_ambigous_seq && read.seq().iter().any(|&c| dna::get_dna_idx_from_u8(c) > 3) {
299        return false;
300    }
301    if read.seq().len().le(&fparas.mini_seq_length) {
302        return false;
303    }
304    true
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use bio::io::fastq::Record;
311
312    #[test]
313    fn test_modify_qual() {
314        fn init_oparas() -> OutArgs {
315            OutArgs {
316                output_qual_shift: 0,
317                fake_fastq_quality: None,
318                output_fasta: false,
319                reverse_complement: false,
320                both_complement: false,
321                trim_header: false,
322                line_len: None,
323            }
324        }
325        let record = Record::with_attrs("SEQ_ID_1", None, b"ATCGATcgACTTG", b"gfryremb[trdg");
326        let mut oparas = init_oparas();
327
328        // [01] without modify
329        let out_qual = modify_qual(&record, &oparas, false);
330        assert_eq!(&out_qual, b"gfryremb[trdg");
331
332        // [02] check output_qual_shift
333        oparas.output_qual_shift = 10;
334        let out_qual = modify_qual(&record, &oparas, false);
335        assert_eq!(&out_qual, b"]\\hoh[cXQjhZ]");
336
337        // [02] check output_qual_shift
338        oparas = init_oparas();
339        oparas.fake_fastq_quality = Some('T');
340        let out_qual = modify_qual(&record, &oparas, false);
341        assert_eq!(&out_qual, b"TTTTTTTTTTTTT");
342    }
343
344    #[test]
345    fn test_modify_seq() {
346        fn init_mparas() -> MaskParas {
347            MaskParas {
348                mask_char: None,
349                uppercases: false,
350                lowercases_to_char: false,
351                q_low: 0,
352                q_high: 255,
353                mask_regions: None,
354                mask_complement_region: false,
355            }
356        }
357        let record = Record::with_attrs("SEQ_ID_1", None, b"ATCGATcgACTTG", b"!(*AAAABbbaaz");
358        let mut bed_map: BedMap = BedMap::new();
359        let mut mparas = init_mparas();
360
361        // [01] without modify
362        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
363        assert_eq!(&out_seq, b"ATCGATcgACTTG", "[err01]");
364
365        // [02] check mask_char + lowrcases_to_char
366        mparas.mask_char = Some('N');
367        mparas.lowercases_to_char = true;
368        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
369        assert_eq!(&out_seq, b"ATCGATNNACTTG", "[err02]");
370
371        // [03] check uppercases
372        mparas = init_mparas();
373        mparas.uppercases = true;
374        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
375        assert_eq!(&out_seq, b"ATCGATCGACTTG", "[err03]");
376
377        // [04] check q_low and q_high
378        mparas = init_mparas();
379        mparas.q_low = 20 + 33;
380        mparas.q_high = 255;
381        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
382        assert_eq!(&out_seq, b"atcGATcgACTTG", "[err04]");
383        mparas.q_low = 0;
384        mparas.q_high = 85 + 33;
385        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
386        assert_eq!(&out_seq, b"ATCGATcgACTTg", "[err05-1]");
387
388        // [05] check q_low and q_high + mask_char and lowercases_to_char
389        mparas = init_mparas();
390        mparas.mask_char = Some('N');
391        mparas.q_low = 20 + 33;
392        mparas.q_high = 85 + 33;
393        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
394        assert_eq!(&out_seq, b"NNNGATcgACTTN", "[err05-2]");
395        mparas.lowercases_to_char = true;
396        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
397        assert_eq!(&out_seq, b"NNNGATNNACTTN", "[err05-3]");
398
399        // [06] check q_low and q_high + uppercases
400        mparas = init_mparas();
401        mparas.uppercases = true;
402        mparas.q_low = 20 + 33;
403        mparas.q_high = 85 + 33;
404        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
405        assert_eq!(&out_seq, b"atcGATCGACTTg", "[err06]");
406
407        // [07] check mask by bed
408        let bed_path = Some("test.bed".to_string());
409        bed_map.add("SEQ_ID_1".to_string(), BedPos(5, 10));
410        mparas = init_mparas();
411        mparas.mask_regions = bed_path.clone();
412        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
413        assert_eq!(&out_seq, b"ATCGAtcgacTTG", "[err07]");
414
415        // [08] checl mask by complement region
416        mparas.mask_complement_region = true;
417        let out_seq = modify_seq(&record, &mparas, &bed_map, false);
418        assert_eq!(&out_seq, b"atcgaTcgACttg", "[err08]");
419    }
420
421    #[test]
422    fn test_is_filtered() {
423        fn init_fparas() -> FilterParas {
424            FilterParas {
425                mini_seq_length: 0,
426                drop_ambigous_seq: false,
427                output_odd_reads: false,
428                output_even_reads: false,
429            }
430        }
431        let mut fparas = init_fparas();
432        let record = Record::with_attrs("@SEQ_ID_1", None, b"ATCGATCGACTTG", b"!!<AAAABbbaab");
433
434        // [01] check mini_seq_length
435        fparas.mini_seq_length = 10;
436        assert_eq!(is_pass(0, &record, &fparas), true);
437        fparas.mini_seq_length = 50;
438        assert_eq!(is_pass(0, &record, &fparas), false);
439
440        // [02] check drop ambiguous seq
441        fparas = init_fparas();
442        fparas.drop_ambigous_seq = true;
443        assert_eq!(is_pass(0, &record, &fparas), true);
444        let record2 = Record::with_attrs("@SEQ_ID_2", None, b"ANCGATCGACTTG", b"!!<AAAABbbaab");
445        assert_eq!(is_pass(0, &record2, &fparas), false);
446
447        // [03] output odd read
448        fparas.output_odd_reads = true;
449        assert_eq!(is_pass(0, &record, &fparas), true);
450        assert_eq!(is_pass(1, &record, &fparas), false);
451
452        // [04] output even read
453        fparas = init_fparas();
454        fparas.output_even_reads = true;
455        assert_eq!(is_pass(3, &record, &fparas), true);
456        assert_eq!(is_pass(4, &record, &fparas), false);
457    }
458    #[test]
459    fn test_add_newlines() {
460        let mut seq = b"aaaaabbbbbcccccdddddeeeeefffff".to_vec();
461        add_newlines(&mut seq, 5);
462        assert_eq!(seq, b"aaaaa\nbbbbb\nccccc\nddddd\neeeee\nfffff");
463    }
464    #[test]
465    fn test_revcomp() {
466        let mut seq = b"AATTCCGG".to_vec();
467        let mut qual = b"<<((vv++".to_vec();
468
469        revcomp(&mut seq, &mut qual);
470
471        assert_eq!(seq, b"CCGGAATT");
472        assert_eq!(qual, b"++vv((<<");
473    }
474}