Add extension traits for percentile and sorting

Author: LukeMathWalker

Cargo.toml

@@ -6,6 +6,8 @@ authors = ["Jim Turner <rust@turner.link>"]
 [dependencies]
 ndarray = "0.12"
 noisy_float = "0.1"
+num-traits = "0.2"
+rand = "0.5"
 
 [dev-dependencies]
 quickcheck = "0.7"

src/lib.rs

@@ -1,13 +1,19 @@
 #[macro_use(s)]
 extern crate ndarray;
 extern crate noisy_float;
+extern crate num_traits;
+extern crate rand;
 
 #[cfg(test)]
 #[macro_use(quickcheck)]
 extern crate quickcheck;
 
 pub use maybe_nan::MaybeNan;
 pub use min_max::MinMaxExt;
+pub use percentile::{interpolate, PercentileExt};
+pub use sort::Sort1dExt;
 
 mod maybe_nan;
 mod min_max;
+mod percentile;
+mod sort;

src/percentile.rs

@@ -0,0 +1,233 @@
+use interpolate::Interpolate;
+use ndarray::prelude::*;
+use ndarray::{Data, DataMut, RemoveAxis};
+use sort::Sort1dExt;
+
+pub mod interpolate {
+    use ndarray::prelude::*;
+    use num_traits::FromPrimitive;
+    use std::ops::{Add, Div, Mul, Sub};
+
+    /// Used to provide an interpolation strategy to [`percentile_axis_mut`].
+    ///
+    /// [`percentile_axis_mut`]: ../trait.PercentileExt.html#tymethod.percentile_axis_mut
+    pub trait Interpolate<T> {
+        fn float_percentile_index(q: f64, len: usize) -> f64 {
+            ((len - 1) as f64) * q
+        }
+        fn lower_index(q: f64, len: usize) -> usize {
+            Self::float_percentile_index(q, len).floor() as usize
+        }
+        fn upper_index(q: f64, len: usize) -> usize {
+            Self::float_percentile_index(q, len).ceil() as usize
+        }
+        fn float_percentile_index_fraction(q: f64, len: usize) -> f64 {
+            Self::float_percentile_index(q, len) - (Self::lower_index(q, len) as f64)
+        }
+        fn needs_lower(q: f64, len: usize) -> bool;
+        fn needs_upper(q: f64, len: usize) -> bool;
+        fn interpolate<D>(
+            lower: Option<Array<T, D>>,
+            upper: Option<Array<T, D>>,
+            q: f64,
+            len: usize,
+        ) -> Array<T, D>
+        where
+            D: Dimension;
+    }
+
+    pub struct Upper;
+    pub struct Lower;
+    pub struct Nearest;
+    pub struct Midpoint;
+    pub struct Linear;
+
+    impl<T> Interpolate<T> for Upper {
+        fn needs_lower(_q: f64, _len: usize) -> bool {
+            false
+        }
+        fn needs_upper(_q: f64, _len: usize) -> bool {
+            true
+        }
+        fn interpolate<D>(
+            _lower: Option<Array<T, D>>,
+            upper: Option<Array<T, D>>,
+            _q: f64,
+            _len: usize,
+        ) -> Array<T, D> {
+            upper.unwrap()
+        }
+    }
+
+    impl<T> Interpolate<T> for Lower {
+        fn needs_lower(_q: f64, _len: usize) -> bool {
+            true
+        }
+        fn needs_upper(_q: f64, _len: usize) -> bool {
+            false
+        }
+        fn interpolate<D>(
+            lower: Option<Array<T, D>>,
+            _upper: Option<Array<T, D>>,
+            _q: f64,
+            _len: usize,
+        ) -> Array<T, D> {
+            lower.unwrap()
+        }
+    }
+
+    impl<T> Interpolate<T> for Nearest {
+        fn needs_lower(q: f64, len: usize) -> bool {
+            let lower = <Self as Interpolate<T>>::lower_index(q, len);
+            ((lower as f64) - <Self as Interpolate<T>>::float_percentile_index(q, len)) <= 0.
+        }
+        fn needs_upper(q: f64, len: usize) -> bool {
+            !<Self as Interpolate<T>>::needs_lower(q, len)
+        }
+        fn interpolate<D>(
+            lower: Option<Array<T, D>>,
+            upper: Option<Array<T, D>>,
+            q: f64,
+            len: usize,
+        ) -> Array<T, D> {
+            if <Self as Interpolate<T>>::needs_lower(q, len) {
+                lower.unwrap()
+            } else {
+                upper.unwrap()
+            }
+        }
+    }
+
+    impl<T> Interpolate<T> for Midpoint
+    where
+        T: Add<T, Output = T> + Div<T, Output = T> + Clone + FromPrimitive,
+    {
+        fn needs_lower(_q: f64, _len: usize) -> bool {
+            true
+        }
+        fn needs_upper(_q: f64, _len: usize) -> bool {
+            true
+        }
+        fn interpolate<D>(
+            lower: Option<Array<T, D>>,
+            upper: Option<Array<T, D>>,
+            _q: f64,
+            _len: usize,
+        ) -> Array<T, D>
+        where
+            D: Dimension,
+        {
+            let denom = T::from_u8(2).unwrap();
+            (lower.unwrap() + upper.unwrap()).mapv_into(|x| x / denom.clone())
+        }
+    }
+
+    impl<T> Interpolate<T> for Linear
+    where
+        T: Add<T, Output = T> + Sub<T, Output = T> + Mul<T, Output = T> + Clone + FromPrimitive,
+    {
+        fn needs_lower(_q: f64, _len: usize) -> bool {
+            true
+        }
+        fn needs_upper(_q: f64, _len: usize) -> bool {
+            true
+        }
+        fn interpolate<D>(
+            lower: Option<Array<T, D>>,
+            upper: Option<Array<T, D>>,
+            q: f64,
+            len: usize,
+        ) -> Array<T, D>
+        where
+            D: Dimension,
+        {
+            let fraction = T::from_f64(<Self as Interpolate<T>>::float_percentile_index_fraction(
+                q, len,
+            )).unwrap();
+            let a = lower.unwrap().mapv_into(|x| x * fraction.clone());
+            let b = upper
+                .unwrap()
+                .mapv_into(|x| x * (T::from_u8(1).unwrap() - fraction.clone()));
+            a + b
+        }
+    }
+}
+
+pub trait PercentileExt<A, S, D>
+where
+    S: Data<Elem = A>,
+    D: Dimension,
+{
+    /// Return the qth percentile of the data along the specified axis.
+    ///
+    /// `q` needs to be a float between 0 and 1, bounds included.
+    /// The qth percentile for a 1-dimensional lane of length `N` is defined
+    /// as the element that would be indexed as `(N-1)q` if the lane were to be sorted
+    /// in increasing order.
+    /// If `(N-1)q` is not an integer the desired percentile lies between
+    /// two data points: we return the lower, nearest, higher or interpolated
+    /// value depending on the type `Interpolate` bound `I`.
+    ///
+    /// Some examples:
+    /// - `q=0.` returns the minimum along each 1-dimensional lane;
+    /// - `q=0.5` returns the median along each 1-dimensional lane;
+    /// - `q=1.` returns the maximum along each 1-dimensional lane.
+    /// (`q=0` and `q=1` are considered improper percentiles)
+    ///
+    /// The array is shuffled **in place** along each 1-dimensional lane in
+    /// order to produce the required percentile without allocating a copy
+    /// of the original array. Each 1-dimensional lane is shuffled independently
+    /// from the others.
+    /// No assumptions should be made on the ordering of the array elements
+    /// after this computation.
+    ///
+    /// Complexity ([quickselect](https://en.wikipedia.org/wiki/Quickselect)):
+    /// - average case: O(`m`);
+    /// - worst case: O(`m`^2);
+    /// where `m` is the number of elements in the array.
+    ///
+    /// **Panics** if `axis` is out of bounds or if `q` is not between
+    /// `0.` and `1.` (inclusive).
+    fn percentile_axis_mut<I>(&mut self, axis: Axis, q: f64) -> Array<A, D::Smaller>
+    where
+        D: RemoveAxis,
+        A: Ord + Clone,
+        S: DataMut,
+        I: Interpolate<A>;
+}
+
+impl<A, S, D> PercentileExt<A, S, D> for ArrayBase<S, D>
+where
+    S: Data<Elem = A>,
+    D: Dimension,
+{
+    fn percentile_axis_mut<I>(&mut self, axis: Axis, q: f64) -> Array<A, D::Smaller>
+    where
+        D: RemoveAxis,
+        A: Ord + Clone,
+        S: DataMut,
+        I: Interpolate<A>,
+    {
+        assert!((0. <= q) && (q <= 1.));
+        let mut lower = None;
+        let mut upper = None;
+        let axis_len = self.len_of(axis);
+        if I::needs_lower(q, axis_len) {
+            let lower_index = I::lower_index(q, axis_len);
+            lower = Some(self.map_axis_mut(axis, |mut x| x.sorted_get_mut(lower_index)));
+            if I::needs_upper(q, axis_len) {
+                let upper_index = I::upper_index(q, axis_len);
+                let relative_upper_index = upper_index - lower_index;
+                upper = Some(self.map_axis_mut(axis, |mut x| {
+                    x.slice_mut(s![lower_index..])
+                        .sorted_get_mut(relative_upper_index)
+                }));
+            };
+        } else {
+            upper = Some(
+                self.map_axis_mut(axis, |mut x| x.sorted_get_mut(I::upper_index(q, axis_len))),
+            );
+        };
+        I::interpolate(lower, upper, q, axis_len)
+    }
+}

src/sort.rs

@@ -0,0 +1,123 @@
+use ndarray::prelude::*;
+use ndarray::{Data, DataMut};
+use rand::prelude::*;
+use rand::thread_rng;
+
+pub trait Sort1dExt<A, S>
+where
+    S: Data<Elem = A>,
+{
+    /// Return the element that would occupy the `i`-th position if
+    /// the array were sorted in increasing order.
+    ///
+    /// The array is shuffled **in place** to retrieve the desired element:
+    /// no copy of the array is allocated.
+    /// After the shuffling, all elements with an index smaller than `i`
+    /// are smaller than the desired element, while all elements with
+    /// an index greater or equal than `i` are greater than or equal
+    /// to the desired element.
+    ///
+    /// No other assumptions should be made on the ordering of the
+    /// elements after this computation.
+    ///
+    /// Complexity ([quickselect](https://en.wikipedia.org/wiki/Quickselect)):
+    /// - average case: O(`n`);
+    /// - worst case: O(`n`^2);
+    /// where n is the number of elements in the array.
+    ///
+    /// **Panics** if `i` is greater than or equal to `n`.
+    fn sorted_get_mut(&mut self, i: usize) -> A
+    where
+        A: Ord + Clone,
+        S: DataMut;
+
+    /// Return the index of `self[partition_index]` if `self` were to be sorted
+    /// in increasing order.
+    ///
+    /// `self` elements are rearranged in such a way that `self[partition_index]`
+    /// is in the position it would be in an array sorted in increasing order.
+    /// All elements smaller than `self[partition_index]` are moved to its
+    /// left and all elements equal or greater than `self[partition_index]`
+    /// are moved to its right.
+    /// The ordering of the elements in the two partitions is undefined.
+    ///
+    /// `self` is shuffled **in place** to operate the desired partition:
+    /// no copy of the array is allocated.
+    ///
+    /// The method uses Hoare's partition algorithm.
+    /// Complexity: O(`n`), where `n` is the number of elements in the array.
+    /// Average number of element swaps: n/6 - 1/3 (see
+    /// [link](https://cs.stackexchange.com/questions/11458/quicksort-partitioning-hoare-vs-lomuto/11550))
+    ///
+    /// **Panics** if `partition_index` is greater than or equal to `n`.
+    fn partition_mut(&mut self, pivot_index: usize) -> usize
+    where
+        A: Ord + Clone,
+        S: DataMut;
+}
+
+impl<A, S> Sort1dExt<A, S> for ArrayBase<S, Ix1>
+where
+    S: Data<Elem = A>,
+{
+    fn sorted_get_mut(&mut self, i: usize) -> A
+    where
+        A: Ord + Clone,
+        S: DataMut,
+    {
+        let n = self.len();
+        if n == 1 {
+            self[0].clone()
+        } else {
+            let mut rng = thread_rng();
+            let pivot_index = rng.gen_range(0, n);
+            let partition_index = self.partition_mut(pivot_index);
+            if i < partition_index {
+                self.slice_mut(s![..partition_index]).sorted_get_mut(i)
+            } else if i == partition_index {
+                self[i].clone()
+            } else {
+                self.slice_mut(s![partition_index + 1..])
+                    .sorted_get_mut(i - (partition_index + 1))
+            }
+        }
+    }
+
+    fn partition_mut(&mut self, pivot_index: usize) -> usize
+    where
+        A: Ord + Clone,
+        S: DataMut,
+    {
+        let pivot_value = self[pivot_index].clone();
+        self.swap(pivot_index, 0);
+        let n = self.len();
+        let mut i = 1;
+        let mut j = n - 1;
+        loop {
+            loop {
+                if i > j {
+                    break;
+                }
+                if self[i] >= pivot_value {
+                    break;
+                }
+                i += 1;
+            }
+            while pivot_value <= self[j] {
+                if j == 1 {
+                    break;
+                }
+                j -= 1;
+            }
+            if i >= j {
+                break;
+            } else {
+                self.swap(i, j);
+                i += 1;
+                j -= 1;
+            }
+        }
+        self.swap(0, i - 1);
+        i - 1
+    }
+}