ndarray_stats/quantile/
mod.rs

1use self::interpolate::{higher_index, lower_index, Interpolate};
2use super::sort::get_many_from_sorted_mut_unchecked;
3use crate::errors::QuantileError;
4use crate::errors::{EmptyInput, MinMaxError, MinMaxError::UndefinedOrder};
5use crate::{MaybeNan, MaybeNanExt};
6use ndarray::prelude::*;
7use ndarray::{RemoveAxis, Zip};
8use noisy_float::types::N64;
9use std::cmp;
10
11/// Quantile methods for `ArrayRef`.
12pub trait QuantileExt<A, D>
13where
14    D: Dimension,
15{
16    /// Finds the index of the minimum value of the array.
17    ///
18    /// Returns `Err(MinMaxError::UndefinedOrder)` if any of the pairwise
19    /// orderings tested by the function are undefined. (For example, this
20    /// occurs if there are any floating-point NaN values in the array.)
21    ///
22    /// Returns `Err(MinMaxError::EmptyInput)` if the array is empty.
23    ///
24    /// Even if there are multiple (equal) elements that are minima, only one
25    /// index is returned. (Which one is returned is unspecified and may depend
26    /// on the memory layout of the array.)
27    ///
28    /// # Example
29    ///
30    /// ```
31    /// use ndarray::array;
32    /// use ndarray_stats::QuantileExt;
33    ///
34    /// let a = array![[1., 3., 5.],
35    ///                [2., 0., 6.]];
36    /// assert_eq!(a.argmin(), Ok((1, 1)));
37    /// ```
38    fn argmin(&self) -> Result<D::Pattern, MinMaxError>
39    where
40        A: PartialOrd;
41
42    /// Finds the index of the minimum value of the array skipping NaN values.
43    ///
44    /// Returns `Err(EmptyInput)` if the array is empty or none of the values in the array
45    /// are non-NaN values.
46    ///
47    /// Even if there are multiple (equal) elements that are minima, only one
48    /// index is returned. (Which one is returned is unspecified and may depend
49    /// on the memory layout of the array.)
50    ///
51    /// # Example
52    ///
53    /// ```
54    /// use ndarray::array;
55    /// use ndarray_stats::QuantileExt;
56    ///
57    /// let a = array![[::std::f64::NAN, 3., 5.],
58    ///                [2., 0., 6.]];
59    /// assert_eq!(a.argmin_skipnan(), Ok((1, 1)));
60    /// ```
61    fn argmin_skipnan(&self) -> Result<D::Pattern, EmptyInput>
62    where
63        A: MaybeNan,
64        A::NotNan: Ord;
65
66    /// Finds the elementwise minimum of the array.
67    ///
68    /// Returns `Err(MinMaxError::UndefinedOrder)` if any of the pairwise
69    /// orderings tested by the function are undefined. (For example, this
70    /// occurs if there are any floating-point NaN values in the array.)
71    ///
72    /// Returns `Err(MinMaxError::EmptyInput)` if the array is empty.
73    ///
74    /// Even if there are multiple (equal) elements that are minima, only one
75    /// is returned. (Which one is returned is unspecified and may depend on
76    /// the memory layout of the array.)
77    fn min(&self) -> Result<&A, MinMaxError>
78    where
79        A: PartialOrd;
80
81    /// Finds the elementwise minimum of the array, skipping NaN values.
82    ///
83    /// Even if there are multiple (equal) elements that are minima, only one
84    /// is returned. (Which one is returned is unspecified and may depend on
85    /// the memory layout of the array.)
86    ///
87    /// **Warning** This method will return a NaN value if none of the values
88    /// in the array are non-NaN values. Note that the NaN value might not be
89    /// in the array.
90    fn min_skipnan(&self) -> &A
91    where
92        A: MaybeNan,
93        A::NotNan: Ord;
94
95    /// Finds the index of the maximum value of the array.
96    ///
97    /// Returns `Err(MinMaxError::UndefinedOrder)` if any of the pairwise
98    /// orderings tested by the function are undefined. (For example, this
99    /// occurs if there are any floating-point NaN values in the array.)
100    ///
101    /// Returns `Err(MinMaxError::EmptyInput)` if the array is empty.
102    ///
103    /// Even if there are multiple (equal) elements that are maxima, only one
104    /// index is returned. (Which one is returned is unspecified and may depend
105    /// on the memory layout of the array.)
106    ///
107    /// # Example
108    ///
109    /// ```
110    /// use ndarray::array;
111    /// use ndarray_stats::QuantileExt;
112    ///
113    /// let a = array![[1., 3., 7.],
114    ///                [2., 5., 6.]];
115    /// assert_eq!(a.argmax(), Ok((0, 2)));
116    /// ```
117    fn argmax(&self) -> Result<D::Pattern, MinMaxError>
118    where
119        A: PartialOrd;
120
121    /// Finds the index of the maximum value of the array skipping NaN values.
122    ///
123    /// Returns `Err(EmptyInput)` if the array is empty or none of the values in the array
124    /// are non-NaN values.
125    ///
126    /// Even if there are multiple (equal) elements that are maxima, only one
127    /// index is returned. (Which one is returned is unspecified and may depend
128    /// on the memory layout of the array.)
129    ///
130    /// # Example
131    ///
132    /// ```
133    /// use ndarray::array;
134    /// use ndarray_stats::QuantileExt;
135    ///
136    /// let a = array![[::std::f64::NAN, 3., 5.],
137    ///                [2., 0., 6.]];
138    /// assert_eq!(a.argmax_skipnan(), Ok((1, 2)));
139    /// ```
140    fn argmax_skipnan(&self) -> Result<D::Pattern, EmptyInput>
141    where
142        A: MaybeNan,
143        A::NotNan: Ord;
144
145    /// Finds the elementwise maximum of the array.
146    ///
147    /// Returns `Err(MinMaxError::UndefinedOrder)` if any of the pairwise
148    /// orderings tested by the function are undefined. (For example, this
149    /// occurs if there are any floating-point NaN values in the array.)
150    ///
151    /// Returns `Err(EmptyInput)` if the array is empty.
152    ///
153    /// Even if there are multiple (equal) elements that are maxima, only one
154    /// is returned. (Which one is returned is unspecified and may depend on
155    /// the memory layout of the array.)
156    fn max(&self) -> Result<&A, MinMaxError>
157    where
158        A: PartialOrd;
159
160    /// Finds the elementwise maximum of the array, skipping NaN values.
161    ///
162    /// Even if there are multiple (equal) elements that are maxima, only one
163    /// is returned. (Which one is returned is unspecified and may depend on
164    /// the memory layout of the array.)
165    ///
166    /// **Warning** This method will return a NaN value if none of the values
167    /// in the array are non-NaN values. Note that the NaN value might not be
168    /// in the array.
169    fn max_skipnan(&self) -> &A
170    where
171        A: MaybeNan,
172        A::NotNan: Ord;
173
174    /// Return the qth quantile of the data along the specified axis.
175    ///
176    /// `q` needs to be a float between 0 and 1, bounds included.
177    /// The qth quantile for a 1-dimensional lane of length `N` is defined
178    /// as the element that would be indexed as `(N-1)q` if the lane were to be sorted
179    /// in increasing order.
180    /// If `(N-1)q` is not an integer the desired quantile lies between
181    /// two data points: we return the lower, nearest, higher or interpolated
182    /// value depending on the `interpolate` strategy.
183    ///
184    /// Some examples:
185    /// - `q=0.` returns the minimum along each 1-dimensional lane;
186    /// - `q=0.5` returns the median along each 1-dimensional lane;
187    /// - `q=1.` returns the maximum along each 1-dimensional lane.
188    /// (`q=0` and `q=1` are considered improper quantiles)
189    ///
190    /// The array is shuffled **in place** along each 1-dimensional lane in
191    /// order to produce the required quantile without allocating a copy
192    /// of the original array. Each 1-dimensional lane is shuffled independently
193    /// from the others.
194    /// No assumptions should be made on the ordering of the array elements
195    /// after this computation.
196    ///
197    /// Complexity ([quickselect](https://en.wikipedia.org/wiki/Quickselect)):
198    /// - average case: O(`m`);
199    /// - worst case: O(`m`^2);
200    /// where `m` is the number of elements in the array.
201    ///
202    /// Returns `Err(EmptyInput)` when the specified axis has length 0.
203    ///
204    /// Returns `Err(InvalidQuantile(q))` if `q` is not between `0.` and `1.` (inclusive).
205    ///
206    /// **Panics** if `axis` is out of bounds.
207    fn quantile_axis_mut<I>(
208        &mut self,
209        axis: Axis,
210        q: N64,
211        interpolate: &I,
212    ) -> Result<Array<A, D::Smaller>, QuantileError>
213    where
214        D: RemoveAxis,
215        A: Ord + Clone,
216        I: Interpolate<A>;
217
218    /// A bulk version of [`quantile_axis_mut`], optimized to retrieve multiple
219    /// quantiles at once.
220    ///
221    /// Returns an `Array`, where subviews along `axis` of the array correspond
222    /// to the elements of `qs`.
223    ///
224    /// See [`quantile_axis_mut`] for additional details on quantiles and the algorithm
225    /// used to retrieve them.
226    ///
227    /// Returns `Err(EmptyInput)` when the specified axis has length 0.
228    ///
229    /// Returns `Err(InvalidQuantile(q))` if any `q` in `qs` is not between `0.` and `1.` (inclusive).
230    ///
231    /// **Panics** if `axis` is out of bounds.
232    ///
233    /// [`quantile_axis_mut`]: #tymethod.quantile_axis_mut
234    ///
235    /// # Example
236    ///
237    /// ```rust
238    /// use ndarray::{array, aview1, Axis};
239    /// use ndarray_stats::{QuantileExt, interpolate::Nearest};
240    /// use noisy_float::types::n64;
241    ///
242    /// let mut data = array![[3, 4, 5], [6, 7, 8]];
243    /// let axis = Axis(1);
244    /// let qs = &[n64(0.3), n64(0.7)];
245    /// let quantiles = data.quantiles_axis_mut(axis, &aview1(qs), &Nearest).unwrap();
246    /// for (&q, quantile) in qs.iter().zip(quantiles.axis_iter(axis)) {
247    ///     assert_eq!(quantile, data.quantile_axis_mut(axis, q, &Nearest).unwrap());
248    /// }
249    /// ```
250    fn quantiles_axis_mut<I>(
251        &mut self,
252        axis: Axis,
253        qs: &ArrayRef<N64, Ix1>,
254        interpolate: &I,
255    ) -> Result<Array<A, D>, QuantileError>
256    where
257        D: RemoveAxis,
258        A: Ord + Clone,
259        I: Interpolate<A>;
260
261    /// Return the `q`th quantile of the data along the specified axis, skipping NaN values.
262    ///
263    /// See [`quantile_axis_mut`](#tymethod.quantile_axis_mut) for details.
264    fn quantile_axis_skipnan_mut<I>(
265        &mut self,
266        axis: Axis,
267        q: N64,
268        interpolate: &I,
269    ) -> Result<Array<A, D::Smaller>, QuantileError>
270    where
271        D: RemoveAxis,
272        A: MaybeNan,
273        A::NotNan: Clone + Ord,
274        I: Interpolate<A::NotNan>;
275
276    private_decl! {}
277}
278
279impl<A, D> QuantileExt<A, D> for ArrayRef<A, D>
280where
281    D: Dimension,
282{
283    fn argmin(&self) -> Result<D::Pattern, MinMaxError>
284    where
285        A: PartialOrd,
286    {
287        let mut current_min = self.first().ok_or(EmptyInput)?;
288        let mut current_pattern_min = D::zeros(self.ndim()).into_pattern();
289
290        for (pattern, elem) in self.indexed_iter() {
291            if elem.partial_cmp(current_min).ok_or(UndefinedOrder)? == cmp::Ordering::Less {
292                current_pattern_min = pattern;
293                current_min = elem
294            }
295        }
296
297        Ok(current_pattern_min)
298    }
299
300    fn argmin_skipnan(&self) -> Result<D::Pattern, EmptyInput>
301    where
302        A: MaybeNan,
303        A::NotNan: Ord,
304    {
305        let mut pattern_min = D::zeros(self.ndim()).into_pattern();
306        let min = self.indexed_fold_skipnan(None, |current_min, (pattern, elem)| {
307            Some(match current_min {
308                Some(m) if (m <= elem) => m,
309                _ => {
310                    pattern_min = pattern;
311                    elem
312                }
313            })
314        });
315        if min.is_some() {
316            Ok(pattern_min)
317        } else {
318            Err(EmptyInput)
319        }
320    }
321
322    fn min(&self) -> Result<&A, MinMaxError>
323    where
324        A: PartialOrd,
325    {
326        let first = self.first().ok_or(EmptyInput)?;
327        self.fold(Ok(first), |acc, elem| {
328            let acc = acc?;
329            match elem.partial_cmp(acc).ok_or(UndefinedOrder)? {
330                cmp::Ordering::Less => Ok(elem),
331                _ => Ok(acc),
332            }
333        })
334    }
335
336    fn min_skipnan(&self) -> &A
337    where
338        A: MaybeNan,
339        A::NotNan: Ord,
340    {
341        let first = self.first().and_then(|v| v.try_as_not_nan());
342        A::from_not_nan_ref_opt(self.fold_skipnan(first, |acc, elem| {
343            Some(match acc {
344                Some(acc) => acc.min(elem),
345                None => elem,
346            })
347        }))
348    }
349
350    fn argmax(&self) -> Result<D::Pattern, MinMaxError>
351    where
352        A: PartialOrd,
353    {
354        let mut current_max = self.first().ok_or(EmptyInput)?;
355        let mut current_pattern_max = D::zeros(self.ndim()).into_pattern();
356
357        for (pattern, elem) in self.indexed_iter() {
358            if elem.partial_cmp(current_max).ok_or(UndefinedOrder)? == cmp::Ordering::Greater {
359                current_pattern_max = pattern;
360                current_max = elem
361            }
362        }
363
364        Ok(current_pattern_max)
365    }
366
367    fn argmax_skipnan(&self) -> Result<D::Pattern, EmptyInput>
368    where
369        A: MaybeNan,
370        A::NotNan: Ord,
371    {
372        let mut pattern_max = D::zeros(self.ndim()).into_pattern();
373        let max = self.indexed_fold_skipnan(None, |current_max, (pattern, elem)| {
374            Some(match current_max {
375                Some(m) if m >= elem => m,
376                _ => {
377                    pattern_max = pattern;
378                    elem
379                }
380            })
381        });
382        if max.is_some() {
383            Ok(pattern_max)
384        } else {
385            Err(EmptyInput)
386        }
387    }
388
389    fn max(&self) -> Result<&A, MinMaxError>
390    where
391        A: PartialOrd,
392    {
393        let first = self.first().ok_or(EmptyInput)?;
394        self.fold(Ok(first), |acc, elem| {
395            let acc = acc?;
396            match elem.partial_cmp(acc).ok_or(UndefinedOrder)? {
397                cmp::Ordering::Greater => Ok(elem),
398                _ => Ok(acc),
399            }
400        })
401    }
402
403    fn max_skipnan(&self) -> &A
404    where
405        A: MaybeNan,
406        A::NotNan: Ord,
407    {
408        let first = self.first().and_then(|v| v.try_as_not_nan());
409        A::from_not_nan_ref_opt(self.fold_skipnan(first, |acc, elem| {
410            Some(match acc {
411                Some(acc) => acc.max(elem),
412                None => elem,
413            })
414        }))
415    }
416
417    fn quantiles_axis_mut<I>(
418        &mut self,
419        axis: Axis,
420        qs: &ArrayRef<N64, Ix1>,
421        interpolate: &I,
422    ) -> Result<Array<A, D>, QuantileError>
423    where
424        D: RemoveAxis,
425        A: Ord + Clone,
426        I: Interpolate<A>,
427    {
428        // Minimize number of type parameters to avoid monomorphization bloat.
429        fn quantiles_axis_mut<A, D, I>(
430            mut data: ArrayViewMut<'_, A, D>,
431            axis: Axis,
432            qs: ArrayView1<'_, N64>,
433            _interpolate: &I,
434        ) -> Result<Array<A, D>, QuantileError>
435        where
436            D: RemoveAxis,
437            A: Ord + Clone,
438            I: Interpolate<A>,
439        {
440            for &q in qs {
441                if !((q >= 0.) && (q <= 1.)) {
442                    return Err(QuantileError::InvalidQuantile(q));
443                }
444            }
445
446            let axis_len = data.len_of(axis);
447            if axis_len == 0 {
448                return Err(QuantileError::EmptyInput);
449            }
450
451            let mut results_shape = data.raw_dim();
452            results_shape[axis.index()] = qs.len();
453            if results_shape.size() == 0 {
454                return Ok(Array::from_shape_vec(results_shape, Vec::new()).unwrap());
455            }
456
457            let mut searched_indexes = Vec::with_capacity(2 * qs.len());
458            for &q in &qs {
459                if I::needs_lower(q, axis_len) {
460                    searched_indexes.push(lower_index(q, axis_len));
461                }
462                if I::needs_higher(q, axis_len) {
463                    searched_indexes.push(higher_index(q, axis_len));
464                }
465            }
466            searched_indexes.sort();
467            searched_indexes.dedup();
468
469            let mut results = Array::from_elem(results_shape, data.first().unwrap().clone());
470            Zip::from(results.lanes_mut(axis))
471                .and(data.lanes_mut(axis))
472                .for_each(|mut results, mut data| {
473                    let index_map =
474                        get_many_from_sorted_mut_unchecked(&mut data, &searched_indexes);
475                    for (result, &q) in results.iter_mut().zip(qs) {
476                        let lower = if I::needs_lower(q, axis_len) {
477                            Some(index_map[&lower_index(q, axis_len)].clone())
478                        } else {
479                            None
480                        };
481                        let higher = if I::needs_higher(q, axis_len) {
482                            Some(index_map[&higher_index(q, axis_len)].clone())
483                        } else {
484                            None
485                        };
486                        *result = I::interpolate(lower, higher, q, axis_len);
487                    }
488                });
489            Ok(results)
490        }
491
492        quantiles_axis_mut(self.view_mut(), axis, qs.view(), interpolate)
493    }
494
495    fn quantile_axis_mut<I>(
496        &mut self,
497        axis: Axis,
498        q: N64,
499        interpolate: &I,
500    ) -> Result<Array<A, D::Smaller>, QuantileError>
501    where
502        D: RemoveAxis,
503        A: Ord + Clone,
504        I: Interpolate<A>,
505    {
506        self.quantiles_axis_mut(axis, &aview1(&[q]), interpolate)
507            .map(|a| a.index_axis_move(axis, 0))
508    }
509
510    fn quantile_axis_skipnan_mut<I>(
511        &mut self,
512        axis: Axis,
513        q: N64,
514        interpolate: &I,
515    ) -> Result<Array<A, D::Smaller>, QuantileError>
516    where
517        D: RemoveAxis,
518        A: MaybeNan,
519        A::NotNan: Clone + Ord,
520        I: Interpolate<A::NotNan>,
521    {
522        if !((q >= 0.) && (q <= 1.)) {
523            return Err(QuantileError::InvalidQuantile(q));
524        }
525
526        if self.len_of(axis) == 0 {
527            return Err(QuantileError::EmptyInput);
528        }
529
530        let quantile = self.map_axis_mut(axis, |lane| {
531            let mut not_nan = A::remove_nan_mut(lane);
532            A::from_not_nan_opt(if not_nan.is_empty() {
533                None
534            } else {
535                Some(
536                    not_nan
537                        .quantile_axis_mut::<I>(Axis(0), q, interpolate)
538                        .unwrap()
539                        .into_scalar(),
540                )
541            })
542        });
543        Ok(quantile)
544    }
545
546    private_impl! {}
547}
548
549/// Quantile methods for 1-D arrays.
550pub trait Quantile1dExt<A> {
551    /// Return the qth quantile of the data.
552    ///
553    /// `q` needs to be a float between 0 and 1, bounds included.
554    /// The qth quantile for a 1-dimensional array of length `N` is defined
555    /// as the element that would be indexed as `(N-1)q` if the array were to be sorted
556    /// in increasing order.
557    /// If `(N-1)q` is not an integer the desired quantile lies between
558    /// two data points: we return the lower, nearest, higher or interpolated
559    /// value depending on the `interpolate` strategy.
560    ///
561    /// Some examples:
562    /// - `q=0.` returns the minimum;
563    /// - `q=0.5` returns the median;
564    /// - `q=1.` returns the maximum.
565    /// (`q=0` and `q=1` are considered improper quantiles)
566    ///
567    /// The array is shuffled **in place** in order to produce the required quantile
568    /// without allocating a copy.
569    /// No assumptions should be made on the ordering of the array elements
570    /// after this computation.
571    ///
572    /// Complexity ([quickselect](https://en.wikipedia.org/wiki/Quickselect)):
573    /// - average case: O(`m`);
574    /// - worst case: O(`m`^2);
575    /// where `m` is the number of elements in the array.
576    ///
577    /// Returns `Err(EmptyInput)` if the array is empty.
578    ///
579    /// Returns `Err(InvalidQuantile(q))` if `q` is not between `0.` and `1.` (inclusive).
580    fn quantile_mut<I>(&mut self, q: N64, interpolate: &I) -> Result<A, QuantileError>
581    where
582        A: Ord + Clone,
583        I: Interpolate<A>;
584
585    /// A bulk version of [`quantile_mut`], optimized to retrieve multiple
586    /// quantiles at once.
587    ///
588    /// Returns an `Array`, where the elements of the array correspond to the
589    /// elements of `qs`.
590    ///
591    /// Returns `Err(EmptyInput)` if the array is empty.
592    ///
593    /// Returns `Err(InvalidQuantile(q))` if any `q` in
594    /// `qs` is not between `0.` and `1.` (inclusive).
595    ///
596    /// See [`quantile_mut`] for additional details on quantiles and the algorithm
597    /// used to retrieve them.
598    ///
599    /// [`quantile_mut`]: #tymethod.quantile_mut
600    fn quantiles_mut<I>(
601        &mut self,
602        qs: &ArrayRef<N64, Ix1>,
603        interpolate: &I,
604    ) -> Result<Array1<A>, QuantileError>
605    where
606        A: Ord + Clone,
607        I: Interpolate<A>;
608
609    private_decl! {}
610}
611
612impl<A> Quantile1dExt<A> for ArrayRef<A, Ix1> {
613    fn quantile_mut<I>(&mut self, q: N64, interpolate: &I) -> Result<A, QuantileError>
614    where
615        A: Ord + Clone,
616        I: Interpolate<A>,
617    {
618        Ok(self
619            .quantile_axis_mut(Axis(0), q, interpolate)?
620            .into_scalar())
621    }
622
623    fn quantiles_mut<I>(
624        &mut self,
625        qs: &ArrayRef<N64, Ix1>,
626        interpolate: &I,
627    ) -> Result<Array1<A>, QuantileError>
628    where
629        A: Ord + Clone,
630        I: Interpolate<A>,
631    {
632        self.quantiles_axis_mut(Axis(0), qs, interpolate)
633    }
634
635    private_impl! {}
636}
637
638pub mod interpolate;