How to use np.sum with multiple arrays and axis?

Numba Support: How do I do ...?
1

I’m quite new to numba. I would have thought the following basic numpy code would compile.

@nb.jit(nopython=True)
def vccf_multi(v: np.ndarray,mode: int) → np.ndarray:
n_molecules = v.shape[0]
res = np.zeros(v.shape[2])

if mode == 1:
    for i in range(n_molecules):
        for j in range(i, n_molecules):
            if i != j:
                delta_x = np.convolve(v[i,0], np.flipud(v[j,0]))
                delta_y = np.convolve(v[i,1], np.flipud(v[j,1]))
                delta_z = np.convolve(v[i,2], np.flipud(v[j,2]))
                vccf_delta = np.sum((delta_x[v.shape[2]-1::-1], \
                                     delta_y[v.shape[2]-1::-1],delta_z[v.shape[2]-1::-1]),axis=0)
                res += vccf_delta
                vccf_delta = np.sum((delta_x[v.shape[2]-1:], \
                                      delta_y[v.shape[2]-1:],delta_z[v.shape[2]-1:]),axis=0)
                res += vccf_delta
    
return res

The goal of the function is to calculate a correlation function, after np.convolve, all three parts of correlation function will be added by column.
I failed in compilation and got the error followed:

TypingError: No implementation of function Function(<function sum at 0x0000022495400EE0>) found for signature:

sum(UniTuple(array(float64, 1d, A) x 3), axis=Literalint)

There are 2 candidate implementations:
- Of which 2 did not match due to:
Overload of function ‘sum’: File: numba\core\typing\npydecl.py: Line 379.
With argument(s): ‘(UniTuple(array(float64, 1d, A) x 3), axis=Literalint)’:
No match.

During: resolving callee type: Function(<function sum at 0x0000022495400EE0>)

So, it’s that means numba don’t support np.sum with multiple array?
I also try a simple test for np.sum like this:

@nb.njit()
def s_u_m() → np.ndarray:

a = np.array([1,2,3])
b = np.array([3,5,3])

res = np.sum((a,b))

return res

s_u_m()

it also reported the same error:
No implementation of function Function(<function sum at 0x0000021A46842DC0>) found for signature:

sum(UniTuple(array(int64, 1d, C) x 2))

There are 2 candidate implementations:

  • Of which 2 did not match due to:
    Overload of function ‘sum’: File: numba\core\typing\npydecl.py: Line 379.
    With argument(s): ‘(UniTuple(array(int64, 1d, C) x 2))’:
    No match.

During: resolving callee type: Function(<function sum at 0x0000021A46842DC0>)