Python: add axis for an axis specified

Question

import numpy as np
x=np.zeros((2,3,4,5))

For the x, is there simpler way to perform (x[:,1,:,:])[:,np.newaxis, :, :]?

Answer 1

If you have to do the slicing and expanding dynamically, you an use something like:

import numpy as np
x=np.zeros((2,3,4,5))
element_idx = 1
axis_idx = 1
sliced_array = np.take(x, element_idx, axis=axis_idx)
result = np.expand_dims(sliced_array, axis_idx)

If, on the other hand, the axis is fixed (often true) then the solution suggested by hpaulj will work nicely:

result = x[:, 1][:, None]