diff --git a/docs-site/src/gen-reads/fastq-shuffling.md b/docs-site/src/gen-reads/fastq-shuffling.md index 8ce2bad4..1ed4a0d3 100644 --- a/docs-site/src/gen-reads/fastq-shuffling.md +++ b/docs-site/src/gen-reads/fastq-shuffling.md @@ -1,29 +1,48 @@ # FASTQ Shuffling -`eidolon` writes reads in genome order: contig by contig, by position within each contig. A -real sequencer's FASTQ is in effectively random order. Variant calling does not depend on this, -since the BAM is coordinate-sorted before calling. Tools that read the FASTQ in batches, -aligners in particular, do see the difference: a batch can be filled entirely with reads from -one region, such as a long satellite array. Measured on a 21x simulation of T2T-CHM13 (#870), -bwa-mem2 peaked at 166 GB on the reads as written and at 90 GB on the same reads shuffled, -with the same alignments either way. Shuffle before aligning a large genome. -#870 tracks shuffling in `eidolon` itself. - -With `seqkit`, shuffle each mate separately **with the same seed**. seqkit has no paired mode, -but its permutation depends only on the seed and the number of records, and R1 and R2 have the -same count. So the two files come out in the same order and mates stay in register: +`eidolon` writes its FASTQs in **random order** by default (`shuffle_fastq: true`), as a +sequencer does. Mates stay in register: line *n* of R1 and line *n* of R2 are the two reads of +one pair. -```bash -# single-ended -seqkit shuffle -s 11 sample.fastq.gz -o sample_shuffled.fastq.gz +Reads are generated chunk by chunk along the genome, so without shuffling the FASTQ comes out in +genome order: contig by contig, by position within each contig. Variant calling does not depend +on this, because the BAM is coordinate-sorted before calling. Tools that read the FASTQ in +batches do see the difference, aligners in particular: a batch can be filled entirely with reads +from one region, such as a long satellite array. Measured on a 21x simulation of T2T-CHM13 +(#870), bwa-mem2 peaked at 166 GB on reads in genome order and at 90 GB on the same reads +shuffled, with the same alignments either way. + +```yaml +shuffle_fastq: true # default; false writes reads in genome order +``` + +**What it does and does not change.** +- Only the order of FASTQ records changes. The reads themselves, the truth VCF and the golden BAM + are identical with the shuffle on or off. The shuffle draws from its own branch of the random + number generator, so it does not disturb the streams that generate reads. +- Output is deterministic: the same seed gives byte-identical FASTQs at any `num_threads`. +- Turning the shuffle off restores genome order exactly as earlier versions wrote it. + +**Cost.** The shuffle is a two-pass external sort. Records are scattered at random into bucket +files in the run's temp directory, then each bucket is shuffled in memory and written out. Buckets +are shuffled in parallel only as far as they fit in about 4 GB, so memory stays bounded at any +genome size. Temp disk stays close to one copy of the output, because each input is deleted as it +is read. -# paired-ended: same -s on both mates +Measured on T2T-CHM13 chr21 at 20x (about 3M pairs, 8 threads), against `shuffle_fastq: false`: +- wall time 112 s against 99 s; +- peak memory 2.3 GB against 0.5 GB; +- **gzipped FASTQs about 60% larger.** Reads in genome order overlap their neighbors, which gzip + exploits; shuffled reads, like a real sequencer's, compress less well. + +**Shuffling output from older versions, or from `shuffle_fastq: false`.** With `seqkit`, shuffle +each mate separately **with the same seed**. seqkit has no paired mode, but its permutation +depends only on the seed and the number of records, and R1 and R2 have the same count: + +```bash seqkit shuffle -s 11 sample_r1.fastq.gz -o sample_r1_shuffled.fastq.gz seqkit shuffle -s 11 sample_r2.fastq.gz -o sample_r2_shuffled.fastq.gz ``` -**Memory:** `seqkit shuffle` reads every record into memory, so expect memory on the order of -the uncompressed FASTQ, per mate. Its low-memory `-2/--two-pass` mode supports FASTA only. -Do not pass both mates to one `seqkit shuffle` call: they would be read as one input and shuffled -together into a single output. - -(`seqkit` is an open-source FASTQ toolkit: https://github.com/shenwei356/seqkit.) +`seqkit shuffle` reads every record into memory; its low-memory `-2/--two-pass` mode supports +FASTA only. Do not pass both mates to one call: they would be shuffled together into a single +output. (`seqkit`: https://github.com/shenwei356/seqkit.) diff --git a/docs-site/src/gen-reads/parallel-processing.md b/docs-site/src/gen-reads/parallel-processing.md index 4d8c9106..23eb73fb 100644 --- a/docs-site/src/gen-reads/parallel-processing.md +++ b/docs-site/src/gen-reads/parallel-processing.md @@ -1,7 +1,7 @@ # Parallel Processing `eidolon gen-reads` processes contigs in parallel by default using rayon's work-stealing thread pool. Each contig is an independent unit of work — variant generation, fragment sampling, and FASTQ/BAM writing all happen concurrently across contigs — so references with many contigs scale well across cores. -Output is **byte-identical regardless of `num_threads`** (the same seed always produces the same reads in the same order), so you can change the thread count freely without affecting results. +Output is **byte-identical regardless of `num_threads`** (the same seed always produces the same reads in the same order, shuffled or not; see [FASTQ Shuffling](fastq-shuffling.md)), so you can change the thread count freely without affecting results. A note on scaling: read generation is largely **memory-bandwidth bound**, so on references dominated by one or a few large chromosomes the wall-time gain from extra cores is limited — the bottleneck is memory throughput, not CPU. An experimental sub-contig **chunk size** knob (below) can split a single chromosome across cores, but it is **disabled by default** because in our benchmarks it did not improve wall time and occasionally regressed it. diff --git a/eidolon-core/src/file_tools.rs b/eidolon-core/src/file_tools.rs index c9c76b9f..f30a9f3f 100644 --- a/eidolon-core/src/file_tools.rs +++ b/eidolon-core/src/file_tools.rs @@ -6,6 +6,7 @@ pub mod bam_writer; pub mod bed_reader; pub mod block_gz; pub mod fasta_stream; +pub mod fastq_shuffle; pub mod fastq_tools; pub mod file_io; pub mod folder_tools; diff --git a/eidolon-core/src/file_tools/fastq_shuffle.rs b/eidolon-core/src/file_tools/fastq_shuffle.rs new file mode 100644 index 00000000..709471f7 --- /dev/null +++ b/eidolon-core/src/file_tools/fastq_shuffle.rs @@ -0,0 +1,674 @@ +//! Combine per-chunk temp FASTQs into the final output in random order (#870). +//! +//! **Why.** gen-reads generates reads chunk by chunk, so concatenating the temp files gives a +//! FASTQ in genome order, which no sequencer produces. Aligners read FASTQs in batches, and a +//! batch filled with one region's reads behaves differently from a mixed one: bwa-mem2 needed +//! 166 GB on a 21x T2T-CHM13 simulation as written and 90 GB on the same reads shuffled (#870). +//! The original in-memory shuffle was removed because it held every read at once. +//! +//! **How: a two-pass external shuffle with bounded memory.** +//! 1. Scatter. Read the temp files in order, and send each record (each pair, when paired) to one +//! of `K` bucket files chosen at random. Both mates of a pair go to the same bucket, in the +//! same position. Each input file is deleted once read, so temp disk stays near one copy. +//! 2. Shuffle. Load one bucket at a time, permute it in memory, and write it as gzip; then +//! concatenate the buckets (gzip members concatenate into one valid stream). +//! +//! Bucket assignment and then a uniform permutation within each bucket give a uniform +//! permutation of the whole file. `K` is chosen so a bucket fits the memory budget. +//! +//! **Determinism.** The scatter draws from one RNG, sequentially. Each bucket's permutation draws +//! from its own child RNG. Threads only change which bucket is processed when, so output is +//! byte-identical across thread counts. + +use std::fs::{self, File}; +use std::io::{BufRead, BufReader, BufWriter, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use flate2::read::MultiGzDecoder; +use libdeflater::{CompressionLvl, Compressor}; + +use crate::file_tools::block_gz::BlockGzWriter; +use crate::file_tools::fastq_tools::FastqToolsError; +use crate::rng::NeatRng; + +/// Uncompressed bytes per bucket, both mates together. Pass 2 holds this much per thread, +/// roughly doubled by per-record allocation. +const BUCKET_TARGET_BYTES: u64 = 256 * 1024 * 1024; +/// Upper bound on buckets. The scatter holds one buffer per bucket per mate, so this bounds that +/// memory (512 * 2 * `BUCKET_BUFFER_BYTES` = 256 MiB). Past it, buckets grow instead. +const MAX_BUCKETS: usize = 512; +/// Scatter buffer per bucket per mate. A full buffer is compressed as one gzip member and +/// appended to its bucket file, which is then closed again, so the scatter holds at most one +/// bucket file open at any bucket count. One writer per bucket per mate would need +/// 2 * MAX_BUCKETS descriptors and failed under `ulimit -n` 256 or 512. +const BUCKET_BUFFER_BYTES: usize = 256 * 1024; +/// Gzipped FASTQ expands about 3.5-4.5x; overestimating only adds buckets. +const EXPANSION_ESTIMATE: u64 = 5; +/// Pass-2 threads, capped because each holds a whole bucket in memory. +const MAX_SHUFFLE_THREADS: usize = 8; +/// Pass 2 runs at most as many buckets at once as fit in this much memory, judged by the +/// largest bucket. Small runs get full parallelism; at genome scale, where `MAX_BUCKETS` makes +/// buckets larger, concurrency drops so memory stays near this figure. +const SHUFFLE_MEMORY_BUDGET: u64 = 4 * 1024 * 1024 * 1024; + +fn read_err(path: &Path, e: impl std::fmt::Display) -> FastqToolsError { + FastqToolsError::FastqReadError(format!("{}: {e}", path.display())) +} + +fn write_err(path: &Path, e: impl std::fmt::Display) -> FastqToolsError { + FastqToolsError::FastqWriteError(format!("{}: {e}", path.display())) +} + +/// Reads one 4-line FASTQ record into `buf` (appending). `Ok(false)` at a clean end of input; +/// an error if the input ends partway through a record. +fn read_record( + reader: &mut impl BufRead, + buf: &mut Vec, + path: &Path, +) -> Result { + let start = buf.len(); + for line in 0..4 { + let n = reader + .read_until(b'\n', buf) + .map_err(|e| read_err(path, e))?; + if n == 0 { + if line == 0 { + return Ok(false); + } + return Err(read_err( + path, + format!("input ends {line} lines into a FASTQ record"), + )); + } + } + if buf.last() != Some(&b'\n') { + buf.push(b'\n'); + } + if buf[start] != b'@' { + return Err(read_err(path, "a FASTQ record does not start with '@'")); + } + Ok(true) +} + +fn open_gz(path: &Path) -> Result>, FastqToolsError> { + let file = File::open(path).map_err(|e| read_err(path, e))?; + Ok(BufReader::with_capacity(1 << 20, MultiGzDecoder::new(file))) +} + +/// One bucket file being scattered into: records accumulate in `buf` and leave as a gzip member +/// appended to `path`, so the file is open only while a member is written. +struct BucketSink { + path: PathBuf, + buf: Vec, + wrote: bool, +} + +impl BucketSink { + fn new(path: PathBuf) -> Self { + Self { + path, + buf: Vec::new(), + wrote: false, + } + } + + fn push( + &mut self, + rec: &[u8], + c: &mut Compressor, + scratch: &mut Vec, + ) -> Result<(), FastqToolsError> { + self.buf.extend_from_slice(rec); + if self.buf.len() >= BUCKET_BUFFER_BYTES { + self.flush(c, scratch)?; + } + Ok(()) + } + + /// Appends the buffer as one gzip member. An empty member is written only for a bucket that + /// received nothing, so every bucket file is valid gzip. + fn flush(&mut self, c: &mut Compressor, scratch: &mut Vec) -> Result<(), FastqToolsError> { + if self.buf.is_empty() && self.wrote { + return Ok(()); + } + let bound = c.gzip_compress_bound(self.buf.len()); + if scratch.len() < bound { + scratch.resize(bound, 0); + } + let n = c + .gzip_compress(&self.buf, scratch) + .map_err(|e| write_err(&self.path, format!("gzip compression failed: {e:?}")))?; + let mut f = fs::OpenOptions::new() + .append(true) + .create(true) + .open(&self.path) + .map_err(|e| write_err(&self.path, e))?; + f.write_all(&scratch[..n]) + .map_err(|e| write_err(&self.path, e))?; + self.buf.clear(); + self.wrote = true; + Ok(()) + } + + fn finish(&mut self, c: &mut Compressor, scratch: &mut Vec) -> Result<(), FastqToolsError> { + self.flush(c, scratch)?; + self.buf = Vec::new(); + Ok(()) + } +} + +/// Pass-2 concurrency: the requested threads, capped at `MAX_SHUFFLE_THREADS` and at however +/// many of the largest bucket fit in `SHUFFLE_MEMORY_BUDGET`. Always at least one. +fn shuffle_threads(requested: usize, largest_bucket_bytes: u64) -> usize { + let fit = (SHUFFLE_MEMORY_BUDGET / largest_bucket_bytes.max(1)) as usize; + requested.min(MAX_SHUFFLE_THREADS).min(fit).max(1) +} + +/// Number of buckets for inputs totalling `compressed_bytes` on disk. +fn bucket_count(compressed_bytes: u64) -> usize { + let estimate = compressed_bytes.saturating_mul(EXPANSION_ESTIMATE); + (estimate.div_ceil(BUCKET_TARGET_BYTES) as usize).clamp(1, MAX_BUCKETS) +} + +/// Writes `files_r1` (and `files_r2`, record for record, when paired) to `out_r1` (and +/// `out_r2`) in a uniformly random order, keeping mates in register. `rng` is consumed for the +/// scatter; each bucket's permutation uses `rng.derive_child(bucket)`. Temp files go in +/// `work_dir`, and every input file is deleted once read. +pub fn combine_temp_fastqs_shuffled( + files_r1: Vec, + files_r2: Vec, + out_r1: &Path, + out_r2: Option<&Path>, + rng: &mut NeatRng, + work_dir: &Path, + threads: usize, +) -> Result<(), FastqToolsError> { + let compressed: u64 = files_r1 + .iter() + .chain(files_r2.iter()) + .map(|p| fs::metadata(p).map(|m| m.len()).unwrap_or(0)) + .sum(); + shuffle_into_buckets( + files_r1, + files_r2, + out_r1, + out_r2, + rng, + work_dir, + threads, + bucket_count(compressed), + ) +} + +/// The shuffle with an explicit bucket count, so tests can exercise the multi-bucket path on +/// inputs small enough that `bucket_count` would choose one. +#[allow(clippy::too_many_arguments)] +fn shuffle_into_buckets( + files_r1: Vec, + files_r2: Vec, + out_r1: &Path, + out_r2: Option<&Path>, + rng: &mut NeatRng, + work_dir: &Path, + threads: usize, + k: usize, +) -> Result<(), FastqToolsError> { + let paired = out_r2.is_some(); + if paired && files_r1.len() != files_r2.len() { + return Err(FastqToolsError::FastqWriteError(format!( + "shuffle: {} read-1 temp files but {} read-2", + files_r1.len(), + files_r2.len() + ))); + } + let bucket_dir = work_dir.join("shuffle_buckets"); + fs::create_dir_all(&bucket_dir).map_err(|e| write_err(&bucket_dir, e))?; + let bucket_path = |b: usize, mate: u8, stage: &str| { + bucket_dir.join(format!("{stage}_{b:04}_r{mate}.fastq.gz")) + }; + + // ── Pass 1: scatter ─────────────────────────────────────────────────────────────────────── + let mates: Vec = if paired { vec![1, 2] } else { vec![1] }; + let mut sinks: Vec> = (0..k) + .map(|b| { + mates + .iter() + .map(|&m| BucketSink::new(bucket_path(b, m, "bucket"))) + .collect() + }) + .collect(); + let mut compressor = Compressor::new(CompressionLvl::fastest()); + let mut scratch = Vec::new(); + let mut counts = vec![0u64; k]; + let mut bytes = vec![0u64; k]; + let (mut rec1, mut rec2) = (Vec::with_capacity(1024), Vec::with_capacity(1024)); + for (i, p1) in files_r1.iter().enumerate() { + let mut in1 = open_gz(p1)?; + let mut in2 = if paired { + Some((open_gz(&files_r2[i])?, &files_r2[i])) + } else { + None + }; + loop { + rec1.clear(); + let more1 = read_record(&mut in1, &mut rec1, p1)?; + let more2 = match in2.as_mut() { + Some((r, p)) => { + rec2.clear(); + read_record(r, &mut rec2, p)? + } + None => more1, + }; + if more1 != more2 { + return Err(read_err( + p1, + "read-1 and read-2 temp files hold different numbers of records", + )); + } + if !more1 { + break; + } + let b = ((rng.random()? * k as f64) as usize).min(k - 1); + sinks[b][0].push(&rec1, &mut compressor, &mut scratch)?; + if paired { + sinks[b][1].push(&rec2, &mut compressor, &mut scratch)?; + } + counts[b] += 1; + bytes[b] += (rec1.len() + if paired { rec2.len() } else { 0 }) as u64; + } + drop(in1); + drop(in2); + let _ = fs::remove_file(p1); + if paired { + let _ = fs::remove_file(&files_r2[i]); + } + } + for sink in sinks.iter_mut().flatten() { + sink.finish(&mut compressor, &mut scratch)?; + } + drop(sinks); + + // ── Pass 2: shuffle each bucket ─────────────────────────────────────────────────────────── + let next = AtomicUsize::new(0); + let first_error: Mutex> = Mutex::new(None); + let base = *rng; + let shuffle_bucket = |b: usize| -> Result<(), FastqToolsError> { + // One contiguous buffer per mate, with record start offsets: a Vec per record would + // cost an allocation and its slack for every read. + let mut records: Vec<(Vec, Vec)> = Vec::with_capacity(mates.len()); + for &m in &mates { + let p = bucket_path(b, m, "bucket"); + let mut r = open_gz(&p)?; + let mut data = Vec::new(); + let mut starts = Vec::with_capacity(counts[b] as usize + 1); + loop { + starts.push(data.len()); + if !read_record(&mut r, &mut data, &p)? { + break; + } + } + if starts.len() as u64 - 1 != counts[b] { + return Err(read_err( + &p, + format!( + "bucket holds {} records, {} were written", + starts.len() - 1, + counts[b] + ), + )); + } + drop(r); + let _ = fs::remove_file(&p); + records.push((data, starts)); + } + let mut order: Vec = (0..counts[b] as usize).collect(); + base.derive_child(b as u64).shuffle_in_place(&mut order)?; + for (j, &m) in mates.iter().enumerate() { + let p = bucket_path(b, m, "shuffled"); + let f = File::create(&p).map_err(|e| write_err(&p, e))?; + let mut w = BlockGzWriter::new(BufWriter::new(f)); + let (data, starts) = &records[j]; + for &i in &order { + w.write_all(&data[starts[i]..starts[i + 1]]) + .map_err(|e| write_err(&p, e))?; + } + w.finish().map_err(|e| write_err(&p, e))?; + } + Ok(()) + }; + std::thread::scope(|s| { + for _ in 0..shuffle_threads(threads, bytes.iter().copied().max().unwrap_or(0)) { + s.spawn(|| { + loop { + let b = next.fetch_add(1, Ordering::SeqCst); + if b >= k || first_error.lock().unwrap().is_some() { + break; + } + if let Err(e) = shuffle_bucket(b) { + first_error.lock().unwrap().get_or_insert(e); + break; + } + } + }); + } + }); + if let Some(e) = first_error.into_inner().unwrap() { + return Err(e); + } + + // ── Concatenate, in bucket order ────────────────────────────────────────────────────────── + for (j, &m) in mates.iter().enumerate() { + let out = if j == 0 { + out_r1 + } else { + out_r2.expect("paired") + }; + let mut w = BufWriter::new( + fs::OpenOptions::new() + .append(true) + .create(true) + .open(out) + .map_err(|e| write_err(out, e))?, + ); + for b in 0..k { + let p = bucket_path(b, m, "shuffled"); + let mut f = BufReader::new(File::open(&p).map_err(|e| read_err(&p, e))?); + std::io::copy(&mut f, &mut w).map_err(|e| write_err(out, e))?; + drop(f); + let _ = fs::remove_file(&p); + } + w.flush().map_err(|e| write_err(out, e))?; + } + let _ = fs::remove_dir(&bucket_dir); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Read; + + /// Writes `n` records named by index (and their mates) as `chunks` temp-file pairs, in order. + fn temp_inputs( + dir: &Path, + n: usize, + chunks: usize, + paired: bool, + ) -> (Vec, Vec) { + let (mut f1, mut f2) = (Vec::new(), Vec::new()); + let per = n.div_ceil(chunks); + for c in 0..chunks { + let p1 = dir.join(format!("t{c}_r1.fastq.gz")); + let p2 = dir.join(format!("t{c}_r2.fastq.gz")); + let mut w1 = BlockGzWriter::new(File::create(&p1).unwrap()); + let mut w2 = BlockGzWriter::new(File::create(&p2).unwrap()); + for i in c * per..((c + 1) * per).min(n) { + write!(w1, "@read_{i:08}/1\nACGT{i}\n+\nIIII{i}\n").unwrap(); + write!(w2, "@read_{i:08}/2\nTGCA{i}\n+\nJJJJ{i}\n").unwrap(); + } + w1.finish().unwrap(); + w2.finish().unwrap(); + f1.push(p1); + if paired { + f2.push(p2); + } + } + (f1, f2) + } + + fn names(path: &Path) -> Vec { + let mut s = String::new(); + MultiGzDecoder::new(File::open(path).unwrap()) + .read_to_string(&mut s) + .unwrap(); + s.lines().step_by(4).map(str::to_string).collect() + } + + fn index_of(name: &str) -> usize { + name[6..14].parse().unwrap() + } + + /// Shuffles `n` records from `chunks` temp-file pairs through `k` buckets. Tests pass k > 1 + /// explicitly: at these sizes `bucket_count` would always choose 1, leaving the scatter untested. + fn run_k( + n: usize, + chunks: usize, + seed: &str, + threads: usize, + k: usize, + ) -> (Vec, Vec) { + let dir = tempfile::tempdir().unwrap(); + let (f1, f2) = temp_inputs(dir.path(), n, chunks, true); + let (o1, o2) = ( + dir.path().join("o1.fastq.gz"), + dir.path().join("o2.fastq.gz"), + ); + let mut rng = NeatRng::new_from_seed(&vec![seed.to_string()]).unwrap(); + shuffle_into_buckets(f1, f2, &o1, Some(&o2), &mut rng, dir.path(), threads, k).unwrap(); + (names(&o1), names(&o2)) + } + + /// The output files' raw bytes, for comparing gzip streams rather than their contents. + fn raw_k(n: usize, chunks: usize, seed: &str, threads: usize, k: usize) -> (Vec, Vec) { + let dir = tempfile::tempdir().unwrap(); + let (f1, f2) = temp_inputs(dir.path(), n, chunks, true); + let (o1, o2) = ( + dir.path().join("o1.fastq.gz"), + dir.path().join("o2.fastq.gz"), + ); + let mut rng = NeatRng::new_from_seed(&vec![seed.to_string()]).unwrap(); + shuffle_into_buckets(f1, f2, &o1, Some(&o2), &mut rng, dir.path(), threads, k).unwrap(); + (fs::read(&o1).unwrap(), fs::read(&o2).unwrap()) + } + + fn run(n: usize, chunks: usize, seed: &str, threads: usize) -> (Vec, Vec) { + run_k(n, chunks, seed, threads, 16) + } + + #[test] + fn output_is_a_permutation_with_mates_in_register() { + let n = 20_000; + let (r1, r2) = run(n, 7, "perm", 4); + assert_eq!(r1.len(), n); + let mut idx: Vec = r1.iter().map(|s| index_of(s)).collect(); + for (a, b) in r1.iter().zip(&r2) { + assert_eq!( + a.trim_end_matches("/1"), + b.trim_end_matches("/2"), + "mates out of register" + ); + } + idx.sort_unstable(); + assert_eq!(idx, (0..n).collect::>(), "every record exactly once"); + } + + #[test] + fn output_order_is_random_not_the_input_order() { + // Adjacent records: in input order every neighbor is the next index. Shuffled, a neighbor + // is within 100 indexes with probability ~ 2 * 100 / n = 0.01. + let n = 20_000; + let (r1, _) = run(n, 7, "order", 4); + let idx: Vec = r1.iter().map(|s| index_of(s) as i64).collect(); + let near = idx + .windows(2) + .filter(|w| (w[1] - w[0]).abs() <= 100) + .count() as f64 + / (n - 1) as f64; + assert!( + near < 0.03, + "adjacent records within 100 indexes: {near:.4}; the output is not shuffled" + ); + assert_ne!( + idx[..100], + (0..100).collect::>()[..], + "starts in input order" + ); + } + + #[test] + fn same_seed_same_output_across_thread_counts_and_different_seed_differs() { + let n = 5_000; + let a = run(n, 3, "det", 1); + assert_eq!(a, run(n, 3, "det", 8), "thread count changed the output"); + // The gzip bytes too: member boundaries or headers that varied with scheduling would + // pass the record comparison above. + assert!( + raw_k(n, 3, "det", 1, 16) == raw_k(n, 3, "det", 8, 16), + "thread count changed the gzip bytes" + ); + assert_ne!( + a.0, + run(n, 3, "det2", 1).0, + "a different seed gave the same order" + ); + } + + #[test] + fn single_ended_and_empty_inputs() { + let dir = tempfile::tempdir().unwrap(); + let (f1, _) = temp_inputs(dir.path(), 1_000, 2, false); + let o1 = dir.path().join("o1.fastq.gz"); + let mut rng = NeatRng::new_from_seed(&vec!["se".to_string()]).unwrap(); + combine_temp_fastqs_shuffled(f1, vec![], &o1, None, &mut rng, dir.path(), 2).unwrap(); + assert_eq!(names(&o1).len(), 1_000); + + let (e1, e2) = temp_inputs(dir.path(), 0, 1, true); + let (o1, o2) = ( + dir.path().join("e1.fastq.gz"), + dir.path().join("e2.fastq.gz"), + ); + combine_temp_fastqs_shuffled(e1, e2, &o1, Some(&o2), &mut rng, dir.path(), 2).unwrap(); + assert!(names(&o1).is_empty() && names(&o2).is_empty()); + } + + #[test] + fn unequal_mate_files_are_an_error() { + let dir = tempfile::tempdir().unwrap(); + let (f1, f2) = temp_inputs(dir.path(), 100, 1, true); + // Drop the last read-2 record. + let mut s = String::new(); + MultiGzDecoder::new(File::open(&f2[0]).unwrap()) + .read_to_string(&mut s) + .unwrap(); + let kept: Vec<&str> = s.lines().collect(); + let mut w = BlockGzWriter::new(File::create(&f2[0]).unwrap()); + for l in &kept[..kept.len() - 4] { + writeln!(w, "{l}").unwrap(); + } + w.finish().unwrap(); + let (o1, o2) = ( + dir.path().join("o1.fastq.gz"), + dir.path().join("o2.fastq.gz"), + ); + let mut rng = NeatRng::new_from_seed(&vec!["x".to_string()]).unwrap(); + let err = combine_temp_fastqs_shuffled(f1, f2, &o1, Some(&o2), &mut rng, dir.path(), 1) + .unwrap_err(); + assert!( + err.to_string().contains("different numbers of records"), + "{err}" + ); + } + + #[test] + fn one_bucket_and_many_buckets_both_give_a_shuffled_permutation() { + // k = 1 is all in-bucket shuffle; k = 64 is mostly scatter. Both must permute and mix. + for k in [1, 64] { + let n = 10_000; + let (r1, r2) = run_k(n, 5, "k", 3, k); + let mut idx: Vec = r1.iter().map(|s| index_of(s) as i64).collect(); + let near = idx + .windows(2) + .filter(|w| (w[1] - w[0]).abs() <= 100) + .count() as f64 + / (n - 1) as f64; + assert!(near < 0.05, "k={k}: adjacent within 100: {near:.4}"); + assert!( + r1.iter() + .zip(&r2) + .all(|(a, b)| a.trim_end_matches("/1") == b.trim_end_matches("/2")), + "k={k}" + ); + idx.sort_unstable(); + assert_eq!(idx, (0..n as i64).collect::>(), "k={k}"); + } + } + + #[test] + fn shuffle_threads_fit_the_memory_budget() { + let gib = 1024 * 1024 * 1024; + assert_eq!( + shuffle_threads(64, 256 * 1024 * 1024), + MAX_SHUFFLE_THREADS, + "small buckets: full cap" + ); + assert_eq!(shuffle_threads(64, gib), 4, "1 GiB buckets: 4 fit in 4 GiB"); + assert_eq!( + shuffle_threads(64, 10 * gib), + 1, + "a bucket over budget still runs, alone" + ); + assert_eq!(shuffle_threads(2, 1), 2, "never more than requested"); + assert_eq!(shuffle_threads(0, 0), 1, "at least one"); + } + + const FD_CHILD_ENV: &str = "EIDOLON_SHUFFLE_FD_CHILD"; + /// A file-descriptor limit well under what one writer per bucket per mate would need at the + /// cap (2 * MAX_BUCKETS), and well over what the shuffle needs. + const FD_LIMIT: usize = 48; + + /// Runs only as the child of `cap_path_runs_under_a_low_file_descriptor_limit`, which sets + /// the limit; a no-op otherwise. + #[test] + fn cap_path_child() { + if std::env::var(FD_CHILD_ENV).is_err() { + return; + } + let n = 3_000; + let (r1, r2) = run_k(n, 4, "fd", MAX_SHUFFLE_THREADS, MAX_BUCKETS); + assert_eq!(r1.len(), n); + assert!( + r1.iter() + .zip(&r2) + .all(|(a, b)| a.trim_end_matches("/1") == b.trim_end_matches("/2")) + ); + println!("CAP-PATH-OK {n} pairs through {MAX_BUCKETS} buckets"); + } + + #[cfg(target_os = "linux")] + #[test] + fn cap_path_runs_under_a_low_file_descriptor_limit() { + // Paired, at MAX_BUCKETS, under `ulimit -n 48`. The child prints a marker so a filter + // that matched nothing cannot pass for a run that succeeded. + let exe = std::env::current_exe().unwrap(); + let out = std::process::Command::new("bash") + .arg("-c") + .arg(format!( + "ulimit -n {FD_LIMIT} && exec \"$0\" --exact file_tools::fastq_shuffle::tests::cap_path_child --nocapture --test-threads=1" + )) + .arg(&exe) + .env(FD_CHILD_ENV, "1") + .output() + .unwrap(); + let (stdout, stderr) = ( + String::from_utf8_lossy(&out.stdout), + String::from_utf8_lossy(&out.stderr), + ); + assert!( + out.status.success() && stdout.contains("CAP-PATH-OK"), + "shuffle at {MAX_BUCKETS} buckets failed under ulimit -n {FD_LIMIT}:\n{stdout}\n{stderr}" + ); + } + + #[test] + fn bucket_count_scales_with_input_and_is_capped() { + assert_eq!(bucket_count(0), 1); + assert_eq!(bucket_count(1_000), 1); + assert_eq!( + bucket_count(BUCKET_TARGET_BYTES), + EXPANSION_ESTIMATE as usize + ); + assert_eq!(bucket_count(u64::MAX), MAX_BUCKETS); + } +} diff --git a/eidolon/src/gen_reads/utils/config.rs b/eidolon/src/gen_reads/utils/config.rs index f68c8b8a..ad08b9e9 100644 --- a/eidolon/src/gen_reads/utils/config.rs +++ b/eidolon/src/gen_reads/utils/config.rs @@ -151,6 +151,10 @@ pub struct RunConfiguration { // set it on its own to get the adapter-free short-insert arm. Default off → // short fragments discarded, output byte-identical to before. pub keep_short_fragments: bool, + // Write the final FASTQs in random order rather than genome order (#870). Default on: a real + // sequencer's FASTQ is unordered, and genome order made bwa-mem2 need 166 GB instead of 90 GB + // on a 21x CHM13 simulation. Off restores the byte-for-byte concatenation of chunk outputs. + pub shuffle_fastq: bool, // Scale applied to the SvModel's per-base SV rate at generation time. // 0.0 (default) disables de novo SV generation entirely — only SVs // from `input_vcf` flow through. Values > 0 enable sampling from @@ -230,6 +234,7 @@ impl Default for RunConfiguration { chunk_size: None, long_reads: false, keep_short_fragments: false, + shuffle_fastq: true, sv_rate_scale: 0.0, sv_max_length_fraction: 0.25, adapters: AdapterConfig::default(), @@ -273,6 +278,7 @@ pub const KNOWN_KEYS: &[&str] = &[ "reference", "rng_seed", "sequence_error_model", + "shuffle_fastq", "sv_max_length_fraction", "sv_rate_scale", "target_bed", @@ -720,6 +726,14 @@ impl RunConfiguration { ) })?; } + "shuffle_fastq" => { + config.shuffle_fastq = value.as_bool().ok_or_else(|| { + GenerateReadsError::ConfigReadError( + "shuffle_fastq".to_string(), + "boolean".to_string(), + ) + })?; + } "sv_rate_scale" => { config.sv_rate_scale = value.as_f64().ok_or_else(|| { GenerateReadsError::ConfigReadError( @@ -813,6 +827,9 @@ impl RunConfiguration { info!("\t>mutation rate: {}", rate) } info!("\t>paired ended: {}", config.paired_ended); + if config.produce_fastq { + info!("\t>shuffle FASTQ order: {}", config.shuffle_fastq); + } if config.long_reads { info!("\t>long reads mode: enabled (short fragments produce truncated reads)"); } @@ -1043,6 +1060,7 @@ mod tests { chunk_size: None, long_reads: false, keep_short_fragments: false, + shuffle_fastq: true, sv_rate_scale: 0.0, sv_max_length_fraction: 0.25, adapters: AdapterConfig::default(), diff --git a/eidolon/src/gen_reads/utils/runner.rs b/eidolon/src/gen_reads/utils/runner.rs index 17b7e769..f160f5ee 100644 --- a/eidolon/src/gen_reads/utils/runner.rs +++ b/eidolon/src/gen_reads/utils/runner.rs @@ -17,6 +17,7 @@ use crate::{ bam_writer::{BamBodyWriter, BamContext, BamRecordStager, concat_temp_bams}, bed_reader::read_bed, fasta_stream::{FastaStream, map_buffer, resolve_iupac_bases}, + fastq_shuffle::combine_temp_fastqs_shuffled, fastq_tools::{ HaplotypeContext, PlacedFragment, Strand, combine_temp_fastqs, generate_read, generate_reverse_mate, r2_quality_model, reverse_complement, write_block_fastq, @@ -84,6 +85,43 @@ struct ContigContext<'a> { contig_order: Arc>, } +/// Writes the final FASTQs from the per-chunk temp files: shuffled by default (#870), or +/// concatenated in contig order when `shuffle_fastq: false`. The shuffle draws from its own child +/// of the run's RNG, so it consumes nothing from the streams that generate reads; the truth VCF +/// and golden BAM are the same either way. +fn combine_fastqs( + config: &RunConfiguration, + ctx: &ContigContext<'_>, + all_r1: Vec, + all_r2: Vec, + filename1: &PathBuf, + filename2: Option<&PathBuf>, +) -> Result<(), GenerateReadsError> { + if !config.shuffle_fastq { + combine_temp_fastqs(all_r1, all_r2, filename1, filename2)?; + return Ok(()); + } + info!("Shuffling FASTQ order (shuffle_fastq: true)"); + let mut rng = ctx.base_rng.derive_child(SHUFFLE_RNG_CHILD); + let threads = config + .num_threads + .unwrap_or_else(|| std::thread::available_parallelism().map_or(1, |n| n.get())); + combine_temp_fastqs_shuffled( + all_r1, + all_r2, + filename1, + filename2.map(PathBuf::as_path), + &mut rng, + ctx.working_dir, + threads, + )?; + Ok(()) +} + +/// Child index of the run RNG reserved for the FASTQ shuffle. Contigs use their index and the +/// chimeric pass uses 999999; this is far from both. +const SHUFFLE_RNG_CHILD: u64 = 0x5348_5546_464c_4521; + pub fn run_neat( config: &RunConfiguration, rng: &mut NeatRng, @@ -545,7 +583,14 @@ pub fn run_neat( match &config.output_fastq_2 { Some(filename2) => { create_output_file(filename2, config.overwrite_output)?; - combine_temp_fastqs(all_r1, all_r2, filename1, Some(filename2))?; + combine_fastqs( + config, + &ctx, + all_r1, + all_r2, + filename1, + Some(filename2), + )?; } None => { return Err(GenerateReadsError::ConfigError(format!( @@ -554,7 +599,7 @@ pub fn run_neat( } } } else { - combine_temp_fastqs(all_r1, vec![], filename1, None)?; + combine_fastqs(config, &ctx, all_r1, vec![], filename1, None)?; } } None => { diff --git a/eidolon/tests/common/mod.rs b/eidolon/tests/common/mod.rs index 7f5a6fb2..c65b9433 100644 --- a/eidolon/tests/common/mod.rs +++ b/eidolon/tests/common/mod.rs @@ -96,6 +96,8 @@ pub struct GenReadsConfig { pub mutation_rate: Option, pub mutation_model: Option, pub sv_rate_scale: Option, + /// `None` leaves gen-reads' default, which shuffles the FASTQs (#870). + pub shuffle_fastq: Option, } impl GenReadsConfig { @@ -121,6 +123,7 @@ impl GenReadsConfig { mutation_rate: None, mutation_model: None, sv_rate_scale: None, + shuffle_fastq: None, } } @@ -169,6 +172,10 @@ impl GenReadsConfig { Some(s) => format!("sv_rate_scale: {s}\n"), None => String::new(), }; + let shuffle_section = match self.shuffle_fastq { + Some(b) => format!("shuffle_fastq: {b}\n"), + None => String::new(), + }; write!( f, "reference: {ref_}\n\ @@ -181,7 +188,7 @@ impl GenReadsConfig { output_filename: {name}\n\ overwrite_output: true\n\ rng_seed: {seed}\n\ - {pair}{model}{qmodel}{threads}{chunksz}{ivcf}{mrate}{mmodel}{svscale}", + {pair}{model}{qmodel}{threads}{chunksz}{ivcf}{mrate}{mmodel}{svscale}{shuffle}", ref_ = self.reference.display(), rl = self.read_len, cov = self.coverage, @@ -198,6 +205,7 @@ impl GenReadsConfig { chunksz = chunk_size_section, ivcf = input_vcf_section, mrate = mutation_rate_section, + shuffle = shuffle_section, mmodel = mutation_model_section, svscale = sv_rate_scale_section, ) diff --git a/eidolon/tests/determinism.rs b/eidolon/tests/determinism.rs index 2428545b..db0b0590 100644 --- a/eidolon/tests/determinism.rs +++ b/eidolon/tests/determinism.rs @@ -41,11 +41,23 @@ fn run_gen_reads_cfg( name: &str, threads: Option, sv_rate_scale: Option, +) { + run_gen_reads_full(seed, output_dir, name, threads, sv_rate_scale, None) +} + +fn run_gen_reads_full( + seed: &str, + output_dir: &std::path::Path, + name: &str, + threads: Option, + sv_rate_scale: Option, + shuffle_fastq: Option, ) { let mut config = GenReadsConfig::new(h1n1_reference(), output_dir.to_path_buf(), name); config.rng_seed = seed.to_string(); config.num_threads = threads; config.sv_rate_scale = sv_rate_scale; + config.shuffle_fastq = shuffle_fastq; let yaml = config.write_yaml(); eidolon() .args(["gen-reads", "-c"]) @@ -213,10 +225,21 @@ fn byte_identical_output_with_svs_across_thread_counts() { /// them and that section is legitimately not sorted by contig. Only the first section is /// asserted here; chimeric block ordering is left unpinned rather than pinned to something /// this test would have to guess at. +/// +/// Since #870 the default output is shuffled, so this pins the `shuffle_fastq: false` path. +/// That order still matters with the shuffle on: it is the input order the shuffle's +/// determinism rests on. #[test] fn ordinary_reads_are_emitted_in_reference_file_order() { let (_g, dir) = fresh_workdir(); - run_gen_reads_cfg("contig order seed", &dir, "order", Some(1), Some(50.0)); + run_gen_reads_full( + "contig order seed", + &dir, + "order", + Some(1), + Some(50.0), + Some(false), + ); let lines = read_gzip_fastq_lines(&dir.join("order_r1.fastq.gz")); let reference_order = [ diff --git a/eidolon/tests/fastq_shuffle_output.rs b/eidolon/tests/fastq_shuffle_output.rs new file mode 100644 index 00000000..a5f47e73 --- /dev/null +++ b/eidolon/tests/fastq_shuffle_output.rs @@ -0,0 +1,165 @@ +//! gen-reads writes its FASTQs in random order by default (#870). +//! +//! Genome order made bwa-mem2 need 166 GB instead of 90 GB on a 21x T2T-CHM13 simulation, with +//! identical alignments either way. A real sequencer's FASTQ is unordered, so that is the default. +//! +//! **Known answer.** Read names carry each read's origin (`…___/1`). In genome +//! order nearly every read's neighbor starts within a few hundred bases of it; in random order on +//! H1N1 (13.5 kb), a neighbor on the same contig within 200 bp is rare. **Must not fire:** +//! `shuffle_fastq: false` keeps genome order, and the shuffle changes nothing but order. The same +//! records come out, mates stay in register, and the truth VCF and golden BAM are byte-identical, +//! because the shuffle draws from its own RNG child. + +mod common; + +use common::{eidolon, h1n1_reference}; +use noodles::bam; +use std::io::Read; +use std::path::Path; + +fn gz_text(path: &Path) -> String { + let mut s = String::new(); + flate2::read::MultiGzDecoder::new(std::fs::File::open(path).unwrap()) + .read_to_string(&mut s) + .unwrap(); + s +} + +fn run(dir: &Path, name: &str, shuffle: Option, threads: usize) { + let cfg = dir.join(format!("{name}.yml")); + let shuffle_line = shuffle + .map(|s| format!("shuffle_fastq: {s}\n")) + .unwrap_or_default(); + std::fs::write( + &cfg, + format!( + "reference: {ref}\nread_len: 101\ncoverage: 30\npaired_ended: true\n\ + fragment_mean: 300\nfragment_st_dev: 30\n\ + produce_fastq: true\nproduce_vcf: true\nproduce_bam: true\n\ + overwrite_output: true\noutput_dir: {out}\noutput_filename: {name}\n\ + rng_seed: fastq shuffle\nnum_threads: {threads}\n{shuffle_line}", + ref = h1n1_reference().display(), + out = dir.display(), + ), + ) + .unwrap(); + eidolon() + .args(["gen-reads", "-c"]) + .arg(&cfg) + .assert() + .success(); +} + +/// (contig, start) parsed from a read name, counting fields from the right because contig names +/// may contain underscores. +fn origin(name: &str) -> (String, i64) { + let core = name + .trim_start_matches('@') + .trim_start_matches("EIDOLON_generated_"); + let core = core.rsplit_once('/').map_or(core, |(c, _)| c); + let parts: Vec<&str> = core.rsplitn(4, '_').collect(); + (parts[3].to_string(), parts[2].parse().unwrap()) +} + +fn names(fastq: &str) -> Vec { + fastq.lines().step_by(4).map(str::to_string).collect() +} + +/// Fraction of adjacent reads on the same contig with starts within 200 bp. +fn near_neighbors(names: &[String]) -> f64 { + let o: Vec<(String, i64)> = names.iter().map(|n| origin(n)).collect(); + let near = o + .windows(2) + .filter(|w| w[0].0 == w[1].0 && (w[0].1 - w[1].1).abs() <= 200) + .count(); + near as f64 / (o.len() - 1) as f64 +} + +#[test] +fn default_output_is_shuffled_and_off_keeps_genome_order() { + let tmp = tempfile::tempdir().unwrap(); + let d = tmp.path(); + run(d, "on", None, 2); + run(d, "off", Some(false), 2); + + let on1 = names(&gz_text(&d.join("on_r1.fastq.gz"))); + let on2 = names(&gz_text(&d.join("on_r2.fastq.gz"))); + let off1 = names(&gz_text(&d.join("off_r1.fastq.gz"))); + // Random order on H1N1 puts a neighbor on the same contig ~15% of the time and within 200 bp + // of it ~23% of that (contigs ~1.7 kb): ~0.035 expected, sigma ~0.006 at 1000 reads. So 1000 + // reads keep the < 0.1 assertion below far from noise. + assert!( + on1.len() > 1000, + "only {} reads; the order statistics would be weak", + on1.len() + ); + + let (near_on, near_off) = (near_neighbors(&on1), near_neighbors(&off1)); + eprintln!("adjacent reads within 200 bp: shuffled {near_on:.4}, genome order {near_off:.4}"); + assert!( + near_off > 0.9, + "shuffle_fastq: false should keep genome order: {near_off:.4}" + ); + assert!( + near_on < 0.1, + "the default output is not shuffled: {near_on:.4}" + ); + + // Same records, mates in register. + for (a, b) in on1.iter().zip(&on2) { + assert_eq!( + a.trim_end_matches("/1"), + b.trim_end_matches("/2"), + "mates out of register" + ); + } + let (mut s_on, mut s_off) = (on1.clone(), off1.clone()); + s_on.sort(); + s_off.sort(); + assert_eq!( + s_on, s_off, + "the shuffle must change order only, not which reads exist" + ); + + // The shuffle draws from its own RNG child: the truth VCF and the golden BAM are unaffected. + assert_eq!( + gz_text(&d.join("on.vcf.gz")), + gz_text(&d.join("off.vcf.gz")), + "truth VCF changed" + ); + let bam_records = |p: &Path| -> Vec { + let mut r = bam::io::Reader::new(std::fs::File::open(p).unwrap()); + r.read_header().unwrap(); + r.records().map(|x| format!("{:?}", x.unwrap())).collect() + }; + assert_eq!( + bam_records(&d.join("on.bam")), + bam_records(&d.join("off.bam")), + "golden BAM changed" + ); +} + +#[test] +fn shuffled_output_is_identical_across_thread_counts() { + let tmp = tempfile::tempdir().unwrap(); + let d = tmp.path(); + run(d, "t1", None, 1); + run(d, "t4", None, 4); + for mate in ["r1", "r2"] { + let (a, b) = ( + d.join(format!("t1_{mate}.fastq.gz")), + d.join(format!("t4_{mate}.fastq.gz")), + ); + assert_eq!( + gz_text(&a), + gz_text(&b), + "{mate} records differ between 1 and 4 threads" + ); + // The files themselves, not just their decompressed text: gzip member boundaries or + // headers that varied with thread count would pass the line above. + assert!( + std::fs::read(&a).unwrap() == std::fs::read(&b).unwrap(), + "{mate} gzip bytes differ between 1 and 4 threads" + ); + } +} diff --git a/template_config/gen_reads_template.yml b/template_config/gen_reads_template.yml index e6225675..7aaba07e 100644 --- a/template_config/gen_reads_template.yml +++ b/template_config/gen_reads_template.yml @@ -55,6 +55,12 @@ fragment_st_dev: . # At least one of produce_fastq, produce_vcf, or produce_bam must be true. produce_fastq: true +# Write the FASTQ(s) in random order, as a sequencer does, rather than genome order. Mates stay +# in register. Only record order changes: the truth VCF and BAM are identical either way, and the +# same seed gives the same output at any thread count. Genome order made bwa-mem2 need ~1.8x the +# memory on a T2T-CHM13 simulation (#870). (default: true) +shuffle_fastq: true + # Produce a gzipped VCF file recording all inserted variants. (default: false) produce_vcf: false