根据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.fromfunction：

np.fromfunction(f,(3,4,7,8),dtype=int)

只要f 是“numpy 感知”的，它就可以工作，即可以处理数组参数。

Answer 2

你的f，至少在这个例子中，可以处理数组，因为它使用了 numpy 矢量化运算符+ 和*。

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

使用np.ix_ 生成可以相互广播的索引，我可以生成一个与您的迭代A 匹配的数组：

In [14]: A1 = f(*np.ix_(np.arange(3),np.arange(4),np.arange(7),np.arange(8)))                                   
In [15]: np.allclose(A,A1)                                                                                      
Out[15]: True

现在如果f 可以使用数组，我们就不能这样做了。一个经典的例子是使用math.sin 而不是np.sin。或者是if 测试。

ix_ 的方式很快，因为它只调用一次f。我还可以使用np.ndindex 将所有索引生成为元组，并多次调用f：

In [24]: A2 = np.zeros_like(A)                                                                                  
In [25]: for ijk in np.ndindex(*A2.shape): A2[ijk]=f(*ijk)                                                      
In [26]: np.allclose(A,A2)                                                                                      
Out[26]: True

速度将与您的A 评估相似。

In [29]: %%timeit  
    ...: A2 = np.zeros_like(A) 
    ...: for ijk in np.ndindex(*A2.shape): A2[ijk]=f(*ijk)  
    ...:  
    ...:                                                                                                        
758 µs ± 11.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [30]: timeit A1 = f(*np.ix_(np.arange(3),np.arange(4),np.arange(7),np.arange(8)))                            
54.3 µs ± 2.44 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [31]: %%timeit 
    ...: 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 
 ....
    ...:                 A[i1, i2, i3, i4] = v 
    ...:                                                                                                        
506 µs ± 2.03 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

想想看，我们可以用更典型的 Python range 循环来清理你的循环：

In [35]: %%timeit 
    ...: A = np.zeros((3,4,7,8), dtype=int) 
    ...: for i1 in range(A.shape[0]): 
    ...:     for i2 in range(A.shape[1]): 
    ...:         for i3 in range(A.shape[2]): 
    ...:             for i4 in range(A.shape[3]): 
    ...:                 v = f(i1, i2, i3, i4) 
    ...:                 A[i1, i2, i3, i4] = v 
    ...:                                                                                                        
343 µs ± 17.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 3

用np.apply_along_axis方式解决问题

我把它分解成一个更简单的例子来解释我的想法背后的逻辑，然后继续前进。

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

例子：

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

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

请注意，函数 f 的输入已更改，因为它的输入是形状为 (4,) 的 np.array，而不是 4 元素元组。

我们将使用

np.indices(A.shape)

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

np.apply_along_axis(f, axis=0, arr=np.indices(A.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

现在让我们转到 OP 给出的示例。代码是一样的，我们只是交换f和A

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

A = np.zeros((3,4,7,8), dtype=int)
np.apply_along_axis(f, axis=0, arr=np.indices(A.shape))

由于结果的形状为(3, 4, 7, 8)，我不会在这里发布。

计时结果

np.apply_along_axis 被证明比手动迭代慢很多，即使对于相当大的数组也是如此。

以下是一些计时结果：

shape            manual            np.apply_along_axis
100,100,100,10  4.18 s ± 39.7 ms     34.8 s ± 298 ms
100,100,10,10   412 ms ± 6.23 ms     3.49 s ± 13.4 ms
3,4,7,8         285 µs ± 1.79 µs     2.43 ms ± 35.6 µs