The Easy "Hard" Way: Cythonizing

In this notebook, we'll build on the previous work where we used SymPy's code printers to generate code for evaluating expressions numerically. As a layer of abstraction on top of C code printers, which generate snippets of code we can copy into a C program, we can generate a fully compilable C library. On top of this, we will see how to use Cython to compile such a library into a Python extension module so its computational routines can be called directly from Python.

Learning Objectives

After this lesson, you will be able to:

1. write a simple Cython function and run it a Jupyter notebook using the %%cython magic command
2. use the SymPy codegen function to output compilable C code
3. wrap codegen-generated code with Cython, compile it into an extension module, and call it from Python
4. use SymPy's autowrap function to do all of this behind the scenes
5. pass a custom code printer to autowrap to make use of an external C library

1. A Quick Introduction to Cython

Cython is a compiler and a programming language used to generate C extension modules for Python.

The Cython language is a Python/C creole which is essentially Python with some additional keywords for specifying static data types. It looks something like this:

def cython_sum(int n):
    cdef float s = 0.0
    cdef int i
    for i in range(n):
        s += i
    return s

The Cython compiler transforms this code into a "flavor" of C specific to Python extension modules. This C code is then compiled into a binary file that can be imported and used just like a regular Python module -- the difference being that the functions you use from that module can potentially be much faster and more efficient than an equivalent pure Python implementation.

Aside from writing Cython code for computations, Cython is commonly used for writing wrappers around existing C code so that the functions therein can be made available in an extension module as described above. We will use this technique to make the SymPy-generated C code accessible to Python for use in SciPy's odeint.

Example

As a quick demonstration of what Cython can offer, we'll walk through a simple example of generating numbers in the Fibonacci sequence. If you're not familiar with it already, the sequence is initialized with $F_0 = 0$ and $F_1 = 1$, then the remaining terms are defined recursively by:

Our objective is to write a function that computes the $n$-th Fibonacci number. Let's start by writing a simple iterative solution in pure Python.

Let's see how long it takes to compute the 100th Fibonacci number.

Now let's implement the same thing with Cython. Since Cython is essentially "Python with types," it is often fairly easy to make the move and see improvements in speed. It does come at the cost, however, of a separate compilation step.

There are several ways to ways to go about the compilation process, and in many cases, Cython's tooling makes it fairly simple. For example, Jupyter notebooks can make use of a %%cython magic command that will do all of compilation in the background for us. To make use of it, we need to load the cython extension.

Now we can write a Cython function.

Note: the --annotate (or -a) flag of the %%cython magic command will produce an interactive annotated printout of the Cython code, allowing us to see the C code that is generated.

To see a bit more about writing Cython and its potential performance benefits, see this Cython examples notebook.

Even better, check out Kurt Smith's Cython tutorial which is happening at the same time as this tutorial.

2. Generating C Code with SymPy's codegen()

Our main goal in using Cython is to wrap SymPy-generated C code into a Python extension module so that we can call the fast compiled numerical routines from Python.

SymPy's codegen function takes code printing a step further: it wraps a snippet of code that numerically evaluates an expression with a function, and puts that function into the context of a file that is fully ready-to-compile code.

Here we'll revisit the water radiolysis system, with the aim of numerically computing the right hand side of the system of ODEs and integrating using SciPy's odeint.

Recall that this system looks like:

where we are representing our state variables $y_0,\,y_1,\dots,y_{13}$ as a vector $\mathbf{y}(t)$ that we called states in our code, and the collection of functions on the right hand side $\mathbf{f}(\mathbf{y}(t))$ we called rhs_of_odes.

Start by importing the system of ODEs and the matrix of state variables.

Now we'll use codegen (under sympy.utilities.codegen) to output C source and header files which can compute the right hand side (RHS) of the ODEs numerically, given the current values of our state variables. Here we'll import it and show the documentation:

We just have one expression we're interested in computing, and that is the matrix expression representing the derivatives of our state variables with respect to time: rhs_of_odes. What we want codegen to do is create a C function that takes in the current values of the state variables and gives us back each of the derivatives.

Note that we've just expanded the outputs into individual variables so we can access the generated code easily. codegen gives us back the .c filename and its source code in a tuple, and the .h filename and its source in another tuple. Let's print the source code.

There are several things here worth noting:

• the state variables are passed in individually
• the state variables in the function signature are out of order
• the output array is passed in as a pointer like in our Fibonacci sequence example, but it has an auto-generated name

Let's address the first issue first. Similarly to what we did in the C printing exercises, let's use a MatrixSymbol to represent our state vector instead of a matrix of individual state variable symbols (i.e. y[0] instead of y0). First, create the MatrixSymbol object that is the same shape as our states matrix.

Now we need to replace the use of y0, y1, etc. in our rhs_of_odes matrix with the elements of our new state vector (e.g. y[0], y[1], etc.). We saw how to do this already in the previous notebook. Start by forming a mapping from y0 -> y[0, 0], y1 -> y[1, 0], etc.

Now replace the symbols in rhs_of_odes according to the mapping. We'll call it rhs_of_odes_ind and use that from now on.

Exercise: use codegen again, but this time with rhs_of_odes_ind which makes use of a state vector rather than a container of symbols. Check out the resulting code. What is different about the function signature?

[(cf, cs), (hf, hs)] = codegen(???)

Solution

|
|
|
|
|
|
|
|
|
|
v

So by re-writing our expression in terms of a MatrixSymbol rather than individual symbols, the function signature of the generated code is cleaned up greatly.

However, we still have the issue of the auto-generated output variable name. To fix this, we can form a matrix equation rather than an expression. The name given to the symbol on the left hand side of the equation will then be used for our output variable name.

We'll start by defining a new MatrixSymbol that will represent the left hand side of our equation -- the derivatives of each state variable.

Exercise: form an equation using sym.Eq to equate the two sides of our system of differential equations, then use this as the expression in codegen. Print out just the header source to see the function signature. What is the output argument called now?

ode_eq = sym.Eq(???)
[(cf, cs), (hf, hs)] = codegen(???)
print(???)

Solution

|
|
|
|
|
|
|
|
|
|
v

Now we see that the c_odes function signature is nice and clean. We pass it a pointer to an array representing the current values of all of our state variables and a pointer to an array that we want to fill with the derivatives of each of those state variables.

If you're not familiar with C and pointers, you just need to know that it is idiomatic in C to preallocate a block of memory representing an array, then pass the location of that memory (and usually the number of elements it can hold), rather than passing the array itself to/from a function. For our purposes, this is as complicated as pointers will get.

Just so we can compile this code and use it, we'll re-use the codegen call above with to_files=True so the .c and .h files are actually written to the filesystem, rather than having their contents returned in a string.

3. Wrapping the Generated Code with Cython

Now we want to wrap the function that was generated c_odes with a Cython function so we can generate an extension module and call that function from Python. Wrapping a set of C functions involves writing a Cython script that specifies the Python interface to the C functions. This script must do two things:

1. specify the function signatures as found in the C source
2. implement the Python interface to the C functions by wrapping them

The build system of Cython is able to take the Cython wrapper source code as well as the C library source code and compile/link things together into a Python extension module. We will write our wrapper code in a cell making use of the magic command %%cython_pyximport, which does a few things for us:

• writes the contents of the cell to a Cython source file (modname.pyx)
• looks for a modname.pyxbld file for instructions on how to build things
• builds everything into an extension module
• imports the extension module, making the functions declared there available in the notebook

So, it works similarly to the %%cython magic command we saw at the very beginning, but things are a bit more complicated now because we have this external library c_odes that needs to be compiled as well.

Note: The pyxbld file contains code similar to what would be found in the setup.py file of a package making use of Cython code for wrapping C libraries.

In either case, all that's needed is to tell setuptools/Cython:

• the name of the extension module we want to make
• the location of the Cython and C source files to be built
• the location of headers needed during compilation -- both our C library's headers as well as NumPy's headers

We will call our extension module cy_odes, so here we'll generate a cy_odes.pyxbld file to specify how to build the module.

Now we can write our wrapper code.

To write the wrapper, we first write the function signature as specified by the C library. Then, we create a wrapper function that makes use of the C implementation and returns the result. This wrapper function becomes the interface to the compiled code, and it does not need to be identical to the C function signature. In fact, we'll make our wrapper function compliant with the odeint interface (i.e. takes a 1-dimensional array of state variable values and the current time).

Exercise: use np.random.randn to generate random state variable values and evaluate the right-hand-side of our ODEs with those values.

random_vals = np.random.randn(???)
???

Solution

|
|
|
|
|
|
|
|
|
|
v

Now we can use odeint to integrate the equations and plot the results to check that it worked. First we need to import odeint.

A couple convenience functions are provided in the scipy2017codegen package which give some reasonable initial conditions for the system and plot the state trajectories, respectively. Start by grabbing some initial values for our state variables and time values.

Finally we can integrate the equations using our Cython-wrapped C function and plot the results.

4. Generating and Compiling a C Extension Module Automatically

As yet another layer of abstraction on top of codegen, SymPy provides an autowrap function that can automatically generate a Cython wrapper for the generated C code. This greatly simplifies the process of going from a SymPy expression to numerical computation, but as we'll see, we lose a bit of flexibility compared to manually creating the Cython wrapper.

Let's start by importing the autowrap function and checking out its documentation.

So autowrap takes in a SymPy expression and gives us back a binary callable which evaluates the expression numerically. Let's use the Equality formed earlier to generate a function we can call to evaluate the right hand side of our system of ODEs.

Exercise: use the main Jupyter notebook tab to head to the temporary directory autowrap just created. Take a look at some of the files it contains. Can you map everything we did manually to the files generated?

Solution

|
|
|
|
|
|
|
|
|
|
v

autowrap generates quite a few files, but we'll explicitly list a few here:

• wrapped_code_#.c: the same thing codegen generated
• wrapper_module_#.pyx: the Cython wrapper code
• wrapper_module_#.c: the cythonized code
• setup.py: specification of the Extension for how to build the extension module

Exercise: just like we did before, generate some random values for the state variables and use auto_odes to compute the derivatives. Did it work like before?

Hint: take a look at wrapper_module_#.pyx to see the types of the arrays being passed in / created.

random_vals = np.random.randn(???)
auto_odes(???)

Solution

|
|
|
|
|
|
|
|
|
|
v

One advantage to wrapping the generated C code manually is that we get fine control over how the function is used from Python. That is, in our hand-written Cython wrapper we were able to specify that from the Python side, the input to our wrapper function and its return value are both 1-dimensional ndarray objects. We were also able to add in the extra argument t for the current time, making the wrapper function fully compatible with odeint.

However, autowrap just sees that we have a matrix equation where each side is a 2-dimensional array with shape (14, 1). The function returned then expects the input array to be 2-dimensional and it returns a 2-dimensional array.

This won't work with odeint, so we can write a simple wrapper that massages the input and output and adds an extra argument for t.

Now a 1-dimensional input works.

Exercise: As we have seen previously, we can analytically evaluate the Jacobian of our system of ODEs, which can be helpful in numerical integration. Compute the Jacobian of rhs_of_odes_ind with respect to y, then use autowrap to generate a function that evaluates the Jacobian numerically. Finally, write a Python wrapper called auto_jac_wrapper to make it compatible with odeint.

# compute jacobian of rhs_of_odes_ind with respect to y
???

# generate a function that computes the jacobian
auto_jac = autowrap(???)

def auto_jac_wrapper(y, t):
    return ???

Test your wrapper by passing in the random_vals array from above. The shape of the result should be (14, 14).

Solution

|
|
|
|
|
|
|
|
|
|
v

Finally, we can use our two wrapped functions in odeint and compare to our manually-written Cython wrapper result.

5. Using a Custom Printer and an External Library with autowrap

As of SymPy 1.1, autowrap accepts a custom CodeGen object, which is responsible for generating the code. The CodeGen object in turn accepts a custom CodePrinter object, meaning we can use these two points of flexibility to make use of customized code printing in an autowrapped function. The following example is somewhat contrived, but the concept in general is powerful.

In our set of ODEs, there are quite a few instances of $y_i^2$, where $y_i$ is one of the 14 state variables. As an example, here's the equation for $\frac{dy_3(t)}{dt}$:

There is a library called fastapprox that provides computational routines things like powers, exponentials, logarithms, and a few others. These routines provide limited precision with respect to something like math.h's equivalent functions, but they offer potentially faster computation.

The fastapprox library provides a function called fastpow, with the signature fastpow(float x, float p). It it follows the interface of pow from math.h. In the previous notebook, we saw how to turn instances of $x^3$ into x*x*x, which is potentially quicker than pow(x, 3). Here, let's just use fastpow instead.

Exercise: implement a CustomPrinter class that inherits from C99CodePrinter and overrides the _print_Pow function to make use of fastpow. Test it by instantiating the custom printer and printing a SymPy expression $x^3$.

Hint: it may be helpful to run C99CodePrinter._print_Pow?? to see how it works

from sympy.printing.ccode import C99CodePrinter

class CustomPrinter(C99CodePrinter):
    def _print_Pow(self, expr):
        ???

printer = CustomPrinter()
x = sym.symbols('x')
# print x**3 using the custom printer
???

Solution

|
|
|
|
|
|
|
|
|
|
v

Now we can create a C99CodeGen object that will make use of this printer. This object will be passed in to autowrap with the code_gen keyword argument, and autowrap will use it in the code generation process.

However, for our generated code to use the fastpow function, it needs to have a #include "fastpow.h" preprocessor statement at the top. The code gen object supports this by allowing us to append preprocessor statements to its preprocessor_statements attribute.

One final issue remains, and that is telling autowrap where to find the fastapprox library headers. These header files have just been downloaded from GitHub and placed in the scipy2017codegen package, so it should be installed with the conda environment. We can find it by looking for the package directory.

Finally we're ready to call autowrap. We'll just use ode_eq, the Equality we created before, pass in the custom CodeGen object, and tell autowrap where the fastapprox headers are located.

If we navigate to the tmp directory created, we can view the wrapped_code_#.c to see our custom printing in action.

As before, we need a wrapper function for use with odeint, but aside from that, everything should be in place.

Exercise: generate an array of random state variable values, then use this array in the auto_odes_wrapper and auto_odes_fastpow_wrapper functions. Compare their outputs.

Solution

|
|
|
|
|
|
|
|
|
|
v