1use crate::{
2 core::prelude::*,
3 errors::prelude::*,
4 extensions::prelude::*,
5 numeric::prelude::*,
6 validators::prelude::*,
7};
8
9pub trait ArraySumProdDiff<N: NumericOps> where Self: Sized + Clone {
11
12 fn prod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
31
32 fn sum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
51
52 fn nanprod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
71
72 fn nansum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
91
92 fn cumprod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
111
112 fn cumsum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
131
132 fn nancumprod(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
151
152 fn nancumsum(&self, axis: Option<isize>) -> Result<Array<N>, ArrayError>;
171
172 fn diff(&self, n: usize, axis: Option<isize>, prepend: Option<Array<N>>, append: Option<Array<N>>) -> Result<Array<N>, ArrayError>;
194
195 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}