diff options
Diffstat (limited to 'boltzgen/utility/optimizations.py')
| -rw-r--r-- | boltzgen/utility/optimizations.py | 12 |
1 files changed, 10 insertions, 2 deletions
diff --git a/boltzgen/utility/optimizations.py b/boltzgen/utility/optimizations.py index 93dad09..6dc23e9 100644 --- a/boltzgen/utility/optimizations.py +++ b/boltzgen/utility/optimizations.py @@ -2,9 +2,17 @@ from sympy import * from sympy.codegen.rewriting import ReplaceOptim -expand_square = ReplaceOptim( +expand_pos_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 +expand_neg_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_pos_square, expand_pos_square), + (expand_neg_square, expand_neg_square) +] + cse_main.basic_optimizations |