Slicing multi-dimensional array with another array

Question

edited with a clearer example, and included solution

I'd like to slice an arbitrary dimensional array, where I pin the first n dimensions and keep the remaining dimensions. In addition, I'd like to be able to store the n pinning dimensions in a variable. For example

Q = np.arange(24).reshape(2, 3, 4) # array to be sliced
# array([[[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]],
#       [[12, 13, 14, 15],
#        [16, 17, 18, 19],
#        [20, 21, 22, 23]]])

Q[0, 1, ...]  # this is what I want manually
# array([4, 5, 6, 7])

# but programmatically:
s = np.array([0, 1])
Q[s, ...]  # this doesn't do what I want: it uses both s[0] and s[1] along the 0th dimension of Q
# array([[[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]],
#       [[12, 13, 14, 15],
#        [16, 17, 18, 19],
#        [20, 21, 22, 23]]])

np.take(Q, s)  # this unravels the indices and takes the s[i]th elements of Q
# array([0, 1])

Q[tuple(s)]  # this works! Thank you kwin
# array([4, 5, 6, 7])

Is there a clean way to do this?

Answer 1

You could do this:

Q[tuple(s)]

Or this:

np.take(Q, s)

Both of these yield array([0.58383736, 0.80486868]) .

I'm afraid I don't have a great intuition for exactly why the tuple version of s works differently from indexing with s itself. The other thing I intuitively tried is Q[*s] but that's a syntax error.

Answer 2

I am not sure what output you want but there are several things you can do.

If you want the output to be like this:

array([[[0.46988733, 0.19062458],
        [0.69307707, 0.80242129],
        [0.36212295, 0.2927196 ],
        [0.34043998, 0.87408959],
        [0.5096636 , 0.37797475]],

       [[0.98322049, 0.00572271],
        [0.06374176, 0.98195354],
        [0.63195656, 0.44767722],
        [0.61140211, 0.58889763],
        [0.18344186, 0.9587247 ]]])

Q[list(s)] should work. np.array([Q[i] for i in s]) also works.

If you want the output to be like this:

array([0.58383736, 0.80486868])

Then as @kwinkunks mentioned you could use Q[tuple(s)] or np.take(Q, s)