直接在scipy.weave.inline中对复杂的numpy数组进行FFTW3

Question

I am trying to implement an FFT based subpixel shifting (translation) algorithm in Python . The Fourier shift theorem allows an array to be translated by a subpixel amount by: 1. Forward FFT array 2. Multiply array by linear phase ramp in Fourier space 3. Inverse FFT array

This algorithm is easy to implement in python using numpy/scipy but its incredibly slow (~10msec) per shift for 256**2 array. I am trying to speed this up by calling c code directly from python using scipy.weave.inline.

I'm having trouble however in passing complex numpy arrays to FFTW. The c code looks like:

    #include <fftw3.h>
    #include <stdlib.h>

    #define INVERSE +1
    #define FORWARD -1


    fftw_complex *i, *o;
    int n, m;
    fftw_plan pf, pi;
    #line 22 "test_scipy_weave.py" 

    i = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * xdim*ydim);
    o = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * xdim*ydim);

    pf = fftw_plan_dft_2d(xdim, ydim, i, o, -1, FFTW_PATIENT);
    pi = fftw_plan_dft_2d(xdim, ydim, o, i,  1, FFTW_PATIENT);

    # Copy data to fftw_complex array. How to use python arrays directly
    for (n=0; n<xdim;n++){
        for (m=0; m<ydim; m++){
            i[n*xdim+m][0]=a[n*xdim+m].real();
            i[n*xdim+m][1]=a[n*xdim+m].imag();
        }
    }

    fftw_execute(pf);

    /* Mult by linear phase ramp here */

    fftw_execute(pi);

    for (n=0; n<xdim;n++){
        for (m=0; m<ydim; m++){
            b[n*xdim+m] = std::complex<double>([in*xdim+m][0], i[n*xdim+m][1]);
        }
    }

    fftw_destroy_plan(p);

So you can see I have to copy the data stored in the numpy array "a" into the fftw_complex array "i". And again at the end I have to copy the result "i" into the output numpy array "b". It would be much more efficient to use the numpy arrays "a" and "b" directly in the fftw but I cannot get this to work.

Does anyone have an idea on how to get fftw to use complex numpy arrays directly in scipy.weave.inline ?

Thanks

Answer 1

According to the fftw manual , you can import complex.h before fftw.h , which will guarantee that fftw_complex will correspond to the native C data type. I'm pretty sure that numpy data types are also guaranteed to be (or in practice are likely to be) compatible with native C data types.

In this case you can access a pointer to the array data as a.data_as(ctypes.c_void_p) . Unfortunately ctypes doesn't recognise complex types, but hopefully casting to a void pointer will do the trick.

When doing this, you have to be careful that your array a is stored in C-contiguous fashion, specified by the parameter order='C' when creating the array.