将值分配给 2D numpy 数组作为索引函数的有效方法是什么

Question

这可能是一个愚蠢的问题，但我找不到类似的问题（目前）。 例如，我定义为函数 f(x,y)

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

现在我想输出一个 2D numpy 数组，一个元素的值等于它的索引总和，例如，如果我想要一个 2x2 数组：

arr = [[0, 1],
       [1, 2]]

如果我想要一个 3x3 数组，那么输出应该是：

arr = [[0, 1, 2], 
       [1, 2, 3],
       [2, 3, 4]]

逐个赋值效率不高，尤其是数组大小很大的时候，比如10000*10000，这也是对numpy速度快的浪费。虽然这听起来很基本，但我想不出一个简单快捷的解决方案。最常见和最有效的方法是什么？

顺便说一下，求和指数只是一个例子。我希望该方法也可以推广到任意函数，例如，

def f(x,y):
    return np.cos(x)+np.sin(y)

甚至是更高维的数组，例如 4x4 数组。

Answer 1

您可以使用numpy.indices，它返回一个表示网格索引的数组；您只需要沿 0 轴求和即可：

>>> a = np.random.random((2,2))
>>> np.indices(a.shape).sum(axis=0) # array([[0, 1], [1, 2]])
>>> a = np.random.random((3,3))
>>> np.indices((3,3)).sum(axis=0) #array([[0, 1, 2], [1, 2, 3], [2, 3, 4]])