numpy - 在点网格上评估函数

Question

有什么好的方法可以生成一个numpy数组，这个数组包含在n维点网格上计算函数的值呢？

比如说，我想计算一个函数，这个函数的定义是：

def func(x, y):
    return <some function of x and y>

假设我想在一个二维的点数组上进行计算，x的值从0到4，分成十个步骤，而y的值从-1到1，分成二十个步骤。用numpy怎么做比较好呢？

附注：这个问题在StackOverflow上以各种形式被问过很多次，但我找不到一个简洁明了的问题和答案。我发这个帖子是为了提供一个简单明了的解决方案（见下文）。

Answer 1

我来分享一下我的看法：

    import numpy as np

    x = np.linspace(0, 4, 10)
    y = np.linspace(-1, 1, 20)

    [X, Y] = np.meshgrid(x, y, indexing = 'ij', sparse = 'true')

    def func(x, y):
        return x*y/(x**2 + y**2 + 4)
        # I have defined a function of x and y.

   func(X, Y)

Answer 2

如果你的函数实际上是接收一个包含 d 个元素的元组，比如说 f((x1,x2,x3,...xd))（例如 scipy.stats.multivariate_normal 函数），而你想要在 N^d 的组合/网格上评估 f，你也可以这样做（以二维为例）：

x=np.arange(-1,1,0.2)   # each variable is instantiated N=10 times
y=np.arange(-1,1,0.2)
Z=f(np.dstack(np.meshgrid(x,y)))    # result is an NxN (10x10) matrix, whose entries are f((xi,yj))

这里的 np.dstack(np.meshgrid(x,y)) 创建了一个 10x10 的“矩阵”（严格来说是一个 10x10x2 的 numpy 数组），这个矩阵里的每个元素都是要被 f 评估的二维元组。

Answer 3

我用这个函数来准备 X、Y、Z 的值，以便绘图：

def npmap2d(fun, xs, ys, doPrint=False):
  Z = np.empty(len(xs) * len(ys))
  i = 0
  for y in ys:
    for x in xs:
      Z[i] = fun(x, y)
      if doPrint: print([i, x, y, Z[i]])
      i += 1
  X, Y = np.meshgrid(xs, ys)
  Z.shape = X.shape
  return X, Y, Z

用法：

def f(x, y): 
  # ...some function that can't handle numpy arrays

X, Y, Z = npmap2d(f, np.linspace(0, 0.5, 21), np.linspace(0.6, 0.4, 41))

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_wireframe(X, Y, Z)

也可以通过使用 map 来达到同样的效果：

xs = np.linspace(0, 4, 10)
ys = np.linspace(-1, 1, 20)
X, Y = np.meshgrid(xs, ys)
Z = np.fromiter(map(f, X.ravel(), Y.ravel()), X.dtype).reshape(X.shape)

Answer 4

import numpy as np

def func(x, y):
    return np.sin(y * x)

xaxis = np.linspace(0, 4, 10)
yaxis = np.linspace(-1, 1, 20)
x, y = np.meshgrid(xaxis, yaxis)
result = func(x, y)

当然可以！请把你想要翻译的内容发给我，我会帮你用简单易懂的语言解释清楚。

Answer 5

更简短、更快速、更清晰的回答，避免使用meshgrid：

import numpy as np

def func(x, y):
    return np.sin(y * x)

xaxis = np.linspace(0, 4, 10)
yaxis = np.linspace(-1, 1, 20)
result = func(xaxis[:,None], yaxis[None,:])

如果你用像x^2+y这样的函数，这样在内存中会更快，因为x^2是在一维数组上计算的（而不是二维数组），只有在你进行“+”运算时才会增加维度。而使用meshgrid时，x^2是在二维数组上计算的，实际上每一行都是一样的，这会导致时间大幅增加。

编辑：这个“x[:,None]”会把x变成一个二维数组，但第二个维度是空的。这个“None”跟用“x[:,numpy.newaxis]”是一样的。Y也是这样处理的，只不过是把第一个维度设为空。

编辑：在三维情况下：

def func2(x, y, z):
    return np.sin(y * x)+z

xaxis = np.linspace(0, 4, 10)
yaxis = np.linspace(-1, 1, 20)
zaxis = np.linspace(0, 1, 20)
result2 = func2(xaxis[:,None,None], yaxis[None,:,None],zaxis[None,None,:])

这样的话，如果你想的话，可以很容易地扩展到n维，使用尽可能多的None或:，每个:代表一个维度，每个None代表一个“空”维度。下一个例子会更详细地展示这些空维度是如何工作的。你会看到，如果使用None，形状会发生变化，显示在下一个例子中这是一个三维对象，但空维度只有在你和一个在这些维度上有实际内容的对象相乘时才会被填充（听起来复杂，但下一个例子会说明我的意思）

In [1]: import numpy

In [2]: a = numpy.linspace(-1,1,20)

In [3]: a.shape
Out[3]: (20,)

In [4]: a[None,:,None].shape 
Out[4]: (1, 20, 1)

In [5]: b = a[None,:,None] # this is a 3D array, but with the first and third dimension being "empty"
In [6]: c = a[:,None,None] # same, but last two dimensions are "empty" here

In [7]: d=b*c 

In [8]: d.shape # only the last dimension is "empty" here
Out[8]: (20, 20, 1)

编辑：不需要自己输入None

def ndm(*args):
    return [x[(None,)*i+(slice(None),)+(None,)*(len(args)-i-1)] for i, x in enumerate(args)]


x2,y2,z2  = ndm(xaxis,yaxis,zaxis)
result3 = func2(x2,y2,z2)

这样，你可以通过将你给ndm的第一个参数作为第一个完整维度，第二个作为第二个完整维度，等等，来自动创建额外的空维度，这样做的效果和之前手动输入的None类型语法是一样的。

简单解释：执行x2, y2, z2 = ndm(xaxis, yaxis, zaxis)和执行

x2 = xaxis[:,None,None]
y2 = yaxis[None,:,None]
z2 = zaxis[None,None,:]

是一样的，但ndm方法也应该适用于更多维度，而不需要像刚才那样在多行中硬编码None切片。这在numpy 1.8之前的版本中也能工作，而numpy.meshgrid只有在numpy 1.8或更高版本中才能处理超过2维的情况。