假设我想为网格上的每个点计算一个值。我将定义一些函数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
一起使用来创建热图类型的绘图。
我会用
numpy.vectorize
来“矢量化”你的函数。请注意,尽管有这个名称,vectorize
并不是为了让代码运行得更快——只是稍微简化一下。下面是一些例子:
用你的代码,用
np.vectorize
来装饰func
就足够了,然后你可以称之为func(X, Y)
——不需要ravel
或reshape
必要的:相关问题 更多 >
编程相关推荐