Bulk quantiles (#26)

* Promoted module to directory

* Moved interpolate to separate file

* Re-implemented quantile_axis_mut to get closer to something we can use for bulk computation

* Use a set instead of a vec to avoid repeating computations

* Use bulk method for single quantile

* Implement bulk method to get sorted

* Refactored quantiles_axis_mut to use sorted_get_many_mut

* Avoid recomputing index value

* Add quantiles_mut to 1d trait

* Return hashmaps from bulk methods

* Fixed tests

* Use IndexSet to preserve insertion order

* Fix indentation

* IndexMap provides a more intuitive behaviour

* Remove prints

* Renamed methods

* Docs for get_many_from_sorted_mut

* Added docs for private free function

* Docs for quantiles_mut

* Fixed several typos in docs

* More robust test

* Added test for quantiles

* Test quantiles_axis_mut

* Add comments

* Return options when the lane we are computing against is empty

* Fixed docs

* Fixed tests

* Move *index* functions out of Interpolate trait

The behavior of these functions should be independent of the
interpolation strategy.

* Reduce indentation in quantiles_axis_mut

* Reduce indentation in quantile_axis_skipnan_mut

* Use .into_scalar() method

* Improve docs of partition_mut

* Reformat quantiles_axis_mut

* Cargo fmt

* Fmt

* Formatting

* Log version works

* Refactor

* Fix indexes

* Working implementation

* Shorter syntax

* Formatting

* Better docs

* Comments

* Typo

* Don't lose pieces after rebase

* Fmt

* Reduce code duplication

* Fmt

* Clarify docs of get_many_from_sorted_mut_unchecked

* Add get_many_from_sorted_mut benchmark

* Simplify get_many_from_sorted_mut_unchecked

* Eliminate allocations from _get_many_from_sorted_mut_unchecked

This improves performance, especially for small arrays.

* Add get_from_sorted_mut benchmark

* Call slice_axis_mut instead of slice_mut

This has slightly lower overhead.

* Replace iter::repeat with vec!

This is a bit cleaner.

* Fix typo in comment

* Remove unnecessary type annotation

* Simplify quantiles tests

For me, the tests are easier to understand when they don't collect
into a `Vec`.

* Check keys in test_sorted_get_many_mut

* Simplify sort tests

* Improve sort and quantiles docs

* Make Interpolate::interpolate operate elementwise

* Make quantiles_* return Array instead of IndexMap

* Add interpolate parameter to quantile*

This has a few advantages:

* It's now possible to save the interpolation strategy in a variable
  and easily re-use it.

* We can now freely add more type parameters to the `quantile*`
  methods as needed without making them more difficult to call.

* We now have the flexibility to add more advanced interpolation
  strategies in the future (e.g. one that wraps a closure).

* Calling the `quantile*` methods is now slightly more compact because
  a turbofish isn't necessary.

* Make get_many_from_sorted_mut take array of indexes

This is slightly more versatile because `ArrayBase` allows arbitrary
strides.

* Make quantiles* take array instead of slice

* Remove unnecessary IndexSet

* Return EmptyInput instead of None

* Fix tests

* Match output type for argmin/max_skipnan

* Fix tests

* Fmt

* Update src/lib.rs

Co-Authored-By: LukeMathWalker <[email protected]>

* Add quantile error

* Renamed InvalidFraction to InvalidQuantile

* Return QuantileError

* Fix tests

* Fix docs

* Fmt

* Simplify and deduplicate

Author: LukeMathWalker

Cargo.toml

@@ -21,8 +21,15 @@ num-integer = "0.1"
 num-traits = "0.2"
 rand = "0.6"
 itertools = { version = "0.7.0", default-features = false }
+indexmap = "1.0"
 
 [dev-dependencies]
-quickcheck = "0.7"
+criterion = "0.2"
+quickcheck = { version = "0.8.1", default-features = false }
 ndarray-rand = "0.9"
 approx = "0.3"
+quickcheck_macros = "0.8"
+
+[[bench]]
+name = "sort"
+harness = false

benches/sort.rs

@@ -0,0 +1,67 @@
+extern crate criterion;
+extern crate ndarray;
+extern crate ndarray_stats;
+extern crate rand;
+
+use criterion::{
+    black_box, criterion_group, criterion_main, AxisScale, BatchSize, Criterion,
+    ParameterizedBenchmark, PlotConfiguration,
+};
+use ndarray::prelude::*;
+use ndarray_stats::Sort1dExt;
+use rand::prelude::*;
+
+fn get_from_sorted_mut(c: &mut Criterion) {
+    let lens = vec![10, 100, 1000, 10000];
+    let benchmark = ParameterizedBenchmark::new(
+        "get_from_sorted_mut",
+        |bencher, &len| {
+            let mut rng = StdRng::seed_from_u64(42);
+            let mut data: Vec<_> = (0..len).collect();
+            data.shuffle(&mut rng);
+            let indices: Vec<_> = (0..len).step_by(len / 10).collect();
+            bencher.iter_batched(
+                || Array1::from(data.clone()),
+                |mut arr| {
+                    for &i in &indices {
+                        black_box(arr.get_from_sorted_mut(i));
+                    }
+                },
+                BatchSize::SmallInput,
+            )
+        },
+        lens,
+    )
+    .plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic));
+    c.bench("get_from_sorted_mut", benchmark);
+}
+
+fn get_many_from_sorted_mut(c: &mut Criterion) {
+    let lens = vec![10, 100, 1000, 10000];
+    let benchmark = ParameterizedBenchmark::new(
+        "get_many_from_sorted_mut",
+        |bencher, &len| {
+            let mut rng = StdRng::seed_from_u64(42);
+            let mut data: Vec<_> = (0..len).collect();
+            data.shuffle(&mut rng);
+            let indices: Vec<_> = (0..len).step_by(len / 10).collect();
+            bencher.iter_batched(
+                || Array1::from(data.clone()),
+                |mut arr| {
+                    black_box(arr.get_many_from_sorted_mut(&indices));
+                },
+                BatchSize::SmallInput,
+            )
+        },
+        lens,
+    )
+    .plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic));
+    c.bench("get_many_from_sorted_mut", benchmark);
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default();
+    targets = get_from_sorted_mut, get_many_from_sorted_mut
+}
+criterion_main!(benches);

src/errors.rs

@@ -1,4 +1,5 @@
 //! Custom errors returned from our methods and functions.
+use noisy_float::types::N64;
 use std::error::Error;
 use std::fmt;
 
@@ -112,3 +113,31 @@ impl From<ShapeMismatch> for MultiInputError {
         MultiInputError::ShapeMismatch(err)
     }
 }
+
+/// An error computing a quantile.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum QuantileError {
+    /// The input was empty.
+    EmptyInput,
+    /// The `q` was not between `0.` and `1.` (inclusive).
+    InvalidQuantile(N64),
+}
+
+impl fmt::Display for QuantileError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            QuantileError::EmptyInput => write!(f, "Empty input."),
+            QuantileError::InvalidQuantile(q) => {
+                write!(f, "{:} is not between 0. and 1. (inclusive).", q)
+            }
+        }
+    }
+}
+
+impl Error for QuantileError {}
+
+impl From<EmptyInput> for QuantileError {
+    fn from(_: EmptyInput) -> QuantileError {
+        QuantileError::EmptyInput
+    }
+}

src/histogram/strategies.rs

@@ -24,6 +24,7 @@ use super::errors::BinsBuildError;
 use super::{Bins, Edges};
 use ndarray::prelude::*;
 use ndarray::Data;
+use noisy_float::types::n64;
 use num_traits::{FromPrimitive, NumOps, Zero};
 
 /// A trait implemented by all strategies to build [`Bins`]
@@ -334,8 +335,8 @@ where
         }
 
         let mut a_copy = a.to_owned();
-        let first_quartile = a_copy.quantile_mut::<Nearest>(0.25).unwrap();
-        let third_quartile = a_copy.quantile_mut::<Nearest>(0.75).unwrap();
+        let first_quartile = a_copy.quantile_mut(n64(0.25), &Nearest).unwrap();
+        let third_quartile = a_copy.quantile_mut(n64(0.75), &Nearest).unwrap();
         let iqr = third_quartile - first_quartile;
 
         let bin_width = FreedmanDiaconis::compute_bin_width(n_points, iqr);

src/lib.rs

@@ -25,6 +25,7 @@
 //! [`NumPy`]: https://docs.scipy.org/doc/numpy-1.14.1/reference/routines.statistics.html
 //! [`StatsBase.jl`]: https://juliastats.github.io/StatsBase.jl/latest/
 
+extern crate indexmap;
 extern crate itertools;
 extern crate ndarray;
 extern crate noisy_float;

src/quantile/interpolate.rs

@@ -0,0 +1,138 @@
+//! Interpolation strategies.
+use noisy_float::types::N64;
+use num_traits::{Float, FromPrimitive, NumOps, ToPrimitive};
+
+fn float_quantile_index(q: N64, len: usize) -> N64 {
+    q * ((len - 1) as f64)
+}
+
+/// Returns the fraction that the quantile is between the lower and higher indices.
+///
+/// This ranges from 0, where the quantile exactly corresponds the lower index,
+/// to 1, where the quantile exactly corresponds to the higher index.
+fn float_quantile_index_fraction(q: N64, len: usize) -> N64 {
+    float_quantile_index(q, len).fract()
+}
+
+/// Returns the index of the value on the lower side of the quantile.
+pub(crate) fn lower_index(q: N64, len: usize) -> usize {
+    float_quantile_index(q, len).floor().to_usize().unwrap()
+}
+
+/// Returns the index of the value on the higher side of the quantile.
+pub(crate) fn higher_index(q: N64, len: usize) -> usize {
+    float_quantile_index(q, len).ceil().to_usize().unwrap()
+}
+
+/// Used to provide an interpolation strategy to [`quantile_axis_mut`].
+///
+/// [`quantile_axis_mut`]: ../trait.QuantileExt.html#tymethod.quantile_axis_mut
+pub trait Interpolate<T> {
+    /// Returns `true` iff the lower value is needed to compute the
+    /// interpolated value.
+    #[doc(hidden)]
+    fn needs_lower(q: N64, len: usize) -> bool;
+
+    /// Returns `true` iff the higher value is needed to compute the
+    /// interpolated value.
+    #[doc(hidden)]
+    fn needs_higher(q: N64, len: usize) -> bool;
+
+    /// Computes the interpolated value.
+    ///
+    /// **Panics** if `None` is provided for the lower value when it's needed
+    /// or if `None` is provided for the higher value when it's needed.
+    #[doc(hidden)]
+    fn interpolate(lower: Option<T>, higher: Option<T>, q: N64, len: usize) -> T;
+}
+
+/// Select the higher value.
+pub struct Higher;
+/// Select the lower value.
+pub struct Lower;
+/// Select the nearest value.
+pub struct Nearest;
+/// Select the midpoint of the two values (`(lower + higher) / 2`).
+pub struct Midpoint;
+/// Linearly interpolate between the two values
+/// (`lower + (higher - lower) * fraction`, where `fraction` is the
+/// fractional part of the index surrounded by `lower` and `higher`).
+pub struct Linear;
+
+impl<T> Interpolate<T> for Higher {
+    fn needs_lower(_q: N64, _len: usize) -> bool {
+        false
+    }
+    fn needs_higher(_q: N64, _len: usize) -> bool {
+        true
+    }
+    fn interpolate(_lower: Option<T>, higher: Option<T>, _q: N64, _len: usize) -> T {
+        higher.unwrap()
+    }
+}
+
+impl<T> Interpolate<T> for Lower {
+    fn needs_lower(_q: N64, _len: usize) -> bool {
+        true
+    }
+    fn needs_higher(_q: N64, _len: usize) -> bool {
+        false
+    }
+    fn interpolate(lower: Option<T>, _higher: Option<T>, _q: N64, _len: usize) -> T {
+        lower.unwrap()
+    }
+}
+
+impl<T> Interpolate<T> for Nearest {
+    fn needs_lower(q: N64, len: usize) -> bool {
+        float_quantile_index_fraction(q, len) < 0.5
+    }
+    fn needs_higher(q: N64, len: usize) -> bool {
+        !<Self as Interpolate<T>>::needs_lower(q, len)
+    }
+    fn interpolate(lower: Option<T>, higher: Option<T>, q: N64, len: usize) -> T {
+        if <Self as Interpolate<T>>::needs_lower(q, len) {
+            lower.unwrap()
+        } else {
+            higher.unwrap()
+        }
+    }
+}
+
+impl<T> Interpolate<T> for Midpoint
+where
+    T: NumOps + Clone + FromPrimitive,
+{
+    fn needs_lower(_q: N64, _len: usize) -> bool {
+        true
+    }
+    fn needs_higher(_q: N64, _len: usize) -> bool {
+        true
+    }
+    fn interpolate(lower: Option<T>, higher: Option<T>, _q: N64, _len: usize) -> T {
+        let denom = T::from_u8(2).unwrap();
+        let lower = lower.unwrap();
+        let higher = higher.unwrap();
+        lower.clone() + (higher.clone() - lower.clone()) / denom.clone()
+    }
+}
+
+impl<T> Interpolate<T> for Linear
+where
+    T: NumOps + Clone + FromPrimitive + ToPrimitive,
+{
+    fn needs_lower(_q: N64, _len: usize) -> bool {
+        true
+    }
+    fn needs_higher(_q: N64, _len: usize) -> bool {
+        true
+    }
+    fn interpolate(lower: Option<T>, higher: Option<T>, q: N64, len: usize) -> T {
+        let fraction = float_quantile_index_fraction(q, len).to_f64().unwrap();
+        let lower = lower.unwrap();
+        let higher = higher.unwrap();
+        let lower_f64 = lower.to_f64().unwrap();
+        let higher_f64 = higher.to_f64().unwrap();
+        lower.clone() + T::from_f64(fraction * (higher_f64 - lower_f64)).unwrap()
+    }
+}