1use mscore::data::peptide::PeptideProductIonSeriesCollection;
8use mscore::data::spectrum::{IndexedMzSpectrum, MsType, MzSpectrum};
9use mscore::timstof::collision::{TimsTofCollisionEnergy, TimsTofCollisionEnergyDIA};
10use mscore::timstof::frame::TimsFrame;
11use mscore::timstof::quadrupole::{IonTransmission, TimsTransmissionDDA, TimsTransmissionDIA};
12use mscore::timstof::spectrum::TimsSpectrum;
13use std::collections::{BTreeMap, HashSet};
14use std::path::Path;
15use std::sync::Arc;
16
17use rayon::prelude::*;
18
19use crate::sim::containers::{FragmentIonSim, FramesSim, IonSim, PeptidesSim, ScansSim};
20use crate::sim::dda::TimsTofSyntheticsFrameBuilderDDA;
21use crate::sim::handle::TimsTofSyntheticsDataHandle;
22use crate::sim::precursor::TimsTofSyntheticsPrecursorFrameBuilder;
23use crate::sim::projector::DistributionSource;
24
25pub struct TimsTofLazyFrameBuilderDIA {
34 pub db_path: String,
36 pub frames: Vec<FramesSim>,
38 pub scans: Vec<ScansSim>,
40 pub precursor_frame_id_set: HashSet<u32>,
42 pub frame_to_rt: BTreeMap<u32, f32>,
44 pub scan_to_mobility: BTreeMap<u32, f32>,
46 pub transmission_settings: TimsTransmissionDIA,
48 pub fragmentation_settings: TimsTofCollisionEnergyDIA,
50 pub num_threads: usize,
52 pub source: DistributionSource,
54}
55
56impl TimsTofLazyFrameBuilderDIA {
57 pub fn new(path: &Path, num_threads: usize) -> rusqlite::Result<Self> {
71 Self::new_with_source(path, num_threads, DistributionSource::Columns)
72 }
73
74 pub fn new_with_source(
77 path: &Path,
78 num_threads: usize,
79 source: DistributionSource,
80 ) -> rusqlite::Result<Self> {
81 let handle = TimsTofSyntheticsDataHandle::new(path)?;
82 handle
84 .read_prediction_set()?
85 .assert_render_compatible()
86 .map_err(|_| rusqlite::Error::InvalidQuery)?;
87
88 let frames = handle.read_frames()?;
89 let scans = handle.read_scans()?;
90
91 let precursor_frame_id_set = TimsTofSyntheticsDataHandle::build_precursor_frame_id_set(&frames);
92 let frame_to_rt = TimsTofSyntheticsDataHandle::build_frame_to_rt(&frames);
93 let scan_to_mobility = TimsTofSyntheticsDataHandle::build_scan_to_mobility(&scans);
94
95 let transmission_settings = handle.get_transmission_dia();
96 let fragmentation_settings = handle.get_collision_energy_dia();
97
98 Ok(Self {
99 db_path: path.to_str().unwrap().to_string(),
100 frames,
101 scans,
102 precursor_frame_id_set,
103 frame_to_rt,
104 scan_to_mobility,
105 transmission_settings,
106 fragmentation_settings,
107 num_threads,
108 source,
109 })
110 }
111
112 fn load_data_for_frame_range(
116 &self,
117 frame_min: u32,
118 frame_max: u32,
119 ) -> rusqlite::Result<(Vec<PeptidesSim>, Vec<IonSim>, Vec<FragmentIonSim>)> {
120 let path = Path::new(&self.db_path);
121 let handle = TimsTofSyntheticsDataHandle::new(path)?;
122
123 let peptides = handle.read_peptides_for_frame_range_with_source(frame_min, frame_max, &self.source)?;
125
126 if peptides.is_empty() {
127 return Ok((Vec::new(), Vec::new(), Vec::new()));
128 }
129
130 let peptide_ids: Vec<u32> = peptides.iter().map(|p| p.peptide_id).collect();
132
133 let ions = handle.read_ions_for_peptides_with_source(&peptide_ids, &self.source)?;
135 let fragment_ions = handle.read_fragment_ions_for_peptides(&peptide_ids)?;
136
137 Ok((peptides, ions, fragment_ions))
138 }
139
140 pub fn build_frames_lazy(
160 &self,
161 frame_ids: Vec<u32>,
162 fragmentation: bool,
163 mz_noise_precursor: bool,
164 uniform: bool,
165 precursor_noise_ppm: f64,
166 mz_noise_fragment: bool,
167 fragment_noise_ppm: f64,
168 right_drag: bool,
169 ) -> Vec<TimsFrame> {
170 if frame_ids.is_empty() {
171 return Vec::new();
172 }
173
174 let frame_min = *frame_ids.iter().min().unwrap();
176 let frame_max = *frame_ids.iter().max().unwrap();
177
178 let (peptides, ions, fragment_ions) = match self.load_data_for_frame_range(frame_min, frame_max) {
180 Ok(data) => data,
181 Err(_) => return Vec::new(),
182 };
183
184 let peptide_map = TimsTofSyntheticsDataHandle::build_peptide_map(&peptides);
186 let peptide_to_ions = TimsTofSyntheticsDataHandle::build_peptide_to_ions(&ions);
187 let frame_to_abundances = TimsTofSyntheticsDataHandle::build_frame_to_abundances(&peptides);
188 let peptide_to_events = TimsTofSyntheticsDataHandle::build_peptide_to_events(&peptides);
189
190 let fragment_ions_map = if fragmentation {
192 Some(TimsTofSyntheticsDataHandle::build_fragment_ions(
193 &peptide_map,
194 &fragment_ions,
195 self.num_threads,
196 ))
197 } else {
198 None
199 };
200
201 let pool = rayon::ThreadPoolBuilder::new()
203 .num_threads(self.num_threads)
204 .build()
205 .unwrap();
206
207 pool.install(|| {
208 let mut tims_frames: Vec<TimsFrame> = Vec::with_capacity(frame_ids.len());
209 unsafe { tims_frames.set_len(frame_ids.len()); }
210
211 frame_ids.par_iter().enumerate().for_each(|(idx, frame_id)| {
212 let frame = self.build_single_frame(
213 *frame_id,
214 fragmentation,
215 mz_noise_precursor,
216 uniform,
217 precursor_noise_ppm,
218 mz_noise_fragment,
219 fragment_noise_ppm,
220 right_drag,
221 &peptide_map,
222 &peptide_to_ions,
223 &frame_to_abundances,
224 &peptide_to_events,
225 &fragment_ions_map,
226 );
227 unsafe {
228 let ptr = tims_frames.as_ptr() as *mut TimsFrame;
229 std::ptr::write(ptr.add(idx), frame);
230 }
231 });
232
233 tims_frames
234 })
235 }
236
237 #[allow(clippy::too_many_arguments)]
239 fn build_single_frame(
240 &self,
241 frame_id: u32,
242 fragmentation: bool,
243 mz_noise_precursor: bool,
244 uniform: bool,
245 precursor_noise_ppm: f64,
246 mz_noise_fragment: bool,
247 fragment_noise_ppm: f64,
248 right_drag: bool,
249 _peptide_map: &BTreeMap<u32, PeptidesSim>,
250 peptide_to_ions: &BTreeMap<u32, (Vec<f32>, Vec<Vec<u32>>, Vec<Vec<f32>>, Vec<i8>, Vec<MzSpectrum>)>,
251 frame_to_abundances: &BTreeMap<u32, (Vec<u32>, Vec<f32>)>,
252 peptide_to_events: &BTreeMap<u32, f32>,
253 fragment_ions_map: &Option<BTreeMap<(u32, i8, i32), (PeptideProductIonSeriesCollection, Vec<MzSpectrum>)>>,
254 ) -> TimsFrame {
255 let is_precursor = self.precursor_frame_id_set.contains(&frame_id);
257
258 if is_precursor {
259 self.build_precursor_frame(
260 frame_id,
261 mz_noise_precursor,
262 uniform,
263 precursor_noise_ppm,
264 right_drag,
265 peptide_to_ions,
266 frame_to_abundances,
267 peptide_to_events,
268 )
269 } else {
270 self.build_fragment_frame(
271 frame_id,
272 fragmentation,
273 mz_noise_fragment,
274 uniform,
275 fragment_noise_ppm,
276 right_drag,
277 peptide_to_ions,
278 frame_to_abundances,
279 peptide_to_events,
280 fragment_ions_map,
281 )
282 }
283 }
284
285 #[allow(clippy::too_many_arguments)]
287 fn build_precursor_frame(
288 &self,
289 frame_id: u32,
290 mz_noise_precursor: bool,
291 uniform: bool,
292 precursor_noise_ppm: f64,
293 right_drag: bool,
294 peptide_to_ions: &BTreeMap<u32, (Vec<f32>, Vec<Vec<u32>>, Vec<Vec<f32>>, Vec<i8>, Vec<MzSpectrum>)>,
295 frame_to_abundances: &BTreeMap<u32, (Vec<u32>, Vec<f32>)>,
296 peptide_to_events: &BTreeMap<u32, f32>,
297 ) -> TimsFrame {
298 let ms_type = MsType::Precursor;
299 let rt = *self.frame_to_rt.get(&frame_id).unwrap_or(&0.0) as f64;
300
301 let Some((peptide_ids, abundances)) = frame_to_abundances.get(&frame_id) else {
303 return TimsFrame::new(frame_id as i32, ms_type, rt, vec![], vec![], vec![], vec![], vec![]);
304 };
305
306 let estimated_capacity = peptide_ids.len() * 4;
308 let mut tims_spectra: Vec<TimsSpectrum> = Vec::with_capacity(estimated_capacity);
309
310 for (peptide_id, abundance) in peptide_ids.iter().zip(abundances.iter()) {
311 let Some((ion_abundances, scan_occurrences, scan_abundances, _, spectra)) =
312 peptide_to_ions.get(peptide_id)
313 else {
314 continue;
315 };
316
317 let total_events = *peptide_to_events.get(peptide_id).unwrap_or(&1.0);
319
320 for (index, ion_abundance) in ion_abundances.iter().enumerate() {
321 let scan_occurrence = &scan_occurrences[index];
322 let scan_abundance = &scan_abundances[index];
323 let spectrum = &spectra[index];
324
325 for (scan, scan_abu) in scan_occurrence.iter().zip(scan_abundance.iter()) {
326 let abundance_factor = abundance * ion_abundance * scan_abu * total_events;
327 let scaled_spec: MzSpectrum = spectrum.clone() * abundance_factor as f64;
328
329 let mz_spectrum = if mz_noise_precursor {
330 if uniform {
331 scaled_spec.add_mz_noise_uniform(precursor_noise_ppm, right_drag)
332 } else {
333 scaled_spec.add_mz_noise_normal(precursor_noise_ppm)
334 }
335 } else {
336 scaled_spec
337 };
338
339 let scan_mobility = *self.scan_to_mobility.get(scan).unwrap_or(&0.0) as f64;
340 let spectrum_len = mz_spectrum.mz.len();
341
342 tims_spectra.push(TimsSpectrum::new(
343 frame_id as i32,
344 *scan as i32,
345 rt,
346 scan_mobility,
347 ms_type.clone(),
348 IndexedMzSpectrum::from_mz_spectrum(
349 vec![0; spectrum_len],
350 mz_spectrum,
351 ),
352 ));
353 }
354 }
355 }
356
357 let mut filtered = TimsFrame::from_tims_spectra_filtered(
358 tims_spectra, 0.0, 10000.0, 0, 2000, 0.0, 10.0, 1.0, 1e9,
359 );
360
361 let intensities_rounded: Vec<f64> = filtered
363 .ims_frame
364 .intensity
365 .iter()
366 .map(|x| x.round())
367 .collect();
368 filtered.ims_frame.intensity = Arc::new(intensities_rounded);
369
370 filtered
371 }
372
373 #[allow(clippy::too_many_arguments)]
375 fn build_fragment_frame(
376 &self,
377 frame_id: u32,
378 fragmentation: bool,
379 mz_noise_fragment: bool,
380 uniform: bool,
381 fragment_noise_ppm: f64,
382 right_drag: bool,
383 peptide_to_ions: &BTreeMap<u32, (Vec<f32>, Vec<Vec<u32>>, Vec<Vec<f32>>, Vec<i8>, Vec<MzSpectrum>)>,
384 frame_to_abundances: &BTreeMap<u32, (Vec<u32>, Vec<f32>)>,
385 peptide_to_events: &BTreeMap<u32, f32>,
386 fragment_ions_map: &Option<BTreeMap<(u32, i8, i32), (PeptideProductIonSeriesCollection, Vec<MzSpectrum>)>>,
387 ) -> TimsFrame {
388 let ms_type = MsType::FragmentDia;
389 let rt = *self.frame_to_rt.get(&frame_id).unwrap_or(&0.0) as f64;
390
391 if !fragmentation || fragment_ions_map.is_none() {
392 let precursor_frame = self.build_precursor_frame(
394 frame_id,
395 mz_noise_fragment,
396 uniform,
397 fragment_noise_ppm,
398 right_drag,
399 peptide_to_ions,
400 frame_to_abundances,
401 peptide_to_events,
402 );
403 let mut frame = self.transmission_settings.transmit_tims_frame(&precursor_frame, None);
404 frame.ms_type = MsType::FragmentDia;
405 return frame;
406 }
407
408 let fragment_ions = fragment_ions_map.as_ref().unwrap();
409
410 let Some((peptide_ids, frame_abundances)) = frame_to_abundances.get(&frame_id) else {
412 return TimsFrame::new(frame_id as i32, ms_type, rt, vec![], vec![], vec![], vec![], vec![]);
413 };
414
415 let estimated_capacity = peptide_ids.len() * 4;
417 let mut tims_spectra: Vec<TimsSpectrum> = Vec::with_capacity(estimated_capacity);
418
419 for (peptide_id, frame_abundance) in peptide_ids.iter().zip(frame_abundances.iter()) {
420 let Some((ion_abundances, scan_occurrences, scan_abundances, charges, spectra)) =
421 peptide_to_ions.get(peptide_id)
422 else {
423 continue;
424 };
425
426 let total_events = *peptide_to_events.get(peptide_id).unwrap_or(&1.0);
428
429 for (index, ion_abundance) in ion_abundances.iter().enumerate() {
430 let all_scan_occurrence = &scan_occurrences[index];
431 let all_scan_abundance = &scan_abundances[index];
432 let spectrum = &spectra[index];
433 let charge_state = charges[index];
434
435 for (scan, scan_abundance) in all_scan_occurrence.iter().zip(all_scan_abundance.iter()) {
436 if !self.transmission_settings.any_transmitted(
438 frame_id as i32,
439 *scan as i32,
440 &spectrum.mz,
441 None,
442 ) {
443 continue;
444 }
445
446 let fraction_events = frame_abundance * scan_abundance * ion_abundance * total_events;
448
449 let collision_energy = self.fragmentation_settings.get_collision_energy(
451 frame_id as i32,
452 *scan as i32,
453 );
454 let Some(collision_energy_quantized) = crate::sim::handle::resolve_fragment_ce_key(
457 fragment_ions, *peptide_id, charge_state, collision_energy,
458 ) else {
459 if crate::sim::handle::fragment_prefix_exists(fragment_ions, *peptide_id, charge_state) {
460 panic!(
461 "lazy DIA fragment lookup miss: peptide {} charge {} applied CE {:.4} eV \
462 has predicted fragments, but none within 0.1 eV — the prediction set \
463 does not cover this instrument's collision energy",
464 *peptide_id, charge_state, collision_energy,
465 );
466 }
467 continue;
468 };
469 let (_, fragment_series_vec) = fragment_ions
470 .get(&(*peptide_id, charge_state, collision_energy_quantized))
471 .expect("resolve_fragment_ce_key returned a present key");
472
473 let scan_mobility = *self.scan_to_mobility.get(scan).unwrap_or(&0.0) as f64;
475
476 for fragment_ion_series in fragment_series_vec.iter() {
477 let scaled_spec = fragment_ion_series.clone() * fraction_events as f64;
478
479 let mz_spectrum = if mz_noise_fragment {
480 if uniform {
481 scaled_spec.add_mz_noise_uniform(fragment_noise_ppm, right_drag)
482 } else {
483 scaled_spec.add_mz_noise_normal(fragment_noise_ppm)
484 }
485 } else {
486 scaled_spec
487 };
488
489 let spectrum_len = mz_spectrum.mz.len();
490 tims_spectra.push(TimsSpectrum::new(
491 frame_id as i32,
492 *scan as i32,
493 rt,
494 scan_mobility,
495 ms_type.clone(),
496 IndexedMzSpectrum::from_mz_spectrum(
497 vec![0; spectrum_len],
498 mz_spectrum,
499 ).filter_ranged(100.0, 1700.0, 1.0, 1e9),
500 ));
501 }
502 }
503 }
504 }
505
506 if tims_spectra.is_empty() {
507 return TimsFrame::new(frame_id as i32, ms_type, rt, vec![], vec![], vec![], vec![], vec![]);
508 }
509
510 let mut filtered = TimsFrame::from_tims_spectra_filtered(
511 tims_spectra, 100.0, 1700.0, 0, 1000, 0.0, 10.0, 1.0, 1e9,
512 );
513
514 let intensities_rounded: Vec<f64> = filtered
516 .ims_frame
517 .intensity
518 .iter()
519 .map(|x| x.round())
520 .collect();
521 filtered.ims_frame.intensity = Arc::new(intensities_rounded);
522
523 filtered
524 }
525
526 pub fn num_frames(&self) -> usize {
528 self.frames.len()
529 }
530
531 pub fn frame_ids(&self) -> Vec<u32> {
533 self.frames.iter().map(|f| f.frame_id).collect()
534 }
535
536 pub fn precursor_frame_ids(&self) -> Vec<u32> {
538 self.precursor_frame_id_set.iter().cloned().collect()
539 }
540
541 pub fn fragment_frame_ids(&self) -> Vec<u32> {
543 self.frames
544 .iter()
545 .filter(|f| !self.precursor_frame_id_set.contains(&f.frame_id))
546 .map(|f| f.frame_id)
547 .collect()
548 }
549}
550
551impl TimsTofCollisionEnergy for TimsTofLazyFrameBuilderDIA {
552 fn get_collision_energy(&self, frame_id: i32, scan_id: i32) -> f64 {
553 self.fragmentation_settings.get_collision_energy(frame_id, scan_id)
554 }
555}
556
557pub struct TimsTofLazyFrameBuilderDDA {
566 pub db_path: String,
568 pub frames: Vec<FramesSim>,
570 pub scans: Vec<ScansSim>,
572 pub precursor_frame_id_set: HashSet<u32>,
574 pub frame_to_rt: BTreeMap<u32, f32>,
576 pub scan_to_mobility: BTreeMap<u32, f32>,
578 pub transmission_settings: TimsTransmissionDDA,
580 pub num_threads: usize,
582 pub source: DistributionSource,
584}
585
586impl TimsTofLazyFrameBuilderDDA {
587 pub fn new(path: &Path, num_threads: usize) -> rusqlite::Result<Self> {
601 Self::new_with_source(path, num_threads, DistributionSource::Columns)
602 }
603
604 pub fn new_with_source(
607 path: &Path,
608 num_threads: usize,
609 source: DistributionSource,
610 ) -> rusqlite::Result<Self> {
611 let handle = TimsTofSyntheticsDataHandle::new(path)?;
612 handle
614 .read_prediction_set()?
615 .assert_render_compatible()
616 .map_err(|_| rusqlite::Error::InvalidQuery)?;
617
618 let frames = handle.read_frames()?;
619 let scans = handle.read_scans()?;
620
621 let precursor_frame_id_set = TimsTofSyntheticsDataHandle::build_precursor_frame_id_set(&frames);
622 let frame_to_rt = TimsTofSyntheticsDataHandle::build_frame_to_rt(&frames);
623 let scan_to_mobility = TimsTofSyntheticsDataHandle::build_scan_to_mobility(&scans);
624
625 let transmission_settings = handle.get_transmission_dda();
626
627 Ok(Self {
628 db_path: path.to_str().unwrap().to_string(),
629 frames,
630 scans,
631 precursor_frame_id_set,
632 frame_to_rt,
633 scan_to_mobility,
634 transmission_settings,
635 num_threads,
636 source,
637 })
638 }
639
640 fn load_data_for_frame_range(
644 &self,
645 frame_min: u32,
646 frame_max: u32,
647 ) -> rusqlite::Result<(Vec<PeptidesSim>, Vec<IonSim>, Vec<FragmentIonSim>)> {
648 let path = Path::new(&self.db_path);
649 let handle = TimsTofSyntheticsDataHandle::new(path)?;
650
651 let peptides = handle.read_peptides_for_frame_range_with_source(frame_min, frame_max, &self.source)?;
653
654 if peptides.is_empty() {
655 return Ok((Vec::new(), Vec::new(), Vec::new()));
656 }
657
658 let peptide_ids: Vec<u32> = peptides.iter().map(|p| p.peptide_id).collect();
660
661 let ions = handle.read_ions_for_peptides_with_source(&peptide_ids, &self.source)?;
663 let fragment_ions = handle.read_fragment_ions_for_peptides(&peptide_ids)?;
664
665 Ok((peptides, ions, fragment_ions))
666 }
667
668 pub fn build_frames_lazy(
688 &self,
689 frame_ids: Vec<u32>,
690 fragmentation: bool,
691 mz_noise_precursor: bool,
692 uniform: bool,
693 precursor_noise_ppm: f64,
694 mz_noise_fragment: bool,
695 fragment_noise_ppm: f64,
696 right_drag: bool,
697 ) -> Vec<TimsFrame> {
698 if frame_ids.is_empty() {
699 return Vec::new();
700 }
701
702 let frame_min = *frame_ids.iter().min().unwrap();
704 let frame_max = *frame_ids.iter().max().unwrap();
705
706 let (peptides, ions, fragment_ions) = match self.load_data_for_frame_range(frame_min, frame_max) {
708 Ok(data) => data,
709 Err(_) => return Vec::new(),
710 };
711
712 let precursor_builder = TimsTofSyntheticsPrecursorFrameBuilder::from_entities(
719 ions,
720 peptides,
721 self.scans.clone(),
722 self.frames.clone(),
723 );
724 let dda_builder = TimsTofSyntheticsFrameBuilderDDA::from_entities(
725 precursor_builder,
726 self.transmission_settings.clone(),
727 fragment_ions,
728 None,
729 fragmentation,
730 self.num_threads,
731 );
732
733 dda_builder.build_frames(
734 frame_ids,
735 fragmentation,
736 mz_noise_precursor,
737 uniform,
738 precursor_noise_ppm,
739 mz_noise_fragment,
740 fragment_noise_ppm,
741 right_drag,
742 self.num_threads,
743 )
744 }
745
746 pub fn get_collision_energy(&self, frame_id: i32, scan_id: i32) -> f64 {
748 self.transmission_settings.get_collision_energy(frame_id, scan_id).unwrap_or(0.0)
749 }
750
751 pub fn num_frames(&self) -> usize {
753 self.frames.len()
754 }
755
756 pub fn frame_ids(&self) -> Vec<u32> {
758 self.frames.iter().map(|f| f.frame_id).collect()
759 }
760
761 pub fn precursor_frame_ids(&self) -> Vec<u32> {
763 self.precursor_frame_id_set.iter().cloned().collect()
764 }
765
766 pub fn fragment_frame_ids(&self) -> Vec<u32> {
768 self.frames
769 .iter()
770 .filter(|f| !self.precursor_frame_id_set.contains(&f.frame_id))
771 .map(|f| f.frame_id)
772 .collect()
773 }
774}