如何在 numba.njit 中进行离散傅里叶变换（FFT）？

Question

如何在 numba.njit 中进行离散傅里叶变换（FFT）？

Art*_*nov 8 python numpy fft scipy numba

各位程序员大家好

我正在尝试与装饰器一起制作discrete Fourier transform一个：minimal working examplenumba.njit

import numba
import numpy as np
import scipy
import scipy.fftpack

@numba.njit
def main():
    wave = [[[0.09254795,  0.10001078,  0.10744892, 0.07755555,  0.08506225, 0.09254795],
          [0.09907245,  0.10706145,  0.11502401,  0.08302302,  0.09105898, 0.09907245],
          [0.09565098,  0.10336405,  0.11105158,  0.08015589,  0.08791429, 0.09565098],
          [0.00181467,  0.001961,    0.00210684,  0.0015207,   0.00166789, 0.00181467]],
         [[-0.45816267, - 0.46058367, - 0.46289091, - 0.45298182, - 0.45562851, -0.45816267],
          [-0.49046506, - 0.49305676, - 0.49552669, - 0.48491893, - 0.48775223, -0.49046506],
          [-0.47352483, - 0.47602701, - 0.47841162, - 0.46817027, - 0.4709057, -0.47352483],
          [-0.00898358, - 0.00903105, - 0.00907629, - 0.008882, - 0.00893389, -0.00898358]],
         [[0.36561472,  0.36057289,  0.355442,  0.37542627,  0.37056626, 0.36561472],
          [0.39139261,  0.38599531,  0.38050268,  0.40189591,  0.39669325, 0.39139261],
          [0.37787385,  0.37266296,  0.36736003,  0.38801438,  0.38299141, 0.37787385],
          [0.00716892,  0.00707006,  0.00696945,  0.0073613,  0.00726601, 0.00716892]]]

    new_fft = scipy.fftpack.fft(wave)


if __name__ == '__main__':
    main()

Run Code Online (Sandbox Code Playgroud)

输出：

C:\Users\Artur\Anaconda\python.exe C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py
Traceback (most recent call last):
  File "C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py", line 25, in <module>
    main()
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'fft' of type Module(<module 'scipy.fftpack' from 'C:\\Users\\Artur\\Anaconda\\lib\\site-packages\\scipy\\fftpack\\__init__.py'>)

File "test2.py", line 21:
def main():
    <source elided>

    new_fft = scipy.fftpack.fft(wave)
    ^

[1] During: typing of get attribute at C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py (21)

File "test2.py", line 21:
def main():
    <source elided>

    new_fft = scipy.fftpack.fft(wave)
    ^


Process finished with exit code 1

Run Code Online (Sandbox Code Playgroud)

不幸的是scipy.fftpack.fft，这似乎是一个不支持的遗留功能numba。所以我寻找替代方案。我找到了两个：

1. scipy.fft(wave)这是上述遗留功能的更新版本。它产生以下错误输出：

C:\Users\Artur\Anaconda\python.exe C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py
Traceback (most recent call last):
  File "C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py", line 25, in <module>
    main()
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Module(<module 'scipy.fft' from 'C:\\Users\\Artur\\Anaconda\\lib\\site-packages\\scipy\\fft\\__init__.py'>) with parameters (list(list(list(float64))))
No type info available for Module(<module 'scipy.fft' from 'C:\\Users\\Artur\\Anaconda\\lib\\site-packages\\scipy\\fft\\__init__.py'>) as a callable.
[1] During: resolving callee type: Module(<module 'scipy.fft' from 'C:\\Users\\Artur\\Anaconda\\lib\\site-packages\\scipy\\fft\\__init__.py'>)
[2] During: typing of call at C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py (21)


File "test2.py", line 21:
def main():
    <source elided>

    new_fft = scipy.fft(wave)
    ^


Process finished with exit code 1

Run Code Online (Sandbox Code Playgroud)

2. np.fft.fft(wave)似乎支持但也会产生错误：

C:\Users\Artur\Anaconda\python.exe C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py
Traceback (most recent call last):
  File "C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py", line 25, in <module>
    main()
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'fft' of type Module(<module 'numpy.fft' from 'C:\\Users\\Artur\\Anaconda\\lib\\site-packages\\numpy\\fft\\__init__.py'>)

File "test2.py", line 21:
def main():
    <source elided>

    new_fft = np.fft.fft(wave)
    ^

[1] During: typing of get attribute at C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py (21)

File "test2.py", line 21:
def main():
    <source elided>

    new_fft = np.fft.fft(wave)
    ^


Process finished with exit code 1

Run Code Online (Sandbox Code Playgroud)

您知道fft与numba.njit装饰器一起使用的函数吗？

Answer 1

nor*_*ok2 5

如果您对 1D DFT 感到满意，您也可以使用 FFT。\n这里报告了一个适用于fft_1d()任意输入大小的 Numba 友好实现：

\n\n

import cmath\nimport numpy as np\nimport numba as nb\n\n\n@nb.jit\ndef ilog2(n):\n    result = -1\n    if n < 0:\n        n = -n\n    while n > 0:\n        n >>= 1\n        result += 1\n    return result\n\n\n@nb.njit(fastmath=True)\ndef reverse_bits(val, width):\n    result = 0\n    for _ in range(width):\n        result = (result << 1) | (val & 1)\n        val >>= 1\n    return result\n\n\n@nb.njit(fastmath=True)\ndef fft_1d_radix2_rbi(arr, direct=True):\n    arr = np.asarray(arr, dtype=np.complex128)\n    n = len(arr)\n    levels = ilog2(n)\n    e_arr = np.empty_like(arr)\n    coeff = (-2j if direct else 2j) * cmath.pi / n\n    for i in range(n):\n        e_arr[i] = cmath.exp(coeff * i)\n    result = np.empty_like(arr)\n    for i in range(n):\n        result[i] = arr[reverse_bits(i, levels)]\n    # Radix-2 decimation-in-time FFT\n    size = 2\n    while size <= n:\n        half_size = size // 2\n        step = n // size\n        for i in range(0, n, size):\n            k = 0\n            for j in range(i, i + half_size):\n                temp = result[j + half_size] * e_arr[k]\n                result[j + half_size] = result[j] - temp\n                result[j] += temp\n                k += step\n        size *= 2\n    return result\n\n\n@nb.njit(fastmath=True)\ndef fft_1d_arb(arr, fft_1d_r2=fft_1d_radix2_rbi):\n    """1D FFT for arbitrary inputs using chirp z-transform"""\n    arr = np.asarray(arr, dtype=np.complex128)\n    n = len(arr)\n    m = 1 << (ilog2(n) + 2)\n    e_arr = np.empty(n, dtype=np.complex128)\n    for i in range(n):\n        e_arr[i] = cmath.exp(-1j * cmath.pi * (i * i) / n)\n    result = np.zeros(m, dtype=np.complex128)\n    result[:n] = arr * e_arr\n    coeff = np.zeros_like(result)\n    coeff[:n] = e_arr.conjugate()\n    coeff[-n + 1:] = e_arr[:0:-1].conjugate()\n    return fft_convolve(result, coeff, fft_1d_r2)[:n] * e_arr / m\n\n\n@nb.njit(fastmath=True)\ndef fft_convolve(a_arr, b_arr, fft_1d_r2=fft_1d_radix2_rbi):\n    return fft_1d_r2(fft_1d_r2(a_arr) * fft_1d_r2(b_arr), False)\n\n\n@nb.njit(fastmath=True)\ndef fft_1d(arr):\n    n = len(arr)\n    if not n & (n - 1):\n        return fft_1d_radix2_rbi(arr)\n    else:\n        return fft_1d_arb(arr)\n

Run Code Online (Sandbox Code Playgroud)\n\n

与 na\xc3\xafve DFT 算法（与此dft_1d()基本相同）相比，您获得了数量级的提升，但通常仍然比慢很多。np.fft.fft()

\n\n

相对速度根据输入大小而变化很大。\n对于2 的幂输入，这通常在的一个数量级内np.fft.fft()。

\n\n

对于非2 的幂，这通常在的两个数量级内np.fft.fft()。

\n\n

对于最坏的情况（素数左右，这里是 2 + 1 的幂），这是的 1 倍np.fft.fft()。

\n\n

FFT 时序的非线性行为是由于需要更复杂的算法来处理非2 的幂的任意输入大小的结果。这会影响此实现和来自的实现np.fft.fft()，但np.fft.fft()包含更多优化，使其平均性能更好。

\n\n

此处显示了 2 次幂 FFT 的替代实现。

\n

Answer 2

Art*_*nov 2

我找到了解决方法。现在，请记住，像这样的函数numpy.fft.fft有很多方便的操作，所以如果你不像我一样陷入困境，你应该使用它们。

以下njit函数对 a 执行discrete fourier transforma 操作one dimensional array：

import numba
import numpy as np
import cmath

def dft(wave=None):
    dft = np.fft.fft(wave)
    return dft

@numba.njit
def dft_njit(wave=None):
    N = len(wave)
    dft_njit = np.zeros(N, dtype=np.complex128)
    for i in range(N):
        series_element = 0
        for n in range(N):
            series_element += wave[n] * cmath.exp(-2j * cmath.pi * i * n * (1 / N))
        dft_njit[i] = series_element
    return dft_njit

if __name__ == '__main__':

    wave = [1,2,3,4,5]
    wave = np.array(wave)

    print(f' dft: \n{dft(wave=wave)}')
    print(f' dft_njit: \n{dft_njit(wave=wave)}')

Run Code Online (Sandbox Code Playgroud)

输出：

 dft: 
[15. +0.j         -2.5+3.4409548j  -2.5+0.81229924j -2.5-0.81229924j
 -2.5-3.4409548j ]
 dft_njit: 
[15. +0.j         -2.5+3.4409548j  -2.5+0.81229924j -2.5-0.81229924j
 -2.5-3.4409548j ]

Run Code Online (Sandbox Code Playgroud)

归档时间：	5 年，5 月前
查看次数：	6317 次
最近记录：	4 年，7 月前