Slicing n-dimensional numpy array using list of indices

Question

Say I have a 3 dimensional numpy array:

np.random.seed(1145)
A = np.random.random((5,5,5))

and I have two lists of indices corresponding to the 2nd and 3rd dimensions:

second = [1,2]
third = [3,4]

and I want to select the elements in the numpy array corresponding to

A[:][second][third]

so the shape of the sliced array would be (5,2,2) and

A[:][second][third].flatten()

would be equivalent to to:

In [226]:

for i in range(5):
    for j in second:
        for k in third:
            print A[i][j][k]

0.556091074129
0.622016249651
0.622530505868
0.914954716368
0.729005532319
0.253214472335
0.892869371179
0.98279375528
0.814240066639
0.986060321906
0.829987410941
0.776715489939
0.404772469431
0.204696635072
0.190891168574
0.869554447412
0.364076117846
0.04760811817
0.440210532601
0.981601369658

Is there a way to slice a numpy array in this way? So far when I try A[:][second][third] I get IndexError: index 3 is out of bounds for axis 0 with size 2 because the [:] for the first dimension seems to be ignored.

Answer 1

One way would be to use np.ix_:

>>> out = A[np.ix_(range(A.shape[0]),second, third)]
>>> out.shape
(5, 2, 2)
>>> manual = [A[i,j,k] for i in range(5) for j in second for k in third]
>>> (out.ravel() == manual).all()
True

Downside is that you have to specify the missing coordinate ranges explicitly, but you could wrap that into a function.