concision_core/utils/pad/
impl_pad.rs

1/*
2    Appellation: impl_pad <module>
3    Created At: 2025.11.26:16:12:28
4    Contrib: @FL03
5*/
6use super::{Pad, PadAction, PadMode};
7use concision_traits::ArrayLike;
8use ndarray::{Array, ArrayBase, AxisDescription, DataOwned, Dimension, Slice};
9use num_traits::{FromPrimitive, Num};
10
11fn reader(ndim: usize, pad: &[[usize; 2]]) -> Option<Vec<[usize; 2]>> {
12    debug_assert!(pad.len() == ndim, "Inconsistent dimensions for padding");
13    if pad.len() != ndim {
14        return None;
15    }
16    Some(pad.to_vec())
17}
18
19fn apply_padding<A, S, D>(
20    data: &ArrayBase<S, D, A>,
21    pad: &[[usize; 2]],
22    mode: PadMode<A>,
23    output: &mut Array<A, D>,
24) -> Option<bool>
25where
26    A: Copy + FromPrimitive + Num,
27    D: Dimension,
28    S: DataOwned<Elem = A>,
29{
30    let pad = reader(data.ndim(), pad)?;
31
32    // Select portion of padded array that needs to be copied from the original array.
33    output
34        .slice_each_axis_mut(|ad| {
35            let AxisDescription { axis, len, .. } = ad;
36            let pad = pad[axis.index()];
37            Slice::from(pad[0]..len - pad[1])
38        })
39        .assign(data);
40
41    match mode.into_pad_action() {
42        PadAction::StopAfterCopy => {
43            // Do nothing
44            Some(true)
45        }
46        _ => unimplemented!(),
47    }
48}
49
50pub fn pad<A, S, D>(
51    data: &ArrayBase<S, D, A>,
52    padding: &[[usize; 2]],
53    mode: PadMode<A>,
54) -> Array<A, D>
55where
56    A: Copy + FromPrimitive + Num,
57    D: Dimension,
58    S: DataOwned<Elem = A>,
59{
60    let pad = reader(data.ndim(), padding).expect("Inconsistent dimensions for padding");
61    let mut dim = data.raw_dim();
62    for (ax, (&ax_len, pad)) in data.shape().iter().zip(pad.iter()).enumerate() {
63        dim[ax] = ax_len + pad[0] + pad[1];
64    }
65
66    let mut padded = data.array_like(dim, mode.init()).to_owned();
67    apply_padding(data, &pad, mode, &mut padded).expect("Failed to apply padding");
68    padded
69}
70
71impl<A, S, D> Pad<A> for ArrayBase<S, D, A>
72where
73    A: Copy + FromPrimitive + Num,
74    D: Dimension,
75    S: DataOwned<Elem = A>,
76{
77    type Output = Array<A, D>;
78
79    fn pad(&self, mode: PadMode<A>, padding: &[[usize; 2]]) -> Self::Output {
80        pad(self, padding, mode)
81    }
82}
concision_core/utils/pad/impl_pad.rs

concision_core/utils/pad/
impl_pad.rs
