Exposing array indices programmatically

Programming Diary
1

Often times, we need to index certain parts of an array.

import numpy as np

array = 0.1 * np.arange(24).reshape(4, 6)
print(array)

[[0.  0.1 0.2 0.3 0.4 0.5]
 [0.6 0.7 0.8 0.9 1.  1.1]
 [1.2 1.3 1.4 1.5 1.6 1.7]
 [1.8 1.9 2.  2.1 2.2 2.3]]

To index these arrays, we use can use slicing. E.g., to index the first element in the first row:

array[0, 0]

0.0

and to index the entire first row:

array[0, :]

array([0. , 0.1, 0.2, 0.3, 0.4, 0.5])

Now, this becomes a bit more complicated when want to expose indexing in a method

def index_array(array, indices):
    return array[indices]

E.g. indexing the first element does not yield the desired result:

index_array(array, [0, 0])

array([[0. , 0.1, 0.2, 0.3, 0.4, 0.5],
       [0. , 0.1, 0.2, 0.3, 0.4, 0.5]])

And indexing entire dimensions using index notation raises an SyntaxError


index_array(array, [0, :])

  File "<stdin>", line 2
    index_array(array, [0, :])
                           ^
SyntaxError: invalid syntax

Instead, we need to pass a slice:

help(slice)

Help on class slice in module builtins:

class slice(object)
 |  slice(stop)
 |  slice(start, stop[, step])
 |  
 |  Create a slice object.  This is used for extended slicing (e.g. a[0:10:2]).
 |  
 |  Methods defined here:
 |  
 |  __eq__(self, value, /)
 |      Return self==value.
 |  
 |  __ge__(self, value, /)
 |      Return self>=value.
 |  
 |  __getattribute__(self, name, /)
 |      Return getattr(self, name).
 |  
 |  __gt__(self, value, /)
 |      Return self>value.
 |  
 |  __le__(self, value, /)
 |      Return self<=value.
 |  
 |  __lt__(self, value, /)
 |      Return self<value.
 |  
 |  __ne__(self, value, /)
 |      Return self!=value.
 |  
 |  __reduce__(...)
 |      Return state information for pickling.
 |  
 |  __repr__(self, /)
 |      Return repr(self).
 |  
 |  indices(...)
 |      S.indices(len) -> (start, stop, stride)
 |      
 |      Assuming a sequence of length len, calculate the start and stop
 |      indices, and the stride length of the extended slice described by
 |      S. Out of bounds indices are clipped in a manner consistent with the
 |      handling of normal slices.
 |  
 |  ----------------------------------------------------------------------
 |  Static methods defined here:
 |  
 |  __new__(*args, **kwargs) from builtins.type
 |      Create and return a new object.  See help(type) for accurate signature.
 |  
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |  
 |  start
 |  
 |  step
 |  
 |  stop
 |  
 |  ----------------------------------------------------------------------
 |  Data and other attributes defined here:
 |  
 |  __hash__ = None

In this case, we would need to write:

sl = slice(None)
index_array(array, (0, sl))

array([0. , 0.1, 0.2, 0.3, 0.4, 0.5])

While this works, it’s not really pretty.
Instead, we can use numpy’s s_ functionality.

index_array(array, np.s_[0, :])

array([0. , 0.1, 0.2, 0.3, 0.4, 0.5])
1 Like