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

Question

我有一段numpy代码我知道可以工作。我知道这一点，因为我已经在我的通用案例中成功地进行了测试。然而，经过两个小时的反复参考文档和反复试验，我终于找到了解决方案。我无法理解我是如何凭直觉知道要做这个

设置:

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

目标:设置为1的指标0,1的轴1,0,1的轴2，所有轴3和指标3,4的轴1,3,4的轴2，所有轴3

更清晰的目标:设置block 1和block 2的前两行为1,block 3和block 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.]])