Improve bulk quantiles

Cargo.toml

@@ -23,7 +23,12 @@ itertools = { version = "0.7.0", default-features = false }
 indexmap = "1.0"
 
 [dev-dependencies]
+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/histogram/strategies.rs

@@ -308,8 +308,8 @@ where
         let n_points = a.len();
 
         let mut a_copy = a.to_owned();
-        let first_quartile = a_copy.quantile_mut::<Nearest>(n64(0.25)).unwrap();
-        let third_quartile = a_copy.quantile_mut::<Nearest>(n64(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/quantile/interpolate.rs

@@ -1,6 +1,4 @@
 //! Interpolation strategies.
-use ndarray::azip;
-use ndarray::prelude::*;
 use noisy_float::types::N64;
 use num_traits::{Float, FromPrimitive, NumOps, ToPrimitive};
 
@@ -45,14 +43,7 @@ pub trait Interpolate<T> {
     /// **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<D>(
-        lower: Option<Array<T, D>>,
-        higher: Option<Array<T, D>>,
-        q: N64,
-        len: usize,
-    ) -> Array<T, D>
-    where
-        D: Dimension;
+    fn interpolate(lower: Option<T>, higher: Option<T>, q: N64, len: usize) -> T;
 }
 
 /// Select the higher value.
@@ -75,12 +66,7 @@ impl<T> Interpolate<T> for Higher {
     fn needs_higher(_q: N64, _len: usize) -> bool {
         true
     }
-    fn interpolate<D>(
-        _lower: Option<Array<T, D>>,
-        higher: Option<Array<T, D>>,
-        _q: N64,
-        _len: usize,
-    ) -> Array<T, D> {
+    fn interpolate(_lower: Option<T>, higher: Option<T>, _q: N64, _len: usize) -> T {
         higher.unwrap()
     }
 }
@@ -92,12 +78,7 @@ impl<T> Interpolate<T> for Lower {
     fn needs_higher(_q: N64, _len: usize) -> bool {
         false
     }
-    fn interpolate<D>(
-        lower: Option<Array<T, D>>,
-        _higher: Option<Array<T, D>>,
-        _q: N64,
-        _len: usize,
-    ) -> Array<T, D> {
+    fn interpolate(lower: Option<T>, _higher: Option<T>, _q: N64, _len: usize) -> T {
         lower.unwrap()
     }
 }
@@ -109,12 +90,7 @@ impl<T> Interpolate<T> for Nearest {
     fn needs_higher(q: N64, len: usize) -> bool {
         !<Self as Interpolate<T>>::needs_lower(q, len)
     }
-    fn interpolate<D>(
-        lower: Option<Array<T, D>>,
-        higher: Option<Array<T, D>>,
-        q: N64,
-        len: usize,
-    ) -> Array<T, D> {
+    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 {
@@ -133,24 +109,11 @@ where
     fn needs_higher(_q: N64, _len: usize) -> bool {
         true
     }
-    fn interpolate<D>(
-        lower: Option<Array<T, D>>,
-        higher: Option<Array<T, D>>,
-        _q: N64,
-        _len: usize,
-    ) -> Array<T, D>
-    where
-        D: Dimension,
-    {
+    fn interpolate(lower: Option<T>, higher: Option<T>, _q: N64, _len: usize) -> T {
         let denom = T::from_u8(2).unwrap();
-        let mut lower = lower.unwrap();
+        let lower = lower.unwrap();
         let higher = higher.unwrap();
-        azip!(
-            mut lower, ref higher in {
-                *lower = lower.clone() + (higher.clone() - lower.clone()) / denom.clone()
-            }
-        );
-        lower
+        lower.clone() + (higher.clone() - lower.clone()) / denom.clone()
     }
 }
 
@@ -164,23 +127,12 @@ where
     fn needs_higher(_q: N64, _len: usize) -> bool {
         true
     }
-    fn interpolate<D>(
-        lower: Option<Array<T, D>>,
-        higher: Option<Array<T, D>>,
-        q: N64,
-        len: usize,
-    ) -> Array<T, D>
-    where
-        D: Dimension,
-    {
+    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 mut a = lower.unwrap();
-        let b = higher.unwrap();
-        azip!(mut a, ref b in {
-            let a_f64 = a.to_f64().unwrap();
-            let b_f64 = b.to_f64().unwrap();
-            *a = a.clone() + T::from_f64(fraction * (b_f64 - a_f64)).unwrap();
-        });
-        a
+        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()
     }
 }