From 1a595b3a67c592f7dedd7d8769165db4706f044a Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 27 Mar 2012 00:40:28 -0400
Subject: [PATCH] Add Substitution and Derivative primitives, for ease of
 transcribing sympy expressions.

---
 pymbolic/mapper/__init__.py    | 27 +++++++++++++++++++--------
 pymbolic/mapper/stringifier.py | 15 +++++++++++++++
 pymbolic/primitives.py         | 29 ++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/pymbolic/mapper/__init__.py b/pymbolic/mapper/__init__.py
index a38e9c2..9703993 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 e40bad7..801bbd4 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 09c51a7..6674e35 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):
-- 
GitLab