Reverse-over-forward AD: SupermeshProjectBlock

Author: jrmaddison

firedrake/adjoint_utils/blocks/solving.py

@@ -842,6 +842,7 @@ def __init__(self, source, target_space, target, bcs=[], **kwargs):
             mesh = target_space.mesh()
         self.source_space = source.function_space()
         self.target_space = target_space
+        self._kwargs = dict(kwargs)
         self.projector = firedrake.Projector(source, target_space, **kwargs)
 
         # Assemble mixed mass matrix
@@ -922,6 +923,15 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
                                              prepared)
         return dJdm
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        dep, = self.get_dependencies()
+        if dep.tlm_value is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = firedrake.Function(x.output.function_space())
+            firedrake.project(dep.tlm_value, x.tlm_value, **self._kwargs)
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx,
                                    relevant_dependencies, prepared=None):

tests/regression/test_adjoint_reverse_over_forward.py

@@ -86,22 +86,44 @@ def test_function_assignment():
 
 @pytest.mark.skipcomplex
 def test_project():
-    mesh = UnitIntervalMesh(10)
+    mesh = UnitSquareMesh(10, 10)
     X = SpatialCoordinate(mesh)
-    space = FunctionSpace(mesh, "Lagrange", 1)
-    test = TestFunction(space)
+    space_a = FunctionSpace(mesh, "Lagrange", 1)
+    space_b = FunctionSpace(mesh, "Discontinuous Lagrange", 0)
+    test_a = TestFunction(space_a)
 
-    u = Function(space, name="u").interpolate(X[0] - 0.5)
-    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+    zeta = Function(space_a, name="tlm_u").interpolate(X[0])
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
-    space_0 = FunctionSpace(mesh, "Discontinuous Lagrange", 0)
+    with reverse_over_forward():
+        v = Function(space_b, name="v").project(u)
+        J = assemble(v * v * 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(Function(space_b).project(zeta), test_a) * dx).dat.data_ro)
+
+
+@pytest.mark.skipcomplex
+def test_supermesh_project():
+    mesh_a = UnitSquareMesh(10, 10)
+    mesh_b = UnitSquareMesh(5, 20)
+    X_a = SpatialCoordinate(mesh_a)
+    space_a = FunctionSpace(mesh_a, "Lagrange", 1)
+    space_b = FunctionSpace(mesh_b, "Discontinuous Lagrange", 0)
+    test_a = TestFunction(space_a)
+
+    u = Function(space_a, name="u").interpolate(X_a[0] - 0.5)
+    zeta = Function(space_a, name="tlm_u").interpolate(X_a[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
     with reverse_over_forward():
-        v = Function(space_0, name="v").project(u)
+        v = Function(space_b, name="v").project(u)
         J = assemble(v * v * 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(Function(space_0).project(zeta), test) * dx).dat.data_ro)
+                       assemble(2 * inner(Function(space_a).project(Function(space_b).project(zeta)), test_a) * dx).dat.data_ro)