3
\$\begingroup\$

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.

\$\endgroup\$

1 Answer 1

1
\$\begingroup\$
  • 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)
\$\endgroup\$
2
  • \$\begingroup\$ Thank you very much, this is really helpful! Just one more clarification: why eval is so bad? Thank you \$\endgroup\$
    – Aelius
    Aug 9, 2021 at 7:51
  • 1
    \$\begingroup\$ It greatly increases the surface area for bugs and security flaws. \$\endgroup\$
    – Reinderien
    Aug 9, 2021 at 12:51

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.