How do I efficiently fill a numpy ndarray with a function called at each row and column?

Question

I want to create a numpy ndarray by specifying a function of row and column to define the values.

For example, I'd like to be able to do something like this (completely incorrect syntax, but you should get the idea):

>>> np.ndarray((2,3), lambda r,c: 3*r+c)
[[ 0  1  2]
 [ 3  4  5]]

Searching online has yielded nothing, though I've had trouble thinking of exactly how to search for it...

Right now I've got (the equivalent of) the following code, but it's horribly inefficient:

def ndarrayFuncFill(size, func):
    z = np.empty(size)
    for r in np.arange(size[0]):
        for c in np.arange(size[1]):
            z[r][c] = func(r,c)
    return z

>>> ndarrayFuncFill((2,3), lambda r,c: 3*r+c)
array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.]])

Unfortunately, the function I particularly want to use this with right now is not something I can easily rewrite as a ufunc or anything like that. I pretty much have to treat it as a black box.

The function I'm actually interested in using this with (not something so simple as the above lambda), is not one I have permission to post. However, it essentially does interpolation on a lookup table. So you give it a row and column, and then it translates that to indices in a lookup table -- but there's some tricky stuff going on where it's not just a one-to-one lookup, it sometimes does a combination of 'nearby' values, and that sort of thing. So it's not the most efficient function either, but I'd rather not have too many other silly sources of waste like nested for-loops.

Any suggestions?

Answer 1

You could try using index arrays . For your simple example, using np.indices you could do something like:

import numpy as np
r, c = 2, 3
a = np.empty((r, c))
b = np.indices((r, c))
a[b[0], b[1]] = 3 * b[0] + b[1]

So then we have:

>>> a
array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.]])

Answer 2

The fastest solution for your particular example is np.arange(6).reshape(3, 2) . In general you could use np.vectorize for 1D arrays and then reshape if necessary, but this isn't optimized ("The implementation is essentially a for loop.").