Merge pull request geodynamics#222 from tjhei/performance_2

- jit some functions
- add misc/performance.py with several checks

burnman/eos/debye.py

@@ -111,6 +111,7 @@ def debye_fn_cheb(x):
         return ((val_infinity/x)/x)/x;
 
 
+@jit
 def thermal_energy(T, debye_T, n):
     """
     calculate the thermal energy of a substance.  Takes the temperature,
@@ -122,6 +123,7 @@ def thermal_energy(T, debye_T, n):
     E_th = 3.*n*constants.gas_constant*T * debye_fn_cheb(debye_T/T)
     return E_th
 
+@jit
 def heat_capacity_v(T,debye_T,n):
     """
     Heat capacity at constant volume.  In J/K/mol
@@ -132,6 +134,7 @@ def heat_capacity_v(T,debye_T,n):
     C_v = 3.0*n*constants.gas_constant* ( 4.0*debye_fn_cheb(x) - 3.0*x/(np.exp(x)-1.0) )
     return C_v
 
+@jit
 def helmholtz_free_energy(T, debye_T, n):
     """
     Helmholtz free energy of lattice vibrations in the Debye model.

burnman/eos/slb.py

@@ -7,12 +7,45 @@
 import scipy.optimize as opt
 import warnings
 
+# 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
+
+
 from . import birch_murnaghan as bm
 from . import debye
 from . import equation_of_state as eos
 from ..tools import bracket
 
 
+@jit
+def _grueneisen_parameter_fast(V_0, volume, gruen_0, q_0):
+    """global function with plain parameters so jit will work"""
+    x = V_0 / volume
+    f = 1./2. * (pow(x, 2./3.) - 1.)
+    a1_ii = 6. * gruen_0 # EQ 47
+    a2_iikk = -12.*gruen_0 + 36.*gruen_0*gruen_0 - 18.*q_0*gruen_0 # EQ 47
+    nu_o_nu0_sq = 1.+ a1_ii*f + (1./2.)*a2_iikk * f*f # EQ 41
+    return 1./6./nu_o_nu0_sq * (2.*f+1.) * ( a1_ii + a2_iikk*f )
+
+
+@jit
+def _delta_pressure(x,pressure,temperature,V_0,T_0,Debye_0,n,a1_ii,a2_iikk,b_iikk,b_iikkmm):
+
+    f= 0.5*(pow(V_0/x,2./3.)-1.)
+    debye_temperature =  Debye_0 * np.sqrt(1. + a1_ii * f + 1./2. * a2_iikk*f*f)
+    E_th =  debye.thermal_energy(temperature, debye_temperature, n) #thermal energy at temperature T
+    E_th_ref = debye.thermal_energy(T_0, debye_temperature, n) #thermal energy at reference temperature
+    nu_o_nu0_sq = 1.+ a1_ii*f + (1./2.)*a2_iikk * f*f # EQ 41
+    gr = 1./6./nu_o_nu0_sq * (2.*f+1.) * ( a1_ii + a2_iikk*f )
+
+    return (1./3.)*(pow(1.+2.*f,5./2.))*((b_iikk*f)+(0.5*b_iikkmm*f*f)) \
+            + gr*(E_th - E_th_ref)/x - pressure #EQ 21
 
 class SLBBase(eos.EquationOfState):
     """
@@ -71,17 +104,7 @@ def _thermal_pressure(self,T,V, params):
         P_th = gr * debye.thermal_energy(T,Debye_T, params['n'])/V
         return P_th
 
-    def _delta_pressure(self,x,pressure,temperature,V_0,T_0,Debye_0,n,a1_ii,a2_iikk,b_iikk,b_iikkmm):
-
-        f= 0.5*(pow(V_0/x,2./3.)-1.)
-        debye_temperature =  Debye_0 * np.sqrt(1. + a1_ii * f + 1./2. * a2_iikk*f*f)
-        E_th =  debye.thermal_energy(temperature, debye_temperature, n) #thermal energy at temperature T
-        E_th_ref = debye.thermal_energy(T_0, debye_temperature, n) #thermal energy at reference temperature
-        nu_o_nu0_sq = 1.+ a1_ii*f + (1./2.)*a2_iikk * f*f # EQ 41
-        gr = 1./6./nu_o_nu0_sq * (2.*f+1.) * ( a1_ii + a2_iikk*f )
-
-        return (1./3.)*(pow(1.+2.*f,5./2.))*((b_iikk*f)+(0.5*b_iikkmm*f*f)) \
-                + gr*(E_th - E_th_ref)/x - pressure #EQ 21
+
 
     def volume(self, pressure, temperature, params):
         """
@@ -102,10 +125,10 @@ def volume(self, pressure, temperature, params):
         # conservative guess:
         args = (pressure,temperature,V_0,T_0,Debye_0,n,a1_ii,a2_iikk,b_iikk,b_iikkmm)
         try:
-            sol = bracket(self._delta_pressure, params['V_0'], 1.e-2*params['V_0'], args)
+            sol = bracket(_delta_pressure, params['V_0'], 1.e-2*params['V_0'], args)
         except:
             raise Exception('Cannot find a volume, perhaps you are outside of the range of validity for the equation of state?')
-        return opt.brentq(self._delta_pressure, sol[0], sol[1] ,args=args)
+        return opt.brentq(_delta_pressure, sol[0], sol[1] ,args=args)
 
     def pressure( self, temperature, volume, params):
         """
@@ -124,18 +147,11 @@ def pressure( self, temperature, volume, params):
 
         return P
 
-
     def grueneisen_parameter(self, pressure, temperature, volume, params):
         """
         Returns grueneisen parameter :math:`[unitless]` 
         """
-        x = params['V_0'] / volume
-        f = 1./2. * (pow(x, 2./3.) - 1.)
-        gruen_0 = params['grueneisen_0']
-        a1_ii = 6. * gruen_0 # EQ 47
-        a2_iikk = -12.*gruen_0 + 36.*gruen_0*gruen_0 - 18.*params['q_0']*gruen_0 # EQ 47
-        nu_o_nu0_sq = 1.+ a1_ii*f + (1./2.)*a2_iikk * f*f # EQ 41
-        return 1./6./nu_o_nu0_sq * (2.*f+1.) * ( a1_ii + a2_iikk*f )
+        return _grueneisen_parameter_fast(params['V_0'], volume, params['grueneisen_0'], params['q_0'])
 
     def isothermal_bulk_modulus(self, pressure,temperature, volume, params):
         """

misc/performance.py

@@ -0,0 +1,75 @@
+from __future__ import absolute_import
+from __future__ import print_function
+# This file is part of BurnMan - a thermoelastic and thermodynamic toolkit for the Earth and Planetary Sciences
+# Copyright (C) 2012 - 2015 by the BurnMan team, released under the GNU GPL v2 or later.
+
+import os, sys, numpy as np, matplotlib.pyplot as plt
+#hack to allow scripts to be placed in subdirectories next to burnman:
+if not os.path.exists('burnman') and os.path.exists('../burnman'):
+    sys.path.insert(1,os.path.abspath('..'))
+
+import burnman
+from burnman import minerals
+
+
+
+
+import timeit
+
+if True:
+    def test_grueneisen():
+        for aa in range(0,500):
+            a = burnman.eos.SLB3()
+            m = minerals.SLB_2011.fayalite()
+            x = 0
+            for i in range(0,10000):
+                x += a.grueneisen_parameter(1e9,1e4,1.1,m.params)
+        return x
+
+    test_grueneisen()
+    #r=timeit.timeit('__main__.test_grueneisen()', setup="import __main__", number=3)
+    #print("test_grueneisen", r)
+
+if True:
+    seismic_model = burnman.seismic.PREM() # pick from .prem() .slow() .fast() (see code/seismic.py)
+    number_of_points = 10 #set on how many depth slices the computations should be done
+    depths = np.linspace(1000e3,2500e3, number_of_points)
+    seis_p, seis_rho, seis_vp, seis_vs, seis_vphi = seismic_model.evaluate(['pressure','density','v_p','v_s','v_phi'],depths)
+
+
+    temperature = burnman.geotherm.brown_shankland(seis_p)
+
+    def calc_velocities(a,b,c):
+        amount_perovskite = a
+        pv = minerals.SLB_2011.mg_fe_perovskite([b, 1.0-b, 0.0])
+        fp = minerals.SLB_2011.ferropericlase([c,1.0-c])
+        rock = burnman.Composite([pv, fp], [amount_perovskite, 1.0-amount_perovskite])
+
+
+        mat_rho, mat_vp, mat_vs = rock.evaluate(['density','v_phi','v_s'], seis_p,temperature)
+        return mat_vp, mat_vs, mat_rho
+
+    def error(a,b,c):
+        mat_vp, mat_vs, mat_rho = calc_velocities(a,b,c)
+
+        vs_err = burnman.l2(depths, mat_vs, seis_vs) /1e9
+        vp_err = burnman.l2(depths, mat_vp, seis_vp) /1e9
+        den_err = burnman.l2(depths, mat_rho, seis_rho) /1e9
+
+        #print vs_err, vp_err, den_err
+
+        return vs_err + vp_err + den_err
+
+    def run():
+        n=5
+        m=1e10
+        for a in np.linspace(0.0, 1.0, n):
+            for b in np.linspace(0.0, 1.0, n):
+                for c in np.linspace(0.0, 1.0, n):
+                    m=min(m,error(a,b,c))
+        print(m)
+        return m
+
+    run()
+    #r=timeit.timeit('__main__.run()', setup="import __main__", number=3)
+    #print("test_seismic_velocities", r)