给定索引列表，将任意维度的 numpy 切片到一维数组

Question

I have a numpy ndarray arr and also indices , a list of indices specifying a particular entry. For concreteness let's take:

arr = np.arange(2*3*4).reshape((2,3,4))
indices= [1,0,3]

I have code to take 1d slices through arr observing all but one index n :

arr[:, indices[1], indices[2]]  # n = 0
arr[indices[0], :, indices[2]]  # n = 1
arr[indices[0], indices[1], :]  # n = 2

I would like to change my code to loop over n and support an arr of arbitrary dimension.

I've had a look at the indexing routines entry in the documentation and found information about slice() and np.s_() . I was able to hack together something that works like what I want:

def make_custom_slice(n, indices):
    s = list()
    for i, idx in enumerate(indices):
        if i == n:
            s.append(slice(None))
        else:
            s.append(slice(idx, idx+1))
    return tuple(s)


for n in range(arr.ndim):
    np.squeeze(arr[make_custom_slice(n, indices)])

Where np.squeeze is used to remove the axes of length 1. Without this, the array this produced has shape (arr.shape[n],1,1,...) rather than (arr.shape[n],) .

Is there a more idiomatic way to accomplish this task?

Answer 1

Some improvements to the solution above (there may still be a one-liner or a more performant solution):

def make_custom_slice(n, indices):
    s = indices.copy()
    s[n] = slice(None)
    return tuple(s)


for n in range(arr.ndim):
    print(arr[make_custom_slice(n, indices)])

An integer value idx can be used to replace the slice object slice(idx, idx+1) . Because most indices are copied over directly, start with a copy of indices rather than building the list from scratch.

When built in this way, the result of arr[make_custom_slice(n, indices) has the expected dimension and np.squeeze is unnecessary.