numpy - 在 numpy 中优雅地生成结果数组

Question

我有我的 X 和 Y numpy 数组:

X = np.array([0,1,2,3])
Y = np.array([0,1,2,3])

我的函数将 x,y 值映射到 Z 点:

def z(x,y):
    return x+y

我希望生成 3D 绘图所需的明显内容:对应 Z 值的二维 numpy 数组。我相信它应该看起来像:

Z = np.array([[0, 1, 2, 3],
              [1, 2, 3, 4],
              [2, 3, 4, 5],
              [3, 4, 5, 6]])

我可以用几行来完成此操作，但我正在寻找最简短、最优雅的代码段。

Answer 1

对于支持数组的函数，使用开放网格更为经济:

>>> import numpy as np
>>> 
>>> X = np.array([0,1,2,3])
>>> Y = np.array([0,1,2,3])
>>> 
>>> def z(x,y):
...     return x+y
... 
>>> XX, YY = np.ix_(X, Y)
>>> XX, YY
(array([[0],
       [1],
       [2],
       [3]]), array([[0, 1, 2, 3]]))
>>> z(XX, YY)
array([[0, 1, 2, 3],
       [1, 2, 3, 4],
       [2, 3, 4, 5],
       [3, 4, 5, 6]])

如果您的网格轴是范围，您可以使用np.ogrid直接创建网格

>>> XX, YY = np.ogrid[:4, :4]
>>> XX, YY
(array([[0],
       [1],
       [2],
       [3]]), array([[0, 1, 2, 3]]))

如果该函数不支持数组，您可以使用np.vectorize来实现:

>>> def f(x, y):
...    if x > y:
...        return x
...    else:
...        return -x
... 
>>> np.vectorize(f)(*np.ogrid[-3:4, -3:4])
array([[ 3,  3,  3,  3,  3,  3,  3],
       [-2,  2,  2,  2,  2,  2,  2],
       [-1, -1,  1,  1,  1,  1,  1],
       [ 0,  0,  0,  0,  0,  0,  0],
       [ 1,  1,  1,  1, -1, -1, -1],
       [ 2,  2,  2,  2,  2, -2, -2],
       [ 3,  3,  3,  3,  3,  3, -3]])