如何使用函数创建numpy数组？

Question

我正在使用np.fromfunction根据函数创建特定大小的数组。它看起来像这样：

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)

但是，我收到以下错误：

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

Answer 1

该函数需要处理numpy数组。一个简单的方法是：

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返回f的矢量化版本，它将正确处理数组。