Python 如何加快应用两个二维数组的 scipy 集成

发布于05月17日

我正在应用一个积分函数，使用scipy.integrate到两个2Darray.以下是示例:

from scipy import integrate
import numpy as np

def integrate_lno2(top, bottom, peak_height, peak_width):
    return integrate.quad(lambda x: np.exp(-np.power(x - peak_height, 2)/(2*np.power(peak_width, 2))), top, bottom)[0]

# change row and col to test speed
row = 100; col = 100; peak_height=300; peak_width=60

top = np.linspace(100, 200, row*col).reshape(row, col)
bottom = np.linspace(800, 900, row*col).reshape(row, col)

res = np.zeros((row, col))

for i in range(row):
    for j in range(col):
        res[i, j] = integrate_lno2(top[i, j], bottom[i, j], peak_height, peak_width)

如果二维数组的形状增加，则for循环可能会变慢.我找到了Nuba integrand example，但它不接受上限和下限.

推荐答案

与前面的this个答案一样，您可以使用Numba来加速lambda调用，因为Numpy的开销很大(Numpy没有优化为在scalar上操作，所以执行起来非常慢).更好的是:您可以告诉Numba生成一个C函数，该函数可以直接从Scipy调用，开销非常小(因为它几乎完全消除了缓慢的CPython解释器的开销).您还可以预先计算变量的除法，并将其转换为乘法(更快).

以下是生成的代码:

import numba as nb
import numpy as np
import scipy as sp

factor = -1.0 / (2 * np.power(peak_width, 2))

# change row and col to test speed
row = 100; col = 100; peak_height=300; peak_width=60

@nb.cfunc('float64(float64)')
def compute_numba(x):
    return np.exp(np.power(x - peak_height, 2) * factor)

compute_c = sp.LowLevelCallable(compute_numba.ctypes)

def integrate_lno2(top, bottom):
    return sp.integrate.quad(compute_c, top, bottom)[0]

top = np.linspace(100, 200, row*col).reshape(row, col)
bottom = np.linspace(800, 900, row*col).reshape(row, col)

res = np.zeros((row, col))

for i in range(row):
    for j in range(col):
        res[i, j] = integrate_lno2(top[i, j], bottom[i, j])

在我的机器上，计算循环是roughly 100 times faster.