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::{Param, ParamDescribed, 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///
587/// Because [`BinaryArrayMap3D`] spreads [`DataArray`] over multiple ion mobility
588/// bins, it cannot share a single [`ParamList`] for each of the dimensions. Instead,
589/// all split [`ArrayType`] have their [`ParamList`] moved to separate storage accessed
590/// using [`params_of`] and its siblins. These [`ParamList`] will be propagated by [`Self::stack`]
591/// and [`BinaryArrayMap3D::unstack`] as appropriate.
592#[derive(Debug, Default, Clone)]
593#[cfg_attr(feature = "serde", serde_with::serde_as)]
594#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
595pub struct BinaryArrayMap3D {
596    pub ion_mobility_dimension: Vec<f64>,
597    pub ion_mobility_type: ArrayType,
598    pub ion_mobility_unit: Unit,
599    pub arrays: Vec<BinaryArrayMap>,
600    pub additional_arrays: BinaryArrayMap,
601    #[cfg_attr(feature = "serde", serde_as(as = "Vec<(_, _)>"))]
602    ion_mobility_index: HashMap<NonNaNF64, usize>,
603    parameters_of: HashMap<ArrayType, mzdata_param::ParamList>
604}
605
606impl BinaryArrayMap3D {
607    /// Create an empty [`BinaryArrayMap3D`] with a defined ion mobility dimension
608    /// but no content at each ion mobility point.
609    ///
610    /// # Panics
611    /// If the ion mobility dimension contains a NaN value
612    pub fn from_ion_mobility_dimension(
613        ion_mobility_dimension: Vec<f64>,
614        ion_mobility_type: ArrayType,
615        ion_mobility_unit: Unit,
616    ) -> BinaryArrayMap3D {
617        let ion_mobility_index: HashMap<_, _> = ion_mobility_dimension
618            .iter()
619            .copied()
620            .enumerate()
621            .map(|(i, f)| {
622                (
623                    NonNaNF64::wrap(f)
624                        .unwrap_or_else(|| panic!("Expected non-NaN value for ion mobility")),
625                    i,
626                )
627            })
628            .collect();
629        let mut arrays = Vec::new();
630        arrays.resize_with(ion_mobility_dimension.len(), BinaryArrayMap::default);
631        Self {
632            ion_mobility_dimension,
633            ion_mobility_type,
634            ion_mobility_unit,
635            arrays,
636            ion_mobility_index,
637            additional_arrays: Default::default(),
638            parameters_of: Default::default(),
639        }
640    }
641
642    /// Construct a [`BinaryArrayMap3D`] from a defined ion mobility dimension and
643    /// a [`Vec`] of [`BinaryArrayMap`] at each point.
644    ///
645    /// # Panics
646    /// If the ion mobility dimension contains a NaN value or `ion_mobility_dimension.len() != arrays.len()`
647    pub fn from_ion_mobility_dimension_and_arrays(
648        ion_mobility_dimension: Vec<f64>,
649        ion_mobility_type: ArrayType,
650        ion_mobility_unit: Unit,
651        arrays: Vec<BinaryArrayMap>,
652    ) -> BinaryArrayMap3D {
653        let ion_mobility_index: HashMap<_, _> = ion_mobility_dimension
654            .iter()
655            .copied()
656            .enumerate()
657            .map(|(i, f)| {
658                (
659                    NonNaNF64::wrap(f)
660                        .unwrap_or_else(|| panic!("Expected non-NaN value for ion mobility")),
661                    i,
662                )
663            })
664            .collect();
665
666        assert_eq!(
667            ion_mobility_dimension.len(),
668            arrays.len(),
669            "The number of ion mobility points does not match the number of point-level maps!"
670        );
671
672        Self {
673            ion_mobility_dimension,
674            ion_mobility_type,
675            ion_mobility_unit,
676            arrays,
677            ion_mobility_index,
678            additional_arrays: Default::default(),
679            parameters_of: Default::default(),
680        }
681    }
682
683    /// Get a reference to the [`ParamList`] for the specified [`ArrayType`], if one exists.
684    ///
685    /// # See also
686    /// - [`Self::params_of_mut`] to mutate existing [`ParamList`] for extant [`ArrayType`]
687    /// - [`Self::insert_params_of`] to create or overwrite the [`ParamList`] for an [`ArrayType`]
688    pub fn params_of(&self, array_type: &ArrayType) -> Option<&Vec<Param>> {
689        self.parameters_of.get(array_type)
690    }
691
692    /// Get a mutable reference to the [`ParamList`] for the specified [`ArrayType`], if one exists.
693    ///
694    /// # See also
695    /// - [`Self::params_of`] to read existing [`ParamList`] for extant [`ArrayType`]
696    /// - [`Self::insert_params_of`] to create or overwrite the [`ParamList`] for an [`ArrayType`]
697    /// - [`Self::params_of_entry`] to create or modify the [`ParamList`] using a lower
698    /// level Entry API
699    pub fn params_of_mut(&mut self, array_type: &ArrayType) -> Option<&mut Vec<Param>> {
700        self.parameters_of.get_mut(array_type)
701    }
702
703    /// A lower level create or update version of [`Self::insert_params_of`] that uses the
704    /// underlying [`HashMap`]'s Entry API
705    pub fn params_of_entry(&mut self, array_type: ArrayType) -> std::collections::hash_map::Entry<'_, ArrayType, Vec<Param>> {
706        self.parameters_of.entry(array_type)
707    }
708
709    /// Set the [`ParamList`] for the specified [`ArrayType`].
710    ///
711    /// Returns any previous [`ParamList`] associated with that [`ArrayType`].
712    ///
713    /// # See also
714    /// - [`Self::params_of`] to read existing [`ParamList`] for extant [`ArrayType`]
715    /// - [`Self::params_of_mut`] to mutate existing [`ParamList`] for extant [`ArrayType`]
716    /// - [`Self::params_of_entry`] to create or modify the [`ParamList`] using a lower
717    /// level Entry API
718    pub fn insert_params_of(&mut self, array_type: ArrayType, params: Vec<Param>) -> Option<Vec<Param>> {
719        self.parameters_of.insert(array_type, params)
720    }
721
722    /// Get the associated arrays at the requested ion mobility, if they exist
723    pub fn get_ion_mobility(&self, ion_mobility: f64) -> Option<&BinaryArrayMap> {
724        if let Some(i) = NonNaNF64::wrap(ion_mobility) {
725            if let Some(i) = self.ion_mobility_index.get(&i) {
726                self.arrays.get(*i)
727            } else {
728                None
729            }
730        } else {
731            None
732        }
733    }
734
735    /// Perform a binary search on the ion mobility dimension and find the nearest
736    /// point to `ion_mobility` within `error_tolerance` absolute distance.
737    ///
738    /// Returns the [`BinaryArrayMap`] at that point and the distance from `ion_mobility`
739    /// matched.
740    pub fn search_ion_mobility(
741        &self,
742        ion_mobility: f64,
743        error_tolerance: f64,
744    ) -> Option<(&BinaryArrayMap, f64)> {
745        match self
746            .ion_mobility_dimension
747            .binary_search_by(|x: &f64| x.total_cmp(&ion_mobility))
748        {
749            Ok(i) => {
750                let delta = ion_mobility - self.ion_mobility_dimension[i];
751                if delta.abs() <= error_tolerance {
752                    self.arrays.get(i).map(|a| (a, delta))
753                } else {
754                    None
755                }
756            }
757            Err(i) => {
758                if self.arrays.is_empty() {
759                    return None;
760                }
761                let delta = ion_mobility - self.ion_mobility_dimension[i];
762                if delta.abs() <= error_tolerance {
763                    self.arrays.get(i).map(|a| (a, delta))
764                } else {
765                    None
766                }
767            }
768        }
769    }
770
771    /// Perform a binary search on the ion mobility dimension and find the nearest
772    /// point to `ion_mobility` within `error_tolerance` absolute distance.
773    ///
774    /// Unlike [`Self::search_ion_mobility`], this returns the index in the ion mobility
775    /// dimension that matched and the ion mobility value at that index.
776    pub fn search_ion_mobility_index(&self, ion_mobility: f64, error_tolerance: f64) -> Option<(usize, f64)> {
777        match self
778            .ion_mobility_dimension
779            .binary_search_by(|x: &f64| x.total_cmp(&ion_mobility))
780        {
781            Ok(i) => {
782                let delta = ion_mobility - self.ion_mobility_dimension[i];
783                if delta.abs() <= error_tolerance {
784                    Some((i, self.ion_mobility_dimension[i]))
785                } else {
786                    None
787                }
788            }
789            Err(i) => {
790                if self.arrays.is_empty() {
791                    return None;
792                }
793                let delta = ion_mobility - self.ion_mobility_dimension[i];
794                if delta.abs() <= error_tolerance {
795                    Some((i, self.ion_mobility_dimension[i]))
796                } else {
797                    None
798                }
799            }
800        }
801    }
802
803    /// Get the a mutable reference to the associated arrays at the requested ion mobility, if they exist.
804    ///
805    /// The ion mobility value must be an *exact* match
806    pub fn get_ion_mobility_mut(&mut self, ion_mobility: f64) -> Option<&mut BinaryArrayMap> {
807        if let Some(i) = NonNaNF64::wrap(ion_mobility) {
808            if let Some(i) = self.ion_mobility_index.get(&i) {
809                self.arrays.get_mut(*i)
810            } else {
811                None
812            }
813        } else {
814            None
815        }
816    }
817
818    /// Iterate over the ion mobility dimension and associated arrays
819    /// at each point.
820    pub fn iter(&self) -> impl Iterator<Item = (f64, &BinaryArrayMap)> {
821        self.ion_mobility_dimension
822            .iter()
823            .copied()
824            .zip(self.arrays.iter())
825    }
826
827    /// Iterate over the ion mobility dimension and a mutable reference to the associated arrays
828    /// at each point.
829    pub fn iter_mut(&mut self) -> impl Iterator<Item = (f64, &mut BinaryArrayMap)> {
830        self.ion_mobility_dimension
831            .iter()
832            .copied()
833            .zip(self.arrays.iter_mut())
834    }
835
836    /// Flatten this array into a single [`BinaryArrayMap`].
837    ///
838    /// The flattened map is sorted by m/z and ion mobility. If no m/z array is found,
839    /// *no* sorting is done, and the flattened order is left in roughly ion mobility sorted
840    /// order.
841    ///
842    /// # Errors
843    /// [`ArrayRetrievalError`] errors related to array decoding occur if
844    /// any [`DataArray`] cannot be decoded, or if an expected array is absent.
845    pub fn unstack(&self) -> Result<BinaryArrayMap, ArrayRetrievalError> {
846        let mut destination = self.additional_arrays.clone();
847
848        let mut im_dim =
849            DataArray::from_name_and_type(&self.ion_mobility_type, BinaryDataArrayType::Float64);
850        im_dim.unit = self.ion_mobility_unit;
851        let mut sizes = Vec::new();
852        let mut mz_size = None;
853        for (layer, im) in self.arrays.iter().zip(self.ion_mobility_dimension.iter()) {
854            sizes.clear();
855            for (key, array) in layer.iter() {
856                match destination.get_mut(key) {
857                    Some(sink) => {
858                        sink.extend_raw(&array.data).inspect_err(|e| {
859                            log::error!("Failed to extend {key:?}: {e}");
860                        })?;
861                    }
862                    None => {
863                        destination.add(array.clone());
864                    }
865                }
866                if matches!(key, ArrayType::MZArray) {
867                    mz_size = Some(array.data_len()?);
868                } else {
869                    sizes.push((key, array.data_len()?));
870                }
871            }
872            if let Some(mz_size) = mz_size {
873                im_dim.extend_iter(std::iter::repeat_n(*im, mz_size))?;
874            } else if let Some((_, size)) = sizes.first() {
875                im_dim.extend_iter(std::iter::repeat_n(*im, *size))?;
876            }
877        }
878
879        let sorter = if let Some(mz_array) = destination.get(&ArrayType::MZArray) {
880            let mzs = mz_array.to_f64()?;
881            let ims = im_dim.to_f64()?;
882
883            let mut indices: Vec<usize> = Vec::with_capacity(mzs.len());
884            indices.extend(0..mzs.len());
885            indices.sort_by(|i, j| {
886                mzs[*i]
887                    .total_cmp(&mzs[*j])
888                    .then_with(|| ims[*i].total_cmp(&ims[*j]))
889            });
890            Some(indices)
891        } else {
892            None
893        };
894
895        let final_size = im_dim.data_len()?;
896
897        destination.add(im_dim);
898
899        if let Some(sorter) = sorter {
900            destination.sort_from_indices(sorter)?;
901        } else if final_size > 0 {
902            log::debug!("Unsorted unstack");
903        }
904
905        for (k, v) in destination.iter_mut(){
906            if let Some(pars) = self.parameters_of.get(k) {
907                v.params_mut().extend(pars.iter().cloned());
908            }
909        }
910        Ok(destination)
911    }
912
913    /// Convert a [`BinaryArrayMap`] into a [`BinaryArrayMap3D`] if it has an ion mobility dimension.
914    ///
915    /// Any arrays that aren't the same length as the ion mobility dimension will be in
916    /// [`BinaryArrayMap3D::additional_arrays`].
917    ///
918    /// This uses *exact* ion mobility values as keys, bit-level equality. If ion mobility
919    /// values are computed, they may vary at trivial bits and be equivalent, but not **equal**.
920    ///
921    /// # Errors
922    /// [`ArrayRetrievalError`] errors related to array decoding occur if any [`DataArray`]
923    /// cannot be decoded, or if an expected array is absent.
924    pub fn stack(source: &BinaryArrayMap) -> Result<Self, ArrayRetrievalError> {
925        let mut this = Self::default();
926        if !source.has_ion_mobility() {
927            return Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray));
928        }
929        let (im_dim, im_type) = source.ion_mobility()?;
930        this.ion_mobility_unit = source.get(&im_type).unwrap().unit;
931        this.ion_mobility_type = im_type;
932        if im_dim.is_empty() {
933            return Ok(this);
934        }
935        let mut im_dim: Vec<(usize, f64)> = im_dim.iter().copied().enumerate().collect();
936        im_dim.sort_by(|(_, va), (_, vb)| va.total_cmp(vb));
937
938        let mut im_axis = Vec::with_capacity(200);
939        let mut last_v = im_dim.first().unwrap().1 - 1.0;
940        let mut index_map = HashMap::new();
941        for (_, v) in im_dim.iter() {
942            if v.total_cmp(&last_v).is_gt() {
943                last_v = *v;
944                index_map.insert(NonNaNF64::from(*v), im_axis.len());
945                im_axis.push(*v);
946            }
947        }
948
949        let mut array_bins: Vec<BinaryArrayMap> = Vec::with_capacity(im_axis.len());
950        array_bins.resize(im_axis.len(), BinaryArrayMap::default());
951        for (array_type, array) in source.iter() {
952            if array_type.is_ion_mobility() {
953                continue;
954            }
955            if array.data_len()? != im_dim.len() {
956                this.additional_arrays.add(array.clone());
957                continue;
958            }
959            if !array.params().is_empty() {
960                this.parameters_of.insert(array_type.clone(), array.params().to_vec());
961            }
962            match array.dtype() {
963                BinaryDataArrayType::Unknown => {
964                    panic!("Cannot re-sort opaque or unknown dimension data types")
965                }
966                BinaryDataArrayType::Float64 => {
967                    let view = array.to_f64()?;
968                    _populate_stacked_array_from!(
969                        im_dim, view, index_map, array_bins, array_type, array
970                    );
971                }
972                BinaryDataArrayType::Float32 => {
973                    let view = array.to_f32()?;
974                    _populate_stacked_array_from!(
975                        im_dim, view, index_map, array_bins, array_type, array
976                    );
977                }
978                BinaryDataArrayType::Int64 => {
979                    let view = array.to_i64()?;
980                    _populate_stacked_array_from!(
981                        im_dim, view, index_map, array_bins, array_type, array
982                    );
983                }
984                BinaryDataArrayType::Int32 => {
985                    let view = array.to_i32()?;
986                    _populate_stacked_array_from!(
987                        im_dim, view, index_map, array_bins, array_type, array
988                    );
989                }
990                BinaryDataArrayType::ASCII => {
991                    let view = array.decode()?;
992                    _populate_stacked_array_from!(
993                        im_dim, view, index_map, array_bins, array_type, array
994                    );
995                }
996            }
997        }
998        this.ion_mobility_dimension = im_axis;
999        this.ion_mobility_index = index_map;
1000        this.arrays = array_bins;
1001        Ok(this)
1002    }
1003}
1004
1005impl TryFrom<BinaryArrayMap> for BinaryArrayMap3D {
1006    type Error = ArrayRetrievalError;
1007
1008    /// A wrapper around [`BinaryArrayMap3D::stack`]
1009    fn try_from(value: BinaryArrayMap) -> Result<Self, Self::Error> {
1010        Self::stack(&value)
1011    }
1012}
1013
1014/// A [`BinaryArrayMap3D`] aliases [`TryFrom`] on [`BinaryArrayMap`] with [`BinaryArrayMap3D::stack`]
1015impl TryFrom<&BinaryArrayMap> for BinaryArrayMap3D {
1016    type Error = ArrayRetrievalError;
1017
1018    fn try_from(value: &BinaryArrayMap) -> Result<Self, Self::Error> {
1019        Self::stack(value)
1020    }
1021}
1022
1023impl TryFrom<BinaryArrayMap3D> for BinaryArrayMap {
1024    type Error = ArrayRetrievalError;
1025
1026    fn try_from(value: BinaryArrayMap3D) -> Result<Self, Self::Error> {
1027        value.unstack()
1028    }
1029}
1030
1031/// A [`BinaryArrayMap`] aliases [`TryFrom`] on [`BinaryArrayMap3D`] with [`BinaryArrayMap3D::unstack`]
1032impl TryFrom<&BinaryArrayMap3D> for BinaryArrayMap {
1033    type Error = ArrayRetrievalError;
1034
1035    fn try_from(value: &BinaryArrayMap3D) -> Result<Self, Self::Error> {
1036        value.unstack()
1037    }
1038}
1039
1040
1041
1042#[cfg(test)]
1043mod test {
1044    use crate::BinaryDataArrayType;
1045
1046    use super::*;
1047    use std::fs;
1048    use std::io::{self, prelude::*};
1049
1050    fn make_array_from_file() -> io::Result<DataArray> {
1051        let mut fh = fs::File::open("../../test/data/mz_f64_zlib_bas64.txt")?;
1052        let mut buf = String::new();
1053        fh.read_to_string(&mut buf)?;
1054        let bytes: Vec<u8> = buf.into();
1055        let mut da = DataArray::wrap(&ArrayType::MZArray, BinaryDataArrayType::Float64, bytes);
1056        da.compression = BinaryCompressionType::Zlib;
1057        Ok(da)
1058    }
1059
1060    #[test]
1061    fn test_construction() -> io::Result<()> {
1062        let da = make_array_from_file()?;
1063        let mut map = BinaryArrayMap::new();
1064        assert!(!map.has_array(&ArrayType::MZArray));
1065        map.add(da);
1066        assert!(map.has_array(&ArrayType::MZArray));
1067        Ok(())
1068    }
1069
1070    #[test]
1071    fn test_decode() -> io::Result<()> {
1072        let da = make_array_from_file()?;
1073        let mut map = BinaryArrayMap::new();
1074        map.add(da);
1075        assert_eq!(
1076            map.get(&ArrayType::MZArray).unwrap().compression,
1077            BinaryCompressionType::Zlib
1078        );
1079        map.decode_all_arrays()?;
1080        assert_eq!(
1081            map.get(&ArrayType::MZArray).unwrap().compression,
1082            BinaryCompressionType::Decoded
1083        );
1084        Ok(())
1085    }
1086}
mzdata_bindata/map.rs

mzdata_bindata/
map.rs
