更快的numpy笛卡尔到球坐标转换？

Question

我有一个包含300万个数据点的数组，这些数据点来自一个三轴加速度计（也就是X、Y、Z三个方向的加速度）。我想在这个数组里添加3列，分别存放对应的球坐标（r、theta、phi）。我写的代码可以实现这个功能，但感觉运行得太慢了。我该怎么做才能提高效率呢？

import numpy as np
import math as m

def cart2sph(x,y,z):
    XsqPlusYsq = x**2 + y**2
    r = m.sqrt(XsqPlusYsq + z**2)               # r
    elev = m.atan2(z,m.sqrt(XsqPlusYsq))     # theta
    az = m.atan2(y,x)                           # phi
    return r, elev, az

def cart2sphA(pts):
    return np.array([cart2sph(x,y,z) for x,y,z in pts])

def appendSpherical(xyz):
    np.hstack((xyz, cart2sphA(xyz)))

Answer 1

注意！上面的代码还有错误，而且这个内容在谷歌搜索中排名很高。
简单来说：我用VPython测试过，发现用atan2来计算theta（高度角）是错的，应该用acos！对于phi（方位角）来说是正确的。
我推荐使用sympy1.0的acos函数（它甚至不会对acos(z/r)在r=0时发出警告）。

http://mathworld.wolfram.com/SphericalCoordinates.html

如果我们把这个转换成物理系统中的坐标（r, theta, phi）= (r, elev, azimuth)，我们可以得到：

r = sqrt(x*x + y*y + z*z)
phi = atan2(y,x)
theta = acos(z,r)

这段代码虽然没有经过优化，但对于右手坐标系来说是正确的：

from sympy import *
def asCartesian(rthetaphi):
    #takes list rthetaphi (single coord)
    r       = rthetaphi[0]
    theta   = rthetaphi[1]* pi/180 # to radian
    phi     = rthetaphi[2]* pi/180
    x = r * sin( theta ) * cos( phi )
    y = r * sin( theta ) * sin( phi )
    z = r * cos( theta )
    return [x,y,z]

def asSpherical(xyz):
    #takes list xyz (single coord)
    x       = xyz[0]
    y       = xyz[1]
    z       = xyz[2]
    r       =  sqrt(x*x + y*y + z*z)
    theta   =  acos(z/r)*180/ pi #to degrees
    phi     =  atan2(y,x)*180/ pi
    return [r,theta,phi]

你可以用类似下面的函数自己测试一下：

test = asCartesian(asSpherical([-2.13091326,-0.0058279,0.83697319]))

这里还有一些其他象限的测试数据：

[[ 0.          0.          0.        ]
 [-2.13091326 -0.0058279   0.83697319]
 [ 1.82172775  1.15959835  1.09232283]
 [ 1.47554111 -0.14483833 -1.80804324]
 [-1.13940573 -1.45129967 -1.30132008]
 [ 0.33530045 -1.47780466  1.6384716 ]
 [-0.51094007  1.80408573 -2.12652707]]

我额外使用了VPython来方便地可视化向量：

test   = v.arrow(pos = (0,0,0), axis = vis_ori_ALA , shaftwidth=0.05, color=v.color.red)

Answer 2

这里有一段我写的简单Cython代码：

cdef extern from "math.h":
    long double sqrt(long double xx)
    long double atan2(long double a, double b)

import numpy as np
cimport numpy as np
cimport cython

ctypedef np.float64_t DTYPE_t

@cython.boundscheck(False)
@cython.wraparound(False)
def appendSpherical(np.ndarray[DTYPE_t,ndim=2] xyz):
    cdef np.ndarray[DTYPE_t,ndim=2] pts = np.empty((xyz.shape[0],6))
    cdef long double XsqPlusYsq
    for i in xrange(xyz.shape[0]):
        pts[i,0] = xyz[i,0]
        pts[i,1] = xyz[i,1]
        pts[i,2] = xyz[i,2]
        XsqPlusYsq = xyz[i,0]**2 + xyz[i,1]**2
        pts[i,3] = sqrt(XsqPlusYsq + xyz[i,2]**2)
        pts[i,4] = atan2(xyz[i,2],sqrt(XsqPlusYsq))
        pts[i,5] = atan2(xyz[i,1],xyz[i,0])
    return pts

我用300万个数据点测试，运行时间从62.4秒减少到了1.22秒。这个效果还不错。我相信还有其他可以改进的地方。

Answer 3

这段内容和Justin Peel的回答有点像，不过这里只用了numpy，并且利用了它自带的向量化功能：

import numpy as np

def appendSpherical_np(xyz):
    ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
    xy = xyz[:,0]**2 + xyz[:,1]**2
    ptsnew[:,3] = np.sqrt(xy + xyz[:,2]**2)
    ptsnew[:,4] = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
    #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
    ptsnew[:,5] = np.arctan2(xyz[:,1], xyz[:,0])
    return ptsnew

需要注意的是，正如评论中提到的，我改变了你原来函数中对仰角的定义。在我的电脑上，用pts = np.random.rand(3000000, 3)进行测试时，时间从76秒减少到了3.3秒。我没有Cython，所以没法和那个方案的时间进行比较。