symbolic/optimizations.py


1
2
3
4
5
6
7
8
9
10
11

from sympy import *

from sympy.codegen.rewriting import ReplaceOptim
from sympy.simplify import cse_main

expand_square = ReplaceOptim(
    lambda e: e.is_Pow and e.exp.is_integer and e.exp == 2,
    lambda p: UnevaluatedExpr(Mul(p.base, p.base, evaluate = False))
)

custom = [ (expand_square, expand_square) ] + cse_main.basic_optimizations