How to get a fast lambda function from an sympy expression in 3 dimensions?

Question

I am using sympy to generate different expressions for cfd-simulations. Mostly these expressions are of the kind exp = f(x,y,z) for example f(x,y,z) = sin(x)*cos(y)*sin(z). To get values on a grid I use simpy.lambdify. For example:

import numpy as np
import sympy as sp
from sympy.abc import x,y,z

xg, yg, zg = np.mgrid[0:1:50*1j, 0:1:50*1j, 0:1:50*1j]
f  = sp.sin(x)*sp.cos(y)*sp.sin(z)

lambda_f = sp.lambdify([x,y,z], f, "numpy")
fn = lambda_f(xg, yg, zg)

print fn

This seems to work pretty good but unfortunately my expressions are getting more and more complex and the grid computation takes a lot of time. My Idea was that it is maybe possible to use the uFuncify method (see http://docs.sympy.org/latest/modules/numeric-computation.html ) to speed the computations up, im not sure if this is the right way ? And im also not sure how to get ufunctify to work for 3d grids ? Thanks for any Suggestions

Answer 1

In previous releases (0.7.5 and prior), ufuncify only worked on single dimension arrays for the first argument (not very exciting). As of 0.7.6 (not released yet, but should be in a week!) ufuncify creates actual instances of numpy.ufunc by default (wraps C code in numpy api). Your code above only needs a small change to make it work.

In [1]: import numpy as np

In [2]: from sympy import sin, cos, lambdify

In [3]: from sympy.abc import x,y,z

In [4]: from sympy.utilities.autowrap import ufuncify

In [5]: from sympy.printing.theanocode import theano_function

In [6]: xg, yg, zg = np.mgrid[0:1:50*1j, 0:1:50*1j, 0:1:50*1j]

In [7]: f  = sym.sin(x)*sym.cos(y)*sym.sin(z)

In [8]: ufunc_f = ufuncify([x,y,z], f)

In [9]: theano_f = theano_function([x, y, z], f, dims={x: 3, y: 3, z: 3})

In [10]: lambda_f = lambdify([x, y, z], f)

In [11]: type(ufunc_f)
Out[11]: numpy.ufunc

In [12]: type(theano_f)
Out[12]: theano.compile.function_module.Function

In [13]: type(lambda_f)
Out[13]: function

In [14]: %timeit ufunc_f(xg, yg, zg)
10 loops, best of 3: 21 ms per loop

In [15]: %timeit theano_f(xg, yg, zg)
10 loops, best of 3: 20.7 ms per loop

In [16]: %timeit lambda_f(xg, yg, zg)
10 loops, best of 3: 22.3 ms per loop

ufuncify and theano_function are comparable, and slightly faster than lambdify for this simple expression. The difference is greater using the more complicated expression given below:

In [17]: f = sin(x)*cos(y)*sin(z) + sin(4*(x - y**2*sin(z)))

In [18]: ufunc_f = ufuncify([x,y,z], f)

In [19]: theano_f = theano_function([x, y, z], f, dims={x: 3, y: 3, z: 3})

In [20]: lambda_f = lambdify([x, y, z], f)

In [21]: %timeit ufunc_f(xg, yg, zg)
10 loops, best of 3: 29.2 ms per loop

In [22]: %timeit theano_f(xg, yg, zg)
10 loops, best of 3: 29.2 ms per loop

In [23]: %timeit lambda_f(xg, yg, zg)
10 loops, best of 3: 42.1 ms per loop

This is much faster than using the python version, as no intermediate arrays are created, the loop is traversed and the calculation ran in C. Theano produces equivalent speeds, as they also compile to native code. For the large expressions that I see when doing multibody dynamics, ufuncify (and the related autowrap) perform significantly faster than lambdify. I don't have much experience with theano, so I can't say how well their approach scales either, but I'd assume it would be similar.

As I said above, this is only available in sympy 0.7.6 and up. Should be released soon, but until then you can grab the source from github. Docs on ufuncify new behavior can be found here