salmon_model/
libdetect.rs1use salmon_core::{LibraryFormat, ReadOrientation, ReadStrandedness, ReadType};
11use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicU8, Ordering};
12
13pub const DEFAULT_SAMPLES_NEEDED: i64 = 50_000;
15
16const UNSET_FORMAT: u8 = 0xFF;
19
20#[derive(Debug)]
22pub struct LibraryTypeDetector {
23 active: AtomicBool,
25 read_type: ReadType,
26 samples_needed: AtomicI64,
27 counts: Vec<AtomicU64>,
28 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 pub fn can_guess(&self) -> bool {
52 self.samples_needed.load(Ordering::Relaxed) <= 0
53 }
54
55 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 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 Err(existing) => Some(LibraryFormat::from_format_id(existing)),
92 }
93 }
94
95 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 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
126pub 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 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
210fn 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 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 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 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 assert!(d.resolved_format().is_none());
294 assert!(d.is_active());
295 for _ in 0..DEFAULT_SAMPLES_NEEDED {
297 d.add_sample(isr);
298 }
299 assert!(d.can_guess());
300 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 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()); 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 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}