ndarray-examples

ndarray-examples copied to clipboard

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.

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 :grin:

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!

I believe we can do much more in ndarray now and the examples repository should be updated. This repo uses version 0.13 but 0.16.1 has been released now. For example, with the addition of windows and windows_with_stride, we can do convolution more easily and on higher dimensions too. I'd be happy to help implement it!

Hi @tanmay2004, I'd be very happy for you to work on this. Feel free to add to my PR, and let me know if you'd like help with anything. I haven't used ndarray for quite a while now, so I'd be interested to see how much better the code could look.