Clever Geek Handbook

Indexing tensor with binary matrix in numpy

I have a tensor A such that A.shape = (32, 19, 2) and a binary matrix B such that B.shape = (32, 19). Is there a one-line operation I can do to get the matrix C where C.shape = (32, 19) and C (i, j) = A [i, j, B [i, j]]?

Basically, I want to use B as an indexing matrix, where if B [i, j] = 1, then take A [i, j, 1] to form C (i, j).

+3


source to share


3 answers


np.where

for help. The same principle as mtrw's

answers:

In [344]: A=np.arange(4*3*2).reshape(4,3,2)

In [345]: B=np.zeros((4,3),dtype=int)

In [346]: B[[0,1,1,2,3],[0,0,1,2,2]]=1

In [347]: B
Out[347]: 
array([[1, 0, 0],
       [1, 1, 0],
       [0, 0, 1],
       [0, 0, 1]])

In [348]: np.where(B,A[:,:,1],A[:,:,0])
Out[348]: 
array([[ 1,  2,  4],
       [ 7,  9, 10],
       [12, 14, 17],
       [18, 20, 23]])

np.choose

can be used if the last size is greater than 2 (but less than 32). ( choose

works in list or first dimension, therefore rollaxis

.

In [360]: np.choose(B,np.rollaxis(A,2))
Out[360]: 
array([[ 1,  2,  4],
       [ 7,  9, 10],
       [12, 14, 17],
       [18, 20, 23]])

B

can also be used directly as an index. The trick is to specify the other dimensions in such a way that they translate into the same form.



In [373]: A[np.arange(A.shape[0])[:,None], np.arange(A.shape[1])[None,:], B]
Out[373]: 
array([[ 1,  2,  4],
       [ 7,  9, 10],
       [12, 14, 17],
       [18, 20, 23]])

This latter approach can be modified to work if it B

does not fit the 1st 2 dimensions A

.

np.ix_

can make this indexing easier

I, J = np.ix_(np.arange(4),np.arange(3))
A[I, J, B]
+2


source


You can do this using a list comprehension:



C = np.array([[A[i, j, B[i, j]] for j in range(A.shape[1])] for i in range(A.shape[0])])
0


source


C = A[:,:,0]*(B==0) + A[:,:,1]*(B==1)

must work. You can generalize this as np.sum([A[:,:,k]*(B==k) for k in np.arange(A.shape[-1])], axis=0)

if you need to index more planes.

0


source






More articles:


All Articles