如何将某个函数应用于Python的网格?

13 浏览
0 Comments

如何将某个函数应用于Python的网格?

假设我想要计算网格上每个点的值。我会定义一个函数func,该函数以两个参数xy作为输入,并返回第三个值。在下面的示例中,计算这个值需要在外部字典中查找。然后,我会生成一个点的网格,并在每个点上评估func以获得所需的结果。

下面的代码确实实现了这一点,但有点绕弯子。首先,我将X和Y坐标矩阵重新整形为一维数组,计算所有的值,然后将结果重新整形为矩阵。我的问题是,是否可以以更优雅的方式完成这个任务?

import collections as c
# 一些任意的查找表
a = c.defaultdict(int)
a[1] = 2
a[2] = 3
a[3] = 2
a[4] = 3
def func(x,y):
    # 一些任意的函数
    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

这段代码的目的是生成一组值,我可以将其与pcolor在matplotlib中一起使用,创建一个热力图类型的图表。

0