如何实现基于 1D RFFT 的 2D RFFT 算法？

Question

我正在尝试实现 NumPy 的 rfft2()，即支持二维数组的 RFFT 函数，方法是对每一行执行 1D RFFT，然后在每个列上再次执行 1D RFFT以前的结果。

这种方法可以很好地实现 2D FFT 函数，如前所述 on this post，但它似乎不适用于 2D RFFT。

这是一个实现自定义 2D FFT 函数的脚本，它遵循这个想法，使用 NumPy 的 FFT 的 1D 版本作为基础，然后将其结果与 NumPy 的实际 2D 版本进行比较：

import cmath
import numpy as np
import math

def my_fft2d(matrix):
    fft_rows = [np.fft.fft(row) for row in matrix]
    return np.transpose([np.fft.fft(row) for row in np.transpose(fft_rows)])


# initialize test data
img = np.array([[0,0,0,0], [0,1,0,0], [0,0,0,0], [0,0,0,0]])
print('img shape=', img.shape)

# perform custom FFT2D and print result
custom_result = my_fft2d(img)
print('\ncustom_result shape=', custom_result.shape)
for row in custom_result:
   print(', '.join(['%.3f + %.3fi' % (x.real, x.imag) for x in row]))

# perform numpy FFT2D and print result
numpy_result = np.fft.fft2(img)
print('\nnumpy_result shape=', numpy_result.shape)
for row in numpy_result:
   print(', '.join(['%.3f + %.3fi' % (x.real, x.imag) for x in row]))

# compare results
print('\nAre the results equivalent to NumPy?', np.allclose(custom_result, custom_result))
print('ASSERT(assert_array_almost_equal):', np.testing.assert_array_almost_equal(custom_result, custom_result))

输出：

img shape= (4, 4)

custom_result shape= (4, 4)
1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i
0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i
-1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i
0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i

numpy_result shape= (4, 4)
1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i
0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i
-1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i
0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i

Are the results equivalent to NumPy? True
ASSERT(assert_array_almost_equal): None

脚本的输出显示my_fft2d() 实现与np.fft.fft2() 兼容。

但是，当应用相同的逻辑来实现 RFFT 版本的变换时，生成的数组具有不同的形状，如下面的脚本所示：

def my_rfft2d(matrix):
    fft_rows = [np.fft.rfft(row) for row in matrix]
    return np.transpose([np.fft.rfft(row) for row in np.transpose(fft_rows)])


# initialize test data
img = np.array([[0,0,0,0], [0,1,0,0], [0,0,0,0], [0,0,0,0]])
print('img shape=', img.shape)

# perform custom FFT2D and print result
custom_result = my_rfft2d(img)
print('\ncustom_result shape=', custom_result.shape)
for row in custom_result:
   print(', '.join(['%.3f + %.3fi' % (x.real, x.imag) for x in row]))

# perform numpy FFT2D and print results
numpy_result = np.fft.rfft2(img)
print('\nnumpy_result shape=', numpy_result.shape)
for row in numpy_result:
   print(', '.join(['%.3f + %.3fi' % (x.real, x.imag) for x in row]))

输出：

img shape= (4, 4)
C:\Users\username\AppData\Roaming\Python\Python37\site-packages\numpy\fft\_pocketfft.py:77: ComplexWarning: Casting complex values to real discards the imaginary part
  r = pfi.execute(a, is_real, is_forward, fct)

custom_result shape= (3, 3)
1.000 + 0.000i, 0.000 + 0.000i, -1.000 + 0.000i
0.000 + -1.000i, 0.000 + 0.000i, 0.000 + 1.000i
-1.000 + 0.000i, 0.000 + 0.000i, 1.000 + 0.000i

numpy_result shape= (4, 3)
1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i
0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i
-1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i
0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i

如你所见，输出有两个问题：

• 来自 numpy 的警告抱怨我不完全确定如何解决的问题；
• 2D RFFT 的自定义实现返回的结果行数少于np.fft.rfft2() 返回的行数；

如何解决此问题并使my_rfft2d() 与np.fft.rfft2() 兼容？

Answer 1

正如我在评论中所说，在取行的rfft 之后，您应该取fft 而不是rfft，因为rfft 结果通常很复杂。

我不知道你为什么要变得真实，但是如果你真的想变得真实，你应该使用 DCT（离散余弦变换）而不是 FFT，因为 DCT 输出是真实的。您可以采用与上面计算 2D FFT 相同的方法，因为您可以以类似的方式分解 2D DCT。

Answer 2

就像评论者所说的那样。你应该第二次服用fft。这是因为行的 rfft 的输出很复杂。这解决了复杂到真实的错误，以及形状问题。

import numpy as np

def my_rfft2d(matrix):
    fft_rows = [np.fft.rfft(row) for row in matrix]
    return np.transpose([np.fft.fft(row) for row in np.transpose(fft_rows)])


# initialize test data
img = np.array([[0,0,0,0], [0,1,0,0], [0,0,0,0], [0,0,0,0]])
print('img shape=', img.shape)

# perform custom FFT2D and print result
custom_result = my_rfft2d(img)
print('\ncustom_result shape=', custom_result.shape)
for row in custom_result:
   print(', '.join(['%.3f + %.3fi' % (x.real, x.imag) for x in row]))

# perform numpy FFT2D and print results
numpy_result = np.fft.rfft2(img)
print('\nnumpy_result shape=', numpy_result.shape)
for row in numpy_result:
   print(', '.join(['%.3f + %.3fi' % (x.real, x.imag) for x in row]))

输出：

custom_result shape= (4, 3)
1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i
0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i
-1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i
0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i

numpy_result shape= (4, 3)
1.000 + 0.000i, 0.000 + -1.000i, -1.000 + 0.000i
0.000 + -1.000i, -1.000 + 0.000i, 0.000 + 1.000i
-1.000 + 0.000i, 0.000 + 1.000i, 1.000 + 0.000i
0.000 + 1.000i, 1.000 + 0.000i, 0.000 + -1.000i