用python实现2D FFT

前言

用python实现2D FFT，并绘制结果

代码

直接上代码

import numpy as np
from scipy.fftpack import fftshift,fftn
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator

N = 128                           # 采样点数
sample_freq = 120                 # 采样频率 120 Hz, 大于两倍的最高频率
sample_interval = 1/sample_freq   # 采样间隔
signal_len=N*sample_interval      # 信号长度
t=np.arange(0, signal_len, sample_interval)


signal = 5 + 2 * np.sin(2 * np.pi * 15 * t) + 3 * np.sin(2 * np.pi * 30 * t) + 4 * np.sin(2 * np.pi * 40 * t)  # 采集的信号
mat2D = np.zeros((N, N))
for i in range(N):
    mat2D[i, :] = signal

mat2D=np.transpose(mat2D, [0, 1])
plt.matshow(mat2D)
plt.title('Original data')
plt.show()

fft_data_2D=fftn(mat2D)
fft_amp0_2D = np.array(np.abs(fft_data_2D)/N/N*2)

direct_2D=fft_amp0_2D[:, 0]
fft_amp0_2D[:, 0]=0.5*direct_2D
fft_amp0_2D_shift=fftshift(fft_amp0_2D)
list0_shift_2d = np.array(range(0, N))
freq0_shift_2d=sample_freq*list0_shift_2d/N-sample_freq/2  # 零频移动后的频率轴


plt.matshow(fft_amp0_2D_shift)
plt.tick_params(axis='both', which='major',labelsize=14)

x_major_locator=MultipleLocator(40)
y_major_locator=MultipleLocator(40)

ax=plt.gca()
ax.xaxis.set_major_locator(x_major_locator)
ax.yaxis.set_major_locator(y_major_locator)
plt.colorbar()
plt.title('Spectrum 2D')
plt.show()
print('---------------------------------------------------')

xx, yy = np.meshgrid(freq0_shift_2d, freq0_shift_2d )
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.set_top_view()
ax.plot_surface(xx, yy, fft_amp0_2D_shift, rstride=1, cstride=1, cmap='rainbow')
plt.title('Spectrum show in 3D')
plt.show()

结果

原始2D数据：


求得的2DFFT结果


3D视图

总结

2D FFT作为1D FFT的拓展，求取操作与1D FFT 类似。

参考链接：


https://blog.csdn.net/weixin_43537379/article/details/119636757?spm=1001.2014.3001.5501