LMFIT EMCEE SPEED

[Sbongakonke Zondo]

Greetings,

I hope this message finds you well. I am currently using LMFIT EMCEE to find the confidence regions of my optimized KP spectral age model parameters (k0, B and T). The code is working fine but I am struggling when it comes to speeding it up. I was hoping that someone is ware of how I can get it to go faster. Changing burn from 300 to 0, thin from 20 to 1 and including nwalkers and workers equal to 16 and 8 respectively managed to reduce the time it takes for one step to complete to 6 minutes. However the recommended steps is 1000 and from my estimate this would take at least four days to complete. I would kindly appreciate any assistance.

This is the code I am using alongside the results.

# <examples/doc_fitting_emcee.py>

import numpy as np

import multiprocessing

import math as m

import lmfit

import corner

try:

    import matplotlib.pyplot as plt

    HASPYLAB = True

except ImportError:

    HASPYLAB = False

try:

    import corner

    HASCORNER = True

except ImportError:

    HASCORNER = False

xdata = np.array([74.0, 330.0, 1000, 1100, 1330, 1485], dtype=np.float32)

ydata = np.array([91.444, 16.802, 3.181, 2.682, 1.821, 1.494], dtype=np.float32)/1e8

yerr = np.array([9.14, 1.68, 0.018672, 0.018989, 0.009504, 0.010558], dtype=np.float32)/1e8

lnxdata = np.log10(xdata)

lnydata = np.log10(ydata)

lnyerr = (1/np.log(10)) * np.float32(yerr)/np.float32(ydata)

def fgrande(t):

    res=1.78*(t**0.3)*m.exp(-t)

    return res

def integra(x, y):

    n=x.size

    res=0.0

    for i in range(0, n-1):

        a=(x[i+1]-x[i])

        b=(y[i]+y[i+1])/2.

        res=res+a*b

    return res

def KP_radio_data(params, x, lnydata, lnyerr):

    k0 = params['k0'].value

    t = params['t'].value

    bmu = params['bmu'].value

#pro pachol_radio_data,k0,s,t,bmu,xdata,emiss

#;k0: normalizzazione iniziale

#;s: indice spettrale iniziale E^(-s)

#;t: tempo finale in Myr

#;bmu: campo magnetico in microG

#;xdata: asse delle frequenze in MHz

    tfin=t*1e6*3.156e7  #; in secondi

    s= 2.0

    gelm=1000

    assegel=np.zeros(gelm+1)

    for i in range(0, gelm):

        assegel[i]=10**(1.0+5.0*i/gelm)

    ngel0=np.zeros(gelm+1)

    for i in range(0, gelm):

        ngel0[i]=(assegel[i])**(-s)

    ngel=np.zeros(gelm+1)

    thm=1000

    asseth=np.zeros(thm+1)

    for j in range(0, thm+1):

        #asseth[j]=m.pi/18.0+(m.pi/2.0-m.pi/18.0)*j/thm

        asseth[j]=0.+(m.pi/2.)*j/thm

    z=0.055     #redshift

    intth=np.zeros(thm+1)

    assenu= 10**x

    num=assenu.size

    neterm=1e-6#  ;cm^-3

    r0=2.82e-13#    ;cm

    me=0.511e-3#    ;GeV

    c=3.0e10#     ;cm/s

    nu0=2.8*bmu*1e-6# ;MHz

    nup=8980.*neterm**0.5*1e-6#   ;MHz

    a=2.0*m.pi*m.sqrt(3.)*r0*me/c*1e6*nu0

    integel=np.zeros(gelm+1)

    emiss=np.zeros(num)

    for k in range(0, num):

        nu=assenu[k]

        for i in range(0, gelm):

            gel=assegel[i]

            x_val=2.*(nu/(3.*nu0*gel**2.))#*(1.+(gel*nup/nu)**2.)**(3./2.)

            for j in range(0, thm):

                theta=asseth[j]

                if (theta==0):

                    intth1=0.0

                else:

                    intth1=m.sin(theta)*fgrande(x_val/m.sin(theta))*m.sin(theta)*0.5

                bic=(1.37e-20)*(1.0+z)**4 # (1.0e-6) to delete

                bsyn=(1.94e-21)*(bmu*m.sin(theta))**2

                btot=bsyn+bic

                gelmax=1./(btot*tfin)

                if (gel<gelmax):

                    intth2=np.log10(k0)*(gel**(-s))*(1.0-btot*tfin*gel)**(s-2.0)

                else:

                    intth2=0.

                intth[j]=intth1*intth2

            integel[i]=a*integra(asseth,intth)

        emiss[k]=integra(assegel,integel)

    model_flux= np.log10(np.float32(emiss))  +  0.003718963823114596 + np.log10(1.821/1e8) + 28.31056415

    return (model_flux - lnydata) / lnyerr

params = lmfit.Parameters()

params.add('k0', value=10)

params.add('t', value=8.75)

params.add('bmu', value=37.13)

mi = lmfit.minimize(KP_radio_data, params, args=(lnxdata, lnydata, lnyerr), method='leastsq', nan_policy='omit')

lmfit.printfuncs.report_fit(mi.params, min_correl=0.5)

res = lmfit.minimize(KP_radio_data, method='emcee', nwalkers=16, args=(lnxdata, lnydata, lnyerr), nan_policy='omit', burn=0,

                     steps=15, thin=1, workers=8, params=mi.params, is_weighted=True,

                     progress=True)

if HASPYLAB and HASCORNER:

    emcee_corner = corner.corner(res.flatchain, labels=res.var_names,

                                 truths=list(res.params.valuesdict().values()))

    plt.show()

print('median of posterior probability distribution')

print('--------------------------------------------')

lmfit.report_fit(res.params)

highest_prob = np.argmax(res.lnprob)

hp_loc = np.unravel_index(highest_prob, res.lnprob.shape)

mle_soln = res.chain[hp_loc]

for i, par in enumerate(params):

    params[par].value = mle_soln[i]

print('\nMaximum Likelihood Estimation from emcee       ')

print('-------------------------------------------------')

print('Parameter  MLE Value   Median Value   Uncertainty')

fmt = '  {:5s}  {:11.5f} {:11.5f}   {:11.5f}'.format

for name, param in params.items():

    print(fmt(name, param.value, res.params[name].value,

              res.params[name].stderr))

print('\nError estimates from emcee:')

print('------------------------------------------------------')

print('Parameter  -2sigma  -1sigma   median  +1sigma  +2sigma')

for name in params.keys():

    quantiles = np.percentile(res.flatchain[name],

                              [2.275, 15.865, 50, 84.135, 97.275])

    median = quantiles[2]

    err_m2 = quantiles[0] - median

    err_m1 = quantiles[1] - median

    err_p1 = quantiles[3] - median

    err_p2 = quantiles[4] - median

    fmt = '  {:5s}   {:8.4f} {:8.4f} {:8.4f} {:8.4f} {:8.4f}'.format

    print(fmt(name, err_m2, err_m1, median, err_p1, err_p2))

[[Variables]] k0: 10.0089626 +/- 1.98115509 (19.79%) (init = 10) t: 8.74997171 +/- 0.29256361 (3.34%) (init = 8.75) bmu: 37.1072007 +/- 1.43361323 (3.86%) (init = 37.13) [[Correlations]] (unreported correlations are < 0.500) C(k0, t) = +0.9974 C(k0, bmu) = -0.9778 C(t, bmu) = -0.9717

100%|██████████| 15/15 [1:41:54<00:00, 407.63s/it]

The chain is shorter than 50 times the integrated autocorrelation time for 3 parameter(s). Use this estimate with caution and run a longer chain! N/50 = 0; tau: [1.45414783 1.35584182 1.48205188]

WARNING:root:Pandas support in corner is deprecated; use ArviZ directly

median of posterior probability distribution -------------------------------------------- [[Variables]] k0: 10.0098822 +/- 0.01490645 (0.15%) (init = 10.00896) t: 8.75023394 +/- 0.00561892 (0.06%) (init = 8.749972) bmu: 37.1058780 +/- 0.05500376 (0.15%) (init = 37.1072) [[Correlations]] (unreported correlations are < 0.100) C(k0, t) = +0.9399 C(k0, bmu) = -0.8924 C(t, bmu) = -0.8477 Maximum Likelihood Estimation from emcee ------------------------------------------------- Parameter MLE Value Median Value Uncertainty k0 10.00904 10.00988 0.01491 t 8.74996 8.75023 0.00562 bmu 37.10723 37.10588 0.05500 Error estimates from emcee: ------------------------------------------------------ Parameter -2sigma -1sigma median +1sigma +2sigma k0 -0.0387 -0.0120 10.0099 0.0178 0.0579 t -0.0223 -0.0030 8.7502 0.0083 0.0371 bmu -0.1646 -0.0568 37.1059 0.0532 0.1221

[Jeremy DeJournett]

Hi Sbongakonke,

Looking at the code, the obvious places are the for loops where you're performing matrix operations; encoding those as actual ndarray operations will allow you to take full advantage of the numpy backend.

For any sort of speed profiling, it's important to back up your approach to optimization with data, and the tool I find the most useful is flamegraphs.

Here's how I personally use them:

1. say your script is in my_script.py; change directory into the same directory as it

2. Activate your virtual environment and install snakeviz (one of many flamegraph visualizers, works with the native cProfile module): pip install snakeviz 

3. Run your program under cProfile to generate a profile file: python -m cProfile -o flamegraph.prof -m my_script

4. Run snakeviz on it: snakeviz flamegraph.prof

This will open a browser window with an interactive flamegraph. The longer the horizontal bar, the more time that function takes overall. You might see some interpreter functions at the outset, just click into your main function to focus your efforts on the things you can change. Within those bars you can start to piece together the "hottest" loops and start your refsctoring with them. It might help to move some loops into dedicated functions so you can parse out which loops are more time consuming.

Hope this helps, best of luck!

Jeremy 

--
You received this message because you are subscribed to the Google Groups "lmfit-py" group.
To unsubscribe from this group and stop receiving emails from it, send an email to lmfit-py+u...@googlegroups.com.
To view this discussion visit https://groups.google.com/d/msgid/lmfit-py/dcc01c48-8e78-4111-819d-d460bf904396n%40googlegroups.com.

[Jeremy DeJournett]

I'll add a comment on a useful technique for comparing different implementations you're thinking about: 

Let's say you want to test my prior comment about migration of specific loops to numpy native data type operations. We've identified a specific loop as being "hot" and consuming inordinate computation time (for me, this is any loop consuming more than 1%).

it might look something like:

for i in range(num):

   d[i] = a[i] * exp(b[i] + c[i])

You could rewrite this in numpy like so:

d = np.multiply(a, np.exp(np.add(b, c)))

If we put them in functions, hot_loop_simple and hot_loop_numpy, you could write a test function like so:

test_data = ...

for i in range(10000):

  if random.random()%2:

    hot_loop_simple(test_data)

  else:

    hot_loop_numpy(test_data)

Running a flamegraph on this it will be clear to see which version takes longest. You can generalize to more implementations with a simple dict approach, if you like.

[matt.n...@gmail.com]

Hi Zondo,

 

I agree with Jeremy’s suggestions. 

If runtime is a concern, then you really need to first convert loops to numpy ufuncs.   This is not really about lmfit or even emcee.  You need to speed up your `KP_radio_data` function, ideally getting rid of the double loop over array indices and replacing your `integra` with an integration function.     Like,

 

asseth=np.zeros(thm+1)

for j in range(0, thm+1):

    asseth[j]=0.+(m.pi/2.)*j/thm

 

could be

         asseth = m.pi/(2*thm) * np.arange(thm+1)

 

That’s the simplest case.  There are several more to look at.

 

I cannot recommend coercing data to float32.  That will almost certainly make the analysis worse and might actually be slower.   I would also recommend against using so many constants with different scales in your code.  And (especially if you do feel that you are somehow obligated to use float32), I would try to work in units so that the numerical values you are passing out of your function are more on the magnitude scale of 1e-7 to 1.e+7.    When you start going far outside these ranges, floating point precisions become strained when comparing small differences.

 

--Matt

median of posterior probability distribution -------------------------------------------- [[Variables]] k0: 10.0098822 +/- 0.01490645 (0.15%) (init = 10.00896) t: 8.75023394 +/- 0.00561892 (0.06%) (init = 8.749972) bmu: 37.1058780 +/- 0.05500376 (0.15%) (init = 37.1072) [[Correlations]] (unreported correlations are < 0.100) C(k0, t) = +0.9399 C(k0, bmu) = -0.8924 C(t, bmu) = -0.8477 Maximum Likelihood Estimation from emcee ------------------------------------------------- Parameter MLE Value Median Value Uncertainty k0 10.00904 10.00988 0.01491 t 8.74996 8.75023 0.00562 bmu 37.10723 37.10588 0.05500 Error estimates from emcee: ------------------------------------------------------ Parameter -2sigma -1sigma median +1sigma +2sigma k0 -0.0387 -0.0120 10.0099 0.0178 0.0579 t -0.0223 -0.0030 8.7502 0.0083 0.0371 bmu -0.1646 -0.0568 37.1059 0.0532 0.1221

 

 

 

 

--

[Sbongakonke Zondo]

Thank you so much for the assistance.