Start to move Hessian evaluation into NonlinearVariationalSolveBlock

Author: angus-g

firedrake/adjoint_utils/blocks/solving.py

@@ -30,6 +30,7 @@ class Solver(Enum):
     FORWARD = 0
     ADJOINT = 1
     TLM = 2
+    HESSIAN = 3
 
 
 class GenericSolveBlock(Block):
@@ -221,6 +222,9 @@ def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
 
         return adj_sol, adj_sol_bdy
 
+    def _hessian_solve(self, *args):
+        return self._assemble_and_solve_adj_eq(*args)
+
     def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
         adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol))
         return adj_sol_bdy.riesz_representation("l2")
@@ -379,8 +383,7 @@ def _assemble_and_solve_soa_eq(self, dFdu_form, adj_sol, hessian_input,
         b = self._assemble_soa_eq_rhs(dFdu_form, adj_sol, hessian_input,
                                       d2Fdu2)
         dFdu_form = firedrake.adjoint(dFdu_form)
-        adj_sol2, adj_sol2_bdy = self._assemble_and_solve_adj_eq(dFdu_form, b,
-                                                                 compute_bdy)
+        adj_sol2, adj_sol2_bdy = self._hessian_solve(dFdu_form, b, compute_bdy)
         if self.adj2_cb is not None:
             self.adj2_cb(adj_sol2)
         if self.adj2_bdy_cb is not None and compute_bdy:
@@ -679,6 +682,22 @@ def _adjoint_solve(self, dJdu, compute_bdy):
                 u_sol, adj_sol_bdy, jac_adj, dJdu_copy)
         return u_sol, adj_sol_bdy
 
+    def _hessian_solve(self, adj_form, rhs, compute_bdy):
+        # self._ad_solver_replace_forms(Solver.HESSIAN)
+        # self._ad_solvers["hessian_lvs"].invalidate_jacobian()
+        self._ad_solvers["hessian_lvs"]._problem.F._components[1].assign(rhs)
+        self._ad_solvers["hessian_lvs"].solve()
+        u_sol = self._ad_solvers["hessian_lvs"]._problem.u
+
+        adj_sol_bdy = None
+        if compute_bdy:
+            jac_adj = self._ad_solvers["hessian_lvs"]._problem.J
+            adj_sol_bdy = self._compute_adj_bdy(
+                u_sol, adj_sol_bdy, jac_adj, rhs.copy()
+            )
+
+        return u_sol, adj_sol_bdy
+
     def _ad_assign_map(self, form, solver):
         if solver == Solver.FORWARD:
             count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
@@ -697,8 +716,10 @@ def _ad_assign_map(self, form, solver):
                            firedrake.Cofunction)):
                 coeff_count = coeff.count()
                 if coeff_count in form_ad_count_map:
-                    assign_map[form_ad_count_map[coeff_count]] = \
-                        block_variable.saved_output
+                    if solver == Solver.HESSIAN:
+                        assign_map[form_ad_count_map[coeff_count]] = block_variable.tlm_value
+                    else:
+                        assign_map[form_ad_count_map[coeff_count]] = block_variable.saved_output
 
         if (
             solver == Solver.ADJOINT
@@ -709,6 +730,7 @@ def _ad_assign_map(self, form, solver):
             if coeff_count in form_ad_count_map:
                 assign_map[form_ad_count_map[coeff_count]] = \
                     block_variable.saved_output
+
         return assign_map
 
     def _ad_assign_coefficients(self, form, solver):
@@ -728,6 +750,10 @@ def _ad_solver_replace_forms(self, solver=Solver.FORWARD):
             self._ad_assign_coefficients(
                 self._ad_solvers["tlm_lvs"]._problem.J, solver
             )
+        elif solver == Solver.HESSIAN:
+            self._ad_assign_coefficients(
+                self._ad_solvers["hessian_lvs"]._problem.J, solver
+            )
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
@@ -851,11 +877,6 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
 
         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):

firedrake/adjoint_utils/variational_solver.py

@@ -4,6 +4,7 @@
 from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
 from firedrake.ufl_expr import derivative, adjoint
 from ufl import replace
+from ufl.algorithms import expand_derivatives
 
 
 class NonlinearVariationalProblemMixin:
@@ -17,6 +18,7 @@ def wrapper(self, *args, **kwargs):
             self._ad_u = self.u_restrict
             self._ad_bcs = self.bcs
             self._ad_J = self.J
+
             try:
                 # Some forms (e.g. SLATE tensors) are not currently
                 # differentiable.
@@ -27,8 +29,10 @@ def wrapper(self, *args, **kwargs):
                     # Try again without expanding derivatives,
                     # as dFdu might have been simplied to an empty Form
                     self._ad_adj_F = adjoint(dFdu, derivatives_expanded=True)
+
             except (TypeError, NotImplementedError):
                 self._ad_adj_F = None
+
             self._ad_kwargs = {'Jp': self.Jp, 'form_compiler_parameters': self.form_compiler_parameters, 'is_linear': self.is_linear}
             self._ad_count_map = {}
         return wrapper
@@ -49,7 +53,8 @@ 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, "tlm_lvs": None}
+                                "recompute_count": 0, "tlm_lvs": None,
+                                "hessian_lvs": None}
             self._ad_adj_cache = {}
 
         return wrapper
@@ -100,6 +105,12 @@ def wrapper(self, **kwargs):
                         if self._ad_problem._constant_jacobian:
                             self._ad_solvers["update_adjoint"] = False
 
+                if not self._ad_solvers["hessian_lvs"]:
+                    with stop_annotating():
+                        self._ad_solvers["hessian_lvs"] = LinearVariationalSolver(
+                            self._ad_hessian_lvs_problem(block, problem._ad_adj_F),
+                        )
+
                 if not self._ad_solvers["tlm_lvs"]:
                     with stop_annotating():
                         self._ad_solvers["tlm_lvs"] = LinearVariationalSolver(
@@ -168,6 +179,27 @@ def _ad_adj_lvs_problem(self, block, adj_F):
         lvp._ad_count_map_update(_ad_count_map)
         return lvp
 
+    @no_annotations
+    def _ad_hessian_lvs_problem(self, block, adj_dFdu):
+        from firedrake import Function, Cofunction, LinearVariationalProblem
+
+        bcs = block._homogenize_bcs()
+        adj_sol = Function(block.function_space)
+        right_hand_side = Cofunction(block.function_space.dual())
+        tmp_problem = LinearVariationalProblem(
+            adj_dFdu, right_hand_side, adj_sol, bcs=bcs,
+            constant_jacobian=self._ad_problem._constant_jacobian)
+
+        _ad_count_map, J_replace_map, _ = self._build_count_map(
+            adj_dFdu, 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