基于索引进行NumPy映射计算

Question

我有一个形状为 (高度, 宽度) 的 NumPy 数组。现在我想计算这个数组中所有元素的值。这个值是根据元素的位置 [x,y] 来决定的。

假设我想给所有元素赋值为 (x**2+y**2)/2。我希望能不使用 for 循环来实现这个。有没有什么办法可以做到这一点呢？

Answer 1

应该像这样：

numpy.fromfunction(lambda i, j: (i**2+j**2)/2, (3, 3), dtype=int)

更多信息请查看： http://docs.scipy.org/doc/numpy/reference/generated/numpy.fromfunction.html#numpy.fromfunction

Answer 2

你也可以看看 meshgrid、mgrid 和/或 indices：

>>> H, W = 4,5
>>> x, y = np.indices([H, W])
>>> m
array([[  0. ,   0.5,   2. ,   4.5,   8. ],
       [  0.5,   1. ,   2.5,   5. ,   8.5],
       [  2. ,   2.5,   4. ,   6.5,  10. ],
       [  4.5,   5. ,   6.5,   9. ,  12.5]])

之所以这样做是因为 x 和 y 是包含合适的 x 和 y 坐标的数组：

>>> x
array([[0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3]])
>>> y
array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])

meshgrid 和 mgrid 让你可以更精细地控制，比如：

>>> x, y = np.meshgrid(np.linspace(0, 1, 5), np.linspace(0, 10, 3))
>>> x
array([[ 0.  ,  0.25,  0.5 ,  0.75,  1.  ],
       [ 0.  ,  0.25,  0.5 ,  0.75,  1.  ],
       [ 0.  ,  0.25,  0.5 ,  0.75,  1.  ]])
>>> y
array([[  0.,   0.,   0.,   0.,   0.],
       [  5.,   5.,   5.,   5.,   5.],
       [ 10.,  10.,  10.,  10.,  10.]])