Reverse-over-forward AD: FunctionAssignBlock. Minor test edits.

Author: jrmaddison

firedrake/adjoint_utils/blocks/function.py

@@ -129,6 +129,25 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
 
         return dudm
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        expr = self.expr
+
+        tlm_rhs = 0
+        for block_variable in self.get_dependencies():
+            dep = block_variable.output
+            tlm_dep = block_variable.tlm_value
+            if tlm_dep is not None:
+                tlm_rhs = tlm_rhs + ufl.derivative(expr, dep, tlm_dep)
+                
+        x.tlm_value = None
+        if isinstance(tlm_rhs, int) and tlm_rhs == 0:
+            return
+        tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+        if isinstance(tlm_rhs, ufl.constantvalue.Zero):
+            return
+        x.tlm_value = firedrake.Function(x.output.function_space()).assign(tlm_rhs)
+
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
         return self.prepare_evaluate_adj(inputs, hessian_inputs,

tests/regression/test_adjoint_reverse_over_forward.py

@@ -34,17 +34,17 @@ def test_assembly():
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
 
-    u = Function(space, name="u").interpolate(Constant(1.0))
+    u = Function(space, name="u").interpolate(X[0])
     zeta = Function(space, name="tlm_u").interpolate(X[0])
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
     with reverse_over_forward():
-        J = assemble(u * u * dx)
+        J = assemble(u * u.dx(0) * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)
+                       assemble(inner(zeta.dx(0), test) * dx + inner(zeta, test.dx(0)) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -57,7 +57,28 @@ def test_constant_assignment():
 
     assert float(b.block_variable.tlm_value) == -2.0
 
-    # Minimal test that the TLM operations are on the tape
+    # Minimal test that the TLM operation is on the tape
     _ = compute_gradient(b.block_variable.tlm_value, Control(a.block_variable.tlm_value))
     adj_value = a.block_variable.tlm_value.block_variable.adj_value
     assert float(adj_value) == 1.0
+
+
+@pytest.mark.skipcomplex
+def test_function_assignment():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    u = Function(space, name="u").interpolate(X[0] - 0.5)
+    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+    with reverse_over_forward():
+        v = Function(space, name="v").assign(-3 * u)
+        J = assemble(v * v.dx(0) * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(9 * inner(zeta.dx(0), test) * dx + 9 * inner(zeta, test.dx(0)) * dx).dat.data_ro)