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)