NLVS tape block stores adj_sol per block not per solver

Author: JHopeCollins

firedrake/adjoint_utils/blocks/solving.py

@@ -56,6 +56,7 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         # Solution function
         self.func = func
         self.function_space = self.func.function_space()
+        self.adj_state_buf = self.func.copy(deepcopy=True)
         # Boundary conditions
         self.bcs = []
         if bcs is not None:
@@ -193,6 +194,8 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         if self.adj_bdy_cb is not None and compute_bdy:
             self.adj_bdy_cb(adj_sol_bdy)
 
+        self.adj_state_buf.assign(adj_sol)
+
         r = {}
         r["form"] = F_form
         r["adj_sol"] = adj_sol
@@ -399,6 +402,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
         if tlm_output is None:
             return
 
+        self.adj_state.assign(self.adj_state_buf)
+
         F_form = self._create_F_form()
 
         # Using the equation Form derive dF/du, d^2F/du^2 * du/dm * direction.
@@ -727,6 +732,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         )
         adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy)
         self.adj_state = adj_sol
+        self.adj_state_buf.assign(adj_sol)
         if self.adj_cb is not None:
             self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:

tests/firedrake/adjoint/test_hessian.py

@@ -37,7 +37,7 @@ def test_simple_solve(rg):
     mesh = IntervalMesh(10, 0, 1)
     V = FunctionSpace(mesh, "Lagrange", 1)
 
-    f = Function(V).assign(2)
+    f = Function(V).assign(2.)
 
     u = TrialFunction(V)
     v = TestFunction(V)
@@ -76,10 +76,10 @@ def test_mixed_derivatives(rg):
     mesh = IntervalMesh(10, 0, 1)
     V = FunctionSpace(mesh, "Lagrange", 1)
 
-    f = Function(V).assign(2)
+    f = Function(V).assign(2.)
     control_f = Control(f)
 
-    g = Function(V).assign(3)
+    g = Function(V).assign(3.)
     control_g = Control(g)
 
     u = TrialFunction(V)
@@ -126,7 +126,7 @@ def test_function(rg):
     R = FunctionSpace(mesh, "R", 0)
     c = Function(R, val=4)
     control_c = Control(c)
-    f = Function(V).assign(3)
+    f = Function(V).assign(3.)
     control_f = Control(f)
 
     u = Function(V)
@@ -139,14 +139,14 @@ def test_function(rg):
     J = assemble(c ** 2 * u ** 2 * dx)
 
     Jhat = ReducedFunctional(J, [control_c, control_f])
-    dJdc, dJdf = compute_gradient(J, [control_c, control_f], apply_riesz=True)
+    dJdc, dJdf = compute_derivative(J, [control_c, control_f], apply_riesz=True)
 
     # Step direction for derivatives and convergence test
     h_c = Function(R, val=1.0)
     h_f = rg.uniform(V, 0, 10)
 
     # Total derivative
-    dJdc, dJdf = compute_gradient(J, [control_c, control_f], apply_riesz=True)
+    dJdc, dJdf = compute_derivative(J, [control_c, control_f], apply_riesz=True)
     dJdm = assemble(dJdc * h_c * dx + dJdf * h_f * dx)
 
     # Hessian
@@ -163,7 +163,7 @@ def test_nonlinear(rg):
     mesh = UnitSquareMesh(10, 10)
     V = FunctionSpace(mesh, "Lagrange", 1)
     R = FunctionSpace(mesh, "R", 0)
-    f = Function(V).assign(5)
+    f = Function(V).assign(5.)
 
     u = Function(V)
     v = TestFunction(V)
@@ -201,11 +201,11 @@ def test_dirichlet(rg):
     mesh = UnitSquareMesh(10, 10)
     V = FunctionSpace(mesh, "Lagrange", 1)
 
-    f = Function(V).assign(30)
+    f = Function(V).assign(30.)
 
     u = Function(V)
     v = TestFunction(V)
-    c = Function(V).assign(1)
+    c = Function(V).assign(1.)
     bc = DirichletBC(V, c, "on_boundary")
 
     F = inner(grad(u), grad(v)) * dx + u**4*v*dx - f**2 * v * dx
@@ -249,24 +249,25 @@ def Dt(u, u_, timestep):
     pr = project(sin(2*pi*x), V, annotate=False)
     ic = Function(V).assign(pr)
 
-    u_ = Function(V)
-    u = Function(V)
+    u_ = Function(V).assign(ic)
+    u = Function(V).assign(ic)
     v = TestFunction(V)
 
     nu = Constant(0.0001)
 
-    timestep = Constant(1.0/n)
+    dt = 0.01
+    nt = 20
 
     params = {
         'snes_rtol': 1e-10,
         'ksp_type': 'preonly',
         'pc_type': 'lu',
     }
 
-    F = (Dt(u, ic, timestep)*v
+    F = (Dt(u, u_, dt)*v
          + u*u.dx(0)*v + nu*u.dx(0)*v.dx(0))*dx
+
     bc = DirichletBC(V, 0.0, "on_boundary")
-    t = 0.0
 
     if solve_type == "nlvs":
         use_nlvs = True
@@ -285,21 +286,14 @@ def Dt(u, u_, timestep):
     else:
         solve(F == 0, u, bc, solver_parameters=params)
     u_.assign(u)
-    t += float(timestep)
 
-    F = (Dt(u, u_, timestep)*v
-         + u*u.dx(0)*v + nu*u.dx(0)*v.dx(0))*dx
-
-    end = 0.2
-    while (t <= end):
+    for _ in range(nt):
         if use_nlvs:
             solver.solve()
         else:
             solve(F == 0, u, bc, solver_parameters=params)
         u_.assign(u)
 
-        t += float(timestep)
-
     J = assemble(u_*u_*dx + ic*ic*dx)
 
     Jhat = ReducedFunctional(J, Control(ic))