arr_rs/core/operations/
tiling.rs1use crate::{
2 core::prelude::*,
3 errors::prelude::*,
4 extensions::prelude::*,
5};
6use crate::prelude::Numeric;
7
8pub trait ArrayTiling<T: ArrayElement> where Self: Sized + Clone {
10
11 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}