slicing numpy array along an arbitrary dimension

Question

say that I have a (40,20,30) numpy array and that I have a function that after some work will return half of the input array along a selected input axis. Is there an automatic way to do so ? I would like to avoid such an ugly code:

def my_function(array,axis=0):

    ...

    if axis == 0:
        return array[:array.shape[0]/2,:,:] --> (20,20,30) array
    elif axis = 1:
        return array[:,:array.shape[1]/2,:] --> (40,10,30) array
    elif axis = 2: 
        return array[:,:,:array.shape[2]/2] --> (40,20,15) array

thanks for your help

Eric

DSM · Accepted Answer · 2013-05-24 14:56:42Z

8

I think you can use np.split for this [docs], and simply take the first or second element returned, depending on which one you want. For example:

>>> a = np.random.random((40,20,30))
>>> np.split(a, 2, axis=0)[0].shape
(20, 20, 30)
>>> np.split(a, 2, axis=1)[0].shape
(40, 10, 30)
>>> np.split(a, 2, axis=2)[0].shape
(40, 20, 15)
>>> (np.split(a, 2, axis=0)[0] == a[:a.shape[0]/2, :,:]).all()
True

answered May 24, 2013 at 14:56

DSM

355k67 gold badges606 silver badges504 bronze badges

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

1 Comment

mhsmith Over a year ago

FYI: split() also takes a tuple specifying arbitrary split points.

Eurydice · Accepted Answer · 2013-05-24 15:19:33Z

4

thanks for your help, DSM. I will use your approach.

In the meantime, I found a (dirty ?) hack:

>>> a = np.random.random((40,20,30))
>>> s = [slice(None),]*a.ndim
>>> s[axis] = slice(f,l,s)
>>> a1 = a[s]

Perhaps a bit more general than np.split but much less elegant !

answered May 24, 2013 at 15:19

Eurydice

8,5095 gold badges29 silver badges44 bronze badges

1 Comment

Mad Physicist Over a year ago

This is not less elegant. This just removes some of the syntactic sugar: : is slice(None), a:b:c is slice(a, b, c), etc.

Bi Rico · Accepted Answer · 2013-05-24 18:38:25Z

2

numpy.rollaxis is a good tool for this:

def my_func(array, axis=0):
    array = np.rollaxis(array, axis)
    out = array[:array.shape[0] // 2]
    # Do stuff with array and out knowing that the axis of interest is now 0
    ...

    # If you need to restore the order of the axes
    if axis == -1:
        axis = out.shape[0] - 1
    out = np.rollaxis(out, 0, axis + 1)

edited May 24, 2013 at 18:38

answered May 24, 2013 at 17:09

Bi Rico

25.9k3 gold badges57 silver badges75 bronze badges

Collectives™ on Stack Overflow

slicing numpy array along an arbitrary dimension

3 Answers 3

1 Comment

1 Comment

Comments

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

1 Comment

Comments

Related