Simplify numpy expression

Question

How can I simplify this:

import numpy as np
ex = np.arange(27).reshape(3, 3, 3)

def get_plane(axe, index):
    return ex.swapaxes(axe, 0)[index]  # is there a better way ?

I cannot find a numpy function to get a plane in a higher dimensional array, is there one?

EDIT

The ex.take(index, axis=axe) method is great, but it copies the array instead of giving a view, what I originally wanted.

So what is the shortest way to index (without copying) a n-th dimensional array to get a 2d slice of it, with index and axis?

Think you are after `np.take` : `np.take(ex, index, axis=axe)`. — Divakar, Jun 05 '20 at 19:13

score 2 · Accepted Answer · answered Jun 05 '20 at 18:49

2

Inspired by this answer, you can do something like this:

def get_plane(axe, index):
    slices = [slice(None)]*len(ex.shape)
    slices[axe]=index
    return ex[tuple(slices)]

get_plane(1,1)

output:

array([[ 3,  4,  5],
       [12, 13, 14],
       [21, 22, 23]])

answered Jun 05 '20 at 18:49

Quang Hoang

146,074
10
56
74

great answer: it is really understandable, and it gives a view of the array instead of copying, like `np.take` – rambi Jun 05 '20 at 21:09

score 1 · Answer 2 · answered Jun 05 '20 at 19:17

What do you mean by a 'plane'?

In [16]: ex = np.arange(27).reshape(3, 3, 3)

Names like plane, row, and column, are arbitrary conventions, not formally defined in numpy. The default display of this array looks like 3 'planes' or 'blocks', each with rows and columns:

In [17]: ex                                                                     
Out[17]: 
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])

Standard indexing lets us view any 2d block, in any dimension:

In [18]: ex[0]                                                                  
Out[18]: 
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
In [19]: ex[0,:,:]                                                              
Out[19]: 
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
In [20]: ex[:,0,:]                                                              
Out[20]: 
array([[ 0,  1,  2],
       [ 9, 10, 11],
       [18, 19, 20]])
In [21]: ex[:,:,0]                                                              
Out[21]: 
array([[ 0,  3,  6],
       [ 9, 12, 15],
       [18, 21, 24]])

There are ways of saying I want block 0 in dimension 1 etc, but first make sure you understand this indexing. This is the core numpy functionality.

In [23]: np.take(ex, 0, 1)                                                      
Out[23]: 
array([[ 0,  1,  2],
       [ 9, 10, 11],
       [18, 19, 20]])

In [24]: idx = (slice(None), 0, slice(None))     # also np.s_[:,0,:]                                
In [25]: ex[idx]                                                                
Out[25]: 
array([[ 0,  1,  2],
       [ 9, 10, 11],
       [18, 19, 20]])

And yes you can swap axes (or transpose), it that suits your needs.

Simplify numpy expression

EDIT

2 Answers2