Skip to main content

msn_target_mapping/
msn_target_mapping.rs

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                    // eprintln!("Added {p:?}");
141                    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                    // Keep targets which have not ended by this time point
231                    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}