Python use array axis to select slice

Question

I have a function that accepts an multi-dimensional array, an axis number and the index I would like to get:

def get_slice(my_array, dimension, index):
    if dimension == 0:
        slice = my_array[index, :, :]
    elif dimension == 1:
        slice = my_array[:, index, :]
    else:
        slice = my_array[:, :, index]
    return np.squeeze(slice)

But, I now have to change the code to accept 4 dimensions and I was wondering if there is a more general way to do this in python?

So, I am looking for a function that accepts an general n-dimension array, a dimension (axis) and an index to select on that dimension/axis and returns me the entire slice of that index on that dimension.

mgilson · Accepted Answer · 2014-03-27 15:42:16Z

1

Sure, it's not too difficult actually:

def get_slice(my_array, dimension, index):
    items = [slice(None, None, None)] * my_array.ndim
    items[dimension] = index
    array_slice = my_array[tuple(items)]
    return np.squeeze(array_slice)

Although I don't think it helps you here, if you ever have a function and you want to slice along the first dimension or the last dimension, you can use Ellipsis:

def get_slice_along_first_dim(array_with_arbitrary_dimensions, idx):
    return array_with_arbitrary_dimensions[idx, ...]

def get_slice_along_last_dim(array_with_arbitrary_dimensions, idx):
    return array_with_arbitrary_dimensions[..., idx]

You can even do stuff like:

arr[..., 3, :, 8]
arr[1, ..., 6]

if I recall correctly. the Ellipsis object just fills in as many empty slices as it needs to in order to fill out the dimensionality.

What you can't do is have more than 1 ellipsis object passed to __getitem__:

arr[..., 1, ...]  # Fail.

edited Mar 27, 2014 at 15:42

answered Mar 27, 2014 at 15:36

mgilson

312k70 gold badges656 silver badges722 bronze badges

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

1 Comment

Robbert Harms Over a year ago

This is great! Thank you for sharing. I use both your answer on top aswell as the function get_slice_along_last_dim() which I also needed.

Collectives™ on Stack Overflow

Python use array axis to select slice

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related