non adjacent slicing of numpy multidimensional array in python

Question

I have a multidimensional array a:

a = np.random.uniform(1,10,(2,4,2,3,10,10))

For dimensions 4-6, I have 3 lists which contain the indexes for slicing that dimension of array 'a'

dim4 = [0,2]
dim5 = [3,5,9]
dim6 = [1,2,7,8]

How do I slice out array 'a' such that i get:

b = a[0,:,0,dim4,dim5,dim6]

So b should be an array with shape (4,2,3,4), and containing elements from the corresponding dimensions of a. When I try the code above, I get an error saying that different shapes can't be broadcast together for axis 4-6, but if I were to do:

b = a[0,:,0:2,0:3,0:4]

then it does work, even though the slicing lists all have different lengths. So how do you slice multidimensional arrays with non adjacent indexes?

Answer 1

You can use the numpy.ix_ function to construct complex indexing like this. It takes a sequence of array_like, and makes an "open mesh" from them. The example from the docstring is pretty clear:

Using ix_ one can quickly construct index arrays that will index the cross product. a[np.ix_([1,3],[2,5])] returns the array [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

So, for your data, you'd do:

>>> indices = np.ix_((0,), np.arange(a.shape[1]), (0,), dim4, dim5, dim6)
>>> a[indices].shape
(1, 4, 1, 2, 3, 4)

Get rid of the size-1 dimensions with np.squeeze:

>>> np.squeeze(a[indices]).shape
(4, 2, 3, 4)