Numba：矢量化标准 SciPy ufunc 和 numpy.sum() 语法错误

Question

我对使用numba 比较陌生，我想用它来使我的数组计算尽可能高效。有问题的函数是 numba 文档中几个概念的组合。

我在 Scipy 库中使用单一函数

scipy.special.eval_laguerre(n, x, out=None) = <ufunc 'eval_laguerre'>

在点 n 处计算拉盖尔多项式 L_n(x)。

问题 1： Numba 文档清楚地说明了如何使用装饰器 @vectorize 来优化用户编写的 ufunc。 http://numba.pydata.org/numba-doc/0.12/ufuncs.html#generalized-ufuncs

是否有标准程序可以使用 python 库提供的 ufunc 执行此操作？

问题 2：我想为矩阵的每个条目评估 L_n(x)，对于数组中的 n 个值的数组。然后我必须使用以下表达式对这些值求和：

result = np.sum( [eval_laguerre(n, matrix) for n in array], axis=0)

我用过import numpy as np的地方。

如果我要使用广播，我会改为评估：

result = np.sum( eval_laguerre( array[:, None, None], matrix ), axis=0)

axis=0 表示要求和的维度。

我想使用'@jit' 来编译这个部分，但我不确定'numpy.sum() 的过程是什么。目前，上面带有@jit 表达式的表达式给出了语法错误。

result = np.sum( eval_laguerre( array[:, None, None], matrix ), axis=0)
                                                                  ^
SyntaxError: invalid syntax

@jit和np.sum()的正确使用方法是什么？

编辑：回应@hpaulj：

我的想法是numba 可以优化for循环，即

for n in array: 
    eval_laguerre(n, matrix)

这可能吗？如果不是numba，那是什么？ Pythran?

Answer 1

让我们更具体一点：

一个示例数组，我将用于n 和x（您可以选择更实际的值）：

In [782]: A=np.arange(12.).reshape(3,4)

版本，充分利用ufunc的广播能力

In [790]: special.eval_laguerre(A[:,None,:],A[None,:,:]).shape
Out[790]: (3, 3, 4)

或求和：

In [784]: np.sum(special.eval_laguerre(A[:,None,:],A[None,:,:]),0)
Out[784]: 
array([[  3.00000000e+00,  -1.56922399e-01,  -4.86843034e-01,
          7.27719156e-02],
       [  1.37460317e+00,  -4.47492284e+00,   5.77714286e+00,
         -9.71780654e-01],
       [ -1.76222222e+01,   7.00178571e+00,   5.55396825e+01,
         -1.32810866e+02]])

相当于sum内的列表压缩：

In [785]: np.sum([special.eval_laguerre(n,A) for n in A],0)
Out[785]: 
array([[  3.00000000e+00,  -1.56922399e-01,  -4.86843034e-01,
          7.27719156e-02],
       [  1.37460317e+00,  -4.47492284e+00,   5.77714286e+00,
         -9.71780654e-01],
       [ -1.76222222e+01,   7.00178571e+00,   5.55396825e+01,
         -1.32810866e+02]])

或显式循环：

In [786]: x=np.zeros_like(A)    
In [787]: for n in A:
    x += special.eval_laguerre(n, A)

最后一个版本有机会用numba编译。

在简单的时间测试中，ufunc 广播更快：

In [791]: timeit np.sum([special.eval_laguerre(n,A) for n in A],axis=0)
10000 loops, best of 3: 84.8 µs per loop

In [792]: timeit np.sum(special.eval_laguerre(A[:,None,:],A[None,:,:]),0)
10000 loops, best of 3: 43.9 µs per loop

我的猜测是 numba 版本会改进理解版本和显式循环，但可能不会比广播版本更快。