Faster numpy cartesian to spherical coordinate conversion?

前端 未结 5 1584
南旧
南旧 2020-12-04 15:57

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, the

5条回答
  •  旧巷少年郎
    2020-12-04 16:20

    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
    

提交回复
热议问题