Numpy：访问具有多个索引列表的子ndarrays

Question

我正在尝试使用多个索引列表访问（更准确地说：添加到）numpy 数组的子集。 然而，与切片完美配合的内容，在使用任意索引列表时却无法按预期工作：

import numpy as np

# this one works:

A_4th_order = np.zeros( (4,2,4,2) )

sub_a = np.ones( (3,1,3,1) )

i = k = slice(0,3)
j = l = slice(1,2)

A_4th_order[i,j,k,l] += sub_a


# this one doesn't:

B_4th_order = np.zeros( (4,2,4,2) )

sub_b = np.ones( (3,1,3,1) )

i = k = np.array([0,1,2], dtype=int)
j = l = np.array([1], dtype=int)

B_4th_order[i,j,k,l] += sub_b

如何以高效且可读的方式实现此操作？

Answer 1

如果您在执行就地更新之前打印形状。您会看到以下内容：

>>> print(A_4th_order[i, j, k, l].shape)
(3, 1, 3, 1)

并且：

>>> print(B_4th_order[i, j, k, l].shape)
(3,)

第一种情况下子数组是 4d，而第二种情况下是 1d。这是因为您的坐标数组

i, j, k, l

是一维数组。一旦您将数组转换为开放网格（例如使用

np.ix_

），它就会按预期工作：

import numpy as np

# this one works:

A_4th_order = np.zeros((4, 2, 4, 2))

sub_a = np.ones((3, 1, 3, 1))

i = k = slice(0, 3)
j = l = slice(1, 2)

print(A_4th_order[i, j, k, l].shape)
A_4th_order[i, j, k, l] += sub_a

# this one doesn't:

B_4th_order = np.zeros((4, 2, 4, 2))

sub_b = np.ones((3, 1, 3, 1))

i = k = np.array([0, 1, 2], dtype=int)
j = l = np.array([1,], dtype=int)

# transform into an open meshgrid
i, j, k, l = np.ix_(i, j, k, l)

print(B_4th_order[i, j, k, l].shape)
B_4th_order[i, j, k, l] += sub_b

np.allclose(A_4th_order, B_4th_order)

应用

np.ix_

后，

i, j, k, l

坐标数组具有与其各自维度相对应的形状，它们也适用：

>>> print(i.shape)
>>> print(j.shape)
>>> print(k.shape)
>>> print(l.shape)
(3, 1, 1, 1)
(1, 1, 1, 1)
(1, 1, 3, 1)
(1, 1, 1, 1)

我希望这有帮助！