np.ufunc.at 和 np.ix_ 的 Numpy 多维索引答案

【问题标题】：Numpy multidimensional indexing for np.ufunc.at and np.ix_np.ufunc.at 和 np.ix_ 的 Numpy 多维索引
【发布时间】：2020-11-18 11:08:47
【问题描述】：

我想知道如何从一个数组中获取索引并与另一个数组相乘。我有两个 4d 数组和一个 2d 索引数组：

base = np.ones((2, 3, 5, 5))
to_multiply = np.arange(120).reshape(2, 3, 4, 5)
index = np.array([[0, 2, 4, 2], [0, 3, 3, 2]])

索引数组的行索引对应base和to_multiply的第1维，索引数组的值对应base的第3维。我想根据索引从base中取出切片并与to_multiply相乘。

使用 for 循环和 np.multiply.at（因为我可能有相同的索引）我可以通过以下方式实现它：

for i, x in enumerate(index):
    np.multiply.at(base[i, :, :, :], np.s_[:, x, :], to_multiply[i, :, :, :])

上面的正确性可以通过以下方式验证：

to_multiply[0, 0, :, 0]
array([ 0,  5, 10, 15])

base[0, 0, :, 0]
array([ 0.,  1., 75.,  1., 10.])

但是，我想知道是否有使用 np.multiply.at 和 np.ix_ 的单行解决方案

我尝试使用 np.ix_ 但我对此感到非常困惑，因为在这种情况下它是多维的。

【问题讨论】：

标签： python numpy numpy-ndarray numpy-slicing numpy-ufunc

【解决方案1】：

不能用ix_ 完成。来自其文档：

这个函数接受 N 个一维序列并返回 N 个输出，其中 N 每个维度，

您的index 是 2d。

但是，我们可以做等效的“手工”：

In [196]: np.multiply.at(base1, (np.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:]), to_
     ...: multiply)                                                                                  
In [197]: np.allclose(base,base1)                                                                    
Out[197]: True

目标是使 3 个数组一起广播以匹配 to_multiply（最后一个大小为 5 的维度除外）。

即 (2,1,1), (1,3,1) 和 (2,1,4) => (2,3,4)

In [199]: np.broadcast_arrays(np.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:])        
Out[199]: 
[array([[[0, 0, 0, 0],
         [0, 0, 0, 0],
         [0, 0, 0, 0]],
 
        [[1, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]]]),
 array([[[0, 0, 0, 0],
         [1, 1, 1, 1],
         [2, 2, 2, 2]],
 
        [[0, 0, 0, 0],
         [1, 1, 1, 1],
         [2, 2, 2, 2]]]),
 array([[[0, 2, 4, 2],
         [0, 2, 4, 2],
         [0, 2, 4, 2]],
 
        [[0, 3, 3, 2],
         [0, 3, 3, 2],
         [0, 3, 3, 2]]])]

虽然我对自己想去的地方有了大致的了解，但我必须先尝试很多想法。

【讨论】：