Start to move TLM evaluation into NonlinearVariationalSolveBlock

Author: angus-g

firedrake/adjoint_utils/blocks/solving.py

@@ -29,6 +29,7 @@ class Solver(Enum):
     """Enum for solver types."""
     FORWARD = 0
     ADJOINT = 1
+    TLM = 2
 
 
 class GenericSolveBlock(Block):
@@ -681,8 +682,11 @@ def _adjoint_solve(self, dJdu, compute_bdy):
     def _ad_assign_map(self, form, solver):
         if solver == Solver.FORWARD:
             count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
-        else:
+        elif solver == Solver.ADJOINT:
             count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map
+        elif solver == Solver.TLM:
+            count_map = self._ad_solvers["tlm_lvs"]._problem._ad_count_map
+
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -717,9 +721,13 @@ def _ad_solver_replace_forms(self, solver=Solver.FORWARD):
             problem = self._ad_solvers["forward_nlvs"]._problem
             self._ad_assign_coefficients(problem.F, solver)
             self._ad_assign_coefficients(problem.J, solver)
-        else:
+        elif solver == Solver.ADJOINT:
             self._ad_assign_coefficients(
                 self._ad_solvers["adjoint_lvs"]._problem.J, solver)
+        elif solver == Solver.TLM:
+            self._ad_assign_coefficients(
+                self._ad_solvers["tlm_lvs"]._problem.J, solver
+            )
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
@@ -796,6 +804,59 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
 
         return dFdm
 
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
+                               prepared=None):
+        F_form = prepared["form"]
+        dFdu = prepared["dFdu"]
+
+        bcs = []
+        dFdm = 0.
+        for block_variable in self.get_dependencies():
+            tlm_value = block_variable.tlm_value
+            c = block_variable.output
+            c_rep = block_variable.saved_output
+
+            if isinstance(c, firedrake.DirichletBC):
+                if tlm_value is None:
+                    bcs.append(c.reconstruct(g=0))
+                else:
+                    bcs.append(tlm_value)
+                continue
+            elif isinstance(c, firedrake.MeshGeometry):
+                X = firedrake.SpatialCoordinate(c)
+                c_rep = X
+
+            if tlm_value is None:
+                continue
+
+            if c == self.func and not self.linear:
+                continue
+
+            dFdm += firedrake.derivative(-F_form, c_rep, tlm_value)
+
+        if isinstance(dFdm, float):
+            v = dFdu.arguments()[0]
+            dFdm = firedrake.inner(
+                firedrake.Constant(numpy.zeros(v.ufl_shape)), v
+            ) * firedrake.dx
+
+        dFdm = ufl.algorithms.expand_derivatives(dFdm)
+        dFdm = firedrake.assemble(dFdm)
+
+        # XXX I dunno how this works
+        self._ad_solver_replace_forms(Solver.TLM)
+        self._ad_solvers["tlm_lvs"].invalidate_jacobian()
+        # update RHS
+        self._ad_solvers["tlm_lvs"]._problem.F._components[1].assign(dFdm)
+
+        self._ad_solvers["tlm_lvs"].solve()
+        return self._ad_solvers["tlm_lvs"]._problem.u
+        # return self._assemble_and_solve_tlm_eq(
+        #     firedrake.assemble(dFdu, bcs=bcs, **self.assemble_kwargs),
+        #     dFdm, dudm, bcs
+        # )
+
+
 
 class ProjectBlock(SolveVarFormBlock):
     def __init__(self, v, V, output, bcs=[], *args, **kwargs):

firedrake/adjoint_utils/variational_solver.py

@@ -49,7 +49,7 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_args = args
             self._ad_kwargs = kwargs
             self._ad_solvers = {"forward_nlvs": None, "adjoint_lvs": None,
-                                "recompute_count": 0}
+                                "recompute_count": 0, "tlm_lvs": None}
             self._ad_adj_cache = {}
 
         return wrapper
@@ -100,6 +100,14 @@ def wrapper(self, **kwargs):
                         if self._ad_problem._constant_jacobian:
                             self._ad_solvers["update_adjoint"] = False
 
+                if not self._ad_solvers["tlm_lvs"]:
+                    with stop_annotating():
+                        self._ad_solvers["tlm_lvs"] = LinearVariationalSolver(
+                            self._ad_tlm_lvs_problem(block, problem.F, problem.u_restrict)
+                        )
+                        if self._ad_problem._constant_jacobian:
+                            self._ad_solvers["update_tlm"] = False
+
                 block._ad_solvers = self._ad_solvers
 
                 tape.add_block(block)
@@ -151,14 +159,33 @@ def _ad_adj_lvs_problem(self, block, adj_F):
         # linear variational problem is created with a deep copy of the
         # `block.adj_F` coefficients.
         _ad_count_map, J_replace_map, _ = self._build_count_map(
-            adj_F, block._dependencies)
+            adj_F, block._dependencies,
+        )
         lvp = LinearVariationalProblem(
             replace(tmp_problem.J, J_replace_map), right_hand_side, adj_sol,
             bcs=tmp_problem.bcs,
             constant_jacobian=self._ad_problem._constant_jacobian)
         lvp._ad_count_map_update(_ad_count_map)
         return lvp
 
+    @no_annotations
+    def _ad_tlm_lvs_problem(self, block, F, u):
+        from firedrake import Function, Cofunction, LinearVariationalProblem
+
+        lhs = derivative(F, u)
+        _ad_count_map, F_replace_map, _ = self._build_count_map(lhs, block._dependencies)
+        sol = Function(block.function_space)
+        rhs = Cofunction(block.function_space.dual())
+        lvp = LinearVariationalProblem(
+            replace(lhs, F_replace_map),
+            rhs,
+            sol,
+            bcs=block._homogenize_bcs(),
+            constant_jacobian=self._ad_problem._constant_jacobian,
+        )
+        lvp._ad_count_map_update(_ad_count_map)
+        return lvp
+
     def _build_count_map(self, J, dependencies, F=None):
         from firedrake import Function