Skip to main content

rustdf/data/
dataset.rs

1use crate::data::acquisition::AcquisitionMode;
2use crate::data::handle::{IndexConverter, TimsData, TimsDataLoader};
3use crate::data::meta::{read_global_meta_sql, read_meta_data_sql};
4use mscore::timstof::frame::{RawTimsFrame, TimsFrame};
5use mscore::timstof::slice::TimsSlice;
6
7pub struct TimsDataset {
8    pub loader: TimsDataLoader,
9}
10
11impl TimsDataset {
12    pub fn new(
13        bruker_lib_path: &str,
14        data_path: &str,
15        in_memory: bool,
16        use_bruker_sdk: bool,
17    ) -> Self {
18        // TODO: error handling
19        let global_meta_data = read_global_meta_sql(data_path).unwrap();
20        let meta_data = read_meta_data_sql(data_path).unwrap();
21
22        let scan_max_index = meta_data.iter().map(|x| x.num_scans).max().unwrap() as u32;
23        let im_lower = global_meta_data.one_over_k0_range_lower;
24        let im_upper = global_meta_data.one_over_k0_range_upper;
25
26        let tof_max_index = global_meta_data.tof_max_index;
27        let mz_lower = global_meta_data.mz_acquisition_range_lower;
28        let mz_upper = global_meta_data.mz_acquisition_range_upper;
29
30        let loader = match in_memory {
31            true => TimsDataLoader::new_in_memory(
32                bruker_lib_path,
33                data_path,
34                use_bruker_sdk,
35                scan_max_index,
36                im_lower,
37                im_upper,
38                tof_max_index,
39                mz_lower,
40                mz_upper,
41            ),
42            false => TimsDataLoader::new_lazy(
43                bruker_lib_path,
44                data_path,
45                use_bruker_sdk,
46                scan_max_index,
47                im_lower,
48                im_upper,
49                tof_max_index,
50                mz_lower,
51                mz_upper,
52            ),
53        };
54
55        TimsDataset { loader }
56    }
57
58    /// Create a new dataset with pre-computed ion mobility calibration lookup table.
59    ///
60    /// This enables accurate ion mobility calibration with fast parallel extraction.
61    /// The im_lookup table should be pre-computed using the Bruker SDK.
62    ///
63    /// # Arguments
64    /// * `data_path` - Path to the .d folder
65    /// * `in_memory` - Whether to load all data into memory
66    /// * `bruker_lib_path` - Path to the Bruker SDK shared library; used to
67    ///   derive an accurate m/z calibration. Pass "NO_SDK" (or an empty
68    ///   string) to skip and use the 2-point boundary m/z model.
69    /// * `im_lookup` - Pre-computed scan→1/K0 lookup table from Bruker SDK
70    ///
71    /// # Returns
72    /// A new TimsDataset with LookupIndexConverter
73    pub fn new_with_calibration(
74        data_path: &str,
75        in_memory: bool,
76        bruker_lib_path: &str,
77        im_lookup: Vec<f64>,
78    ) -> Self {
79        let global_meta_data = read_global_meta_sql(data_path).unwrap();
80
81        let tof_max_index = global_meta_data.tof_max_index;
82        let mz_lower = global_meta_data.mz_acquisition_range_lower;
83        let mz_upper = global_meta_data.mz_acquisition_range_upper;
84
85        let loader = match in_memory {
86            true => TimsDataLoader::new_in_memory_with_calibration(
87                data_path,
88                bruker_lib_path,
89                tof_max_index,
90                mz_lower,
91                mz_upper,
92                im_lookup,
93            ),
94            false => TimsDataLoader::new_lazy_with_calibration(
95                data_path,
96                bruker_lib_path,
97                tof_max_index,
98                mz_lower,
99                mz_upper,
100                im_lookup,
101            ),
102        };
103
104        TimsDataset { loader }
105    }
106
107    /// Check if the Bruker SDK is being used for index conversion.
108    /// Returns false for both Simple and Lookup converters (which are thread-safe).
109    pub fn uses_bruker_sdk(&self) -> bool {
110        self.loader.uses_bruker_sdk()
111    }
112}
113
114impl TimsData for TimsDataset {
115    // Get a frame by its id
116    fn get_frame(&self, frame_id: u32) -> TimsFrame {
117        self.loader.get_frame(frame_id)
118    }
119    // Get a raw frame by its id
120    fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame {
121        self.loader.get_raw_frame(frame_id)
122    }
123    // Get a collection of frames by their ids
124    fn get_slice(&self, frame_ids: Vec<u32>, num_threads: usize) -> TimsSlice {
125        self.loader.get_slice(frame_ids, num_threads)
126    }
127    // Get the acquisition mode, DDA or DIA
128    fn get_acquisition_mode(&self) -> AcquisitionMode {
129        self.loader.get_acquisition_mode().clone()
130    }
131    // Get total number of frames in the dataset
132    fn get_frame_count(&self) -> i32 {
133        self.loader.get_frame_count()
134    }
135    // Get the path to the data
136    fn get_data_path(&self) -> &str {
137        &self.loader.get_data_path()
138    }
139}
140
141impl IndexConverter for TimsDataset {
142    fn tof_to_mz(&self, frame_id: u32, tof_values: &Vec<u32>) -> Vec<f64> {
143        self.loader
144            .get_index_converter()
145            .tof_to_mz(frame_id, tof_values)
146    }
147    // convert m/z values to TOF values given a valid data handle and frame id
148    fn mz_to_tof(&self, frame_id: u32, mz_values: &Vec<f64>) -> Vec<u32> {
149        self.loader
150            .get_index_converter()
151            .mz_to_tof(frame_id, mz_values)
152    }
153    // convert inverse mobility values to scan values given a valid data handle and frame id
154    fn scan_to_inverse_mobility(&self, frame_id: u32, scan_values: &Vec<u32>) -> Vec<f64> {
155        self.loader
156            .get_index_converter()
157            .scan_to_inverse_mobility(frame_id, scan_values)
158    }
159    // convert scan values to inverse mobility values given a valid data handle and frame id
160    fn inverse_mobility_to_scan(
161        &self,
162        frame_id: u32,
163        inverse_mobility_values: &Vec<f64>,
164    ) -> Vec<u32> {
165        self.loader
166            .get_index_converter()
167            .inverse_mobility_to_scan(frame_id, inverse_mobility_values)
168    }
169}