python中grid函数应用_python - 如何将一些函数应用于python meshgrid？

假设我想计算网格上每个点的值。我将定义一些函数func，它将两个值x和y作为参数，并返回第三个值。在下面的示例中，计算此值需要在外部字典中查找。然后，我将生成一个点网格，并对每个点进行评估，以获得我想要的结果。

下面的代码确实做到了这一点，但有点迂回。首先，我将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中与func一起使用这些值来创建heatmap类型的图。

最佳答案：

我将使用numpy.vectorize来“向量化”您的函数。请注意，尽管有这个名称，vectorize并不是为了让代码运行得更快——只是稍微简化一下。

以下是一些例子：>>> import numpy as np

>>> @np.vectorize

... def foo(a, b):

... return a + b

...

>>> foo([1,3,5], [2,4,6])

array([ 3, 7, 11])

>>> foo(np.arange(9).reshape(3,3), np.arange(9).reshape(3,3))

array([[ 0, 2, 4],

[ 6, 8, 10],

[12, 14, 16]])

有了您的代码，就足够用func来修饰np.vectorize了，然后您可能只需要将其称为func(X, Y)——noraveling或reshapeing必要：

import numpy as np

import collections as c

# some arbitrary lookup table

a = c.defaultdict(int)

a[1] = 2

a[2] = 3

a[3] = 2

a[4] = 3

@np.vectorize

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 = func(X, Y)