Slicing multi-dimensional array with another array-CodePudding

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.random.rand(3, 5, 2)  # array to be sliced
s = np.array([0, 1])  # the pinned first n dimensions
Q[0, 1, ...]  # this is what I want manually
Q[s, ...]  # this doesn't do what I want: it uses both s[0] and s[1] along the 0th dimension of Q

Is there a clean way to do this?

CodePudding user response：

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.

CodePudding user response：

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)