Using Sympy’s Common Subexpression Elimination to generate Code

A nifty feature of sympy is its code generator. This allows generating code in many languages from expressions derived in sympy. I’ve been using this for a while and came across the Common Subexpression Elimination (CSE) function (again), recently.

In earlier* versions of sympy I did not find it of great use. But now I was pleasantly surprised how much it matured. (* I cannot say for sure when I last tried using it; it must’ve been a long while ago.)

What is it and what is it good for?

Let’s say we want to generate code for a PI controller in C. The literature says its transfer function (in DIN notation) reads:

$G(s) = K_p + \frac{K_i}{s}$

Then furthermore the representation for a discrete time system with time step $T_s$ becomes

$G(z) = K_p + \frac{K_i T_s}{2} + \frac{\frac{K_i T_s}{2} - K_p}{z}$

In order to efficiently and flexibly calculate the response we want to implement it in a second order IIR (infinite impulse response) filter structure. This structure is basically a representation of $G_z$ as a rational function

$\frac{ b_0 + b_1 z^{-1} + b_2 z^{-2} }{ a_0 + a_1 z^{-1} + a_2 z^{-2} }$

so we need to take the expression apart such that we can determine the coefficients a_1, a_2, b_0, … (by convention, the problem is scaled such that a_0 = 1.

This is all very cumbersome to do manually (and I suck at doing this, especially as I insert quite embarrassing mistakes all the time). So, sympy to the rescue:

import sympy as sy

Kp, Ki, Ts = sy.symbols("K_p K_i T_s", real=True)
z = sy.symbols("z", complex=True)
G_z = Kp + Ki*Ts/2 + (Ki*Ts/2 - Kp)/z

display(G_z)

n, d = G_z.ratsimp().as_numer_denom()
n, d = n.as_poly(z), d.as_poly(z)

display(n, d)

We first define the symbols needed. Then set up G_z according to our equation above. After displaying it for good measure, we do a rational simplification ratsimp() on it and split it into numerator n and denominator d.

Now, we represent n and d as polynomials of z. And this is how we obtain our coefficients a1, a2, and so on. We continue:

scale = d.LC()

n, d = (n.as_expr()/scale).as_poly(z), (d.as_expr()/scale).as_poly(z)
N, D = n.degree(), d.degree()

We find the scaling factor as the leading coefficient (LC()) of the denominator d, scale both accordingly and determine the degree of each (this is mostly to determine for how many terms we have to iterate when printing etc.).

Bs = n.all_coeffs()
As = d.all_coeffs()

display(Bs)
display(As)

Now, basically Bs contains all coefficients [b0 b1] and likewise for As. Here’s a shortcut, see if you can spot it.

And now comes the cse magic:

commons, exprs = sy.cse(Bs)

This returns a tuple. The first element is a list of tuples of arbitrarily assigned place holder variables x0, x1, and so on. The second element are the original expressions, but the placeholders are substituted. In this example:

(
  [
    (x0, K_i*T_s/2)
  ],
  [K_p + x0, -K_p + x0]
)

Which means, instead of calculating K_i*T_s/2 twice, we can do it once, assign it to x0 and use that on the expressions.

You can even concatenate As and Bs for good measure (not really necessary here, since the As only contain 1 and 0, but in other controller topologies this may change).

commons, exprs = sy.cse(As + Bs)

# Printing C code for instance
for com in commons:
    print( sy.ccode(com[1], assign_to=repr(com[0])) )

for expr, name in zip(exprs, "a0 a1 b0 b1".split(" ")):
    print( sy.ccode(expr, assign_to=f"coeffs.{name:s}") )

We obtain the common terms commons and the reduced expressions exprs, then print the common terms assigning them to their respective placeholder variables, and finally print the reduced expressions. This produces this five liner in C:

x0 = (1.0/2.0)*K_i*T_s;
coeffs.a0 = 1;
coeffs.a1 = 0;
coeffs.b0 = K_p + x0;
coeffs.b1 = -K_p + x0;

Now, how is this useful you ask? Well, for this toy example it certainly is a lot of work for five lines I could’ve just typed out.

But if you go into a PID controller with damped differential term, and you want your code to be efficient (i.e. have as few operations as possible), this is very handy. Just by changing the symbols and the expression for G_z we obtain this:

x0 = 2*T_v;
x1 = 1.0/(T_s + x0);
x2 = 4*T_v;
x3 = T_i*T_s;
x4 = 2*x3;
x5 = T_i*x2;
x6 = 1.0/(x4 + x5);
x7 = K_p*x4;
x8 = K_p*T_s*x0;
x9 = pow(T_s, 2);
x10 = 4*K_p*T_d*T_i + K_p*x5;
x11 = K_p*x9 + x10;
c.a0 = 1;
c.a1 = -x1*x2;
c.a2 = x1*(-T_s + 2*T_v);
c.b0 = x6*(x11 + x7 + x8);
c.b1 = (K_p*x9 - x10)/(T_i*x0 + x3);
c.b2 = x6*(x11 - x7 - x8);

I did this manually once, and only obtained five or six of the placeholders. That’s the power of cse().

Oh, you need the code in Python for testing? No problem, just use pycode instead (and work around that assign_to parameter).

So, there you have it.