CODEWITHSUNDEEP
pythonarraysnumpyindexing
Robbert Harms
Robbert Harms

Reputation: 13

Python use array axis to select slice

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.

Upvotes: 1

Views: 377

Answers (1)

mgilson
mgilson

Reputation: 309831

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.

Upvotes: 1

Related Questions