matplotlib pcolormesh在FFT数组中出现问题？

Question

我正在尝试复现维基百科上关于Gabor变换的例子，但我不知道这是个bug还是我漏掉了什么。这个例子是计算一个正弦信号的Gabor变换：

为了绘制频率的排序图，我先创建了一个未排序的坐标轴。然后我使用网格生成（mesh grid）来创建二维坐标轴，并用pcolormesh来绘图。这是代码的一部分：

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridsp

dt = 0.05
x = np.arange(-50.0,50.0,dt)
y = np.sin(2.0 * np.pi * x)
Nx = len(x)
w = np.fft.fftfreq(Nx,dt)

sigma = 1.0 / 3.0 
neg = np.where (x <= 0.0)
pos = np.where (x > 0.0)

T,W = np.meshgrid(x,w)
func = np.zeros(Nx)
tmp = np.zeros(Nx,dtype='complex64')
gabor = np.zeros((Nx,Nx))

func[neg] = np.sin(2.0 * np.pi * x[neg])
func[pos] = np.sin(4.0 * np.pi * x[pos])

for it in range(Nx):    
        tmp[:] = np.fft.fft(func[:] * np.exp( - ( x[it] - x[:] ) * ( x[it] - x[:] ) / 2.0 / sigma / sigma ) ) 

        gabor[:,it] = np.real(np.conj(tmp) * tmp)

fig = plt.figure(figsize=(20,10),facecolor='white')
gs = gridsp.GridSpec(2, 1)

ax1 = plt.subplot(gs[0,0])
ax1.plot(x,func,'r',linewidth=2)
ax1.axis('tight')
ax1.set_xticks(np.arange(min(x),max(x),1.) )
ax1.set_xlabel('time',fontsize=20)
ax1.set_ylabel(r'$\sin{time}$',fontsize=20)
ax1.set_xlim([-6.0,6.0])


ax2 = plt.subplot(gs[1,0])
surf1 = ax2.pcolormesh(T,W,gabor,shading='gouraud')
ax2.axis('tight')
ax2.set_xticks(np.arange(min(x),max(x),2.) )
ax2.set_yticks(np.arange(min(w),max(w),2.) )
ax2.set_xlabel('time',fontsize=20)
ax2.set_ylabel('frequency',fontsize=20)
ax2.set_xlim([-6.0,6.0])
ax2.set_ylim([-4.0,4.0])

gs.tight_layout(fig)
plt.show()

这是我得到的图：

看起来图的上半部分被压缩到零了。如果我在创建变换和坐标轴时使用fftshift，

for it in range(Nx):
        tmp[:] = np.fft.fftshift(np.fft.fft(func[:] * np.exp( - ( x[it] - x[:] ) * ( x[it] - x[:] ) / 2.0 / sigma / sigma ) ) )     


        gabor[:,it] = np.real(np.conj(tmp) * tmp)

T,W = np.meshgrid(x,np.fft.fftshift(w))

那么我得到的图是这样的：

!

似乎pcolormesh这个函数不能像一维图那样翻转数组。有没有人知道为什么会这样？

谢谢，

亚历克斯

Answer 1

问题出在 W 上。其实是 w 出了问题。当我们把 w 画出来的时候：

所以 pcolormesh 收到了不规则的 Y 坐标，导致它搞不清楚状况。如果你查看 pcolor 或 pcolormesh 的说明，就会发现它们无法处理不规则的数据。

不过，你的 gabor 是没问题的：

ax.imshow(gabor)

如你所见：

有几种方法可以解决这个问题。其中一种是把 W 和 gabor 都传给 fftshift，这样频率就会变得规则。如果你想要像上面的图那样（负频率在上面），只需把最大频率加到所有负的 W 值上。

另外，给 pcolormesh 提供 x 和 w 可能会更清晰，而不是 T 和 W。

---

如果你想要更好的性能，使用 imshow 可能更合适（当数据在两个维度上均匀分布时可以使用）。唯一的小问题是计算范围（实际上在问题中可能也有点偏差）。范围告诉我们最高、最低、最左和最右像素的外部边界。然而，像素向量只告诉我们像素的中心位置。

我们需要知道以下几点：

• X 方向的点数 (num_x)
• Y 方向的点数 (num_y)
• 第一个和最后一个 x 采样值 (x0, x1)
• 第一个和最后一个 y 采样值 (y0, y1)

之后我们可以使用 imshow 来正确显示数据：

dx = 1. * (x1 - x0) / (num_x-1)
dy = 1. * (y1 - y0) / (num_y-1)

ax.imshow(img, extent=[x0 - dx/2, x1 + dx/2, y0 - dy/2, y1 + dy/2], origin='lower', interpolation='nearest')

所以，应用到问题的数据上：

gabor_shifted = np.fft.fftshift(gabor, axes=0)
w_shifted = np.fft.fftshift(w)

x0 = x[0]
x1 = x[-1]
w0 = w_shifted[0]
w1 = w_shifted[-1]

dx = 1.*(x1-x0) / (len(x) - 1)
dw = 1.*(w1-w0) / (len(w) - 1)

ax2.imshow(gabor_shifted, extent=[x0-dx/2, x1+dx/2, w0-dw/2, w1+dw/2], interpolation='nearest', origin='lower', aspect='auto')
ax2.grid('on', color='w')
ax2.ylim(-4,4)

结果是：