Best practice for using common subexpression elimination with lambdify in SymPy
Question: I'm currently attempting to use SymPy to generate and numerically evaluate a
function and its gradient. For simplicity, I'll use the following function as
an example (keeping in mind that the real function is much lengthier):

    
    
    import sympy as sp
    def g(x):
        return sp.cos(x) + sp.cos(x)**2 + sp.cos(x)**3
    

It's easy enough to numerically evaluate this function and its derivative:

    
    
    import numpy as np
    g_expr = sp.lambdify(x,g(x),modules='numpy')
    dg_expr = sp.lambdify(x,sp.diff(g(x)),modules='numpy')
    
    print g_expr(np.linspace(0,1,50))
    print dg_expr(np.linspace(0,1,50))
    

For my real function, however, lambdify is slow, both in terms of generating
the numerical function and in its evaluation. As many elements in my function
are similar, I'd like to use common subexpression elimination (cse) within
lambdify to speed this process up. I know that SymPy has a built-in function
to perform cse,

    
    
    >>> print sp.cse(g(x))
    ([(x0, cos(x))], [x0**3 + x0**2 + x0])
    

but don't know what syntax to use in order to utilize this result in my
lambdify function (where I'd still like to use x as my input argument):

    
    
    >>> g_expr_fast = sp.lambdify(x,sp.cse(g(x)),modules='numpy')
    >>> print g_expr_fast(np.linspace(0,1,50))
    Traceback (most recent call last):
      File "test3.py", line 34, in <module>
        print g_expr1(nx1)
      File "<string>", line 1, in <lambda>
    NameError: global name 'x0' is not defined
    

Any help on how to use cse in lambdify would be appreciated. Or, if there are
better ways to speed up my gradient calculations, I'd appreciate hearing those
as well.

In case it's relevant, I'm using Python 2.7.3 and SymPy 0.7.6.


Answer: So this might not be the most optimal way to do it, but for my small example
it works.

The idea of following code is lambdifying each common subexpression and
generate a new function with possibly all arguments. I added a few additional
sin and cos terms to add possible dependencies from previous subexpressions.

    
    
    import sympy as sp
    import sympy.abc
    import numpy as np
    import matplotlib.pyplot as pl
    
    def g(x):
        return sp.cos(x) + sp.cos(x)**2 + sp.cos(x)**3 + sp.sin(sp.cos(x)+sp.sin(x))**4 + sp.sin(x) - sp.cos(3*x) + sp.sin(x)**2
    
    repl, redu=sp.cse(g(sp.abc.x))
    
    funs = []
    syms = [sp.abc.x]
    for i, v in enumerate(repl):
        funs.append(sp.lambdify(syms,v[1],modules='numpy'))
        syms.append(v[0])
    
    glam = sp.lambdify(syms,redu[0],modules='numpy')
    
    x = np.linspace(-1,5,10)
    xs=[x]
    
    for f in funs:
        xs.append(f(*xs))
    
    print glam(*xs)
    glamlam = sp.lambdify(sp.abc.x,g(sp.abc.x),modules='numpy')
    print glamlam(x)
    print np.allclose(glamlam(x),glam(*xs))
    

repl contains:

    
    
    [(x0, cos(x)), (x1, sin(x)), (x2, x0 + x1)]
    

and redu contains

    
    
    [x0**3 + x0**2 + x1**2 + x2 + sin(x2)**4 - cos(3*x)]
    

So `funs` contains all subexpression lambdified and the list `xs` contains
each subexpression evaluated, such one can feed `glam` properly in the end.
`xs` grows with each subexpression and might turn out to be a bottle neck in
the end.

You can do the same approach on the expression of
`sp.cse(sp.diff(g(sp.abc.x)))`.

