Reverse-over-forward AD: ConstantAssignBlock

Author: jrmaddison

firedrake/adjoint_utils/blocks/constant.py

@@ -1,6 +1,6 @@
-from pyadjoint import Block, OverloadedType
+import firedrake
 import numpy
-
+from pyadjoint import Block, OverloadedType
 from pyadjoint.reduced_functional_numpy import gather
 from .block_utils import isconstant
 
@@ -70,6 +70,21 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             values = values.values()
         return constant_from_values(block_variable.output, values)
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        if len(x.output.ufl_shape) == 0:
+            x.tlm_value = firedrake.Constant(0.0)
+        else:
+            x.tlm_value = firedrake.Constant(
+                numpy.reshape(numpy.zeros_like(x.output.values()), x.output.ufl_shape))
+        if self.assigned_list:
+            # Not reachable?
+            raise NotImplementedError
+        else:
+            dep, = self.get_dependencies()
+            if dep.tlm_value is not None:
+                x.tlm_value.assign(dep.tlm_value)
+
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
         return self.prepare_evaluate_adj(inputs, hessian_inputs,

tests/regression/test_adjoint_reverse_over_forward.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import numpy as np
 import pytest
 
@@ -6,7 +7,7 @@
 from firedrake.__future__ import *
 
 
-@pytest.fixture(autouse=True, scope="module")
+@pytest.fixture(autouse=True)
 def _():
     get_working_tape().clear_tape()
     pause_annotation()
@@ -17,6 +18,16 @@ def _():
     pause_reverse_over_forward()
 
 
+@contextmanager
+def reverse_over_forward():
+    continue_annotation()
+    continue_reverse_over_forward()
+    yield
+    pause_annotation()
+    pause_reverse_over_forward()
+
+
+@pytest.mark.skipcomplex
 def test_assembly():
     mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
@@ -27,13 +38,26 @@ def test_assembly():
     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()
+    with reverse_over_forward():
+        J = assemble(u * u * 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(zeta, test) * dx).dat.data_ro)
+
+
+@pytest.mark.skipcomplex
+def test_constant_assignment():
+    a = Constant(2.5)
+    a.block_variable.tlm_value = Constant(-2.0)
+
+    with reverse_over_forward():
+        b = Constant(0.0).assign(a)
+
+    assert float(b.block_variable.tlm_value) == -2.0
+
+    # Minimal test that the TLM operations are on the tape
+    _ = compute_gradient(b.block_variable.tlm_value, Control(a.block_variable.tlm_value))
+    adj_value = a.block_variable.tlm_value.block_variable.adj_value
+    assert float(adj_value) == 1.0