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 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 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
195fn 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 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 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 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 assert!(d.resolved_format().is_none());
279 assert!(d.is_active());
280 for _ in 0..DEFAULT_SAMPLES_NEEDED {
282 d.add_sample(isr);
283 }
284 assert!(d.can_guess());
285 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 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()); 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 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}