Faster numpy cartesian to spherical coordinate conversion?
Question:
I have an array of 3 million data points from a 3-axiz accellerometer (XYZ), and I want to add 3 columns to the array containing the equivalent spherical coordinates (r, theta, phi). The following code works, but seems way too slow. How can I do better?
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)))
Answers:
Here’s a quick Cython code that I wrote up for this:
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
It took the time down from 62.4 seconds to 1.22 seconds using 3,000,000 points for me. That’s not too shabby. I’m sure there are some other improvements that can be made.
This is similar to Justin Peel‘s answer, but using just numpy and taking advantage of its built-in vectorization:
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
Note that, as suggested in the comments, I’ve changed the definition of elevation angle from your original function. On my machine, testing with pts = np.random.rand(3000000, 3), the time went from 76 seconds to 3.3 seconds. I don’t have Cython so I wasn’t able to compare the timing with that solution.
To complete the previous answers, here is a Numexpr implementation (with a possible fallback to Numpy),
import numpy as np
from numpy import arctan2, sqrt
import numexpr as ne
def cart2sph(x,y,z, ceval=ne.evaluate):
""" x, y, z : ndarray coordinates
ceval: backend to use:
- eval : pure Numpy
- numexpr.evaluate: Numexpr """
azimuth = ceval('arctan2(y,x)')
xy2 = ceval('x**2 + y**2')
elevation = ceval('arctan2(z, sqrt(xy2))')
r = eval('sqrt(xy2 + z**2)')
return azimuth, elevation, r
For large array sizes, this allows a factor of 2 speed up compared to pure a Numpy implementation, and would be comparable to C or Cython speeds. The present numpy solution (when used with the ceval=eval argument) is also 25% faster than the appendSpherical_np function in the @mtrw answer for large array sizes,
In [1]: xyz = np.random.rand(3000000,3)
...: x,y,z = xyz.T
In [2]: %timeit -n 1 appendSpherical_np(xyz)
1 loops, best of 3: 397 ms per loop
In [3]: %timeit -n 1 cart2sph(x,y,z, ceval=eval)
1 loops, best of 3: 280 ms per loop
In [4]: %timeit -n 1 cart2sph(x,y,z, ceval=ne.evaluate)
1 loops, best of 3: 145 ms per loop
although for smaller sizes, appendSpherical_np is actually faster,
In [5]: xyz = np.random.rand(3000,3)
...: x,y,z = xyz.T
In [6]: %timeit -n 1 appendSpherical_np(xyz)
1 loops, best of 3: 206 µs per loop
In [7]: %timeit -n 1 cart2sph(x,y,z, ceval=eval)
1 loops, best of 3: 261 µs per loop
In [8]: %timeit -n 1 cart2sph(x,y,z, ceval=ne.evaluate)
1 loops, best of 3: 271 µs per loop
! There is an error still in all the code above.. and this is a top Google result..
TLDR:
I have tested this with VPython, using atan2 for theta (elev) is wrong, use
acos! It is correct for phi (azim).
I recommend the sympy1.0 acos function (it does not even complain about acos(z/r) with r = 0 ) .
http://mathworld.wolfram.com/SphericalCoordinates.html
If we convert that to the physics system (r, theta, phi) = (r, elev, azimuth) we have:
r = sqrt(x*x + y*y + z*z)
phi = atan2(y,x)
theta = acos(z,r)
Non optimized but correct code for right-handed physics system:
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]
you can test it yourself with a function like:
test = asCartesian(asSpherical([-2.13091326,-0.0058279,0.83697319]))
some other test data for some quadrants:
[[ 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]]
I used VPython additionally to easily visualize vectors:
test = v.arrow(pos = (0,0,0), axis = vis_ori_ALA , shaftwidth=0.05, color=v.color.red)
Octave has some built-in functionality for coordinate transformations that can be accessed with the package oct2py to convert numpy arrays in Cartesian coordinates to spherical or polar coordinates (and back):
from oct2py import octave
xyz = np.random.rand(3000000,3)
%timeit thetaphir = octave.cart2sph(xyz)
724 ms ± 206 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
I base my code on mrtw’s answer – I added axis support, the reverse function and base it off a 3-tuple shape:
def from_xyz(xyz, axis=-1):
x, y, z = np.moveaxis(xyz, axis, 0)
lea = np.empty_like(xyz)
pre_selector = ((slice(None),) * lea.ndim)[:axis]
xy_sq = x ** 2 + y ** 2
lea[(*pre_selector, 0)] = np.sqrt(xy_sq + z ** 2)
lea[(*pre_selector, 1)] = np.arctan2(np.sqrt(xy_sq), z)
lea[(*pre_selector, 2)] = np.arctan2(y, x)
return lea
def to_xyz(lea, axis=-1):
l, e, a = np.moveaxis(lea, axis, 0)
xyz = np.empty_like(lea)
pre_selector = ((slice(None),) * xyz.ndim)[:axis]
xyz[(*pre_selector, 0)] = l * np.sin(e) * np.cos(a)
xyz[(*pre_selector, 1)] = l * np.sin(e) * np.sin(a)
xyz[(*pre_selector, 2)] = l * np.cos(e)
return xyz
I have an array of 3 million data points from a 3-axiz accellerometer (XYZ), and I want to add 3 columns to the array containing the equivalent spherical coordinates (r, theta, phi). The following code works, but seems way too slow. How can I do better?
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)))
Here’s a quick Cython code that I wrote up for this:
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
It took the time down from 62.4 seconds to 1.22 seconds using 3,000,000 points for me. That’s not too shabby. I’m sure there are some other improvements that can be made.
This is similar to Justin Peel‘s answer, but using just numpy and taking advantage of its built-in vectorization:
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
Note that, as suggested in the comments, I’ve changed the definition of elevation angle from your original function. On my machine, testing with pts = np.random.rand(3000000, 3), the time went from 76 seconds to 3.3 seconds. I don’t have Cython so I wasn’t able to compare the timing with that solution.
To complete the previous answers, here is a Numexpr implementation (with a possible fallback to Numpy),
import numpy as np
from numpy import arctan2, sqrt
import numexpr as ne
def cart2sph(x,y,z, ceval=ne.evaluate):
""" x, y, z : ndarray coordinates
ceval: backend to use:
- eval : pure Numpy
- numexpr.evaluate: Numexpr """
azimuth = ceval('arctan2(y,x)')
xy2 = ceval('x**2 + y**2')
elevation = ceval('arctan2(z, sqrt(xy2))')
r = eval('sqrt(xy2 + z**2)')
return azimuth, elevation, r
For large array sizes, this allows a factor of 2 speed up compared to pure a Numpy implementation, and would be comparable to C or Cython speeds. The present numpy solution (when used with the ceval=eval argument) is also 25% faster than the appendSpherical_np function in the @mtrw answer for large array sizes,
In [1]: xyz = np.random.rand(3000000,3)
...: x,y,z = xyz.T
In [2]: %timeit -n 1 appendSpherical_np(xyz)
1 loops, best of 3: 397 ms per loop
In [3]: %timeit -n 1 cart2sph(x,y,z, ceval=eval)
1 loops, best of 3: 280 ms per loop
In [4]: %timeit -n 1 cart2sph(x,y,z, ceval=ne.evaluate)
1 loops, best of 3: 145 ms per loop
although for smaller sizes, appendSpherical_np is actually faster,
In [5]: xyz = np.random.rand(3000,3)
...: x,y,z = xyz.T
In [6]: %timeit -n 1 appendSpherical_np(xyz)
1 loops, best of 3: 206 µs per loop
In [7]: %timeit -n 1 cart2sph(x,y,z, ceval=eval)
1 loops, best of 3: 261 µs per loop
In [8]: %timeit -n 1 cart2sph(x,y,z, ceval=ne.evaluate)
1 loops, best of 3: 271 µs per loop
! There is an error still in all the code above.. and this is a top Google result..
TLDR:
I have tested this with VPython, using atan2 for theta (elev) is wrong, use
acos! It is correct for phi (azim).
I recommend the sympy1.0 acos function (it does not even complain about acos(z/r) with r = 0 ) .
http://mathworld.wolfram.com/SphericalCoordinates.html
If we convert that to the physics system (r, theta, phi) = (r, elev, azimuth) we have:
r = sqrt(x*x + y*y + z*z)
phi = atan2(y,x)
theta = acos(z,r)
Non optimized but correct code for right-handed physics system:
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]
you can test it yourself with a function like:
test = asCartesian(asSpherical([-2.13091326,-0.0058279,0.83697319]))
some other test data for some quadrants:
[[ 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]]
I used VPython additionally to easily visualize vectors:
test = v.arrow(pos = (0,0,0), axis = vis_ori_ALA , shaftwidth=0.05, color=v.color.red)
Octave has some built-in functionality for coordinate transformations that can be accessed with the package oct2py to convert numpy arrays in Cartesian coordinates to spherical or polar coordinates (and back):
from oct2py import octave
xyz = np.random.rand(3000000,3)
%timeit thetaphir = octave.cart2sph(xyz)
724 ms ± 206 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
I base my code on mrtw’s answer – I added axis support, the reverse function and base it off a 3-tuple shape:
def from_xyz(xyz, axis=-1):
x, y, z = np.moveaxis(xyz, axis, 0)
lea = np.empty_like(xyz)
pre_selector = ((slice(None),) * lea.ndim)[:axis]
xy_sq = x ** 2 + y ** 2
lea[(*pre_selector, 0)] = np.sqrt(xy_sq + z ** 2)
lea[(*pre_selector, 1)] = np.arctan2(np.sqrt(xy_sq), z)
lea[(*pre_selector, 2)] = np.arctan2(y, x)
return lea
def to_xyz(lea, axis=-1):
l, e, a = np.moveaxis(lea, axis, 0)
xyz = np.empty_like(lea)
pre_selector = ((slice(None),) * xyz.ndim)[:axis]
xyz[(*pre_selector, 0)] = l * np.sin(e) * np.cos(a)
xyz[(*pre_selector, 1)] = l * np.sin(e) * np.sin(a)
xyz[(*pre_selector, 2)] = l * np.cos(e)
return xyz