apply a function to np.ndarray?

Question

I have one numpy array, for instance a (a.shape = (2,20,50,50)). I would like to apply a function over its second axes.

My function is the difference between each elements only along the second axis, ie

res = a[:,i+1,:,:] - a[:,i,:,:] for i in range(20)

I have already tried it with lambda function, but the output is a list. I would like to have the result with the same shape as a. That means I want to have res.shape = (2,20,50,50).

I do appreciate that if someone guide me.

Thank you in advance.

Answer 1

You don't need to apply a function. Just subtract them directly.

res = a[:,1:,:,:]-a[:,0:-1,:,:]

Moreover, you won't get (2,20,50,50) ndarray but (2,19,50,50).