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

b = a[(slice(None),) * k + (i,)]

Construct the indexing tuple manually.

As documented in the Python language reference, an expression of the form

a[:, :, :, :, :, i]

is converted to

a[(slice(None), slice(None), slice(None), slice(None), slice(None), i)]

We can achieve the same effect by building that tuple directly instead of using slicing notation. (There's the minor caveat that building the tuple directly produces a[(i,)] instead of a[i] for k=0, but NumPy handles these the same for scalar i.)

Answer 2

I am not sure if it will work for k dim but it does for 2 dim

a.take(i,axis=k)

Answer 3

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,)]

Answer 4

here is a late entry that can handle negative axis arguments without having to know the shape of its operand beforehand:

def put_at(inds, axis=-1, slc=(slice(None),)):
    return (axis<0)*(Ellipsis,) + axis*slc + (inds,) + (-1-axis)*slc

To be used as in

a[put_at(ind_list,axis=axis)]

ind_list can be a scalar as in your case or something more interesting as well.

Copied from this comment of mine.

Answer 5

I'm not sure if this approach would create an entire copy of the array*, but I would take a slice of the transposed matrix in order to get the kth axis:

import numpy as np

def get_slice(arr, k, i):
    if k >= arr.ndim: #we need at least k dimensions (0 indexed)
        raise ValueError("arr is of lower dimension than {}".format(k))

    axes_reorder = list(range(arr.ndim)) #order of axes for transpose
    axes_reorder.remove(k) #remove original position of k
    axes_reorder.insert(0,k) #insert k at beginning of order

    return arr.transpose(axes_reorder)[i] #k is first axis now

This also has the added bonus of easier checking of the number of dimensions before trying the slice.

* according to the docs, a memory view is created whenever possible.