get the i-th slice of the k-th dimension in a numpy array

Question

I have an n-dimensional numpy array, and I\'d like to get the i-th slice of the k-th dimension. There must be something better than

Answer 1

Basically, you want to be able to programmatically create the tuple :, :, :, :, :, i, ... in order to pass it in as the index of a. Unfortunately, you cannot simply use ordinary tuple multiplication on the colon operator directly (i.e., (:,) * k won't work to generate a tuple of k colon operators). You can, however, get an instance of a "colon slice" by using colon = slice(None). You could then do b = a[(colon,) * k + (i,)], which would effectively index a at the ith column of the kth dimension.

Wrapping this up in a function, you'd get:

def nDimSlice(a, k, i):
    colon = slice(None)
    return a[(colon,) * k + (i,)]