Spaces:
Sleeping
Sleeping
File size: 1,798 Bytes
6a86ad5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
""" 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
|