Skip to main content

salmon_model/
libdetect.rs

1//! Automatic library-type detection.
2//!
3//! Port of salmon's `LibraryTypeDetector`
4//! (`include/.../model/LibraryTypeDetector.hpp`). During the first reads of a
5//! run, the observed [`LibraryFormat`] of each confidently mapped fragment is
6//! tallied; once enough samples are seen, the most likely orientation and
7//! strandedness are inferred from the count ratios using salmon's 30%/70%
8//! thresholds.
9
10use salmon_core::{LibraryFormat, ReadOrientation, ReadStrandedness, ReadType};
11use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicU8, Ordering};
12
13/// Default number of samples to collect before guessing (matches salmon).
14pub const DEFAULT_SAMPLES_NEEDED: i64 = 50_000;
15
16/// Sentinel for `resolved` meaning "not yet locked in" (no valid format has this
17/// id; `MAX_FORMAT_ID` is 11).
18const UNSET_FORMAT: u8 = 0xFF;
19
20/// Accumulates observed library formats and infers the most likely type.
21#[derive(Debug)]
22pub struct LibraryTypeDetector {
23    /// still sampling (gates [`add_sample`]); cleared once the format locks in
24    active: AtomicBool,
25    read_type: ReadType,
26    samples_needed: AtomicI64,
27    counts: Vec<AtomicU64>,
28    /// the locked-in format id once detection completes, else [`UNSET_FORMAT`]
29    resolved: AtomicU8,
30}
31
32impl LibraryTypeDetector {
33    pub fn new(read_type: ReadType) -> Self {
34        let counts = (0..=LibraryFormat::MAX_FORMAT_ID)
35            .map(|_| AtomicU64::new(0))
36            .collect();
37        Self {
38            active: AtomicBool::new(true),
39            read_type,
40            samples_needed: AtomicI64::new(DEFAULT_SAMPLES_NEEDED),
41            counts,
42            resolved: AtomicU8::new(UNSET_FORMAT),
43        }
44    }
45
46    pub fn is_active(&self) -> bool {
47        self.active.load(Ordering::Relaxed)
48    }
49
50    /// True once enough samples have been collected to guess.
51    pub fn can_guess(&self) -> bool {
52        self.samples_needed.load(Ordering::Relaxed) <= 0
53    }
54
55    /// Record one confidently mapped fragment's observed format. Only formats
56    /// matching the detector's read type are counted, and only until the sample
57    /// budget is exhausted. Thread-safe.
58    pub fn add_sample(&self, f: LibraryFormat) {
59        if f.read_type == self.read_type && self.samples_needed.load(Ordering::Relaxed) >= 0 {
60            self.counts[f.format_id() as usize].fetch_add(1, Ordering::Relaxed);
61            self.samples_needed.fetch_sub(1, Ordering::Relaxed);
62        }
63    }
64
65    /// Mid-run resolution (salmon's prefix-detect-then-apply): once enough
66    /// samples have been collected ([`can_guess`]), infer and **lock in** the
67    /// library format (one writer wins the CAS), stop sampling, and return it;
68    /// idempotent thereafter. Returns `None` while still sampling. The caller
69    /// applies the returned format as a strand-compatibility filter for the rest
70    /// of the run.
71    pub fn resolved_format(&self) -> Option<LibraryFormat> {
72        let r = self.resolved.load(Ordering::Acquire);
73        if r != UNSET_FORMAT {
74            return Some(LibraryFormat::from_format_id(r));
75        }
76        if !self.can_guess() {
77            return None;
78        }
79        let f = self.infer_format();
80        match self.resolved.compare_exchange(
81            UNSET_FORMAT,
82            f.format_id(),
83            Ordering::AcqRel,
84            Ordering::Acquire,
85        ) {
86            Ok(_) => {
87                self.active.store(false, Ordering::Release);
88                Some(f)
89            }
90            // Another thread locked in first; use its result.
91            Err(existing) => Some(LibraryFormat::from_format_id(existing)),
92        }
93    }
94
95    /// The final library format to report at end of run: the locked-in format if
96    /// resolution happened mid-run, else inferred from whatever samples were
97    /// collected (recorded so repeat calls agree). Always returns a format.
98    pub fn final_format(&self) -> LibraryFormat {
99        let r = self.resolved.load(Ordering::Acquire);
100        if r != UNSET_FORMAT {
101            return LibraryFormat::from_format_id(r);
102        }
103        let f = self.infer_format();
104        let _ = self.resolved.compare_exchange(
105            UNSET_FORMAT,
106            f.format_id(),
107            Ordering::AcqRel,
108            Ordering::Acquire,
109        );
110        LibraryFormat::from_format_id(self.resolved.load(Ordering::Acquire))
111    }
112
113    /// Pure inference of the most likely library format from the accumulated
114    /// counts (no state change). Falls back to inward/unstranded when there are
115    /// no usable samples.
116    fn infer_format(&self) -> LibraryFormat {
117        let count = |id: u8| self.counts[id as usize].load(Ordering::Relaxed);
118
119        match self.read_type {
120            ReadType::SingleEnd => {
121                let mut nf = 0u64;
122                let mut nr = 0u64;
123                for id in 0..=LibraryFormat::MAX_FORMAT_ID {
124                    let f = LibraryFormat::from_format_id(id);
125                    let c = count(id);
126                    nf += if f.strandedness == ReadStrandedness::S {
127                        c
128                    } else {
129                        0
130                    };
131                    nr += if f.strandedness == ReadStrandedness::A {
132                        c
133                    } else {
134                        0
135                    };
136                }
137                let strandedness = if nf + nr == 0 {
138                    ReadStrandedness::U
139                } else {
140                    // Single-end uses the matching (S/A) encoding, like a
141                    // paired "same"-orientation library.
142                    strandedness_from_fw_ratio(nf as f64 / (nf + nr) as f64, true)
143                };
144                LibraryFormat::new(ReadType::SingleEnd, ReadOrientation::None, strandedness)
145            }
146            ReadType::PairedEnd => {
147                let (mut nsf, mut nsr) = (0u64, 0u64);
148                let (mut ninward, mut noutward, mut nsame) = (0u64, 0u64, 0u64);
149                for id in 0..=LibraryFormat::MAX_FORMAT_ID {
150                    let f = LibraryFormat::from_format_id(id);
151                    let c = count(id);
152                    nsf += matches!(f.strandedness, ReadStrandedness::S | ReadStrandedness::SA)
153                        .then_some(c)
154                        .unwrap_or(0);
155                    nsr += matches!(f.strandedness, ReadStrandedness::A | ReadStrandedness::AS)
156                        .then_some(c)
157                        .unwrap_or(0);
158                    match f.orientation {
159                        ReadOrientation::Toward => ninward += c,
160                        ReadOrientation::Away => noutward += c,
161                        ReadOrientation::Same => nsame += c,
162                        ReadOrientation::None => {}
163                    }
164                }
165
166                let num_orient = ninward + noutward + nsame;
167                if num_orient > 0 && (nsf + nsr) > 0 {
168                    let ratio_in = ninward as f64 / num_orient as f64;
169                    let ratio_out = noutward as f64 / num_orient as f64;
170                    let ratio_same = nsame as f64 / num_orient as f64;
171
172                    let (orientation, same) = if ratio_in >= ratio_out && ratio_in >= ratio_same {
173                        (ReadOrientation::Toward, false)
174                    } else if ratio_out >= ratio_in && ratio_out >= ratio_same {
175                        (ReadOrientation::Away, false)
176                    } else {
177                        (ReadOrientation::Same, true)
178                    };
179
180                    let ratio_fw = nsf as f64 / (nsf + nsr) as f64;
181                    let strandedness = strandedness_from_fw_ratio(ratio_fw, same);
182                    LibraryFormat::new(ReadType::PairedEnd, orientation, strandedness)
183                } else {
184                    LibraryFormat::new(
185                        ReadType::PairedEnd,
186                        ReadOrientation::Toward,
187                        ReadStrandedness::U,
188                    )
189                }
190            }
191        }
192    }
193}
194
195/// Map a forward-strand fraction to a strandedness using salmon's 30%/70%
196/// thresholds. `same` selects between the matching (S/A) and opposite (SA/AS)
197/// stranded encodings for paired-end "same"-orientation libraries; for
198/// single-end pass `false`.
199fn strandedness_from_fw_ratio(ratio_fw: f64, same: bool) -> ReadStrandedness {
200    if ratio_fw < 0.3 {
201        if same {
202            ReadStrandedness::A
203        } else {
204            ReadStrandedness::AS
205        }
206    } else if ratio_fw < 0.7 {
207        ReadStrandedness::U
208    } else if same {
209        ReadStrandedness::S
210    } else {
211        ReadStrandedness::SA
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    #[test]
220    fn single_end_detects_sense() {
221        let d = LibraryTypeDetector::new(ReadType::SingleEnd);
222        let sf = LibraryFormat::parse("SF").unwrap();
223        let sr = LibraryFormat::parse("SR").unwrap();
224        for _ in 0..90 {
225            d.add_sample(sf);
226        }
227        for _ in 0..10 {
228            d.add_sample(sr);
229        }
230        assert_eq!(d.infer_format().canonical(), "SF");
231        // final_format records and returns the same result idempotently
232        assert_eq!(d.final_format().canonical(), "SF");
233        assert_eq!(d.final_format().canonical(), "SF");
234    }
235
236    #[test]
237    fn single_end_detects_unstranded() {
238        let d = LibraryTypeDetector::new(ReadType::SingleEnd);
239        let sf = LibraryFormat::parse("SF").unwrap();
240        let sr = LibraryFormat::parse("SR").unwrap();
241        for _ in 0..50 {
242            d.add_sample(sf);
243            d.add_sample(sr);
244        }
245        assert_eq!(d.infer_format().canonical(), "U");
246    }
247
248    #[test]
249    fn paired_end_detects_isr() {
250        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
251        let isr = LibraryFormat::parse("ISR").unwrap();
252        for _ in 0..100 {
253            d.add_sample(isr);
254        }
255        // ISR: inward + antisense -> toward + AS
256        assert_eq!(d.infer_format().canonical(), "ISR");
257    }
258
259    #[test]
260    fn paired_end_detects_iu() {
261        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
262        let isf = LibraryFormat::parse("ISF").unwrap();
263        let isr = LibraryFormat::parse("ISR").unwrap();
264        for _ in 0..50 {
265            d.add_sample(isf);
266            d.add_sample(isr);
267        }
268        // balanced strandedness -> unstranded, inward -> IU
269        assert_eq!(d.infer_format().canonical(), "IU");
270    }
271
272    #[test]
273    fn resolved_format_locks_in_after_prefix() {
274        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
275        let isr = LibraryFormat::parse("ISR").unwrap();
276        // Before the sample budget is consumed: no resolution yet, so the caller
277        // applies no filter, and the detector keeps sampling.
278        assert!(d.resolved_format().is_none());
279        assert!(d.is_active());
280        // Feed the full prefix budget.
281        for _ in 0..DEFAULT_SAMPLES_NEEDED {
282            d.add_sample(isr);
283        }
284        assert!(d.can_guess());
285        // Now it locks in to the inferred type, stops sampling, and is idempotent.
286        assert_eq!(d.resolved_format().unwrap().canonical(), "ISR");
287        assert!(!d.is_active());
288        assert_eq!(d.resolved_format().unwrap().canonical(), "ISR");
289        assert_eq!(d.final_format().canonical(), "ISR");
290    }
291
292    #[test]
293    fn final_format_without_lockin_infers_from_partial() {
294        // Fewer than the budget: never locks in mid-run, but end-of-run reporting
295        // still returns a best-guess format from the partial samples.
296        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
297        let isf = LibraryFormat::parse("ISF").unwrap();
298        for _ in 0..100 {
299            d.add_sample(isf);
300        }
301        assert!(d.resolved_format().is_none()); // not enough to lock in mid-run
302        assert_eq!(d.final_format().canonical(), "ISF");
303    }
304
305    #[test]
306    fn sample_budget_is_respected() {
307        let d = LibraryTypeDetector::new(ReadType::SingleEnd);
308        assert!(!d.can_guess());
309        let sf = LibraryFormat::parse("SF").unwrap();
310        // exhaust the budget
311        let mut n = DEFAULT_SAMPLES_NEEDED + 5;
312        while n > 0 {
313            d.add_sample(sf);
314            n -= 1;
315        }
316        assert!(d.can_guess());
317    }
318}