使用边界检查索引numpy数组

时间:2017-11-05 13:07:13

标签: python arrays numpy indexing numpy-broadcasting

假设我有一些随机的numpy数组:

>>> x = np.random.randn(10, 4, 4)

因此可以将其视为10个4x4阵列。现在,我也有一些关于这些数组的坐标:

>>> coords.shape
(10, 4, 4, 2)

此数组的特定元素可能如下所示:

>>> coords[i][j][k]
array([ 2,  2])

现在,我想创建一个大小为y的{​​{1}}数组,(10, 4, 4)如果y[i][j][k] = x[coords[i][j][k]]在界限内,则为coords[i][j][k]

我尝试过以下操作:

0

但我不能完全正确地解决广播规则,而且我遇到了错误。

例如,假设我们有,

>>> y = np.where(np.logical_and(
                 np.all(coords >= 0, axis = 3), 
                 np.all(coords <= 3, axis = 3)), 
                 x[tuple(coords)], 0)

用,

>>> x = np.array([
                 [[1., 2.],
                  [4., 5.]],
                 [[6., 7.],
                  [8., 9.]]])

然后,我想以某种方式得到

>>> coords = np.array([
                       [[[0,-1], [0, 0]],
                        [[1, 1], [1, 0]]],
                       [[[1, 1], [1, 0]],
                        [[7, 2], [0, 0]]]
                      ])

我如何在Python中实现这一目标?

1 个答案:

答案 0 :(得分:1)

似乎需要一些adavnced-indexing工作,并且应该适用于通用n-dim案例 -

m,n = x.shape[0], x.shape[-1]
mask = (coords < n).all((-1)) & (coords >=0).all((-1))
v_coords = np.where(mask[:,:,:,None],coords,0)
out = x[np.arange(m)[:,None,None], v_coords[...,0], v_coords[...,1]]
out[~mask] = 0