diff --git a/modkit-core/src/command_utils.rs b/modkit-core/src/command_utils.rs index 29ee4689..e1ca6858 100644 --- a/modkit-core/src/command_utils.rs +++ b/modkit-core/src/command_utils.rs @@ -16,6 +16,23 @@ use crate::threshold_mod_caller::MultipleThresholdModCaller; use crate::thresholds::calc_threshold_from_bam; use crate::util::{create_out_directory, Region}; +pub(crate) fn parse_sampling_fraction(raw: &str) -> Result { + let fraction = raw.parse::().map_err(|_| { + format!( + "sampling fraction must be a finite number in the inclusive \ + range [0, 1]; got '{raw}'" + ) + })?; + if fraction.is_finite() && (0.0..=1.0).contains(&fraction) { + Ok(fraction) + } else { + Err(format!( + "sampling fraction must be a finite number in the inclusive \ + range [0, 1]; got '{raw}'" + )) + } +} + pub fn parse_per_mod_thresholds( raw_per_mod_thresholds: &[String], ) -> anyhow::Result> { @@ -454,3 +471,41 @@ pub(crate) fn parse_raw_motifs( e @ _ => e, } } + +#[cfg(test)] +mod tests { + use super::parse_sampling_fraction; + + #[test] + fn parse_sampling_fraction_matrix() { + for (raw, expected) in [ + ("0", 0.0), + ("-0.0", -0.0), + ("0.25", 0.25), + ("1e-3", 0.001), + ("1", 1.0), + ("1e0", 1.0), + ] { + assert_eq!(parse_sampling_fraction(raw).unwrap(), expected); + } + assert!(parse_sampling_fraction("-0.0").unwrap().is_sign_negative()); + + for raw in [ + "-0.0000001", + "1.0000001", + "NaN", + "nan", + "inf", + "+inf", + "-inf", + "1e309", + "not-a-number", + "", + ] { + assert!( + parse_sampling_fraction(raw).is_err(), + "expected {raw:?} to be rejected" + ); + } + } +} diff --git a/modkit-core/src/extract/subcommand.rs b/modkit-core/src/extract/subcommand.rs index 42b4d6ed..f4fb511a 100644 --- a/modkit-core/src/extract/subcommand.rs +++ b/modkit-core/src/extract/subcommand.rs @@ -17,7 +17,7 @@ use modkit_logging::{init_logging, init_logging_smart}; use crate::command_utils::{ get_serial_reader, parse_edge_filter_input, parse_per_base_thresholds, parse_per_mod_thresholds, parse_raw_thresholds_string_with_default, - parse_thresholds, using_stream, + parse_sampling_fraction, parse_thresholds, using_stream, }; use crate::extract::args::InputArgs; use crate::extract::util::{ @@ -450,12 +450,14 @@ pub struct EntryExtractCalls { /// In practice, 10-100 thousand reads is sufficient to estimate the model /// output distribution and determine the filtering threshold. See /// filtering.md for details on filtering. + /// Must be a finite value in the inclusive range [0, 1]. #[clap(help_heading = "Sampling Options")] #[arg( group = "sampling_options", short = 'f', long, - hide_short_help = true + hide_short_help = true, + value_parser = parse_sampling_fraction )] sampling_frac: Option, /// Sample this many reads when estimating the filtering threshold. If a diff --git a/modkit-core/src/modbam_util/subcommands.rs b/modkit-core/src/modbam_util/subcommands.rs index 509b270b..b06d7fe2 100644 --- a/modkit-core/src/modbam_util/subcommands.rs +++ b/modkit-core/src/modbam_util/subcommands.rs @@ -22,7 +22,8 @@ use crate::command_utils::{ get_bam_writer, get_motif_lookup_from_parts, get_serial_reader, get_threshold_from_options, parse_edge_filter_input, parse_forward_motifs, parse_per_mod_thresholds, parse_raw_motifs, - parse_raw_thresholds_string_with_default, parse_thresholds, using_stream, + parse_raw_thresholds_string_with_default, parse_sampling_fraction, + parse_thresholds, using_stream, }; use crate::errs::{MkError, MkResult}; use crate::interval_chunks::{ @@ -39,6 +40,7 @@ use crate::modbam_util::check_tags::ModTagViews; use crate::monoid::Moniod; use crate::motifs::motif_bed::RegexMotif; use crate::position_filter::StrandedPositionFilter; +use crate::reads_sampler::deterministic_sampler::resolve_master_seed; use crate::reads_sampler::record_sampler::RecordSampler; use crate::reads_sampler::sample_reads_from_interval; use crate::reads_sampler::sampling_schedule::SamplingSchedule; @@ -1012,21 +1014,23 @@ pub struct SampleModBaseProbs { num_reads: Option, /// Instead of using a defined number of reads, specify a fraction of reads /// to sample, for example 0.1 will sample 1/10th of the reads. + /// Must be a finite value in the inclusive range [0, 1]. #[clap(help_heading = "Sampling Options")] #[arg( short = 'f', long, alias = "sample-frac", group = "sampling_options", - conflicts_with = "no_sampling" + conflicts_with = "no_sampling", + value_parser = parse_sampling_fraction )] sampling_frac: Option, /// No sampling, use all of the reads to calculate the filter thresholds. #[clap(help_heading = "Sampling Options")] #[arg(long, alias = "no_filtering", default_value_t = false)] no_sampling: bool, - /// Random seed for deterministic running, the default is - /// non-deterministic, only used when no BAM index is provided. + /// Provide a seed to make fractional read sampling decisions repeatable + /// for indexed and unindexed inputs. #[clap(help_heading = "Sampling Options")] #[arg(short, conflicts_with = "no_sampling", long)] seed: Option, @@ -1247,6 +1251,11 @@ impl SampleModBaseProbs { 0f64, ) }; + let master_seed = if rng_sample { + resolve_master_seed(self.seed) + } else { + self.seed.unwrap_or_default() + }; let mut workers: Vec> = Vec::with_capacity(self.threads); @@ -1273,7 +1282,7 @@ impl SampleModBaseProbs { "collecting base and modification histograms at aligned \ positions" ); - for i in 0..self.threads { + for _ in 0..self.threads { let w = RegionMleProbs::< AlignedBaseAndModArgmaxProbs, ProbsExtractor, @@ -1284,7 +1293,7 @@ impl SampleModBaseProbs { motif_bases, edge_filter.as_ref(), &thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1293,7 +1302,7 @@ impl SampleModBaseProbs { } } else { info!("collecting base-level histograms at aligned positions"); - for i in 0..self.threads { + for _ in 0..self.threads { let w = RegionMleProbs::< AlignedBaseArgmaxProbs, ProbsExtractor, @@ -1304,7 +1313,7 @@ impl SampleModBaseProbs { motif_bases, edge_filter.as_ref(), &thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1318,7 +1327,7 @@ impl SampleModBaseProbs { "collecting base and modification histograms, using all \ read positions" ); - for i in 0..self.threads { + for _ in 0..self.threads { let w = RegionMleProbs::< BaseAndModArgmaxProbs, ProbsExtractor, @@ -1329,7 +1338,7 @@ impl SampleModBaseProbs { motif_bases, edge_filter.as_ref(), &thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1341,7 +1350,7 @@ impl SampleModBaseProbs { "collecting base-level histograms, using all read \ positions" ); - for i in 0..self.threads { + for _ in 0..self.threads { let w = RegionMleProbs::::new( &bam_fp, @@ -1350,7 +1359,7 @@ impl SampleModBaseProbs { motif_bases, edge_filter.as_ref(), &thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1638,8 +1647,14 @@ pub struct ModSummarize { /// Instead of using a defined number of reads, specify a fraction of reads /// to sample when estimating the filter threshold. For example 0.1 will /// sample 1/10th of the reads. + /// Must be a finite value in the inclusive range [0, 1]. #[clap(help_heading = "Sampling Options")] - #[arg(group = "sampling_options", short = 'f', long)] + #[arg( + group = "sampling_options", + short = 'f', + long, + value_parser = parse_sampling_fraction + )] sampling_frac: Option, /// Sets a random seed for deterministic running (when using /// --sample-frac). @@ -2257,15 +2272,17 @@ pub struct CallMods { /// filter-percentile. In practice, 50-100 thousand reads is sufficient /// to estimate the model output distribution and determine the /// filtering threshold. See filtering.md for details on filtering. + /// Must be a finite value in the inclusive range [0, 1]. #[arg( group = "sampling_options", short = 'f', long, - hide_short_help = true + hide_short_help = true, + value_parser = parse_sampling_fraction )] sampling_frac: Option, - /// Set a random seed for deterministic running, the default is - /// non-deterministic, only used when no BAM index is provided. + /// Provide a seed to make fractional read sampling decisions repeatable + /// for indexed and unindexed inputs. #[arg( long, conflicts_with = "num_reads", diff --git a/modkit-core/src/pileup/subcommand.rs b/modkit-core/src/pileup/subcommand.rs index 4503f99e..406082c6 100644 --- a/modkit-core/src/pileup/subcommand.rs +++ b/modkit-core/src/pileup/subcommand.rs @@ -19,8 +19,8 @@ use crate::command_utils::{ get_motif_lookup_from_parts, get_threshold_from_options, parse_edge_filter_input, parse_per_base_thresholds, parse_per_mod_thresholds, parse_raw_motifs, - parse_raw_thresholds_string_with_default, parse_thresholds, - parse_thresholds_values, + parse_raw_thresholds_string_with_default, parse_sampling_fraction, + parse_thresholds, parse_thresholds_values, }; use crate::fasta::MotifLocationsLookup; use crate::interval_chunks::{ @@ -41,6 +41,7 @@ use crate::pileup::pileup_processor::{ }; use crate::pileup::{ModBasePileup2, PileupNumericOptions}; use crate::position_filter::StrandedPositionFilter; +use crate::reads_sampler::deterministic_sampler::resolve_master_seed; use crate::reads_sampler::sampling_schedule::IdxStats; use crate::sample_probs::{ calculate_reads_per_contig, run_extract_probs_workers, @@ -172,12 +173,14 @@ pub struct ModBamPileup { /// In practice, 10-100 thousand reads is sufficient to estimate the model /// output distribution and determine the filtering threshold. See /// filtering.md for details on filtering. + /// Must be a finite value in the inclusive range [0, 1]. #[clap(help_heading = "Sampling Options")] #[arg( group = "sampling_options", short = 'f', long, - hide_short_help = true + hide_short_help = true, + value_parser = parse_sampling_fraction )] sampling_frac: Option, /// Set a random seed for deterministic running, the default is @@ -776,6 +779,11 @@ impl ModBamPileup { self.sampling_interval_size, sampling_region, )?; + let master_seed = if rng_sample { + resolve_master_seed(self.seed) + } else { + self.seed.unwrap_or_default() + }; let chrom_to_counts = chrom_to_counts.map(|x| Arc::new(x)); if let Some(preset) = preset { @@ -806,7 +814,7 @@ impl ModBamPileup { *motif_bases } }; - for i in 0..n_workers { + for _ in 0..n_workers { let worker = RegionMleProbs::< AlignedBaseArgmaxProbs, ProbsExtractor, @@ -817,7 +825,7 @@ impl ModBamPileup { motif_bases, edge_filter, thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -832,7 +840,7 @@ impl ModBamPileup { motif_bases.iter().unique().join(",") ) }); - for i in 0..n_workers { + for _ in 0..n_workers { let worker = RegionMleProbs::< AlignedBaseArgmaxProbs, ProbsExtractor, @@ -843,7 +851,7 @@ impl ModBamPileup { motif_bases, edge_filter, thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -856,7 +864,7 @@ impl ModBamPileup { "calculating threshold value with probabilites from reads", ) }); - for i in 0..n_workers { + for _ in 0..n_workers { let worker = RegionMleProbs::::new( &self.in_bam, @@ -865,7 +873,7 @@ impl ModBamPileup { [DnaBase::A; 4], edge_filter, thread_pool, - self.seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1787,12 +1795,14 @@ pub struct DuplexModBamPileup { /// filter-percentile. In practice, 50-100 thousand reads is sufficient /// to estimate the model output distribution and determine the /// filtering threshold. See filtering.md for details on filtering. + /// Must be a finite value in the inclusive range [0, 1]. #[clap(help_heading = "Sampling Options")] #[arg( group = "sampling_options", short = 'f', long, - hide_short_help = true + hide_short_help = true, + value_parser = parse_sampling_fraction )] sampling_frac: Option, /// Set a random seed for deterministic running, the default is diff --git a/modkit-core/src/read_ids_to_base_mod_probs.rs b/modkit-core/src/read_ids_to_base_mod_probs.rs index 39513aa2..68860361 100644 --- a/modkit-core/src/read_ids_to_base_mod_probs.rs +++ b/modkit-core/src/read_ids_to_base_mod_probs.rs @@ -131,7 +131,14 @@ impl ReadIdsToBaseModProbs { } } if let Some(p) = explicit_prob { - canonical_probs.insert(*dna_base, p); + canonical_probs + .entry(*dna_base) + .and_modify(|current| { + if p > *current { + *current = p; + } + }) + .or_insert(p); } } (calls_per_base, canonical_probs) @@ -276,7 +283,7 @@ impl RecordProcessor for ReadIdsToBaseModProbs { position_filter: Option<&StrandedPositionFilter<()>>, only_mapped: bool, allow_non_primary: bool, - _cut: Option, + cut: Option, _kmer_size: Option, ) -> anyhow::Result { let spinner = if with_progress { @@ -285,6 +292,15 @@ impl RecordProcessor for ReadIdsToBaseModProbs { None }; let mod_base_info_iter = records + .filter(|result| { + result + .as_ref() + .map(|record| { + cut.map(|cut| record.reference_start() >= cut as i64) + .unwrap_or(true) + }) + .unwrap_or(true) + }) .with_mod_base_info() .filter(|(record, _)| { if only_mapped || edge_filter.is_some() { @@ -302,7 +318,7 @@ impl RecordProcessor for ReadIdsToBaseModProbs { }); let mut read_ids_to_mod_base_probs = Self::zero(); for (record, mod_base_info) in mod_base_info_iter { - match record_sampler.ask() { + match record_sampler.ask_record(&record) { Indicator::Use(token) => { let record_name = get_query_name_string(&record); let aligned_pairs = if only_mapped { @@ -1293,10 +1309,36 @@ mod read_ids_to_base_mod_probs_tests { use rust_htslib::bam::{self, Read}; use rustc_hash::{FxHashMap, FxHashSet}; - use crate::mod_bam::filter_records_iter; + use crate::mod_bam::{filter_records_iter, BaseModProbs}; + use crate::mod_base_code::DnaBase; use crate::position_filter::StrandedPositionFilter; use crate::util::get_aligned_pairs_forward; + use super::ReadIdsToBaseModProbs; + + #[test] + fn explicit_canonical_probability_is_global_max_across_reads() { + let pool = + rayon::ThreadPoolBuilder::new().num_threads(1).build().unwrap(); + + for repetition in 0..8 { + let mut reads = ReadIdsToBaseModProbs { inner: HashMap::new() }; + for index in 1..100 { + // The canonical probability is 1 - modified probability, so + // the maximum explicit canonical probability is 0.99. + reads.add_mod_probs_for_read( + &format!("read-{repetition}-{index}"), + DnaBase::C, + vec![BaseModProbs::new_init('m', index as f32 / 100.0)], + ); + } + + let (_, explicit_canonical) = + pool.install(|| reads.mle_probs_per_base()); + assert_eq!(explicit_canonical[&DnaBase::C], 0.99); + } + } + #[test] fn test_seq_pos_base_mod_probs_filter_positions() { let mut reader = bam::Reader::from_path( diff --git a/modkit-core/src/reads_sampler/deterministic_sampler.rs b/modkit-core/src/reads_sampler/deterministic_sampler.rs new file mode 100644 index 00000000..9866c950 --- /dev/null +++ b/modkit-core/src/reads_sampler/deterministic_sampler.rs @@ -0,0 +1,299 @@ +use anyhow::{bail, Result}; +use rust_htslib::bam; + +/// Domain separation and compatibility version for the stable alignment +/// sampler. Changing the identity fields, byte encoding, digest, or threshold +/// rule requires a new version. +const SAMPLER_V1_DOMAIN: &[u8] = b"modkit-alignment-fraction-sampler-v1\0"; + +/// Alignment-identity flag bits: reverse, read1, read2, secondary, and +/// supplementary. Mutable annotation flags such as duplicate, QC failure, and +/// proper pair are intentionally excluded. +const ALIGNMENT_IDENTITY_FLAG_MASK: u16 = 0x09d0; +const DRAW_SPACE: f64 = (1u64 << 53) as f64; +const FNV1A_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325; +const FNV1A_PRIME: u64 = 0x0000_0100_0000_01b3; + +/// The v1 portable hash is FNV-1a-64 followed by the SplitMix64 finalizer. +/// These operations and constants are part of the sampler compatibility +/// contract, rather than an implementation-selected `Hash` algorithm. +struct StableV1Hasher(u64); + +impl StableV1Hasher { + fn new(master_seed: u64) -> Self { + let mut hasher = Self(FNV1A_OFFSET_BASIS); + hasher.update(SAMPLER_V1_DOMAIN); + hasher.update(&master_seed.to_le_bytes()); + hasher + } + + #[inline] + fn update(&mut self, bytes: &[u8]) { + for byte in bytes { + self.0 ^= *byte as u64; + self.0 = self.0.wrapping_mul(FNV1A_PRIME); + } + } + + fn finish(mut self) -> u64 { + self.0 ^= self.0 >> 30; + self.0 = self.0.wrapping_mul(0xbf58_476d_1ce4_e5b9); + self.0 ^= self.0 >> 27; + self.0 = self.0.wrapping_mul(0x94d0_49bb_1331_11eb); + self.0 ^ (self.0 >> 31) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum InclusionThreshold { + None, + Below(u64), + All, +} + +/// A versioned, stateless Bernoulli sampler for BAM alignment records. +/// +/// The v1 sampling unit is one alignment record, not one molecule or QNAME. +/// Its stable identity consists of the raw QNAME, numeric reference ID, +/// alignment start, canonical CIGAR operations, and the masked flag bits in +/// [`ALIGNMENT_IDENTITY_FLAG_MASK`]. This distinguishes ordinary primary, +/// secondary, supplementary, strand, and read-end alignments sharing a QNAME, +/// while ensuring an alignment refetched in multiple genomic intervals gets +/// one decision. Byte-identical duplicate alignment records intentionally +/// share a decision because BAM records do not expose a portable stable +/// ordinal. Numeric reference IDs make this contract stable for refetches of +/// the same BAM/CRAM header; reproducibility after reference-header reordering +/// is not promised. +/// +/// MAPQ, mate fields, SEQ, QUAL, and auxiliary tags (including MM/ML) are not +/// identity. They are mutable annotations or measured content and must not +/// influence whether the observation is sampled. +#[derive(Clone, Copy, Debug)] +pub(crate) struct DeterministicFractionSampler { + master_seed: u64, + threshold: InclusionThreshold, +} + +impl DeterministicFractionSampler { + pub(crate) fn new(master_seed: u64, fraction: f64) -> Result { + if !fraction.is_finite() || !(0.0..=1.0).contains(&fraction) { + bail!( + "sampling fraction must be a finite number in the inclusive \ + range [0, 1]; got '{fraction}'" + ); + } + let threshold = if fraction == 0.0 { + InclusionThreshold::None + } else if fraction == 1.0 { + InclusionThreshold::All + } else { + InclusionThreshold::Below((fraction * DRAW_SPACE).floor() as u64) + }; + Ok(Self { master_seed, threshold }) + } + + #[inline] + pub(crate) fn include(&self, record: &bam::Record) -> bool { + match self.threshold { + InclusionThreshold::None => false, + InclusionThreshold::All => true, + InclusionThreshold::Below(threshold) => { + (self.alignment_score(record) >> 11) < threshold + } + } + } + + /// Return the v1 score used by the inclusion threshold. Integer fields use + /// fixed-width little-endian encoding (signed fields use their two's + /// complement bytes), and variable-length fields have an explicit u64 + /// length. Inclusion compares the score's top 53 bits with + /// `floor(fraction * 2^53)`. + pub(crate) fn alignment_score(&self, record: &bam::Record) -> u64 { + let mut hash = StableV1Hasher::new(self.master_seed); + + let qname = record.qname(); + hash.update(&(qname.len() as u64).to_le_bytes()); + hash.update(qname); + hash.update(&record.tid().to_le_bytes()); + hash.update(&record.pos().to_le_bytes()); + hash.update( + &(record.flags() & ALIGNMENT_IDENTITY_FLAG_MASK).to_le_bytes(), + ); + + let cigar = record.cigar(); + hash.update(&(cigar.len() as u64).to_le_bytes()); + for op in cigar.iter() { + hash.update(&op.len().to_le_bytes()); + hash.update(&[op.char() as u8]); + } + hash.finish() + } +} + +/// Resolve one seed for the whole indexed sampling job. The resolved value +/// must be shared by every worker; resolving per worker would reintroduce +/// scheduling-dependent decisions. An omitted seed remains nondeterministic as +/// documented by the CLI. +pub(crate) fn resolve_master_seed(seed: Option) -> u64 { + seed.unwrap_or_else(rand::random) +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use rust_htslib::bam::{ + self, + record::{Aux, Cigar, CigarString}, + }; + + use super::DeterministicFractionSampler; + + fn record(name: &str, pos: i64, flags: u16) -> bam::Record { + let mut record = bam::Record::new(); + let cigar = CigarString(vec![Cigar::Match(8)]); + record.set(name.as_bytes(), Some(&cigar), b"ACGTACGT", &[20; 8]); + record.set_tid(3); + record.set_pos(pos); + record.set_flags(flags); + record.push_aux(b"MM", Aux::String("C+m?,0;")).unwrap(); + record.push_aux(b"ML", Aux::ArrayU8((&[128][..]).into())).unwrap(); + record + } + + #[test] + fn boundary_fractions_are_exact() { + let record = record("read", 10, 0); + for seed in [0, 7, u64::MAX] { + assert!(!DeterministicFractionSampler::new(seed, 0.0) + .unwrap() + .include(&record)); + assert!(DeterministicFractionSampler::new(seed, 1.0) + .unwrap() + .include(&record)); + } + } + + #[test] + fn rejects_non_finite_or_out_of_range_fractions() { + for fraction in [f64::NEG_INFINITY, -0.1, 1.1, f64::INFINITY, f64::NAN] + { + assert!(DeterministicFractionSampler::new(7, fraction).is_err()); + } + } + + #[test] + fn identity_distinguishes_alignment_records_sharing_qname() { + let sampler = DeterministicFractionSampler::new(7, 0.5).unwrap(); + let mut different_cigar = record("shared", 10, 0); + different_cigar.set( + b"shared", + Some(&CigarString(vec![Cigar::Equal(8)])), + b"ACGTACGT", + &[20; 8], + ); + let records = [ + record("shared", 10, 0), + record("shared", 11, 0), + different_cigar, + record("shared", 10, 0x10), + record("shared", 10, 0x40), + record("shared", 10, 0x80), + record("shared", 10, 0x100), + record("shared", 10, 0x800), + ]; + let scores = records + .iter() + .map(|record| sampler.alignment_score(record)) + .collect::>(); + + assert_eq!(scores.iter().copied().collect::>().len(), 8); + } + + #[test] + fn mutable_annotations_and_measured_content_are_not_identity() { + let sampler = DeterministicFractionSampler::new(7, 0.5).unwrap(); + let original = record("read", 10, 0); + let expected = sampler.alignment_score(&original); + + let mut changed = record("read", 10, 0); + changed.set_mapq(42); + changed.set_mtid(8); + changed.set_mpos(900); + changed.set_insert_size(1234); + changed.set( + b"read", + Some(&CigarString(vec![Cigar::Match(8)])), + b"TTTTTTTT", + &[40; 8], + ); + changed.remove_aux(b"MM").unwrap(); + changed.remove_aux(b"ML").unwrap(); + changed.push_aux(b"MM", Aux::String("T+m?,0;")).unwrap(); + changed.push_aux(b"ML", Aux::ArrayU8((&[250][..]).into())).unwrap(); + changed.set_flags(0x2 | 0x200 | 0x400); + + assert_eq!(sampler.alignment_score(&changed), expected); + } + + #[test] + fn fixed_v1_score_locks_hash_compatibility() { + let sampler = DeterministicFractionSampler::new(7, 0.5).unwrap(); + assert_eq!( + sampler.alignment_score(&record("read", 10, 0x10 | 0x800)), + 6_384_807_673_264_054_871 + ); + } + + #[test] + fn master_seed_and_integer_threshold_are_part_of_v1() { + let record = record("read", 10, 0x10 | 0x800); + let sampler = DeterministicFractionSampler::new(7, 0.5).unwrap(); + assert_ne!( + sampler.alignment_score(&record), + DeterministicFractionSampler::new(8, 0.5) + .unwrap() + .alignment_score(&record) + ); + + let draw = sampler.alignment_score(&record) >> 11; + let at_draw = draw as f64 / super::DRAW_SPACE; + let just_above = (draw + 1) as f64 / super::DRAW_SPACE; + assert!(!DeterministicFractionSampler::new(7, at_draw) + .unwrap() + .include(&record)); + assert!(DeterministicFractionSampler::new(7, just_above) + .unwrap() + .include(&record)); + } + + #[test] + fn fixed_population_is_monotone_and_statistically_representative() { + const POPULATION_SIZE: usize = 10_000; + const INTERIOR_FRACTION: f64 = 0.37; + let low = DeterministicFractionSampler::new(7, 0.1).unwrap(); + let interior = + DeterministicFractionSampler::new(7, INTERIOR_FRACTION).unwrap(); + let high = DeterministicFractionSampler::new(7, 0.9).unwrap(); + let mut included = 0usize; + + for i in 0..POPULATION_SIZE { + let record = record(&format!("population-{i}"), i as i64, 0); + let at_low = low.include(&record); + let at_interior = interior.include(&record); + let at_high = high.include(&record); + assert!(!at_low || at_interior); + assert!(!at_interior || at_high); + included += at_interior as usize; + } + + let expected = POPULATION_SIZE as f64 * INTERIOR_FRACTION; + let standard_deviation = (expected * (1.0 - INTERIOR_FRACTION)).sqrt(); + let conservative_tolerance = 8.0 * standard_deviation; + assert!( + (included as f64 - expected).abs() <= conservative_tolerance, + "fixed population included {included} records; expected about \ + {expected} within {conservative_tolerance}" + ); + } +} diff --git a/modkit-core/src/reads_sampler/mod.rs b/modkit-core/src/reads_sampler/mod.rs index 7dfc6b4b..1ee69a06 100644 --- a/modkit-core/src/reads_sampler/mod.rs +++ b/modkit-core/src/reads_sampler/mod.rs @@ -24,6 +24,9 @@ use crate::util::{ }; use record_sampler::RecordSampler; +use self::deterministic_sampler::resolve_master_seed; + +pub(crate) mod deterministic_sampler; pub mod record_sampler; pub mod sampling_schedule; @@ -46,6 +49,11 @@ where { let use_regions = bam::IndexedReader::from_path(&bam_fp).is_ok(); if use_regions { + // Resolve once for the mapped indexed job so every interval worker + // makes the same identity-based decision. Count sampling, + // passthrough, and the unmapped fallback preserve existing behavior. + let indexed_fraction_seed = + sample_frac.map(|_| resolve_master_seed(seed)); debug!( "found BAM index, sampling reads in {interval_size} base pair \ chunks" @@ -83,6 +91,7 @@ where collapse_method, position_filter, &schedule, + sample_frac.zip(indexed_fraction_seed), only_mapped, suppress_progress, )?; @@ -99,6 +108,9 @@ where let num_reads_unmapped = num_reads.map(|nr| { nr.checked_sub(read_ids_to_base_mod_calls.len()).unwrap_or(0) }); + // Keep the indexed-unmapped fallback on its legacy, sequential + // sampler. Its semantics are tracked separately from mapped + // indexed sampling. let record_sampler = RecordSampler::new_from_options( sample_frac, num_reads_unmapped, @@ -169,6 +181,7 @@ fn sample_reads_base_mod_calls_over_regions( collapse_method: Option<&CollapseMethod>, position_filter: Option<&StrandedPositionFilter<()>>, sampling_schedule: &SamplingSchedule, + indexed_fraction: Option<(f64, u64)>, only_mapped: bool, suppress_progress: bool, ) -> anyhow::Result @@ -189,7 +202,6 @@ where .iter() .map(|rec| (rec.tid, rec.length)) .collect::>(); - let feeder = ReferenceIntervalBatchesFeeder::new( contigs, batch_size, @@ -215,17 +227,64 @@ where let mut aggregator = ::zero(); let mut reads_sampled_per_chr = FxHashMap::default(); + // Indexed fetches return every alignment overlapping an interval, including + // alignments that start before it. Track the preceding interval that was + // actually retained by an include-BED so records which begin in a skipped + // gap remain owned by the next retained interval. + let mut retained_interval_ends = FxHashMap::default(); let feeder = feeder.map(|x| x.unwrap()); for super_batch in feeder { let total_batch_length = super_batch.iter().map(|c| c.total_length()).sum::(); - let super_batch_with_counts = sampling_schedule - .accumulate_sample_counts( - super_batch, - &contig_sizes, - &reads_sampled_per_chr, - batch_size, - ); + let super_batch_with_counts: Vec< + Vec<(ChromCoordinates, CountOrSample, Option)>, + > = if indexed_fraction.is_some() { + // Fractional sampling is decided independently for every record, + // so every requested interval must be visited. The legacy count + // schedule can stop assigning intervals after its target count is + // reached, which would make the result interval-size dependent. + super_batch + .into_iter() + .flat_map(|coordinates| coordinates.0) + .filter_map(|coordinates| { + let retained = position_filter + .map(|pf| { + pf.overlaps_not_stranded( + coordinates.chrom_tid, + coordinates.start_pos as u64, + coordinates.end_pos as u64, + ) + }) + .unwrap_or(true); + retained.then(|| { + let record_start_cut = retained_interval_ends + .insert(coordinates.chrom_tid, coordinates.end_pos); + (coordinates, CountOrSample::All, record_start_cut) + }) + }) + .chunks(batch_size) + .into_iter() + .map(|batch| batch.collect()) + .collect() + } else { + sampling_schedule + .accumulate_sample_counts( + super_batch, + &contig_sizes, + &reads_sampled_per_chr, + batch_size, + ) + .into_iter() + .map(|batch| { + batch + .into_iter() + .map(|(coordinates, count_or_sample)| { + (coordinates, count_or_sample, None) + }) + .collect() + }) + .collect() + }; let (super_batch_result, chrom_counts_for_batch) = super_batch_with_counts .into_par_iter() @@ -234,6 +293,7 @@ where bam_fp, multi_coords, sampling_schedule, + indexed_fraction, collapse_method, edge_filter, position_filter, @@ -258,8 +318,9 @@ where fn run_batch( bam_fp: &PathBuf, - batch: Vec<(ChromCoordinates, CountOrSample)>, + batch: Vec<(ChromCoordinates, CountOrSample, Option)>, sampling_schedule: &SamplingSchedule, + indexed_fraction: Option<(f64, u64)>, collapse_method: Option<&CollapseMethod>, edge_filter: Option<&EdgeFilter>, position_filter: Option<&StrandedPositionFilter<()>>, @@ -273,8 +334,8 @@ where { batch .into_par_iter() - .filter(|(cc, _)| sampling_schedule.chrom_has_reads(cc.chrom_tid)) - .filter(|(cc, _)| { + .filter(|(cc, _, _)| sampling_schedule.chrom_has_reads(cc.chrom_tid)) + .filter(|(cc, _, _)| { position_filter .map(|pf| { pf.overlaps_not_stranded( @@ -285,21 +346,26 @@ where }) .unwrap_or(true) }) - .filter_map(|(cc, counts_or_sample)| { - let record_sampler = match counts_or_sample { - CountOrSample::Count(x) => RecordSampler::new_num_reads(x), - CountOrSample::Sample(x) => { - RecordSampler::new_sample_frac(x as f64, None) + .filter_map(|(cc, counts_or_sample, record_start_cut)| { + let record_sampler = if let Some((frac, master_seed)) = + indexed_fraction + { + RecordSampler::new_deterministic_sample_frac(frac, master_seed) + } else { + match counts_or_sample { + CountOrSample::Count(x) => RecordSampler::new_num_reads(x), + CountOrSample::Sample(x) => { + RecordSampler::new_sample_frac(x as f64, None) + } + CountOrSample::All => RecordSampler::new_passthrough(), } - CountOrSample::All => RecordSampler::new_passthrough(), }; - match sample_reads_from_interval::

( bam_fp, cc.chrom_tid, cc.start_pos, cc.end_pos, - None, + record_start_cut, record_sampler, collapse_method, edge_filter, @@ -389,3 +455,290 @@ fn log_sampled_reads(sampled_reads_per_chr: &FxHashMap) { tab.add_row(row!["total", total]); debug!("final mapped reads sampled:\n{tab}"); } + +#[cfg(test)] +mod tests { + use std::{collections::HashSet, fs, path::PathBuf}; + + use rust_htslib::bam::{ + self, + header::HeaderRecord, + record::{Aux, Cigar, CigarString}, + Read, + }; + + use super::{ + deterministic_sampler::DeterministicFractionSampler, + get_sampled_read_ids_to_base_mod_probs, SamplingSchedule, + }; + use crate::{ + mod_base_code::DnaBase, position_filter::StrandedPositionFilter, + read_ids_to_base_mod_probs::ReadIdsToBaseModProbs, + thresholds::calc_threshold_from_bam, + }; + + fn mapped_mod_record_with_calls( + name: &str, + tid: i32, + start: i64, + length: usize, + mod_offsets: &[usize], + probs: &[u8], + ) -> bam::Record { + assert_eq!(mod_offsets.len(), probs.len()); + assert!(mod_offsets.iter().all(|offset| *offset < length)); + + let mut record = bam::Record::new(); + let cigar = CigarString(vec![Cigar::Match(length as u32)]); + let sequence = vec![b'C'; length]; + let qualities = vec![30; length]; + record.set(name.as_bytes(), Some(&cigar), &sequence, &qualities); + record.set_tid(tid); + record.set_pos(start); + record.set_flags(0); + record.set_mapq(60); + + let mut previous_offset = None; + let deltas = mod_offsets + .iter() + .map(|offset| { + let delta = previous_offset + .map(|previous| offset - previous - 1) + .unwrap_or(*offset); + previous_offset = Some(*offset); + delta.to_string() + }) + .collect::>() + .join(","); + let mm_tag = format!("C+m?,{deltas};"); + record.push_aux(b"MM", Aux::String(&mm_tag)).unwrap(); + record.push_aux(b"ML", Aux::ArrayU8(probs.into())).unwrap(); + record.push_aux(b"MN", Aux::U32(length as u32)).unwrap(); + record + } + + fn mapped_mod_record(name: &str, tid: i32) -> bam::Record { + mapped_mod_record_with_calls(name, tid, 10, 10, &[0], &[200]) + } + + fn sparse_position_fixture() -> ( + tempfile::TempDir, + PathBuf, + StrandedPositionFilter<()>, + Vec, + ) { + let temp_dir = tempfile::tempdir().unwrap(); + let bam_path = temp_dir.path().join("sparse-position-filter.bam"); + let bed_path = temp_dir.path().join("sparse-positions.bed"); + + let mut header = bam::Header::new(); + let mut sq = HeaderRecord::new(b"SQ"); + sq.push_tag(b"SN", "chr1").push_tag(b"LN", 160); + header.push_record(&sq); + + // Retained 20 bp chunks are [20, 40) and [100, 120). The first read + // spans both, while the second starts in the skipped gap and must be + // owned by the later retained chunk. + let records = vec![ + mapped_mod_record_with_calls( + "spans-retained-chunks", + 0, + 5, + 120, + &[20, 100], + &[51, 204], + ), + mapped_mod_record_with_calls( + "starts-in-skipped-gap", + 0, + 60, + 60, + &[45], + &[153], + ), + mapped_mod_record_with_calls( + "starts-in-retained-chunk", + 0, + 100, + 10, + &[5], + &[102], + ), + ]; + let mut writer = + bam::Writer::from_path(&bam_path, &header, bam::Format::Bam) + .unwrap(); + for record in &records { + writer.write(record).unwrap(); + } + drop(writer); + bam::index::build(&bam_path, None, bam::index::Type::Bai, 1).unwrap(); + + fs::write(&bed_path, "chr1\t25\t26\nchr1\t105\t106\n").unwrap(); + let position_filter = StrandedPositionFilter::from_bam_and_bed( + &bam_path, &bed_path, true, + ) + .unwrap(); + + (temp_dir, bam_path, position_filter, records) + } + + fn assert_sparse_fraction_is_interval_invariant( + fraction: f64, + seed: u64, + expected_names: &[&str], + ) { + let (_temp_dir, bam_path, position_filter, _) = + sparse_position_fixture(); + let sample = |threads, interval_size| { + get_sampled_read_ids_to_base_mod_probs::( + &bam_path, + threads, + interval_size, + Some(fraction), + None, + Some(seed), + None, + None, + None, + Some(&position_filter), + true, + true, + ) + .unwrap() + }; + + let small_intervals = sample(2, 20); + let whole_contig = sample(4, 160); + assert_eq!(small_intervals.inner, whole_contig.inner); + + let sampled_names = + small_intervals.inner.keys().cloned().collect::>(); + let expected_names = expected_names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + assert_eq!(sampled_names, expected_names); + + let threshold = |threads, interval_size| { + calc_threshold_from_bam( + &bam_path, + threads, + interval_size, + Some(fraction), + None, + 0.25, + Some(seed), + None, + None, + None, + Some(&position_filter), + true, + true, + ) + .unwrap() + }; + let small_threshold = threshold(2, 20); + let whole_contig_threshold = threshold(4, 160); + assert_eq!(small_threshold, whole_contig_threshold); + assert!(small_threshold.contains_key(&DnaBase::C)); + } + + #[test] + fn indexed_fraction_visits_later_single_read_contig() { + const FRACTION: f64 = 0.01; + + let temp_dir = tempfile::tempdir().unwrap(); + let bam_path = temp_dir.path().join("two-contig.bam"); + let mut header = bam::Header::new(); + for name in ["tiny-0", "tiny-1"] { + let mut sq = HeaderRecord::new(b"SQ"); + sq.push_tag(b"SN", name).push_tag(b"LN", 100); + header.push_record(&sq); + } + let mut writer = + bam::Writer::from_path(&bam_path, &header, bam::Format::Bam) + .unwrap(); + writer.write(&mapped_mod_record("read-tiny-0", 0)).unwrap(); + writer.write(&mapped_mod_record("read-tiny-1", 1)).unwrap(); + drop(writer); + bam::index::build(&bam_path, None, bam::index::Type::Bai, 1).unwrap(); + + let schedule = SamplingSchedule::from_sample_frac( + &bam_path, + FRACTION as f32, + None, + None, + false, + ) + .unwrap(); + assert!(schedule.chrom_has_reads(0)); + assert!(schedule.chrom_has_reads(1)); + + let mut reader = bam::Reader::from_path(&bam_path).unwrap(); + let records = + reader.records().map(Result::unwrap).collect::>(); + let seed = (0_u64..100_000) + .find(|seed| { + let sampler = + DeterministicFractionSampler::new(*seed, FRACTION).unwrap(); + !sampler.include(&records[0]) && sampler.include(&records[1]) + }) + .expect("expected a seed selecting only the later tiny contig"); + + let sampled = + get_sampled_read_ids_to_base_mod_probs::( + &bam_path, + 2, + 20, + Some(FRACTION), + None, + Some(seed), + None, + None, + None, + None, + true, + true, + ) + .unwrap(); + + assert_eq!(sampled.inner.len(), 1); + assert!(sampled.inner.contains_key("read-tiny-1")); + } + + #[test] + fn sparse_position_filter_fraction_one_is_interval_invariant() { + assert_sparse_fraction_is_interval_invariant( + 1.0, + 7, + &[ + "spans-retained-chunks", + "starts-in-skipped-gap", + "starts-in-retained-chunk", + ], + ); + } + + #[test] + fn sparse_position_filter_seeded_fraction_is_interval_invariant() { + const FRACTION: f64 = 0.5; + let (_temp_dir, _bam_path, _position_filter, records) = + sparse_position_fixture(); + let seed = (0_u64..100_000) + .find(|seed| { + let sampler = + DeterministicFractionSampler::new(*seed, FRACTION).unwrap(); + sampler.include(&records[0]) + && sampler.include(&records[1]) + && !sampler.include(&records[2]) + }) + .expect("expected a seed selecting the spanning and gap reads"); + + assert_sparse_fraction_is_interval_invariant( + FRACTION, + seed, + &["spans-retained-chunks", "starts-in-skipped-gap"], + ); + } +} diff --git a/modkit-core/src/reads_sampler/record_sampler.rs b/modkit-core/src/reads_sampler/record_sampler.rs index 7b125906..4dd4202b 100644 --- a/modkit-core/src/reads_sampler/record_sampler.rs +++ b/modkit-core/src/reads_sampler/record_sampler.rs @@ -4,6 +4,9 @@ use indicatif::ProgressBar; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use rust_htslib::bam; + +use super::deterministic_sampler::DeterministicFractionSampler; /// A utility data structure that when used in an interator allows /// to randomly sample either a preset number of reads or a fraction @@ -12,6 +15,7 @@ use rand::{Rng, SeedableRng}; pub struct RecordSampler { pub(crate) num_reads: Option, pub(crate) sample_frac: Option, + deterministic_fraction_sampler: Option, rng: StdRng, reads_sampled: usize, } @@ -21,6 +25,7 @@ impl RecordSampler { Self { num_reads: Some(num_reads), sample_frac: None, + deterministic_fraction_sampler: None, rng: StdRng::from_entropy(), reads_sampled: 0, } @@ -33,15 +38,36 @@ impl RecordSampler { Self { num_reads: None, sample_frac: Some(sample_frac), + deterministic_fraction_sampler: None, rng, reads_sampled: 0, } } + /// Construct a stateless fractional sampler for indexed record fetching. + /// Every worker in the indexed job must receive the same resolved seed. + pub(crate) fn new_deterministic_sample_frac( + sample_frac: f64, + master_seed: u64, + ) -> Self { + let deterministic_fraction_sampler = Some( + DeterministicFractionSampler::new(master_seed, sample_frac) + .expect("sampling fraction should already be validated"), + ); + Self { + num_reads: None, + sample_frac: Some(sample_frac), + deterministic_fraction_sampler, + rng: StdRng::seed_from_u64(master_seed), + reads_sampled: 0, + } + } + pub(crate) fn new_passthrough() -> Self { Self { num_reads: None, sample_frac: None, + deterministic_fraction_sampler: None, rng: StdRng::from_entropy(), reads_sampled: 0, } @@ -86,6 +112,10 @@ impl RecordSampler { } pub(crate) fn ask(&mut self) -> Indicator { + debug_assert!( + self.deterministic_fraction_sampler.is_none(), + "deterministic sampling decisions require a BAM record" + ); match (self.num_reads, self.sample_frac) { (Some(_nr), _) => self.check_num_reads(), (_, Some(_sample_frac)) => self.check_sample_frac(), @@ -93,6 +123,18 @@ impl RecordSampler { } } + pub(crate) fn ask_record(&mut self, record: &bam::Record) -> Indicator { + if let Some(sampler) = self.deterministic_fraction_sampler { + if sampler.include(record) { + Indicator::Use(Token) + } else { + Indicator::Skip + } + } else { + self.ask() + } + } + pub(crate) fn used(&mut self, _token: Token) { self.reads_sampled += 1; } diff --git a/modkit-core/src/sample_probs/mod.rs b/modkit-core/src/sample_probs/mod.rs index a0080af9..d0e2bab3 100644 --- a/modkit-core/src/sample_probs/mod.rs +++ b/modkit-core/src/sample_probs/mod.rs @@ -32,6 +32,9 @@ use crate::{ AlignedBaseModsIterator, BaseModsAdapter, ModState, }, position_filter::StrandedPositionFilter, + reads_sampler::deterministic_sampler::{ + resolve_master_seed, DeterministicFractionSampler, + }, util::{ get_human_readable_table, get_master_progress_bar, get_targets, get_ticker, record_is_not_primary, CheckedAddArr, Region, Strand, @@ -675,7 +678,9 @@ pub(crate) trait ExtractsMleProbs { sample_frac: f64, motif_bases: [DnaBase; 4], edge_filter: Option<&EdgeFilter>, - ) -> Self; + ) -> anyhow::Result + where + Self: Sized; fn process_record( &mut self, record: &bam::Record, @@ -725,6 +730,8 @@ impl> RegionMleProbs { sample_frac: f64, chrom_to_counts: Option>>, ) -> anyhow::Result { + let hist = + H::new(worker_no, sample, sample_frac, motif_bases, edge_filter)?; let mut reader = bam::IndexedReader::from_path(bam_fp)?; reader.set_thread_pool(&thread_pool)?; if reader_is_cram(&reader) { @@ -732,7 +739,8 @@ impl> RegionMleProbs { reader.set_reference(reference_fp)?; reader.set_cram_options( htslib::hts_fmt_option_CRAM_OPT_REQUIRED_FIELDS, - htslib::sam_fields_SAM_FLAG + htslib::sam_fields_SAM_QNAME + | htslib::sam_fields_SAM_FLAG | htslib::sam_fields_SAM_RNAME | htslib::sam_fields_SAM_POS | htslib::sam_fields_SAM_MAPQ @@ -744,8 +752,6 @@ impl> RegionMleProbs { bail!("CRAM input requires reference") } }; - let hist = - H::new(worker_no, sample, sample_frac, motif_bases, edge_filter); Ok(Self { reader, allow_non_primary, @@ -838,14 +844,44 @@ impl + Send> ExtractProbsWorker pub(crate) struct ProbsExtractor { rng: SmallRng, - sample: bool, - sample_frac: f64, + sampler: Option, motif_bases: [DnaBase; 4], edge_filter_start: usize, edge_filter_end: usize, } impl ProbsExtractor { + fn new( + seed: u64, + sample: bool, + sample_frac: f64, + motif_bases: [DnaBase; 4], + edge_filter: Option<&EdgeFilter>, + ) -> anyhow::Result { + let sampler = sample + .then(|| DeterministicFractionSampler::new(seed, sample_frac)) + .transpose()?; + let edge_filter_start = + edge_filter.map(|ef| ef.edge_filter_start).unwrap_or(0usize); + let edge_filter_end = + edge_filter.map(|ef| ef.edge_filter_end).unwrap_or(0usize); + Ok(Self { + rng: SmallRng::seed_from_u64(seed), + sampler, + motif_bases, + edge_filter_start, + edge_filter_end, + }) + } + + #[inline] + fn include_record(&self, record: &bam::Record) -> bool { + self.sampler + .as_ref() + .map(|sampler| sampler.include(record)) + .unwrap_or(true) + } + fn get_aligned_mod_state_iterator<'a>( record: &'a bam::Record, chrom_coords: &'a ChromCoordinates, @@ -1060,16 +1096,8 @@ impl ExtractsMleProbs for ProbsExtractor { sample_frac: f64, motif_bases: [DnaBase; 4], _edge_filter: Option<&EdgeFilter>, - ) -> Self { - let rng = SmallRng::seed_from_u64(seed); - Self { - rng, - sample, - sample_frac, - motif_bases, - edge_filter_start: 0usize, - edge_filter_end: 0usize, - } + ) -> anyhow::Result { + Self::new(seed, sample, sample_frac, motif_bases, None) } #[inline] @@ -1083,10 +1111,8 @@ impl ExtractsMleProbs for ProbsExtractor { _mods_hists: &mut Vec, records_with_base_mods: &mut [u32; 4], ) -> anyhow::Result { - if self.sample { - if !self.rng.gen_bool(self.sample_frac) { - return Ok(false); - } + if !self.include_record(record) { + return Ok(false); } if !chrom_coords.final_interval { @@ -1127,16 +1153,8 @@ impl ExtractsMleProbs for ProbsExtractor { sample_frac: f64, motif_bases: [DnaBase; 4], _edge_filter: Option<&EdgeFilter>, - ) -> Self { - let rng = SmallRng::seed_from_u64(seed); - Self { - rng, - sample, - sample_frac, - motif_bases, - edge_filter_start: 0usize, - edge_filter_end: 0usize, - } + ) -> anyhow::Result { + Self::new(seed, sample, sample_frac, motif_bases, None) } #[inline] @@ -1150,10 +1168,8 @@ impl ExtractsMleProbs for ProbsExtractor { mods_hists: &mut Vec, records_with_base_mods: &mut [u32; 4], ) -> anyhow::Result { - if self.sample { - if !self.rng.gen_bool(self.sample_frac) { - return Ok(false); - } + if !self.include_record(record) { + return Ok(false); } if !chrom_coords.final_interval { let aln_end = record.reference_end(); @@ -1194,20 +1210,8 @@ impl ExtractsMleProbs for ProbsExtractor { sample_frac: f64, motif_bases: [DnaBase; 4], edge_filter: Option<&EdgeFilter>, - ) -> Self { - let rng = SmallRng::seed_from_u64(seed); - let edge_filter_start = - edge_filter.map(|ef| ef.edge_filter_start).unwrap_or(0usize); - let edge_filter_end = - edge_filter.map(|ef| ef.edge_filter_end).unwrap_or(0usize); - Self { - rng, - sample, - sample_frac, - motif_bases, - edge_filter_start, - edge_filter_end, - } + ) -> anyhow::Result { + Self::new(seed, sample, sample_frac, motif_bases, edge_filter) } #[inline] @@ -1221,10 +1225,8 @@ impl ExtractsMleProbs for ProbsExtractor { mods_hists: &mut Vec, records_with_base_mods: &mut [u32; 4], ) -> anyhow::Result { - if self.sample { - if !self.rng.gen_bool(self.sample_frac) { - return Ok(false); - } + if !self.include_record(record) { + return Ok(false); } let start_pos = chrom_coords.start_pos; let end_pos = chrom_coords.end_pos; @@ -1279,20 +1281,8 @@ impl ExtractsMleProbs for ProbsExtractor { sample_frac: f64, motif_bases: [DnaBase; 4], edge_filter: Option<&EdgeFilter>, - ) -> Self { - let rng = SmallRng::seed_from_u64(seed); - let edge_filter_start = - edge_filter.map(|ef| ef.edge_filter_start).unwrap_or(0usize); - let edge_filter_end = - edge_filter.map(|ef| ef.edge_filter_end).unwrap_or(0usize); - Self { - rng, - sample, - sample_frac, - motif_bases, - edge_filter_start, - edge_filter_end, - } + ) -> anyhow::Result { + Self::new(seed, sample, sample_frac, motif_bases, edge_filter) } #[inline] @@ -1306,10 +1296,8 @@ impl ExtractsMleProbs for ProbsExtractor { _mods_hists: &mut Vec, records_with_base_mods: &mut [u32; 4], ) -> anyhow::Result { - if self.sample { - if !self.rng.gen_bool(self.sample_frac) { - return Ok(false); - } + if !self.include_record(record) { + return Ok(false); } let should_count_record = if !chrom_coords.final_interval { let aln_end = record.reference_end(); @@ -1556,6 +1544,11 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( interval_size, sampling_region, )?; + let master_seed = if rng_sample { + resolve_master_seed(seed) + } else { + seed.unwrap_or_default() + }; let chrom_to_counts = chrom_to_counts.map(|x| Arc::new(x)); let motif_lookup = get_motif_lookup_from_parts( motifs, @@ -1572,7 +1565,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( stranded_position_filter.clone(), )?; if let Some(motif_bases) = feeder.get_motif_bases() { - for i in 0..n_workers { + for _ in 0..n_workers { let worker: Box = if collect_mod_histograms { Box::new(RegionMleProbs::< @@ -1585,7 +1578,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( motif_bases, edge_filter, io_threadpool, - seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1601,7 +1594,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( motif_bases, edge_filter, io_threadpool, - seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1610,7 +1603,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( workers.push(worker); } } else if feeder.has_position_filter() { - for i in 0..n_workers { + for _ in 0..n_workers { let worker: Box = if collect_mod_histograms { Box::new(RegionMleProbs::< @@ -1623,7 +1616,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( [DnaBase::A; 4], edge_filter, io_threadpool, - seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1639,7 +1632,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( [DnaBase::A; 4], edge_filter, io_threadpool, - seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1648,7 +1641,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( workers.push(worker); } } else { - for i in 0..n_workers { + for _ in 0..n_workers { let worker: Box = if collect_mod_histograms { Box::new(RegionMleProbs::::new( @@ -1658,7 +1651,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( [DnaBase::A; 4], edge_filter, io_threadpool, - seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1672,7 +1665,7 @@ pub(crate) fn get_base_mods_quals_from_indexed_hts_file( [DnaBase::A; 4], edge_filter, io_threadpool, - seed.unwrap_or(i as u64), + master_seed, rng_sample, sample_frac, chrom_to_counts.clone(), @@ -1759,3 +1752,6 @@ fn byte_to_bool_positions(b: u8, agg: &mut [u32; SIZE]) { agg[i] = count.saturating_add(1u32); } } + +#[cfg(test)] +mod tests; diff --git a/modkit-core/src/sample_probs/tests.rs b/modkit-core/src/sample_probs/tests.rs new file mode 100644 index 00000000..15d2139b --- /dev/null +++ b/modkit-core/src/sample_probs/tests.rs @@ -0,0 +1,197 @@ +use std::collections::BTreeMap; + +use bitvec::{bitvec, order::Lsb0}; +use rust_htslib::bam::{ + self, + record::{Aux, Cigar, CigarString}, +}; + +use super::{ + AlignedBaseAndModArgmaxProbs, AlignedBaseArgmaxProbs, + BaseAndModArgmaxProbs, BaseArgmaxProbs, ExtractsMleProbs, ProbsExtractor, + QualHist, +}; +use crate::{ + interval_chunks::{ChromCoordinates, FocusPositions2}, + mod_base_code::DnaBase, + reads_sampler::deterministic_sampler::DeterministicFractionSampler, +}; + +fn make_record(name: &str, pos: i64) -> bam::Record { + let mut record = bam::Record::new(); + let cigar = CigarString(vec![Cigar::Match(1)]); + record.set(name.as_bytes(), Some(&cigar), b"C", &[255]); + record.set_tid(0); + record.set_pos(pos); + record.push_aux(b"MM", Aux::String("C+m?,0;")).unwrap(); + record.push_aux(b"ML", Aux::ArrayU8((&[128][..]).into())).unwrap(); + record +} + +fn make_long_record(name: &str) -> bam::Record { + const READ_LENGTH: usize = 12; + let mut record = bam::Record::new(); + let cigar = CigarString(vec![Cigar::Match(READ_LENGTH as u32)]); + record.set( + name.as_bytes(), + Some(&cigar), + &vec![b'C'; READ_LENGTH], + &vec![255; READ_LENGTH], + ); + record.set_tid(0); + record.set_pos(0); + let mm = format!("C+m?,{};", vec!["0"; READ_LENGTH].join(",")); + record.push_aux(b"MM", Aux::String(&mm)).unwrap(); + record + .push_aux(b"ML", Aux::ArrayU8((&vec![128; READ_LENGTH]).into())) + .unwrap(); + record +} + +fn process_record( + extractor: &mut ProbsExtractor, + record: &bam::Record, + interval_start: u32, +) -> bool +where + ProbsExtractor: ExtractsMleProbs, +{ + let mut mask = bitvec![usize, Lsb0; 0; 2]; + mask.set(0, true); + let coords = ChromCoordinates::new( + 0, + interval_start, + interval_start + 1, + FocusPositions2::MaskedPositions { mask }, + true, + ); + let mut hist = QualHist::default(); + >::process_record( + extractor, + record, + &coords, + &mut hist.explicit_canonical_probs, + &mut hist.hist, + &mut hist.base_totals, + &mut hist.mods_hists, + &mut hist.num_records_with_base_mods, + ) + .unwrap() +} + +fn sampling_decisions( + records: &[bam::Record], + assignments: &[usize], + worker_count: usize, +) -> BTreeMap, bool> +where + ProbsExtractor: ExtractsMleProbs, +{ + let mut workers = (0..worker_count) + .map(|_| { + >::new( + 7, + true, + 0.5, + [DnaBase::A; 4], + None, + ) + .unwrap() + }) + .collect::>(); + + records + .iter() + .zip(assignments) + .map(|(record, worker)| { + ( + record.qname().to_vec(), + process_record::( + &mut workers[*worker], + record, + record.pos() as u32, + ), + ) + }) + .collect() +} + +#[test] +fn seeded_fraction_sampling_is_independent_of_worker_assignment() { + let records = (0..12) + .map(|i| make_record(&format!("read-{i}"), i)) + .collect::>(); + fn assert_stable(records: &[bam::Record]) + where + ProbsExtractor: ExtractsMleProbs, + { + let serial = sampling_decisions::(records, &[0; 12], 1); + let scheduled = sampling_decisions::( + records, + &[0, 1, 2, 0, 2, 1, 1, 0, 2, 2, 0, 1], + 3, + ); + assert_eq!(serial, scheduled); + } + + assert_stable::(&records); + assert_stable::(&records); + assert_stable::(&records); + assert_stable::(&records); +} + +#[test] +fn seeded_fraction_sampling_is_consistent_across_interval_refetches() { + let sampler = DeterministicFractionSampler::new(7, 0.5).unwrap(); + for should_include in [false, true] { + let record = (0..100) + .map(|i| make_long_record(&format!("long-read-{i}"))) + .find(|record| sampler.include(record) == should_include) + .unwrap(); + let mut extractor = >::new(7, true, 0.5, [DnaBase::A; 4], None) + .unwrap(); + let mut hist = QualHist::default(); + let mut counted = Vec::new(); + + for (start, final_interval) in [(0, false), (4, false), (8, true)] { + let mut mask = bitvec![usize, Lsb0; 0; 8]; + for position in (0..8).step_by(2) { + mask.set(position, true); + } + let coords = ChromCoordinates::new( + 0, + start, + start + 4, + FocusPositions2::MaskedPositions { mask }, + final_interval, + ); + counted.push( + >::process_record( + &mut extractor, + &record, + &coords, + &mut hist.explicit_canonical_probs, + &mut hist.hist, + &mut hist.base_totals, + &mut hist.mods_hists, + &mut hist.num_records_with_base_mods, + ) + .unwrap(), + ); + } + + if should_include { + assert_eq!(counted, [false, false, true]); + assert_eq!(hist.base_totals[DnaBase::C as usize], 12); + assert_eq!(hist.num_records_with_base_mods[DnaBase::C as usize], 1); + } else { + assert_eq!(counted, [false, false, false]); + assert_eq!(hist.base_totals, [0; 4]); + assert_eq!(hist.num_records_with_base_mods, [0; 4]); + } + } +} diff --git a/modkit/tests/test_legacy_threshold_sampling_determinism.rs b/modkit/tests/test_legacy_threshold_sampling_determinism.rs new file mode 100644 index 00000000..b927110f --- /dev/null +++ b/modkit/tests/test_legacy_threshold_sampling_determinism.rs @@ -0,0 +1,249 @@ +use std::{fs, path::Path, process::Command}; + +fn sam_record_body(path: &Path) -> Vec { + let sam = String::from_utf8(fs::read(path).unwrap()).unwrap(); + sam.lines() + .filter(|line| !line.starts_with('@')) + .collect::>() + .join("\n") + .into_bytes() +} + +fn run_call_mods( + out_dir: &Path, + name: &str, + input: &str, + seed: u64, + threads: usize, + sampling_interval_size: u32, +) -> Vec { + let output_path = out_dir.join(format!("{name}.sam")); + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args([ + "call-mods", + input, + output_path.to_str().unwrap(), + "--output-sam", + "--sampling-frac", + "0.5", + "--seed", + &seed.to_string(), + "--threads", + &threads.to_string(), + "--sampling-interval-size", + &sampling_interval_size.to_string(), + "--filter-percentile", + "0.25", + "--suppress-progress", + ]) + .output() + .expect("failed to run modkit call-mods"); + assert!( + output.status.success(), + "call-mods failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + sam_record_body(&output_path) +} + +fn run_pileup_hemi( + out_dir: &Path, + name: &str, + threads: usize, + sampling_interval_size: u32, +) -> (Vec, String) { + let output_path = out_dir.join(format!("{name}.bed")); + let log_path = out_dir.join(format!("{name}.log")); + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args([ + "pileup-hemi", + "../tests/resources/duplex_modcalls_sort.bam", + "--out-bed", + output_path.to_str().unwrap(), + "--ref", + "../tests/resources/GRCh38_chr20.fa", + "--cpg", + "--region", + "chr20:22,613,835-22,640,468", + "--sampling-frac", + "0.5", + "--seed", + "7", + "--threads", + &threads.to_string(), + "--sampling-interval-size", + &sampling_interval_size.to_string(), + "--suppress-progress", + "--log-filepath", + log_path.to_str().unwrap(), + ]) + .output() + .expect("failed to run modkit pileup-hemi"); + assert!( + output.status.success(), + "pileup-hemi failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + + let log = fs::read_to_string(log_path).unwrap(); + let threshold = log + .lines() + .find_map(|line| { + line.find("estimated pass threshold") + .map(|start| line[start..].to_string()) + }) + .expect("pileup-hemi log did not contain an estimated threshold"); + (fs::read(output_path).unwrap(), threshold) +} + +fn run_sparse_pileup_hemi( + out_dir: &Path, + name: &str, + sampling_frac: &str, + threads: usize, + sampling_interval_size: u32, +) -> (Vec, String) { + let output_path = out_dir.join(format!("{name}.bed")); + let log_path = out_dir.join(format!("{name}.log")); + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args([ + "pileup-hemi", + "../tests/resources/bc_anchored_10_reads.sorted.bam", + "--out-bed", + output_path.to_str().unwrap(), + "--ref", + "../tests/resources/CGI_ladder_3.6kb_ref.fa", + "--cpg", + "--region", + "oligo_1512_adapters", + "--include-bed", + "../tests/resources/include-pos-1-site.bed", + "--sampling-frac", + sampling_frac, + "--seed", + "7", + "--threads", + &threads.to_string(), + "--sampling-interval-size", + &sampling_interval_size.to_string(), + "--suppress-progress", + "--log-filepath", + log_path.to_str().unwrap(), + ]) + .output() + .expect("failed to run sparse include-BED pileup-hemi"); + assert!( + output.status.success(), + "sparse include-BED pileup-hemi failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + + let log = fs::read_to_string(log_path).unwrap(); + let threshold = log + .lines() + .find_map(|line| { + line.find("estimated pass threshold") + .map(|start| line[start..].to_string()) + }) + .expect( + "sparse include-BED log did not contain an estimated threshold", + ); + (fs::read(output_path).unwrap(), threshold) +} + +#[test] +fn call_mods_seeded_fraction_is_stable_across_workers_and_intervals() { + let temp_dir = tempfile::tempdir().unwrap(); + let input = "../tests/resources/bc_anchored_10_reads.sorted.bam"; + let baseline = run_call_mods(temp_dir.path(), "baseline", input, 7, 1, 20); + + for (name, threads, interval_size) in [ + ("repeat", 1, 20), + ("more_workers", 8, 20), + ("larger_intervals", 8, 1_000), + ] { + assert_eq!( + run_call_mods( + temp_dir.path(), + name, + input, + 7, + threads, + interval_size + ), + baseline, + "call-mods output changed for {threads} workers and interval size \ + {interval_size}" + ); + } + + assert_ne!( + run_call_mods(temp_dir.path(), "different_seed", input, 8, 3, 20), + baseline, + "call-mods ignored the explicit fractional-sampling seed" + ); +} + +#[test] +fn call_mods_unindexed_seeded_streaming_behavior_remains_repeatable() { + let temp_dir = tempfile::tempdir().unwrap(); + let input = temp_dir.path().join("unindexed.bam"); + fs::copy("../tests/resources/bc_anchored_10_reads.sorted.bam", &input) + .unwrap(); + let input = input.to_str().unwrap(); + let baseline = + run_call_mods(temp_dir.path(), "unindexed_a", input, 7, 1, 20); + + assert_eq!( + run_call_mods(temp_dir.path(), "unindexed_b", input, 7, 8, 1_000), + baseline + ); + assert_ne!( + run_call_mods(temp_dir.path(), "unindexed_seed8", input, 8, 3, 20), + baseline + ); +} + +#[test] +fn duplex_threshold_sampling_is_stable_across_workers_and_intervals() { + let temp_dir = tempfile::tempdir().unwrap(); + let baseline = run_pileup_hemi(temp_dir.path(), "baseline", 1, 1_000); + + for (name, threads, interval_size) in + [("repeat", 1, 1_000), ("matrix", 8, 100_000)] + { + assert_eq!( + run_pileup_hemi(temp_dir.path(), name, threads, interval_size), + baseline, + "pileup-hemi threshold or output changed for {threads} workers \ + and interval size {interval_size}" + ); + } +} + +#[test] +fn duplex_sparse_include_bed_sampling_is_interval_invariant() { + let temp_dir = tempfile::tempdir().unwrap(); + + for sampling_frac in ["1", "0.5"] { + let small_intervals = run_sparse_pileup_hemi( + temp_dir.path(), + &format!("sparse_{sampling_frac}_small"), + sampling_frac, + 1, + 20, + ); + let large_intervals = run_sparse_pileup_hemi( + temp_dir.path(), + &format!("sparse_{sampling_frac}_large"), + sampling_frac, + 4, + 1_000, + ); + assert_eq!( + small_intervals, large_intervals, + "sparse include-BED output or sampled threshold changed at \ + fraction {sampling_frac}" + ); + } +} diff --git a/modkit/tests/test_sample_probs_determinism.rs b/modkit/tests/test_sample_probs_determinism.rs new file mode 100644 index 00000000..04f95a1d --- /dev/null +++ b/modkit/tests/test_sample_probs_determinism.rs @@ -0,0 +1,331 @@ +use std::{collections::BTreeMap, fs, path::Path, process::Command}; + +use rust_htslib::bam::{self, ext::BamRecordExtensions, Read}; + +fn run_sample_probs( + out_dir: &Path, + prefix: &str, + input: &str, + reference: Option<&str>, + fraction: Option<&str>, + seed: u64, + threads: usize, + interval_size: u32, +) -> (Vec, Vec) { + let mut args = vec![ + "sample-probs".to_string(), + input.to_string(), + "--threads".to_string(), + threads.to_string(), + "--interval-size".to_string(), + interval_size.to_string(), + "--hist".to_string(), + "--out-dir".to_string(), + out_dir.to_str().unwrap().to_string(), + "--prefix".to_string(), + prefix.to_string(), + "--suppress-progress".to_string(), + ]; + if let Some(reference) = reference { + args.extend(["--reference".to_string(), reference.to_string()]); + } + if let Some(fraction) = fraction { + args.extend([ + "--sample-frac".to_string(), + fraction.to_string(), + "--seed".to_string(), + seed.to_string(), + ]); + } else { + args.push("--no-sampling".to_string()); + } + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args(args) + .output() + .expect("failed to run modkit sample-probs"); + assert!( + output.status.success(), + "sample-probs failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + + let thresholds = + fs::read(out_dir.join(format!("{prefix}_thresholds.tsv"))).unwrap(); + let probabilities = + fs::read(out_dir.join(format!("{prefix}_probabilities.tsv"))).unwrap(); + (thresholds, probabilities) +} + +fn run_summary(threads: usize, interval_size: u32) -> BTreeMap { + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args([ + "summary", + "../tests/resources/bc_anchored_10_reads.sorted.bam", + "--sampling-frac", + "0.5", + "--seed", + "7", + "--threads", + &threads.to_string(), + "--interval-size", + &interval_size.to_string(), + "--tsv", + "--suppress-progress", + ]) + .output() + .expect("failed to run modkit summary"); + assert!( + output.status.success(), + "summary failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + // Summary currently has a separate, known nondeterministic row-order issue + // (U46), so compare its complete key/value row map while rejecting duplicate + // keys instead of weakening this test to a sorted bag of lines. + let stdout = String::from_utf8(output.stdout).unwrap(); + let mut rows = BTreeMap::new(); + for line in stdout.lines() { + let (key, value) = line + .split_once('\t') + .unwrap_or_else(|| panic!("invalid summary TSV row: {line}")); + assert!( + rows.insert(key.to_string(), value.to_string()).is_none(), + "duplicate summary TSV key: {key}" + ); + } + for required_key in ["mod_bases", "total_reads_used"] { + assert!( + rows.contains_key(required_key), + "summary TSV is missing metadata key: {required_key}" + ); + } + rows +} + +fn run_pileup( + out_dir: &Path, + name: &str, + threads: usize, + sampling_interval_size: u32, +) -> Vec { + let output_path = out_dir.join(format!("{name}.bed")); + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args([ + "pileup", + "-i", + "25", + "--sampling-frac", + "0.5", + "--filter-percentile", + "0.25", + "--seed", + "7", + "--threads", + &threads.to_string(), + "--sampling-threads", + &threads.to_string(), + "--sampling-interval-size", + &sampling_interval_size.to_string(), + "--suppress-progress", + "../tests/resources/bc_anchored_10_reads.sorted.bam", + output_path.to_str().unwrap(), + ]) + .output() + .expect("failed to run modkit pileup"); + assert!( + output.status.success(), + "pileup failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + fs::read(output_path).unwrap() +} + +fn help_text(args: &[&str]) -> String { + let output = Command::new(env!("CARGO_BIN_EXE_modkit")) + .args(args) + .output() + .expect("failed to run modkit help"); + assert!( + output.status.success(), + "help command failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + String::from_utf8(output.stdout) + .unwrap() + .split_whitespace() + .collect::>() + .join(" ") +} + +#[test] +fn seeded_indexed_sampling_is_stable_across_worker_and_interval_matrix() { + let temp_dir = tempfile::tempdir().unwrap(); + let input = "../tests/resources/bc_anchored_10_reads.sorted.bam"; + let mut reader = bam::Reader::from_path(input).unwrap(); + assert!(reader.records().any(|result| { + let record = result.unwrap(); + record.reference_end() - record.pos() > 20 + })); + let baseline = run_sample_probs( + temp_dir.path(), + "w1_i20", + input, + None, + Some("0.5"), + 7, + 1, + 20, + ); + + for (prefix, threads, interval_size) in [ + ("w1_i20_repeat", 1, 20), + ("w2_i20", 2, 20), + ("w3_i20", 3, 20), + ("w8_i20", 8, 20), + ("w8_i20_repeat", 8, 20), + ("w1_i100", 1, 100), + ("w3_i100", 3, 100), + ("w8_i1000", 8, 1_000), + ] { + let observed = run_sample_probs( + temp_dir.path(), + prefix, + input, + None, + Some("0.5"), + 7, + threads, + interval_size, + ); + assert!( + observed == baseline, + "seeded threshold or probability output changed for {threads} \ + workers and interval size {interval_size}" + ); + } + + let distinct_seed = run_sample_probs( + temp_dir.path(), + "seed8", + input, + None, + Some("0.5"), + 8, + 3, + 20, + ); + assert!( + distinct_seed != baseline, + "different explicit seeds unexpectedly produced identical outputs" + ); +} + +#[test] +fn seeded_indexed_sampling_matches_bam_and_selective_cram_decoding() { + let temp_dir = tempfile::tempdir().unwrap(); + let bam = run_sample_probs( + temp_dir.path(), + "bam", + "../tests/resources/bc_anchored_10_reads.sorted.bam", + None, + Some("0.5"), + 7, + 3, + 20, + ); + let cram = run_sample_probs( + temp_dir.path(), + "cram", + "../tests/resources/bc_anchored_10_reads.sorted.cram", + Some("../tests/resources/CGI_ladder_3.6kb_ref.fa"), + Some("0.5"), + 7, + 3, + 20, + ); + + assert!(cram == bam, "seeded BAM and CRAM outputs differed"); +} + +#[test] +fn boundary_fractions_exclude_none_or_all_exactly() { + let temp_dir = tempfile::tempdir().unwrap(); + let input = "../tests/resources/bc_anchored_10_reads.sorted.bam"; + let all = run_sample_probs( + temp_dir.path(), + "all", + input, + None, + Some("1"), + 7, + 1, + 20, + ); + let no_sampling = run_sample_probs( + temp_dir.path(), + "no_sampling", + input, + None, + None, + 0, + 8, + 1_000, + ); + assert!( + all == no_sampling, + "fraction 1 did not match the all-records control" + ); + + let none_serial = run_sample_probs( + temp_dir.path(), + "none_serial", + input, + None, + Some("0"), + 7, + 1, + 20, + ); + let none_parallel = run_sample_probs( + temp_dir.path(), + "none_parallel", + input, + None, + Some("0"), + 7, + 8, + 1_000, + ); + assert!( + none_serial == none_parallel, + "fraction 0 changed across scheduling" + ); + assert_eq!( + String::from_utf8(none_serial.1).unwrap().lines().count(), + 1, + "fraction 0 emitted probability observations" + ); +} + +#[test] +fn seeded_summary_and_pileup_consumers_are_schedule_independent() { + assert_eq!(run_summary(1, 20), run_summary(8, 1_000)); + + let temp_dir = tempfile::tempdir().unwrap(); + assert_eq!( + run_pileup(temp_dir.path(), "serial", 1, 20), + run_pileup(temp_dir.path(), "parallel", 8, 1_000) + ); +} + +#[test] +fn indexed_sampling_seed_help_describes_repeatable_decisions() { + for args in [["sample-probs", "--help"], ["call-mods", "--help"]] { + let help = help_text(&args); + assert!(help.contains( + "Provide a seed to make fractional read sampling decisions \ + repeatable for indexed and unindexed inputs" + )); + assert!(!help.contains("only used when no BAM index is provided")); + } +} diff --git a/modkit/tests/test_sampling_fraction.rs b/modkit/tests/test_sampling_fraction.rs new file mode 100644 index 00000000..de37be47 --- /dev/null +++ b/modkit/tests/test_sampling_fraction.rs @@ -0,0 +1,175 @@ +use std::process::{Command, Output}; + +fn run_modkit(args: &[String]) -> Output { + Command::new(env!("CARGO_BIN_EXE_modkit")) + .args(args) + .output() + .expect("failed to run modkit") +} + +fn assert_sampling_fraction_parse_error(args: Vec) { + let output = run_modkit(&args); + let stderr = String::from_utf8(output.stderr).unwrap(); + + assert_eq!( + output.status.code(), + Some(2), + "expected argument parsing to fail for {args:?}, stderr: {stderr}" + ); + assert!( + stderr.contains( + "sampling fraction must be a finite number in the inclusive \ + range [0, 1]" + ), + "unexpected stderr for {args:?}: {stderr}" + ); +} + +#[test] +fn invalid_sampling_fraction_fails_during_argument_parsing() { + let temp_dir = tempfile::tempdir().unwrap(); + let missing_bam = temp_dir.path().join("input-that-does-not-exist.bam"); + let missing_reference = + temp_dir.path().join("reference-that-does-not-exist.fa"); + let out_bed = temp_dir.path().join("pileup.bed"); + let out_calls = temp_dir.path().join("calls.tsv"); + let out_bam = temp_dir.path().join("calls.bam"); + let missing_bam = missing_bam.to_str().unwrap().to_string(); + let missing_reference = missing_reference.to_str().unwrap().to_string(); + + let cases = [ + vec![ + "pileup".to_string(), + missing_bam.clone(), + out_bed.to_str().unwrap().to_string(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "pileup-hemi".to_string(), + missing_bam.clone(), + "--reference".to_string(), + missing_reference, + "--cpg".to_string(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "extract".to_string(), + "calls".to_string(), + missing_bam.clone(), + out_calls.to_str().unwrap().to_string(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "sample-probs".to_string(), + missing_bam.clone(), + "--sample-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "modbam".to_string(), + "sample-probs".to_string(), + missing_bam.clone(), + "--sample-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "summary".to_string(), + missing_bam.clone(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "modbam".to_string(), + "summary".to_string(), + missing_bam.clone(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "call-mods".to_string(), + missing_bam.clone(), + out_bam.to_str().unwrap().to_string(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + vec![ + "modbam".to_string(), + "call-mods".to_string(), + missing_bam, + out_bam.to_str().unwrap().to_string(), + "--sampling-frac".to_string(), + "1.1".to_string(), + ], + ]; + + for args in cases { + assert_sampling_fraction_parse_error(args); + } + + for output_path in [&out_bed, &out_calls, &out_bam] { + assert!( + !output_path.exists(), + "command created {output_path:?} before argument validation" + ); + } +} + +#[test] +fn boundary_sampling_fractions_reach_runtime_validation() { + let temp_dir = tempfile::tempdir().unwrap(); + let missing_bam = temp_dir.path().join("missing.bam"); + let missing_bam = missing_bam.to_str().unwrap().to_string(); + + for fraction in ["0", "1"] { + let args = vec![ + "summary".to_string(), + missing_bam.clone(), + "--sampling-frac".to_string(), + fraction.to_string(), + ]; + let output = run_modkit(&args); + let stderr = String::from_utf8(output.stderr).unwrap(); + + assert_eq!( + output.status.code(), + Some(1), + "expected {fraction} to parse and reach missing-input handling, \ + stderr: {stderr}" + ); + assert!( + !stderr.contains("sampling fraction must"), + "boundary value {fraction} failed parsing: {stderr}" + ); + } +} + +#[test] +fn sampling_fraction_help_states_allowed_range() { + let cases = [ + vec!["pileup", "--help"], + vec!["pileup-hemi", "--help"], + vec!["extract", "calls", "--help"], + vec!["sample-probs", "--help"], + vec!["summary", "--help"], + vec!["call-mods", "--help"], + ]; + + for args in cases { + let args = args.into_iter().map(String::from).collect::>(); + let output = run_modkit(&args); + let stdout = String::from_utf8(output.stdout).unwrap(); + let normalized = + stdout.split_whitespace().collect::>().join(" "); + + assert!(output.status.success(), "help failed for {args:?}"); + assert!( + normalized.contains( + "Must be a finite value in the inclusive range [0, 1]" + ), + "sampling range missing from help for {args:?}: {stdout}" + ); + } +}