diff --git a/pymbolic/mapper/__init__.py b/pymbolic/mapper/__init__.py index a38e9c296d2c56a4f503924524b9e1d9475dda9e..97039938e616e654799fc9f5c8bef0d001af989c 100644 --- a/pymbolic/mapper/__init__.py +++ b/pymbolic/mapper/__init__.py @@ -268,17 +268,28 @@ class IdentityMapperBase(object): expr.prefix, **expr.get_extra_properties()) - def map_if_positive(self, expr): + def map_substitution(self, expr, *args): return type(expr)( - self.rec(expr.criterion), - self.rec(expr.then), - self.rec(expr.else_)) + self.rec(expr.child, *args), + expr.variables, + tuple(self.rec(v, *args) for v in expr.values)) - def map_if(self, expr): + def map_derivative(self, expr, *args): + return type(expr)( + self.rec(expr.child, *args), + expr.variables) + + def map_if_positive(self, expr, *args): + return type(expr)( + self.rec(expr.criterion, *args), + self.rec(expr.then, *args), + self.rec(expr.else_, *args)) + + def map_if(self, expr, *args): return type(expr)( - self.rec(expr.condition), - self.rec(expr.then), - self.rec(expr.else_)) + self.rec(expr.condition, *args), + self.rec(expr.then, *args), + self.rec(expr.else_, *args)) diff --git a/pymbolic/mapper/stringifier.py b/pymbolic/mapper/stringifier.py index e40bad7f16981deae44b49a7161f387e84f980c5..801bbd49387d4b3ccc9a8e7e0d7b5ac21952f026 100644 --- a/pymbolic/mapper/stringifier.py +++ b/pymbolic/mapper/stringifier.py @@ -200,6 +200,21 @@ class StringifyMapper(pymbolic.mapper.RecursiveMapper): map_max = map_min + def map_derivative(self, expr, enclosing_prec): + derivs = " ".join( + "d/d%s" % v + for v in expr.variables) + + return "%s %s" % ( + derivs, self.rec(expr.child, PREC_PRODUCT)) + + def map_substitution(self, expr, enclosing_prec): + substs = ", ".join( + "%s=%s" % (name, self.rec(val, PREC_NONE)) + for name, val in zip(expr.variables, expr.values)) + + return "[%s]{%s}" % (self.rec(expr.child, PREC_NONE), substs) + def __call__(self, expr, prec=PREC_NONE): return pymbolic.mapper.RecursiveMapper.__call__(self, expr, prec) diff --git a/pymbolic/primitives.py b/pymbolic/primitives.py index 09c51a78d205822d40e37cb54caeb832f1763eaf..6674e351cdcb664a61d0c331158bc9f19520e5e5 100644 --- a/pymbolic/primitives.py +++ b/pymbolic/primitives.py @@ -690,11 +690,38 @@ class CommonSubexpression(Expression): mapper_method = intern("map_common_subexpression") -# }}} + + +class Substitution(Expression): + """Work-alike of sympy's Subs.""" + + def __init__(self, child, variables, values): + self.child = child + self.variables = variables + self.values = values + + def __getinitargs__(self): + return (self.child, self.variables, self.values) + + mapper_method = intern("map_substitution") +class Derivative(Expression): + """Work-alike of sympy's Derivative.""" + + def __init__(self, child, variables): + self.child = child + self.variables = variables + + def __getinitargs__(self): + return (self.child, self.variables) + + mapper_method = intern("map_derivative") + +# }}} + # intelligent factory functions ---------------------------------------------- def make_variable(var_or_string): if not isinstance(var_or_string, Expression):