Skip to main content

rsomics_fastq_downsample/
lib.rs

1use std::io::{BufWriter, Write};
2use std::path::Path;
3
4use rand::SeedableRng;
5use rand::distributions::{Distribution, Uniform};
6use rand::rngs::StdRng;
7use rsomics_common::{Result, RsomicsError};
8
9pub fn downsample(
10    input: &Path,
11    fraction: f64,
12    seed: u64,
13    output: &mut dyn Write,
14) -> Result<(u64, u64)> {
15    let mut reader = needletail::parse_fastx_file(input)
16        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", input.display())))?;
17
18    let mut out = BufWriter::with_capacity(256 * 1024, output);
19    let mut rng = StdRng::seed_from_u64(seed);
20    let dist = Uniform::new(0.0f64, 1.0);
21    let mut total: u64 = 0;
22    let mut kept: u64 = 0;
23
24    while let Some(result) = reader.next() {
25        let record =
26            result.map_err(|e| RsomicsError::InvalidInput(format!("reading record: {e}")))?;
27        total += 1;
28
29        if dist.sample(&mut rng) < fraction {
30            let name = record.id();
31            let seq = record.seq();
32            let qual = record.qual();
33
34            out.write_all(b"@").map_err(RsomicsError::Io)?;
35            out.write_all(name).map_err(RsomicsError::Io)?;
36            out.write_all(b"\n").map_err(RsomicsError::Io)?;
37            out.write_all(&seq).map_err(RsomicsError::Io)?;
38            out.write_all(b"\n+\n").map_err(RsomicsError::Io)?;
39            if let Some(q) = qual {
40                out.write_all(q.as_ref()).map_err(RsomicsError::Io)?;
41            }
42            out.write_all(b"\n").map_err(RsomicsError::Io)?;
43            kept += 1;
44        }
45    }
46
47    out.flush().map_err(RsomicsError::Io)?;
48    Ok((total, kept))
49}