ndarray_ndimage/filters/
uniform.rs1use ndarray::{s, Array, Array1, ArrayRef, Axis, Dimension, Zip};
2use num_traits::{FromPrimitive, Num};
3
4use crate::{array_like, pad_to, BorderMode};
5
6pub 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 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 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
52pub 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 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}