Reverse-over-forward AD: AssembleBlock

Author: jrmaddison

firedrake/adjoint/__init__.py

@@ -19,7 +19,9 @@
 
 from pyadjoint.tape import Tape, set_working_tape, get_working_tape, \
                             pause_annotation, continue_annotation, \
-                            stop_annotating, annotate_tape  # noqa F401
+                            stop_annotating, annotate_tape, \
+                            pause_reverse_over_forward, continue_reverse_over_forward, \
+                            stop_reverse_over_forward  # noqa F401
 from pyadjoint.reduced_functional import ReducedFunctional  # noqa F401
 from firedrake.adjoint_utils.checkpointing import \
     enable_disk_checkpointing, pause_disk_checkpointing, \

firedrake/adjoint_utils/blocks/assembly.py

@@ -4,7 +4,7 @@
 from ufl.formatting.ufl2unicode import ufl2unicode
 from pyadjoint import Block, AdjFloat, create_overloaded_object
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
-from .block_utils import isconstant
+from .block_utils import isconstant, restored_outputs
 
 
 class AssembleBlock(Block):
@@ -145,6 +145,32 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             dform = firedrake.assemble(dform)
         return dform
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        form = self.form
+
+        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:
+                if isinstance(dep, firedrake.MeshGeometry):
+                    dep = firedrake.SpatialCoordinate(dep)
+                    tlm_rhs = tlm_rhs + firedrake.derivative(
+                        form, dep, tlm_dep)
+                else:
+                    tlm_rhs = tlm_rhs + firedrake.action(
+                        firedrake.derivative(form, dep), tlm_dep)
+
+        x.tlm_value = None
+        if isinstance(tlm_rhs, int) and tlm_rhs == 0:
+            return
+        tau_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+        if tau_rhs.empty():
+            return
+        with restored_outputs(x, restore=lambda x: x in form.coefficients()):
+            x.tlm_value = firedrake.assemble(tau_rhs)
+
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
         return self.prepare_evaluate_adj(inputs, adj_inputs,

firedrake/adjoint_utils/blocks/block_utils.py

@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 import firedrake
 
 
@@ -8,3 +10,65 @@ def isconstant(expr):
     if isinstance(expr, firedrake.Constant):
         raise ValueError("Firedrake Constant requires a domain to work with pyadjoint")
     return isinstance(expr, firedrake.Function) and expr.ufl_element().family() == "Real"
+
+
+@contextmanager
+def restored_outputs(*X, restore=None):
+    """Construct a context manager which can be used to temporarily restore
+    block variable outputs to saved values.
+
+    Parameters
+    ----------
+    X : tuple[BlockVariable]
+        Block variables to temporarily restore.
+    restore : callable
+        Can be used to exclude variables. Only inputs for which
+        `restore(x.output)` is true have their outputs temporarily restored.
+
+    Returns
+    -------
+
+    The context manager.
+
+    Notes
+    -----
+
+    A forward operation is allowed to modify the original variable, e.g. in
+
+    .. code-block:: python3
+
+        solve(inner(trial, test) * dx
+              == inner(x * x, test) * dx,
+              x)
+
+    `x` has two versions: the input and the output. Reverse-over-forward AD
+    requires that we use the symbolic representation `x`, but with input value
+    `x.block_variable.saved_output`. A context manager can be used to
+    temporarily restore the value of `x` so that we can perform and annotate
+    a tangent-linear operation,
+
+    .. code-block:: python3
+
+        with restored_outputs(x):
+            # The value of x is now x.block_variable.saved_output
+            solve(inner(trial, test) * dx
+                  == 2 * inner(x * x.block_variable.tlm_value, test) * dx,
+                  x.block_variable.tlm_value)
+        # The value of x is again the output from the forward solve(...)
+    """
+
+    if restore is None:
+        def restore(x):
+            return True
+
+    X = tuple(x for x in X if restore(x.output))
+    X_old = tuple(x.output._ad_copy(x) for x in X)
+    for x in X:
+        # Ideally would use a generic _ad_assign here
+        x.output.assign(x.output.block_variable.saved_output)
+    try:
+        yield
+    finally:
+        for x, x_old in zip(X, X_old):
+            # Ideally would use a generic _ad_assign here
+            x.output.assign(x_old)

tests/regression/test_adjoint_reverse_over_forward.py

@@ -0,0 +1,39 @@
+import numpy as np
+import pytest
+
+from firedrake import *
+from firedrake.adjoint import *
+from firedrake.__future__ import *
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _():
+    get_working_tape().clear_tape()
+    pause_annotation()
+    pause_reverse_over_forward()
+    yield
+    get_working_tape().clear_tape()
+    pause_annotation()
+    pause_reverse_over_forward()
+
+
+def test_assembly():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    u = Function(space, name="u").interpolate(Constant(1.0))
+    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+    continue_annotation()
+    continue_reverse_over_forward()
+    J = assemble(u * u * dx)
+    pause_annotation()
+    pause_reverse_over_forward()
+
+    _ = 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)