concision_core/ops/
pad.rs

1/*
2   Appellation: pad <mod>
3   Contrib: FL03 <jo3mccain@icloud.com>
4*/
5pub use self::{action::PadAction, mode::PadMode, utils::*};
6
7pub(crate) mod action;
8pub(crate) mod mode;
9
10use nd::{Array, ArrayBase, DataOwned, Dimension};
11use num::traits::{FromPrimitive, Num};
12
13pub trait Pad<T> {
14    type Output;
15
16    fn pad(&self, mode: PadMode<T>, pad: &[[usize; 2]]) -> Self::Output;
17}
18
19impl<A, S, D> Pad<A> for ArrayBase<S, D>
20where
21    A: Copy + FromPrimitive + Num,
22    D: Dimension,
23    S: DataOwned<Elem = A>,
24{
25    type Output = Array<A, D>;
26
27    fn pad(&self, mode: PadMode<A>, pad: &[[usize; 2]]) -> Self::Output {
28        self::utils::pad(self, pad, mode)
29    }
30}
31
32pub struct Padding<T> {
33    pub(crate) action: PadAction,
34    pub(crate) mode: PadMode<T>,
35    pub(crate) pad: Vec<[usize; 2]>,
36    pub(crate) padding: usize,
37}
38
39impl<T> Padding<T> {
40    pub fn new() -> Self {
41        Self {
42            action: PadAction::default(),
43            mode: PadMode::default(),
44            pad: Vec::new(),
45            padding: 0,
46        }
47    }
48
49    pub fn pad(&self) -> &[[usize; 2]] {
50        &self.pad
51    }
52
53    pub fn with_action(mut self, action: PadAction) -> Self {
54        self.action = action;
55        self
56    }
57
58    pub fn with_mode(mut self, mode: PadMode<T>) -> Self {
59        self.mode = mode;
60        self
61    }
62
63    pub fn with_padding(mut self, padding: usize) -> Self {
64        self.padding = padding;
65        self
66    }
67}
68
69mod utils {
70    #![cfg(any(feature = "std", feature = "alloc"))]
71    use super::{PadAction, PadMode};
72    use crate::traits::ArrayLike;
73    use nd::{Array, ArrayBase, AxisDescription, Data, DataOwned, Dimension, Slice};
74    use num::{FromPrimitive, Num};
75
76    #[cfg(all(feature = "alloc", no_std))]
77    use alloc::borrow::Cow;
78    #[cfg(feature = "std")]
79    use std::borrow::Cow;
80
81    fn reader(nb_dim: usize, pad: &[[usize; 2]]) -> Cow<[[usize; 2]]> {
82        if pad.len() == 1 && pad.len() < nb_dim {
83            // The user provided a single padding for all dimensions
84            Cow::from(vec![pad[0]; nb_dim])
85        } else if pad.len() == nb_dim {
86            Cow::from(pad)
87        } else {
88            panic!("Inconsistant number of dimensions and pad arrays");
89        }
90    }
91
92    pub fn pad<A, S, D>(data: &ArrayBase<S, D>, pad: &[[usize; 2]], mode: PadMode<A>) -> Array<A, D>
93    where
94        A: Copy + FromPrimitive + Num,
95        D: Dimension,
96        S: DataOwned<Elem = A>,
97    {
98        let pad = reader(data.ndim(), pad);
99        let mut new_dim = data.raw_dim();
100        for (ax, (&ax_len, pad)) in data.shape().iter().zip(pad.iter()).enumerate() {
101            new_dim[ax] = ax_len + pad[0] + pad[1];
102        }
103
104        // let mut padded = array_like(&data, new_dim, mode.init());
105        let mut padded = data.array_like(new_dim, mode.init()).to_owned();
106        pad_to(data, &pad, mode, &mut padded);
107        padded
108    }
109
110    pub fn pad_to<A, S, D>(
111        data: &ArrayBase<S, D>,
112        pad: &[[usize; 2]],
113        mode: PadMode<A>,
114        output: &mut Array<A, D>,
115    ) where
116        A: Copy + FromPrimitive + Num,
117        D: Dimension,
118        S: Data<Elem = A>,
119    {
120        let pad = reader(data.ndim(), pad);
121
122        // Select portion of padded array that needs to be copied from the original array.
123        output
124            .slice_each_axis_mut(|ad| {
125                let AxisDescription { axis, len, .. } = ad;
126                let pad = pad[axis.index()];
127                Slice::from(pad[0]..len - pad[1])
128            })
129            .assign(data);
130
131        match mode.action() {
132            PadAction::StopAfterCopy => { /* Nothing */ }
133            _ => unimplemented!(),
134        }
135    }
136}