How to unpack a list of colons and NumPy Nones to index an array?

Question

I am writing a program that will have an arbitrary number of : and None in arbitrary locations of an n-dimensional NumPy array. Therefore, I want a way to unpack these : and None axis operators into the [] that indexes an array and auto-populates certain axes according to where the : and None are. According to Pylance:

Unpack operator in subscript requires Python 3.11 or newerPylance

However, while using Python 3.11, I get the following error:

Traceback (most recent call last):
  File "/home/.../quant.py", line 261, in <module>
    print(arr[*lhs_axes] + arr2[None,None,:])
          ~~~^^^^^^^^^^^
IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

Current code:

import numpy as np
if __name__ == "__main__":
    lhs_ind, rhs_ind = 'ij', 'k'
        
    lhs_axes = [':' for i in lhs_ind]
    lhs_axes.append(None)

    arr1 = np.ones((2,2))
    arr2 = np.ones(2)
    print(arr1[*lhs_axes] + arr2[None,None,:])

Note that I tried using the string 'None' instead of None, but it returned the same error. — ExactPlace441, Oct 27 '22 at 17:52

score 1 · Answer 1 · answered Oct 27 '22 at 18:26

A ':' string and the : are very different things. : is equivalent to slice(None).

Use a slice and tuple:

import numpy as np
if __name__ == "__main__":
    lhs_ind, rhs_ind = 'ij', 'k'
 
    lhs_axes = [slice(None) for i in lhs_ind]
    lhs_axes.append(None)
    arr1 = np.ones((2,2))
    arr2 = np.ones(2)
    print(arr1[tuple(lhs_axes)] + arr2[None,None,:])

Output:

[[[2. 2.]
  [2. 2.]]

 [[2. 2.]
  [2. 2.]]]

score 0 · Answer 2 · answered Oct 27 '22 at 18:54

If I make a list with slice and None objects:

In [44]: idx = [slice(None),slice(None),None]
In [45]: idx
Out[45]: [slice(None, None, None), slice(None, None, None), None]

I can make a tuple from the list with:

In [46]: tuple(idx)
Out[46]: (slice(None, None, None), slice(None, None, None), None)

or with unpacking - in a tuple context (a list context as well):

In [47]: (*idx,)
Out[47]: (slice(None, None, None), slice(None, None, None), None)

For indexing, the tuple works fine:

In [48]: arr = np.ones((2,3,4),int)    

In [50]: arr[tuple(idx)].shape
Out[50]: (2, 3, 1, 4)

The unpacking that you try to use is a 3.11 addition:

In [51]: arr[*idx].shape
  Input In [51]
    arr[*idx].shape
            ^
SyntaxError: invalid syntax

But it isn't needed if you provide a tuple. arr[1,2,3] is the same as arr[(1,2,3)]. arr[[1,2,3]] can be problematic. Strictly speaking that's indexing the first dimension, same as arr[np.array([1,2,3])]. But for historical reasons, sometimes such a list is treated as a tuple. In which case you'll get a future warning:

FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; 
use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted 
as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.

How to unpack a list of colons and NumPy Nones to index an array?

2 Answers2