假设我想要在一个网格的每个点上计算一个值。我会定义一个名为func
的函数,该函数将作为参数接收两个值x
和y
并返回第三个值。在下面的示例中,计算这个值需要查找外部字典。然后我会生成一个点的网格,并对每个点评估func
以获得所需的结果。
下面的代码正是这样做的,但它的方式有些绕。首先,我将X和Y坐标矩阵重塑为一维数组,计算所有的值,然后再将结果重塑回矩阵。我的问题是,是否可以用更优雅的方式完成这个任务?
import collections as c
# some arbitrary lookup table
a = c.defaultdict(int)
a[1] = 2
a[2] = 3
a[3] = 2
a[4] = 3
def func(x,y):
# some arbitrary function
return a[x] + a[y]
X,Y = np.mgrid[1:3, 1:4]
X = X.T
Y = Y.T
Z = np.array([func(x,y) for (x,y) in zip(X.ravel(), Y.ravel())]).reshape(X.shape)
print Z
这段代码的目的是生成一组数值,我可以将其与matplotlib中的pcolor
结合使用,创建一种热图类型的图形。
X.reshape(X.size)
和X.ravel()
是一样的。 - mgilsonnumpy.vectorize
。 - mgilson