如何将某些功能应用于 python 网格？

Question

假设我想为网格上的每个点计算一个值。我会定义一些函数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 一起创建热图类型的绘图。

Answer 1

我会使用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]])

使用您的代码，用np.vectorize 装饰func 就足够了，然后您可以将其称为func(X, Y) -- 不需要raveling 或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)