Saturday, April 27, 2013

Fun with Numba, NumPy and F2PY

I've started using Numba to speed up MD simulations for our course in Molecular Statistics, which I teach together with +Jan Jensen and +Jimmy Charnley Kromann.

One exercise is using the Velo-Verlet algorithm to simulate a Lennard-Jones gas/liquid with periodic boundary conditions. Since we do everything in Python in this course, running the actual simulation is quite slow. What we do is we supply the students with a compiled FORTRAN module, compiled with F2PY which they can use when they've written their own Lennard-Jones gradient code.

The FORmula TRANslator module is extremely fast, compared to Python code. I rewrote the code from last year's course to a function I could use with Numba and compare directly to the F2PY module

For 100 particles:

  • Numba: 0.146 ms/iteration (466x speed up)
  • F2PY: 0.236 ms/iteration (289x speed up)
  • Python: 68.06 ms/iteration
Now, I'm pretty sure a pure FORTRAN implementation would be faster than the F2PY, but I am very impressed with the Numba, since it's standard Python code and MUCH more simple to read than FORTRAN -- in theory all you have to do is write @autojit before a function definition. This is of course in theory ... I didn't do anything to improve on the readability of the Lennard-Jones code here. Before I posted, I had actually inserted a comment saying """This code block is unreadable""".

Now, there were a few issues I discovered:

1) @autojit vs. @jit(argtypes=[double[:,:], double[:]])

Since Numba has to decide what arguments a function takes before it's being compiled, it needs to know what possible type of arguments the function takes. If you're lucky, Numba decides on the right type and you can get away with @autojit, but sometimes autojitting makes the code SLOWER than standard, interpreted Python. In one case I got a factor of 10x slower with autojit, but explicitly stating the function argument types with @jit(argtypes = .... ) I got a speed up of 20x on the same code compared to interpreted Python.

2) Returning tuples in compiled code block

Numba does not allow a tuple to be returned inside a compiled code block. So don't do this.

3) Use Numpy properly

  • Numba is NumPy aware, so code ran faster when numbers were stored in numpy.array types rather than just regular Python lists.
  • At first I was numpy.round() to round to nearest integer in the periodic boundary condition code. Switching three numpy.round() calls to numpy.rint() gave a speed up of around 100x on the code execution.
  • Use U[i, j] instead of U[i][j] on NumPy arrays. Not much of a difference in the vanilla Python code, but MASSIVE speed gains in compiled code.

Python Lennard-Jones code:

F2PY Lennard-Jones code:

Python Velo-Verlet solver:


  1. thanks for posting.valuable information about numpy and other stuffs

  2. computational biochemistry in a effective way.Thanks for blogging

    Python Training in chennai | Python Classes in Chennai

  3. Aimore Technologies is the best Selenium training institute in Chennai with 6+ years of experience. We are offering online and classroom training.
    Visit Us: Selenium Training in Chennai

  4. Much obliged for sharing this brilliant substance. its extremely fascinating. Numerous web journals I see these days don't actually give whatever pulls in others however the manner in which you have plainly clarified everything it's truly awesome. There are loads of posts But your method of Writing is so Good and Knowledgeable. continue to post such helpful data and view my site too...
    how to make a paper airplane eagle | how to make a boomerang airplane | the eagle paper airplane | best paper airplane design for distance and accuracy | best paper airplanes for distance and speed | what is the best paper airplane design for distance | paper airplane that flies far and straight | nakamura lock paper airplane instructions | paper airplane templates for distance | science behind paper airplanes

  5. Get your Career in IT on track with our Career Accelerator Program by Premium Learnings

    Check our Youtube channel for our training videos -

    For more info do visit us:

  6. Hello,
    Can I get some questions answers because I don't have any dating sitge account but I am getting dating messages.. Why am I getting Spam Emails for Dating Sites