diff --git a/src/maybe_nan/impl_not_none.rs b/src/maybe_nan/impl_not_none.rs index 22703084..0fdd8b4e 100644 --- a/src/maybe_nan/impl_not_none.rs +++ b/src/maybe_nan/impl_not_none.rs @@ -2,7 +2,7 @@ use super::NotNone; use num_traits::{FromPrimitive, ToPrimitive}; use std::cmp; use std::fmt; -use std::ops::{Add, Deref, DerefMut, Div, Mul, Sub}; +use std::ops::{Add, Deref, DerefMut, Div, Mul, Sub, Rem}; impl Deref for NotNone { type Target = T; @@ -96,6 +96,14 @@ impl Div for NotNone { } } +impl Rem for NotNone { + type Output = NotNone; + #[inline] + fn rem(self, rhs: Self) -> Self::Output { + self.map(|v| v.rem(rhs.unwrap())) + } +} + impl ToPrimitive for NotNone { #[inline] fn to_isize(&self) -> Option { diff --git a/src/quantile.rs b/src/quantile.rs index 9188aa2f..f92cdff7 100644 --- a/src/quantile.rs +++ b/src/quantile.rs @@ -8,8 +8,7 @@ use {MaybeNan, MaybeNanExt, Sort1dExt}; pub mod interpolate { use ndarray::azip; use ndarray::prelude::*; - use num_traits::{FromPrimitive, ToPrimitive}; - use std::ops::{Add, Div}; + use num_traits::{FromPrimitive, ToPrimitive, NumOps}; /// Used to provide an interpolation strategy to [`quantile_axis_mut`]. /// @@ -116,7 +115,7 @@ pub mod interpolate { impl Interpolate for Midpoint where - T: Add + Div + Clone + FromPrimitive, + T: NumOps + Clone + FromPrimitive, { fn needs_lower(_q: f64, _len: usize) -> bool { true @@ -134,13 +133,20 @@ pub mod interpolate { D: Dimension, { let denom = T::from_u8(2).unwrap(); - (lower.unwrap() + higher.unwrap()).mapv_into(|x| x / denom.clone()) + let mut lower = lower.unwrap(); + let higher = higher.unwrap(); + azip!( + mut lower, ref higher in { + *lower = lower.clone() + (higher.clone() - lower.clone()) / denom.clone() + } + ); + lower } } impl Interpolate for Linear where - T: Add + Clone + FromPrimitive + ToPrimitive, + T: NumOps + Clone + FromPrimitive + ToPrimitive, { fn needs_lower(_q: f64, _len: usize) -> bool { true diff --git a/tests/quantile.rs b/tests/quantile.rs index ea5ddaa7..ae85d277 100644 --- a/tests/quantile.rs +++ b/tests/quantile.rs @@ -6,6 +6,7 @@ use ndarray::prelude::*; use ndarray_stats::{ interpolate::{Higher, Linear, Lower, Midpoint, Nearest}, QuantileExt, + Quantile1dExt, }; #[test] @@ -148,3 +149,13 @@ fn test_quantile_axis_skipnan_mut_linear_opt_i32() { assert_eq!(q[0], Some(3)); assert!(q[1].is_none()); } + +#[test] +fn test_midpoint_overflow() { + // Regression test + // This triggered an overflow panic with a naive Midpoint implementation: (a+b)/2 + let mut a: Array1 = array![129, 130, 130, 131]; + let median = a.quantile_mut::(0.5).unwrap(); + let expected_median = 130; + assert_eq!(median, expected_median); +} \ No newline at end of file