diff --git a/benchmarks/cse.py b/benchmarks/cse.py index 112753a..751895a 100644 --- a/benchmarks/cse.py +++ b/benchmarks/cse.py @@ -1,9 +1,29 @@ +"""Benchmarks for common subexpression elimination (CSE).""" + from functools import reduce from operator import add -import sympy as sp + +from sympy import cse, exp, sin, symbols, tan, Matrix + def _get_args_exprs(nexprs, nterms): - x = sp.symbols('x:%d' % nterms) + """Construct expressions with a number of symbols based on chained additions + and subtractions. + + Examples + ======== + + >>> _get_args_exprs(1, 1) + ((x0,), [x0]) + >>> _get_args_exprs(1, 3) + ((x0, x1, x2), [x0 - x1 + x2]) + >>> _get_args_exprs(3, 1) + ((x0,), [x0, -x0, x0]) + >>> _get_args_exprs(3, 3) + ((x0, x1, x2), [x0 - x1 + x2, -x0 + x1 - x2, x0 - x1 + x2]) + + """ + x = symbols('x:%d' % nterms) exprs = [ reduce(add, [x[j]*(-1)**(i+j) for j in range(nterms)]) for i in range(nexprs) ] @@ -12,13 +32,208 @@ def _get_args_exprs(nexprs, nterms): class TimeCSE: - params = (( - (2, 8), (3, 8), (4, 8), - (2, 9), (2, 10), (2, 11) - ),) + params = (((2, 8), (3, 8), (4, 8), (2, 9), (2, 10), (2, 11), ), ) def setup(self, nexprs_nterms): self.args, self.exprs = _get_args_exprs(*nexprs_nterms) def time_cse(self, nexprs_nterms): - sp.cse(self.exprs) + cse(self.exprs) + + +class GriewankBabyExampleCSE: + """Simple test function with multiple repeated subexpressions. + + Function is the "baby example" taken from Griewank, A., & Walther, A. + (2008). Evaluating derivatives: principles and techniques of algorithmic + differentiation. SIAM. + + The baby example function expressed as a tree structure is as follows: + + (MUL) + _______|______ + / \ + (SUB) (SUB) + ___|___ __|__ + / \ / \ + (ADD) (EXP) (DIV) (EXP) + ___|___ | _|_ | + / \ | / \ | + (SIN) (DIV) (x2) (x1) (x2) (x2) + | _|_ + | / \ + (DIV) (x1) (x2) + _|_ + / \ + (x1) (x2) + + It involves 1 ADD, 3 DIV, 2 EXP, 1 MUL, 1 SIN, and 2 SUB operations. + + When CSE is conducted on the expression, the following intermediate terms + are introduced: + + @0: EXP[x2] + @1: DIV[x1, x2] + @2: SUB[@1, @0] + @3: SIN[@1] + @4: ADD[@2, @3] + @5: MUL[@2, @4] + + @0, @3, and @4 are operands in only one operation each so can be collapsed. + @1 and @2 are operands in two operations each so are required common + subexpressions. @1 will become the first common subexpression (x0, x1 / x2) + and @2 will become the second common subexpression (x3, x0 - exp(x2)). @6 is + the expression itself so can used directly there resulting in the + substituted expression x3 * (x3 + sin(x0)). + + This example highlights an important optimisation around operation ordering + in CSE. The @4 operation could instead be constructed as + ADD[@1, SUB[@3, @0]]. This is suboptimal CSE as it would require the + introduction of an additional unnecessary operation. + + """ + + def setup(self): + """Create the required symbols (x1, x2), the expression (y), and its + Jacobian (G). + + G is the 1x2 matrix that contains the derivatives of y with respect to + x1 and x2 as the two columns respectively. + + """ + x1, x2 = symbols("x1, x2") + x = Matrix([x1, x2]) + self.y = (sin(x1 / x2) + (x1 / x2) - exp(x2)) * ((x1 / x2) - exp(x2)) + self.G = self.y.diff(x) + + def test_function_cse(self): + """Expected result from CSE on the baby example.""" + x0 = sym.Symbol("x0") + x1 = sym.Symbol("x1") + x2 = sym.Symbol("x2") + x3 = sym.Symbol("x3") + + cse = [ + (x0, x1 / x2), + (x3, x0 - exp(x2)), + ] + expr = [ + x3 * (x3 + sin(x0)), + ] + + assert cse(self.y) == (cse, expr) + + def time_function_cse(self): + """Time CSE on the baby example.""" + cse(self.y) + + def time_jacobian_cse(self): + """Time CSE on the baby example's Jacobian.""" + cse(self.G) + + def time_combined_cse(self): + """Time simultaneous CSE on the baby example and its Jacobian.""" + cse([self.y, self.G]) + + +class GriewankLighthouseExampleCSE: + """Simple matrix test function with multiple repeated subexpressions. + + Function is the "lighthouse example" taken from Griewank, A., & Walther, A. + (2008). Evaluating derivatives: principles and techniques of algorithmic + differentiation. SIAM. + + The lighthouse example function expressed as a tree structure is as follows: + + (MATRIX) + __________|___________ + / \ + (MUL) (MUL) + __|__ _|_ + / \ / \ + (MUL) (gamma) (DIV) (nu) + _|_ ______|_____ + / \ / \ + (DIV) (nu) (TAN) (SUB) + ______|_____ | __|__ + / \ | / \ + (TAN) (SUB) (MUL) (gamma) (TAN) + | __|__ _|_ | + | / \ / \ | + (MUL) (gamma) (TAN) (omega) (t) (MUL) + _|_ | _|_ + / \ | / \ + (omega) (t) (MUL) (omega) (t) + _|_ + / \ + (omega) (t) + + It involves 2 DIV, 7 MUL, 2 SUB, and 4 TAN operations. Additionally, one + matrix population is required. + + When CSE is conducted on the expression, the following intermediate terms + are introduced: + + @0: MUL[omega, t] + @1: TAN[@0] + @2: SUB[gamma, @1] + @3: DIV[@1, @2] + @4: MUL[nu, @3] + @5: MUL[gamma, @4] + + @0, @2, and @3 are operands in only one operation each so can be collapsed. + @1 is an operand in four operations and @4 is both a required expression and + an operand in an expression so both are required common subexpressions. @1 + will become the first common subexpression (x0, tan(omega * t)) and @4 will + become the second common subexpression (x1, nu * x0 / (gamma - x0)). @5 is + one of the required expressions in the matrix so can used directly there + resulting in the substituted matrix expression Matrix([x1, gamma * x1]). + + """ + + def setup(self): + """Create the required symbols (nu, gamma, omega, t), the matrix of + expressions (y), and its Jacobian (G). + + G is the 2x4 matrix that contains the derivatives of y with respect to + nu, gamma, omega, and t as the four columns respectively. + + """ + nu, gamma, omega, t = symbols("nu, gamma, omega, t") + x = Matrix([nu, gamma, omega, t]) + self.y = Matrix([ + nu * tan(omega * t) / (gamma - tan(omega * t)), + nu * gamma * tan(omega * t) / (gamma - tan(omega * t)), + ]) + self.G = self.y.jacobian(x) + + def test_function_cse(self): + """Expected result from CSE on the lighthouse example.""" + nu = sym.Symbol("nu") + gamma = sym.Symbol("gamma") + omega = sym.Symbol("omega") + t = sym.Symbol("t") + x0 = sym.Symbol("x0") + x1 = sym.Symbol("x1") + + cse = [ + (x0, tan(omega * t)), + (x1, nu * x0 / (gamma - x0)), + ] + expr = [ + Matrix([x1, gamma * x1]), + ] + + assert cse(self.y) == (cse, expr) + + def time_function_cse(self): + """Time CSE on the lighthouse example.""" + cse(self.y) + + def time_jacobian_cse(self): + """Time CSE on the lighthouse example's Jacobian.""" + cse(self.G) + + def time_combined_cse(self): + """Time simultaneous CSE on the lighthouse example and its Jacobian.""" + cse([self.y, self.G])