1use mscore::data::peptide::PeptideIon;
2use mscore::data::spectrum::{IndexedMzSpectrum, MsType, MzSpectrum};
3use mscore::simulation::annotation::{
4 MzSpectrumAnnotated, TimsFrameAnnotated, TimsSpectrumAnnotated,
5};
6use mscore::timstof::frame::TimsFrame;
7use mscore::timstof::spectrum::TimsSpectrum;
8use rusqlite::Result;
9use std::collections::{BTreeMap, HashSet};
10use std::path::Path;
11
12use crate::sim::containers::{FramesSim, IonSim, PeptidesSim, ScansSim};
13use crate::sim::handle::TimsTofSyntheticsDataHandle;
14use crate::sim::projector::{IntensityStage, MzCoordSpace, RenderedEvent, RenderedSpectrum};
15use crate::sim::scheme::DataMode;
16use rayon::prelude::*;
17
18pub struct TimsTofSyntheticsPrecursorFrameBuilder {
19 pub ions: BTreeMap<u32, Vec<IonSim>>,
20 pub peptides: BTreeMap<u32, PeptidesSim>,
21 pub scans: Vec<ScansSim>,
22 pub frames: Vec<FramesSim>,
23 pub precursor_frame_id_set: HashSet<u32>,
24 pub frame_to_abundances: BTreeMap<u32, (Vec<u32>, Vec<f32>)>,
25 pub peptide_to_ions: BTreeMap<
26 u32,
27 (
28 Vec<f32>,
29 Vec<Vec<u32>>,
30 Vec<Vec<f32>>,
31 Vec<i8>,
32 Vec<MzSpectrum>,
33 ),
34 >,
35 pub frame_to_rt: BTreeMap<u32, f32>,
36 pub scan_to_mobility: BTreeMap<u32, f32>,
37 pub peptide_to_events: BTreeMap<u32, f32>,
38 pub ion_id_to_peptide_charge: BTreeMap<u32, (u32, i8)>,
40}
41
42impl TimsTofSyntheticsPrecursorFrameBuilder {
43 pub fn new(path: &Path) -> Result<Self> {
54 Self::from_source(path, &crate::sim::projector::DistributionSource::Columns)
55 }
56
57 pub fn from_source(
62 path: &Path,
63 source: &crate::sim::projector::DistributionSource,
64 ) -> Result<Self> {
65 let handle = TimsTofSyntheticsDataHandle::new(path)?;
66 let ions = handle.read_ions_with_source(source)?;
67 let peptides = handle.read_peptides_with_source(source)?;
68 let scans = handle.read_scans()?;
69 let frames = handle.read_frames()?;
70
71 Ok(Self::from_entities(ions, peptides, scans, frames))
72 }
73
74 pub fn from_entities(
81 ions: Vec<IonSim>,
82 peptides: Vec<PeptidesSim>,
83 scans: Vec<ScansSim>,
84 frames: Vec<FramesSim>,
85 ) -> Self {
86 let mut ion_id_to_peptide_charge: BTreeMap<u32, (u32, i8)> = BTreeMap::new();
88 for ion in &ions {
89 ion_id_to_peptide_charge.insert(ion.ion_id, (ion.peptide_id, ion.charge));
90 }
91
92 Self {
93 ions: TimsTofSyntheticsDataHandle::build_peptide_to_ion_map(&ions),
94 peptides: TimsTofSyntheticsDataHandle::build_peptide_map(&peptides),
95 scans: scans.clone(),
96 frames: frames.clone(),
97 precursor_frame_id_set: TimsTofSyntheticsDataHandle::build_precursor_frame_id_set(
98 &frames,
99 ),
100 frame_to_abundances: TimsTofSyntheticsDataHandle::build_frame_to_abundances(&peptides),
101 peptide_to_ions: TimsTofSyntheticsDataHandle::build_peptide_to_ions(&ions),
102 frame_to_rt: TimsTofSyntheticsDataHandle::build_frame_to_rt(&frames),
103 scan_to_mobility: TimsTofSyntheticsDataHandle::build_scan_to_mobility(&scans),
104 peptide_to_events: TimsTofSyntheticsDataHandle::build_peptide_to_events(&peptides),
105 ion_id_to_peptide_charge,
106 }
107 }
108
109 pub fn precursor_frame_contributions(&self, frame_id: u32) -> Vec<(i32, MzSpectrum)> {
117 let mut out: Vec<(i32, MzSpectrum)> = Vec::new();
118 let Some((peptide_ids, abundances)) = self.frame_to_abundances.get(&frame_id) else {
119 return out;
120 };
121 out.reserve(peptide_ids.len() * 4);
122 for (peptide_id, abundance) in peptide_ids.iter().zip(abundances.iter()) {
123 let Some((ion_abundances, scan_occurrences, scan_abundances, _, spectra)) =
124 self.peptide_to_ions.get(peptide_id)
125 else {
126 continue;
127 };
128 let total_events = *self.peptide_to_events.get(peptide_id).unwrap();
129 for (index, ion_abundance) in ion_abundances.iter().enumerate() {
130 let scan_occurrence = &scan_occurrences[index];
131 let scan_abundance = &scan_abundances[index];
132 let spectrum = &spectra[index];
133 for (scan, scan_abu) in scan_occurrence.iter().zip(scan_abundance.iter()) {
134 let abundance_factor = abundance * ion_abundance * scan_abu * total_events;
135 out.push((*scan as i32, spectrum.clone() * abundance_factor as f64));
136 }
137 }
138 }
139 out
140 }
141
142 pub fn precursor_scan_spectrum(&self, frame_id: u32) -> MzSpectrum {
156 let specs: Vec<MzSpectrum> = self
157 .precursor_frame_contributions(frame_id)
158 .into_iter()
159 .map(|(_scan, spectrum)| spectrum)
160 .collect();
161 MzSpectrum::from_collection(specs)
162 }
163
164 pub fn precursor_scan_marginal_spectrum(&self, frame_id: u32) -> MzSpectrum {
177 let mut specs: Vec<MzSpectrum> = Vec::new();
178 let Some((peptide_ids, abundances)) = self.frame_to_abundances.get(&frame_id) else {
179 return MzSpectrum::from_collection(specs);
180 };
181 specs.reserve(peptide_ids.len() * 4);
182 for (peptide_id, frame_abundance) in peptide_ids.iter().zip(abundances.iter()) {
183 let Some((ion_abundances, _scan_occurrences, _scan_abundances, _, spectra)) =
184 self.peptide_to_ions.get(peptide_id)
185 else {
186 continue;
187 };
188 let total_events = *self.peptide_to_events.get(peptide_id).unwrap();
189 for (index, ion_abundance) in ion_abundances.iter().enumerate() {
190 let abundance_factor = frame_abundance * ion_abundance * total_events;
192 specs.push(spectra[index].clone() * abundance_factor as f64);
193 }
194 }
195 MzSpectrum::from_collection(specs)
196 }
197
198 pub fn render_precursor_scan(&self, frame_id: u32, data_mode: DataMode) -> RenderedEvent {
214 let rt = *self.frame_to_rt.get(&frame_id).unwrap() as f64;
215 let spectrum = self.precursor_scan_marginal_spectrum(frame_id);
216 RenderedEvent::Scan {
217 ms_level: 1,
218 retention_time_s: rt,
219 isolation: None,
220 spectrum: RenderedSpectrum {
221 mz: (*spectrum.mz).clone(),
222 intensity: (*spectrum.intensity).clone(),
223 coords: MzCoordSpace::Physical,
224 mode: data_mode,
225 detector_applied: false,
226 stage: IntensityStage::Mobility,
227 },
228 }
229 }
230
231 pub fn build_precursor_frame(
241 &self,
242 frame_id: u32,
243 mz_noise_precursor: bool,
244 uniform: bool,
245 precursor_noise_ppm: f64,
246 right_drag: bool,
247 ) -> TimsFrame {
248 let ms_type = if self.precursor_frame_id_set.contains(&frame_id) {
250 MsType::Precursor
251 } else {
252 MsType::Unknown
253 };
254 let rt = *self.frame_to_rt.get(&frame_id).unwrap() as f64;
255
256 let contributions = self.precursor_frame_contributions(frame_id);
261 let mut tims_spectra: Vec<TimsSpectrum> = Vec::with_capacity(contributions.len());
262 for (scan, scaled_spec) in contributions {
263 let mz_spectrum = if mz_noise_precursor {
264 if uniform {
265 scaled_spec.add_mz_noise_uniform(precursor_noise_ppm, right_drag)
266 } else {
267 scaled_spec.add_mz_noise_normal(precursor_noise_ppm)
268 }
269 } else {
270 scaled_spec
271 };
272
273 let scan_mobility = *self.scan_to_mobility.get(&(scan as u32)).unwrap() as f64;
274 let spectrum_len = mz_spectrum.mz.len();
275
276 tims_spectra.push(TimsSpectrum::new(
277 frame_id as i32,
278 scan,
279 rt,
280 scan_mobility,
281 ms_type.clone(),
282 IndexedMzSpectrum::from_mz_spectrum(vec![0; spectrum_len], mz_spectrum),
283 ));
284 }
285
286 if tims_spectra.is_empty() {
294 return TimsFrame::new(frame_id as i32, ms_type, rt, vec![], vec![], vec![], vec![], vec![]);
295 }
296 let tims_frame = TimsFrame::from_tims_spectra(tims_spectra);
297 tims_frame.filter_ranged(0.0, 10000.0, 0, 2000, 0.0, 10.0, 1.0, 1e9, 0, i32::MAX)
298 }
299
300 pub fn build_precursor_frames(
312 &self,
313 frame_ids: Vec<u32>,
314 mz_noise_precursor: bool,
315 uniform: bool,
316 precursor_noise_ppm: f64,
317 right_drag: bool,
318 num_threads: usize,
319 ) -> Vec<TimsFrame> {
320 let pool = rayon::ThreadPoolBuilder::new()
321 .num_threads(num_threads)
322 .build()
323 .unwrap();
324
325 pool.install(|| {
326 let mut tims_frames: Vec<TimsFrame> = Vec::with_capacity(frame_ids.len());
328 unsafe { tims_frames.set_len(frame_ids.len()); }
329
330 frame_ids.par_iter().enumerate().for_each(|(idx, frame_id)| {
331 let frame = self.build_precursor_frame(
332 *frame_id,
333 mz_noise_precursor,
334 uniform,
335 precursor_noise_ppm,
336 right_drag,
337 );
338 unsafe {
339 let ptr = tims_frames.as_ptr() as *mut TimsFrame;
340 std::ptr::write(ptr.add(idx), frame);
341 }
342 });
343
344 tims_frames
345 })
346 }
347
348 pub fn build_precursor_frame_annotated(
349 &self,
350 frame_id: u32,
351 mz_noise_precursor: bool,
352 uniform: bool,
353 precursor_noise_ppm: f64,
354 right_drag: bool,
355 ) -> TimsFrameAnnotated {
356 let ms_type = if self.precursor_frame_id_set.contains(&frame_id) {
358 MsType::Precursor
359 } else {
360 MsType::Unknown
361 };
362 let rt = *self.frame_to_rt.get(&frame_id).unwrap_or(&0.0) as f64;
363
364 let Some((peptide_ids, abundances)) = self.frame_to_abundances.get(&frame_id) else {
366 return TimsFrameAnnotated::new(frame_id as i32, rt, ms_type, vec![], vec![], vec![], vec![], vec![], vec![]);
367 };
368
369 let estimated_capacity = peptide_ids.len() * 4;
371 let mut tims_spectra: Vec<TimsSpectrumAnnotated> = Vec::with_capacity(estimated_capacity);
372
373 for (peptide_id, abundance) in peptide_ids.iter().zip(abundances.iter()) {
374 let Some((ion_abundances, scan_occurrences, scan_abundances, charges, _)) =
376 self.peptide_to_ions.get(peptide_id)
377 else {
378 continue;
379 };
380
381 let total_events = *self.peptide_to_events.get(peptide_id).unwrap();
383 let peptide = self.peptides.get(peptide_id).unwrap();
384
385 for (index, ion_abundance) in ion_abundances.iter().enumerate() {
386 let scan_occurrence = &scan_occurrences[index];
387 let scan_abundance = &scan_abundances[index];
388 let charge = charges[index];
389
390 let ion = PeptideIon::new(
391 peptide.sequence.sequence.clone(),
392 charge as i32,
393 *ion_abundance as f64,
394 Some(*peptide_id as i32),
395 );
396 let spectrum = ion.calculate_isotopic_spectrum_annotated(1e-3, 1e-8, 200, 1e-4);
398
399 for (scan, scan_abu) in scan_occurrence.iter().zip(scan_abundance.iter()) {
400 let abundance_factor = abundance * ion_abundance * scan_abu * total_events;
401 let scaled_spec: MzSpectrumAnnotated = spectrum.clone() * abundance_factor as f64;
402
403 let mz_spectrum = if mz_noise_precursor {
404 if uniform {
405 scaled_spec.add_mz_noise_uniform(precursor_noise_ppm, right_drag)
406 } else {
407 scaled_spec.add_mz_noise_normal(precursor_noise_ppm)
408 }
409 } else {
410 scaled_spec
411 };
412
413 let scan_mobility = *self.scan_to_mobility.get(scan).unwrap() as f64;
415 let spectrum_len = mz_spectrum.mz.len();
416
417 tims_spectra.push(TimsSpectrumAnnotated::new(
418 frame_id as i32,
419 *scan,
420 rt,
421 scan_mobility,
422 ms_type.clone(),
423 vec![0; spectrum_len],
424 mz_spectrum,
425 ));
426 }
427 }
428 }
429
430 if tims_spectra.is_empty() {
434 return TimsFrameAnnotated::new(
435 frame_id as i32, rt, ms_type, vec![], vec![], vec![], vec![], vec![], vec![],
436 );
437 }
438 let tims_frame = TimsFrameAnnotated::from_tims_spectra_annotated(tims_spectra);
439 let filtered_frame = tims_frame.filter_ranged(0.0, 2000.0, 0.0, 2.0, 0, 1000, 1.0, 1e9);
440
441 TimsFrameAnnotated {
442 frame_id: filtered_frame.frame_id,
443 retention_time: filtered_frame.retention_time,
444 ms_type: filtered_frame.ms_type,
445 tof: filtered_frame.tof,
446 mz: filtered_frame.mz,
447 scan: filtered_frame.scan,
448 inv_mobility: filtered_frame.inv_mobility,
449 intensity: filtered_frame.intensity,
450 annotations: filtered_frame
451 .annotations
452 .into_iter()
453 .map(|mut x| {
454 x.contributions.sort_by(|a, b| {
455 a.intensity_contribution
456 .partial_cmp(&b.intensity_contribution)
457 .unwrap()
458 });
459 x
460 })
461 .collect(),
462 }
463 }
464
465 pub fn build_precursor_frames_annotated(
466 &self,
467 frame_ids: Vec<u32>,
468 mz_noise_precursor: bool,
469 uniform: bool,
470 precursor_noise_ppm: f64,
471 right_drag: bool,
472 num_threads: usize,
473 ) -> Vec<TimsFrameAnnotated> {
474 let pool = rayon::ThreadPoolBuilder::new()
475 .num_threads(num_threads)
476 .build()
477 .unwrap();
478
479 pool.install(|| {
480 let mut tims_frames: Vec<TimsFrameAnnotated> = Vec::with_capacity(frame_ids.len());
482 unsafe { tims_frames.set_len(frame_ids.len()); }
483
484 frame_ids.par_iter().enumerate().for_each(|(idx, frame_id)| {
485 let frame = self.build_precursor_frame_annotated(
486 *frame_id,
487 mz_noise_precursor,
488 uniform,
489 precursor_noise_ppm,
490 right_drag,
491 );
492 unsafe {
493 let ptr = tims_frames.as_ptr() as *mut TimsFrameAnnotated;
494 std::ptr::write(ptr.add(idx), frame);
495 }
496 });
497
498 tims_frames
499 })
500 }
501}
502
503#[cfg(test)]
504mod tests {
505 use super::*;
506
507 fn one_frame_builder() -> TimsTofSyntheticsPrecursorFrameBuilder {
513 let mut frame_to_abundances = BTreeMap::new();
514 frame_to_abundances.insert(1u32, (vec![10u32, 11u32], vec![0.8f32, 0.8f32]));
515
516 let mut peptide_to_ions = BTreeMap::new();
517 peptide_to_ions.insert(
518 10u32,
519 (
520 vec![1.0f32],
521 vec![vec![100u32, 101u32]],
522 vec![vec![0.3f32, 0.2f32]], vec![1i8],
524 vec![MzSpectrum::new(vec![500.0], vec![1.0])],
525 ),
526 );
527 peptide_to_ions.insert(
528 11u32,
529 (
530 vec![1.0f32],
531 vec![vec![]], vec![vec![]],
533 vec![1i8],
534 vec![MzSpectrum::new(vec![700.0], vec![1.0])],
535 ),
536 );
537
538 let mut peptide_to_events = BTreeMap::new();
539 peptide_to_events.insert(10u32, 1000.0f32);
540 peptide_to_events.insert(11u32, 1000.0f32);
541
542 let mut frame_to_rt = BTreeMap::new();
543 frame_to_rt.insert(1u32, 60.0f32);
544
545 TimsTofSyntheticsPrecursorFrameBuilder {
546 ions: BTreeMap::new(),
547 peptides: BTreeMap::new(),
548 scans: Vec::new(),
549 frames: Vec::new(),
550 precursor_frame_id_set: HashSet::new(),
551 frame_to_abundances,
552 peptide_to_ions,
553 frame_to_rt,
554 scan_to_mobility: BTreeMap::new(),
555 peptide_to_events,
556 ion_id_to_peptide_charge: BTreeMap::new(),
557 }
558 }
559
560 fn total_intensity(s: &MzSpectrum) -> f64 {
561 s.intensity.iter().sum()
562 }
563
564 #[test]
565 fn astral_marginal_uses_full_mobility_and_recovers_grid_edge_ions() {
566 let b = one_frame_builder();
567
568 let grid = b.precursor_scan_spectrum(1);
571 assert!((total_intensity(&grid) - 400.0).abs() < 1e-2, "grid total {}", total_intensity(&grid));
572 assert_eq!(grid.mz.len(), 1, "grid drops the edge ion (no captured scans)");
573 assert!((grid.mz[0] - 500.0).abs() < 1e-9);
574
575 let marginal = b.precursor_scan_marginal_spectrum(1);
580 assert!((total_intensity(&marginal) - 1600.0).abs() < 1e-2, "marginal total {}", total_intensity(&marginal));
581 assert_eq!(marginal.mz.len(), 2, "marginal keeps both ions");
582 assert!(marginal.mz.iter().any(|&m| (m - 700.0).abs() < 1e-9), "edge ion recovered");
583
584 assert!(total_intensity(&marginal) > total_intensity(&grid));
588
589 let RenderedEvent::Scan { ms_level, isolation, spectrum, .. } =
591 b.render_precursor_scan(1, DataMode::Profile)
592 else {
593 panic!("MS1 render must be a Scan");
594 };
595 assert_eq!(ms_level, 1);
596 assert!(isolation.is_none());
597 assert_eq!(spectrum.intensity, *marginal.intensity);
598 assert_eq!(spectrum.mz, *marginal.mz);
599 }
600}