diff --git a/src/chainer.rs b/src/chainer.rs index fb524e5d..9a3e64de 100644 --- a/src/chainer.rs +++ b/src/chainer.rs @@ -187,7 +187,7 @@ impl Chainer { let mut n_anchors = 0; let mut time_chaining = 0.0; let mut chains = vec![]; - for is_revcomp in orientations { + for &is_revcomp in &orientations { let hits_timer = Instant::now(); let mut anchors = hits_to_anchors(&hits[is_revcomp], index); time_find_hits += hits_timer.elapsed().as_secs_f64(); @@ -241,6 +241,7 @@ impl Chainer { time_chaining, time_rescue: 0.0, time_sort_nams: 0f64, + both_orientations: orientations.len() > 1, }; (details, chains) diff --git a/src/details.rs b/src/details.rs index d11c7dff..d7c86406 100644 --- a/src/details.rs +++ b/src/details.rs @@ -23,6 +23,9 @@ pub struct NamDetails { pub time_chaining: f64, pub time_rescue: f64, pub time_sort_nams: f64, + + /// Whether both orientations were tested + pub both_orientations: bool, } impl ops::AddAssign for NamDetails { diff --git a/src/maponly.rs b/src/maponly.rs index e37f2e30..81376d48 100644 --- a/src/maponly.rs +++ b/src/maponly.rs @@ -1,15 +1,16 @@ use fastrand::Rng; use crate::chainer::Chainer; -use crate::details::Details; +use crate::details::{Details, NamDetails}; use crate::index::StrobemerIndex; use crate::insertsize::InsertSizeDistribution; use crate::io::fasta::RefSequence; use crate::io::paf::PafRecord; use crate::io::record::{End, SequenceRecord}; -use crate::mapper::{NamPair, get_best_scoring_nam_pairs, mapping_quality}; +use crate::mapper::mapping_quality; +use crate::math::normal_pdf; use crate::mcsstrategy::McsStrategy; -use crate::nam::{Nam, get_nams_by_chaining}; +use crate::nam::{Nam, get_nams_by_chaining, sort_nams}; /// Map a single-end read to the reference and return PAF records /// @@ -23,14 +24,14 @@ pub fn map_single_end_read( chainer: &Chainer, rng: &mut Rng, ) -> (Vec, Details) { - let (nam_details, nams) = get_nams_by_chaining( + let (mut nam_details, mut nams) = get_nams_by_chaining( &record.sequence, index, chainer, rescue_distance, mcs_strategy, - rng, ); + nam_details.time_sort_nams = sort_nams(&mut nams, rng); if nams.is_empty() { (vec![], nam_details.into()) @@ -62,14 +63,15 @@ pub fn abundances_single_end_read( chainer: &Chainer, rng: &mut Rng, ) { - let (_, nams) = get_nams_by_chaining( + let (_, mut nams) = get_nams_by_chaining( &record.sequence, index, chainer, rescue_distance, mcs_strategy, - rng, ); + sort_nams(&mut nams, rng); + let n_best = nams .iter() .take_while(|nam| nam.score == nams[0].score) @@ -120,59 +122,55 @@ pub fn map_paired_end_read( chainer: &Chainer, rng: &mut Rng, ) -> (Vec, Details) { - let (mut nam_details1, nams1) = get_nams_by_chaining( - &r1.sequence, - index, - chainer, - rescue_distance, - mcs_strategy, - rng, - ); - let (nam_details2, nams2) = get_nams_by_chaining( - &r2.sequence, - index, - chainer, - rescue_distance, - mcs_strategy, - rng, - ); + let (mut nam_details1, mut nams1) = + get_nams_by_chaining(&r1.sequence, index, chainer, rescue_distance, mcs_strategy); + let (mut nam_details2, mut nams2) = + get_nams_by_chaining(&r2.sequence, index, chainer, rescue_distance, mcs_strategy); - let nam_pairs = get_best_scoring_nam_pairs( - &nams1, - &nams2, + if nams1.is_empty() && nams2.is_empty() { + nam_details1 += nam_details2; + return (vec![], nam_details1.into()); + } + + let nam_pairs = get_nam_pairs( + &mut nams1, + &mut nams2, insert_size_distribution.mu, insert_size_distribution.sigma, + &nam_details1, + &nam_details2, ); - let mapped_nam = - get_best_paired_map_location(&nam_pairs, &nams1, &nams2, insert_size_distribution); - let mut records = vec![]; - match mapped_nam { - MappedNams::Individual(nam1, nam2) => { - if let Some(nam) = nam1 { + nam_details1.time_sort_nams = sort_nams(&mut nams1, rng); + nam_details2.time_sort_nams = sort_nams(&mut nams2, rng); + + let mut records = vec![]; + match get_best_paired_mapping_location(&nam_pairs, &nams1, &nams2, insert_size_distribution) { + MappedNams::Individual(best1, best2) => { + if let Some(nam1) = best1 { records.push(paf_record_from_nam( - &nam, + nam1, &r1.name, references, r1.sequence.len(), None, End::One, - )) + )); } - if let Some(nam) = nam2 { + if let Some(nam2) = best2 { records.push(paf_record_from_nam( - &nam, + nam2, &r2.name, references, r2.sequence.len(), None, End::Two, - )) + )); } } - MappedNams::Pair(nam1, nam2) => { + MappedNams::Pair(nam1, nam2, _) => { records.push(paf_record_from_nam( - &nam1, + nam1, &r1.name, references, r1.sequence.len(), @@ -180,7 +178,7 @@ pub fn map_paired_end_read( End::One, )); records.push(paf_record_from_nam( - &nam2, + nam2, &r2.name, references, r2.sequence.len(), @@ -188,8 +186,8 @@ pub fn map_paired_end_read( End::Two, )); } - MappedNams::Unmapped => {} } + nam_details1 += nam_details2; (records, nam_details1.into()) } @@ -208,56 +206,28 @@ pub fn abundances_paired_end_read( chainer: &Chainer, rng: &mut Rng, ) { - let nams1 = get_nams_by_chaining( - &r1.sequence, - index, - chainer, - rescue_distance, - mcs_strategy, - rng, - ) - .1; - let nams2 = get_nams_by_chaining( - &r2.sequence, - index, - chainer, - rescue_distance, - mcs_strategy, - rng, - ) - .1; + let (nam_details1, mut nams1) = + get_nams_by_chaining(&r1.sequence, index, chainer, rescue_distance, mcs_strategy); + let (nam_details2, mut nams2) = + get_nams_by_chaining(&r2.sequence, index, chainer, rescue_distance, mcs_strategy); - let nam_pairs = get_best_scoring_nam_pairs( - &nams1, - &nams2, + if nams1.is_empty() && nams2.is_empty() { + return; + } + + let nam_pairs = get_nam_pairs( + &mut nams1, + &mut nams2, insert_size_distribution.mu, insert_size_distribution.sigma, + &nam_details1, + &nam_details2, ); - let mapped_nam = - get_best_paired_map_location(&nam_pairs, &nams1, &nams2, insert_size_distribution); - match mapped_nam { - MappedNams::Pair(nam1, nam2) => { - let joint_score = nam1.score + nam2.score; - let n_best = nam_pairs - .iter() - .take_while(|nam_pair| { - nam_pair.nam1.as_ref().map_or(0.0, |nam| nam.score) - + nam_pair.nam2.as_ref().map_or(0.0, |nam| nam.score) - == joint_score - }) - .count(); - let weight_r1 = r1.sequence.len() as f64 / n_best as f64; - let weight_r2 = r2.sequence.len() as f64 / n_best as f64; - for nam_pair in &nam_pairs[..n_best] { - if let Some(nam) = &nam_pair.nam1 { - abundances[nam.ref_id] += weight_r1; - } - if let Some(nam) = &nam_pair.nam2 { - abundances[nam.ref_id] += weight_r2; - } - } - } + sort_nams(&mut nams1, rng); + sort_nams(&mut nams2, rng); + + match get_best_paired_mapping_location(&nam_pairs, &nams1, &nams2, insert_size_distribution) { MappedNams::Individual(_, _) => { for (nams, read_len) in [(&nams1, r1.sequence.len()), (&nams2, r2.sequence.len())] { let n_best = nams @@ -270,66 +240,219 @@ pub fn abundances_paired_end_read( } } } - MappedNams::Unmapped => {} + MappedNams::Pair(_, _, joint_score) => { + let n_best = nam_pairs + .iter() + .take_while(|nam_pair| nam_pair.score == joint_score) + .count(); + let weight_r1 = r1.sequence.len() as f64 / n_best as f64; + let weight_r2 = r2.sequence.len() as f64 / n_best as f64; + for NamPair { + nam1, + nam2, + score: _, + } in &nam_pairs[..n_best] + { + abundances[nam1.ref_id] += weight_r1; + abundances[nam2.ref_id] += weight_r2; + } + } } } -enum MappedNams { - Individual(Option, Option), - Pair(Nam, Nam), - Unmapped, +enum MappedNams<'a> { + /// Two independent best NAMs (one per read) + Individual(Option<&'a Nam>, Option<&'a Nam>), + /// A proper paired NAMs (nam1, nam2, pairing score) + Pair(&'a Nam, &'a Nam, f64), } -/// Given two lists of NAMs from R1 and R2, find the best location (preferably a proper pair). -/// This is used for mapping-only (PAF) mode and abundances output -fn get_best_paired_map_location( - nam_pairs: &[NamPair], - nams1: &[Nam], - nams2: &[Nam], +/// Choose between: +/// - the best proper pair of mappings +/// - the best individual mappings +/// +/// Also updates the insert size distribution using confident pairs. +/// +/// For paired-end mapping and abundance estimation modes only +fn get_best_paired_mapping_location<'a>( + nam_pairs: &'a [NamPair], + nams1: &'a [Nam], + nams2: &'a [Nam], insert_size_distribution: &mut InsertSizeDistribution, -) -> MappedNams { - if nam_pairs.is_empty() && nams1.is_empty() && nams2.is_empty() { - return MappedNams::Unmapped; +) -> MappedNams<'a> { + let best_nam1 = nams1.first(); + let best_nam2 = nams2.first(); + + // Score if reads are treated independently. + let individual_score = best_nam1.map_or(0.0, |nam| nam.score as f64) + + best_nam2.map_or(0.0, |nam| nam.score as f64); + + // Prefer a proper pair only if it beats a penalized individual mapping. + // Divisor 2 is penalty for being mapped individually + if let Some(NamPair { nam1, nam2, score }) = nam_pairs.first() + && *score >= individual_score / 2.0 + { + // Update insert size using confident proper pairs. + if insert_size_distribution.sample_size < 400 { + insert_size_distribution.update(nam1.ref_start.abs_diff(nam2.ref_start)); + } + + MappedNams::Pair(nam1, nam2, *score) + } else { + MappedNams::Individual(best_nam1, best_nam2) } +} - // Find first NAM pair that is a proper pair. - // The first one is also the one with the highest score - // since nam_pairs is sorted descending by score - let best_joint_pair = nam_pairs - .iter() - .find(|&nam_pair| nam_pair.nam1.is_some() && nam_pair.nam2.is_some()); +#[derive(Debug)] +pub struct NamPair { + pub nam1: Nam, + pub nam2: Nam, + pub score: f64, +} + +/// Build all plausible forward/revcomp mapping pairings +fn get_nam_pairs( + nams1: &mut [Nam], + nams2: &mut [Nam], + mu: f32, + sigma: f32, + details1: &NamDetails, + details2: &NamDetails, +) -> Vec { + let mut nam_pairs = vec![]; + if nams1.is_empty() || nams2.is_empty() { + return nam_pairs; + } - let joint_score = if let Some(nam_pair) = best_joint_pair { - nam_pair.nam1.as_ref().map_or(0.0, |nam| nam.score) - + nam_pair.nam2.as_ref().map_or(0.0, |nam| nam.score) + let (fwd1, rev1): (&mut [Nam], &mut [Nam]) = + split_nams_by_orientation_checked(nams1, details1.both_orientations); + let (fwd2, rev2): (&mut [Nam], &mut [Nam]) = + split_nams_by_orientation_checked(nams2, details2.both_orientations); + + if !fwd1.is_empty() && !rev2.is_empty() { + fwd1.sort_unstable_by_key(|nam| (nam.ref_id, nam.projected_ref_start())); + rev2.sort_unstable_by_key(|nam| (nam.ref_id, nam.projected_ref_start())); + nam_pairs.extend(find_pairs(fwd1, rev2, mu, sigma, false)); + } + if !fwd2.is_empty() && !rev1.is_empty() { + fwd2.sort_unstable_by_key(|nam| (nam.ref_id, nam.projected_ref_start())); + rev1.sort_unstable_by_key(|nam| (nam.ref_id, nam.projected_ref_start())); + nam_pairs.extend(find_pairs(fwd2, rev1, mu, sigma, true)); + } + + nam_pairs.sort_unstable_by(|a, b| b.score.total_cmp(&a.score)); + nam_pairs +} + +/// Split nams into (forward, revcomp), +/// if only 1 orientation exists, returns it and a empty slice +fn split_nams_by_orientation_checked(nams: &mut [Nam], both: bool) -> (&mut [Nam], &mut [Nam]) { + if both { + split_nams_by_orientation(nams) + } else if nams[0].is_revcomp { + (&mut [], nams) } else { - 0.0 - }; + (nams, &mut []) + } +} - // Get individual best scores. - // nams1 and nams2 are also sorted descending by score. - let best_individual_nam1 = nams1.first(); - let best_individual_nam2 = nams2.first(); +/// In-place partition of NAMs by orientation: +/// forward on the left, revcomp on the right. +/// Returns two slices separating (forward, revcomp) +fn split_nams_by_orientation(nams: &mut [Nam]) -> (&mut [Nam], &mut [Nam]) { + let mut left = 0; + let mut right = nams.len(); - let individual_score = best_individual_nam1.map_or(0.0, |nam| nam.score) - + best_individual_nam2.map_or(0.0, |nam| nam.score); + while left < right { + if nams[left].is_revcomp { + right -= 1; + nams.swap(left, right); + } else { + left += 1; + } + } - // Divisor 2 is penalty for being mapped individually - if joint_score > individual_score / 2.0 { - let best_joint_pair = best_joint_pair.unwrap(); - let best = (best_joint_pair.nam1.clone(), best_joint_pair.nam2.clone()); - if insert_size_distribution.sample_size < 400 { - insert_size_distribution.update( - best.0 - .as_ref() - .unwrap() - .ref_start - .abs_diff(best.1.as_ref().unwrap().ref_start), - ); + nams.split_at_mut(left) +} + +/// Find most forward/revcomp pairs using a two-pointer scan. +/// Assumes both slices are sorted by (ref_id, projected_ref_start). +fn find_pairs(fwd: &[Nam], rev: &[Nam], mu: f32, sigma: f32, swap_order: bool) -> Vec { + let mut out = Vec::new(); + let max_dist = (mu + 10.0 * sigma).ceil() as usize; // distance cutoff from insert size distribution + let mut rev_ptr = 0; + let mut last_paired = None; + + for f in fwd { + // Advance revcomp pointer to the first possible candidate + while rev_ptr < rev.len() + && (rev[rev_ptr].ref_id < f.ref_id + || rev[rev_ptr].ref_id == f.ref_id + && rev[rev_ptr].projected_ref_start() < f.projected_ref_start()) + { + rev_ptr += 1; + } + if rev_ptr == rev.len() { + break; + } + if rev[rev_ptr].ref_id > f.ref_id { + continue; } - MappedNams::Pair(best.0.unwrap(), best.1.unwrap()) - } else { - MappedNams::Individual(best_individual_nam1.cloned(), best_individual_nam2.cloned()) + // Scan window of revcomp nams within distance limit. + let mut best = None; + let mut i = rev_ptr; + while i < rev.len() + && rev[i].ref_id == f.ref_id + && (rev[i].projected_ref_start() - f.projected_ref_start()) <= max_dist + { + let r = &rev[i]; + // The pairing score gets a bonus based on the reference distance of the two chosen nam + // paired and from our current knowledge of the reference distance distribution + let x = f.ref_start.abs_diff(r.ref_start); + let score = f.score as f64 + + r.score as f64 + + 0.001f64.max((normal_pdf(x as f32, mu, sigma) + 1.0).ln() as f64); + + if best.is_none_or(|(_, highest_score)| score > highest_score) { + best = Some((i, score)); + } + i += 1; + } + + // Highest scoring candidate + let Some((best_id, score)) = best else { + continue; + }; + let r = &rev[best_id]; + + // If the same revcomp nam was paired previously, keep only the better scoring pair. + if let Some((last_id, prev_score)) = last_paired + && last_id == best_id + { + if score <= prev_score { + continue; // keep the previous better pair + } + out.pop(); // replace it + } + + out.push(if swap_order { + NamPair { + nam1: r.clone(), + nam2: f.clone(), + score, + } + } else { + NamPair { + nam1: f.clone(), + nam2: r.clone(), + score, + } + }); + + last_paired = Some((best_id, score)); + rev_ptr = best_id; } + + out } diff --git a/src/mapper.rs b/src/mapper.rs index a1fc552d..f25278c3 100644 --- a/src/mapper.rs +++ b/src/mapper.rs @@ -21,7 +21,7 @@ use crate::io::sam::{ }; use crate::math::normal_pdf; use crate::mcsstrategy::McsStrategy; -use crate::nam::{Nam, get_nams_by_chaining, reverse_nam_if_needed}; +use crate::nam::{Nam, get_nams_by_chaining, reverse_nam_if_needed, sort_nams}; use crate::piecewisealigner::remove_spurious_anchors; use crate::read::Read; use crate::revcomp::reverse_complement; @@ -330,14 +330,14 @@ pub fn align_single_end_read( aligner: &Aligner, rng: &mut Rng, ) -> (Vec, Details) { - let (nam_details, mut nams) = get_nams_by_chaining( + let (mut nam_details, mut nams) = get_nams_by_chaining( &record.sequence, index, chainer, mapping_parameters.rescue_distance, mapping_parameters.mcs_strategy, - rng, ); + nam_details.time_sort_nams = sort_nams(&mut nams, rng); let mut details: Details = nam_details.into(); let timer = Instant::now(); @@ -600,14 +600,14 @@ pub fn align_paired_end_read( for is_r1 in [0, 1] { let record = if is_r1 == 0 { r1 } else { r2 }; - let (nam_details, nams) = get_nams_by_chaining( + let (mut nam_details, mut nams) = get_nams_by_chaining( &record.sequence, index, chainer, mapping_parameters.rescue_distance, mapping_parameters.mcs_strategy, - rng, ); + nam_details.time_sort_nams = sort_nams(&mut nams, rng); details[is_r1].nam = nam_details; nams_pair[is_r1] = nams; } @@ -1122,7 +1122,7 @@ fn is_proper_pair( } } -fn is_proper_nam_pair(nam1: &Nam, nam2: &Nam, mu: f32, sigma: f32) -> bool { +pub fn is_proper_nam_pair(nam1: &Nam, nam2: &Nam, mu: f32, sigma: f32) -> bool { if nam1.ref_id != nam2.ref_id || nam1.is_revcomp == nam2.is_revcomp { return false; } diff --git a/src/nam.rs b/src/nam.rs index 4d3e5022..334f6fb6 100644 --- a/src/nam.rs +++ b/src/nam.rs @@ -15,7 +15,7 @@ use crate::read::Read; use crate::seeding::randstrobes_query; /// Non-overlapping approximate match -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub struct Nam { pub nam_id: usize, pub ref_start: usize, @@ -127,15 +127,12 @@ pub fn reverse_nam_if_needed( } /// Obtain NAMs for a sequence record, doing rescue if needed. -/// -/// NAMs are returned sorted by decreasing score pub fn get_nams_by_chaining( sequence: &[u8], index: &StrobemerIndex, chainer: &Chainer, rescue_distance: usize, mcs_strategy: McsStrategy, - rng: &mut Rng, ) -> (NamDetails, Vec) { let timer = Instant::now(); let query_randstrobes = randstrobes_query(sequence, &index.parameters); @@ -147,20 +144,23 @@ pub fn get_nams_by_chaining( query_randstrobes[1].len() ); - let (mut nam_details, mut nams) = + let (mut nam_details, nams) = chainer.get_chains(&query_randstrobes, index, rescue_distance, mcs_strategy); - let timer = Instant::now(); + nam_details.time_randstrobes = time_randstrobes; + (nam_details, nams) +} + +pub fn sort_nams(nams: &mut [Nam], rng: &mut Rng) -> f64 { + let timer = Instant::now(); nams.sort_by(|a, b| b.score.total_cmp(&a.score)); - shuffle_top_nams(&mut nams, rng); - nam_details.time_sort_nams = timer.elapsed().as_secs_f64(); - nam_details.time_randstrobes = time_randstrobes; + shuffle_top_nams(nams, rng); if log::log_enabled!(Trace) { trace!("Found {} NAMs", nams.len()); let mut printed = 0; - for nam in &nams { + for nam in nams.iter() { if nam.n_matches > 1 || printed < 10 { trace!("- {}", nam); printed += 1; @@ -171,7 +171,7 @@ pub fn get_nams_by_chaining( } } - (nam_details, nams) + timer.elapsed().as_secs_f64() } /// Shuffle the top-scoring NAMs. Input must be sorted by score.