1use crate::{Peak, PeakReader, TimsResult, spectrum_reader::TimsCentroidError};
2use rayon::prelude::*;
3use rustc_hash::FxHashMap;
4use std::sync::Arc;
5use std::sync::atomic;
6use timsrust_core::FrameInfo;
7use timsrust_core::utils::reader::ParIterableReader;
8use timsrust_core::utils::reader::Reader;
9use timsrust_core::{
10 Charge, Converter, FrameIndex, Mz, Precursor, Rt, ScanIndex, Spectrum,
11 TofIndex,
12};
13use timsrust_core::{Im, InvertibleConverter};
14
15pub struct NarrowSpectrumReader<
16 IonReader,
17 InfoReader,
18 ImConverter: InvertibleConverter<ScanIndex, Im>,
19 MzConverter: Converter<TofIndex, Mz>,
20> {
21 peak_reader: PeakReader<IonReader, InfoReader>,
22 im_converter: Arc<ImConverter>,
23 mz_converter: Arc<MzConverter>,
24 spectrum_id: atomic::AtomicUsize,
25 precursor_id: atomic::AtomicUsize,
26 min_spectrum_size: usize,
27 monoisotopic_only: bool,
28 highest_charge_state_only: bool,
29 charges: Vec<u8>,
30}
31
32impl<
33 IonReader: timsrust_core::utils::reader::Reader<timsrust_core::FrameIons>
34 + Sync
35 + Send,
36 InfoReader: timsrust_core::utils::reader::Reader<timsrust_core::FrameInfo>
37 + timsrust_core::utils::reader::IndexedReader<timsrust_core::FrameInfo>
38 + Sync
39 + Send,
40 ImConverter: InvertibleConverter<ScanIndex, Im> + Sync + Send,
41 MzConverter: Converter<TofIndex, Mz> + Sync + Send,
42> NarrowSpectrumReader<IonReader, InfoReader, ImConverter, MzConverter>
43{
44 pub fn new(
45 peak_reader: PeakReader<IonReader, InfoReader>,
46 min_spectrum_size: usize,
47 im_converter: Arc<ImConverter>,
48 mz_converter: Arc<MzConverter>,
49 ) -> TimsResult<Self> {
50 let spectrum_id = atomic::AtomicUsize::new(0);
51 let precursor_id = atomic::AtomicUsize::new(0);
52 let result = Self {
53 peak_reader,
54 mz_converter,
55 im_converter,
56 spectrum_id,
57 precursor_id,
58 min_spectrum_size,
59 monoisotopic_only: true,
60 highest_charge_state_only: true,
61 charges: (1..6).collect(),
62 };
63 Ok(result)
64 }
65
66 pub fn set_charges(&mut self, charges: Vec<u8>) {
67 self.charges = charges;
68 }
69
70 pub fn set_monoisotopic_only(&mut self, value: bool) {
71 self.monoisotopic_only = value;
72 }
73
74 pub fn set_highest_charge_state_only(&mut self, value: bool) {
75 self.highest_charge_state_only = value;
76 }
77
78 pub fn charges(&self) -> &[u8] {
79 &self.charges
80 }
81
82 pub fn monoisotopic_only(&self) -> bool {
83 self.monoisotopic_only
84 }
85
86 pub fn highest_charge_state_only(&self) -> bool {
87 self.highest_charge_state_only
88 }
89
90 pub fn min_spectrum_size(&self) -> usize {
91 self.min_spectrum_size
92 }
93
94 pub fn get_spectral_chunk_from_frame(
95 &self,
96 index: usize,
97 ) -> Option<SpectralChunk> {
98 if let Ok(frame) = self
99 .peak_reader
100 .frame_reader()
101 .get_partial_frame_without_ions(index)
102 {
103 if frame.info().ms_level() != timsrust_core::MSLevel::MS1 {
104 return None;
105 }
106 if let Ok(ms1_peaks) = self.peak_reader.get_peaks_from_frame(index)
107 {
108 let precursors = deisotope(
109 ms1_peaks,
110 &frame,
111 self.im_converter.as_ref(),
112 self.mz_converter.as_ref(),
113 &self.precursor_id,
114 self.peak_reader.scan_fwhm(),
115 self.monoisotopic_only,
116 self.highest_charge_state_only,
117 &self.charges,
118 );
119 if precursors.is_empty() {
120 return None;
121 }
122 let mut spectral_chunk = SpectralChunk {
123 peaks: FxHashMap::default(),
124 precursors,
125 };
126 for ms2_frame_index in find_ms2_frames(
127 index,
128 self.peak_reader.frame_reader().info_reader(),
129 ) {
130 if let Ok(mut ms2_peaks) =
131 self.peak_reader.get_peaks_from_frame(ms2_frame_index)
132 {
133 if ms2_peaks.is_empty() {
134 continue;
135 }
136 ms2_peaks.sort_by(|a, b| a.scan.cmp(&b.scan));
137 spectral_chunk.peaks.insert(ms2_frame_index, ms2_peaks);
138 }
139 }
140 return Some(spectral_chunk);
141 }
142 }
143 None
144 }
145
146 pub fn get_spectra_from_frame(
147 &self,
148 index: usize,
149 ) -> Vec<timsrust_core::Spectrum> {
150 match self.get_spectral_chunk_from_frame(index) {
151 Some(chunk) => chunk
152 .peaks
153 .iter()
154 .flat_map(|(ms2_frame_index, ms2_peaks)| {
155 create_spectra_from_ms2_peaks(
156 ms2_peaks,
157 &chunk.precursors,
158 self.peak_reader.scan_fwhm(),
159 &self.spectrum_id,
160 self.min_spectrum_size,
161 self.peak_reader
162 .frame_reader()
163 .get_partial_frame_without_ions(*ms2_frame_index)
164 .expect("Known to exist")
165 .info()
166 .quadrupole_settings(),
167 )
168 })
169 .collect(),
170 None => vec![],
171 }
172 }
173
174 pub fn _par_iter(
175 &self,
176 ) -> impl ParallelIterator<Item = timsrust_core::Spectrum> + '_ {
177 (0..self.frame_count())
178 .into_par_iter()
179 .flat_map(|index| self.get_spectra_from_frame(index))
180 }
181
182 pub(crate) fn len(&self) -> usize {
183 self.spectrum_id.load(atomic::Ordering::Relaxed)
184 }
185
186 pub fn scan_fwhm(&self) -> usize {
187 self.peak_reader.scan_fwhm()
188 }
189
190 pub fn tof_fwhm(&self) -> usize {
191 self.peak_reader.tof_fwhm()
192 }
193
194 pub fn frame_count(&self) -> usize {
195 self.peak_reader.frame_count()
196 }
197
198 pub fn spectrum_id(&self) -> &atomic::AtomicUsize {
199 &self.spectrum_id
200 }
201}
202
203#[derive(Debug)]
204pub struct SpectralChunk {
205 pub peaks: FxHashMap<usize, Vec<Peak>>,
206 pub precursors: Vec<timsrust_core::Precursor>,
207}
208
209impl<
210 'a,
211 IonReader: timsrust_core::utils::reader::Reader<timsrust_core::FrameIons>
212 + Sync
213 + Send,
214 InfoReader: timsrust_core::utils::reader::Reader<timsrust_core::FrameInfo>
215 + timsrust_core::utils::reader::IndexedReader<timsrust_core::FrameInfo>
216 + Sync
217 + Send,
218 ImConverter: InvertibleConverter<ScanIndex, Im> + Sync + Send,
219 MzConverter: Converter<TofIndex, Mz> + Sync + Send,
220> ParIterableReader<'a, Spectrum>
221 for NarrowSpectrumReader<IonReader, InfoReader, ImConverter, MzConverter>
222{
223 type Error = TimsCentroidError;
224
225 fn par_iter(
226 &'a self,
227 ) -> impl ParallelIterator<Item = Result<Spectrum, Self::Error>> {
228 self._par_iter().map(Ok)
229 }
230}
231
232#[allow(clippy::too_many_arguments)]
233fn deisotope(
234 mut peaks: Vec<Peak>,
235 frame: &timsrust_core::Frame,
236 im_converter: impl Converter<ScanIndex, Im>,
237 mz_converter: impl Converter<TofIndex, Mz>,
238 precursor_id: &atomic::AtomicUsize,
239 scan_fwhm: usize,
240 try_monoisotopic_only: bool,
241 highest_charge_state_only: bool,
242 charges: &[u8],
243) -> Vec<timsrust_core::Precursor> {
244 peaks.sort_by_key(|p| p.scan);
245 const ISOTOPE_MASS: f64 = 1.0033548378;
247 const MAX_DELTA_MZ: f64 = 0.01;
248 let mut result = vec![];
249 let mzs = peaks
250 .iter()
251 .map(|p| mz_converter.convert(TofIndex::try_from(p.tof).unwrap()))
252 .collect::<Vec<_>>();
253 let mut lower = 0;
254 let mut upper = 0;
255 for (i, peak) in peaks.iter().enumerate() {
256 while (lower < peaks.len())
257 && (peaks[lower].scan
258 < (peak.scan.saturating_sub(scan_fwhm as u32 / 2)))
259 {
260 lower += 1;
261 }
262 while (upper < peaks.len())
263 && (peaks[upper].scan < (peak.scan + scan_fwhm as u32 / 2))
264 {
265 upper += 1;
266 }
267 let mz = f64::from(mzs[i]);
268 for charge in charges.iter().rev() {
269 if try_monoisotopic_only {
270 let previous_isotope_mz = mz - ISOTOPE_MASS / *charge as f64;
271 if mzs[lower..upper].iter().any(|&x| {
272 (f64::from(x) - previous_isotope_mz).abs() < MAX_DELTA_MZ
273 }) {
274 continue;
275 }
276 }
277 let next_isotope_mz = mz + ISOTOPE_MASS / *charge as f64;
278 if mzs[lower..upper]
279 .iter()
280 .any(|&x| (f64::from(x) - next_isotope_mz).abs() < MAX_DELTA_MZ)
281 {
282 let id = precursor_id.fetch_add(1, atomic::Ordering::Relaxed);
284 let scan = ScanIndex::try_from(peak.scan).unwrap();
285 let precursor = timsrust_core::Precursor::new(
286 Mz::from(mz as f32),
287 im_converter.convert(scan),
288 Rt::from(frame.info().rt_in_seconds()),
289 scan,
290 Some(Charge::try_from(*charge as usize).unwrap()),
291 Some(peak.apex_intensity as f64),
292 id,
293 FrameIndex::try_from(frame.index() as u32).unwrap(),
294 );
295 result.push(precursor);
296 if highest_charge_state_only {
297 break;
298 }
299 }
300 }
301 }
302 result
303}
304
305fn find_ms2_frames(
306 mut frame_index: usize,
307 frame_reader: &impl Reader<FrameInfo>,
308) -> impl Iterator<Item = usize> {
309 std::iter::from_fn(move || {
310 frame_index += 1;
311 frame_reader.get(frame_index).ok().and_then(|frame_info| {
312 if frame_info.ms_level() == timsrust_core::MSLevel::MS2 {
313 Some(frame_index)
314 } else {
315 None
316 }
317 })
318 })
319}
320
321fn create_spectra_from_ms2_peaks(
322 peaks: &[Peak],
323 precursors: &[timsrust_core::Precursor],
324 scan_fwhm: usize,
325 spectrum_id: &atomic::AtomicUsize,
326 min_spectrum_size: usize,
327 quadrupole_settings: &timsrust_core::QuadrupoleSettings,
328) -> Vec<timsrust_core::Spectrum> {
329 assert!(precursors.is_sorted_by(|a, b| a.scan_index() <= b.scan_index()));
330 assert!(peaks.is_sorted_by(|a, b| a.scan <= b.scan));
331 split_peaks(peaks, scan_fwhm, precursors)
332 .map(|(_precursor_id, lower_id, upper_id, scan)| {
333 (
334 precursors[_precursor_id].clone(),
335 &peaks[lower_id..upper_id],
336 scan,
337 )
338 })
339 .filter_map(|(precursor, subpeaks, scan)| {
340 to_spectrum(
341 &precursor,
342 subpeaks,
343 quadrupole_settings,
344 min_spectrum_size,
345 scan,
346 spectrum_id,
347 )
348 })
349 .collect()
350}
351
352fn to_spectrum(
353 precursor: &Precursor,
354 subpeaks: &[Peak],
355 quadrupole_settings: &timsrust_core::QuadrupoleSettings,
356 min_spectrum_size: usize,
357 scan: usize,
358 spectrum_id: &atomic::AtomicUsize,
359) -> Option<timsrust_core::Spectrum> {
360 if subpeaks.len() < min_spectrum_size {
361 return None;
362 }
363 let quad_info = QuadInfo::new(quadrupole_settings, scan);
364 if !quad_info.is_valid_for_precursor(precursor) {
365 return None;
366 }
367 let mut subpeaks = subpeaks.to_vec();
368 subpeaks.sort_by(|a, b| a.tof.cmp(&b.tof));
369 let intensity_values = subpeaks
370 .iter()
371 .map(|p| p.apex_intensity as f32)
372 .collect::<Vec<_>>();
373 let id = spectrum_id.fetch_add(1, atomic::Ordering::Relaxed);
374 let isolation_window = timsrust_core::IsolationWindow::new_from_center(
375 Mz::from(quad_info.isolation_mz),
376 Mz::from(quad_info.isolation_width),
377 quad_info.ce,
378 );
379 let spectrum = timsrust_core::Spectrum::new(
380 intensity_values.into_iter().map(|x| x.into()).collect(),
381 id,
382 Some(precursor.clone()),
383 subpeaks
384 .iter()
385 .map(|p| TofIndex::try_from(p.tof).unwrap())
386 .collect(),
387 isolation_window,
388 );
389 Some(spectrum)
401}
402
403#[derive(Debug)]
404pub struct QuadInfo {
405 pub isolation_mz: f64,
406 pub isolation_width: f64,
407 pub ce: f64,
408}
409
410impl QuadInfo {
411 pub fn new(
412 quadrupole_settings: &timsrust_core::QuadrupoleSettings,
413 scan: usize,
414 ) -> Self {
415 let index = quadrupole_settings
416 .scan_starts
417 .iter()
418 .position(|&s| s > scan)
419 .unwrap_or(quadrupole_settings.len())
420 .max(1)
421 - 1;
422 let isolation_mz =
423 f64::from(quadrupole_settings.isolation_windows[index].center());
424 let isolation_width =
425 f64::from(quadrupole_settings.isolation_windows[index].width())
426 / 2.0;
427 let ce =
428 quadrupole_settings.isolation_windows[index].collision_energy();
429 Self {
430 isolation_mz,
431 isolation_width,
432 ce,
433 }
434 }
435
436 pub fn is_valid_for_precursor(&self, precursor: &Precursor) -> bool {
437 (f64::from(precursor.mz()) >= self.isolation_mz - self.isolation_width)
438 && (f64::from(precursor.mz())
439 <= self.isolation_mz + self.isolation_width)
440 }
441}
442
443pub fn split_peaks<'a>(
444 peaks: &'a [Peak],
445 scan_fwhm: usize,
446 precursors: &'a [timsrust_core::Precursor],
447 ) -> impl Iterator<Item = (usize, usize, usize, usize)> + 'a {
449 assert!(peaks.is_sorted_by(|a, b| a.scan <= b.scan));
450 let mut results = Vec::new();
451 let mut lower_id = 0;
452 let mut upper_id = 0;
453 for (precursor_id, precursor) in precursors.iter().enumerate() {
454 let scan = usize::from(precursor.scan_index());
455 let start = scan.saturating_sub(scan_fwhm / 2) as u32;
456 let end = (scan + scan_fwhm / 2 + 1) as u32;
457 while lower_id < peaks.len() && peaks[lower_id].scan < start {
458 lower_id += 1;
459 }
460 while upper_id < peaks.len() && peaks[upper_id].scan < end {
461 upper_id += 1;
462 }
463 if upper_id > lower_id {
464 results.push((precursor_id, lower_id, upper_id, scan));
466 }
467 }
468 results.into_iter()
469}