mzdata_bindata/
map.rs

1use std::borrow::Cow;
2use std::cmp::Ordering;
3use std::collections::HashMap;
4use std::collections::hash_map::{Iter, IterMut};
5use std::convert::TryFrom;
6
7#[cfg(feature = "parallelism")]
8use rayon::prelude::*;
9
10use mzpeaks::Tolerance;
11
12use mzdata_param::Unit;
13
14use super::BinaryDataArrayType;
15use super::array::DataArray;
16use super::encodings::{ArrayRetrievalError, ArrayType, BinaryCompressionType};
17use super::traits::{ByteArrayView, ByteArrayViewMut};
18
19/// A collection of [`DataArray`]s that are identified by name.
20///
21/// # Examples
22/// ```
23/// use mzdata_bindata::{ArrayType, BinaryArrayMap, BinaryDataArrayType, DataArray};
24///
25/// let mut map = BinaryArrayMap::new();
26/// assert!(!map.has_array(&ArrayType::MZArray));
27///
28/// let array = DataArray::from_name_and_type(&ArrayType::MZArray, BinaryDataArrayType::Float64);
29/// map.add(array);
30///
31/// assert!(map.has_array(&ArrayType::MZArray));
32/// assert_eq!(map.len(), 1);
33/// ```
34#[derive(Debug, Default, Clone)]
35#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
36pub struct BinaryArrayMap {
37    pub byte_buffer_map: HashMap<ArrayType, DataArray>,
38}
39
40impl BinaryArrayMap {
41    pub fn new() -> BinaryArrayMap {
42        BinaryArrayMap {
43            ..Default::default()
44        }
45    }
46
47    /// Get the number of arrays in the map
48    pub fn len(&self) -> usize {
49        self.byte_buffer_map.len()
50    }
51
52    pub fn is_empty(&self) -> bool {
53        self.byte_buffer_map.is_empty()
54    }
55
56    /// Check if there is an ion mobility array present
57    pub fn has_ion_mobility(&self) -> bool {
58        self.byte_buffer_map.keys().any(|a| a.is_ion_mobility())
59    }
60
61    /// Iterate over references to the key-value pairs of this map
62    pub fn iter(&self) -> Iter<'_, ArrayType, DataArray> {
63        self.byte_buffer_map.iter()
64    }
65
66    #[cfg(feature = "parallelism")]
67    pub fn par_iter(&self) -> rayon::collections::hash_map::Iter<'_, ArrayType, DataArray> {
68        self.byte_buffer_map.par_iter()
69    }
70
71    /// Iterate over mutable references to the key-value pairs of this map
72    pub fn iter_mut(&mut self) -> IterMut<'_, ArrayType, DataArray> {
73        self.byte_buffer_map.iter_mut()
74    }
75
76    #[cfg(feature = "parallelism")]
77    pub fn par_iter_mut(
78        &mut self,
79    ) -> rayon::collections::hash_map::IterMut<'_, ArrayType, DataArray> {
80        self.byte_buffer_map.par_iter_mut()
81    }
82
83    /// Compress a specific [`DataArray`] with the `compression` scheme provided, if it is present.
84    ///
85    /// This method may fail if the compression fails or if the array type is missing.
86    ///
87    /// ## See Also
88    /// [`DataArray::store_compressed`]
89    pub fn encode_array(
90        &mut self,
91        array_type: &ArrayType,
92        compression: BinaryCompressionType,
93    ) -> Result<(), ArrayRetrievalError> {
94        if let Some(arr) = self.get_mut(array_type) {
95            arr.store_compressed(compression)?;
96            Ok(())
97        } else {
98            Err(ArrayRetrievalError::NotFound(array_type.clone()))
99        }
100    }
101
102    /// Decode all [`DataArray`] in this map if they have not been decoded already so
103    /// that they are ready for use. If there are many arrays and the `parallelism` feature
104    /// is enabled, arrays may be decoded in parallel.
105    ///
106    /// # Examples
107    /// ```
108    /// use mzdata_bindata::{ArrayType, BinaryArrayMap, BinaryCompressionType, BinaryDataArrayType, DataArray};
109    ///
110    /// let mut mz_array = DataArray::from_name_and_type(&ArrayType::MZArray, BinaryDataArrayType::Float64);
111    /// mz_array.extend(&[1.0f64, 2.0, 3.0]).unwrap();
112    /// mz_array.store_compressed(BinaryCompressionType::Zlib).unwrap();
113    ///
114    /// let mut map = BinaryArrayMap::new();
115    /// map.add(mz_array);
116    /// assert_eq!(map.get(&ArrayType::MZArray).unwrap().compression, BinaryCompressionType::Zlib);
117    ///
118    /// map.decode_all_arrays().unwrap();
119    /// assert_eq!(map.get(&ArrayType::MZArray).unwrap().compression, BinaryCompressionType::Decoded);
120    /// ```
121    pub fn decode_all_arrays(&mut self) -> Result<(), ArrayRetrievalError> {
122        #[cfg(not(feature = "parallelism"))]
123        {
124            self._decode_all_arrays()
125        }
126        #[cfg(feature = "parallelism")]
127        {
128            if self.len() > 2 {
129                self._decode_all_arrays_parallel()
130            } else {
131                self._decode_all_arrays()
132            }
133        }
134    }
135
136    fn _decode_all_arrays(&mut self) -> Result<(), ArrayRetrievalError> {
137        for (_key, value) in self.iter_mut() {
138            match value.compression {
139                BinaryCompressionType::Decoded => {}
140                _ => {
141                    value.decode_and_store()?;
142                }
143            }
144        }
145        Ok(())
146    }
147
148    #[cfg(feature = "parallelism")]
149    fn _decode_all_arrays_parallel(&mut self) -> Result<(), ArrayRetrievalError> {
150        let res: Result<(), ArrayRetrievalError> = self
151            .iter_mut()
152            .par_bridge()
153            .map(|(_key, value)| {
154                match value.compression {
155                    BinaryCompressionType::Decoded => {}
156                    _ => {
157                        value.decode_and_store()?;
158                    }
159                }
160                Ok(())
161            })
162            .collect::<Result<(), ArrayRetrievalError>>();
163        res
164    }
165
166    /// Decode a specific [`DataArray`] if it is present.
167    ///
168    /// This method may fail if decoding fails or if the array type is missing.
169    ///
170    /// See [`DataArray::decode_and_store`]
171    pub fn decode_array(&mut self, array_type: &ArrayType) -> Result<(), ArrayRetrievalError> {
172        if let Some(array) = self.get_mut(array_type) {
173            array.decode_and_store()?;
174            Ok(())
175        } else {
176            Err(ArrayRetrievalError::NotFound(array_type.clone()))
177        }
178    }
179
180    /// Add a [`DataArray`] to the map by its [`ArrayType`] name
181    pub fn add(&mut self, array: DataArray) {
182        self.byte_buffer_map.insert(array.name.clone(), array);
183    }
184
185    /// Get a reference to a specific [`DataArray`] if present
186    pub fn get(&self, array_type: &ArrayType) -> Option<&DataArray> {
187        self.byte_buffer_map.get(array_type)
188    }
189
190    /// Get a mutable reference to a specific [`DataArray`] if present
191    pub fn get_mut(&mut self, array_type: &ArrayType) -> Option<&mut DataArray> {
192        self.byte_buffer_map.get_mut(array_type)
193    }
194
195    /// Check whether a specific [`ArrayType`] is present
196    pub fn has_array(&self, array_type: &ArrayType) -> bool {
197        self.byte_buffer_map.contains_key(array_type)
198    }
199
200    /// Clear the map, discarding any array data
201    pub fn clear(&mut self) {
202        self.byte_buffer_map.clear();
203    }
204
205    /// Search for a specific m/z within a specified [`Tolerance`].
206    ///
207    /// This effectively calls [`BinaryArrayMap::mzs`] to get the m/z array and then
208    /// performs a binary search to find the best index within range, if one exists.
209    pub fn search(&self, query: f64, error_tolerance: Tolerance) -> Option<usize> {
210        if let Ok(mzs) = self.mzs() {
211            let (lower, _upper) = error_tolerance.bounds(query);
212            match mzs[..].binary_search_by(|m| m.partial_cmp(&lower).unwrap()) {
213                Ok(i) => {
214                    let mut best_error = error_tolerance.call(query, mzs[i]).abs();
215                    let mut best_index = i;
216                    let mut index = i + 1;
217                    while index < mzs.len() {
218                        let error = error_tolerance.call(query, mzs[index]).abs();
219                        if error < best_error {
220                            best_index = index;
221                            best_error = error;
222                        }
223                        index += 1;
224                    }
225                    if best_error < error_tolerance.tol() {
226                        return Some(best_index);
227                    }
228                    None
229                }
230                Err(_err) => None,
231            }
232        } else {
233            None
234        }
235    }
236
237    /// Get a reference to the m/z array if it is present
238    ///
239    /// # Examples
240    /// ```
241    /// use mzdata_bindata::{ArrayType, BinaryArrayMap, BinaryDataArrayType, DataArray};
242    ///
243    /// let mut mz_array = DataArray::from_name_and_type(&ArrayType::MZArray, BinaryDataArrayType::Float64);
244    /// mz_array.extend(&[100.0f64, 200.0, 300.0]).unwrap();
245    ///
246    /// let mut map = BinaryArrayMap::new();
247    /// map.add(mz_array);
248    ///
249    /// assert_eq!(&*map.mzs().unwrap(), &[100.0f64, 200.0, 300.0][..]);
250    /// ```
251    pub fn mzs(&'_ self) -> Result<Cow<'_, [f64]>, ArrayRetrievalError> {
252        let mz_array = self
253            .get(&ArrayType::MZArray)
254            .ok_or(ArrayRetrievalError::NotFound(ArrayType::MZArray))?
255            .to_f64()
256            .inspect_err(|e| log::error!("Failed to decode m/z array: {e}"))?;
257        Ok(mz_array)
258    }
259
260    /// Re-order all arrays of matching length with respect to the provided index `mask` operating
261    /// on the decoded native value states.
262    ///
263    /// The mask would be created using an indirect sort from another array:
264    /// ```notest
265    /// let vals = self.mzs().unwrap();
266    /// let n = vals.len();
267    /// let mut mask: Vec<usize> = (0..n).into_iter().collect();
268    /// mask.sort_by(|i, j| {
269    ///     let a = vals[*i];
270    ///     let b = vals[*j];
271    ///     a.total_cmp(&b)
272    /// });
273    /// drop(vals);
274    /// self.sort_from_indices(mask)
275    /// ```
276    pub fn sort_from_indices(&mut self, mut mask: Vec<usize>) -> Result<(), ArrayRetrievalError> {
277        let n = mask.len();
278        const TOMBSTONE: usize = usize::MAX;
279        for idx in 0..n {
280            if mask[idx] != TOMBSTONE {
281                let mut current_idx = idx;
282                loop {
283                    let next_idx = mask[current_idx];
284                    mask[current_idx] = TOMBSTONE;
285                    if mask[next_idx] == TOMBSTONE {
286                        break;
287                    }
288                    for (_, v) in self.iter_mut() {
289                        if v.data_len()? != n {
290                            continue;
291                        }
292                        match v.dtype {
293                            BinaryDataArrayType::Float64 => {
294                                let view = v.coerce_mut::<f64>()?;
295                                view.swap(current_idx, next_idx);
296                            }
297                            BinaryDataArrayType::Float32 => {
298                                let view = v.coerce_mut::<f32>()?;
299                                view.swap(current_idx, next_idx);
300                            }
301                            BinaryDataArrayType::Int64 => {
302                                let view = v.coerce_mut::<i64>()?;
303                                view.swap(current_idx, next_idx);
304                            }
305                            BinaryDataArrayType::Int32 => {
306                                let view = v.coerce_mut::<i32>()?;
307                                view.swap(current_idx, next_idx);
308                            }
309                            BinaryDataArrayType::ASCII => todo!(),
310                            BinaryDataArrayType::Unknown => todo!(),
311                        }
312                    }
313                    current_idx = next_idx;
314                }
315            }
316        }
317        Ok(())
318    }
319
320    /// Sort all the arrays in the map by `name` if they are the same length. This operates
321    /// on the decoded state directly.
322    ///
323    /// This builds the sorting mask and then leaves [`Self::sort_from_indices`] to perform
324    /// the re-ordering.
325    ///
326    /// Not supported for [`BinaryDataArrayType::ASCII`] and [`BinaryDataArrayType::Unknown`].
327    pub fn sort_by_array(&mut self, name: &ArrayType) -> Result<(), ArrayRetrievalError> {
328        let query_axis = self
329            .get(name)
330            .ok_or_else(|| ArrayRetrievalError::NotFound(name.clone()))?;
331        macro_rules! sort_mask {
332            ($conv:expr, $cmp:expr) => {{
333                let vals = $conv?;
334                if vals.is_sorted() {
335                    return Ok(());
336                }
337                let n = vals.len();
338                let mut mask: Vec<usize> = (0..n).into_iter().collect();
339                mask.sort_by(|i, j| {
340                    let a = vals[*i];
341                    let b = vals[*j];
342                    $cmp(&a, &b)
343                });
344                (mask, n)
345            }};
346        }
347        let (mask, _) = match query_axis.dtype() {
348            BinaryDataArrayType::Float64 => {
349                sort_mask!(query_axis.to_f64(), f64::total_cmp)
350            }
351            BinaryDataArrayType::Float32 => {
352                sort_mask!(query_axis.to_f32(), f32::total_cmp)
353            }
354            BinaryDataArrayType::Int64 => {
355                sort_mask!(query_axis.to_i64(), i64::cmp)
356            }
357            BinaryDataArrayType::Int32 => {
358                sort_mask!(query_axis.to_i32(), i32::cmp)
359            }
360            BinaryDataArrayType::ASCII => todo!(),
361            BinaryDataArrayType::Unknown => todo!(),
362        };
363
364        self.sort_from_indices(mask)
365    }
366
367    /// Get a mutable reference to the m/z array if it is present.
368    ///
369    /// This method obtains the m/z array, calls [`DataArray::decode_and_store`] to force
370    /// the data to be stored natively in memory, and forces the storage to be
371    /// [`BinaryDataArrayType::Float64`] if not already. Finally it returns a mutable
372    /// view of that decoded dataas `&mut [f64]`
373    pub fn mzs_mut(&mut self) -> Result<&mut [f64], ArrayRetrievalError> {
374        if let Some(mz_array) = self.get_mut(&ArrayType::MZArray) {
375            mz_array
376                .decode_and_store()
377                .inspect_err(|e| log::error!("Failed to decode m/z array: {e}"))?;
378            mz_array.store_as(BinaryDataArrayType::Float64)?;
379            mz_array.coerce_mut()
380        } else {
381            Err(ArrayRetrievalError::NotFound(ArrayType::MZArray))
382        }
383    }
384
385    /// Get a reference to the intensity array if it is present
386    pub fn intensities(&'_ self) -> Result<Cow<'_, [f32]>, ArrayRetrievalError> {
387        let intensities = self
388            .get(&ArrayType::IntensityArray)
389            .ok_or(ArrayRetrievalError::NotFound(ArrayType::IntensityArray))?
390            .to_f32()
391            .inspect_err(|e| log::error!("Failed to decode intensity array: {e:?}"))?;
392        Ok(intensities)
393    }
394
395    /// Get a mutable reference to the intensity array if it is present.
396    ///
397    /// See [`BinaryArrayMap::mzs_mut`] for an explanation of the intermediate steps,
398    /// save that the forced storage type is [`BinaryArrayDataType::Float32`].
399    pub fn intensities_mut(&mut self) -> Result<&mut [f32], ArrayRetrievalError> {
400        if let Some(arr) = self.get_mut(&ArrayType::IntensityArray) {
401            arr
402                .decode_and_store()
403                .inspect_err(|e| log::error!("Failed to decode intensity array: {e}"))?;
404            arr
405                .store_as(BinaryDataArrayType::Float32)
406                .inspect_err(|e| log::error!("Failed to decode intensity array: {e}"))?;
407            arr.coerce_mut()
408        } else {
409            Err(ArrayRetrievalError::NotFound(ArrayType::IntensityArray))
410        }
411    }
412
413    /// Get a reference to the charge array if it is present
414    pub fn charges(&'_ self) -> Result<Cow<'_, [i32]>, ArrayRetrievalError> {
415        match self.get(&ArrayType::ChargeArray) {
416            Some(data_array) => data_array.to_i32(),
417            None => Err(ArrayRetrievalError::NotFound(ArrayType::ChargeArray)),
418        }
419    }
420
421    /// Get a mutable reference to the charge array if it is present.
422    ///
423    /// See [`BinaryArrayMap::mzs_mut`] for an explanation of the intermediate steps,
424    /// save that the forced storage type is [`BinaryArrayDataType::Int32`].
425    pub fn charge_mut(&mut self) -> Result<&mut [i32], ArrayRetrievalError> {
426        if let Some(arr) = self.get_mut(&ArrayType::ChargeArray) {
427            arr.decode_and_store()?;
428            arr
429                .store_as(BinaryDataArrayType::Int32)
430                .inspect_err(|e| log::error!("Failed to decode charge array: {e}"))?;
431            arr.coerce_mut()
432        } else {
433            Err(ArrayRetrievalError::NotFound(ArrayType::ChargeArray))
434        }
435    }
436
437    /// Get a reference to the ion mobility array if it is present
438    pub fn ion_mobility(&self) -> Result<(Cow<'_, [f64]>, ArrayType), ArrayRetrievalError> {
439        if let Some((array_type, data_array)) = self
440            .byte_buffer_map
441            .iter()
442            .find(|(a, _)| a.is_ion_mobility())
443        {
444            Ok((
445                data_array
446                    .to_f64()
447                    .inspect_err(|e| log::error!("Failed to decode ion mobility array: {e}"))?,
448                array_type.clone(),
449            ))
450        } else {
451            Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray))
452        }
453    }
454
455    /// Get a mutable reference to the ion mobility array if it is present.
456    ///
457    /// See [`BinaryArrayMap::mzs_mut`] for an explanation of the intermediate steps.
458    pub fn ion_mobility_mut(&mut self) -> Result<(&mut [f64], ArrayType), ArrayRetrievalError> {
459        if let Some((array_type, data_array)) = self
460            .byte_buffer_map
461            .iter_mut()
462            .find(|(a, _)| a.is_ion_mobility())
463        {
464            data_array
465                .decode_and_store()
466                .inspect_err(|e| log::error!("Failed to decode ion mobility array: {e}"))?;
467            data_array.store_as(BinaryDataArrayType::Float32)?;
468            Ok((data_array.coerce_mut()?, array_type.clone()))
469        } else {
470            Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray))
471        }
472    }
473
474    /// Split an array map along the ion mobility dimension if it is present.
475    ///
476    /// This function will fail if there is no ion mobility dimension.
477    pub fn stack_ion_mobility(self) -> Result<BinaryArrayMap3D, ArrayRetrievalError> {
478        BinaryArrayMap3D::try_from(self)
479    }
480}
481
482impl IntoIterator for BinaryArrayMap {
483    type Item = (ArrayType, DataArray);
484
485    type IntoIter = <HashMap<ArrayType, DataArray> as IntoIterator>::IntoIter;
486
487    fn into_iter(self) -> Self::IntoIter {
488        self.byte_buffer_map.into_iter()
489    }
490}
491
492#[cfg(feature = "mzsignal")]
493mod mzsignal_impl {
494    use super::*;
495
496    use crate::to_bytes;
497    use mzsignal::ArrayPair;
498
499    impl From<ArrayPair<'_>> for BinaryArrayMap {
500        fn from(value: ArrayPair<'_>) -> Self {
501            let mz_array = DataArray::wrap(
502                &ArrayType::MZArray,
503                BinaryDataArrayType::Float64,
504                to_bytes(&value.mz_array),
505            );
506
507            let intensity_array = DataArray::wrap(
508                &ArrayType::IntensityArray,
509                BinaryDataArrayType::Float32,
510                to_bytes(&value.intensity_array),
511            );
512
513            let mut array_map = BinaryArrayMap::new();
514            array_map.add(mz_array);
515            array_map.add(intensity_array);
516            array_map
517        }
518    }
519}
520
521#[derive(Debug, Default, Clone, Copy, PartialEq)]
522#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
523struct NonNaNF64(f64);
524
525impl std::hash::Hash for NonNaNF64 {
526    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
527        self.0.to_bits().hash(state);
528    }
529}
530
531impl From<f64> for NonNaNF64 {
532    fn from(value: f64) -> Self {
533        Self::wrap(value)
534            .unwrap_or_else(|| panic!("Expected an order-able f64 value, but found {}", value))
535    }
536}
537
538impl NonNaNF64 {
539    fn wrap(value: f64) -> Option<Self> {
540        if value.is_nan() {
541            None
542        } else {
543            Some(Self(value))
544        }
545    }
546}
547
548impl PartialOrd for NonNaNF64 {
549    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
550        Some(self.cmp(other))
551    }
552}
553
554impl Eq for NonNaNF64 {}
555
556impl Ord for NonNaNF64 {
557    fn cmp(&self, other: &Self) -> Ordering {
558        self.0.total_cmp(&other.0)
559    }
560}
561
562macro_rules! _populate_stacked_array_from {
563    ($im_dim:ident, $view:ident, $index_map:ident, $array_bins:ident, $array_type:ident, $array:ident) => {
564        for (i_im, im) in $im_dim.iter() {
565            let v = $view[*i_im];
566            let i_axis = $index_map[&NonNaNF64(*im)];
567            let bin = &mut $array_bins[i_axis];
568            if let Some(bin_array) = bin.get_mut($array_type) {
569                bin_array.push(v)?;
570            } else {
571                let mut bin_array = DataArray::from_name_and_type($array_type, $array.dtype());
572                *bin_array.unit_mut() = $array.unit();
573                bin_array.push(v)?;
574                bin.add(bin_array);
575            }
576        }
577    };
578}
579
580/// Represent a set of [`BinaryArrayMap`] that has been split across an
581/// ion mobility dimension.
582///
583/// This effectively removes an ion mobility [`DataArray`] and for each unique
584/// value in that array, creates a new [`BinaryArrayMap`] with subsets of each
585/// other [`DataArray`] that co-occur at the same ion mobility coordinate.
586#[derive(Debug, Default, Clone)]
587#[cfg_attr(feature = "serde", serde_with::serde_as)]
588#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
589pub struct BinaryArrayMap3D {
590    pub ion_mobility_dimension: Vec<f64>,
591    pub ion_mobility_type: ArrayType,
592    pub ion_mobility_unit: Unit,
593    pub arrays: Vec<BinaryArrayMap>,
594    pub additional_arrays: BinaryArrayMap,
595    #[cfg_attr(feature = "serde", serde_as(as = "Vec<(_, _)>"))]
596    ion_mobility_index: HashMap<NonNaNF64, usize>,
597}
598
599impl BinaryArrayMap3D {
600    /// Create an empty [`BinaryArrayMap3D`] with a defined ion mobility dimension
601    /// but no content at each ion mobility point.
602    ///
603    /// # Panics
604    /// If the ion mobility dimension contains a NaN value
605    pub fn from_ion_mobility_dimension(
606        ion_mobility_dimension: Vec<f64>,
607        ion_mobility_type: ArrayType,
608        ion_mobility_unit: Unit,
609    ) -> BinaryArrayMap3D {
610        let ion_mobility_index: HashMap<_, _> = ion_mobility_dimension
611            .iter()
612            .copied()
613            .enumerate()
614            .map(|(i, f)| {
615                (
616                    NonNaNF64::wrap(f)
617                        .unwrap_or_else(|| panic!("Expected non-NaN value for ion mobility")),
618                    i,
619                )
620            })
621            .collect();
622        let mut arrays = Vec::new();
623        arrays.resize_with(ion_mobility_dimension.len(), BinaryArrayMap::default);
624        Self {
625            ion_mobility_dimension,
626            ion_mobility_type,
627            ion_mobility_unit,
628            arrays,
629            ion_mobility_index,
630            additional_arrays: Default::default(),
631        }
632    }
633
634    /// Construct a [`BinaryArrayMap3D`] from a defined ion mobility dimension and
635    /// a [`Vec`] of [`BinaryArrayMap`] at each point.
636    ///
637    /// # Panics
638    /// If the ion mobility dimension contains a NaN value or `ion_mobility_dimension.len() != arrays.len()`
639    pub fn from_ion_mobility_dimension_and_arrays(
640        ion_mobility_dimension: Vec<f64>,
641        ion_mobility_type: ArrayType,
642        ion_mobility_unit: Unit,
643        arrays: Vec<BinaryArrayMap>,
644    ) -> BinaryArrayMap3D {
645        let ion_mobility_index: HashMap<_, _> = ion_mobility_dimension
646            .iter()
647            .copied()
648            .enumerate()
649            .map(|(i, f)| {
650                (
651                    NonNaNF64::wrap(f)
652                        .unwrap_or_else(|| panic!("Expected non-NaN value for ion mobility")),
653                    i,
654                )
655            })
656            .collect();
657
658        assert_eq!(
659            ion_mobility_dimension.len(),
660            arrays.len(),
661            "The number of ion mobility points does not match the number of point-level maps!"
662        );
663
664        Self {
665            ion_mobility_dimension,
666            ion_mobility_type,
667            ion_mobility_unit,
668            arrays,
669            ion_mobility_index,
670            additional_arrays: Default::default(),
671        }
672    }
673
674    /// Get the associated arrays at the requested ion mobility, if they exist
675    pub fn get_ion_mobility(&self, ion_mobility: f64) -> Option<&BinaryArrayMap> {
676        if let Some(i) = NonNaNF64::wrap(ion_mobility) {
677            if let Some(i) = self.ion_mobility_index.get(&i) {
678                self.arrays.get(*i)
679            } else {
680                None
681            }
682        } else {
683            None
684        }
685    }
686
687    /// Perform a binary search on the ion mobility dimension and find the nearest
688    /// point to `ion_mobility` within `error_tolerance` absolute distance.
689    ///
690    /// Returns the [`BinaryArrayMap`] at that point and the distance from `ion_mobility`
691    /// matched.
692    pub fn search_ion_mobility(
693        &self,
694        ion_mobility: f64,
695        error_tolerance: f64,
696    ) -> Option<(&BinaryArrayMap, f64)> {
697        match self
698            .ion_mobility_dimension
699            .binary_search_by(|x: &f64| x.total_cmp(&ion_mobility))
700        {
701            Ok(i) => {
702                let delta = ion_mobility - self.ion_mobility_dimension[i];
703                if delta.abs() <= error_tolerance {
704                    self.arrays.get(i).map(|a| (a, delta))
705                } else {
706                    None
707                }
708            }
709            Err(i) => {
710                if self.arrays.is_empty() {
711                    return None;
712                }
713                let delta = ion_mobility - self.ion_mobility_dimension[i];
714                if delta.abs() <= error_tolerance {
715                    self.arrays.get(i).map(|a| (a, delta))
716                } else {
717                    None
718                }
719            }
720        }
721    }
722
723    /// Perform a binary search on the ion mobility dimension and find the nearest
724    /// point to `ion_mobility` within `error_tolerance` absolute distance.
725    ///
726    /// Unlike [`Self::search_ion_mobility`], this returns the index in the ion mobility
727    /// dimension that matched and the ion mobility value at that index.
728    pub fn search_ion_mobility_index(&self, ion_mobility: f64, error_tolerance: f64) -> Option<(usize, f64)> {
729        match self
730            .ion_mobility_dimension
731            .binary_search_by(|x: &f64| x.total_cmp(&ion_mobility))
732        {
733            Ok(i) => {
734                let delta = ion_mobility - self.ion_mobility_dimension[i];
735                if delta.abs() <= error_tolerance {
736                    Some((i, self.ion_mobility_dimension[i]))
737                } else {
738                    None
739                }
740            }
741            Err(i) => {
742                if self.arrays.is_empty() {
743                    return None;
744                }
745                let delta = ion_mobility - self.ion_mobility_dimension[i];
746                if delta.abs() <= error_tolerance {
747                    Some((i, self.ion_mobility_dimension[i]))
748                } else {
749                    None
750                }
751            }
752        }
753    }
754
755    /// Get the a mutable reference to the associated arrays at the requested ion mobility, if they exist.
756    ///
757    /// The ion mobility value must be an *exact* match
758    pub fn get_ion_mobility_mut(&mut self, ion_mobility: f64) -> Option<&mut BinaryArrayMap> {
759        if let Some(i) = NonNaNF64::wrap(ion_mobility) {
760            if let Some(i) = self.ion_mobility_index.get(&i) {
761                self.arrays.get_mut(*i)
762            } else {
763                None
764            }
765        } else {
766            None
767        }
768    }
769
770    /// Iterate over the ion mobility dimension and associated arrays
771    /// at each point.
772    pub fn iter(&self) -> impl Iterator<Item = (f64, &BinaryArrayMap)> {
773        self.ion_mobility_dimension
774            .iter()
775            .copied()
776            .zip(self.arrays.iter())
777    }
778
779    /// Iterate over the ion mobility dimension and a mutable reference to the associated arrays
780    /// at each point.
781    pub fn iter_mut(&mut self) -> impl Iterator<Item = (f64, &mut BinaryArrayMap)> {
782        self.ion_mobility_dimension
783            .iter()
784            .copied()
785            .zip(self.arrays.iter_mut())
786    }
787
788    /// Flatten this array into a single [`BinaryArrayMap`].
789    ///
790    /// The flattened map is sorted by m/z and ion mobility. If no m/z array is found,
791    /// *no* sorting is done, and the flattened order is left in roughly ion mobility sorted
792    /// order.
793    ///
794    /// # Errors
795    /// [`ArrayRetrievalError`] errors related to array decoding occur if
796    /// any [`DataArray`] cannot be decoded, or if an expected array is absent.
797    pub fn unstack(&self) -> Result<BinaryArrayMap, ArrayRetrievalError> {
798        let mut destination = self.additional_arrays.clone();
799
800        let mut im_dim =
801            DataArray::from_name_and_type(&self.ion_mobility_type, BinaryDataArrayType::Float64);
802        im_dim.unit = self.ion_mobility_unit;
803        let mut sizes = Vec::new();
804        let mut mz_size = None;
805        for (layer, im) in self.arrays.iter().zip(self.ion_mobility_dimension.iter()) {
806            sizes.clear();
807            for (key, array) in layer.iter() {
808                match destination.get_mut(key) {
809                    Some(sink) => {
810                        sink.extend_raw(&array.data).inspect_err(|e| {
811                            log::error!("Failed to extend {key:?}: {e}");
812                        })?;
813                    }
814                    None => {
815                        destination.add(array.clone());
816                    }
817                }
818                if matches!(key, ArrayType::MZArray) {
819                    mz_size = Some(array.data_len()?);
820                } else {
821                    sizes.push((key, array.data_len()?));
822                }
823            }
824            if let Some(mz_size) = mz_size {
825                im_dim.extend_iter(std::iter::repeat_n(*im, mz_size))?;
826            } else if let Some((_, size)) = sizes.first() {
827                im_dim.extend_iter(std::iter::repeat_n(*im, *size))?;
828            }
829        }
830
831        let sorter = if let Some(mz_array) = destination.get(&ArrayType::MZArray) {
832            let mzs = mz_array.to_f64()?;
833            let ims = im_dim.to_f64()?;
834
835            let mut indices: Vec<usize> = Vec::with_capacity(mzs.len());
836            indices.extend(0..mzs.len());
837            indices.sort_by(|i, j| {
838                mzs[*i]
839                    .total_cmp(&mzs[*j])
840                    .then_with(|| ims[*i].total_cmp(&ims[*j]))
841            });
842            Some(indices)
843        } else {
844            None
845        };
846
847        let final_size = im_dim.data_len()?;
848
849        destination.add(im_dim);
850
851        if let Some(sorter) = sorter {
852            destination.sort_from_indices(sorter)?;
853        } else if final_size > 0 {
854            log::debug!("Unsorted unstack");
855        }
856        Ok(destination)
857    }
858
859    /// Convert a [`BinaryArrayMap`] into a [`BinaryArrayMap3D`] if it has an ion mobility dimension.
860    ///
861    /// Any arrays that aren't the same length as the ion mobility dimension will be in
862    /// [`BinaryArrayMap3D::additional_arrays`].
863    ///
864    /// This uses *exact* ion mobility values as keys, bit-level equality. If ion mobility
865    /// values are computed, they may vary at trivial bits and be equivalent, but not **equal**.
866    ///
867    /// # Errors
868    /// [`ArrayRetrievalError`] errors related to array decoding occur if any [`DataArray`]
869    /// cannot be decoded, or if an expected array is absent.
870    pub fn stack(source: &BinaryArrayMap) -> Result<Self, ArrayRetrievalError> {
871        let mut this = Self::default();
872        if !source.has_ion_mobility() {
873            return Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray));
874        }
875        let (im_dim, im_type) = source.ion_mobility()?;
876        this.ion_mobility_unit = source.get(&im_type).unwrap().unit;
877        this.ion_mobility_type = im_type;
878        if im_dim.is_empty() {
879            return Ok(this);
880        }
881        let mut im_dim: Vec<(usize, f64)> = im_dim.iter().copied().enumerate().collect();
882        im_dim.sort_by(|(_, va), (_, vb)| va.total_cmp(vb));
883
884        let mut im_axis = Vec::with_capacity(200);
885        let mut last_v = im_dim.first().unwrap().1 - 1.0;
886        let mut index_map = HashMap::new();
887        for (_, v) in im_dim.iter() {
888            if v.total_cmp(&last_v).is_gt() {
889                last_v = *v;
890                index_map.insert(NonNaNF64::from(*v), im_axis.len());
891                im_axis.push(*v);
892            }
893        }
894
895        let mut array_bins: Vec<BinaryArrayMap> = Vec::with_capacity(im_axis.len());
896        array_bins.resize(im_axis.len(), BinaryArrayMap::default());
897        for (array_type, array) in source.iter() {
898            if array_type.is_ion_mobility() {
899                continue;
900            }
901            if array.data_len()? != im_dim.len() {
902                this.additional_arrays.add(array.clone());
903                continue;
904            }
905            match array.dtype() {
906                BinaryDataArrayType::Unknown => {
907                    panic!("Cannot re-sort opaque or unknown dimension data types")
908                }
909                BinaryDataArrayType::Float64 => {
910                    let view = array.to_f64()?;
911                    _populate_stacked_array_from!(
912                        im_dim, view, index_map, array_bins, array_type, array
913                    );
914                }
915                BinaryDataArrayType::Float32 => {
916                    let view = array.to_f32()?;
917                    _populate_stacked_array_from!(
918                        im_dim, view, index_map, array_bins, array_type, array
919                    );
920                }
921                BinaryDataArrayType::Int64 => {
922                    let view = array.to_i64()?;
923                    _populate_stacked_array_from!(
924                        im_dim, view, index_map, array_bins, array_type, array
925                    );
926                }
927                BinaryDataArrayType::Int32 => {
928                    let view = array.to_i32()?;
929                    _populate_stacked_array_from!(
930                        im_dim, view, index_map, array_bins, array_type, array
931                    );
932                }
933                BinaryDataArrayType::ASCII => {
934                    let view = array.decode()?;
935                    _populate_stacked_array_from!(
936                        im_dim, view, index_map, array_bins, array_type, array
937                    );
938                }
939            }
940        }
941        this.ion_mobility_dimension = im_axis;
942        this.ion_mobility_index = index_map;
943        this.arrays = array_bins;
944        Ok(this)
945    }
946}
947
948impl TryFrom<BinaryArrayMap> for BinaryArrayMap3D {
949    type Error = ArrayRetrievalError;
950
951    fn try_from(value: BinaryArrayMap) -> Result<Self, Self::Error> {
952        Self::stack(&value)
953    }
954}
955
956/// A [`BinaryArrayMap3D`] aliases [`TryFrom`] on [`BinaryArrayMap`] with [`BinaryArrayMap3D::stack`]
957impl TryFrom<&BinaryArrayMap> for BinaryArrayMap3D {
958    type Error = ArrayRetrievalError;
959
960    fn try_from(value: &BinaryArrayMap) -> Result<Self, Self::Error> {
961        Self::stack(value)
962    }
963}
964
965impl TryFrom<BinaryArrayMap3D> for BinaryArrayMap {
966    type Error = ArrayRetrievalError;
967
968    fn try_from(value: BinaryArrayMap3D) -> Result<Self, Self::Error> {
969        value.unstack()
970    }
971}
972
973/// A [`BinaryArrayMap`] aliases [`TryFrom`] on [`BinaryArrayMap3D`] with [`BinaryArrayMap3D::unstack`]
974impl TryFrom<&BinaryArrayMap3D> for BinaryArrayMap {
975    type Error = ArrayRetrievalError;
976
977    fn try_from(value: &BinaryArrayMap3D) -> Result<Self, Self::Error> {
978        value.unstack()
979    }
980}
981
982
983
984#[cfg(test)]
985mod test {
986    use crate::BinaryDataArrayType;
987
988    use super::*;
989    use std::fs;
990    use std::io::{self, prelude::*};
991
992    fn make_array_from_file() -> io::Result<DataArray> {
993        let mut fh = fs::File::open("../../test/data/mz_f64_zlib_bas64.txt")?;
994        let mut buf = String::new();
995        fh.read_to_string(&mut buf)?;
996        let bytes: Vec<u8> = buf.into();
997        let mut da = DataArray::wrap(&ArrayType::MZArray, BinaryDataArrayType::Float64, bytes);
998        da.compression = BinaryCompressionType::Zlib;
999        Ok(da)
1000    }
1001
1002    #[test]
1003    fn test_construction() -> io::Result<()> {
1004        let da = make_array_from_file()?;
1005        let mut map = BinaryArrayMap::new();
1006        assert!(!map.has_array(&ArrayType::MZArray));
1007        map.add(da);
1008        assert!(map.has_array(&ArrayType::MZArray));
1009        Ok(())
1010    }
1011
1012    #[test]
1013    fn test_decode() -> io::Result<()> {
1014        let da = make_array_from_file()?;
1015        let mut map = BinaryArrayMap::new();
1016        map.add(da);
1017        assert_eq!(
1018            map.get(&ArrayType::MZArray).unwrap().compression,
1019            BinaryCompressionType::Zlib
1020        );
1021        map.decode_all_arrays()?;
1022        assert_eq!(
1023            map.get(&ArrayType::MZArray).unwrap().compression,
1024            BinaryCompressionType::Decoded
1025        );
1026        Ok(())
1027    }
1028}
mzdata_bindata/map.rs

mzdata_bindata/
map.rs
