arr_rs/math/operations/
sum_prod_diff.rs

1use crate::{
2    core::prelude::*,
3    errors::prelude::*,
4    extensions::prelude::*,
5    numeric::prelude::*,
6    validators::prelude::*,
7};
8
9/// `ArrayTrait` - Array Sum, Product, Diff functions
10pub trait ArraySumProdDiff<N: NumericOps> where Self: Sized + Clone {
11
12    /// Multiplication of array elements
13    ///
14    /// # Arguments
15    ///
16    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
17    ///
18    /// # Examples
19    ///
20    /// ```
21    /// use arr_rs::prelude::*;
22    ///
23    /// let arr = Array::flat(vec![1, 2, 3, 4]);
24    /// assert_eq!(Array::single(24), arr.prod(None));
25    /// ```
26    ///
27    /// # Errors
28    ///
29    /// may returns `ArrayError`
30    fn prod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
31
32    /// Sum of array elements
33    ///
34    /// # Arguments
35    ///
36    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
37    ///
38    /// # Examples
39    ///
40    /// ```
41    /// use arr_rs::prelude::*;
42    ///
43    /// let arr = Array::flat(vec![1, 2, 3, 4]);
44    /// assert_eq!(Array::single(10), arr.sum(None));
45    /// ```
46    ///
47    /// # Errors
48    ///
49    /// may returns `ArrayError`
50    fn sum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
51
52    /// Sum of array elements treating NaN as one
53    ///
54    /// # Arguments
55    ///
56    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
57    ///
58    /// # Examples
59    ///
60    /// ```
61    /// use arr_rs::prelude::*;
62    ///
63    /// let arr = Array::flat(vec![1., 2., 3., 4., f64::NAN]);
64    /// assert_eq!(Array::single(24.), arr.nanprod(None));
65    /// ```
66    ///
67    /// # Errors
68    ///
69    /// may returns `ArrayError`
70    fn nanprod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
71
72    /// Sum of array elements treating NaN as zero
73    ///
74    /// # Arguments
75    ///
76    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
77    ///
78    /// # Examples
79    ///
80    /// ```
81    /// use arr_rs::prelude::*;
82    ///
83    /// let arr = Array::flat(vec![1., 2., 3., 4., f64::NAN]);
84    /// assert_eq!(Array::single(10.), arr.nansum(None));
85    /// ```
86    ///
87    /// # Errors
88    ///
89    /// may returns `ArrayError`
90    fn nansum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
91
92    /// Cumulative product of array elements
93    ///
94    /// # Arguments
95    ///
96    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
97    ///
98    /// # Examples
99    ///
100    /// ```
101    /// use arr_rs::prelude::*;
102    ///
103    /// let arr = Array::flat(vec![1, 2, 3, 4]);
104    /// assert_eq!(Array::flat(vec![1, 2, 6, 24]), arr.cumprod(None));
105    /// ```
106    ///
107    /// # Errors
108    ///
109    /// may returns `ArrayError`
110    fn cumprod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
111
112    /// Cumulative sum of array elements
113    ///
114    /// # Arguments
115    ///
116    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
117    ///
118    /// # Examples
119    ///
120    /// ```
121    /// use arr_rs::prelude::*;
122    ///
123    /// let arr = Array::flat(vec![1, 2, 3, 4]);
124    /// assert_eq!(Array::flat(vec![1, 3, 6, 10]), arr.cumsum(None));
125    /// ```
126    ///
127    /// # Errors
128    ///
129    /// may returns `ArrayError`
130    fn cumsum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
131
132    /// Cumulative product of array elements
133    ///
134    /// # Arguments
135    ///
136    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
137    ///
138    /// # Examples
139    ///
140    /// ```
141    /// use arr_rs::prelude::*;
142    ///
143    /// let arr = Array::flat(vec![1., 2., 3., 4., f64::NAN]);
144    /// assert_eq!(Array::flat(vec![1., 2., 6., 24., 24.]), arr.nancumprod(None));
145    /// ```
146    ///
147    /// # Errors
148    ///
149    /// may returns `ArrayError`
150    fn nancumprod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
151
152    /// Cumulative sum of array elements
153    ///
154    /// # Arguments
155    ///
156    /// * `axis` - the axis along which to execute the function. optional. if negative, counts from last to first axis. if None, array is raveled
157    ///
158    /// # Examples
159    ///
160    /// ```
161    /// use arr_rs::prelude::*;
162    ///
163    /// let arr = Array::flat(vec![1., 2., 3., 4., f64::NAN]);
164    /// assert_eq!(Array::flat(vec![1., 3., 6., 10., 10.]), arr.nancumsum(None));
165    /// ```
166    ///
167    /// # Errors
168    ///
169    /// may returns `ArrayError`
170    fn nancumsum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
171
172    /// The differences between consecutive elements of an array
173    ///
174    /// # Arguments
175    ///
176    /// * `n` - number of times values are differenced
177    /// * `axis` - the axis along which to execute the function. optional. defaults to last axis
178    /// * `append` - number(s) to append at the end along axis prior to performing the difference
179    /// * `prepend` - number(s) to append at the beginning along axis prior to performing the difference
180    ///
181    /// # Examples
182    ///
183    /// ```
184    /// use arr_rs::prelude::*;
185    ///
186    /// let arr = Array::flat(vec![1., 2., 4., 4., 7.]);
187    /// assert_eq!(Array::flat(vec![1., 2., 0., 3.]), arr.diff(1, None, None, None));
188    /// ```
189    ///
190    /// # Errors
191    ///
192    /// may returns `ArrayError`
193    fn diff(&self, n: usize, axis: Option<isize>, prepend: Option<Array<N>>, append: Option<Array<N>>) -> Result<Array<N>, ArrayError>;
194
195    /// The differences between consecutive elements of an array
196    ///
197    /// # Arguments
198    ///
199    /// * `to_end` - number(s) to append at the end of the returned differences
200    /// * `to_begin` - number(s) to append at the beginning of the returned differences
201    ///
202    /// # Examples
203    ///
204    /// ```
205    /// use arr_rs::prelude::*;
206    ///
207    /// let arr = Array::flat(vec![1., 2., 4., 4., 7.]);
208    /// assert_eq!(Array::flat(vec![1., 2., 0., 3.]), arr.ediff1d(None, None));
209    /// ```
210    ///
211    /// # Errors
212    ///
213    /// may returns `ArrayError`
214    fn ediff1d(&self, to_end: Option<Array<N>>, to_begin: Option<Array<N>>) -> Result<Array<N>, ArrayError>;
215}
216
217impl <N: NumericOps> ArraySumProdDiff<N> for Array<N> {
218
219    fn prod(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
220        if let Some(axis) = axis {
221            let axis = self.normalize_axis(axis);
222            let result = self.apply_along_axis(axis, |arr| arr.prod(None));
223            result.reshape(&result.get_shape()?.remove_at_if(axis, result.ndim()? > 1))
224        } else {
225            Self::single(self.elements.iter().fold(N::one(), |acc, &x| acc * x))
226        }
227    }
228
229    fn sum(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
230        if let Some(axis) = axis {
231            let axis = self.normalize_axis(axis);
232            let result = self.apply_along_axis(axis, |arr| arr.sum(None));
233            result.reshape(&result.get_shape()?.remove_at_if(axis, result.ndim()? > 1))
234        } else {
235            Self::single(self.elements.iter().fold(N::zero(), |acc, &x| acc + x))
236        }
237    }
238
239    fn nanprod(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
240        if let Some(axis) = axis {
241            let axis = self.normalize_axis(axis);
242            let result = self.apply_along_axis(axis, |arr| arr.nanprod(None));
243                result.reshape(&result.get_shape()?.remove_at_if(axis, result.ndim()? > 1))
244        } else {
245            Self::single(self.elements.iter().fold(N::one(), |acc, &x|
246                acc * if x.to_f64().is_nan() { N::one() } else { x }
247            ))
248        }
249    }
250
251    fn nansum(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
252        if let Some(axis) = axis {
253            let axis = self.normalize_axis(axis);
254            let result = self.apply_along_axis(axis, |arr| arr.nansum(None));
255                result.reshape(&result.get_shape()?.remove_at_if(axis, result.ndim()? > 1))
256        } else {
257            Self::single(self.elements.iter().fold(N::zero(), |acc, &x|
258                acc + if x.to_f64().is_nan() { N::zero() } else { x }
259            ))
260        }
261    }
262
263    fn cumprod(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
264        if let Some(axis) = axis {
265            let axis = self.normalize_axis(axis);
266            self.apply_along_axis(axis, |arr| arr.cumprod(None))
267        } else {
268            let mut acc = N::one();
269            self.ravel()?.map(|&x| {
270                acc *= x;
271                acc
272            })
273        }
274    }
275
276    fn cumsum(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
277        if let Some(axis) = axis {
278            let axis = self.normalize_axis(axis);
279            self.apply_along_axis(axis, |arr| arr.cumsum(None))
280        } else {
281            let mut acc = N::zero();
282            self.ravel()?.map(|&x| {
283                acc += x;
284                acc
285            })
286        }
287    }
288
289    fn nancumprod(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
290        if let Some(axis) = axis {
291            let axis = self.normalize_axis(axis);
292            self.apply_along_axis(axis, |arr| arr.nancumprod(None))
293        } else {
294            let mut acc = N::one();
295            self.ravel()?.map(|&x| {
296                acc *= if x.to_f64().is_nan() { N::one() } else { x };
297                acc
298            })
299        }
300    }
301
302    fn nancumsum(&self, axis: Option<isize>) -> Result<Self, ArrayError> {
303        if let Some(axis) = axis {
304            let axis = self.normalize_axis(axis);
305            self.apply_along_axis(axis, |arr| arr.nancumsum(None))
306        } else {
307            let mut acc = N::zero();
308            self.ravel()?.map(|&x| {
309                acc += if x.to_f64().is_nan() { N::zero() } else { x };
310                acc
311            })
312        }
313    }
314
315    fn diff(&self, n: usize, axis: Option<isize>, prepend: Option<Self>, append: Option<Self>) -> Result<Self, ArrayError> {
316        if n == 0 {
317            Self::empty()
318        } else if self.ndim()? == 1 {
319            let mut elements = prepend.unwrap_or(Self::empty()?).get_elements()?;
320            elements.extend_from_slice(&self.get_elements()?);
321            elements.extend_from_slice(&append.unwrap_or(Self::empty()?).get_elements()?);
322            for _ in 0..n { elements = Self::flat(elements.clone()).ediff1d(None, None).get_elements()? }
323            Self::flat(elements)
324        } else {
325            fn diff_extend_partial<N: Numeric>(array: &Array<N>, partial: Vec<Array<N>>, other: Option<Array<N>>, axis: usize, rev: bool) -> Result<Vec<Array<N>>, ArrayError> {
326                if other.is_none() {return Ok(partial) }
327                let other = other.unwrap();
328                array.ndim()?.is_equal(&other.ndim()?)?;
329                array.get_shape()?.remove_at(axis).is_equal(&other.get_shape()?.remove_at(axis))?;
330                let p_partial = other
331                    .moveaxis(vec![axis.to_isize()], vec![array.ndim()?.to_isize()])
332                    .ravel().split(other.get_shape()?.remove_at(axis).into_iter().product(), None)?;
333                let mut tmp_v = vec![partial, p_partial];
334                if rev { tmp_v.reverse() };
335                let result = tmp_v[0].clone().into_iter().zip(&tmp_v[1]).map(|(arr, other)| {
336                    let mut elements = other.elements.clone();
337                    elements.extend_from_slice(&arr.elements);
338                    Array::flat(elements).unwrap()
339                }).collect::<Vec<Array<N>>>();
340                Ok(result)
341            }
342
343            let axis = axis.unwrap_or(-1);
344            let axis = self.normalize_axis(axis);
345
346            let parts = self.get_shape()?.remove_at(axis).into_iter().product();
347            let mut partial = self
348                .moveaxis(vec![axis.to_isize()], vec![self.ndim()?.to_isize()])?
349                .ravel().split(parts, None)?;
350
351            partial = diff_extend_partial(self, partial, prepend.clone(), axis, false)?;
352            partial = diff_extend_partial(self, partial, append.clone(), axis, true)?;
353
354            let mut new_shape = self.get_shape()?;
355            if let Some(p) = prepend { new_shape[axis] += p.get_shape()?[axis] };
356            if let Some(a) = append { new_shape[axis] += a.get_shape()?[axis] };
357
358            let array = partial.into_iter()
359                .flatten()
360                .collect::<Self<>>()
361                .reshape(&new_shape.swap_ext(axis, self.ndim()? - 1));
362            let array =
363                if axis == 0 { array.transpose(None) }
364                else { array.moveaxis(vec![axis.to_isize()], vec![self.ndim()?.to_isize()]) };
365
366            array.apply_along_axis(axis, |arr| arr.diff(n, None, None, None))
367        }
368    }
369
370    fn ediff1d(&self, to_end: Option<Self>, to_begin: Option<Self>) -> Result<Self, ArrayError> {
371        let array = self.ravel()?;
372        let (to_end, to_begin) = (to_end.unwrap_or(Self::empty()?), to_begin.unwrap_or(Self::empty()?));
373        let diffs = (1..array.len()?).map(|i| array[i] - array[i - 1]).collect::<Vec<N>>();
374        let mut result = to_begin.get_elements()?;
375        result.extend_from_slice(&diffs);
376        result.extend_from_slice(&to_end.get_elements()?);
377        Self::flat(result)
378    }
379}
380
381impl <N: NumericOps> ArraySumProdDiff<N> for Result<Array<N>, ArrayError> {
382
383    fn prod(&self, axis: Option<isize>) -> Self {
384        self.clone()?.prod(axis)
385    }
386
387    fn sum(&self, axis: Option<isize>) -> Self {
388        self.clone()?.sum(axis)
389    }
390
391    fn nanprod(&self, axis: Option<isize>) -> Self {
392        self.clone()?.nanprod(axis)
393    }
394
395    fn nansum(&self, axis: Option<isize>) -> Self {
396        self.clone()?.nansum(axis)
397    }
398
399    fn cumprod(&self, axis: Option<isize>) -> Self {
400        self.clone()?.cumprod(axis)
401    }
402
403    fn cumsum(&self, axis: Option<isize>) -> Self {
404        self.clone()?.cumsum(axis)
405    }
406
407    fn nancumprod(&self, axis: Option<isize>) -> Self {
408        self.clone()?.nancumprod(axis)
409    }
410
411    fn nancumsum(&self, axis: Option<isize>) -> Self {
412        self.clone()?.nancumsum(axis)
413    }
414
415    fn diff(&self, n: usize, axis: Option<isize>, prepend: Option<Array<N>>, append: Option<Array<N>>) -> Self {
416        self.clone()?.diff(n, axis, prepend, append)
417    }
418
419    fn ediff1d(&self, to_end: Option<Array<N>>, to_begin: Option<Array<N>>) -> Self {
420        self.clone()?.ediff1d(to_end, to_begin)
421    }
422}