计算 Python 中大型复杂数组的指数 [exp()] 函数的最快方法

Question

我正在开发使用 scipy 的 complex_ode 对 ODE 进行积分的代码，其中被积函数包括傅里叶变换和作用于大量复数值的指数运算符。

为了优化性能，我对此进行了分析，发现主要瓶颈（在使用 PyFFTW 等优化 FFT 之后）在以下行中：

val = np.exp(float_value * arr)

我目前正在使用 numpy，我知道它调用 C 代码 - 因此应该很快。请问有什么办法可以进一步提高性能吗？

我研究过使用 Numba，但由于我的主循环也包含 FFT，我认为它无法编译（nopython=True 标志会导致错误），因此，我怀疑它没有任何好处。

这是我要优化的代码的测试示例：

arr = np.random.rand(2**14) + 1j *np.random.rand(2**14)
float_value = 0.5
%timeit np.exp(float_value * arr)

任何建议都欢迎谢谢。

Answer 1

我们可以利用 numexpr module，它对涉及先验操作的大数据非常有效 -

In [91]: arr = np.random.rand(2**14) + 1j *np.random.rand(2**14)
    ...: float_value = 0.5
    ...: 

In [92]: %timeit np.exp(float_value * arr)
1000 loops, best of 3: 739 µs per loop

In [94]: import numexpr as ne

In [95]: %timeit ne.evaluate('exp(float_value*arr)')
1000 loops, best of 3: 241 µs per loop

这似乎与文档中所述的 expected performance 一致。

计算 Python 中大型复杂数组的指数 [exp()] 函数的最快方法

Fastest way to calculate exponential [exp()] function of large complex array in Python

python

optimization

numpy

ode

numba