How can I combine fancy indexing with ...
in numpy?
Consider some test data:
# Some test data of high dimensionality
shape = (3,5,2,8)
data = np.arange(np.prod(shape)).reshape(shape)
Maybe I have a tuple of indices. I cannot unpack these values during indexing:
# Arbitrary indices
indices = (0, 4)
# Try unpacking
data[*indices]
# >>> SyntaxError: invalid syntax
Not to worry, I get the same result if I pass the indices individually or as a tuple — nice!
np.all(data[indices[0], indices[1]] == data[indices])
# >>> True
So unpacking is not needed, right? This answer to a related question seems to imply so.
Python's indexing already has built-in support for tuples in general [...] there is no need for iterator unpacking in indexing. Simply convert your iterator to a tuple
At first glance, it even seems to work with fancy indexing:
# New arbitrary indices
new_indices = (np.array([0,1]), np.array([4,4]))
np.all(data[new_indices] == data[new_indices[0], new_indices[1]])
# >>> True
But if I try to use ellipsis
it breaks down. For example if I want to do fancy indexing on trailing axes:
# Passing unpacked values produces the desired result
data[..., new_indices[0], new_indices[1]] # shape=(3,5,2)
# Passing as a tuple does NOT produce the desired result
data[..., new_indices] # shape=(3,5,2,2,2)
Of course if my data is of variable dimensionality, I cannot write ..., new_indices[0], new_indices[1]
. So it seems that in this case unpacking is in fact needed.
Okay, well maybe I can just do this:
data.T[indices[::-1]].T.shape # shape=(3,5,2)
It feels a bit hacky, but it works. Except if we want to have fancy indexing at both leading and trailing axes with ...
in between, it seems we are truly out of luck:
# Some test data of even higher dimensionality
data_shape = (3,5,2,8,12,25)
data = np.arange(np.prod(data_shape)).reshape(data_shape)
index_A = (np.array([0,0,1]), np.array([0,2,1]))
index_B = (np.array([0,0,0]), np.array([1,1,1]))
# Desired result, shape=(3,2)
data[index_A[0], index_A[1], ..., index_B[0], index_B[1]]
# Without unpacking, shape=(2,3,5,2,8)
data[index_A, ..., index_B]
Maybe this is an edge case, but it seems there should be a way to do it. How?
P.S. Why should iterable unpacking raise a SyntaxError in this situation, anyway?
CodePudding user response:
tuple
join does the trick:
In [142]: data[(...,) indices]
Out[142]:
array([[ 4, 20, 36, 52, 68],
[ 84, 100, 116, 132, 148],
[164, 180, 196, 212, 228]])
In [143]: data[...,0,4]
Out[143]:
array([[ 4, 20, 36, 52, 68],
[ 84, 100, 116, 132, 148],
[164, 180, 196, 212, 228]])
In [144]: (...,) indices
Out[144]: (Ellipsis, 0, 4)
unpacking inside a tuple also works:
data[(...,*indices)]