如何将某个函数应用于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结合使用，创建一种热图类型的图形。

- juniper-

---

codereview.stackoverflow.com 可能是一个更好的地方。 - Joran Beasley

我相信 X.reshape(X.size) 和 X.ravel() 是一样的。 - mgilson

你可能还想了解一下 numpy.vectorize。 - mgilson

Answer 1

我会使用numpy.vectorize来"向量化"您的函数。请注意，尽管名称中包含"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) - 无需ravel或reshape：

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)

- mgilson

---

我们也可以直接将vectorize作为函数应用。例如，对内置的chr(..)函数进行向量化：np.vectorize(chr)(np.arange(97, 107))会得到array(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'], dtype='|S1')。 - Evgeni Sergeev

没错 - 但你只是在揭开装饰器语法糖的表层。缺点是每次执行该行时都会创建一个新函数，这是一项相当昂贵的操作。通常情况下，使用装饰器语法时，该行仅在导入时执行一次，因此不会引起太大的担忧。 - mgilson