Monday, April 22, 2013

Numba -- It's like Python on steroids ... on steroids.

I'm teaching a course in molecular simulations this year in which the students have to program simulations of various simple systems. We're trying to keep things as simple as possible and everything is done in Python.

The task of the second week of the course is to program a simple simulation of a 2D Lennard-Jones gas. The computationally intensive step in this exercise is calculating the energy and gradient. In the past we've (and by we, I mean +Casper Steinmann) written a FORTRAN module and used F2PY to exploit the awesomeness of the good old FORmula TRANslator and have the students link their Python scripts with a precompiled .so module for this to run at an acceptable speed

This solution is of course completely fine, but not very general. I've been looking into Cython, pypy and various other methods to speed things up a bit. But they all seemed quite elaborate an unPythonic, and you might as well do the whole thing in C then.

Yesterday Numba was brought to my attention on this blog. In short, Numba is a way to compile (at execution time) Python functions into C which normally makes your program run substantially faster. The blog has a very nice example where the code gets a 1000x speedup just by adding one function decorator. And the syntax couldn't be simpler.

Let's try Numba on our Lennard-Jones program. Here is the naive function they have to implement:

The code is then executed via Python as one would normally execute Python code, specifying in the script that the calc_energy_and_gradient() function is to be compiled via Numba. This can be done simply by writing the the @autojit decorator before the function definition. 

You can also specify exactly which type of input the just-in-time compilation takes takes (an example is commented out). This makes the compilation of the function be type specific, but a bit faster.  Make sure to import numba and the specific type of argument (here the arguments are of the type double) and the relevant decorators (jit or autojit).

Conclusion: 10000 gradient and energy evaluations take 16 seconds with the standard Python implementation. 2 seconds with @autojit and 1.5 seconds with the explicit @jit. Not too bad for something that doesn't take ANY code changes.




8 comments:

  1. Do you know the performance using native f2py?

    ReplyDelete
  2. Did the timings between the Fortran module and the Numba implementation. I had to rewrite the entire thing to not use classes. The implementation of the gradient in python works in exactly the same way as the Fortran code, so no tricks are used to speed up the Python gradient code.

    The potential energies er numerically identical and the gradient norms have a difference of around 10^-16.

    Timing for 10 particles and 108 particles, respectively. 20,000 iteration steps

    Numba: 0.56 sec; 3.31 sec
    Fortran: 0.76 sec; 5.31 sec

    Python runs faster than Fortran? O_O

    I will make a blawg pawst when I get a bit more experience with Numba

    ReplyDelete
  3. Your blog provided us with valuable information to work with. Each & every tips of your post are awesome. Thanks a lot for sharing. Keep blogging, bodybuilding site

    ReplyDelete
  4. Much obliged to you for another extraordinary article. Where else would anyone be able to get that sort of data in such an impeccable method for composing? I have a presentation one week from now, and I am on the search for such data. buy phen375

    ReplyDelete
  5. Morbi sit amet ante quis lorem ultricies rhoncus vel quis sapien http://www.hcgultradietbuy.com

    ReplyDelete
  6. I am Dr Ogudugu a Traditional healer, herbalist and spell caster. I use real magical powers to fight demons, evil spirits and super villains. Are you suffering in silence? does everything in your life seems to be going wrong no matter how hard you try to make things better? have Questions about love, or career let me help you am a born gifted traditional healer and master psychic. I 've experience in exploration of the past, present and future. I can help you with any of those problems that are making your life difficult and miserable,do you have many enemies some enemies that you are not even aware of? are you having many obstacles in your life? Is your love,relationship,marriage life falling apart? is your life facing financial ruin ? i specialize in reuniting.Email me about your problems and I'll be of great help to you on my my personal email address: GREATOGUDUGU@GMAIL.COM Website: http://greatdrogudugusolutiontemple.webs.com/

    ReplyDelete