CODEWITHSUNDEEP
numpy
user42298
user42298

Reputation: 1911

how to set domain of 2d fourier transformin in numpy

I want to use numpy FFT2 module.

For example, I have the values of function in discrete grid x=(0,1,2,3,4,5), y=(0,1,2,3,4,5), z=exp(x+y) So my code will be

import numpy as np
import matplotlib.pyplot as plt

x=np.linspace(0,5,6)
y=np.linspace(0,5,6)
xmesh,ymesh=np.meshgrid(x,y)
z=np.exp(xmesh+ymesh)
plt.scatter(xmesh.ravel(),ymesh.ravel(),z.ravel())
plt.show()

This shows the image such as enter image description here Now, I want to transform this into k-space, using Fourier transform. For z value, I just used

fft_z=np.fft.fft2(z)

But how do I determine its domain? To plot scatter plot like above(In fact, I'll plot heatmap), I have to determine where these fft_z values lies on, and make meshgrid in k-space like above code.

Upvotes: 1

Views: 1751

Answers (1)

mrcl
mrcl

Reputation: 2190

You can use np.fft.fftfreq to generate the frequency domain automatically. Just as a reminder, the relationship between the frequency and time domains is df = 1/(N*dt) and 1/dt is the full bandwidth (BW) of the frequency domain.

It is important to notice the fftfreq generates an array from -BW/2 to BW/2, however in an unconventional way. Starting from 0 to BW/2, -BW/2 to -df. This is because of the periodicity of the discrete Fourier transform, so the data will be shown centred around zero. You can use np.fft.fftshift to realing the data before plotting. You can see an example bellow.

fft_z = np.fft.fftshift(np.fft.fft2(z))
fx = np.fft.fftshift(np.fft.fftfreq(x.shape[0],x[1]-x[0]))
fy = np.fft.fftshift(np.fft.fftfreq(y.shape[0],y[1]-y[0]))

you can plot the result using:

plt.pcolormesh(fx,fy,fft_z.real)

OBS: fft_z.real to take the real part of fft_z, which is a complex number. The same can be used to the imaginary part. plt.pcolormesh(fx,fy,fft_z.imag)

Hope it helps

Upvotes: 1

Related Questions