Python Numpy将三维数组索引为二维数组
我有一个以下结构的三维数组:Python Numpy将三维数组索引为二维数组,python,arrays,numpy,multidimensional-array,Python,Arrays,Numpy,Multidimensional Array,我有一个以下结构的三维数组: x = np.array([[[1,2], [3,4]], [[5,6], [7,8]]], dtype=np.double) 此外,我还有一个索引数组 idx = np.array([[0,1],[1,3]], dtype=np.int) idx的每一行定义了行/列索引,用于沿x中的0轴将每个子数组放置到初始化为 K = np.zeros((4,4), dtype=np.d
x = np.array([[[1,2],
[3,4]],
[[5,6],
[7,8]]], dtype=np.double)
此外,我还有一个索引数组
idx = np.array([[0,1],[1,3]], dtype=np.int)
idx
的每一行定义了行/列索引,用于沿x
中的0
轴将每个子数组放置到初始化为
K = np.zeros((4,4), dtype=np.double)
我想使用奇特的索引/广播来执行索引,而不使用for
循环。我目前是这样做的:
for i, id in enumerate(idx):
idx_grid = np.ix_(id,id)
K[idx_grid] += x[i]
结果是:
>>> K = array([[ 1., 2., 0., 0.],
[ 3., 9., 0., 6.],
[ 0., 0., 0., 0.],
[ 0., 7., 0., 8.]])
这可能与花哨的索引有关吗?我认为这是不可能的,因为循环中有
+=
。这意味着,您必须将数组idx
放大一个维度,然后利用np.sum(x[…],axis=…)
再次缩小它。
一个次要的优化是:
import numpy as np
xx = np.array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]], dtype=np.double)
idx = np.array([[0, 1], [1, 3]], dtype=np.int)
K0, K1 = np.zeros((4, 4), dtype=np.double), np.zeros((4, 4), dtype=np.double)
for k, i in enumerate(idx):
idx_grid = np.ix_(i, i)
K0[idx_grid] += xx[k]
for x, i in zip(xx, idx):
K1[np.ix_(i, i)] += x
print("K1 == K0:", np.allclose(K1, K0)) # prints: K1 == K0: True
注意:不要使用
id
作为变量名,因为它是Python关键字。这里有一种替代方法。在您的问题中定义了x
、idx
和K
:
indices = (idx[:,None] + K.shape[1]*idx).ravel('f')
np.add.at(K.ravel(), indices, x.ravel())
然后我们有:
>>> K
array([[ 1., 2., 0., 0.],
[ 3., 9., 0., 6.],
[ 0., 0., 0., 0.],
[ 0., 7., 0., 8.]])
要在NumPy阵列上执行无缓冲就地添加,您需要使用(避免在
for
循环中使用+=
)
但是,将2D索引数组的列表以及要在这些索引处添加的相应数组传递到np.add.at
有点困难。这是因为函数将这些数组列表解释为高维数组,并引发索引器
在一维数组中传递要简单得多。您可以临时拆分
K
和x
以获得一维零数组和添加到这些零的一维值数组。唯一棘手的部分是从idx
构建相应的1D索引数组,在该数组中添加值。如上图所示,这可以通过使用算术运算符进行广播,然后进行散播来完成。预期的操作是将x
中的值累加到idx
索引的位置。您可以将这些idx
位置视为直方图数据的bin
,将x
值视为需要为这些bin累积的权重。现在,要执行这样的装箱操作,可以使用。这里有一个这样的实现-
# Get size info of expected output
N = idx.max()+1
# Extend idx to cover two axes, equivalent to `np.ix_`
idx1 = idx[:,None,:] + N*idx[:,:,None]
# "Accumulate" values from x into places indexed by idx1
K = np.bincount(idx1.ravel(),x.ravel()).reshape(N,N)
运行时测试- 1) 创建输入:
In [361]: # Create x and idx, with idx having unique elements in each row of idx,
...: # as otherwise the intended operation is not clear
...:
...: nrows = 100
...: max_idx = 100
...: ncols_idx = 2
...:
...: x = np.random.rand(nrows,ncols_idx,ncols_idx)
...: idx = np.random.randint(0,max_idx,(nrows,ncols_idx))
...:
...: valid_mask = ~np.any(np.diff(np.sort(idx,axis=1),axis=1)==0,axis=1)
...:
...: x = x[valid_mask]
...: idx = idx[valid_mask]
...:
2) 定义功能:
In [362]: # Define the original and proposed (bincount based) approaches
...:
...: def org_approach(x,idx):
...: N = idx.max()+1
...: K = np.zeros((N,N), dtype=np.double)
...: for i, id in enumerate(idx):
...: idx_grid = np.ix_(id,id)
...: K[idx_grid] += x[i]
...: return K
...:
...:
...: def bincount_approach(x,idx):
...: N = idx.max()+1
...: idx1 = idx[:,None,:] + N*idx[:,:,None]
...: return np.bincount(idx1.ravel(),x.ravel()).reshape(N,N)
...:
3) 最后,给他们计时:
In [363]: %timeit org_approach(x,idx)
100 loops, best of 3: 2.13 ms per loop
In [364]: %timeit bincount_approach(x,idx)
10000 loops, best of 3: 32 µs per loop