python - numpy meshgrid操作问题

Question

Y, X = np.mgrid[-3:-3:10j, -3:3:10j]

我注意到，当像上面那样对 meshgrids 应用某些操作时，我会收到一个错误，因为这些操作可能与 numpy 不兼容。有时可能有一个 numpy 函数替代 sin、cos 但不是所有函数都像 scipy.integrate 中的 quad 函数。

如何解决这个问题？我需要对整个网格应用操作。

Answer 1

您的问题(带有后续评论)至少可以通过两种不同的方式来理解:

1. 您有一个包含多个参数的函数，并且您希望能够以语法上类似于 numpy native 支持的广播调用的方式调用该函数。性能不是问题，函数的调用语法才是问题。

2. 您有一个包含多个参数的函数，该函数将在一系列 numpy 数组上进行计算，但该函数的实现方式无法利用 numpy 数组的连续内存布局。性能是问题；您会很乐意遍历 numpy 数组并以无聊、普通的老式 for 循环样式调用您的函数，只是这样做太慢了。

对于第 1 项，numpy 提供了一个名为 vectorize 的便捷函数，它接受一个常规可调用对象并返回一个可调用对象，该可调用对象可以使用 numpy 数组作为参数调用，并且将遵守 numpy 的广播规则。

考虑这个人为的例子:

def my_func(x, y):
    return x + 2*y

现在假设我需要在二维网格中的任何地方计算这个函数。这是普通的老式无聊方式:

Y, X  =  np.mgrid[0:10:1, 0:10:1]
Z = np.zeros_like(Y)

for i in range(Y.shape[0]):
    for j in range(Y.shape[1]):
        Z[i,j] = my_func(X[i,j], Y[i,j])

如果我们有几个不同的函数，比如 my_func，那么将这个过程概括为一个函数可能会很好，该函数将给定函数“映射”到二维数组上。

import itertools
def array_map(some_func, *arg_arrays):
    output = np.zeros_like(arg_arrays[0])
    coordinates = itertools.imap(range, output.shape)
    for coord in itertools.product(coordinates):
        args = [arg_array[coord] for arg_array in arg_arrays]
        output[coord] = some_func(*args)
    return output

现在我们可以看到 array_map(my_func, X, Y) 的行为就像嵌套的 for 循环:

In [451]: array_map(my_func, X, Y)
Out[451]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
       [ 4,  5,  6,  7,  8,  9, 10, 11, 12, 13],
       [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       [ 8,  9, 10, 11, 12, 13, 14, 15, 16, 17],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
       [12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
       [14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
       [16, 17, 18, 19, 20, 21, 22, 23, 24, 25],
       [18, 19, 20, 21, 22, 23, 24, 25, 26, 27]])

现在，如果我们可以调用 array_map(my_func) 并省去额外的数组参数不是很好吗？取而代之的是取回一个正在等待执行所需 for 循环的新函数。

我们可以使用 functools.partial 来做到这一点——这样我们就可以像这样编写一个方便的小矢量化程序:

import functools
def vectorizer(regular_function):
    awesome_function = functools.partial(array_map, regular_function)
    return awesome_function

并进行测试:

In [453]: my_awesome_func = vectorizer(my_func)

In [454]: my_awesome_func(X, Y)
Out[454]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
       [ 4,  5,  6,  7,  8,  9, 10, 11, 12, 13],
       [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       [ 8,  9, 10, 11, 12, 13, 14, 15, 16, 17],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
       [12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
       [14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
       [16, 17, 18, 19, 20, 21, 22, 23, 24, 25],
       [18, 19, 20, 21, 22, 23, 24, 25, 26, 27]])

现在 my_awesome_func 的行为就好像您可以直接在 ndarrays 上调用它一样!

我忽略了许多额外的小性能细节、边界检查等，同时制作这个名为 vectorizer 的玩具版本......但幸运的是在 numpy 中有 vectorize这已经做到了!

In [455]: my_vectorize_func = np.vectorize(my_func)

In [456]: my_vectorize_func(X, Y)
Out[456]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
       [ 4,  5,  6,  7,  8,  9, 10, 11, 12, 13],
       [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       [ 8,  9, 10, 11, 12, 13, 14, 15, 16, 17],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
       [12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
       [14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
       [16, 17, 18, 19, 20, 21, 22, 23, 24, 25],
       [18, 19, 20, 21, 22, 23, 24, 25, 26, 27]])

再一次，正如我之前对 OP 的评论和 vectorize 文档中所强调的那样——这不是速度优化。事实上，额外的函数调用开销在某些情况下会比直接编写 for 循环更慢。但是，对于速度不是问题的情况，此方法确实允许您使您的自定义函数遵循与 numpy 相同的调用约定——这可以提高您的库接口(interface)的统一性，并使代码更加一致和可读性更强。

关于第 2 项已经写了很多其他的东西。如果你的问题是你需要优化你的函数以利用连续的内存块并绕过重复的动态类型检查(numpy 数组添加的主要功能到 Python 列表)，那么这里有一些链接可能对您有帮助:

1. <http://pandas.pydata.org/pandas-docs/stable/enhancingperf.html>
2. <http://csl.name/C-functions-from-Python/>
3. <https://jakevdp.github.io/blog/2014/05/09/why-python-is-slow>