arr_rs/core/operations/
tiling.rs

1use crate::{
2    core::prelude::*,
3    errors::prelude::*,
4    extensions::prelude::*,
5};
6use crate::prelude::Numeric;
7
8/// `ArrayTrait` - Array Tiling functions
9pub trait ArrayTiling<T: ArrayElement> where Self: Sized + Clone {
10
11    /// Repeat each element of an array after themselves
12    ///
13    /// # Arguments
14    ///
15    /// * `repeats` - number of repetitions for each element, broadcasted to fit the shape of the given axis
16    /// * `axis` - the axis along which to repeat. optional, if None, array is flattened
17    ///
18    /// # Examples
19    /// ```
20    /// use arr_rs::prelude::*;
21    ///
22    /// let arr = Array::<i32>::single(3);
23    /// assert_eq!(array![i32, 3, 3, 3, 3], arr.repeat(&vec![4], None));
24    ///
25    /// let arr = Array::<i32>::new(vec![1, 2, 3, 4], vec![2, 2]);
26    /// assert_eq!(array!(i32, [[1, 2], [3, 4], [3, 4]]), arr.repeat(&vec![1, 2], Some(0)));
27    /// ```
28    ///
29    /// # Errors
30    ///
31    /// may returns `ArrayError`
32    fn repeat(&self, repeats: &[usize], axis: Option<usize>) -> Result<Array<T>, ArrayError>;
33}
34
35impl <T: ArrayElement> ArrayTiling<T> for Array<T> {
36
37    fn repeat(&self, repeats: &[usize], axis: Option<usize>) -> Result<Self, ArrayError> {
38        if let Some(axis) = axis {
39            let repeats = repeats.to_vec().to_array()?.broadcast_to(vec![self.get_shape()?[axis]]).get_elements()?;
40            let new_axis_len = repeats.clone().into_iter().sum();
41            let new_shape = self.get_shape()?.update_at(axis, new_axis_len);
42            let tmp_shape = new_shape.clone().swap_ext(0, axis);
43            let partial = self.split(self.get_shape()?[axis], Some(axis))?.into_iter()
44                .zip(&repeats)
45                .flat_map(|(el, &rep)| vec![el; rep])
46                .flatten()
47                .collect::<Self>();
48            partial.reshape(&tmp_shape)
49                .moveaxis(vec![0], vec![axis.to_isize()])
50                .reshape(&new_shape)
51        } else {
52            let result = self.get_elements()?.into_iter()
53                .zip(&repeats.to_vec().to_array()?.broadcast_to(self.get_shape()?).get_elements()?)
54                .flat_map(|(el, &rep)| vec![el; rep])
55                .collect();
56            Self::flat(result)
57        }
58    }
59}
60
61impl <T: ArrayElement> ArrayTiling<T> for Result<Array<T>, ArrayError> {
62
63    fn repeat(&self, repeats: &[usize], axis: Option<usize>) -> Self {
64        self.clone()?.repeat(repeats, axis)
65    }
66}