How to convert N-D slice to indexes in NumPy?

Question

Given any N-tuple of slices (aka N-D slice) in NumPy how to convert it to corresponding indexes of N-D array represented as tuple of 1D arrays (indexes along each axes)? E.g. if we have np.nd_slice_to_indexes next code:

import numpy as np
print(np.nd_slice_to_indexes(np.s_[1 : 3]))
print(np.nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))

should print

(array([1, 2]),)
(array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))

It is common for NumPy to represent indexes of N-D array as N-tuple of 1-D arrays of same length (each element of k-th array in tuple represents next index along k-th dimension). E.g. np.nonzero returns such N-tuple in code

print(np.nonzero([[0, 1, 1], [1, 1, 0]])) # Non-zero elements in 2D array.
# (array([0, 0, 1, 1], dtype=int64), array([1, 2, 0, 1], dtype=int64))

Same behavior should be achieved like in Pythonic function below, but in a more efficient (performant) way:

Try it online!

import numpy as np

def nd_slice_to_indexes(nd_slice):
    assert type(nd_slice) in [tuple, slice], type(nd_slice)
    if type(nd_slice) is not tuple:
        nd_slice = (nd_slice,)
    def iter_slices(slices):
        if len(slices) == 0:
            yield ()
        else:
            for i in range(slices[0].start, slices[0].stop, slices[0].step or 1):
                for r in iter_slices(slices[1:]):
                    yield (i,) + r
    *res, = np.vstack(list(iter_slices(nd_slice))).T
    return tuple(res)

print(nd_slice_to_indexes(np.s_[1 : 3]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2, 8 : 14 : 3]))
# (array([1, 2]),)
# (array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))
# (array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]), array([5, 5, 7, 7, 9, 9, 5, 5, 7, 7, 9, 9]), array([ 8, 11,  8, 11,  8, 11,  8, 11,  8, 11,  8, 11]))

@hpaulj Thanks! Works! Added your suggestion regarding np.mgrid as this answer. — Arty
– Arty, Commented Sep 28, 2020 at 7:24
@hpaulj Although looking at code of np.mgrid it could be the case that it is not very performant implementation. — Arty
– Arty, Commented Sep 28, 2020 at 8:42

Arty · Accepted Answer · 2020-09-28 08:29:36Z

1

Thanks to suggestion of @hpaulj solved task efficiently using np.mgrid.

Try it online!

import numpy as np

def nd_slice_to_indexes(nd_slice):
    grid = np.mgrid[{tuple: nd_slice, slice: (nd_slice,)}[type(nd_slice)]]
    return tuple(grid[i].ravel() for i in range(grid.shape[0]))
    
print(nd_slice_to_indexes(np.s_[1 : 3]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2, 8 : 14 : 3]))
# (array([1, 2]),)
# (array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))
# (array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]), array([5, 5, 7, 7, 9, 9, 5, 5, 7, 7, 9, 9]), array([ 8, 11,  8, 11,  8, 11,  8, 11,  8, 11,  8, 11]))

edited Sep 28, 2020 at 8:29

answered Sep 28, 2020 at 7:23

Arty

17k6 gold badges60 silver badges96 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Arty Over a year ago

Might be not efficient, didn't do timings, not sure if np.mgrid is impelemented to be really fast.

Alex van Houten Over a year ago

How should this code be adjusted for processing a list of slices, e.g. as returned from scipy.ndimage.find_objects?

Collectives™ on Stack Overflow

How to convert N-D slice to indexes in NumPy?

1 Answer 1

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related