Numpy: Combine list of arrays by another array (np.choose alternative)

Question

I have a list of numpy arrays, each of the same shape. Let's say:

a = [np.array([[1, 2, 3],
               [4, 5, 6],
               [7, 8, 9]]),
     np.array([[11, 12, 13],
               [14, 15, 16],
               [17, 18, 19]]),
     np.array([[99, 98, 97],
               [96, 95, 94],
               [93, 92, 91]])]

And I have another array of the same shape that gives the list indices I want to take the elements from:

b = np.array([[0, 0, 1],
              [2, 1, 0],
              [2, 1, 2]])

What I want to get is the following:

np.array([[1, 2, 13],
          [96, 15, 6],
          [93, 18, 91]])

There was a simple solution that worked fine:

np.choose(b, a)

But this is limited to 32 arrays at most. But in my case, I have to combine more arrays (more than 100). So I need another way to do so.

I guess, it has to be something about advances indexing or maybe the np.take method. So probably, the first step is a = np.array(a) and then something like a[np.arange(a.shape[0]), b]. But I do not get it working.

Can somebody help? :)

score 3 · Answer 1 · answered Mar 23 '19 at 01:51

3

You can try using np.ogrid. Based on this answer. Of course you will have to convert a to a NumPy array first

i, j = np.ogrid[0:3, 0:3]
print (a[b, i, j])

# array([[ 1,  2, 13],
#        [96, 15,  6],
#        [93, 18, 91]])

answered Mar 23 '19 at 01:51

Sheldore

37,862
7
57
71

hpaulj · Answer 2 · 2019-03-23T02:21:05.320

In [129]: a = [np.array([[1, 2, 3], 
     ...:                [4, 5, 6], 
     ...:                [7, 8, 9]]), 
     ...:      np.array([[11, 12, 13], 
     ...:                [14, 15, 16], 
     ...:                [17, 18, 19]]), 
     ...:      np.array([[99, 98, 97], 
     ...:                [96, 95, 94], 
     ...:                [93, 92, 91]])]                                        
In [130]: b = np.array([[0, 0, 1], 
     ...:               [2, 1, 0], 
     ...:               [2, 1, 2]])                                             
In [131]:                                                                       
In [131]: A = np.array(a)                                                       
In [132]: A.shape                                                               
Out[132]: (3, 3, 3)

You want to use b to index the first dimension. For the other dimensions you need a indices that broadcast with b, i.e. a column vector and a row vector:

In [133]: A[b, np.arange(3)[:,None], np.arange(3)]                              
Out[133]: 
array([[ 1,  2, 13],
       [96, 15,  6],
       [93, 18, 91]])

there are various convenience functions for creating these arrays, e.g.

In [134]: np.ix_(range(3),range(3))                                             
Out[134]: 
(array([[0],
        [1],
        [2]]), array([[0, 1, 2]]))

and ogrid as mentioned in the other answer.

Here's a relatively new function that also does the job:

In [138]: np.take_along_axis(A, b[None,:,:], axis=0)                              
Out[138]: 
array([[[ 1,  2, 13],
        [96, 15,  6],
        [93, 18, 91]]])

I had to think a bit before I got the adjustment to b right.

I had to think a lot as well. For trying and for understand these answers :D The new function `np.take_along_axis` is in my favour because it looks to be the most readable one. — Max16hr, Mar 23 '19 at 02:35

Numpy: Combine list of arrays by another array (np.choose alternative)

2 Answers2