numpy將3維數組索引為2維數組

Question

我有以下結構的三維數組：

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 ，該數組初始化為

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.]])

這可能與花式索引有關嗎？

Answer 1

這是另一種方法。 使用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數組上執行無緩沖的就地加法，您需要使用np.add.at （以避免在for循環中使用+= ）。

但是，將2D索引數組的列表以及要在這些索引處添加的對應數組傳遞給np.add.at 。 這是因為該函數將這些數組列表解釋為高維數組，並且會引發IndexErrors。

傳遞一維數組要簡單得多。 您可以暫時拉開K和x以得到一維零的數組和一維值數組以添加到這些零。 唯一的麻煩是從idx構造一個相應的索引一維數組，在該數組中添加值。 如上所示，這可以通過與算術運算符進行廣播然后進行破壞來完成。

Answer 2

預期的操作是x到由idx索引的位置中的值的accumulation之一。 您可以將這些idx位置視為直方圖數據的bins ，將x值視為需要為這些bin累加的權重。 現在，要執行這種裝箱操作，可以使用np.bincount 。 這是一個這樣的實現-

# 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

Answer 3

我認為這不可能有效，因為循環中有+= 。 這意味着，您將不得不將數組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

PS：請勿將id用作變量名，因為它是Python關鍵字。