Operation by indexing only last axis

Question

I have an array of 3 dimensional vectors. The dimension of the array is arbitrary: it could be a single (N×3), double (M×N×3), triple (K×M×N×3) etc. I need to operate on two components of the vector while preserving the other dimensions.

For example, if I know it is three dimensionsional, I could do the following:

R = numpy.arctan2(A[:,:,:,1], A[:,:,:,0])

which gives me a three dimensional array of scalar values.

Now, to be able to do this on arbitrary number of dimensions. I need to slice over all other dimensions except the the last. So far, I'm able to do it with this:

s = [numpy.s_[:]] * (len(A.shape)-1)
R = numpy.arctan2(A[s+[1]], A[s+[0]])

which works even for single vectors. Is there a more numpythonic way of achieving the above?

Answer 1

I found an even nicer way. This here works for me

R = numpy.arctan2(A[...,1],A[...,0])