[cig-commits] [commit] master: Test for numba performance improvements (7c318a9)

[cig_noreply at geodynamics.org]

Repository : https://github.com/geodynamics/burnman

On branch  : master
Link       : https://github.com/geodynamics/burnman/compare/8d345bc80879c34ccb9b2fad2290c947ac5c942a...e5f8442ac9f68970dec1854b19853d3a00292db7

>---------------------------------------------------------------

commit 7c318a9e9f0507917e6d13a1b0b82bce76ad492c
Author: Ian Rose <ian.r.rose at gmail.com>
Date:   Wed Oct 22 13:50:44 2014 -0700

    Test for numba performance improvements


>---------------------------------------------------------------

7c318a9e9f0507917e6d13a1b0b82bce76ad492c
 burnman/debye.py | 48 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/burnman/debye.py b/burnman/debye.py
index c0f6452..3df9290 100644
--- a/burnman/debye.py
+++ b/burnman/debye.py
@@ -3,8 +3,17 @@
 # Released under GPL v2 or later.
 
 import numpy as np
+
+# Try to import the jit from numba.  If it is
+# not available, just go with the standard 
+# python interpreter
+try:
+    from numba import jit
+except ImportError:
+    def jit(fn):
+        return fn
+
 import scipy.integrate as integrate
-from numpy.polynomial.chebyshev import Chebyshev
 
 import constants
 
@@ -15,6 +24,35 @@ Birch-Murnaghan to get a full EOS
 """
 
 
+chebyshev_representation = np.array( [ 2.707737068327440945/2.0, 0.340068135211091751, -0.12945150184440869e-01, \
+                                     0.7963755380173816e-03, -0.546360009590824e-04, 0.39243019598805e-05, \
+                                    -0.2894032823539e-06, 0.217317613962e-07, -0.16542099950e-08, \
+                                     0.1272796189e-09, -0.987963460e-11, 0.7725074e-12, -0.607797e-13, \
+                                     0.48076e-14, -0.3820e-15, 0.305e-16, -0.24e-17] )
+ at jit
+def _chebval(x, c):
+    """
+    Evaluate a Chebyshev series at points x.
+    This is just a lightly modified copy/paste job from the numpy
+    implementation of the same function, copied over here to put a
+    jit wrapper around it.
+    """
+    if len(c) == 1 :
+        c0 = c[0]
+        c1 = 0
+    elif len(c) == 2 :
+        c0 = c[0]
+        c1 = c[1]
+    else :
+        x2 = 2*x
+        c0 = c[-2]
+        c1 = c[-1]
+        for i in range(3, len(c) + 1) :
+            tmp = c0
+            c0 = c[-i] - c1
+            c1 = tmp + c1*x2
+    return c0 + c1*x
+
 
 def debye_fn(x):
     """
@@ -26,15 +64,11 @@ def debye_fn(x):
 
 
 
-chebyshev_representation = Chebyshev( [ 2.707737068327440945/2.0, 0.340068135211091751, -0.12945150184440869e-01, \
-                                     0.7963755380173816e-03, -0.546360009590824e-04, 0.39243019598805e-05, \
-                                    -0.2894032823539e-06, 0.217317613962e-07, -0.16542099950e-08, \
-                                     0.1272796189e-09, -0.987963460e-11, 0.7725074e-12, -0.607797e-13, \
-                                     0.48076e-14, -0.3820e-15, 0.305e-16, -0.24e-17] )
 eps = np.finfo(np.float).eps
 sqrt_eps = np.sqrt(np.finfo(np.float).eps)
 log_eps = np.log(np.finfo(np.float).eps)
 
+ at jit
 def debye_fn_cheb(x):
     """
     Evaluate the Debye function using a Chebyshev series expansion coupled with
@@ -51,7 +85,7 @@ def debye_fn_cheb(x):
         return 1.0 - 3.0*x/8.0 + x*x/20.0;
     elif x <= 4.0 :
         t = x*x/8.0 - 1.0;
-        c = chebyshev_representation(t)
+        c = _chebval(t, chebyshev_representation)
         return c - 0.375*x;
     elif x < -(np.log(2.0) + log_eps ):
         nexp = int(np.floor(xcut/x));