FFT convolution implementation - incorrect results potentially due to normalization

[rhysnewell]

Hello,

Thanks for making this crate. I'm hoping that you'll be able to help me out and perhaps provide some answers on the results I'm getting when using rustfft. I'm attempting to implement both 1D and 2D fft convolution as is done in scipy.signal.convolve and scipy.signal.correlate. The idea being that I'll be able to speed up some large pairwise matrix correlation calculations that I need to perform. Here is my current attempt at a 1-dimensional convolution algorithm which is the easy case:

// convolve two 1-dimensional arrays using FFT
pub fn fftconvolve1d(in1: &Array1<f64>, in2: &Array1<f64>) -> Result<Array1<f64>, Box<dyn Error>> {
    if in1.shape() != in2.shape() {
        return Err("Input arrays must have the same shape.".into());
    }

    let shape = in1.len() + in2.len() - 1;
    println!("shape: {}", shape);
    // pad in1 and in2 with zeros
    // the number of zeros to pad should be the difference between the original length and shape
    // zeroes should be added equally to each end
    let in1 = pad_with_zeros(in1, vec![[(shape - in1.len()) / 2, (shape - in1.len()) / 2]], 0.0);
    let in2 = pad_with_zeros(in2, vec![[(shape - in2.len()) / 2, (shape - in2.len()) / 2]], 0.0);
    println!("in1: {:?}", in1);
    println!("in2: {:?}", in2);
    // let fft = planner.plan_fft_forward(max);
    let mut in1 = in1.mapv(|x| Complex::new(x, 0.0)).to_vec();
    let mut in2 = in2.mapv(|x| Complex::new(x, 0.0)).to_vec();

    let mut scratch = vec![Complex::new(0.0, 0.0); shape];
    let mut planner = FftPlanner::<f64>::new();
    let fft = planner.plan_fft_forward(shape);
    fft.process_with_scratch(in1.as_mut_slice(), scratch.as_mut_slice());
    fft.process_with_scratch(in2.as_mut_slice(), scratch.as_mut_slice());
    println!("{:?}", &in1);
    println!("{:?}", &in2);
    // Multiply the FFTs.
    let shape_sqrt = (shape as f64).sqrt();
    let mut out = in1.into_par_iter().zip(in2.into_par_iter()).map(|(a, b)| (a / shape_sqrt)  * (b / shape_sqrt)).collect::<Vec<_>>();

    // Perform the inverse FFT.
    // let mut planner = FftPlanner::new();
    // let mut scratch = vec![Complex::new(0.0, 0.0); shape];
    let fft = planner.plan_fft_inverse(shape);
    fft.process_with_scratch(out.as_mut_slice(), scratch.as_mut_slice());
    println!("{:?}", &out);
    Ok(Array::from_shape_vec((out.len()), out.into_iter().map(|c: Complex<f64>| c.re).collect())?)
}

/// Pad the edges of an array with zeros.
///
/// `pad_width` specifies the length of the padding at the beginning
/// and end of each axis.
///
/// **Panics** if `arr.ndim() != pad_width.len()`.
fn pad_with_zeros<A, S, D>(arr: &ArrayBase<S, D>, pad_width: Vec<[usize; 2]>, const_value: A) -> Array<A, D>
where
    A: Clone,
    S: Data<Elem = A>,
    D: Dimension,
{
    assert_eq!(
        arr.ndim(),
        pad_width.len(),
        "Array ndim must match length of `pad_width`."
    );

    // Compute shape of final padded array.
    let mut padded_shape = arr.raw_dim();
    for (ax, (&ax_len, &[pad_lo, pad_hi])) in arr.shape().iter().zip(&pad_width).enumerate() {
        padded_shape[ax] = ax_len + pad_lo + pad_hi;
    }

    let mut padded = Array::from_elem(padded_shape, const_value);
        let padded_dim = padded.raw_dim();
    {
        // Select portion of padded array that needs to be copied from the
        // original array.
        let mut orig_portion = padded.view_mut();
        for (ax, &[pad_lo, pad_hi]) in pad_width.iter().enumerate() {
            orig_portion.slice_axis_inplace(
                Axis(ax),
                Slice::from(pad_lo as isize..padded_dim[ax] as isize - (pad_hi as isize)),
            );
        }
        // Copy the data from the original array.
        orig_portion.assign(arr);
    }
    padded
}

I've got this simple test case:

#[test]
fn test_fftconvolve1d() {
    let in1 = Array1::<f64>::range(1.0, 4.0, 1.0);
    let in2 = Array1::<f64>::range(6.0, 3.0, -1.0);
    let out = fftconvolve1d(&in1, &in2).unwrap();
    println!("{}", out);
    panic!();
}

this should produce [ 4, 13, 28, 27, 18] but instead produces [23, 12.000000000000005, 6.000000000000002, 17, 31.999999999999993] ignore the floating point imprecision. At this point, I can't really tell where I am going wrong. I feel like it has something to do with the normalization, but I think I'm doing that right so I'm really not sure. Any help you could provide would be great.

Also, I feel like a convolution algorithm is a logical next step for rustfft so I was hoping if I could get this working we could add it into a pull request or something if you would like

Cheers, Rhys

[WalterSmuts]

EDIT: Oops, just saw you using the range function and not just declaring an array.

this should produce [ 4, 13, 28, 27, 18] but instead produces [23, 12.000000000000005, 6.000000000000002, 17, 31.999999999999993]

AFAICT it should produce 6 17 32 23 12.

octave:14> a = [1, 2, 3]
a =

   1   2   3

octave:15> b = [6, 5, 4]
b =

   6   5   4

octave:16> conv(a,b)
ans =

    6   17   32   23   12

Although not sure why you're declaring the arrays using the range stuff, easy to make mistakes with.

[WalterSmuts]

Your pad_with_zeros function pads at the start and end. You need to only pad at the end.

[rhysnewell]

You were right on both counts, thank you! I must have got mixed up when I was looking at the output from scipy convolve and correlate. Making the padding only occur at the end of the array fixed the issue and made it return the correct result. Thank you!

[rhysnewell]

Thanks for the help @WalterSmuts, N-dimensional implementations of convolve and correlate can be found here: https://github.com/rhysnewell/fftconvolve I'll publish as crate at some point soon