13

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?

Daniel Widdis
  • 8,424
  • 13
  • 41
  • 63
Andrew Schwartz
  • 4,440
  • 3
  • 25
  • 58
  • 2
    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](http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html) mentions `slice` saying that it "can be useful for constructing generic code that works on arrays of arbitrary dimension." – Bakuriu Jun 26 '14 at 14:15
  • 1
    Yup, that's the answer, thanks. The None input argument to slice wasn't clear to me, but this is great. – Andrew Schwartz Jun 26 '14 at 15:20
  • excellent question and answers, exactly what I was looking for – pcko1 Apr 15 '19 at 14:30

3 Answers3

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.

Bakuriu
  • 98,325
  • 22
  • 197
  • 231
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.

Pablo Francisco Pérez Hidalgo
  • 27,044
  • 8
  • 36
  • 62
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)])]
Quazer
  • 353
  • 1
  • 8