Merge pull request #727 from zwicker-group/expression_consts

Fix ScalarExpression and TensorExpression to support constants in differentiation and __getitem__

Author: david-zwicker

pde/tools/expressions.py

@@ -714,6 +714,7 @@ def differentiate(self, var: str) -> ScalarExpression:
             signature=self.vars,
             allow_indexed=self.allow_indexed,
             user_funcs=self.user_funcs,
+            consts=self.consts,
         )
 
     @cached_property()
@@ -732,7 +733,10 @@ def derivatives(self) -> TensorExpression:
 
         grad = sympy.Array([self._sympy_expr.diff(sympy.Symbol(v)) for v in self.vars])
         return TensorExpression(
-            sympy.simplify(grad), signature=self.vars, user_funcs=self.user_funcs
+            sympy.simplify(grad),
+            signature=self.vars,
+            user_funcs=self.user_funcs,
+            consts=self.consts,
         )
 
 
@@ -788,6 +792,10 @@ def __init__(
                 user_funcs = expression.user_funcs
             else:
                 user_funcs.update(expression.user_funcs)
+            if consts is None:
+                consts = expression.consts
+            else:
+                consts.update(expression.consts)
 
         elif isinstance(expression, (np.ndarray, list, tuple)):
             # expression is a constant array
@@ -834,11 +842,17 @@ def __getitem__(self, index):
         expr = self._sympy_expr[index]
         if isinstance(expr, sympy.Array):
             return TensorExpression(
-                expr, signature=self.vars, user_funcs=self.user_funcs
+                expr,
+                signature=self.vars,
+                user_funcs=self.user_funcs,
+                consts=self.consts,
             )
         else:
             return ScalarExpression(
-                expr, signature=self.vars, user_funcs=self.user_funcs
+                expr,
+                signature=self.vars,
+                user_funcs=self.user_funcs,
+                consts=self.consts,
             )
 
     @property
@@ -870,7 +884,9 @@ def differentiate(self, var: str) -> TensorExpression:
             derivative = np.zeros(self.shape)
         else:
             derivative = self._sympy_expr.diff(sympy.Symbol(var))
-        return TensorExpression(derivative, self.vars, user_funcs=self.user_funcs)
+        return TensorExpression(
+            derivative, self.vars, user_funcs=self.user_funcs, consts=self.consts
+        )
 
     @cached_property()
     def derivatives(self) -> TensorExpression:
@@ -885,7 +901,9 @@ def derivatives(self) -> TensorExpression:
             dx = sympy.Array([sympy.Symbol(s) for s in self.vars])
             derivatives = sympy.derive_by_array(self._sympy_expr, dx)
 
-        return TensorExpression(derivatives, self.vars, user_funcs=self.user_funcs)
+        return TensorExpression(
+            derivatives, self.vars, user_funcs=self.user_funcs, consts=self.consts
+        )
 
     def get_compiled_array(
         self, single_arg: bool = True

tests/tools/test_expressions.py

@@ -348,6 +348,23 @@ def test_expression_consts():
     np.testing.assert_allclose(expr(np.array([2, 3])), np.array([3, 5]))
     np.testing.assert_allclose(expr.get_compiled()(np.array([2, 3])), np.array([3, 5]))
 
+    expr = ScalarExpression("a * b", consts={"a": np.array([1, 2])})
+    dexpr_da = expr.differentiate("b")
+    np.testing.assert_allclose(dexpr_da(np.array([2, 3])), np.array([1, 2]))
+    dexpr = expr.derivatives
+    assert dexpr.shape == (1,)
+    np.testing.assert_allclose(dexpr(np.array([2, 3])), np.array([[1, 2]]))
+
+
+def test_tensor_expression_consts():
+    """Test the usage of consts in TensorExpression."""
+    e = TensorExpression("[a, a*b]", consts={"b": 5})
+    assert e[0](2) == 2
+    assert e[1](2) == 10
+    d1 = e.differentiate("a")
+    assert d1[0](2) == 1
+    assert d1[1](2) == 5
+
 
 def test_evaluate_func_scalar():
     """Test the evaluate function with scalar fields."""