Python: Index an array using the colon operator in an arbitrary dimension

Question

I have a numpy nd array. A simplified version of my task is to take a vector from along each axis. To illustrate:

import numpy
x = numpy.array(range(24)).reshape((2,3,4))
x0 = x[0,0,:]
x1 = x[0,:,0]
x2 = x[:,0,0]

However I do not necessarily know the number of dimensions x will have. So the challenge is how to place the colon : indexing operator in a variable location. An example of what such syntax could look like:

n = x.ndim
ind = list(np.zeros(n))
dim = 0
ind[dim] = ':'
y = x[ind]

or

y = indexer.index(x,ind)

for some module indexer. I could write it, but I feel like this must already be solved, I can't be the only one who wants to do this. In MATLAB, for example, you can do this with the subsref() function.

Does any such construct exist in python / numpy / other module?

You can easily construct the tuple at runtime: (0,) * n yields a tuple of n zeros. The : is equivalent to slice(None). Writing a function to put the slice were you want should be quite trivial... In fact numpy's documentation about indexing mentions slice saying that it "can be useful for constructing generic code that works on arrays of arbitrary dimension." — Bakuriu
– Bakuriu, Commented Jun 26, 2014 at 14:15
Yup, that's the answer, thanks. The None input argument to slice wasn't clear to me, but this is great. — Andrew Schwartz
– Andrew Schwartz, Commented Jun 26, 2014 at 15:20
excellent question and answers, exactly what I was looking for — pcko1
– pcko1, Commented Apr 15, 2019 at 14:30

Bakuriu · Accepted Answer · 2014-06-26 15:36:22Z

16

As suggested from numpy's documentation about indexing you can use the slice built-in function and tuple concatenation to create variable indexes.

In fact the : in the subscript is simply the literal notation for a slice literal.

In particular : is equivalent to slice(None) (which, itself, is equivalent to slice(None, None, None) where the arguments are start, stop and step).

For example:

a[(0,) * N + (slice(None),)]

is equivalent to:

a[0, 0, ..., 0, :]   # with N zeros

The : notation for slices can only be used directly inside a subscript. For example this fails:

In [10]: a[(0,0,:)]
  File "<ipython-input-10-f41b33bd742f>", line 1
    a[(0,0,:)]
           ^
SyntaxError: invalid syntax

To allow extracting a slice from an array of arbitrary dimensions you can write a simple function such as:

def make_index(num_dimension, slice_pos):
    zeros = [0] * num_dimension
    zeros[slice_pos] = slice(None)
    return tuple(zeros)

And use it as in:

In [3]: a = np.array(range(24)).reshape((2, 3, 4))

In [4]: a[make_index(3, 2)]
Out[4]: array([0, 1, 2, 3])

In [5]: a[make_index(3, 1)]
Out[5]: array([0, 4, 8])

In [6]: a[make_index(3, 0)]
Out[6]: array([ 0, 12])

You can generalize make_index to do any kind of things. The important thing to remember is that it should, in the end, return a tuple containing either integers or slices.

answered Jun 26, 2014 at 15:36

Bakuriu

103k23 gold badges208 silver badges236 bronze badges

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

1 Comment

pcko1 Over a year ago

a[(0,) * N + (slice(None),)] is a very useful line of code, thanks

Pablo Francisco Pérez Hidalgo · Accepted Answer · 2014-06-26 14:17:11Z

2

You could compose an string with the code selecting the dimension you want and use eval to execute that code string.

An start is:

n = 2
sel = "0,"*(n-1) + ":"
eval('x[' + sel + ']')

To get exactly what you want, thinks are a little bit more complicated (but not so much):

ind = 2
n = 3
sel = "".join([ ("0" if i != ind else ":") + ("," if i < n-1 else "") for i in xrange(n)])
eval('x[' + sel + ']')

It is the same strategy that is used for Dynamic SQL.

edited Jun 26, 2014 at 14:17

answered Jun 26, 2014 at 13:58

Pablo Francisco Pérez Hidalgo

27.8k8 gold badges39 silver badges64 bronze badges

2 Comments

Bakuriu Over a year ago

Why are you using eval? There is no need for that. What you are doing with that eval expression is: x[(0,) * (n-1) + (slice(None),)].

Pablo Francisco Pérez Hidalgo Over a year ago

@Bakuriu Beautiful solution of yours. It is more elegant. Mine is somehow more "literal".

Quazer · Accepted Answer · 2023-01-19 18:43:07Z

0

Create the function subref (as MATLAB your wrote)

def subref(arr, ind): 
    return arr[tuple([slice(None) if i == ind else 0 for i in range(arr.ndim)])]

answered Jan 19, 2023 at 18:43

Quazer

3831 silver badge8 bronze badges

Collectives™ on Stack Overflow

Python: Index an array using the colon operator in an arbitrary dimension

3 Answers 3

1 Comment

2 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related