Dynamically indexing numpy array

Question

I want to create a function that takes a numpy array, an axis and an index of that axis and returns the array with the index on the specified axis fixed. What I thought is to create a string that change dynamically and then is evaluated as an index slicing in the array (as showed in this answer). I came up with the following function:

import numpy as np

def select_slc_ax(arr, slc, axs):

    dim = len(arr.shape)-1
    slc = str(slc)
    slice_str = ":,"*axs+slc+",:"*(dim-axs)
    print(slice_str)
    slice_obj = eval(f'np.s_[{slice_str}]')
    return arr[slice_obj]

Example

>>> arr = np.array([[[0, 0, 0],
                     [0, 0, 0],
                     [0, 0, 0]],
                    [[0, 1, 0],
                     [1, 1, 1],
                     [0, 1, 0]],
                    [[0, 0, 0],
                     [0, 0, 0], 
                     [0, 0, 0]]], dtype='uint8')

>>> select_slc_ax(arr, 2, 1)
:,2,:
array([[0, 0, 0],
       [0, 1, 0],
       [0, 0, 0]], dtype=uint8)

I was wondering if there is a better method to do this.

Answer 1

• Avoid eval at all costs
• Use the slice built-in for your : components, or the numeric index at the selected axis
• Do not abbreviate your variable names
• Type-hint your function signature
• Turn your example into something resembling a unit test with an assert
• Modify the data in your example to add more non-zero entries making it clearer what's happening
• Prefer immutable tuples over mutable lists when passing initialization constants to Numpy
• Prefer the symbolic constants for Numpy types rather than strings

Suggested

import numpy as np


def select_slice_axis(array: np.ndarray, index: int, axis: int) -> np.ndarray:
    slices = tuple(
        index if a == axis else slice(None)
        for a in range(len(array.shape))
    )
    return array[slices]


arr = np.array(
    (
        ((0, 0, 0),
         (0, 0, 0),
         (0, 0, 9)),
        ((0, 1, 0),
         (2, 3, 4),
         (0, 5, 0)),
        ((0, 0, 0),
         (0, 0, 0),
         (8, 0, 0)),
    ),
    dtype=np.uint8,
)

actual = select_slice_axis(arr, 2, 1)
expected = np.array(
    (
        (0, 0, 9),
        (0, 5, 0),
        (8, 0, 0),
    ), dtype=np.uint8
)
assert np.all(expected == actual)