How do I create a numpy array using a function?

Question

I am using np.fromfunction to create an array of a specific sized based on a function. It looks like this:

import numpy as np
test = [[1,0],[0,2]]
f = lambda i, j: sum(test[i])
matrix = np.fromfunction(f, (len(test), len(test)), dtype=int)

However, I receive the following error:

TypeError: only integer arrays with one element can be converted to an index

Solution

The function needs to handle numpy arrays. An easy way to get this working is:

import numpy as np
test = [[1,0],[0,2]]
f = lambda i, j: sum(test[i])
matrix = np.fromfunction(np.vectorize(f), (len(test), len(test)), dtype=int)

np.vectorize returns a vectorized version of f, which will handle the arrays correctly.