1use std::cmp::Ordering;
2use std::collections::VecDeque;
3use std::io;
4use std::ops::Range;
5use std::{env, fs, path};
6
7use mzdata::prelude::*;
8use mzdata::spectrum::{MultiLayerSpectrum, SpectrumGroup, SpectrumGroupingIterator};
9use mzpeaks::{CentroidPeak, DeconvolutedPeak, Tolerance};
10
11#[derive(Default, Debug, Clone, PartialEq)]
12pub struct SelectionTargetSpecification {
13 pub mz: f64,
14 pub charge: Option<i32>,
15 pub time_range: Range<f64>,
16}
17
18impl SelectionTargetSpecification {
19 pub fn new(mz: f64, charge: Option<i32>, time_range: Range<f64>) -> Self {
20 Self {
21 mz,
22 charge,
23 time_range,
24 }
25 }
26
27 pub fn spans(&self, time: f64) -> bool {
28 self.time_range.start <= time && self.time_range.end > time
29 }
30
31 pub fn from_spectrum(spectrum: &MultiLayerSpectrum) -> Self {
32 let prec = spectrum.precursor().unwrap().ion().unwrap();
33 Self {
34 mz: prec.mz(),
35 charge: prec.charge(),
36 time_range: (spectrum.start_time()..spectrum.start_time()),
37 }
38 }
39
40 pub fn observe(&mut self, spectrum: &MultiLayerSpectrum) {
41 let t = spectrum.start_time();
42 if t < self.time_range.start {
43 self.time_range.start = t;
44 }
45 if t > self.time_range.end {
46 self.time_range.end = t;
47 }
48 }
49}
50
51impl PartialOrd for SelectionTargetSpecification {
52 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
53 match self.mz.partial_cmp(&other.mz) {
54 Some(Ordering::Equal) => {}
55 ord => return ord,
56 }
57 match self.charge.partial_cmp(&other.charge) {
58 Some(Ordering::Equal) => {}
59 ord => return ord,
60 }
61 match self.time_range.start.partial_cmp(&other.time_range.start) {
62 Some(Ordering::Equal) => {}
63 ord => return ord,
64 }
65 self.time_range.end.partial_cmp(&other.time_range.end)
66 }
67}
68
69#[derive(Default, Debug, Clone, Copy, PartialEq)]
70pub struct SelectedTarget {
71 pub mz: f64,
72 pub charge: Option<i32>,
73}
74
75impl SelectedTarget {
76 pub fn new(mz: f64, charge: Option<i32>) -> Self {
77 Self { mz, charge }
78 }
79}
80
81pub struct MSnTargetTrackingIterator<R: SpectrumSource> {
82 source: SpectrumGroupingIterator<R, CentroidPeak, DeconvolutedPeak, MultiLayerSpectrum>,
83 time_width: f64,
84 error_tolerance: Tolerance,
85 buffer: VecDeque<(SpectrumGroup, f64)>,
86 pushback_buffer: Option<(SpectrumGroup, f64)>,
87 targets: VecDeque<SelectionTargetSpecification>,
88}
89
90impl<R: SpectrumSource> MSnTargetTrackingIterator<R> {
91 pub fn new(
92 source: SpectrumGroupingIterator<R, CentroidPeak, DeconvolutedPeak, MultiLayerSpectrum>,
93 time_width: f64,
94 error_tolerance: Tolerance,
95 ) -> Self {
96 let mut inst = Self {
97 source,
98 time_width,
99 error_tolerance,
100 buffer: Default::default(),
101 pushback_buffer: Default::default(),
102 targets: Default::default(),
103 };
104 inst.initial_feed();
105 inst
106 }
107
108 fn observe(&mut self, group: &SpectrumGroup) {
109 let error_tolerance = self.error_tolerance;
110 let time_width_half = self.time_width / 2.0;
111 let new_targets: usize = group
112 .products
113 .iter()
114 .map(|s| {
115 let prec = s.precursor().unwrap();
116 let ion = prec.ion().unwrap();
117 let mz = ion.mz();
118 let t = s.start_time();
119 let hits: usize = self
120 .targets
121 .iter_mut()
122 .map(|p| {
123 if error_tolerance.test(p.mz, mz)
124 && (t > p.time_range.end)
125 && (t - p.time_range.end) < time_width_half
126 {
127 p.time_range.end = t;
128 1
129 } else {
130 0
131 }
132 })
133 .sum();
134 if hits == 0 {
135 let p = SelectionTargetSpecification::new(
136 mz,
137 prec.ion().and_then(|i| i.charge()),
138 t - 0.5..t + 0.5,
139 );
140 self.targets.push_back(p);
142 1
143 } else {
144 0
145 }
146 })
147 .sum();
148 if new_targets > 0 {
149 log::debug!("Added {new_targets} new targets");
150 }
151 }
152
153 fn initial_feed(&mut self) {
154 let group = self.source.next();
155 let start_time = group
156 .as_ref()
157 .and_then(|s| s.earliest_spectrum().map(|s| s.start_time()))
158 .unwrap();
159 let end_time = start_time + self.time_width;
160 eprintln!("Initial time window {start_time} to {end_time}");
161 if let Some(g) = group {
162 self.observe(&g);
163 self.buffer.push_back((g, start_time));
164 }
165 while let Some(group) = self.source.next() {
166 let t = group.earliest_spectrum().map(|s| s.start_time()).unwrap();
167 if t < end_time {
168 self.observe(&group);
169 self.buffer.push_back((group, t));
170 } else {
171 self.pushback_buffer = Some((group, t));
172 break;
173 }
174 }
175 eprintln!(
176 "{} targets extracted from buffer size {}",
177 self.targets.len(),
178 self.buffer.len()
179 );
180 }
181
182 fn get_current_window_end(&self) -> f64 {
183 if self.buffer.is_empty() {
184 return f64::NEG_INFINITY;
185 }
186 let start = self.buffer.front().map(|(_, t)| *t).unwrap();
187 let end = self.buffer.back().map(|(_, t)| *t).unwrap();
188 let mid = start + (end - start) / 2.0;
189 mid + self.time_width
190 }
191
192 fn feed_next(&mut self) {
193 let threshold = self.get_current_window_end();
194 let (use_pushback, pushback_populated) = if let Some((_, t)) = self.pushback_buffer.as_ref()
195 {
196 (*t < threshold, true)
197 } else {
198 (false, false)
199 };
200 if use_pushback {
201 let (group, t) = self.pushback_buffer.take().unwrap();
202 self.observe(&group);
203 self.buffer.push_back((group, t));
204 }
205 if !pushback_populated {
206 if let Some(group) = self.source.next() {
207 let t = group.earliest_spectrum().map(|s| s.start_time()).unwrap();
208 if t < threshold {
209 self.observe(&group);
210 self.buffer.push_back((group, t));
211 } else {
212 self.pushback_buffer = Some((group, t));
213 }
214 }
215 }
216 }
217
218 fn step(&mut self) -> Option<(SpectrumGroup, Vec<SelectedTarget>)> {
219 if let Some((group, t)) = self.buffer.pop_front() {
220 let targets: Vec<_> = self
221 .targets
222 .iter()
223 .filter(|p| p.spans(t))
224 .map(|p| SelectedTarget::new(p.mz, p.charge))
225 .collect();
226 self.targets = self
227 .targets
228 .drain(..)
229 .filter(|target| {
230 let cond = target.time_range.end >= t;
232 if !cond {
233 log::debug!("Dropping {target:?} at {t}")
234 }
235 cond
236 })
237 .collect();
238 Some((group, targets))
239 } else {
240 None
241 }
242 }
243}
244
245impl<R: SpectrumSource> Iterator for MSnTargetTrackingIterator<R> {
246 type Item = (SpectrumGroup, Vec<SelectedTarget>);
247
248 fn next(&mut self) -> Option<Self::Item> {
249 let value = self.step();
250 self.feed_next();
251 value
252 }
253}
254
255fn main() -> io::Result<()> {
256 let path = path::PathBuf::from(
257 env::args()
258 .nth(1)
259 .expect("Please pass an MS data file path"),
260 );
261 env_logger::init();
262
263 let reader = mzdata::MzMLReader::new(fs::File::open(path)?);
264 let group_iter = reader.into_groups();
265 let tracking_iter = MSnTargetTrackingIterator::new(group_iter, 2.0, Tolerance::PPM(5.0));
266 for (group, targets) in tracking_iter {
267 if group.precursor.is_some() {
268 let prec = group.precursor.as_ref().unwrap();
269 eprintln!("{} {}: {}", prec.id(), prec.start_time(), targets.len());
270 }
271 }
272 Ok(())
273}