TimeDerivativeRuleDispatcher: handle generic UFL types (#132)

* Fix TimeDerivativeRuleDispatcher

* add a test

Author: pbrubeck

irksome/deriv.py

@@ -79,12 +79,7 @@ def time_derivative(self, o):
         f, = o.ufl_operands
         return map_expr_dag(self.rules, f)
 
-    expr = MultiFunction.reuse_if_untouched
-    terminal = expr
-    derivative = expr
-    grad = expr
-    curl = expr
-    div = expr
+    ufl_type = MultiFunction.reuse_if_untouched
 
 
 def apply_time_derivatives(expression, t=None, timedep_coeffs=None):

tests/test_differentiation.py

@@ -1,7 +1,7 @@
 import pytest
 from ufl.algorithms import expand_derivatives
 from irksome import MeshConstant, Dt, expand_time_derivatives
-from firedrake import Constant, diff, dot, FunctionSpace, Function, sin, UnitIntervalMesh, VectorFunctionSpace
+from firedrake import Cofunction, Constant, diff, dot, dx, FunctionSpace, Function, inner, sin, TestFunction, UnitIntervalMesh, VectorFunctionSpace
 
 
 @pytest.fixture
@@ -79,3 +79,16 @@ def test_expand_second_derivative_product_rule(V):
                 + dot(u, Dt(w, 2)))
     # UFL equality is failing here due to different index numbers
     assert str(expr) == str(expected)
+
+
+def test_cofunction(V):
+    u = Function(V)
+    v = TestFunction(V)
+    c = Cofunction(V.dual())
+    expr = inner(Dt(u), v)*dx + c
+
+    # TimeDerivatives are already expanded
+    # This is just to test that it can handle Cofunction properly
+    expected = expr
+    expr = expand_time_derivatives(expr)
+    assert expr == expected