根据ndarray中的索引设置值的通用方法？

Question

我在下面写下了这个，但对我来说看起来很丑/笨拙。有没有更好的方法来遍历所有元素同时跟踪索引（即跟踪我在哪个地址）并设置一个值，该值是一个函数索引？

我可以想象，如果我有更大的尺寸，这种情况将变得更加难看。应该设计出一种更好的方法来做到这一点。

import numpy as np

def f(i,j,k,l):
    return i + 2 * j+ 3 * k + 4 * l

A = np.zeros((3,4,7,8), dtype=int)

i1 = i2 = i3 = i4 = -1

i1 = -1
for r1 in A:
    i1 += 1

    i2 = -1
    for r2 in r1:
        i2 += 1

        i3 = -1
        for r3 in r2:

            i3 += 1

            i4 = -1
            for r4 in r3:
                i4 += 1

                v = f(i1, i2, i3, i4)

                print ("Setting value {0} at ({1}, {2}, {3}, {4})".format(v, i1, i2, i3, i4))

                A[i1, i2, i3, i4] = v

print(A)

Answer 1

我将其分解为一个更简单的示例，以解释我的想法背后的逻辑并从中继续前进。

我想使用np.indices和np.apply_along_axis的组合来返回给定形状的索引网格，通过它们我们可以沿多维数组的给定轴应用函数。

示例：

def f(arr):
    return arr[0] + 2 * arr[1]

A=np.zeros((2,3))

现在我们使用

indices=np.indices(A.shape)

得到矩阵为A的索引网格。我们可以轻松地将该函数应用为

np.apply_along_axis(f, axis=0, arr=np.indices(np.zeros((2,3)).shape))

导致]

array([[0, 2, 4],    # 0 + 2*0, 0 + 2*1, 0 +2*2
       [1, 3, 5]])   + 1 + 2*0, 1 + 2*1, 1 +2*2