ndarray_stats/summary_statistics/
means.rs

1use super::SummaryStatisticsExt;
2use crate::errors::{EmptyInput, MultiInputError, ShapeMismatch};
3use ndarray::{Array, ArrayBase, Axis, Data, Dimension, Ix1, RemoveAxis};
4use num_integer::IterBinomial;
5use num_traits::{Float, FromPrimitive, Zero};
6use std::ops::{Add, AddAssign, Div, Mul};
7
8impl<A, S, D> SummaryStatisticsExt<A, S, D> for ArrayBase<S, D>
9where
10    S: Data<Elem = A>,
11    D: Dimension,
12{
13    fn mean(&self) -> Result<A, EmptyInput>
14    where
15        A: Clone + FromPrimitive + Add<Output = A> + Div<Output = A> + Zero,
16    {
17        let n_elements = self.len();
18        if n_elements == 0 {
19            Err(EmptyInput)
20        } else {
21            let n_elements = A::from_usize(n_elements)
22                .expect("Converting number of elements to `A` must not fail.");
23            Ok(self.sum() / n_elements)
24        }
25    }
26
27    fn weighted_mean(&self, weights: &Self) -> Result<A, MultiInputError>
28    where
29        A: Copy + Div<Output = A> + Mul<Output = A> + Zero,
30    {
31        return_err_if_empty!(self);
32        let weighted_sum = self.weighted_sum(weights)?;
33        Ok(weighted_sum / weights.sum())
34    }
35
36    fn weighted_sum(&self, weights: &ArrayBase<S, D>) -> Result<A, MultiInputError>
37    where
38        A: Copy + Mul<Output = A> + Zero,
39    {
40        return_err_unless_same_shape!(self, weights);
41        Ok(self
42            .iter()
43            .zip(weights)
44            .fold(A::zero(), |acc, (&d, &w)| acc + d * w))
45    }
46
47    fn weighted_mean_axis(
48        &self,
49        axis: Axis,
50        weights: &ArrayBase<S, Ix1>,
51    ) -> Result<Array<A, D::Smaller>, MultiInputError>
52    where
53        A: Copy + Div<Output = A> + Mul<Output = A> + Zero,
54        D: RemoveAxis,
55    {
56        return_err_if_empty!(self);
57        let mut weighted_sum = self.weighted_sum_axis(axis, weights)?;
58        let weights_sum = weights.sum();
59        weighted_sum.mapv_inplace(|v| v / weights_sum);
60        Ok(weighted_sum)
61    }
62
63    fn weighted_sum_axis(
64        &self,
65        axis: Axis,
66        weights: &ArrayBase<S, Ix1>,
67    ) -> Result<Array<A, D::Smaller>, MultiInputError>
68    where
69        A: Copy + Mul<Output = A> + Zero,
70        D: RemoveAxis,
71    {
72        if self.shape()[axis.index()] != weights.len() {
73            return Err(MultiInputError::ShapeMismatch(ShapeMismatch {
74                first_shape: self.shape().to_vec(),
75                second_shape: weights.shape().to_vec(),
76            }));
77        }
78
79        // We could use `lane.weighted_sum` here, but we're avoiding 2
80        // conditions and an unwrap per lane.
81        Ok(self.map_axis(axis, |lane| {
82            lane.iter()
83                .zip(weights)
84                .fold(A::zero(), |acc, (&d, &w)| acc + d * w)
85        }))
86    }
87
88    fn harmonic_mean(&self) -> Result<A, EmptyInput>
89    where
90        A: Float + FromPrimitive,
91    {
92        self.map(|x| x.recip())
93            .mean()
94            .map(|x| x.recip())
95            .ok_or(EmptyInput)
96    }
97
98    fn geometric_mean(&self) -> Result<A, EmptyInput>
99    where
100        A: Float + FromPrimitive,
101    {
102        self.map(|x| x.ln())
103            .mean()
104            .map(|x| x.exp())
105            .ok_or(EmptyInput)
106    }
107
108    fn weighted_var(&self, weights: &Self, ddof: A) -> Result<A, MultiInputError>
109    where
110        A: AddAssign + Float + FromPrimitive,
111    {
112        return_err_if_empty!(self);
113        return_err_unless_same_shape!(self, weights);
114        let zero = A::from_usize(0).expect("Converting 0 to `A` must not fail.");
115        let one = A::from_usize(1).expect("Converting 1 to `A` must not fail.");
116        assert!(
117            !(ddof < zero || ddof > one),
118            "`ddof` must not be less than zero or greater than one",
119        );
120        inner_weighted_var(self, weights, ddof, zero)
121    }
122
123    fn weighted_std(&self, weights: &Self, ddof: A) -> Result<A, MultiInputError>
124    where
125        A: AddAssign + Float + FromPrimitive,
126    {
127        Ok(self.weighted_var(weights, ddof)?.sqrt())
128    }
129
130    fn weighted_var_axis(
131        &self,
132        axis: Axis,
133        weights: &ArrayBase<S, Ix1>,
134        ddof: A,
135    ) -> Result<Array<A, D::Smaller>, MultiInputError>
136    where
137        A: AddAssign + Float + FromPrimitive,
138        D: RemoveAxis,
139    {
140        return_err_if_empty!(self);
141        if self.shape()[axis.index()] != weights.len() {
142            return Err(MultiInputError::ShapeMismatch(ShapeMismatch {
143                first_shape: self.shape().to_vec(),
144                second_shape: weights.shape().to_vec(),
145            }));
146        }
147        let zero = A::from_usize(0).expect("Converting 0 to `A` must not fail.");
148        let one = A::from_usize(1).expect("Converting 1 to `A` must not fail.");
149        assert!(
150            !(ddof < zero || ddof > one),
151            "`ddof` must not be less than zero or greater than one",
152        );
153
154        // `weights` must be a view because `lane` is a view in this context.
155        let weights = weights.view();
156        Ok(self.map_axis(axis, |lane| {
157            inner_weighted_var(&lane, &weights, ddof, zero).unwrap()
158        }))
159    }
160
161    fn weighted_std_axis(
162        &self,
163        axis: Axis,
164        weights: &ArrayBase<S, Ix1>,
165        ddof: A,
166    ) -> Result<Array<A, D::Smaller>, MultiInputError>
167    where
168        A: AddAssign + Float + FromPrimitive,
169        D: RemoveAxis,
170    {
171        Ok(self
172            .weighted_var_axis(axis, weights, ddof)?
173            .mapv_into(|x| x.sqrt()))
174    }
175
176    fn kurtosis(&self) -> Result<A, EmptyInput>
177    where
178        A: Float + FromPrimitive,
179    {
180        let central_moments = self.central_moments(4)?;
181        Ok(central_moments[4] / central_moments[2].powi(2))
182    }
183
184    fn skewness(&self) -> Result<A, EmptyInput>
185    where
186        A: Float + FromPrimitive,
187    {
188        let central_moments = self.central_moments(3)?;
189        Ok(central_moments[3] / central_moments[2].sqrt().powi(3))
190    }
191
192    fn central_moment(&self, order: u16) -> Result<A, EmptyInput>
193    where
194        A: Float + FromPrimitive,
195    {
196        if self.is_empty() {
197            return Err(EmptyInput);
198        }
199        match order {
200            0 => Ok(A::one()),
201            1 => Ok(A::zero()),
202            n => {
203                let mean = self.mean().unwrap();
204                let shifted_array = self.mapv(|x| x - mean);
205                let shifted_moments = moments(shifted_array, n);
206                let correction_term = -shifted_moments[1];
207
208                let coefficients = central_moment_coefficients(&shifted_moments);
209                Ok(horner_method(coefficients, correction_term))
210            }
211        }
212    }
213
214    fn central_moments(&self, order: u16) -> Result<Vec<A>, EmptyInput>
215    where
216        A: Float + FromPrimitive,
217    {
218        if self.is_empty() {
219            return Err(EmptyInput);
220        }
221        match order {
222            0 => Ok(vec![A::one()]),
223            1 => Ok(vec![A::one(), A::zero()]),
224            n => {
225                // We only perform these operations once, and then reuse their
226                // result to compute all the required moments
227                let mean = self.mean().unwrap();
228                let shifted_array = self.mapv(|x| x - mean);
229                let shifted_moments = moments(shifted_array, n);
230                let correction_term = -shifted_moments[1];
231
232                let mut central_moments = vec![A::one(), A::zero()];
233                for k in 2..=n {
234                    let coefficients =
235                        central_moment_coefficients(&shifted_moments[..=(k as usize)]);
236                    let central_moment = horner_method(coefficients, correction_term);
237                    central_moments.push(central_moment)
238                }
239                Ok(central_moments)
240            }
241        }
242    }
243
244    private_impl! {}
245}
246
247/// Private function for `weighted_var` without conditions and asserts.
248fn inner_weighted_var<A, S, D>(
249    arr: &ArrayBase<S, D>,
250    weights: &ArrayBase<S, D>,
251    ddof: A,
252    zero: A,
253) -> Result<A, MultiInputError>
254where
255    S: Data<Elem = A>,
256    A: AddAssign + Float + FromPrimitive,
257    D: Dimension,
258{
259    let mut weight_sum = zero;
260    let mut mean = zero;
261    let mut s = zero;
262    for (&x, &w) in arr.iter().zip(weights.iter()) {
263        weight_sum += w;
264        let x_minus_mean = x - mean;
265        mean += (w / weight_sum) * x_minus_mean;
266        s += w * x_minus_mean * (x - mean);
267    }
268    Ok(s / (weight_sum - ddof))
269}
270
271/// Returns a vector containing all moments of the array elements up to
272/// *order*, where the *p*-th moment is defined as:
273///
274/// ```text
275/// 1  n
276/// ―  ∑ xᵢᵖ
277/// n i=1
278/// ```
279///
280/// The returned moments are ordered by power magnitude: 0th moment, 1st moment, etc.
281///
282/// **Panics** if `A::from_usize()` fails to convert the number of elements in the array.
283fn moments<A, S, D>(a: ArrayBase<S, D>, order: u16) -> Vec<A>
284where
285    A: Float + FromPrimitive,
286    S: Data<Elem = A>,
287    D: Dimension,
288{
289    let n_elements =
290        A::from_usize(a.len()).expect("Converting number of elements to `A` must not fail");
291    let order = i32::from(order);
292
293    // When k=0, we are raising each element to the 0th power
294    // No need to waste CPU cycles going through the array
295    let mut moments = vec![A::one()];
296
297    if order >= 1 {
298        // When k=1, we don't need to raise elements to the 1th power (identity)
299        moments.push(a.sum() / n_elements)
300    }
301
302    for k in 2..=order {
303        moments.push(a.map(|x| x.powi(k)).sum() / n_elements)
304    }
305    moments
306}
307
308/// Returns the coefficients in the polynomial expression to compute the *p*th
309/// central moment as a function of the sample mean.
310///
311/// It takes as input all moments up to order *p*, ordered by power magnitude - *p* is
312/// inferred to be the length of the *moments* array.
313fn central_moment_coefficients<A>(moments: &[A]) -> Vec<A>
314where
315    A: Float + FromPrimitive,
316{
317    let order = moments.len();
318    IterBinomial::new(order)
319        .zip(moments.iter().rev())
320        .map(|(binom, &moment)| A::from_usize(binom).unwrap() * moment)
321        .collect()
322}
323
324/// Uses [Horner's method] to evaluate a polynomial with a single indeterminate.
325///
326/// Coefficients are expected to be sorted by ascending order
327/// with respect to the indeterminate's exponent.
328///
329/// If the array is empty, `A::zero()` is returned.
330///
331/// Horner's method can evaluate a polynomial of order *n* with a single indeterminate
332/// using only *n-1* multiplications and *n-1* sums - in terms of number of operations,
333/// this is an optimal algorithm for polynomial evaluation.
334///
335/// [Horner's method]: https://en.wikipedia.org/wiki/Horner%27s_method
336fn horner_method<A>(coefficients: Vec<A>, indeterminate: A) -> A
337where
338    A: Float,
339{
340    let mut result = A::zero();
341    for coefficient in coefficients.into_iter().rev() {
342        result = coefficient + indeterminate * result
343    }
344    result
345}