Send array argument to a numba function

Question

I have the following code which sends an array to a numba function as an argument:

import numpy as np
from numba import njit, float64

A = [( 0.0182286178413157, -1.2904019395416308),
 ( 0.5228683581098151,  0.2323207738837293),
 (-0.6056770113345468,  1.5990251249135883),
 (-0.7557841434090988,  1.4641641762952791),
 ( 0.9882455737412416, -1.1838797980930709),
 (-1.2168205368640061,  1.5178083863904257),
 (-0.5566781056044838,  0.2160324328998916),
 ( 0.0671405605855369, -0.4246242749812621),
 ( 0.4806167193998933,  1.0521631181457611),
 ( 0.0563547059786364, -0.8223422191733811)]

A = np.array(A)

@njit(float64(float64[:]))
def distance(a):           
    return a[0]**2 + a[1]**2 + 2*a[0]*a[1]

distance(A)

I can't get the signature string that makes this code run (only works with scalar arguments). Always got this error:

TypeError: No matching definition for argument type(s) array(float64, 2d, C)

Answer 1

You should rewrite the function in this way:

@njit
def distance(A):
    return A[:, 0]**2 + A[:, 1]**2 + 2*A[:, 0]*A[:, 1]