Slicing a numpy array along a dynamically specified axis

Question

I would like to dynamically slice a numpy array along a specific axis. Given this:

axis = 2
start = 5
end = 10

I want to achieve the same result as this:

# m is some matrix
m[:,:,5:10]

Using something like this:

slc = tuple(:,) * len(m.shape)
slc[axis] = slice(start,end)
m[slc]

But the : values can't be put in a tuple, so I can't figure out how to build the slice.

Answer 1

I saw no comments which mentioned the Ellipsis object and so wanted to include another, just slightly different solution, which works for positive and negative axes. This solution has a function which generates the slice along the desired dimension.

The main reason I include this is to make sure that someone who learns about dynamic slicing also learns about using ellipsis.

def _slice_along_axis(slice_inds,axis=-1):
    '''
    Returns a slice such that the 1d slice provided by slice_inds, slices along the dimension provided.
    '''
    from_end=False
    if axis<0: # choosing axis at the end
        from_end = True
        axis = -1-axis
    explicit_inds_slice = axis*(slice(None),) 
    if from_end:
        return (Ellipsis,slice_inds) + explicit_inds_slice
    else:
        return  explicit_inds_slice + (slice_inds,)

In order to use this, one can simply call the slice as usual. Examples:

a = my_array[:, :, :, 8]
b = my_array[:, :, :, 2:7]
c = my_array[...,3] # (equivalent to my_array.take(axis=-1,indices=3)

are equivalent to

a = my_array[_slice_along_axis(8,axis=3)]
b = my_array[_slice_along_axis(slice(2,7),axis=3)]
c = my_array[_slice_along_axis(3,axis=-1)]

One advantage of this approach is that the slice can be generated once and then used in other arrays, even if the number of dimensions in the other arrays is not the same (case of a negative axis).

Answer 2

this is very very late indeed! But I got Leland's answer and expanded it so it works with multiple axis and slice arguments. Here is the verbose version of the function

from numpy import *

def slicer(a, axis=None, slices=None):
    if not hasattr(axis, '__iter__'):
        axis = [axis]
    if not hasattr(slices, '__iter__') or len(slices) != len(axis):
        slices = [slices]
    slices = [ sl if isinstance(sl,slice) else slice(*sl) for sl in slices ]
    mask = []
    fixed_axis = array(axis) % a.ndim
    case = dict(zip(fixed_axis, slices))
    for dim, size in enumerate(a.shape):
        mask.append( case[dim] if dim in fixed_axis else slice(None) )
    return a[tuple(mask)]

it works for variable amount of axes and with tuples of slices as input

>>> a = array( range(10**4) ).reshape(10,10,10,10)
>>> slicer( a, -2, (1,3) ).shape
(10, 10, 2, 10)
>>> slicer( a, axis=(-1,-2,0), slices=((3,), s_[:5], slice(3,None)) ).shape
(7, 10, 5, 3)

---

a slightly more compact version

def slicer2(a, axis=None, slices=None):
    ensure_iter = lambda l: l if hasattr(l, '__iter__') else [l]
    axis = array(ensure_iter(axis)) % a.ndim
    if len(ensure_iter(slices)) != len(axis):
        slices = [slices]
    slice_selector = dict(zip(axis, [ sl if isinstance(sl,slice) else slice(*sl) for sl in ensure_iter(slices) ]))
    element = lambda dim_: slice_selector[dim_] if dim_ in slice_selector.keys() else slice(None)
    return a[( element(dim) for dim in range(a.ndim) )]

Answer 3

This is very late to the party, but I have an alternate slicing function that performs slightly better than those from the other answers:

def array_slice(a, axis, start, end, step=1):
    return a[(slice(None),) * (axis % a.ndim) + (slice(start, end, step),)]

---

Here's a code testing each answer. Each version is labeled with the name of the user who posted the answer:

import numpy as np
from timeit import timeit

def answer_dms(a, axis, start, end, step=1):
    slc = [slice(None)] * len(a.shape)
    slc[axis] = slice(start, end, step)
    return a[slc]

def answer_smiglo(a, axis, start, end, step=1):
    return a.take(indices=range(start, end, step), axis=axis)

def answer_eelkespaak(a, axis, start, end, step=1):
    sl = [slice(None)] * m.ndim
    sl[axis] = slice(start, end, step)
    return a[tuple(sl)]

def answer_clemisch(a, axis, start, end, step=1):
    a = np.moveaxis(a, axis, 0)
    a = a[start:end:step]
    return np.moveaxis(a, 0, axis)

def answer_leland(a, axis, start, end, step=1):
    return a[(slice(None),) * (axis % a.ndim) + (slice(start, end, step),)]

if __name__ == '__main__':
    m = np.arange(2*3*5).reshape((2,3,5))
    axis, start, end = 2, 1, 3
    target = m[:, :, 1:3]
    for answer in (answer_dms, answer_smiglo, answer_eelkespaak,
                   answer_clemisch, answer_leland):
        print(answer.__name__)
        m_copy = m.copy()
        m_slice = answer(m_copy, axis, start, end)
        c = np.allclose(target, m_slice)
        print('correct: %s' %c)
        t = timeit('answer(m, axis, start, end)',
                   setup='from __main__ import answer, m, axis, start, end')
        print('time:    %s' %t)
        try:
            m_slice[0,0,0] = 42
        except:
            print('method:  view_only')
        finally:
            if np.allclose(m, m_copy):
                print('method:  copy')
            else:
                print('method:  in_place')
        print('')

Here are the results:

answer_dms

Warning (from warnings module):
  File "C:\Users\leland.hepworth\test_dynamic_slicing.py", line 7
    return a[slc]
FutureWarning: Using a non-tuple sequence for multidimensional indexing is 
deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be 
interpreted as an array index, `arr[np.array(seq)]`, which will result either in an 
error or a different result.
correct: True
time:    2.2048302
method:  in_place

answer_smiglo
correct: True
time:    5.9013344
method:  copy

answer_eelkespaak
correct: True
time:    1.1219435999999998
method:  in_place

answer_clemisch
correct: True
time:    13.707583699999999
method:  in_place

answer_leland
correct: True
time:    0.9781496999999995
method:  in_place

• DSM's answer includes a few suggestions for improvement in the comments.
• EelkeSpaak's answer applies those improvements, which avoids the warning and is quicker.
• Śmigło's answer involving np.take gives worse results, and while it is not view-only, it does create a copy.
• clemisch's answer involving np.moveaxis takes the longest time to complete, but it surprisingly references back to the previous array's memory location.
• My answer removes the need for the intermediary slicing list. It also uses a shorter slicing index when the slicing axis is toward the beginning. This gives the quickest results, with additional improvements as axis is closer to 0.

I also added a step parameter to each version, in case that is something you need.

Answer 4

There is an elegant way to access an arbitrary axis n of array x: Use numpy.moveaxis¹ to move the axis of interest to the front.

x_move = np.moveaxis(x, n, 0)  # move n-th axis to front
x_move[start:end]              # access n-th axis

The catch is that you likely have to apply moveaxis on other arrays you use with the output of x_move[start:end] to keep axis order consistent. The array x_move is only a view, so every change you make to its front axis corresponds to a change of x in the n-th axis (i.e. you can read/write to x_move).

---

_{1) You could also use swapaxes to not worry about the order of n and 0, contrary to moveaxis(x, n, 0). I prefer moveaxis over swapaxes because it only alters the order concerning n.}

Answer 5

This is a bit late to the party, but the default Numpy way to do this is numpy.take. However, that one always copies data (since it supports fancy indexing, it always assumes this is possible). To avoid that (in many cases you will want a view of the data, not a copy), fallback to the slice(None) option already mentioned in the other answer, possibly wrapping it in a nice function:

def simple_slice(arr, inds, axis):
    # this does the same as np.take() except only supports simple slicing, not
    # advanced indexing, and thus is much faster
    sl = [slice(None)] * arr.ndim
    sl[axis] = inds
    return arr[tuple(sl)]

Answer 6

As it was not mentioned clearly enough (and i was looking for it too):

an equivalent to:

a = my_array[:, :, :, 8]
b = my_array[:, :, :, 2:7]

is:

a = my_array.take(indices=8, axis=3)
b = my_array.take(indices=range(2, 7), axis=3)

Answer 7

I think one way would be to use slice(None):

>>> m = np.arange(2*3*5).reshape((2,3,5))
>>> axis, start, end = 2, 1, 3
>>> target = m[:, :, 1:3]
>>> target
array([[[ 1,  2],
        [ 6,  7],
        [11, 12]],

       [[16, 17],
        [21, 22],
        [26, 27]]])
>>> slc = [slice(None)] * len(m.shape)
>>> slc[axis] = slice(start, end)
>>> np.allclose(m[slc], target)
True

I have a vague feeling I've used a function for this before, but I can't seem to find it now..