CODEWITHSUNDEEP
pythonnumpyscipy
Bálint Aradi
Bálint Aradi

Reputation: 3812

Read elements with given indices from an array with arbitrary number of axis

I'd like to write a routine, which reads out the values of certain elements from an array. The element selection is specified as an array, where each row contains the indices for one element. The routine should work for arrays with arbitrary number of axes.

I could come up with the solution below, but I do not like it, as the conversion to tuple (or list) feels somehow unnecessary, although I need that to prevent advanced indexing. Is there any more numpythonic way to do this?

import numpy as np

def get_elements(aa, inds):
    myinds = tuple(inds.transpose())
    return aa[myinds]

AA = np.arange(6)
AA.shape = ( 3, 2 )
inds = np.array([[ 0, 0 ], [ 2, 1 ]])
data2 = get_elements(AA, inds)   # contains [ AA[0,0], A[2,1] ]

BB = np.arange(12)
BB.shape = ( 2, 3, 2 )
inds = np.array([[ 0, 0, 0], [ 1, 2, 1 ]])
data3 = get_elements(BB, inds)  # contains [ BB[0,0,0], BB[1,2,1] ]

Upvotes: 2

Views: 109

Answers (1)

hpaulj
hpaulj

Reputation: 231335

Your tuple(ind.T) produces the same thing as np.where for the same elements.

In [117]: AA=np.arange(6).reshape(3,2)
In [118]: ind=np.array([[0,0],[2,1]])
In [119]: tuple(ind.T)
Out[119]: (array([0, 2]), array([0, 1]))
In [120]: AA[tuple(ind.T)]
Out[120]: array([0, 5])

Using where to find the indices of these 2 values:

In [121]: np.where((AA==0) + (AA==5))
Out[121]: (array([0, 2]), array([0, 1]))

And copying from the doc for where, another way of finding the [0,5] values:

In [125]: np.in1d(AA.ravel(), [0,5]).reshape(AA.shape)
Out[125]: 
array([[ True, False],
       [False, False],
       [False,  True]], dtype=bool)
In [126]: np.where(np.in1d(AA.ravel(), [0,5]).reshape(AA.shape))
Out[126]: (array([0, 2]), array([0, 1]))

So your tuple is perfectly good numpy code.

From numpy indexing documentation:

Note In Python, x[(exp1, exp2, ..., expN)] is equivalent to x[exp1, exp2, ..., expN]; the latter is just syntactic sugar for the former.

Upvotes: 1

Related Questions