Python - how can I address an array along a given axis?

Question

I have a python script that is reading slices from a 3D array, like this:

def get_from_array(axis, start, end, array):
    if axis == 0:
        slice = array[start:end, :, :]
    elif axis == 1:
        slice = array[:, start:end, :]
    elif axis == 2:
        slice = array[:, :, start:end]
    return slice

I can't help but think there must be a better way of doing this! Any suggestions?

S

Answer 1

You can do something like this:

idx = [slice(None)] * array.ndim
idx[axis] = slice(start, end)
myslice = array[tuple(idx)]

Answer 2

You can also use np.take. Then you can do it in one line more naturally.

a.take(np.arange(start,end), axis=axis)

Notes:

• it won't accept the standard : slice notation but replace that with range
• It returns a copy, not a view

For example:

In [135]: a = np.arange(3*3*3).reshape(3,3,3)

In [136]: a.take(np.arange(1,2), axis=1)
Out[136]: 
array([[[ 3,  4,  5]],

       [[12, 13, 14]],

       [[21, 22, 23]]])