You can do the following:
import numpy as np
A = np.arange(12).reshape((2, 3, 2))
print(A)
x = [1, 1]
print(A[(slice(None), *x)])
You can use slice(None)
instead of :
to build a tuple of slices. The tuple environment allows for value unpacking with the * operator.
Output:
[[[ 0 1]
[ 2 3]
[ 4 5]]
[[ 6 7]
[ 8 9]
[10 11]]]
[3 9]
To verify it matches:
import numpy as np
A = np.arange(12).reshape((2, 3, 2))
x = [1, 1]
s = (slice(None), *x)
print(np.allclose(A[s], A[:, 1, 1])) # True
*This is a modification of answers found here: Slicing a numpy array along a dynamically specified axis
Edit to reflect edit on question and comment:
To clarify, you can unpack any iterable you like in the tuple environment. The * operator functions normally in within the tuple. Order your elements however you like. Mix in different iterables, types, slice(None)
, how ever you want to build your slices, as long as you end up with a valid sequence of values, it will behave as expected.
import numpy as np
A = np.arange(12).reshape((2, 3, 2))
t = [True, False]
x = [1, 1]
print(np.allclose(A[(*t, *x)], A[True, False, 1, 1])) # True
You can also add full lists as well in the tuple:
print(np.allclose(A[(t, *x)], A[[True, False], 1, 1])) # True