ndarray_stats/summary_statistics/
means.rs

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