Numba: vectorize standard SciPy ufunc's and numpy.sum() syntax error

Question

I am relatively new to using numba, and I would like to use it to make my array calculations as efficient as possible. The function in question is a combination of several concepts in the numba documentation.

I am using a unitary function in the Scipy library

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

which evaluates a Laguerre polynomial L_n(x) at a point n.

Question 1: The Numba documentation clearly states how to use the decorator @vectorize to optimize a ufunc the user has written. http://numba.pydata.org/numba-doc/0.12/ufuncs.html#generalized-ufuncs

Is there a standard procedure to do this with ufunc provided by python libraries?

Question 2: I would like to evaluate L_n(x) for each entry of a matrix, for an array of n values in an array. I then must sum these values, using the expression:

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

where I have used import numpy as np.

If I were to use broadcasting, I would instead evaluate:

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

where the axis=0 denotes which dimension to sum.

I would like to use '@jit' to compile this section, but I am unsure what the procedure is for 'numpy.sum(). At the moment, the above expression with the @jit expression gives a syntax error.

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

What is the correct way to use @jit and np.sum()?

EDIT: In response to @hpaulj:

My thought was numba could optimize the for loop, i.e.

for n in array: 
    eval_laguerre(n, matrix)

Is this possible at all? If not with numba, then with what? Pythran?

Answer 1

Let's make this more concrete:

A sample array, which I'll use for both n and x (you can choose more realistic values):

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

The version, making full use of the ufunc broadcasting abilties

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

Or summing:

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]])

equivalent with a list comprension inside the 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]])

Or an explicit loop:

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

The last version has a chance of compiling with numba.

In simple time tests, the ufunc broadcasting is faster:

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

My guess is that a numba version will improve on the comprehension version and the explicit loop, but probably not get faster than the broadcasting one.