I wrote this short and simple python method the other day
def dFFT_1D(f):
k = 2*np.pi * np.fft.fftfreq(f.shape[0])
return np.fft.ifft(1j*k * np.fft.fft(f)).real
which takes a one-dimensional array containing samples of a scalar function (step size is taken as unity) as input argument and returns the derivative using numpy's fast Fourier transform.
My goal is to generalize this method; so that it will take a multidimensional array and return the gradient/derivative along any of its axis.
def dFFT(f, axis):
k = 2*np.pi * np.fft.fftfreq(f.shape[axis])
return np.fft.ifftn( 1j*k * np.fft.fftn(f) ).real
for the sake of simplicity, let us take f to be two-dimensional array with shape (m,n) and say I wish to compute its derivative along its first axis (i.e. axis=0) using the above method. Then, my problem simply is that k has shape (m,) and therefor can't be broadcasted together with a (m,n) array. I have so far been able to come up with to different attempts at resolving this. My first attempt is simply to expand the dimension of k by
k = np.expand_dims(k, axis).T
which ensure that k has shape (n,1). My second attempt is to begin by swapping the axis of f
f = np.swapaxes(f, 0, axis)
However, both of these approaches doesn't seem elegant (as in efficiency and readability). I am also afraid that they don't really extend beyond the two-dimensionally case. I have tried reaching for a solution online, so far unsuccessfully.
Cheers!.