If I have a multi-dimensional numpy array, how can I dynamically refer to a specific slice?

Question

Let's say I have a numpy array such that

import numpy as np
array_shape = (5,6,7,8,9,4)
big_array = np.random.random(array_shape)

If I'd like to call a specific slice, I need to know the axis, and index, i.e.

small_array = big_array[:,:,3,:,:,:]

I'd like to be able to write a function that returns "small_array" when given the argument (2,3), i.e. third axis, fourth index. The issue is I don't see how to properly format the argument.

Thanks for your help!

Answer 1

Use numpy.take. big_array[:,:,3,:,:] would be equivalent to np.take(big_array, 3, axis=2):

(big_array[:,:,3,:,:] == np.take(big_array, 3, axis=2)).all()
# True