aboutsummaryrefslogtreecommitdiff
path: root/boltzgen/utility/optimizations.py
blob: 6dc23e935928e7936b8732e1f9ec41215f5b5285 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from sympy import *

from sympy.codegen.rewriting import 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))
)

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