56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
|
""" Functions to support rewriting of SymPy expressions """
|
||
|
|
|
||
|
|
from sympy.core.expr import Expr
|
||
|
|
from sympy.assumptions import ask
|
||
|
|
from sympy.strategies.tools import subs
|
||
|
|
from sympy.unify.usympy import rebuild, unify
|
||
|
|
|
||
|
|
def rewriterule(source, target, variables=(), condition=None, assume=None):
|
||
|
|
""" Rewrite rule.
|
||
|
|
|
||
|
|
Transform expressions that match source into expressions that match target
|
||
|
|
treating all ``variables`` as wilds.
|
||
|
|
|
||
|
|
Examples
|
||
|
|
========
|
||
|
|
|
||
|
|
>>> from sympy.abc import w, x, y, z
|
||
|
|
>>> from sympy.unify.rewrite import rewriterule
|
||
|
|
>>> from sympy import default_sort_key
|
||
|
|
>>> rl = rewriterule(x + y, x**y, [x, y])
|
||
|
|
>>> sorted(rl(z + 3), key=default_sort_key)
|
||
|
|
[3**z, z**3]
|
||
|
|
|
||
|
|
Use ``condition`` to specify additional requirements. Inputs are taken in
|
||
|
|
the same order as is found in variables.
|
||
|
|
|
||
|
|
>>> rl = rewriterule(x + y, x**y, [x, y], lambda x, y: x.is_integer)
|
||
|
|
>>> list(rl(z + 3))
|
||
|
|
[3**z]
|
||
|
|
|
||
|
|
Use ``assume`` to specify additional requirements using new assumptions.
|
||
|
|
|
||
|
|
>>> from sympy.assumptions import Q
|
||
|
|
>>> rl = rewriterule(x + y, x**y, [x, y], assume=Q.integer(x))
|
||
|
|
>>> list(rl(z + 3))
|
||
|
|
[3**z]
|
||
|
|
|
||
|
|
Assumptions for the local context are provided at rule runtime
|
||
|
|
|
||
|
|
>>> list(rl(w + z, Q.integer(z)))
|
||
|
|
[z**w]
|
||
|
|
"""
|
||
|
|
|
||
|
|
def rewrite_rl(expr, assumptions=True):
|
||
|
|
for match in unify(source, expr, {}, variables=variables):
|
||
|
|
if (condition and
|
||
|
|
not condition(*[match.get(var, var) for var in variables])):
|
||
|
|
continue
|
||
|
|
if (assume and not ask(assume.xreplace(match), assumptions)):
|
||
|
|
continue
|
||
|
|
expr2 = subs(match)(target)
|
||
|
|
if isinstance(expr2, Expr):
|
||
|
|
expr2 = rebuild(expr2)
|
||
|
|
yield expr2
|
||
|
|
return rewrite_rl
|