如何解释 numpy 高级索引解决方案

Question

我有一段 我知道可以工作的 numpy 代码。我知道这一点，因为我已经在我的通用案例中成功地对其进行了测试。但是，经过两个小时来回参考文档和反复试验后，我得出了解决方案。我无法理解我怎么会知道如何直观地。

设置：

a = np.zeros((5,5,3))

目标：设置为 1 轴 1 的索引 0,1、轴 2 的 0,1、轴 3 的所有索引和轴 1 的索引 3,4、轴 2 的索引 3,4、轴 3 的所有索引

更明确的目标：将块 1 和 2 的前两行设置为 1，将块 3 和 4 的最后两行设置为 1

结果：

ax1 =np.array([np.array([0,1]),np.array([3,4])])
ax1 =np.array([x[:,np.newaxis] for x in ax1])
ax2 = np.array([[[0,1]],[[3,4]]])
a[ax1,ax2,:] = 1
a

输出：

array([[[1., 1., 1.],
    [1., 1., 1.],
    [0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.]],

   [[1., 1., 1.],
    [1., 1., 1.],
    [0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.]],

   [[0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.]],

   [[0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.],
    [1., 1., 1.],
    [1., 1., 1.]],

   [[0., 0., 0.],
    [0., 0., 0.],
    [0., 0., 0.],
    [1., 1., 1.],
    [1., 1., 1.]]])

我倾向于相信我应该能够通过查看相关矩阵的形状、索引的形状和索引操作来直观地了解输出。但是，我无法将这个故事放在我的脑海中。比如，它正在改变的子空间的最终形状是什么？ 您如何解释其工作原理？

形状：

input: (5, 5, 3)
ind1: (2, 2, 1)
ind2: (2, 1, 2)
final_op: input[ind1, ind2, :]

Answer 1

有形状

ind1: (2, 2, 1)
ind2: (2, 1, 2)

他们broadcast一起选择一个(2,2,2)空间

In [4]: ax1
Out[4]: 
array([[[0],
        [1]],

       [[3],
        [4]]])
In [5]: ax2
Out[5]: 
array([[[0, 1]],

       [[3, 4]]])

所以对于第一个维度（块），它选择块 0、1、3 和 4。在第二个维度中，它也选择这些行。

这就是前 2 个块的前 2 行和最后 2 个块的最后 2 行。这就是结果中出现 1 的位置。

创建索引数组的更简单方法：

In [7]: np.array([[0,1],[3,4]])[:,:,None]   # (2,2) expanded to (2,2,1)
In [8]: np.array([[0,1],[3,4]])[:,None,:]   # expand to (2,1,2)

这就是广播扩展它们的方式：

In [10]: np.broadcast_arrays(ax1,ax2)
Out[10]: 
[array([[[0, 0],              # block indices
         [1, 1]],
 
        [[3, 3],
         [4, 4]]]),
 array([[[0, 1],              # row indices
         [0, 1]],
 
        [[3, 4],
         [3, 4]]])]

这可能会使模式更清晰：

In [15]: a[ax1,ax2,:] = np.arange(1,5).reshape(2,2,1)
In [16]: a[:,:,0]
Out[16]: 
array([[1., 2., 0., 0., 0.],
       [3., 4., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 1., 2.],
       [0., 0., 0., 3., 4.]])