%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import sympy as sym
sym.init_printing()

"""
Explanation: 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:

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


def python_fib(n):
    a = 0.0
    b = 1.0
    for i in range(n):
        tmp = a
        a = a + b
        b = tmp
    return a

[python_fib(i) for i in range(10)]

"""
Explanation: 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:
cython
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:
$$
F_i = F_{i-1} + F_{i-2}
$$
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.
End of explanation
"""


%timeit python_fib(100)

"""
Explanation: Let's see how long it takes to compute the 100th Fibonacci number.
End of explanation
"""


%load_ext cython

"""
Explanation: 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.
End of explanation
"""


%%cython
def cython_fib(int n):
    cdef double a = 0.0
    cdef double b = 1.0
    cdef double tmp
    for i in range(n):
        tmp = a
        a = a + b
        b = tmp
    return a

%timeit cython_fib(100)

"""
Explanation: 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.
End of explanation
"""


from scipy2017codegen.chem import load_large_ode
rhs_of_odes, states = load_large_ode()
rhs_of_odes[0]

"""
Explanation: 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:
$$
\begin{align}
\frac{dy_0(t)}{dt} &= f_0\left(y_0,\,y_1,\,\dots,\,y_{13},\,t\right) \
&\vdots \
\frac{dy_{13}(t)}{dt} &= f_{13}\left(y_0,\,y_1,\,\dots,\,y_{13},\,t\right)
\end{align}
$$
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.
End of explanation
"""


from sympy.utilities.codegen import codegen
#codegen?

"""
Explanation: 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:
End of explanation
"""


[(cf, cs), (hf, hs)] = codegen(('c_odes', rhs_of_odes), language='c')

"""
Explanation: 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.
End of explanation
"""


print(cs)

"""
Explanation: 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.
End of explanation
"""


y = sym.MatrixSymbol('y', *states.shape)

"""
Explanation: 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.
End of explanation
"""


state_array_map = dict(zip(states, y))
state_array_map

"""
Explanation: 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 -&gt; y[0, 0], y1 -&gt; y[1, 0], etc.
End of explanation
"""


rhs_of_odes_ind = rhs_of_odes.xreplace(state_array_map)
rhs_of_odes_ind[0]

"""
Explanation: 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.
End of explanation
"""


[(cf, cs), (hf, hs)] = codegen(('c_odes', rhs_of_odes_ind), language='c')
print(cs)

"""
Explanation: 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?
python
[(cf, cs), (hf, hs)] = codegen(???)
Solution
|
|
|
|
|
|
|
|
|
|
v
End of explanation
"""


dY = sym.MatrixSymbol('dY', *y.shape)

"""
Explanation: 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.
End of explanation
"""


ode_eq = sym.Eq(dY, rhs_of_odes_ind)
[(cf, cs), (hf, hs)] = codegen(('c_odes', ode_eq), language='c')
print(hs)

"""
Explanation: 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?
python
ode_eq = sym.Eq(???)
[(cf, cs), (hf, hs)] = codegen(???)
print(???)
Solution
|
|
|
|
|
|
|
|
|
|
v
End of explanation
"""


codegen(('c_odes', ode_eq), language='c', to_files=True)

"""
Explanation: 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.
End of explanation
"""


%%writefile cy_odes.pyxbld
import numpy

#            module name specified by `%%cython_pyximport` magic
#            |        just `modname + ".pyx"`
#            |        |
def make_ext(modname, pyxfilename):
    from setuptools.extension import Extension
    return Extension(modname,
                     sources=[pyxfilename, 'c_odes.c'],
                     include_dirs=['.', numpy.get_include()])

"""
Explanation: 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:

specify the function signatures as found in the C source
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.
End of explanation
"""


%%cython_pyximport cy_odes
import numpy as np
cimport numpy as cnp # cimport gives us access to NumPy's C API

# here we just replicate the function signature from the header
cdef extern from "c_odes.h":
    void c_odes(double *y, double *dY)

# here is the "wrapper" signature that conforms to the odeint interface
def cy_odes(cnp.ndarray[cnp.double_t, ndim=1] y, double t):
    # preallocate our output array
    cdef cnp.ndarray[cnp.double_t, ndim=1] dY = np.empty(y.size, dtype=np.double)
    # now call the C function
    c_odes(<double *> y.data, <double *> dY.data)
    # return the result
    return dY

"""
Explanation: 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).
End of explanation
"""


random_vals = np.random.randn(14)
cy_odes(random_vals, 0) # note: any time value will do

"""
Explanation: Exercise: use np.random.randn to generate random state variable values and evaluate the right-hand-side of our ODEs with those values.
python
random_vals = np.random.randn(???)
???
Solution
|
|
|
|
|
|
|
|
|
|
v
End of explanation
"""


from scipy.integrate import odeint

"""
Explanation: Now we can use odeint to integrate the equations and plot the results to check that it worked. First we need to import odeint.
End of explanation
"""


from scipy2017codegen.chem import watrad_init, watrad_plot
y_init, t_vals = watrad_init()

"""
Explanation: 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.
End of explanation
"""


y_vals = odeint(cy_odes, y_init, t_vals)
watrad_plot(t_vals, y_vals)

"""
Explanation: Finally we can integrate the equations using our Cython-wrapped C function and plot the results.
End of explanation
"""


from sympy.utilities.autowrap import autowrap
#autowrap?

"""
Explanation: 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.
End of explanation
"""


auto_odes = autowrap(ode_eq, backend='cython', tempdir='./autowraptmp')

"""
Explanation: 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.
End of explanation
"""


random_vals = np.random.randn(14, 1) # need a 2-dimensional vector
auto_odes(random_vals)

"""
Explanation: 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.
python
random_vals = np.random.randn(???)
auto_odes(???)
Solution
|
|
|
|
|
|
|
|
|
|
v
End of explanation
"""


def auto_odes_wrapper(y, t):
    dY = auto_odes(y[:, np.newaxis])
    return dY.squeeze()

"""
Explanation: 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.
End of explanation
"""


random_vals = np.random.randn(14)
auto_odes_wrapper(random_vals, 0)

"""
Explanation: Now a 1-dimensional input works.
End of explanation
"""


jac = rhs_of_odes_ind.jacobian(y)

auto_jac = autowrap(jac, backend='cython', tempdir='./autowraptmp')

def auto_jac_wrapper(y, t):
    return auto_jac(y[:, np.newaxis])

auto_jac_wrapper(random_vals, 2).shape

"""
Explanation: 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.
```python
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
End of explanation
"""


y_vals = odeint(auto_odes_wrapper, y_init, t_vals, Dfun=auto_jac_wrapper)
watrad_plot(t_vals, y_vals)

"""
Explanation: Finally, we can use our two wrapped functions in odeint and compare to our manually-written Cython wrapper result.
End of explanation
"""


rhs_of_odes[3]

"""
Explanation: 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}$:
End of explanation
"""


from sympy.printing.ccode import C99CodePrinter

class CustomPrinter(C99CodePrinter):
    def _print_Pow(self, expr):
        return "fastpow({}, {})".format(self._print(expr.base),
                                        self._print(expr.exp))

printer = CustomPrinter()
x = sym.symbols('x')
printer.doprint(x**3)

"""
Explanation: 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
```python
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
End of explanation
"""


from sympy.utilities.codegen import C99CodeGen
gen = C99CodeGen(printer=printer)

"""
Explanation: 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.
End of explanation
"""


gen.preprocessor_statements.append('#include "fastpow.h"')

"""
Explanation: 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.
End of explanation
"""


import os
import scipy2017codegen

package_dir = os.path.dirname(scipy2017codegen.__file__)
fastapprox_dir = os.path.join(package_dir, 'fastapprox')

"""
Explanation: 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.
End of explanation
"""


auto_odes_fastpow = autowrap(ode_eq,
                             code_gen=gen,
                             backend='cython',
                             include_dirs=[fastapprox_dir],
                             tempdir='autowraptmp_custom')

"""
Explanation: 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.
End of explanation
"""


def auto_odes_fastpow_wrapper(y, t):
    dY = auto_odes_fastpow(y[:, np.newaxis])
    return dY.squeeze()

y_vals, info = odeint(auto_odes_fastpow_wrapper, y_init, t_vals, full_output=True)
watrad_plot(t_vals, y_vals)

"""
Explanation: 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.
End of explanation
"""


random_vals = np.random.randn(14)
dY1 = auto_odes_wrapper(random_vals, 0)
dY2 = auto_odes_fastpow_wrapper(random_vals, 0)
dY1 - dY2

"""
Explanation: 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
End of explanation
"""


