Lmfit and Numba

[Shankar Dutt]

Hi, 

I am trying to use lmfit with Numba compiled model. The model is as follows:

import numpy as np

from scipy import integrate

from scipy.special import j1

from scipy.special import j0

from numpy import sinEnter code here...

from numpy import cos

from numpy import pi

import time

import numba as nb

from numba import jit

from numba import cfunc

from numba.types import intc, CPointer, float64

from scipy import LowLevelCallable

from matplotlib import pyplot as plt

def HC_angular(q, amp, background, radius, pd, roughness, angle):

   #q = np.linspace(0.03, 1.0, 1000)

   #pd = 0.03

   #radius = 40

   #angle = 1

   length = 20000 #20000nm = 20um

   c = radius

   s = pd * c

   l = length / 2

   

   def jit_integrand_function(integrand_function):

       jitted_function = jit(integrand_function, nopython=True)

       @cfunc(float64(intc, CPointer(float64)), error_model="numpy", fastmath=True)

       def wrapped(n, xx):

           ar = nb.carray(xx, n)

           return jitted_function(ar[0], ar[1], ar[2])

       return LowLevelCallable(wrapped.ctypes)

    @jit_integrand_function

   def f(z,theta,q):

       def ff(theta, z, q):

           qz = q * sin(theta)

           qr = q * cos(theta)

           return (sin(qz * l) * z * j1(qr * z) / (qr * qz)) ** 2

        def g_distribution(z, radius, pd):

           return  np.exp(-0.5 * (((z - radius) / s) ** 2))

        return ff(theta,z,q)*g_distribution(z,radius,pd)* sin(theta)

    va = radius - 4*pd*radius

   if va<0:

       va = 0

   vb = radius + 4*pd*radius

   

   

   def lower_inner(z):

       return va

    def upper_inner(z):

       return vb

    y = np.empty(len(q))

   for n in range(len(q)):

       y[n] = integrate.dblquad(f, 0, angle*pi/180, lower_inner, upper_inner,args=(q[n],),epsabs=1e-2,epsrel=1e-2)[0]

   

   y = (16 * pi**2)*(1 / (np.sqrt(2 * pi) * s)) * 1e7 * (1 / (pi * radius ** 2 * length)) *  y

   

   return y * amp ** 2 * np.exp(-(roughness*q)**2)  + background

When I try to fit it using lmfit, I get the following error:

Traceback (most recent call last):

 File "SD_fit.py", line 512, in <module>

   fit(arguments)

 File "SD_fit.py", line 206, in fit

   fit_hc_a(q, r, p)

 File "SD_fit.py", line 461, in fit_hc_a

   result = minner.minimize()

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/lmfit/minimizer.py", line 2176, in minimize

   return function(**kwargs)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/lmfit/minimizer.py", line 1568, in leastsq

   lsout = scipy_leastsq(self.__residual, variables, **lskws)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/scipy/optimize/minpack.py", line 383, in leastsq

   shape, dtype = _check_func('leastsq', 'func', func, x0, args, n)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/scipy/optimize/minpack.py", line 26, in _check_func

   res = atleast_1d(thefunc(*((x0[:numinputs],) + args)))

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/lmfit/minimizer.py", line 530, in __residual

   out = self.userfcn(params, *self.userargs, **self.userkws)

 File "SD_fit.py", line 456, in fcn2mmin

   model = smdls.HC_angular(q, amp, background, radius, pd, roughness, angle)

 File "/home/shankar/SD_1D_SAXS/models/Special_Models.py", line 35, in HC_angular

   def f(z,theta,q):

 File "/home/shankar/SD_1D_SAXS/models/Special_Models.py", line 29, in jit_integrand_function

   def wrapped(n, xx):

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/decorators.py", line 259, in wrapper

   res.compile()

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler_lock.py", line 32, in _acquire_compile_lock

   return func(*args, **kwargs)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/ccallback.py", line 70, in compile

   cres = self._compile_uncached()

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/ccallback.py", line 84, in _compile_uncached

   cres = self._compiler.compile(sig.args, sig.return_type)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/dispatcher.py", line 81, in compile

   raise retval

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/dispatcher.py", line 91, in _compile_cached

   retval = self._compile_core(args, return_type)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/dispatcher.py", line 104, in _compile_core

   cres = compiler.compile_extra(self.targetdescr.typing_context,

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler.py", line 551, in compile_extra

   return pipeline.compile_extra(func)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler.py", line 331, in compile_extra

   return self._compile_bytecode()

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler.py", line 393, in _compile_bytecode

   return self._compile_core()

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler.py", line 373, in _compile_core

   raise e

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler.py", line 364, in _compile_core

   pm.run(self.state)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler_machinery.py", line 347, in run

   raise patched_exception

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler_machinery.py", line 338, in run

   self._runPass(idx, pass_inst, state)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler_lock.py", line 32, in _acquire_compile_lock

   return func(*args, **kwargs)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler_machinery.py", line 302, in _runPass

   mutated |= check(pss.run_pass, internal_state)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/compiler_machinery.py", line 275, in check

   mangled = func(compiler_state)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/typed_passes.py", line 89, in run_pass

   typemap, return_type, calltypes = type_inference_stage(

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/typed_passes.py", line 67, in type_inference_stage

   infer.propagate(raise_errors=raise_errors)

 File "/home/shankar/anaconda3/envs/Double_Int/lib/python3.8/site-packages/numba/typeinfer.py", line 985, in propagate

   raise errors[0]

numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)

Failed in nopython mode pipeline (step: nopython frontend)

Untyped global name 'radius': cannot determine Numba type of <class 'lmfit.parameter.Parameter'>

File "models/Special_Models.py", line 44:

       def g_distribution(z, radius, pd):

           <source elided>

        return ff(theta,z,q)*g_distribution(z,radius,pd)* sin(theta)

       ^

[1] During: resolving callee type: type(CPUDispatcher(<function HC_angular.<locals>.f at 0x7ff57d998790>))

[2] During: typing of call at /home/shankar/SD_1D_SAXS/models/Special_Models.py (31)

[3] During: resolving callee type: type(CPUDispatcher(<function HC_angular.<locals>.f at 0x7ff57d998790>))

[4] During: typing of call at /home/shankar/SD_1D_SAXS/models/Special_Models.py (31)

File "models/Special_Models.py", line 31:

       def wrapped(n, xx):

           <source elided>

           ar = nb.carray(xx, n)

           return jitted_function(ar[0], ar[1], ar[2])

           ^

Not sure if it is error of my code or I basically cannot use numba with lmfit? Please help

[Matt Newville]

Hi Shankar, 

I expect that using numba within an objective or model function would be challenging for two separate reasons:

a) for many scipy.optimize solvers (including the default `leastsq`), the objective function (python) is called from compiled code using Pythons C library.  I don't know if these can use "numba-compiled" code.  I can believe the answer is and always will be a simple No.  Just to be clear, that would be the case for scipy.optimize without lmfit.

b) lmfit Parameters are objects and can have constraints that use `asteval`, a python library that uses many corners of Python that seem to be allergic to attempts to parallelize or serialize or otherwise "reduce" or "compile" the code.  Could numba support Parameter objects even without asteval? I don't know.  Can it support Parameter objects with asteval?  I'll guess "probably not without real effort".

So, exploring "can lmfit objective/model functions use numba" seems like it would require some real work.   I can believe the answer will simply be No. 

It might be that one could numba-compile a function that does the "heavy lifting"  outside of your objective/model function and call that with carefully prepared values from your objective/model function.   I suspect that could be made to work, but I could also believe that there could be problems.

It seems very strange to me that you went to all the trouble to post some code and traceback but decided to not post the complete code, including deciding to withhold the code that actually sets up and executes your fit.  

--
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 on the web visit https://groups.google.com/d/msgid/lmfit-py/cf6d23f9-e494-443a-9ec4-8b354a49d8a1%40googlegroups.com.

--

--Matt Newville <newville at cars.uchicago.edu> 630-252-0431

[Konstantin Medvedovsky]

I just came across this, and wanted to first thank Matt for all the work on the lmifit library, but also to say that I've had no issues using numba inside an objective function with lmfit. I followed the approach suggested by Matt above, where I call essentially an Lmfit objective scoring function, and that scoring function calls a separate function which is numba-compiled. It works wonderfully and has generated the promised numba speedups of 200x+.

To unsubscribe from this group and stop receiving emails from it, send an email to lmfi...@googlegroups.com.

To view this discussion on the web visit https://groups.google.com/d/msgid/lmfit-py/cf6d23f9-e494-443a-9ec4-8b354a49d8a1%40googlegroups.com.

[Shankar Dutt]

Dear Konstantin,

Thanks a lot for your reply. Can you please explain in a bit detail how you achieved this? I am trying quite a bit but no success. It would be really helpful if you could share a code snippet. 

Cheers,

Shankar

[Matt Newville]

Konstantin, Shankar,

Thanks for the confirmation that calling a numba-compiled function *can* work. 

I wonder -- and perhaps you can confirm -- that it may still be that this works for some solvers (those in "pure python") and not for others.  Specifically, I could believe that it may not work for those solvers (including our default `leastsq` in `minpack` which Shankar's messages shows was being used) that are written in compiled code where the python objective function is called with Python's C API.  

Is it possible that Konstantin is using a pure python solver?

--Matt

[Konstantin Medvedovsky]

I usually work with differential evolution or basinhopping (via Nedler-Mead), so it's possible other solvers wouldn't work as well. I'm not sure if those are 'pure python' or not. That said, I just tried 'leastsq' as well, and it worked fine as well. I have put up a working code example of this implementation here: https://colab.research.google.com/drive/1JTmiq4EsXvydV1eUJiVE0imLfdSbLYlB?usp=sharing

As you can see, I have a numba function which iterates though a loop, and then a wrapper function for lmfit to operate in. In this example (100,000 rows), you get roughly a 150-200x speedup with numba. Apologies for the clunkiness of the code - I am not a natural software developer.

It's open to view and modify, so feel free to experiment with other solvers, or modify as you see fit.

Hope this helps!

[Shankar Dutt]

Dear Konstantin, Matt,

Thanks a lot for your replies. I am so much thankful for having such a supportive community. After looking at Konstantin's code, I went back to basics and defined every part of numba implementation differently and increased the complexity step by step. I found out that in general, numba works completely fine with lmfit(which is a piece of great news). But to make the code more efficient,I was using intc, CPointer, float64 from numba libraries, which doesn't work well with lmfit. As I am a beginner in python, I am not sure why! 

I again want to thank you both for the help. Atleast my code is 170x faster(which is a big improvement for me).

Cheers,

Shankar