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 counts: Vec<u64> = self
118            .counts
119            .iter()
120            .map(|c| c.load(Ordering::Relaxed))
121            .collect();
122        infer_format_from_counts(&counts, self.read_type)
123    }
124}
125
126/// Pure inference of the most likely library format from per-format counts
127/// (indexed by [`LibraryFormat::format_id`]), applying salmon's orientation and
128/// strandedness thresholds. Order-independent — used both by the prefix-sampling
129/// [`LibraryTypeDetector`] and by RAD auto-detection, which tallies *all* unique
130/// fragments rather than a thread-order-dependent prefix. Falls back to
131/// inward/unstranded when there are no usable samples.
132pub fn infer_format_from_counts(counts: &[u64], read_type: ReadType) -> LibraryFormat {
133    let count = |id: u8| counts[id as usize];
134
135    match read_type {
136        ReadType::SingleEnd => {
137            let mut nf = 0u64;
138            let mut nr = 0u64;
139            for id in 0..=LibraryFormat::MAX_FORMAT_ID {
140                let f = LibraryFormat::from_format_id(id);
141                let c = count(id);
142                nf += if f.strandedness == ReadStrandedness::S {
143                    c
144                } else {
145                    0
146                };
147                nr += if f.strandedness == ReadStrandedness::A {
148                    c
149                } else {
150                    0
151                };
152            }
153            let strandedness = if nf + nr == 0 {
154                ReadStrandedness::U
155            } else {
156                // Single-end uses the matching (S/A) encoding, like a
157                // paired "same"-orientation library.
158                strandedness_from_fw_ratio(nf as f64 / (nf + nr) as f64, true)
159            };
160            LibraryFormat::new(ReadType::SingleEnd, ReadOrientation::None, strandedness)
161        }
162        ReadType::PairedEnd => {
163            let (mut nsf, mut nsr) = (0u64, 0u64);
164            let (mut ninward, mut noutward, mut nsame) = (0u64, 0u64, 0u64);
165            for id in 0..=LibraryFormat::MAX_FORMAT_ID {
166                let f = LibraryFormat::from_format_id(id);
167                let c = count(id);
168                nsf += matches!(f.strandedness, ReadStrandedness::S | ReadStrandedness::SA)
169                    .then_some(c)
170                    .unwrap_or(0);
171                nsr += matches!(f.strandedness, ReadStrandedness::A | ReadStrandedness::AS)
172                    .then_some(c)
173                    .unwrap_or(0);
174                match f.orientation {
175                    ReadOrientation::Toward => ninward += c,
176                    ReadOrientation::Away => noutward += c,
177                    ReadOrientation::Same => nsame += c,
178                    ReadOrientation::None => {}
179                }
180            }
181
182            let num_orient = ninward + noutward + nsame;
183            if num_orient > 0 && (nsf + nsr) > 0 {
184                let ratio_in = ninward as f64 / num_orient as f64;
185                let ratio_out = noutward as f64 / num_orient as f64;
186                let ratio_same = nsame as f64 / num_orient as f64;
187
188                let (orientation, same) = if ratio_in >= ratio_out && ratio_in >= ratio_same {
189                    (ReadOrientation::Toward, false)
190                } else if ratio_out >= ratio_in && ratio_out >= ratio_same {
191                    (ReadOrientation::Away, false)
192                } else {
193                    (ReadOrientation::Same, true)
194                };
195
196                let ratio_fw = nsf as f64 / (nsf + nsr) as f64;
197                let strandedness = strandedness_from_fw_ratio(ratio_fw, same);
198                LibraryFormat::new(ReadType::PairedEnd, orientation, strandedness)
199            } else {
200                LibraryFormat::new(
201                    ReadType::PairedEnd,
202                    ReadOrientation::Toward,
203                    ReadStrandedness::U,
204                )
205            }
206        }
207    }
208}
209
210/// Map a forward-strand fraction to a strandedness using salmon's 30%/70%
211/// thresholds. `same` selects between the matching (S/A) and opposite (SA/AS)
212/// stranded encodings for paired-end "same"-orientation libraries; for
213/// single-end pass `false`.
214fn strandedness_from_fw_ratio(ratio_fw: f64, same: bool) -> ReadStrandedness {
215    if ratio_fw < 0.3 {
216        if same {
217            ReadStrandedness::A
218        } else {
219            ReadStrandedness::AS
220        }
221    } else if ratio_fw < 0.7 {
222        ReadStrandedness::U
223    } else if same {
224        ReadStrandedness::S
225    } else {
226        ReadStrandedness::SA
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn single_end_detects_sense() {
236        let d = LibraryTypeDetector::new(ReadType::SingleEnd);
237        let sf = LibraryFormat::parse("SF").unwrap();
238        let sr = LibraryFormat::parse("SR").unwrap();
239        for _ in 0..90 {
240            d.add_sample(sf);
241        }
242        for _ in 0..10 {
243            d.add_sample(sr);
244        }
245        assert_eq!(d.infer_format().canonical(), "SF");
246        // final_format records and returns the same result idempotently
247        assert_eq!(d.final_format().canonical(), "SF");
248        assert_eq!(d.final_format().canonical(), "SF");
249    }
250
251    #[test]
252    fn single_end_detects_unstranded() {
253        let d = LibraryTypeDetector::new(ReadType::SingleEnd);
254        let sf = LibraryFormat::parse("SF").unwrap();
255        let sr = LibraryFormat::parse("SR").unwrap();
256        for _ in 0..50 {
257            d.add_sample(sf);
258            d.add_sample(sr);
259        }
260        assert_eq!(d.infer_format().canonical(), "U");
261    }
262
263    #[test]
264    fn paired_end_detects_isr() {
265        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
266        let isr = LibraryFormat::parse("ISR").unwrap();
267        for _ in 0..100 {
268            d.add_sample(isr);
269        }
270        // ISR: inward + antisense -> toward + AS
271        assert_eq!(d.infer_format().canonical(), "ISR");
272    }
273
274    #[test]
275    fn paired_end_detects_iu() {
276        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
277        let isf = LibraryFormat::parse("ISF").unwrap();
278        let isr = LibraryFormat::parse("ISR").unwrap();
279        for _ in 0..50 {
280            d.add_sample(isf);
281            d.add_sample(isr);
282        }
283        // balanced strandedness -> unstranded, inward -> IU
284        assert_eq!(d.infer_format().canonical(), "IU");
285    }
286
287    #[test]
288    fn resolved_format_locks_in_after_prefix() {
289        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
290        let isr = LibraryFormat::parse("ISR").unwrap();
291        // Before the sample budget is consumed: no resolution yet, so the caller
292        // applies no filter, and the detector keeps sampling.
293        assert!(d.resolved_format().is_none());
294        assert!(d.is_active());
295        // Feed the full prefix budget.
296        for _ in 0..DEFAULT_SAMPLES_NEEDED {
297            d.add_sample(isr);
298        }
299        assert!(d.can_guess());
300        // Now it locks in to the inferred type, stops sampling, and is idempotent.
301        assert_eq!(d.resolved_format().unwrap().canonical(), "ISR");
302        assert!(!d.is_active());
303        assert_eq!(d.resolved_format().unwrap().canonical(), "ISR");
304        assert_eq!(d.final_format().canonical(), "ISR");
305    }
306
307    #[test]
308    fn final_format_without_lockin_infers_from_partial() {
309        // Fewer than the budget: never locks in mid-run, but end-of-run reporting
310        // still returns a best-guess format from the partial samples.
311        let d = LibraryTypeDetector::new(ReadType::PairedEnd);
312        let isf = LibraryFormat::parse("ISF").unwrap();
313        for _ in 0..100 {
314            d.add_sample(isf);
315        }
316        assert!(d.resolved_format().is_none()); // not enough to lock in mid-run
317        assert_eq!(d.final_format().canonical(), "ISF");
318    }
319
320    #[test]
321    fn sample_budget_is_respected() {
322        let d = LibraryTypeDetector::new(ReadType::SingleEnd);
323        assert!(!d.can_guess());
324        let sf = LibraryFormat::parse("SF").unwrap();
325        // exhaust the budget
326        let mut n = DEFAULT_SAMPLES_NEEDED + 5;
327        while n > 0 {
328            d.add_sample(sf);
329            n -= 1;
330        }
331        assert!(d.can_guess());
332    }
333}