Merge pull request #92 from dolfin-adjoint/annotation_assemble_1-form

Extend AssembleBlock to handle 1-forms

dolfin_adjoint_common/blocks/assembly.py

@@ -18,6 +18,49 @@ def __init__(self, form, ad_block_tag=None):
     def __str__(self):
         return f"assemble({ufl2unicode(self.form)})"
 
+    def compute_action_adjoint(self, adj_input, arity_form, form=None, c_rep=None, space=None, dform=None):
+        """This computes the action of the adjoint of the derivative of `form` wrt `c_rep` on `adj_input`.
+           In other words, it returns: `<(dform/dc_rep)*, adj_input>`
+
+           - If `form` has arity 0 => `dform/dc_rep` is a 1-form and `adj_input` a foat, we can simply use
+             the `*` operator.
+
+           - If `form` has arity 1 => `dform/dc_rep` is a 2-form and we can symbolically take its adjoint
+             and then apply the action on `adj_input`, to finally assemble the result.
+        """
+        if arity_form == 0:
+            if dform is None:
+                dc = self.backend.TestFunction(space)
+                dform = self.backend.derivative(form, c_rep, dc)
+            dform_vector = self.compat.assemble_adjoint_value(dform)
+            # Return a Vector scaled by the scalar `adj_input`
+            return dform_vector * adj_input, dform
+        elif arity_form == 1:
+            if dform is None:
+                dc = self.backend.TrialFunction(space)
+                dform = self.backend.derivative(form, c_rep, dc)
+            # Get the Function
+            adj_input = adj_input.function
+            # Symbolic operators such as action/adjoint require derivatives to have been expanded beforehand.
+            # However, UFL doesn't support expanding coordinate derivatives of Coefficients in physical space,
+            # implying that we can't symbolically take the action/adjoint of the Jacobian for SpatialCoordinates.
+            # -> Workaround: Apply action/adjoint numerically (using PETSc).
+            if not isinstance(c_rep, self.backend.SpatialCoordinate):
+                # Symbolically compute: (dform/dc_rep)^* * adj_input
+                adj_output = self.backend.action(self.backend.adjoint(dform), adj_input)
+                adj_output = self.compat.assemble_adjoint_value(adj_output)
+            else:
+                # Get PETSc matrix
+                dform_mat = self.compat.assemble_adjoint_value(dform).petscmat
+                # Action of the adjoint (Hermitian transpose)
+                adj_output = self.backend.Function(space)
+                with adj_input.dat.vec_ro as v_vec:
+                    with adj_output.dat.vec as res_vec:
+                        dform_mat.multHermitian(v_vec, res_vec)
+            return adj_output, dform
+        else:
+            raise ValueError('Forms with arity > 1 are not handled yet!')
+
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         replaced_coeffs = {}
         for block_variable in self.get_dependencies():
@@ -35,6 +78,9 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
         c = block_variable.output
         c_rep = block_variable.saved_output
 
+        from ufl.algorithms.analysis import extract_arguments
+        arity_form = len(extract_arguments(form))
+
         if isinstance(c, self.compat.ExpressionType):
             # Create a FunctionSpace from self.form and Expression.
             # And then make a TestFunction from this space.
@@ -47,24 +93,25 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
             return [[adj_input * output, V]]
 
         if isinstance(c, self.backend.Function):
-            dc = self.backend.TestFunction(c.function_space())
+            space = c.function_space()
         elif isinstance(c, self.backend.Constant):
             mesh = self.compat.extract_mesh_from_form(self.form)
-            dc = self.backend.TestFunction(c._ad_function_space(mesh))
+            space = c._ad_function_space(mesh)
         elif isinstance(c, self.compat.MeshType):
             c_rep = self.backend.SpatialCoordinate(c_rep)
-            dc = self.backend.TestFunction(c._ad_function_space())
+            space = c._ad_function_space()
 
-        dform = self.backend.derivative(form, c_rep, dc)
-        output = self.compat.assemble_adjoint_value(dform)
-        return adj_input * output
+        return self.compute_action_adjoint(adj_input, arity_form, form, c_rep, space)[0]
 
     def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
         return self.prepare_evaluate_adj(inputs, tlm_inputs, self.get_dependencies())
 
     def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
         form = prepared
         dform = 0.
+
+        from ufl.algorithms.analysis import extract_arguments
+        arity_form = len(extract_arguments(form))
         for bv in self.get_dependencies():
             c_rep = bv.saved_output
             tlm_value = bv.tlm_value
@@ -79,6 +126,9 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         if not isinstance(dform, float):
             dform = ufl.algorithms.expand_derivatives(dform)
             dform = self.compat.assemble_adjoint_value(dform)
+            if arity_form == 1 and dform != 0:
+                # Then dform is a Vector
+                dform = dform.function
         return dform
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies):
@@ -90,27 +140,27 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
         hessian_input = hessian_inputs[0]
         adj_input = adj_inputs[0]
 
+        from ufl.algorithms.analysis import extract_arguments
+        arity_form = len(extract_arguments(form))
+
         c1 = block_variable.output
         c1_rep = block_variable.saved_output
 
         if isinstance(c1, self.backend.Function):
-            dc = self.backend.TestFunction(c1.function_space())
+            space = c1.function_space()
         elif isinstance(c1, self.compat.ExpressionType):
             mesh = form.ufl_domain().ufl_cargo()
-            W = c1._ad_function_space(mesh)
-            dc = self.backend.TestFunction(W)
+            space = c1._ad_function_space(mesh)
         elif isinstance(c1, self.backend.Constant):
             mesh = self.compat.extract_mesh_from_form(form)
-            dc = self.backend.TestFunction(c1._ad_function_space(mesh))
+            space = c1._ad_function_space(mesh)
         elif isinstance(c1, self.compat.MeshType):
             c1_rep = self.backend.SpatialCoordinate(c1)
-            dc = self.backend.TestFunction(c1._ad_function_space())
+            space = c1._ad_function_space()
         else:
             return None
 
-        dform = self.backend.derivative(form, c1_rep, dc)
-        dform = ufl.algorithms.expand_derivatives(dform)
-        hessian_outputs = hessian_input * self.compat.assemble_adjoint_value(dform)
+        hessian_outputs, dform = self.compute_action_adjoint(hessian_input, arity_form, form, c1_rep, space)
 
         ddform = 0
         for other_idx, bv in relevant_dependencies:
@@ -129,10 +179,10 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
         if not isinstance(ddform, float):
             ddform = ufl.algorithms.expand_derivatives(ddform)
             if not ddform.empty():
-                hessian_outputs += adj_input * self.compat.assemble_adjoint_value(ddform)
+                hessian_outputs += self.compute_action_adjoint(adj_input, arity_form, dform=ddform)[0]
 
         if isinstance(c1, self.compat.ExpressionType):
-            return [(hessian_outputs, W)]
+            return [(hessian_outputs, space)]
         else:
             return hessian_outputs
 

tests/firedrake_adjoint/test_assemble.py

@@ -1,9 +1,12 @@
-import pytest
-pytest.importorskip("firedrake")
-
 from firedrake import *
 from firedrake_adjoint import *
+
 from numpy.testing import assert_approx_equal
+from numpy.random import rand
+
+import pytest
+pytest.importorskip("firedrake")
+
 
 def test_assemble_0_forms():
     mesh = IntervalMesh(10, 0, 1)
@@ -21,6 +24,7 @@ def test_assemble_0_forms():
     # derivative is: (1+2*u+6*u**2)*dx - summing is equivalent to testing with 1
     assert_approx_equal(rf.derivative().vector().sum(), 1. + 2. * 4 + 6 * 16.)
 
+
 def test_assemble_0_forms_mixed():
     mesh = IntervalMesh(10, 0, 1)
     V = FunctionSpace(mesh, "Lagrange", 1)
@@ -37,3 +41,83 @@ def test_assemble_0_forms_mixed():
     rf = ReducedFunctional(s, Control(u))
     # derivative is: (1+4*u)*dx - summing is equivalent to testing with 1
     assert_approx_equal(rf.derivative().vector().sum(), 1. + 4. * 7)
+
+
+def test_assemble_1_forms_adjoint():
+    mesh = IntervalMesh(10, 0, 1)
+    V = FunctionSpace(mesh, "Lagrange", 1)
+    v = TestFunction(V)
+    x, = SpatialCoordinate(mesh)
+    f = Function(V).interpolate(cos(x))
+
+    def J(f):
+        w1 = assemble(inner(f, v) * dx)
+        w2 = assemble(inner(f**2, v) * dx)
+        w3 = assemble(inner(f**3, v) * dx)
+        return assemble((w1 + w2 + w3)**2 * dx)
+
+    _test_adjoint(J, f)
+
+
+def test_assemble_1_forms_tlm():
+    tape = Tape()
+    set_working_tape(tape)
+
+    mesh = IntervalMesh(10, 0, 1)
+    V = FunctionSpace(mesh, "Lagrange", 1)
+    v = TestFunction(V)
+    f = Function(V).assign(1)
+
+    w1 = assemble(inner(f, v) * dx)
+    w2 = assemble(inner(f**2, v) * dx)
+    w3 = assemble(inner(f**3, v) * dx)
+    J = assemble((w1 + w2 + w3)**2 * dx)
+
+    Jhat = ReducedFunctional(J, Control(f))
+    h = Function(V)
+    h.vector()[:] = rand(h.dof_dset.size)
+    g = f.copy(deepcopy=True)
+    f.block_variable.tlm_value = h
+    tape.evaluate_tlm()
+    assert (taylor_test(Jhat, g, h, dJdm=J.block_variable.tlm_value) > 1.9)
+
+
+def _test_adjoint(J, f):
+    import numpy.random
+    tape = Tape()
+    set_working_tape(tape)
+
+    V = f.function_space()
+    h = Function(V)
+    h.vector()[:] = numpy.random.rand(V.dim())
+
+    eps_ = [0.01 / 2.0**i for i in range(5)]
+    residuals = []
+    for eps in eps_:
+
+        Jp = J(f + eps * h)
+        tape.clear_tape()
+        Jm = J(f)
+        Jm.block_variable.adj_value = 1.0
+        tape.evaluate_adj()
+
+        dJdf = f.block_variable.adj_value
+
+        residual = abs(Jp - Jm - eps * dJdf.inner(h.vector()))
+        residuals.append(residual)
+
+    r = convergence_rates(residuals, eps_)
+    print(r)
+    print(residuals)
+
+    tol = 1E-1
+    assert(r[-1] > 2 - tol)
+
+
+def convergence_rates(E_values, eps_values):
+    from numpy import log
+    r = []
+    for i in range(1, len(eps_values)):
+        r.append(log(E_values[i] / E_values[i - 1]) / log(eps_values[i] / eps_values[i - 1]))
+
+    return r