Store Numpy array index in variable

Question

I want to pass an index slice as an argument to a function:

def myfunction(some_object_from_which_an_array_will_be_made, my_index=[1:5:2,::3]):
    my_array = whatever(some_object_from_which_an_array_will_be_made)
    return my_array[my_index]

Obviously this will not work, and obviously in this particular case there might be other ways to do this, but supposing I really want to do stuff this way, how can I use a variable to slice a numpy array?

hpaulj · Accepted Answer · 2015-04-06T19:56:00.103

np.lib.index_tricks has a number of functions (and classes) that can streamline indexing. np.s_ is one such function. It is actually an instance of a class that has a __get_item__ method, so it uses the [] notation that you want.

An illustration of its use:

In [249]: np.s_[1:5:2,::3]
Out[249]: (slice(1, 5, 2), slice(None, None, 3))

In [250]: np.arange(2*10*4).reshape(2,10,4)[_]
Out[250]: 
array([[[40, 41, 42, 43],
        [52, 53, 54, 55],
        [64, 65, 66, 67],
        [76, 77, 78, 79]]])

In [251]: np.arange(2*10*4).reshape(2,10,4)[1:5:2,::3]
Out[251]: 
array([[[40, 41, 42, 43],
        [52, 53, 54, 55],
        [64, 65, 66, 67],
        [76, 77, 78, 79]]])

Notice that it constructs the same tuple of slices that ajcr did. _ is the temporary variable that IPython uses for the last result.

To pass such a tuple to a function, try:

def myfunction(some_object_from_which_an_array_will_be_made, my_index=np.s_[:,:]):
    my_array = whatever(some_object_from_which_an_array_will_be_made)
    return my_array[my_index]
I = np.s_[1:5:2,::3]
myfunction(obj, my_index=I)

This sounds very interesting - but how could I use it for my example? I suppose `np.s_var` will not work? And even then, declaring var = [2:5] will tell me it's invalid syntax. — TheChymera, Apr 06 '15 at 19:07
`var=np.s[2:5]; foo(obj, my_index=var)`. See my added example. — hpaulj, Apr 06 '15 at 19:45
Amazing! But I believe you meant `var=np.s_[2:5]` in your last comment. Also the way it ended up working was not `my_array[_]` but rather `my_array[var]` — TheChymera, Apr 06 '15 at 22:37

score 4 · Answer 2 · answered Apr 06 '15 at 17:42

One way is to build a slice object (or a tuple of slice objects) and pass it in to the function to use as the index.

For example, the index notation

my_array[1:5:2, ::3]

is equivalent to

my_array[slice(1,5,2), slice(None,None,3)]

So your function could become:

def myfunction(some_object, my_index=(slice(1,5,2), slice(None,None,3))):
    my_array = whatever(some_object)
    return my_array[my_index]

Store Numpy array index in variable

2 Answers2

Linked