Remove restored_output, now handled within pyadjoint

Author: jrmaddison

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, restored_outputs
+from .block_utils import isconstant
 
 
 class AssembleBlock(Block):
@@ -168,8 +168,7 @@ def solve_tlm(self):
         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)
+        x.tlm_value = firedrake.assemble(tau_rhs)
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):

firedrake/adjoint_utils/blocks/block_utils.py

@@ -1,5 +1,3 @@
-from contextlib import contextmanager
-
 import firedrake
 
 
@@ -10,65 +8,3 @@ 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)