Software Engineering; Wave Equation Model
C.5 Migrating Loops to Cython
def assert_no_error(self):
"""Run through mesh and check error"""
Nx = self[’Nx’]
Nt = self.m.Nt
L, T = self.problem[’L T’.split()]
L = L/2 # only half the domain used (symmetry) x = np.linspace(0, L, Nx+1) # Mesh points in space t = np.linspace(0, T, Nt+1) # Mesh points in time for n in range(len(t)):
u_e = self.problem.u_exact(x, t[n]) diff = np.abs(self.f.u[n,:] - u_e).max() print ’diff:’, diff
tol = 1E-13 assert diff < tol
Observe that the solutions from all time steps are stored in the mesh function, which allows error assessment (inassert_no_error) to take place after all solu- tions have been found. Of course, in 2D or 3D, such a strategy may place too high demands on available computer memory, in which case intermediate results could be stored on file.
Runningwave1D_oo.pygives a printout showing that the class-based imple- mentation performs as expected, i.e. that the known exact solution is reproduced (within machine precision).
476 C Software Engineering; Wave Equation Model
functions (or classes) and placing them in one or more separate files with extension .pyx.
C.5.1 Declaring Variables and Annotating the Code
Our starting point is the plainadvance_scalarfunction for a scalar implementa- tion of the updating algorithm for new valuesunC1i;j :
def advance_scalar(u, u_n, u_nm1, f, x, y, t, n, Cx2, Cy2, dt2, V=None, step1=False):
Ix = range(0, u.shape[0]); Iy = range(0, u.shape[1]) if step1:
dt = sqrt(dt2) # save
Cx2 = 0.5*Cx2; Cy2 = 0.5*Cy2; dt2 = 0.5*dt2 # redefine D1 = 1; D2 = 0
else:
D1 = 2; D2 = 1 for i in Ix[1:-1]:
for j in Iy[1:-1]:
u_xx = u_n[i-1,j] - 2*u_n[i,j] + u_n[i+1,j]
u_yy = u_n[i,j-1] - 2*u_n[i,j] + u_n[i,j+1]
u[i,j] = D1*u_n[i,j] - D2*u_nm1[i,j] + \
Cx2*u_xx + Cy2*u_yy + dt2*f(x[i], y[j], t[n]) if step1:
u[i,j] += dt*V(x[i], y[j])
# Boundary condition u=0 j = Iy[0]
for i in Ix: u[i,j] = 0 j = Iy[-1]
for i in Ix: u[i,j] = 0 i = Ix[0]
for j in Iy: u[i,j] = 0 i = Ix[-1]
for j in Iy: u[i,j] = 0 return u
We simply take a copy of this function and put it in a filewave2D_u0_loop_cy.
pyx. The relevant Cython implementation arises from declaring variables with types and adding some important annotations to speed up array computing in Cython. Let us first list the complete code in the.pyxfile:
import numpy as np cimport numpy as np cimport cython
ctypedef np.float64_t DT # data type
@cython.boundscheck(False) # turn off array bounds check
@cython.wraparound(False) # turn off negative indices (u[-1,-1]) cpdef advance(
np.ndarray[DT, ndim=2, mode=’c’] u, np.ndarray[DT, ndim=2, mode=’c’] u_n, np.ndarray[DT, ndim=2, mode=’c’] u_nm1, np.ndarray[DT, ndim=2, mode=’c’] f, double Cx2, double Cy2, double dt2):
cdef:
int Ix_start = 0 int Iy_start = 0
int Ix_end = u.shape[0]-1 int Iy_end = u.shape[1]-1 int i, j
double u_xx, u_yy
for i in range(Ix_start+1, Ix_end):
for j in range(Iy_start+1, Iy_end):
u_xx = u_n[i-1,j] - 2*u_n[i,j] + u_n[i+1,j]
u_yy = u_n[i,j-1] - 2*u_n[i,j] + u_n[i,j+1]
u[i,j] = 2*u_n[i,j] - u_nm1[i,j] + \ Cx2*u_xx + Cy2*u_yy + dt2*f[i,j]
# Boundary condition u=0 j = Iy_start
for i in range(Ix_start, Ix_end+1): u[i,j] = 0 j = Iy_end
for i in range(Ix_start, Ix_end+1): u[i,j] = 0 i = Ix_start
for j in range(Iy_start, Iy_end+1): u[i,j] = 0 i = Ix_end
for j in range(Iy_start, Iy_end+1): u[i,j] = 0 return u
This example may act as a recipe on how to transform array-intensive code with loops into Cython.
1. Variables are declared with types: for example,double vin the argument list instead of justv, andcdef double vfor a variablevin the body of the func- tion. A Python floatobject is declared as double for translation to C by Cython, while anintobject is declared byint.
2. Arrays need a comprehensive type declaration involving the typenp.ndarray,
the data type of the elements, here 64-bit floats, abbreviated asDTthrough ctypedef np.float64_t DT(instead ofDTwe could use the full name of the data type:np.float64_t, which is a Cython-defined type),
the dimensions of the array, herendim=2andndim=1, specification of contiguous memory for the array (mode=’c’).
3. Functions declared withcpdefare translated to C but are also accessible from Python.
4. In addition to the standardnumpyimport we also need a special Cython import ofnumpy:cimport numpy as np, to appearafterthe standard import.
5. By default, array indices are checked to be within their legal limits. To speed up the code one should turn off this feature for a specific function by placing
@cython.boundscheck(False)above the function header.
6. Also by default, array indices can be negative (counting from the end), but this feature has a performance penalty and is therefore here turned off by writing
@cython.wraparound(False)right above the function header.
7. The use of index setsIxandIyin the scalar code cannot be successfully trans- lated to C. One reason is that constructions like Ix[1:-1] involve negative
478 C Software Engineering; Wave Equation Model
indices, and these are now turned off. Another reason is that Cython loops must take the formfor i in xrange orfor i in rangefor being trans- lated into efficient C loops. We have therefore introducedIx_startasIx[0]
andIx_endasIx[-1]to hold the start and end of the values of indexi. Similar variables are introduced for thejindex. A loopfor i in Ixis with these new variables written asfor i in range(Ix_start, Ix_end+1).
Array declaration syntax in Cython
We have used the syntax np.ndarray[DT, ndim=2, mode=’c’]to declare numpyarrays in Cython. There is a simpler, alternative syntax, employingtyped memory views1, where the declaration looks likedouble [:,:]. However, the full support for this functionality is not yet ready, and in this text we use the full array declaration syntax.
C.5.2 Visual Inspection of the C Translation
Cython can visually explain how successfully it translated a code from Python to C.
The command
Terminal
Terminal> cython -a wave2D_u0_loop_cy.pyx
produces an HTML filewave2D_u0_loop_cy.html, which can be loaded into a web browser to illustrate which lines of the code that have been translated to C.
FigureC.1shows the illustrated code. Yellow lines indicate the lines that Cython
Fig. C.1 Visual illustration of Cython’s ability to translate Python to C
1http://docs.cython.org/src/userguide/memoryviews.html
did not manage to translate to efficient C code and that remain in Python. For the present code we see that Cython is able to translate all the loops with array computing to C, which is our primary goal.
You can also inspect the generated C code directly, as it appears in the file wave2D_u0_loop_cy.c. Nevertheless, understanding this C code requires some familiarity with writing Python extension modules in C by hand. Deep down in the file we can see in detail how the compute-intensive statements have been translated into some complex C code that is quite different from what a human would write (at least if a direct correspondence to the mathematical notation was intended).
C.5.3 Building the Extension Module
Cython code must be translated to C, compiled, and linked to form what is known in the Python world as aC extension module. This is usually done by making a setup.pyscript, which is the standard way of building and installing Python soft- ware. For an extension module arising from Cython code, the followingsetup.py script is all we need to build and install the module:
from distutils.core import setup
from distutils.extension import Extension from Cython.Distutils import build_ext cymodule = ’wave2D_u0_loop_cy’
setup(
name=cymodule
ext_modules=[Extension(cymodule, [cymodule + ’.pyx’],)], cmdclass={’build_ext’: build_ext},
)
We run the script by
Terminal
Terminal> python setup.py build_ext --inplace
The–inplaceoption makes the extension module available in the current directory as the filewave2D_u0_loop_cy.so. This file acts as a normal Python module that can be imported and inspected:
>>> import wave2D_u0_loop_cy
>>> dir(wave2D_u0_loop_cy)
[’__builtins__’, ’__doc__’, ’__file__’, ’__name__’,
’__package__’, ’__test__’, ’advance’, ’np’]
The important output from thedirfunction is our Cython functionadvance(the module also features the importednumpymodule under the name npas well as many standard Python objects with double underscores in their names).
480 C Software Engineering; Wave Equation Model
Thesetup.pyfile makes use of thedistutilspackage in Python and Cython’s extension of this package. These tools know how Python was built on the com- puter and will use compatible compiler(s) and options when building other code in Cython, C, or C++. Quite some experience with building large program systems is needed to do the build process manually, so using asetup.pyscript is strongly recommended.
Simplified build of a Cython module
When there is no need to link the C code with special libraries, Cython offers a shortcut for generating and importing the extension module:
import pyximport; pyximport.install()
This makes thesetup.pyscript redundant. However, in thewave2D_u0_adv.py code we do not usepyximportand require an explicit build process of this and many other modules.
C.5.4 Calling the Cython Function from Python
Thewave2D_u0_loop_cymodule contains ouradvancefunction, which we now may call from the Python program for the wave equation:
import wave2D_u0_loop_cy
advance = wave2D_u0_loop_cy.advance ...
for n in It[1:-1]: # time loop
f_a[:,:] = f(xv, yv, t[n]) # precompute, size as u u = advance(u, u_n, u_nm1, f_a, x, y, t, Cx2, Cy2, dt2)
Efficiency For a mesh consisting of120120cells, the scalar Python code requires 1370 CPU time units, the vectorized version requires 5.5, while the Cython version requires only 1! For a smaller mesh with6060cells Cython is about 1000 times faster than the scalar Python code, and the vectorized version is about 6 times slower than the Cython version.