diff options
Diffstat (limited to 'boltzgen/utility/optimizations.py')
| -rw-r--r-- | boltzgen/utility/optimizations.py | 10 |
1 files changed, 10 insertions, 0 deletions
diff --git a/boltzgen/utility/optimizations.py b/boltzgen/utility/optimizations.py new file mode 100644 index 0000000..93dad09 --- /dev/null +++ b/boltzgen/utility/optimizations.py @@ -0,0 +1,10 @@ +from sympy import * + +from sympy.codegen.rewriting import ReplaceOptim + +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 |