Python 使用索引元组切片 numpy 数组

Question

I want to slice an N-dimensional numpy array a using an N-1-dimensional index tuple b . Below is a case with N = 4:

import numpy as np

a = np.random.rand(20,1,5,4)
b1 = (0,0,0)
b2 = (0,3,2)

# want to get the slice c = a[:,b2[0],b2[1],b2[2]] without explicitly writing all N-1 elements of b2

c = a[:,b1] # no error message, but gives a[:,[b1[0],b1[1],b1[2]],:,:]
c = a[:,b2] # error message (because b2[1] = 3 larger than a.shape[1])

What is the correct way to "unpack" the tuple index elements into the different dimensions? I tried using the asterisk operator ( *b ), but that did not seem to work.

Answer 1

a[(..., *b1)]

This boils down to calling the __getitem__ method manually with a variable argument in the index tuple:

In [23]: a.__getitem__((..., *b1))
Out[23]: 
array([0.93049554, 0.63858628, 0.31665933, 0.15048072, 0.17083849,
       0.34621308, 0.66980288, 0.75191251, 0.05844084, 0.16715787,
       0.80409005, 0.64181912, 0.76463963, 0.75616746, 0.44955961,
       0.19034925, 0.8841084 , 0.54450091, 0.02033021, 0.21275405])

In [21]: a[:, 0, 0, 0]
Out[21]: 
array([0.93049554, 0.63858628, 0.31665933, 0.15048072, 0.17083849,
       0.34621308, 0.66980288, 0.75191251, 0.05844084, 0.16715787,
       0.80409005, 0.64181912, 0.76463963, 0.75616746, 0.44955961,
       0.19034925, 0.8841084 , 0.54450091, 0.02033021, 0.21275405])

Be aware though that the usage of the ellipsis ( ... ) has a slightly more general meaning. Otherwise, you have to construct the : slice manually:

In [20]: a.__getitem__((slice(None), *b1))
Out[20]: 
array([0.93049554, 0.63858628, 0.31665933, 0.15048072, 0.17083849,
       0.34621308, 0.66980288, 0.75191251, 0.05844084, 0.16715787,
       0.80409005, 0.64181912, 0.76463963, 0.75616746, 0.44955961,
       0.19034925, 0.8841084 , 0.54450091, 0.02033021, 0.21275405])