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)

• Aelius: Thank you very much, this is really helpful! Just one more clarification: why eval is so bad? Thank you
• Reinderien: It greatly increases the surface area for bugs and security flaws.