Apply function on each row (row-wise) of a NumPy array

tattybojangler picture tattybojangler · Aug 10, 2017 · Viewed 39.4k times · Source

So, I have the function -

def function(x):
    x , y = vector
    return exp(((-x**2/200))-0.5*(y+0.05*(x**2) - 100*0.05)**2)

and let's say that I would like to evaluate it at the following points (first column are the x-values and second column are the y-values) -

array([[-1.56113514,  4.51759732],
       [-2.80261623,  5.068371  ],
       [ 0.7792729 ,  6.0169462 ],
       [-1.35672858,  3.52517478],
       [-1.92074891,  5.79966161],
       [-2.79340321,  4.73430001],
       [-2.79655868,  5.05361163],
       [-2.13637747,  5.39255837],
       [ 0.17341809,  3.60918261],
       [-1.22712921,  4.95327158]])

i.e. I would like to pass the function the first row of values and evaluate, then the second row and evaluate etc. and then the final result would be an array of the values evaluated at these points (so, an array consisting of 10 values).

So, for example, if the function was, say, a bivariate normal distribution -

def function2(x):

function2 = (mvnorm.pdf(x,[0,0],[[1,0],[0,1]]))

return function2

and I passed the above values into this function, I would get -

array([  1.17738907e-05,   1.08383957e-04,   1.69855078e-04,
         5.64757613e-06,   1.37432346e-05,   1.44032800e-04,
         1.33426313e-05,   1.97822328e-06,   6.56121709e-08,
         4.67076770e-05])

So basically, I am looking for a way to rewrite the function so that it can do this. Moreover, I would like to keep the function as a function of one variable only (i.e. only a function of x).

Thank you for your help!

Answer

cs95 picture cs95 · Aug 10, 2017

You can use np.apply_along_axis:

np.apply_along_axis(function, 1, array)

The first argument is the function, the second argument is the axis along which the function is to be applied. In your case, it is the first axis. The last argument is the array, of course.

You should be warned, however, that apply_along_axis is only a convenience function, not a magic bullet. It has a severe speed limitation, since it just hides a loop. You should always try to vectorize your computation, where possible. Here's how I'd do this:

v = array[:, 0] ** 2   # computing just once  
return np.exp((-v / 200) - 0.5 * (array[:, 1] + 0.05 * v - 5) ** 2)