Selecting axis form multidimensional arrays with an array

Question

I am trying to select a subset of a multidimensional array using another array, so for example, if I have:

a=np.linspace(1,30,30)
a=a.reshape(5,3,2)

I would like to take the subset [:,0,1], which I can do by saying

a_subset=a[:,0,1]

but, is there any way to define an array/list specifying that subset and then subtract it? The idea is to do something like:

b=[:,0,1]
a_subset=a[b]

which does not work as ":" is not accepted as item ("SyntaxError: invalid syntax")

Sam van Kampen · Accepted Answer · 2016-05-12 10:21:15Z

4

You can do this using numpy.index_exp (docs) as follows:

import numpy as np

a = np.linspace(1, 30, 30)
a = a.reshape(5, 3, 2)

b = np.index_exp[:,0,1]
a_subset = a[b]

answered May 12, 2016 at 10:21

Sam van Kampen

1,02310 silver badges17 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Equivalent expressions are np.s_[:,0,1] and (slice(None), 0, 1)