在NumPy中沿3D陣列的每個軸應用DFT矩陣嗎？

Question

我首先可以得到給定大小的DFT矩陣，比如說n

import numpy as np
n = 64
D = np.fft.fft(np.eye(n))

FFT當然只是將D應用於向量的一種快速算法：

x = np.random.randn(n)
ft1 = np.dot(D,x)
print( np.abs(ft1 - fft.fft(x)).max() )
# prints near double precision roundoff

可以通過將D應用於矩陣的行和列來獲得2D FFT：

x = np.random.randn(n,n)
ft2 = np.dot(x, D.T) # Apply D to rows.
ft2 = np.dot(D, ft2) # Apply D to cols.
print( np.abs(ft2 - fft.fft2(x)).max() )
# near machine round off again

如何針對3D離散傅立葉變換進行類似計算？

也就是說，

x = np.random.randn(n,n,n)
ft3 = # dot operations using D and x
print( np.abs(ft3 - fft.fftn(x)).max() )
# prints near zero

本質上，我認為我需要將D應用於體積中的每個列向量，然后是體積中的每個行向量，最后是每個“深度向量”。 但是我不確定如何使用dot做到這一點。

Answer 1

您可以使用einsum表達式對每個索引執行轉換：

x = np.random.randn(n, n, n)
ft3 = np.einsum('ijk,im->mjk', x, D)
ft3 = np.einsum('ijk,jm->imk', ft3, D)
ft3 = np.einsum('ijk,km->ijm', ft3, D)
print(np.abs(ft3 - np.fft.fftn(x)).max())
1.25571216554e-12

也可以將其寫為單個NumPy步驟：

ft3 = np.einsum('ijk,im,jn,kl->mnl', ft3, D, D, D, optimize=True)

如果沒有optimize參數（在NumPy 1.12+中可用），它將非常慢。 您也可以使用dot來完成每個步驟，但是這需要一些重塑和轉置。 在NumPy 1.14+中， einsum函數將自動檢測BLAS操作並為您執行此操作。