Convolution

[cjordan]

Hi!

Not sure if there's a better place for this, so please let me know if I'm out of line.

I'm (slowly) trying to port an existing python package of mine to rust. One of the first stumbling blocks was the (apparent) lack of a convolution function available in various rust libraries. Below is my solution.

(Also, I believe that using ffts for convolution is better only when the data is sufficiently big, so I'm interested in the naive method).

/// This function aims to reproduce the behaviour of numpy's convolve with
/// mode='same'.
pub fn convolve(data: ArrayView1<f64>, window: ArrayView1<f64>) -> Array1<f64> {
    let padded = stack![
        Axis(0),
        Array1::zeros(window.len() / 2),
        data,
        Array1::zeros(window.len() / 2)
    ];
    let mut w = window.view();
    w.invert_axis(Axis(0));

    padded
        .windows(w.len())
        .into_iter()
        .map(|x| (&x * &w).sum())
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use float_cmp::*;

    #[test]
    fn convolve_odd_odd() {
        let data = array![1., 2., 3.];
        let window = array![0., 1., 0.5];
        let expected = array![1., 2.5, 4.];

        for (exp, res) in expected.iter().zip(&convolve(data.view(), window.view())) {
            assert!(approx_eq!(f64, *exp, *res, ulps = 2));
        }
    }

    #[test]
    fn convolve_even_odd() {
        let data = array![1., 2., 3., 4.];
        let window = array![0., 1., 0.5];
        let expected = array![1., 2.5, 4., 5.5];

        for (exp, res) in expected.iter().zip(&convolve(data.view(), window.view())) {
            assert!(approx_eq!(f64, *exp, *res, ulps = 2));
        }
    }

    #[test]
    fn convolve_even_even() {
        let data = array![1., 2., 3., 4.];
        let window = array![1., 0.5];
        let expected = array![1., 2.5, 4., 5.5];

        for (exp, res) in expected.iter().zip(&convolve(data.view(), window.view())) {
            assert!(approx_eq!(f64, *exp, *res, ulps = 2));
        }
    }

    #[test]
    fn convolve_odd_even() {
        let data = array![1., 2., 3., 4., 5.];
        let window = array![1., 0.5];
        let expected = array![1., 2.5, 4., 5.5, 7.];

        for (exp, res) in expected.iter().zip(&convolve(data.view(), window.view())) {
            assert!(approx_eq!(f64, *exp, *res, ulps = 2));
        }
    }

    #[test]
    fn convolve_odd_odd2() {
        let data = array![1., 2., 3., 4., 5.];
        let window = array![2., 1., 0., 1., 0.5];
        let result = convolve(data.view(), window.view());
        let expected = array![8., 12., 16.5, 9., 5.5];

        for (exp, res) in expected.iter().zip(&result) {
            assert!(approx_eq!(f64, *exp, *res, ulps = 2));
        }
    }

    #[test]
    fn convolve3() {
        let data = array![1., 2., 3., 4.];
        let window = array![1., 0., 1., 0.5];
        let result = convolve(data.view(), window.view());
        let expected = array![2., 4., 6.5, 4.];

        for (exp, res) in expected.iter().zip(&result) {
            assert!(approx_eq!(f64, *exp, *res, ulps = 2));
        }
    }

This probably could be better; any suggestions are welcome. I have no idea how to generalise this for higher dimensional arrays, so that's also a consideration. I'd be happy to put this into a PR if it's suitable somewhere.

[LukeMathWalker]

It would definitely be nice to have this as an example! Would you mind opening a PR? Then we can discuss the implementation details there, which is better suited to comment on specific lines 😁

I would go as far as saying that if you have a complete port of np.convolve it would probably be a good addition to ndarray itself!