Python使用数组轴来选择切片

Question

我有一个函数，它接受多维数组、轴编号和我想要获取的索引:

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)

但是，我现在必须更改代码以接受 4 个维度，我想知道是否有更通用的方法可以在 python 中执行此操作？

因此，我正在寻找一个函数，它接受一般的 n 维数组、维度(轴)和要在该维度/轴上选择的索引，并返回该维度上该索引的整个切片。

Answer 1

当然，实际上并不太难:

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)

<小时/>

尽管我认为这对您没有帮助，但如果您有一个函数并且想要沿第一个维度或最后一个维度进行切片，您可以使用省略号:

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]

你甚至可以做这样的事情:

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

如果我没记错的话。 Ellipsis 对象只是根据需要填充尽可能多的空切片，以填充维度。

你不能做的就是将超过 1 个省略号对象传递给 __getitem__:

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