Python 如何在一维RFFT的基础上实现二维RFFT算法？

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

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

下面是一个脚本，它实现了一个自定义的2D FFT函数，该函数遵循这个思路，使用NumPy的1D版本的FFT作为基础，然后将其结果与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))

输出：

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

脚本的输出显示我的_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]))

输出：

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

如您所见，输出中存在两个问题：

numpy的警告抱怨了一些我不完全确定如何修复的问题； 2D RFFT的自定义实现返回的结果行数少于np.fft.rfft2返回的结果行数；

---

我如何解决这个问题并使我的rfft2d与np.fft.rfft2兼容？

如评论员所说。你应该第二次做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

---

正如评论员所说。你应该第二次做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

---

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

---

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

正如我在评论中所说，在对行进行rfft之后，您应该采用FFT而不是rfft，因为rfft结果通常很复杂

---

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

在对行进行rfft后，应采用FFT而不是结果的rfft。rfft的结果是复杂的，不一定是真实的，这就是为什么你会得到警告。@dhanushka希望你能将此作为答案发布！非常感谢。哈哈，没关系，很高兴你的评论和其他答案对我有帮助。我刚刚发布了一个可能会引起兴趣的答案，但它不涉及rfft。在获取行的rfft后，应该使用fft而不是结果的rfft。rfft的结果是复杂的，不一定是真实的，这就是为什么你会得到警告。@dhanushka希望你能将此作为答案发布！非常感谢。哈哈，没关系，很高兴你的评论和其他答案对我有帮助。我刚刚发布了一个可能感兴趣的答案，但它不涉及rfft。