Skip to main content

ndarray_ndimage/filters/
uniform.rs

1use ndarray::{s, Array, Array1, ArrayRef, Axis, Dimension, Zip};
2use num_traits::{FromPrimitive, Num};
3
4use crate::{array_like, pad_to, BorderMode};
5
6/// Uniform filter for n-dimensional arrays.
7///
8/// Currently hardcoded with the `PadMode::Reflect` padding mode.
9///
10/// * `data` - The input N-D data.
11/// * `size` - The len
12/// * `mode` - Method that will be used to select the padded values. See the
13///   [`BorderMode`](crate::BorderMode) enum for more information.
14///
15/// **Panics** if `size` is zero, or one of the axis' lengths is lower than `size`.
16pub fn uniform_filter<A, D>(
17    data: &ArrayRef<A, D>,
18    size: usize,
19    mode: BorderMode<A>,
20) -> Array<A, D>
21where
22    A: Copy + Num + FromPrimitive + PartialOrd + 'static,
23    D: Dimension,
24{
25    let half = size / 2;
26
27    // We need 2 buffers because
28    // * We're reading neighbours so we can't read and write on the same location.
29    // * The process is applied for each axis on the result of the previous process.
30    // * It's uglier (using &mut) but much faster than allocating for each axis.
31    let mut data = data.to_owned();
32    let mut output = array_like(&data, data.dim(), A::zero());
33
34    for d in 0..data.ndim() {
35        // TODO This can be made to work if the padding modes (`reflect`, `symmetric`, `wrap`) are
36        // more robust. One just needs to reflect the input data several times if the `weights`
37        // length is greater than the input array. It works in SciPy because they are looping on a
38        // size variable instead of running the algo only once like we do.
39        let n = data.len_of(Axis(d));
40        if half > n {
41            panic!("Data size is too small for the inputs (sigma and truncate)");
42        }
43
44        inner_uniform1d(&data, size, Axis(d), mode, &mut output);
45        if d != data.ndim() - 1 {
46            std::mem::swap(&mut output, &mut data);
47        }
48    }
49    output
50}
51
52/// Uniform filter for 1-dimensional arrays.
53///
54/// * `data` - The input N-D data.
55/// * `size` - Length of the uniform filter.
56/// * `axis` - The axis of input along which to calculate.
57/// * `mode` - Method that will be used to select the padded values. See the
58///   [`BorderMode`](crate::BorderMode) enum for more information.
59///
60/// **Panics** if `size` is zero, or the axis length is lower than `size`.
61pub fn uniform_filter1d<A, D>(
62    data: &ArrayRef<A, D>,
63    size: usize,
64    axis: Axis,
65    mode: BorderMode<A>,
66) -> Array<A, D>
67where
68    A: Copy + Num + FromPrimitive + PartialOrd + 'static,
69    D: Dimension,
70{
71    let mut output = array_like(data, data.dim(), A::zero());
72    inner_uniform1d(data, size, axis, mode, &mut output);
73    output
74}
75
76pub(crate) fn inner_uniform1d<A, D>(
77    data: &ArrayRef<A, D>,
78    size: usize,
79    axis: Axis,
80    mode: BorderMode<A>,
81    output: &mut Array<A, D>,
82) where
83    A: Copy + Num + FromPrimitive + PartialOrd,
84    D: Dimension,
85{
86    let size1 = size / 2;
87    let size2 = size - size1 - 1;
88    let size_as_a = A::from_usize(size).unwrap();
89
90    let mode = mode.to_pad_mode();
91    let n = data.len_of(axis);
92    let pad = vec![[size1, size2]];
93    let mut buffer = Array1::from_elem(n + size - 1, mode.init());
94
95    Zip::from(data.lanes(axis)).and(output.lanes_mut(axis)).for_each(|input, o| {
96        pad_to(&input, &pad, mode, &mut buffer);
97        let mut accumulator = buffer.slice(s![..size - 1]).sum();
98
99        // Optimise the filter by keeping a running total, to which add the newest item entering the
100        // window, and then subtract the element which has fallen out of the window.
101        Zip::from(o).and(buffer.slice(s![size - 1..])).and(buffer.slice(s![..n])).for_each(
102            |o, &leading_edge, &trailing_edge| {
103                accumulator = accumulator + leading_edge;
104                *o = accumulator / size_as_a;
105                accumulator = accumulator - trailing_edge;
106            },
107        );
108    });
109}