python - Efficient evaluation of a function at every cell of a NumPy array

Question

Given a NumPy array A, what is the fastest/most efficient way to apply the same function, f, to every cell?

1. Suppose that we will assign to A(i,j) the f(A(i,j)).

2. The function, f, doesn't have a binary output, thus the mask(ing) operations won't help.

Is the "obvious" double loop iteration (through every cell) the optimal solution?

Answer 1

You could just vectorize the function and then apply it directly to a Numpy array each time you need it:

import numpy as np

def f(x):
    return x * x + 3 * x - 2 if x > 0 else x * 5 + 8

f = np.vectorize(f)  # or use a different name if you want to keep the original f

result_array = f(A)  # if A is your Numpy array

It's probably better to specify an explicit output type directly when vectorizing:

f = np.vectorize(f, otypes=[np.float])