A tool for making Python as fast as C/Fortran¶
In what way is Python "slow"?¶
- It is interpreted*:
pythonis a piece of software (written in C) that executes python code on-the-fly. C/Fortran instead are compiled directly to machine-understadable binary code.
* Actually Python is sorta compiled. But the above is still mostly true.
- In practice this often isn't important because the basic operations are still the same. But numerical
forloops (especially with function calls) can be a problem.
- This is why
numpywas invented - it basically does efficient looping. But sometimes it's not well suited to a problem...
Which is clearer?¶
I want to know the values from values, which are aligned to ids2, corresponding to the IDs in ids1.
sorti = np.argsort(ids2)
idsforvals = np.searchsorted(ids2[sorti], ids1)
new_array = values[sorti][idsforvals]
# there's actually a better way to do this in newer numpy versions...
# but it's a contrived example to start with.
for i in ids1:
for j in ids2:
if i == j:
new_array[i] = values[j]
Sometimes loops are just easier to understand... and they can also be faster if they are at "close-to-metal" speed (i.e., C or Fortran)
What does Cython do about that?¶
- Cython is an "optimising static compiler" for Python
- It can be run on regular Python... But that usually only has a small effect (not worth the trouble), because the C it makes still does everything
pythondoes.
- But you can tell it to behave "like C" by specifying
typesfor variables. This can speed things up by orders of magnitude. Turns Cython into a "hybrid" language mixing syntax of Python with capabilities of C.
- Bonus: you can also call C functions directly
- The price is flexibility (and clarity): you can no longer use Python techniques like freely casting an integer to a float
Now lets do a "hello world" example¶
- First make sure you have cython and pkg-config installed (
conda install cython pkg-configshould work). - Look at the file
hello_world.pyxin the spacetelescope/pylunch repo. - Open up a terminal and go to the repo and do
cython hello_world.pyx. It should create ahello_world.cfile. Now compile that file. This can be tricky to know the right flags. The
pkg-configprogram can help you with that - this should work (Note thatpython3should be changed topythonif you're on python 2.x):clang -bundle $(pkg-config --libs --cflags python3) ./hello_world.c -o hello_world.soIf you have trouble, check out http://cython.readthedocs.io/en/latest/src/reference/compilation.html. If you still have trouble, try the method in http://cython.readthedocs.io/en/latest/src/tutorial/cython_tutorial.html, or just wait for a bit.
Once you've got it compiled, you should have a
hello_world.sofile, which acts like a Python package. So start up python and do:>>> import hello_world >>> hello_world.do_hello() hello world from python hello world from c
There's a much easier way when experimenting: the cython notebook extension!¶
%load_ext Cython
%%cython
from libc.stdio cimport printf
def do_hello():
print("hello world from python")
printf("hello world from c\n")
do_hello()
Note that the C-printed line is missing from the notebook. Watch out for subtle quirks like this!
(In this case it's because the C print function goes straight to the terminal rather than getting caught by the notebook - you should see it in your terminal if that's how you started the notebook server.)
What if you want to use it in multiple notebooks or a releasable package?¶
In a real code that you're releasing or packaging up for your own use over multiple notebooks, you'll want to use a proper setup.py, as then you won't have to manually compile anything - the python building packages will take care of that.
To continue our "hello world" example, you'll see a setup.py file along side the hello_world.pyx. Examine that file with your favorite text editor, and modify it to compile the hello_world.pyx. Then you just do python setup.py build. That will create a build/lib.<a-bunch-of-stuff-that-depends-on-your-computer> directory. cd into that one, start up python and do:
>>> import hello_world
>>> hello_world.do_hello()which should output the exact same thing as we saw above.
You can see a bit more detail about this in the Cython docs, or you can try using the Astropy affiliated package template, which comes out-of-the-box ready to include Cython code along side regular Python.
Why would we want to do this?¶
Consider the two examples below, one in Python and the other in Cython
def fib_py(n):
"""Compute the Fibonacci series up to n."""
a, b = 0, 1
while b < n:
a, b = b, a + b
return b
%%cython
def fib_cy(n):
"""Compute the Fibonacci series up to n."""
a, b = 0, 1
while b < n:
a, b = b, a + b
return b
print('Python version:')
%timeit -n 1 -r 1 fib_py(1000000)
print('Cython version:')
%timeit -n 1 -r 1 fib_cy(1000000)
OK, so Cython gives a 2x speedup...
Aside¶
The `-n 1 -r 1` bit above makes it do just *one* run of the code. That's important here because some machines will cache the results, giving inaccurate timing measurements of the actual algorithm. Sometimes that's *not* what you want, because the caching actually helps in real practical code... But for this example it's deceptive.
How can we do better?¶
cython -ais your friend. It shows where cython produces a lot of python-related function calls, and is a hint that you can add C types to speed things up. The "game" is to try to turn all the yellow to white in performance-sensitive parts.
%%cython -a
def fib_cy(n):
"""Compute the Fibonacci series up to n."""
a, b = 0, 1
while b < n:
a, b = b, a + b
return b
Looks like we spend a lot of time doing the "guts" of the work in the last one. Can we add some C types where we know they should exist?
%%cython -a
def fib_cy2(n):
"""Compute the Fibonacci series up to n."""
cdef int a, b
a, b = 0, 1
while b < n:
a, b = b, a + b
return b
Still spending some time dealing with the n, as it's a python integer. lets cast it in advance before the inner loop.
%%cython -a
def fib_cy3(n):
"""Compute the Fibonacci series up to n."""
cdef int a = 0
cdef int b = 1
cdef int c_n = int(n)
while b < c_n:
a, b = b, a + b
return b
print('Baseline Python:')
%timeit -n 1 -r 1 fib_py(100000000)
print('No C types:')
%timeit -n 1 -r 1 fib_cy(100000000)
print('Some C types:')
%timeit -n 1 -r 1 fib_cy2(100000000)
print('All C types:')
%timeit -n 1 -r 1 fib_cy3(100000000)
So that a reasonable speedup (4-5x), just from adding c types.
Exercise¶
Can you convert this into something that runs faster in Cython? http://cython.readthedocs.io/en/latest/src/reference/language_basics.html might help here. (I managed a ~50x speedup! But maybe you can do even better?)
def primes_py(kmax):
p = [0]*1000
result = []
if kmax > 1000:
kmax = 1000
k = 0
n = 2
while k < kmax:
i = 0
while i < k and n % p[i] != 0:
i = i + 1
if i == k:
p[k] = n
k = k + 1
result.append(n)
n = n + 1
return result
My answer is at the end of the notebook. Don't peek until you're ready!
Numpy and Cython¶
- Where Cython really shines is in combination with numpy. You stick with numpy for general most things, but drop into highly optimized Cython for "hotspots" or places where it's much clearer.
- One major gotcha: if you google "cython numpy", you'll probably come across pages like this: http://cython.readthedocs.io/en/latest/src/tutorial/numpy.html. It's worth reading to get a better understanding, but don't do that first! It's an older, more complex way of working with numpy from Cython. Instead, you want "typed memoryviews": http://cython.readthedocs.io/en/latest/src/userguide/memoryviews.html - those are much easier to use quickly and simply.
Let's see how the "typed memoryview" approach works:¶
def ap2b_py(a ,b):
return a + 2*b
%%cython
import numpy as np
def ap2b_cy(a , b):
# they start out as regular numpy arrays, so they know their own length
cdef int i
cdef int n = len(a)
newarr = np.empty(n)
cdef double[:] a_memview = a
cdef double[:] b_memview = b
cdef double[:] new_memview = newarr
for i in range(n):
new_memview[i] = a_memview[i] + 2.0 * b_memview[i]
return newarr
import numpy as np
a, b = np.random.randn(2, 10000000)
%timeit ap2b_py(a, b)
%timeit ap2b_cy(a, b)
This ~ factor-of-2 speed improvement also comes at a factor-of-2 memory improvement, as there is no intermediate numpy array.
Some more complex (but perhaps a bit more practical) examples of optimized Cython¶
Note that some of these are not the memoryview form described above. But they give you a flavor of the sort of performance improvements you might get
- http://technicaldiscovery.blogspot.com/2011/06/speeding-up-python-numpy-cython-and.html
- http://nealhughes.net/cython1/
- http://www.perrygeo.com/parallelizing-numpy-array-loops-with-cython-and-mpi.html (about Parallelization with OpenMP, suitable for clusters)
- One final tip for working with numpy and Cython. Consider adding this to cython functions:
import cython
@cython.cdivision(True)
@cython.wraparound(False)
@cython.boundscheck(False)
@cython.nonecheck(False)
def my_cy_func(...):
...We won't dwell on the details of this - if you care you can read more at http://cython.readthedocs.io/en/latest/src/reference/compilation.html#compiler-directives. Just know that they will often speed things up a bit more, but at the price of a higher chance of crashing python.
Interfacing with C code¶
- Cython also gives you the capability to call C directly. The main trick is that you need to include something like this in the Cython file:
you can then just usecdef extern from "myheaders.h": double c_function(int arg1, double arg2)c_function(a, b)in a Cython function and it will work.
- The (mental) complexity scales roughly with the complexity of the C code. E.g., functions with scalar arguments are easy,
double *often works (even with numpy arrays), butdouble **: here there be dragons.
- A good starting point is https://github.com/aphearin/cython_c_extension_example
- For the gory details see http://cython.readthedocs.io/en/latest/src/userguide/external_C_code.html
Other options¶
Cython is not the only way to speed numerical code up in python.
- Numba (http://numba.pydata.org/) is an up-and-coming package that claims to do much of what Cython does, but without explicit compilation needed.
- HOPE (http://www.cosmology.ethz.ch/research/software-lab/HOPE.html) is a similar package developed by astronomers (which can be a positive or negative, depending on your POV).
- PyPy (http://pypy.org/) is a re-implementation of python that aims to also do many of the same things (although thus far has not been focused on numerical code).
- Pyston (https://github.com/dropbox/pyston) is yet another effort along similar lines, developed by Dropbox.
And there are probably more. But we only have so much time!
For the remaining time...¶
Explore! If you have a C code of your own, perhaps you can try writing a Python interface to it. If you don't, you could work through some of the earlier examples, or try a couple astronomy-centric exercises - you might start by writing it in Python and then adapting it to Cython...:
- Write a Cython function to simulate a ground-based image of a nearby galaxy, given a star catalog (in the real case this would be from e.g. HST or JWST, but you could just make a random set of x/y positions, e.g. do
x, y = np.random.randn(2, n)). It's probably easiest to assume Gaussian seeing, but you can always try whatever you favorite seeing kernel is. (There's a solution to this at the very end taken from my own science notebooks.) - Write a Cython function that convolves an image (a 2D numpy array) with a small kernel (another 2D numpy array, say 5x5). (You'd probably want to use astropy.convolution for this in a real science situation... But that is, itself implemented in Cython, so you could look at astropy's code for ideas)
Answer to Exercise¶
%%cython
def primes_cy(int kmax):
cdef int n, k, i
cdef int p[1000]
result = []
if kmax > 1000:
kmax = 1000
k = 0
n = 2
while k < kmax:
i = 0
while i < k and n % p[i] != 0:
i = i + 1
if i == k:
p[k] = n
k = k + 1
result.append(n)
n = n + 1
return result
%timeit primes_py(1000)
%timeit primes_cy(1000)
Possible answer to "simulate galaxy" exercise¶
(This is just one of many ways to do it... But I had to do this myself once, so here's how I did it:)
%%cython
import cython
import numpy as np
from libc.math cimport exp
cdef double PI = 3.141592653589793
@cython.cdivision(True)
cdef inline double ngauss2d(double x, double y, double sig):
cdef double tsigsq = 2.0*sig*sig
return exp(-(x*x + y*y)/tsigsq)/(tsigsq * PI)
@cython.wraparound(False)
@cython.boundscheck(False)
@cython.nonecheck(False)
def convolve(data, xgrid, ygrid, fluxes, cenx, ceny, fwhm):
"""
Convolves a star catalog with Gaussian seeing to produce an artificial
ground-based/distant image.
Parameters
----------
data: 2D array
An array to add the simulated image to
xgrid: 2D array
The x-values of each pixel in `data` - must be the same shape as `data`.
ygrid: 2D array
The y-values of each pixel in `data` - must be the same shape as `data`.
fluxes: 1D array
The flux in each star
cenx: 1D array
The x-position of the center of each star - must be the same shape as `fluxes`.
ceny: 1D array
The y-position of the center of each star - must be the same shape as `fluxes`.
fwhm: float
The FWHM of the "seeing". E.g., the gaussian convolution kernel.
(In the same units as `xgrid`, `ygrid`, `cenx`, and `ceny`).
"""
from math import log
import time
cdef double[:, :] data_view = data
cdef double[:, :] xgrid_view = xgrid
cdef double[:, :] ygrid_view = ygrid
cdef double[:] fluxes_view = fluxes
cdef double[:] cenx_view = cenx
cdef double[:] ceny_view = ceny
cdef int I, J, S
cdef double sig = fwhm/(2*(2*log(2))**0.5)
cdef double starflux, starx, stary, dx, dy
I = data.shape[0]
J = data.shape[1]
S = fluxes.shape[0]
if xgrid.shape != (I, J) or ygrid.shape != (I, J):
raise ValueError("xgrid or ygrid don't match data!")
if cenx.shape != (S,) or ceny.shape != (S,):
raise ValueError("cenx or ceny don't match fluxes!")
for s in range(S):
starflux = fluxes_view[s]
starx = cenx_view[s]
stary = ceny_view[s]
for i in range(I):
for j in range(J):
dx = starx + xgrid_view[i, j]
dy = stary + ygrid_view[i, j]
data_view[i,j] += starflux * ngauss2d(dx, dy, sig)