Skip to main content

mzdata_bindata/
map.rs

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