Add expression replacement via dictionary

Author: kburns

dedalus/core/field.py

@@ -212,6 +212,10 @@ def replace(self, old, new):
         """Replace specified operand/operator."""
         raise NotImplementedError()
 
+    def replace_dict(self, subs):
+        """Replace specified operands/operators according to a dictionary."""
+        raise NotImplementedError()
+
     def sym_diff(self, var):
         """Symbolically differentiate with respect to specified operand."""
         raise NotImplementedError()
@@ -284,9 +288,8 @@ def frechet_differential(self, variables, perturbations, backgrounds=None):
         # Compute differential
         epsilon = Field(dist=dist, dtype=dtype)
         # d/dε F(X0 + ε*X1)
-        diff = self
-        for var, pert in zip(variables, perturbations):
-            diff = diff.replace(var, var + epsilon*pert)
+        subs = {var: var + epsilon*pert for var, pert in zip(variables, perturbations)}
+        diff = self.replace_dict(subs)
         diff = diff.sym_diff(epsilon)
         # ε -> 0
         if diff:
@@ -295,8 +298,8 @@ def frechet_differential(self, variables, perturbations, backgrounds=None):
         # Replace variables with backgrounds, if specified
         if diff:
             if backgrounds:
-                for var, bg in zip(variables, backgrounds):
-                    diff = diff.replace(var, bg)
+                subs = {var: bg for var, bg in zip(variables, backgrounds)}
+                diff = diff.replace_dict(subs)
         return diff
 
     @property
@@ -378,6 +381,13 @@ def replace(self, old, new):
         else:
             return self
 
+    def replace_dict(self, subs):
+        """Replace specified operands/operators according to a dictionary."""
+        if self in subs:
+            return subs[self]
+        else:
+            return self
+
     def sym_diff(self, var):
         """Symbolically differentiate with respect to specified operand."""
         if self == var:

dedalus/core/future.py

@@ -128,6 +128,20 @@ def replace(self, old, new):
             args = [arg.replace(old, new) if isinstance(arg, Operand) else arg for arg in self.args]
             return self.new_operands(*args)
 
+    def replace_dict(self, subs):
+        """Replace specified operands/operators according to a dictionary."""
+        # Check for entire expression match
+        if self in subs:
+            return subs[self]
+        # Check base and call with replaced arguments
+        elif type(self) in subs:
+            args = [arg.replace_dict(subs) if isinstance(arg, Operand) else arg for arg in self.args]
+            return subs[type(self)](*args)
+        # Call with replaced arguments
+        else:
+            args = [arg.replace_dict(subs) if isinstance(arg, Operand) else arg for arg in self.args]
+            return self.new_operands(*args)
+
     # def simplify(self, *vars):
     #     """Simplify expression, except subtrees containing specified variables."""
     #     # Simplify arguments if variables are present

dedalus/core/operators.py

@@ -701,6 +701,20 @@ def replace(self, old, new):
             new_operand = self.operand.replace(old, new)
             return self.new_operand(new_operand)
 
+    def replace_dict(self, subs):
+        """Replace specified operands/operators according to a dictionary."""
+        # Check for entire expression match
+        if self in subs:
+            return subs[self]
+        # Check base and call with replaced arguments
+        elif type(self) in subs:
+            new_operand = self.operand.replace_dict(subs)
+            return subs[type(self)](new_operand)
+        # Call with replaced arguments
+        else:
+            new_operand = self.operand.replace_dict(subs)
+            return self.new_operand(new_operand)
+
     def expand(self, *vars):
         """Expand expression over specified variables."""
         from .arithmetic import Add, Multiply
@@ -1611,6 +1625,11 @@ def replace(self, old, new):
         # Replace operand, skipping conversion
         return self.operand.replace(old, new)
 
+    def replace_dict(self, subs):
+        """Replace specified operands/operators according to a dictionary."""
+        # Replace operand, skipping conversion
+        return self.operand.replace_dict(subs)
+
     # def expand(self, *vars):
     #     """Expand expression over specified variables."""
     #     # Expand operand, skipping conversion