Numba - a compiler for Python

Numba is a Python compiler with good support of the Numpy library.

To speed-up linear algebra with 3D vectors I wrote the linalg_3d.py module:

   1 # -*- coding: utf-8 -*-
   2 """ Fast implementation of the following linear algebra functions for 3-dimensional arrays:
   3     cross, dot, norm, normalize.
   4     Average speed-up factor compared to numpy around 8 to 9, heavily machine dependant,
   5     but also dependent on the call context: If these procedures are called from another
   6     procedure, that was also compiled with Numba they can be much faster than when called
   7     from standard Python code. """
   8 from numba import autojit, double
   9 import math
  10 import numpy as np
  11 import numpy.linalg as la
  12 from Timer import Timer
  13 
  14 @autojit
  15 def cross(vec1, vec2):
  16     """ Calculate the dot product of two 3d vectors. """
  17     a1, a2, a3 = double(vec1[0]), double(vec1[1]), double(vec1[2])
  18     b1, b2, b3 = double(vec2[0]), double(vec2[1]), double(vec2[2])
  19     result = np.zeros(3)
  20     result[0] = a2 * b3 - a3 * b2
  21     result[1] = a3 * b1 - a1 * b3
  22     result[2] = a1 * b2 - a2 * b1
  23     return result
  24 
  25 @autojit
  26 def dot(vec1, vec2):
  27     """ Calculate the dot product of two 3d vectors. """
  28     return vec1[0] * vec2[0] + vec1[1] * vec2[1] + vec1[2] * vec2[2]
  29 
  30 @autojit
  31 def norm(vec):
  32     """ Calculate the norm of a 3d vector. """
  33     return math.sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
  34 
  35 @autojit
  36 def normalize(vec):
  37     """ Calculate the normalized vector (norm: one). """
  38     return vec / norm(vec)
  39 
  40 def init():
  41     """ call all functions once to compile them """
  42     vec1, vec2 = np.array((1.0, 2.0, 3.0)), np.array((2.0, 3.0, 4.0))
  43     dot(vec1, vec2)
  44     cross(vec1, vec2)
  45     norm(vec1)
  46     unit_vector = normalize(vec1)
  47     unit_vector *= 2 # suppress warning unused variable
  48 
  49 init()
  50 
  51 if __name__ == '__main__':
  52     vec1 = np.array((1.0, 2.0, 3.0))
  53     vec2 = np.array((2.0, 3.0, 4.0))
  54     assert dot(vec1, vec2) == np.dot(vec1, vec2)
  55     assert (cross(vec1, vec2) == np.cross(vec1, vec2)).all()
  56     with Timer() as t0:
  57         for i in range(10000):
  58             pass
  59     print "time for empty loop ", t0.secs
  60     print
  61     with Timer() as t1:
  62         for i in range(10000):
  63             norm(vec1)
  64     print "time for numba norm [µs]:     ", (t1.secs-t0.secs)  / 10000 * 1e6
  65     with Timer() as t2:
  66         for i in range(10000):
  67             la.norm(vec1)
  68     print "time for linalg norm [µs]:   ", (t2.secs - t0.secs)  / 10000 * 1e6
  69     print "speedup of norm with numba:  ", (t2.secs - t0.secs)  / (t1.secs - t0.secs)
  70     print
  71     with Timer() as t3:
  72         for i in range(10000):
  73             dot(vec1, vec2)
  74     print "time for numba dot [µs]:      ", (t3.secs - t0.secs) / 10000 * 1e6
  75     with Timer() as t4:
  76         for i in range(10000):
  77             np.dot(vec1, vec2)
  78     print "time for numpy dot [µs]:      ", (t4.secs - t0.secs) / 10000 * 1e6
  79     print "speedup of dot with numba:    ", (t4.secs - t0.secs) / (t3.secs - t0.secs)
  80     print
  81     with Timer() as t5:
  82         for i in range(10000):
  83             cross(vec1, vec2)
  84     print "time for numba cross [µs]:    ", (t5.secs - t0.secs) / 10000 * 1e6
  85     with Timer() as t6:
  86         for i in range(10000):
  87             np.cross(vec1, vec2)
  88     print "time for numpy cross [µs]:   ", (t6.secs - t0.secs) / 10000 * 1e6
  89     print "speedup of cross with numba:  ", (t6.secs - t0.secs) / (t5.secs - t0.secs)
  90 
  91 """
  92 Result on Core duo, 2 Ghz:
  93 
  94 time for empty loop  0.000799894332886
  95 
  96 time for numba norm [µs]:      1.14262104034
  97 time for linalg norm [µs]:    20.1265096664
  98 speedup of norm with numba:   17.6143348983
  99 
 100 time for numba dot [µs]:       1.12090110779
 101 time for numpy dot [µs]:       1.58061981201
 102 speedup of dot with numba:     1.41013315183
 103 
 104 time for numba cross [µs]:     3.1898021698
 105 time for numpy cross [µs]:    26.0297060013
 106 speedup of cross with numba:   8.16028851185
 107 """

You will also need the Timer.py module:

   1 # -*- coding: utf-8 -*-
   2 import time
   3 
   4 class Timer(object):
   5     def __init__(self, verbose=False):
   6         self.verbose = verbose
   7 
   8     def __enter__(self):
   9         self.start = time.time()
  10         return self
  11 
  12     def __exit__(self, *args):
  13         self.end = time.time()
  14         self.secs = self.end - self.start
  15         self.msecs = self.secs * 1000  # millisecs
  16         if self.verbose:
  17             print 'elapsed time: %f ms' % self.msecs
  18 
  19 if __name__ == "__main__":
  20     with Timer() as t:
  21         # Code of which the execution time shall be measured
  22         time.sleep(0.1)
  23     print "\n---> Elapsed time: %s s" % t.secs


Back to FrontPage

KielTech: Numba (last edited 2013-06-17 20:43:54 by UweFechner)