通過僅索引最後一軸的操作

Question

我有一個3維向量的數組。陣列的尺寸是任意的：它可以是單個（N×3），雙倍（M×N×3），三倍（K×M×N×3）等等。我需要操作矢量的兩個分量同時保留其他維度。

例如，如果我知道是三dimensionsional，我能做到以下幾點：

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

這給了我標量值的三維陣列。

現在，能夠做到這一點任意數字的維度。除了最後一個之外，我需要對所有其他維度進行切片。到目前爲止，我可以這樣做：

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

它甚至可以用於單個矢量。有沒有更實用的方法來達到上述目的？

Answer 1

我發現了一個更好的方法。這適用於我

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