Function.assign adjoint bugfix

firedrake/adjoint_utils/checkpointing.py

@@ -7,6 +7,7 @@
 import shutil
 import atexit
 from abc import ABC, abstractmethod
+from numbers import Number
 _stop_disk_checkpointing = 1
 _checkpoint_init_data = False
 
@@ -317,6 +318,15 @@ class DelegatedFunctionCheckpoint(CheckpointBase):
     """
     def __init__(self, other):
         self.other = other
+        # Obtain a unique identity for this saved output.
+        self.count = type(other.output)(other.output.function_space()).count()
 
     def restore(self):
-        return self.other.saved_output
+        saved_output = self.other.saved_output
+        if isinstance(saved_output, Number):
+            # Happens if the user calls the ReducedFunctional on a number.
+            return saved_output
+        else:
+            return type(saved_output)(saved_output.function_space(),
+                                      saved_output.dat,
+                                      count=self.count)

tests/regression/test_adjoint_operators.py

@@ -701,3 +701,19 @@ def test_consecutive_nonlinear_solves():
     rf = ReducedFunctional(J, Control(uic))
     h = Constant(0.01, domain=mesh)
     assert taylor_test(rf, uic, h) > 1.9
+
+
+@pytest.mark.skipcomplex
+def test_assign_function():
+    from firedrake.adjoint import ReducedFunctional, Control, taylor_test
+    mesh = UnitSquareMesh(1, 1)
+    V = FunctionSpace(mesh, "CG", 1)
+    uic = Function(V, name="uic").assign(1.0)
+    u0 = Function(V, name="u0")
+    u1 = Function(V, name="u1")
+    u0.assign(uic)
+    u1.assign(2 * u0 + uic)
+    J = assemble(((u1 + Constant(1.0)) ** 2) * dx)
+    rf = ReducedFunctional(J, Control(uic))
+    h = Function(V, name="h").assign(0.01)
+    assert taylor_test(rf, uic, h) > 1.9