concision_core/ops/
pad.rs1pub 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 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 = 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 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 => { }
133 _ => unimplemented!(),
134 }
135 }
136}