Home > Enterprise >  Mixing numpy fancy indexing with ellipsis
Mixing numpy fancy indexing with ellipsis

Time:09-27

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)]
  • Related