Selecting which dimension to index in a numpy array

Question

I am writing a program that is suppose to be able to import numpy arrays of some higher dimension, e.g. something like an array a:

a = numpy.zeros([3,5,7,2])

Further, each dimension will correspond to some physical dimension, e.g. frequency, distance, ... and I will also import arrays with information about these dimensions, e.g. for a above:

freq = [1,2,3]
time = [0,1,2,3,4,5,6]
distance = [0,0,0,4,1]
angle = [0,180]

Clearly from this example and the signature it can be figured out that freq belong to dimension 0, time to dimension 2 and so on. But since this is not known in advance, I can take a frequency slice like

a_f1 = a[1,:,:,:]

since I do not know which dimension the frequency is indexed.

So, what I would like is to have some way to chose which dimension to index with an index; in some Python'ish code something like

a_f1 = a.get_slice([0,], [[1],])

This is suppose to return the slice with index 1 from dimension 0 and the full other dimensions.

Doing

a_p = a[0, 1:, ::2, :-1]

would then correspond to something like

a_p = a.get_slice([0, 1, 2, 3], [[0,], [1,2,3,4], [0,2,4,6], [0,]])

Answer 1

You can fairly easily construct a tuple of indices, using slice objects where needed, and then use this to index into your array. The basic is recipe is this:

indices = {
    0: # put here whatever you want to get on dimension 0,
    1: # put here whatever you want to get on dimension 1,
    # leave out whatever dimensions you want to get all of
}
ix = [indices.get(dim, slice(None)) for dim in range(arr.ndim)]
arr[ix]

Here I have done it with a dictionary since I think that makes it easier to see which dimension goes with which indexer.

So with your example data:

x = np.zeros([3,5,7,2])

We do this:

indices = {0: 1}
ix = [indices.get(dim, slice(None)) for dim in range(x.ndim)]

>>> x[ix].shape
(5L, 7L, 2L)

Because your array is all zeros, I'm just showing the shape of the result to indicate that it is what we want. (Even if it weren't all zeros, it's hard to read a 3D array in text form.)

For your second example:

indices = {
    0: 0,
    1: slice(1, None),
    2: slice(None, None, 2),
    3: slice(None, -1)
}
ix = [indices.get(dim, slice(None)) for dim in range(x.ndim)]

>>> x[ix].shape
(4L, 4L, 1L)

You can see that the shape corresponds to the number of values in your a_p example. One thing to note is that the first dimension is gone, since you only specified one value for that index. The last dimension still exists, but with a length of one, because you specified a slice that happens to just get one element. (This is the same reason that some_list[0] gives you a single value, but some_list[:1] gives you a one-element list.)

Answer 2

You can use advanced indexing to achieve this.

The index for each dimension needs to be shaped appropriately so that the indices will broadcast correctly across the array. For example, the index for the first dimension of a 3-d array needs to be shaped (x, 1, 1) so that it will broadcast across the first dimension. The index for the second dimension of a 3-d array needs to be shaped (1, y, 1) so that it will broadcast across the second dimension.

import numpy as np
a = np.zeros([3,5,7,2])
b = a[0, 1:, ::2, :-1]

indices = [[0,], [1,2,3,4], [0,2,4,6], [0,]]
def get_aslice(a, indices):
    n_dim_ = len(indices)
    index_array = [np.array(thing) for thing in indices]
    idx = []
    # reshape the arrays by adding single-dimensional entries
    # based on the position in the index array 
    for d, thing in enumerate(index_array):
        shape = [1] * n_dim_
        shape[d] = thing.shape[0]
        #print(d, shape)
        idx.append(thing.reshape(shape))

    c = a[idx]

    # to remove leading single-dimensional entries from the shape
    #while c.shape[0] == 1:
    #   c = np.squeeze(c, 0)

    # To remove all single-dimensional entries from the shape
    #c = np.squeeze(c).shape

    return c

For a as an input, it returns an array with shape (1,4,4,1) your a_p example has a shape of (4,4,1). If the extra dimensions need to be removed un-comment the np.squeeze lines in the function.

---

Now I feel silly. While reading the docs slower I noticed numpy has an indexing routine that does what you want - numpy.ix_

>>> a = numpy.zeros([3,5,7,2])
>>> indices = [[0,], [1,2,3,4], [0,2,4,6], [0,]]
>>> index_arrays = np.ix_(*indices)
>>> a_p = a[index_arrays]
>>> a_p.shape
(1, 4, 4, 1)
>>> a_p = np.squeeze(a_p)
>>> a_p.shape
(4, 4)
>>>