nalgebra/linalg/
convolution.rs

1use std::cmp;
2
3use crate::base::allocator::Allocator;
4use crate::base::default_allocator::DefaultAllocator;
5use crate::base::dimension::{Const, Dim, DimAdd, DimDiff, DimSub, DimSum};
6use crate::storage::Storage;
7use crate::{zero, OVector, RealField, Vector, U1};
8
9impl<T: RealField, D1: Dim, S1: Storage<T, D1>> Vector<T, D1, S1> {
10    /// Returns the convolution of the target vector and a kernel.
11    ///
12    /// # Arguments
13    ///
14    /// * `kernel` - A Vector with size > 0
15    ///
16    /// # Errors
17    /// Inputs must satisfy `vector.len() >= kernel.len() > 0`.
18    ///
19    pub fn convolve_full<D2, S2>(
20        &self,
21        kernel: Vector<T, D2, S2>,
22    ) -> OVector<T, DimDiff<DimSum<D1, D2>, U1>>
23    where
24        D1: DimAdd<D2>,
25        D2: DimAdd<D1, Output = DimSum<D1, D2>>,
26        DimSum<D1, D2>: DimSub<U1>,
27        S2: Storage<T, D2>,
28        DefaultAllocator: Allocator<DimDiff<DimSum<D1, D2>, U1>>,
29    {
30        let vec = self.len();
31        let ker = kernel.len();
32
33        if ker == 0 || ker > vec {
34            panic!("convolve_full expects `self.len() >= kernel.len() > 0`, received {} and {} respectively.", vec, ker);
35        }
36
37        let result_len = self
38            .data
39            .shape()
40            .0
41            .add(kernel.shape_generic().0)
42            .sub(Const::<1>);
43        let mut conv = OVector::zeros_generic(result_len, Const::<1>);
44
45        for i in 0..(vec + ker - 1) {
46            let u_i = if i > vec { i - ker } else { 0 };
47            let u_f = cmp::min(i, vec - 1);
48
49            if u_i == u_f {
50                conv[i] += self[u_i].clone() * kernel[i - u_i].clone();
51            } else {
52                for u in u_i..(u_f + 1) {
53                    if i - u < ker {
54                        conv[i] += self[u].clone() * kernel[i - u].clone();
55                    }
56                }
57            }
58        }
59        conv
60    }
61    /// Returns the convolution of the target vector and a kernel.
62    ///
63    /// The output convolution consists only of those elements that do not rely on the zero-padding.
64    /// # Arguments
65    ///
66    /// * `kernel` - A Vector with size > 0
67    ///
68    ///
69    /// # Errors
70    /// Inputs must satisfy `self.len() >= kernel.len() > 0`.
71    ///
72    pub fn convolve_valid<D2, S2>(
73        &self,
74        kernel: Vector<T, D2, S2>,
75    ) -> OVector<T, DimDiff<DimSum<D1, U1>, D2>>
76    where
77        D1: DimAdd<U1>,
78        D2: Dim,
79        DimSum<D1, U1>: DimSub<D2>,
80        S2: Storage<T, D2>,
81        DefaultAllocator: Allocator<DimDiff<DimSum<D1, U1>, D2>>,
82    {
83        let vec = self.len();
84        let ker = kernel.len();
85
86        if ker == 0 || ker > vec {
87            panic!("convolve_valid expects `self.len() >= kernel.len() > 0`, received {} and {} respectively.",vec,ker);
88        }
89
90        let result_len = self
91            .data
92            .shape()
93            .0
94            .add(Const::<1>)
95            .sub(kernel.shape_generic().0);
96        let mut conv = OVector::zeros_generic(result_len, Const::<1>);
97
98        for i in 0..(vec - ker + 1) {
99            for j in 0..ker {
100                conv[i] += self[i + j].clone() * kernel[ker - j - 1].clone();
101            }
102        }
103        conv
104    }
105
106    /// Returns the convolution of the target vector and a kernel.
107    ///
108    /// The output convolution is the same size as vector, centered with respect to the ‘full’ output.
109    /// # Arguments
110    ///
111    /// * `kernel` - A Vector with size > 0
112    ///
113    /// # Errors
114    /// Inputs must satisfy `self.len() >= kernel.len() > 0`.
115    #[must_use]
116    pub fn convolve_same<D2, S2>(&self, kernel: Vector<T, D2, S2>) -> OVector<T, D1>
117    where
118        D2: Dim,
119        S2: Storage<T, D2>,
120        DefaultAllocator: Allocator<D1>,
121    {
122        let vec = self.len();
123        let ker = kernel.len();
124
125        if ker == 0 || ker > vec {
126            panic!("convolve_same expects `self.len() >= kernel.len() > 0`, received {} and {} respectively.",vec,ker);
127        }
128
129        let mut conv = OVector::zeros_generic(self.shape_generic().0, Const::<1>);
130
131        for i in 0..vec {
132            for j in 0..ker {
133                let val = if i + j < 1 || i + j >= vec + 1 {
134                    zero::<T>()
135                } else {
136                    self[i + j - 1].clone()
137                };
138                conv[i] += val * kernel[ker - j - 1].clone();
139            }
140        }
141
142        conv
143    }
144}