rsomics_fastq_downsample/
lib.rs1use 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}