Fast Fourier Transform of exp(-x^2) in NumPy

user_5276943 0
Viewed 663

I have to calculate numerically the 2nd derivative of a Gaussian function: \frac{\partial^2}{\partial x^2} f(x) = \frac{\partial^2}{\partial x^2}e^{-ax^2} = -2axe^{-ax^2}
I've read every question on this topic here, but can't come to a good result. I've chosen NumPy as my tool of choice.

The instructions from our professor :

  1. Get an x array of size N = 128 with step dx = 1. So, -64, -63, ..., 62, 63. Calculate f(x)
  2. Perform FFT on f(x) and receive the transformed array f_m.
  3. Multiply f_m by (jk)^2, where j is imaginary unit, q=2 is degree of derivation and x
  4. Perform an inverse FFT to receive the derivative.
  5. In some FFT implementations you may have to scale by 1/n (but that's the smallest problem now)

Now here is my code, as simple as possible.

import numpy as np

# Set some parameters
n = 128
dx = 1
a = 0.001

# Create x, calculate f(x) and its FFT
x = np.arange(-n/2, n/2) * dx
psi = np.exp(-a * x * x)
f_m = np.fft.fft(psi)

# k_m creation according to professor (point 3. in my instruction)
k_m = np.arange(-n/2, n/2, dtype=float)
k_m[:int(n / 2)] = (2 * np.pi * k_m[:int(n / 2)]) / (n * dx)
k_m[int(n / 2):] = (2 * np.pi * (k_m[int(n / 2):] - n)) / (n * dx)

# Multiply f_m by (j * k_m)^q. For q=2, this is -k_m^2
f_m *= -k_m * k_m
# Inverse FFT on the result to get the second derivative and scale by 1 / n
f_m = np.fft.ifft(f_m) / n

One thing that I can't get is the result still has imaginary part and so something is not right. Can anybody help there?

EDIT: Cris Luengo's answer works.

fft numpy python
2 Answers

This part is wrong:

k_m = np.arange(-n/2, n/2, dtype=float)

The instructions in step 3 talk about m going from 0 to n-1. The code should look like this:

k_m = np.arange(0, n, dtype=float)
half = int(n / 2) + 1;  # notice the + 1 here!
k_m[:half] = (2 * np.pi * k_m[:half]) / (n * dx)
k_m[half:] = (2 * np.pi * (k_m[half:] - n)) / (n * dx)

The FFT produces an output where the first element (index 0) is the 0 frequency, not a frequency -n/2.

Your current version of the k_m array might be correct if you use fftshift to shift the 0 frequency bin to the middle of the array, though I'm not completely sure of that (maybe the -n in the second half should be removed?).

Finally, the division by n is not necessary here:

f_m = np.fft.ifft(f_m) / n

The NumPy IFFT is already normalized.

And remember to plot f_m.real, after verifying that the imaginary component is almost zero (these values should differ from zero only because of numerical rounding errors).

If you make a a little larger, for example a=0.005, then your input Gaussian fits entirely in the input signal, and you won't have ugly edge effects from filtering a signal that is cut off.

user_73287820

You could do with a simpler k, as long as you perform a correct FT shift at some point, which implements under the hood the same as your instructor or @CrisLuengo have written explicitly.

import numpy as np


# Set some parameters
n = 128
dx = 1
a = 0.001

# Create x, calculate f(x) and its FFT
x = np.arange(-n // 2, n // 2) * dx
f_x = np.exp(-a * x ** 2)
dd_f_x = 2 * a * np.exp(-a * x ** 2) * (2 * a * x ** 2 - 1)

f_k = np.fft.fft(f_x)
k = np.fft.ifftshift(np.arange(-n // 2, n // 2))
k = (2 * np.pi * k / (n * dx))
dd_f_k = -k ** 2 * f_k
dd_f_x_ = np.fft.ifft(dd_f_k)

which works as expected:

import matplotlib.pyplot as plt


fig, ax = plt.subplots(1, 1, squeeze=True)
ax.plot(x, dd_f_x_.real, label='∂²/∂x² f(x) with DFT')
ax.plot(x, dd_f_x, label='∂²/∂x² f(x)')
ax.legend()

plot

user_52183540
Related