python - 按维度扩展 np.array()

Question

假设我有一个原始数组

w = np.ones([2, 4, 2, 3])

我正在操作子集[:, 2, :, 1]。在这个子集中，我有一个数组告诉我是否应该更新 w。

update = array([[False,  True], [False,  True]])

基本上，update[0, 1] == True 意味着 w[0, 2, 1, 1] 应该增加:第一个 (0) 维度对应于w中的0维，第二维对应于w中的2维。维度 1 和 3 固定为 (2, 1) (因为这是我正在操作的子集)。

现在，我想影响 update 隐含的维度上的 w。基本上，我正在寻找 index，这样在执行时

w[index] = 0

我设置为 0 个元素 [0, 2, 1, 1] 和 [1, 2, 1, 1]。

如果 w 只是我的 2x2 数组，我可以通过 w[update] = 0 来做到这一点，但我必须考虑考虑中间的附加尺寸。我该如何干净地做到这一点？

更复杂的设置

现在，我不再研究简单的子集 [:, 2, :, 1]，而是研究更复杂的东西。

我有一个名为 secondaryDim 的数组

secondDim = np.array([2, 3])

它包含作为第一维度的函数的第二维度。也就是说，当第一个维度等于 0 时，我在 2 上切片第二个维度。当第一个维度等于 1 时，我在 3 上切片第二个维度。

现在，我正在寻找索引，这样当我设置时

w[index] = 0

我设置为 0 个元素 [0, secondaryDim[0], 1, 1] 和 [1, secondaryDim[1], 1, 1] - 即 [0,2,1,1]和[1,3,1,1]。

Answer 1

第 1 部分

只需 w[:,2,:,1][update] = 0 即可。

第 2 部分

一种方法是存储来自掩码的行、列索引，并将其用于下标索引 -

r,c = np.where(update)
a[r,secondDim,c,1] = 0

示例运行 -

In [213]: a = np.random.randint(11,99,[2, 4, 2, 3])

In [222]: secondDim
Out[222]: array([2, 3])

In [216]: update   # For first and third dims
Out[216]: 
array([[False,  True],
       [False,  True]], dtype=bool)

In [214]: a[0, 2, 1, 1]
Out[214]: 85

In [215]: a[1, 3, 1, 1]
Out[215]: 47

In [217]: r,c = np.where(update)

In [218]: a[r,secondDim,c,1] # access those elems
Out[218]: array([85, 47])

In [219]: a[r,secondDim,c,1] = 0 # assign 0s

In [220]: a[0, 2, 1, 1] # verify
Out[220]: 0

In [221]: a[1, 3, 1, 1] # verify
Out[221]: 0

我们也可以在 part-1 中使用此方法 -

w[r,2,c,1] = 0