aboutsummaryrefslogtreecommitdiff
path: root/symbolic/optimizations.py
diff options
context:
space:
mode:
Diffstat (limited to 'symbolic/optimizations.py')
-rw-r--r--symbolic/optimizations.py10
1 files changed, 10 insertions, 0 deletions
diff --git a/symbolic/optimizations.py b/symbolic/optimizations.py
new file mode 100644
index 0000000..93dad09
--- /dev/null
+++ b/symbolic/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