Function returning elements of a multi-dimensonal list in python

Question

I am trying to define a function that returns elements of a multi-dimensional variable according to a nested list of indices. I am working with variables depending on multiple dimensions like, e.g.: time, altitude, latitude, longitude.

Make it 3D for now (e.g., time, altitude and latitude):

x = np.arange(125).reshape(5,5,5)

If I now want the first 4 time steps, the first 3 altitude levels and the first 2 latitudes I can do several things:

x[[0,1,2,3],:,:][:,[0,1,2],:][:,:,[0,1]]

or

x[ [[[0]],[[1]],[[2]],[[3]]],[[0],[1],[2]],[0,1]]

or

x[np.ix_([0,1,2,3],[0,1,2],[0,1])]

But what I would like to have is a function giving me back the elements like

def get_elements( x, l )

where l is a list of indices

l = [[0,1,2,3],[0,1,2],[0,1]]

How could this function look like? The last alternative comes pretty close but x[np.ix_(l)] gives me an IndexError.

Moreover, it would be great to have the opportunity to leave dimensions untouched. E.g., using all time steps in this pseudo-code:

l = [[:],[0,1,2],[0,1]]

Thanks a lot!

hpaulj · Accepted Answer · 2014-09-10T15:38:51.717

1

Note the signature of np.ix_:

np.ix_(*args)

so you need to 'expand' l:

x[np.ix_(*l)]

Take a look at the code for ix_. Notice it does things like iterate over args, and returns a tuple. Also look at other functions in np.lib.index_tricks.

In general, if a list gives you problems when indexing, check whether you need to a tuple instead. Also familiarize yourself with slice.

edited Sep 10 '14 at 15:38

answered Sep 10 '14 at 15:30

hpaulj

221,503
14
230
353

Great! That was actually exactly what I was looking for. I just didn't know how to search for it. But the 'expand' did the trick. Look also here: [http://stackoverflow.com/questions/7745952/python-expand-list-to-function-arguments] – Lukas Sep 10 '14 at 15:39

score 1 · Answer 2 · edited May 23 '17 at 11:49

To directly answer your question about making a function, using your list as you're constructing it, though, you'll need to use your np.ix_ function:

def get_elements(x, l):
    return x[np.ix_(*l)]

However, I think you want this:

>>> x[:4,:3,:2]
array([[[ 0,  1],
        [ 5,  6],
        [10, 11]],

       [[25, 26],
        [30, 31],
        [35, 36]],

       [[50, 51],
        [55, 56],
        [60, 61]],

       [[75, 76],
        [80, 81],
        [85, 86]]])

It returns the same as this:

x[np.ix_([0,1,2,3],[0,1,2],[0,1])]

There's a great answer on slicing here: https://stackoverflow.com/a/24713353/541136

You can name your slice objects and pass them like this:

>>> first_four_slice = slice(None, 4)
>>> first_three_slice = slice(None, 3)
>>> first_two_slice = slice(None, 2)
>>> x[first_four_slice, first_three_slice, first_two_slice]

You can pack them all together in a single variable too:

>>> slice_tuple = first_four_slice, first_three_slice, first_two_slice
>>> x[slice_tuple]

And get your all "time steps" (i.e., get all from that dimension) like this:

>>> all_slice = slice(None)
>>> x[all_slice, first_three_slice, first_two_slice]

Function returning elements of a multi-dimensonal list in python

2 Answers2