From 6997430e90d728b80c433f35473c5d2792be77b5 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 22 Jul 2024 11:48:57 +0100
Subject: [PATCH 01/46] Add adjoint variational solver.

---
 firedrake/adjoint_utils/variational_solver.py | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index a6811002ac..bcfb524fa1 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -2,6 +2,8 @@
 from functools import wraps
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
 from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
+from firedrake import Cofunction, Function, LinearVariationalProblem, \
+    LinearVariationalSolver, NonlinearVariationalProblem
 from ufl import replace
 
 
@@ -47,6 +49,8 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_kwargs = kwargs
             self._ad_nlvs = None
             self._ad_adj_cache = {}
+            self._ad_adj_varsolver = None
+            self._ad_dJdu = None
 
         return wrapper
 
@@ -58,7 +62,6 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -76,6 +79,7 @@ def wrapper(self, **kwargs):
                                                        solver_kwargs=self._ad_kwargs,
                                                        ad_block_tag=self.ad_block_tag,
                                                        **sb_kwargs)
+                # Forward solver.
                 if not self._ad_nlvs:
                     self._ad_nlvs = type(self)(
                         self._ad_problem_clone(self._ad_problem, block.get_dependencies()),
@@ -83,6 +87,22 @@ def wrapper(self, **kwargs):
                     )
 
                 block._ad_nlvs = self._ad_nlvs
+
+                # Adjoint solver.
+                with stop_annotating():
+                    if not self._ad_adj_varsolver:
+                        self._ad_dJdu = Cofunction(problem._ad_u.function_space().dual())
+                        adj_sol = Function(self._ad_problem._ad_u.function_space())
+                        problem = LinearVariationalProblem(
+                            self._ad_problem._ad_adj_F, self._ad_dJdu, adj_sol)
+                        self._ad_adj_varsolver = LinearVariationalSolver(
+                            problem, solver_parameters=self.parameters)
+                    else:
+                        adj_sol = self._ad_adj_varsolver._ad_problem.u
+                
+                block._ad_dJdu = self._ad_dJdu
+                block._ad_adj_varsolver = self._ad_adj_varsolver
+
                 tape.add_block(block)
 
             with stop_annotating():
@@ -103,7 +123,6 @@ def _ad_problem_clone(self, problem, dependencies):
         affect the user-defined self._ad_problem.F, self._ad_problem.J and self._ad_problem.u
         expressions, we'll instead create clones of them.
         """
-        from firedrake import Function, NonlinearVariationalProblem
         F_replace_map = {}
         J_replace_map = {}
 

From eb52a867608a022d51b72a89d8bfd8e9318218a5 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 22 Jul 2024 16:10:00 +0100
Subject: [PATCH 02/46] wip

---
 firedrake/adjoint_utils/blocks/solving.py     | 71 +++++++++++++------
 firedrake/adjoint_utils/variational_solver.py | 60 +++++++++++-----
 2 files changed, 89 insertions(+), 42 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index ad47834938..3a65fc2e52 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -629,12 +629,28 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         self._ad_nlvs.solve()
         func.assign(self._ad_nlvs._problem.u)
         return func
-
-    def _ad_assign_map(self, form):
-        count_map = self._ad_nlvs._problem._ad_count_map
+    
+    def _adjoint_solve(self, dJdu, adj_sol):
+        # self._ad_adj_lvs_replace_jacobian()
+        self._ad_adj_varsolver.parameters.update(self.solver_params)
+        # Replace right-hand side with dJdu.
+        self._ad_dJdu.dat.data[:] = 0.0
+        self._ad_adj_varsolver.solve()
+        adj_sol.assign(self._ad_adj_varsolver._problem.u)
+        return adj_sol
+
+    def _ad_assign_map(self, form, count_map):
         assign_map = {}
-        form_ad_count_map = dict((count_map[coeff], coeff)
-                                 for coeff in form.coefficients())
+        # form_ad_count_map = dict((count_map[coeff], coeff)
+        #                          for coeff in form.coefficients())
+        form_ad_count_map = dict()
+        for coeff in form.coefficients():
+            if coeff != self._ad_dJdu and coeff != self._ad_adj_varsolver._problem.u:
+                form_ad_count_map[count_map[coeff]] = coeff
+                # dict((count_map[coeff], coeff))
+                # form_ad_count_map[count_map[coeff]] = count_map[coeff]
+                # dict((count_map[coeff], coeff))
+
         for block_variable in self.get_dependencies():
             coeff = block_variable.output
             if isinstance(coeff,
@@ -646,8 +662,12 @@ def _ad_assign_map(self, form):
                         block_variable.saved_output
         return assign_map
 
-    def _ad_assign_coefficients(self, form):
-        assign_map = self._ad_assign_map(form)
+    def _ad_assign_coefficients(self, form, solver_mode="forward"):
+        if solver_mode == "forward":
+            count_map = self._ad_nlvs._problem._ad_count_map
+        else:
+            count_map = self._ad_adj_varsolver._problem._ad_count_map
+        assign_map = self._ad_assign_map(form, count_map)
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
@@ -656,6 +676,10 @@ def _ad_nlvs_replace_forms(self):
         self._ad_assign_coefficients(problem.F)
         self._ad_assign_coefficients(problem.J)
 
+    def _ad_adj_lvs_replace_jacobian(self):
+        problem = self._ad_adj_varsolver._problem
+        self._ad_assign_coefficients(problem.J, solver_mode="adjoint")
+
     def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
         if "dFdu_adj" in self._adj_cache:
             dFdu = self._adj_cache["dFdu_adj"]
@@ -670,30 +694,31 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
 
         F_form = self._create_F_form()
 
-        dFdu_form = self.adj_F
+        # dFdu_form = self.adj_F
         dJdu = dJdu.copy()
 
         # Replace the form coefficients with checkpointed values.
-        replace_map = self._replace_map(dFdu_form)
-        replace_map[self.func] = self.get_outputs()[0].saved_output
-        dFdu_form = replace(dFdu_form, replace_map)
-
-        compute_bdy = self._should_compute_boundary_adjoint(
-            relevant_dependencies
-        )
-        adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
-            dFdu_form, dJdu, compute_bdy
-        )
-        self.adj_sol = adj_sol
+        # replace_map = self._replace_map(dFdu_form)
+        # replace_map[self.func] = self.get_outputs()[0].saved_output
+        # dFdu_form = replace(dFdu_form, replace_map)
+
+        # compute_bdy = self._should_compute_boundary_adjoint(
+        #     relevant_dependencies
+        # )
+        # adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
+        #     dFdu_form, dJdu, compute_bdy
+        # )
+        adj_sol = firedrake.Function(self.function_space)
+        self.adj_sol = self._adjoint_solve(dJdu, 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:
-            self.adj_bdy_cb(adj_sol_bdy)
+        # if self.adj_bdy_cb is not None and compute_bdy:
+        #     self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
         r["form"] = F_form
         r["adj_sol"] = adj_sol
-        r["adj_sol_bdy"] = adj_sol_bdy
+        r["adj_sol_bdy"] = None
         return r
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
@@ -744,7 +769,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         replace_map[self.func] = self.get_outputs()[0].saved_output
         dFdm = replace(dFdm, replace_map)
 
-        dFdm = dFdm * adj_sol
+        dFdm = ufl.Action(dFdm, adj_sol)
         dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs)
 
         return dFdm
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index bcfb524fa1..423ded6049 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -2,8 +2,6 @@
 from functools import wraps
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
 from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
-from firedrake import Cofunction, Function, LinearVariationalProblem, \
-    LinearVariationalSolver, NonlinearVariationalProblem
 from ufl import replace
 
 
@@ -62,6 +60,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
+            from firedrake import LinearVariationalSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -79,7 +78,7 @@ def wrapper(self, **kwargs):
                                                        solver_kwargs=self._ad_kwargs,
                                                        ad_block_tag=self.ad_block_tag,
                                                        **sb_kwargs)
-                # Forward solver.
+                # Forward variational solver.
                 if not self._ad_nlvs:
                     self._ad_nlvs = type(self)(
                         self._ad_problem_clone(self._ad_problem, block.get_dependencies()),
@@ -88,18 +87,14 @@ def wrapper(self, **kwargs):
 
                 block._ad_nlvs = self._ad_nlvs
 
-                # Adjoint solver.
+                # Adjoint variational Solver.
                 with stop_annotating():
                     if not self._ad_adj_varsolver:
-                        self._ad_dJdu = Cofunction(problem._ad_u.function_space().dual())
-                        adj_sol = Function(self._ad_problem._ad_u.function_space())
-                        problem = LinearVariationalProblem(
-                            self._ad_problem._ad_adj_F, self._ad_dJdu, adj_sol)
+                        problem = self._ad_create_lvs_problem(
+                            problem._ad_u.function_space(), block.get_dependencies())
                         self._ad_adj_varsolver = LinearVariationalSolver(
                             problem, solver_parameters=self.parameters)
-                    else:
-                        adj_sol = self._ad_adj_varsolver._ad_problem.u
-                
+
                 block._ad_dJdu = self._ad_dJdu
                 block._ad_adj_varsolver = self._ad_adj_varsolver
 
@@ -123,6 +118,40 @@ def _ad_problem_clone(self, problem, dependencies):
         affect the user-defined self._ad_problem.F, self._ad_problem.J and self._ad_problem.u
         expressions, we'll instead create clones of them.
         """
+        from firedrake import NonlinearVariationalProblem
+        _ad_count_map, F_replace_map, J_replace_map = self._build_count_map(
+            problem, dependencies)
+        nlvp = NonlinearVariationalProblem(replace(problem.F, F_replace_map),
+                                           F_replace_map[problem.u_restrict],
+                                           bcs=problem.bcs,
+                                           J=replace(problem.J, J_replace_map))
+        nlvp._constant_jacobian = problem._constant_jacobian
+        nlvp._ad_count_map_update(_ad_count_map)
+        return nlvp
+
+    @no_annotations
+    def _ad_create_lvs_problem(self, fwd_func_space, dependencies):
+        """Creates a LinearVariationalProblem instance for the adjoint variational problem."""
+        from firedrake import Cofunction, Function, LinearVariationalProblem, NonlinearVariationalProblem
+
+        self._ad_dJdu = Cofunction(fwd_func_space.dual())
+        adj_sol = Function(fwd_func_space)
+        tmp_problem = LinearVariationalProblem(
+            self._ad_problem._ad_adj_F, self._ad_dJdu, adj_sol,
+            bcs=self._ad_problem._ad_bcs)
+        _ad_count_map, F_replace_map, J_replace_map = self._build_count_map(
+            tmp_problem, dependencies)
+        # lvp = LinearVariationalProblem(
+        #     replace(tmp_problem.J, J_replace_map),
+        #     self._ad_dJdu,
+        #     adj_sol,
+        #     bcs=tmp_problem.bcs)
+        # lvp._ad_count_map_update(_ad_count_map)
+        return tmp_problem
+
+    def _build_count_map(self, problem, dependencies):
+        from firedrake import Function
+
         F_replace_map = {}
         J_replace_map = {}
 
@@ -147,11 +176,4 @@ def _ad_problem_clone(self, problem, dependencies):
                 else:
                     J_replace_map[coeff] = coeff.copy()
                 _ad_count_map[J_replace_map[coeff]] = coeff.count()
-
-        nlvp = NonlinearVariationalProblem(replace(problem.F, F_replace_map),
-                                           F_replace_map[problem.u_restrict],
-                                           bcs=problem.bcs,
-                                           J=replace(problem.J, J_replace_map))
-        nlvp._constant_jacobian = problem._constant_jacobian
-        nlvp._ad_count_map_update(_ad_count_map)
-        return nlvp
+        return _ad_count_map, F_replace_map, J_replace_map

From 6b98f55b4bc1403a284c57fc3ededfee43c2bfdd Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 22 Jul 2024 16:30:47 +0100
Subject: [PATCH 03/46] Update and Jacobian coefficients

---
 firedrake/adjoint_utils/blocks/solving.py     | 13 ++-----------
 firedrake/adjoint_utils/variational_solver.py | 17 ++++++++---------
 2 files changed, 10 insertions(+), 20 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 3a65fc2e52..4b221b0af0 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -631,10 +631,10 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         return func
     
     def _adjoint_solve(self, dJdu, adj_sol):
-        # self._ad_adj_lvs_replace_jacobian()
+        self._ad_adj_lvs_replace_jacobian()
         self._ad_adj_varsolver.parameters.update(self.solver_params)
         # Replace right-hand side with dJdu.
-        self._ad_dJdu.dat.data[:] = 0.0
+        self._ad_dJdu.assign(dJdu)
         self._ad_adj_varsolver.solve()
         adj_sol.assign(self._ad_adj_varsolver._problem.u)
         return adj_sol
@@ -694,20 +694,11 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
 
         F_form = self._create_F_form()
 
-        # dFdu_form = self.adj_F
         dJdu = dJdu.copy()
 
-        # Replace the form coefficients with checkpointed values.
-        # replace_map = self._replace_map(dFdu_form)
-        # replace_map[self.func] = self.get_outputs()[0].saved_output
-        # dFdu_form = replace(dFdu_form, replace_map)
-
         # compute_bdy = self._should_compute_boundary_adjoint(
         #     relevant_dependencies
         # )
-        # adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
-        #     dFdu_form, dJdu, compute_bdy
-        # )
         adj_sol = firedrake.Function(self.function_space)
         self.adj_sol = self._adjoint_solve(dJdu, adj_sol)
         if self.adj_cb is not None:
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 423ded6049..5f5c45387a 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -137,17 +137,16 @@ def _ad_create_lvs_problem(self, fwd_func_space, dependencies):
         self._ad_dJdu = Cofunction(fwd_func_space.dual())
         adj_sol = Function(fwd_func_space)
         tmp_problem = LinearVariationalProblem(
-            self._ad_problem._ad_adj_F, self._ad_dJdu, adj_sol,
-            bcs=self._ad_problem._ad_bcs)
+            self._ad_problem._ad_adj_F, self._ad_dJdu, adj_sol)
         _ad_count_map, F_replace_map, J_replace_map = self._build_count_map(
             tmp_problem, dependencies)
-        # lvp = LinearVariationalProblem(
-        #     replace(tmp_problem.J, J_replace_map),
-        #     self._ad_dJdu,
-        #     adj_sol,
-        #     bcs=tmp_problem.bcs)
-        # lvp._ad_count_map_update(_ad_count_map)
-        return tmp_problem
+        lvp = LinearVariationalProblem(
+            replace(tmp_problem.J, J_replace_map),
+            self._ad_dJdu,
+            adj_sol,
+            bcs=tmp_problem.bcs)
+        lvp._ad_count_map_update(_ad_count_map)
+        return lvp
 
     def _build_count_map(self, problem, dependencies):
         from firedrake import Function

From 0a83c85a174dbcbf9cfb8d0c9ca638f72857b788 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 22 Jul 2024 16:48:41 +0100
Subject: [PATCH 04/46] wip

---
 firedrake/adjoint_utils/blocks/solving.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 4b221b0af0..8ae8373724 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -631,7 +631,7 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         return func
     
     def _adjoint_solve(self, dJdu, adj_sol):
-        self._ad_adj_lvs_replace_jacobian()
+        # self._ad_adj_lvs_replace_jacobian()
         self._ad_adj_varsolver.parameters.update(self.solver_params)
         # Replace right-hand side with dJdu.
         self._ad_dJdu.assign(dJdu)

From 9387d48be114a7315756bbc254ea1e9884fd712f Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 22 Jul 2024 17:15:09 +0100
Subject: [PATCH 05/46] wip

---
 firedrake/adjoint_utils/blocks/solving.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 8ae8373724..dda6437951 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -631,7 +631,7 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         return func
     
     def _adjoint_solve(self, dJdu, adj_sol):
-        # self._ad_adj_lvs_replace_jacobian()
+        self._ad_adj_lvs_replace_jacobian()
         self._ad_adj_varsolver.parameters.update(self.solver_params)
         # Replace right-hand side with dJdu.
         self._ad_dJdu.assign(dJdu)
@@ -641,15 +641,12 @@ def _adjoint_solve(self, dJdu, adj_sol):
 
     def _ad_assign_map(self, form, count_map):
         assign_map = {}
-        # form_ad_count_map = dict((count_map[coeff], coeff)
-        #                          for coeff in form.coefficients())
         form_ad_count_map = dict()
         for coeff in form.coefficients():
-            if coeff != self._ad_dJdu and coeff != self._ad_adj_varsolver._problem.u:
+            if (
+                coeff != self._ad_dJdu
+                and coeff != self._ad_adj_varsolver._problem.u):
                 form_ad_count_map[count_map[coeff]] = coeff
-                # dict((count_map[coeff], coeff))
-                # form_ad_count_map[count_map[coeff]] = count_map[coeff]
-                # dict((count_map[coeff], coeff))
 
         for block_variable in self.get_dependencies():
             coeff = block_variable.output
@@ -677,6 +674,7 @@ def _ad_nlvs_replace_forms(self):
         self._ad_assign_coefficients(problem.J)
 
     def _ad_adj_lvs_replace_jacobian(self):
+        # Is this the correct way?
         problem = self._ad_adj_varsolver._problem
         self._ad_assign_coefficients(problem.J, solver_mode="adjoint")
 

From be9046dc11b56800fe4fbfb2cc42d8bad5db7a8f Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sun, 4 Aug 2024 23:07:49 +0100
Subject: [PATCH 06/46] Update adjoint solver coefficients correctly.

---
 firedrake/adjoint_utils/blocks/solving.py     | 119 ++++++++++--------
 firedrake/adjoint_utils/variational_solver.py |  24 ++--
 test_0.py                                     |  72 +++++++++++
 tests/regression/test_adjoint_operators.py    |  17 +--
 4 files changed, 158 insertions(+), 74 deletions(-)
 create mode 100644 test_0.py

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index dda6437951..f0ba78ea10 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -6,7 +6,8 @@
 from pyadjoint import Block, stop_annotating
 from pyadjoint.enlisting import Enlist
 import firedrake
-from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
+from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint, \
+    DelegatedFunctionCheckpoint
 from .block_utils import isconstant
 
 
@@ -203,15 +204,17 @@ def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
 
         adj_sol_bdy = None
         if compute_bdy:
-            adj_sol_bdy = firedrake.Function(
-                self.function_space.dual(),
-                dJdu_copy.dat - firedrake.assemble(
-                    firedrake.action(dFdu_adj_form, adj_sol)
-                ).dat
-            )
+            adj_sol_bdy = self.compute_adj_bdy(
+                adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
 
         return adj_sol, adj_sol_bdy
 
+    def compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
+        adj_sol_bdy = firedrake.Function(
+            self.function_space.dual(), dJdu.dat - firedrake.assemble(
+                firedrake.action(dFdu_adj_form, adj_sol)).dat)
+        return adj_sol_bdy
+
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                                prepared=None):
         if not self.linear and self.func == block_variable.output:
@@ -630,24 +633,49 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         func.assign(self._ad_nlvs._problem.u)
         return func
     
-    def _adjoint_solve(self, dJdu, adj_sol):
-        self._ad_adj_lvs_replace_jacobian()
+    def _adjoint_solve(self, dJdu):
+        # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
+        bcs = self._homogenize_bcs()
+        for bc in bcs:
+            bc.apply(dJdu)
+        problem = self._ad_adj_varsolver._problem
+        for block_variable in self.get_dependencies():
+            if block_variable.output in self.adj_F.coefficients():
+                index = 0
+                for coeff in self.adj_F.coefficients():
+                    if coeff == block_variable.output:
+                        if isinstance(
+                            block_variable.checkpoint, (
+                                firedrake.Function, firedrake.Constant,
+                                firedrake.Cofunction)):
+                            problem.J.coefficients()[index].assign(
+                                block_variable.checkpoint)
+                    elif isinstance(
+                        block_variable.checkpoint, DelegatedFunctionCheckpoint
+                    ):
+                        problem.J.coefficients()[index].assign(
+                            block_variable.checkpoint.restore())
+
+        bv = self.get_outputs()[0]
+        if bv.output in self.adj_F.coefficients():
+            index = 0
+            for coeff in self.adj_F.coefficients():
+                if coeff == bv.output:
+                    problem.J.coefficients()[index].assign(
+                            bv.checkpoint)
+                index += 1
+
         self._ad_adj_varsolver.parameters.update(self.solver_params)
-        # Replace right-hand side with dJdu.
+        # Update the right hand side of the adjoint equation.
         self._ad_dJdu.assign(dJdu)
         self._ad_adj_varsolver.solve()
-        adj_sol.assign(self._ad_adj_varsolver._problem.u)
-        return adj_sol
+        return self._ad_adj_varsolver._problem.u
 
-    def _ad_assign_map(self, form, count_map):
+    def _ad_assign_map(self, form):
+        count_map = self._ad_nlvs._problem._ad_count_map
         assign_map = {}
-        form_ad_count_map = dict()
-        for coeff in form.coefficients():
-            if (
-                coeff != self._ad_dJdu
-                and coeff != self._ad_adj_varsolver._problem.u):
-                form_ad_count_map[count_map[coeff]] = coeff
-
+        form_ad_count_map = dict((count_map[coeff], coeff)
+                                 for coeff in form.coefficients())
         for block_variable in self.get_dependencies():
             coeff = block_variable.output
             if isinstance(coeff,
@@ -659,12 +687,8 @@ def _ad_assign_map(self, form, count_map):
                         block_variable.saved_output
         return assign_map
 
-    def _ad_assign_coefficients(self, form, solver_mode="forward"):
-        if solver_mode == "forward":
-            count_map = self._ad_nlvs._problem._ad_count_map
-        else:
-            count_map = self._ad_adj_varsolver._problem._ad_count_map
-        assign_map = self._ad_assign_map(form, count_map)
+    def _ad_assign_coefficients(self, form):
+        assign_map = self._ad_assign_map(form)
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
@@ -673,11 +697,6 @@ def _ad_nlvs_replace_forms(self):
         self._ad_assign_coefficients(problem.F)
         self._ad_assign_coefficients(problem.J)
 
-    def _ad_adj_lvs_replace_jacobian(self):
-        # Is this the correct way?
-        problem = self._ad_adj_varsolver._problem
-        self._ad_assign_coefficients(problem.J, solver_mode="adjoint")
-
     def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
         if "dFdu_adj" in self._adj_cache:
             dFdu = self._adj_cache["dFdu_adj"]
@@ -688,26 +707,27 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
         return dFdu
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
-        dJdu = adj_inputs[0]
-
+        compute_bdy = self._should_compute_boundary_adjoint(
+            relevant_dependencies
+        )
         F_form = self._create_F_form()
 
-        dJdu = dJdu.copy()
-
-        # compute_bdy = self._should_compute_boundary_adjoint(
-        #     relevant_dependencies
-        # )
-        adj_sol = firedrake.Function(self.function_space)
-        self.adj_sol = self._adjoint_solve(dJdu, adj_sol)
+        dJdu = adj_inputs[0].copy()
+        self.adj_sol = self._adjoint_solve(dJdu)
+        adj_sol_bdy = None
+        if compute_bdy:
+            adj_sol_bdy = self.compute_adj_bdy(
+                self.adj_sol, adj_sol_bdy, self._ad_adj_varsolver._problem.F,
+                dJdu)
         if self.adj_cb is not None:
-            self.adj_cb(adj_sol)
-        # if self.adj_bdy_cb is not None and compute_bdy:
-        #     self.adj_bdy_cb(adj_sol_bdy)
+            self.adj_cb(self.adj_sol)
+        if self.adj_bdy_cb is not None and compute_bdy:
+            self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
         r["form"] = F_form
-        r["adj_sol"] = adj_sol
-        r["adj_sol_bdy"] = None
+        r["adj_sol"] = self.adj_sol
+        r["adj_sol_bdy"] = adj_sol_bdy
         return r
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
@@ -745,13 +765,8 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
             dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs)
             return dFdm
 
-        # dFdm_cache works with original variables, not block saved outputs.
-        if c in self._dFdm_cache:
-            dFdm = self._dFdm_cache[c]
-        else:
-            dFdm = -firedrake.derivative(self.lhs, c, trial_function)
-            dFdm = firedrake.adjoint(dFdm)
-            self._dFdm_cache[c] = dFdm
+        dFdm = -firedrake.derivative(self.lhs, c, trial_function)
+        dFdm = firedrake.adjoint(dFdm)
 
         # Replace the form coefficients with checkpointed values.
         replace_map = self._replace_map(dFdm)
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 5f5c45387a..0cb6f0c52c 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -87,11 +87,11 @@ def wrapper(self, **kwargs):
 
                 block._ad_nlvs = self._ad_nlvs
 
-                # Adjoint variational Solver.
+                # Adjoint variational solver.
                 with stop_annotating():
                     if not self._ad_adj_varsolver:
                         problem = self._ad_create_lvs_problem(
-                            problem._ad_u.function_space(), block.get_dependencies())
+                            problem._ad_u.function_space(), block)
                         self._ad_adj_varsolver = LinearVariationalSolver(
                             problem, solver_parameters=self.parameters)
 
@@ -104,6 +104,11 @@ def wrapper(self, **kwargs):
                 out = solve(self, **kwargs)
 
             if annotate:
+                # Eagerly checkpoint the outputs. This is necessary for the
+                # for the adjoint solver to work correctly. The adjoint solver
+                # will always use the checkpoint forward coefficients.
+                # tape._eagerly_checkpoint_outputs = True
+                tape._eagerly_checkpoint_outputs = True
                 block.add_output(self._ad_problem._ad_u.create_block_variable())
 
             return out
@@ -125,27 +130,32 @@ def _ad_problem_clone(self, problem, dependencies):
                                            F_replace_map[problem.u_restrict],
                                            bcs=problem.bcs,
                                            J=replace(problem.J, J_replace_map))
+        nlvp.is_linear = problem.is_linear
         nlvp._constant_jacobian = problem._constant_jacobian
         nlvp._ad_count_map_update(_ad_count_map)
         return nlvp
 
     @no_annotations
-    def _ad_create_lvs_problem(self, fwd_func_space, dependencies):
+    def _ad_create_lvs_problem(self, fwd_func_space, block):
         """Creates a LinearVariationalProblem instance for the adjoint variational problem."""
-        from firedrake import Cofunction, Function, LinearVariationalProblem, NonlinearVariationalProblem
+        from firedrake import Cofunction, Function, LinearVariationalProblem
 
         self._ad_dJdu = Cofunction(fwd_func_space.dual())
         adj_sol = Function(fwd_func_space)
+        # Homogeneous boundary conditions for the adjoint problem
+        # when Dirichlet boundary conditions are applied.
+        bcs = block._homogenize_bcs()
         tmp_problem = LinearVariationalProblem(
-            self._ad_problem._ad_adj_F, self._ad_dJdu, adj_sol)
-        _ad_count_map, F_replace_map, J_replace_map = self._build_count_map(
-            tmp_problem, dependencies)
+            block.adj_F, self._ad_dJdu, adj_sol, bcs=bcs)
+        _ad_count_map, _, J_replace_map = self._build_count_map(
+            tmp_problem, block._dependencies)
         lvp = LinearVariationalProblem(
             replace(tmp_problem.J, J_replace_map),
             self._ad_dJdu,
             adj_sol,
             bcs=tmp_problem.bcs)
         lvp._ad_count_map_update(_ad_count_map)
+        del tmp_problem
         return lvp
 
     def _build_count_map(self, problem, dependencies):
diff --git a/test_0.py b/test_0.py
new file mode 100644
index 0000000000..1cd90ca1ef
--- /dev/null
+++ b/test_0.py
@@ -0,0 +1,72 @@
+# import finat
+# from firedrake import *
+# from firedrake.adjoint import *
+# from firedrake.__future__ import Interpolator, interpolate
+# import gc
+# import numpy as np
+# from checkpoint_schedules import Revolve
+# continue_annotation()
+
+# tape = get_working_tape()
+# total_steps = 10
+# tape.enable_checkpointing(Revolve(total_steps, 2))
+
+# num_sources = 1
+# source_number = 0
+# Lx, Lz = 1.0, 1.0
+# mesh = UnitSquareMesh(50, 50)
+# V = FunctionSpace(mesh, "CG", 1)
+
+
+# def Dt(u, u_, timestep):
+#     return (u - u_)/timestep
+
+# def test_memory_burger():
+#     u_ = Function(V)
+#     u = Function(V)
+#     v = TestFunction(V)
+#     timestep = Constant(0.0000001)
+#     x,_ = SpatialCoordinate(mesh)
+#     ic = project(sin(2.*pi*x), V)
+#     u_.assign(ic)
+#     nu = Constant(0.001)
+#     F = (Dt(u, u_, timestep)*v + u*u.dx(0)*v + nu*u.dx(0)*v.dx(0))*dx
+#     bc = DirichletBC(V, 0.0, "on_boundary")
+#     t = 0.0
+#     problem = NonlinearVariationalProblem(F, u, bcs=bc)
+#     solver = NonlinearVariationalSolver(problem)
+#     t += float(timestep)
+#     for t in tape.timestepper(iter(range(total_steps))):
+#         print("step = ", t, "revolve ")
+#         solver.solve()
+#         u_.assign(u)
+#     J = assemble(u*u*dx)
+#     J_hat = ReducedFunctional(J, Control(ic))
+#     J_hat(ic)
+#     taylor_test(J_hat, ic, Function(V).assign(0.1, annotate=False))
+#     assert np.allclose(J_hat(ic), J)
+#     print(max(u_.dat.data_ro[:]))
+#     VTKFile("burger.pvd").write(u_)
+
+
+# test_memory_burger()
+# # tape.visualise()
+# print("done")
+from firedrake import *
+from pyadjoint import *
+from firedrake.adjoint import *
+
+continue_annotation()
+mesh = UnitSquareMesh(10, 10)
+V = FunctionSpace(mesh, "CG", 1)
+u = TrialFunction(V)
+v = TestFunction(V)
+a = u*v*dx
+L = Constant(1)*v*dx
+u0 = Function(V)
+solve(a == L, u0)
+j = assemble(u0 * u0 * dx)
+g = exp(j)
+dg_dtheta = compute_gradient(g, Control(u0))
+# print(rank, dg_dtheta.dat.data_ro[:])
+# pause_annotation()
\ No newline at end of file
diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 961bcbfb1f..332f2bc004 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -950,8 +950,7 @@ def test_riesz_representation_for_adjoints():
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.parametrize("constant_jacobian", [False, True])
-def test_lvs_constant_jacobian(constant_jacobian):
+def test_lvs_constant_jacobian():
     mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
     space = FunctionSpace(mesh, "Lagrange", 1)
@@ -963,21 +962,9 @@ def test_lvs_constant_jacobian(constant_jacobian):
         u_ref = u.copy(deepcopy=True)
     v = Function(space, name="v")
     problem = LinearVariationalProblem(
-        inner(trial, test) * dx, inner(u, test) * dx, v,
-        constant_jacobian=constant_jacobian)
+        inner(trial, test) * dx, inner(u, test) * dx, v)
     solver = LinearVariationalSolver(problem)
     solver.solve()
     J = assemble(v * v * dx)
-
-    assert "dFdu_adj" not in solver._ad_adj_cache
-
     dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
-
-    cached_dFdu_adj = solver._ad_adj_cache.get("dFdu_adj", None)
-    assert (cached_dFdu_adj is None) == (not constant_jacobian)
-    assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
-
-    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
-
-    assert cached_dFdu_adj is solver._ad_adj_cache.get("dFdu_adj", None)
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)

From 81da699f7169595baa4d9b6073d8cfcefb475699 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sun, 4 Aug 2024 23:12:05 +0100
Subject: [PATCH 07/46] dd

---
 firedrake/adjoint_utils/variational_solver.py |  1 -
 test_0.py                                     | 72 -------------------
 2 files changed, 73 deletions(-)
 delete mode 100644 test_0.py

diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 0cb6f0c52c..4b2f46a909 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -108,7 +108,6 @@ def wrapper(self, **kwargs):
                 # for the adjoint solver to work correctly. The adjoint solver
                 # will always use the checkpoint forward coefficients.
                 # tape._eagerly_checkpoint_outputs = True
-                tape._eagerly_checkpoint_outputs = True
                 block.add_output(self._ad_problem._ad_u.create_block_variable())
 
             return out
diff --git a/test_0.py b/test_0.py
deleted file mode 100644
index 1cd90ca1ef..0000000000
--- a/test_0.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# import finat
-# from firedrake import *
-# from firedrake.adjoint import *
-# from firedrake.__future__ import Interpolator, interpolate
-# import gc
-# import numpy as np
-# from checkpoint_schedules import Revolve
-# continue_annotation()
-
-# tape = get_working_tape()
-# total_steps = 10
-# tape.enable_checkpointing(Revolve(total_steps, 2))
-
-# num_sources = 1
-# source_number = 0
-# Lx, Lz = 1.0, 1.0
-# mesh = UnitSquareMesh(50, 50)
-# V = FunctionSpace(mesh, "CG", 1)
-
-
-# def Dt(u, u_, timestep):
-#     return (u - u_)/timestep
-
-# def test_memory_burger():
-#     u_ = Function(V)
-#     u = Function(V)
-#     v = TestFunction(V)
-#     timestep = Constant(0.0000001)
-#     x,_ = SpatialCoordinate(mesh)
-#     ic = project(sin(2.*pi*x), V)
-#     u_.assign(ic)
-#     nu = Constant(0.001)
-#     F = (Dt(u, u_, timestep)*v + u*u.dx(0)*v + nu*u.dx(0)*v.dx(0))*dx
-#     bc = DirichletBC(V, 0.0, "on_boundary")
-#     t = 0.0
-#     problem = NonlinearVariationalProblem(F, u, bcs=bc)
-#     solver = NonlinearVariationalSolver(problem)
-#     t += float(timestep)
-#     for t in tape.timestepper(iter(range(total_steps))):
-#         print("step = ", t, "revolve ")
-#         solver.solve()
-#         u_.assign(u)
-#     J = assemble(u*u*dx)
-#     J_hat = ReducedFunctional(J, Control(ic))
-#     J_hat(ic)
-#     taylor_test(J_hat, ic, Function(V).assign(0.1, annotate=False))
-#     assert np.allclose(J_hat(ic), J)
-#     print(max(u_.dat.data_ro[:]))
-#     VTKFile("burger.pvd").write(u_)
-
-
-# test_memory_burger()
-# # tape.visualise()
-# print("done")
-from firedrake import *
-from pyadjoint import *
-from firedrake.adjoint import *
-
-continue_annotation()
-mesh = UnitSquareMesh(10, 10)
-V = FunctionSpace(mesh, "CG", 1)
-u = TrialFunction(V)
-v = TestFunction(V)
-a = u*v*dx
-L = Constant(1)*v*dx
-u0 = Function(V)
-solve(a == L, u0)
-j = assemble(u0 * u0 * dx)
-g = exp(j)
-dg_dtheta = compute_gradient(g, Control(u0))
-# print(rank, dg_dtheta.dat.data_ro[:])
-# pause_annotation()
\ No newline at end of file

From 49681211c796fd143fc436d8a334aeaa00c5a13c Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sun, 4 Aug 2024 23:15:52 +0100
Subject: [PATCH 08/46] linting

---
 firedrake/adjoint_utils/blocks/solving.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index f0ba78ea10..ed8a121526 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -632,7 +632,7 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         self._ad_nlvs.solve()
         func.assign(self._ad_nlvs._problem.u)
         return func
-    
+
     def _adjoint_solve(self, dJdu):
         # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
         bcs = self._homogenize_bcs()
@@ -662,7 +662,7 @@ def _adjoint_solve(self, dJdu):
             for coeff in self.adj_F.coefficients():
                 if coeff == bv.output:
                     problem.J.coefficients()[index].assign(
-                            bv.checkpoint)
+                        bv.checkpoint)
                 index += 1
 
         self._ad_adj_varsolver.parameters.update(self.solver_params)

From 8e482d8b70789317cdf649635deb922ecb525859 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sun, 4 Aug 2024 23:57:09 +0100
Subject: [PATCH 09/46] fix condition identation

---
 firedrake/adjoint_utils/blocks/solving.py     | 11 ++++++-----
 firedrake/adjoint_utils/function.py           |  1 +
 firedrake/adjoint_utils/solving.py            |  1 +
 firedrake/adjoint_utils/variational_solver.py |  2 +-
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index ed8a121526..a974b4c0b6 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -650,11 +650,12 @@ def _adjoint_solve(self, dJdu):
                                 firedrake.Cofunction)):
                             problem.J.coefficients()[index].assign(
                                 block_variable.checkpoint)
-                    elif isinstance(
-                        block_variable.checkpoint, DelegatedFunctionCheckpoint
-                    ):
-                        problem.J.coefficients()[index].assign(
-                            block_variable.checkpoint.restore())
+                        elif isinstance(
+                            block_variable.checkpoint, DelegatedFunctionCheckpoint
+                        ):
+                            problem.J.coefficients()[index].assign(
+                                block_variable.checkpoint.restore())
+                    index += 1
 
         bv = self.get_outputs()[0]
         if bv.output in self.adj_F.coefficients():
diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py
index dc034c7155..42228b787f 100644
--- a/firedrake/adjoint_utils/function.py
+++ b/firedrake/adjoint_utils/function.py
@@ -47,6 +47,7 @@ def wrapper(self, b, *args, **kwargs):
                 output = project(self, b, *args, **kwargs)
 
             if annotate:
+                tape._eagerly_checkpoint_outputs = True
                 block.add_output(output.create_block_variable())
 
             return output
diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py
index e9556e06ea..42c5319430 100644
--- a/firedrake/adjoint_utils/solving.py
+++ b/firedrake/adjoint_utils/solving.py
@@ -61,6 +61,7 @@ def wrapper(*args, **kwargs):
                 block_variable = args[1].create_block_variable()
             else:
                 block_variable = args[1].function.create_block_variable()
+            tape._eagerly_checkpoint_outputs = True
             block.add_output(block_variable)
 
         return output
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 4b2f46a909..e41260ea1a 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -107,7 +107,7 @@ def wrapper(self, **kwargs):
                 # Eagerly checkpoint the outputs. This is necessary for the
                 # for the adjoint solver to work correctly. The adjoint solver
                 # will always use the checkpoint forward coefficients.
-                # tape._eagerly_checkpoint_outputs = True
+                tape._eagerly_checkpoint_outputs = True
                 block.add_output(self._ad_problem._ad_u.create_block_variable())
 
             return out

From 449e567a0eb5b2dc084d8bdb7786201dc95dbe25 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 5 Aug 2024 13:15:48 +0100
Subject: [PATCH 10/46] More docs; code enhacing.

---
 firedrake/adjoint_utils/blocks/solving.py     | 46 +++++++++----------
 firedrake/adjoint_utils/function.py           |  1 -
 firedrake/adjoint_utils/solving.py            |  1 -
 firedrake/adjoint_utils/variational_solver.py | 29 ++++++------
 4 files changed, 36 insertions(+), 41 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index a974b4c0b6..2921fafb66 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -638,39 +638,35 @@ def _adjoint_solve(self, dJdu):
         bcs = self._homogenize_bcs()
         for bc in bcs:
             bc.apply(dJdu)
-        problem = self._ad_adj_varsolver._problem
+        problem = self._ad_adj_lvs._problem
         for block_variable in self.get_dependencies():
+            # The self.adj_F coefficients hold the forward output
+            # (user-defined variables) references.
             if block_variable.output in self.adj_F.coefficients():
-                index = 0
-                for coeff in self.adj_F.coefficients():
-                    if coeff == block_variable.output:
-                        if isinstance(
-                            block_variable.checkpoint, (
-                                firedrake.Function, firedrake.Constant,
-                                firedrake.Cofunction)):
-                            problem.J.coefficients()[index].assign(
+                index = self.adj_F.coefficients().index(block_variable.output)
+                if isinstance(
+                        block_variable.checkpoint, (
+                            firedrake.Function, firedrake.Constant,
+                            firedrake.Cofunction)):
+                    problem.J.coefficients()[index].assign(
                                 block_variable.checkpoint)
-                        elif isinstance(
-                            block_variable.checkpoint, DelegatedFunctionCheckpoint
-                        ):
-                            problem.J.coefficients()[index].assign(
-                                block_variable.checkpoint.restore())
-                    index += 1
+                elif isinstance(
+                        block_variable.checkpoint, DelegatedFunctionCheckpoint
+                ):
+                    problem.J.coefficients()[index].assign(
+                        block_variable.checkpoint.restore())
 
         bv = self.get_outputs()[0]
         if bv.output in self.adj_F.coefficients():
-            index = 0
-            for coeff in self.adj_F.coefficients():
-                if coeff == bv.output:
-                    problem.J.coefficients()[index].assign(
-                        bv.checkpoint)
-                index += 1
+            index = self.adj_F.coefficients().index(bv.output)
+            problem.J.coefficients()[index].assign(
+                bv.checkpoint)
 
-        self._ad_adj_varsolver.parameters.update(self.solver_params)
+        self._ad_adj_lvs.parameters.update(self.solver_params)
         # Update the right hand side of the adjoint equation.
         self._ad_dJdu.assign(dJdu)
-        self._ad_adj_varsolver.solve()
-        return self._ad_adj_varsolver._problem.u
+        self._ad_adj_lvs.solve()
+        return self._ad_adj_lvs._problem.u
 
     def _ad_assign_map(self, form):
         count_map = self._ad_nlvs._problem._ad_count_map
@@ -718,7 +714,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self.compute_adj_bdy(
-                self.adj_sol, adj_sol_bdy, self._ad_adj_varsolver._problem.F,
+                self.adj_sol, adj_sol_bdy, self._ad_adj_lvs._problem.F,
                 dJdu)
         if self.adj_cb is not None:
             self.adj_cb(self.adj_sol)
diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py
index 42228b787f..dc034c7155 100644
--- a/firedrake/adjoint_utils/function.py
+++ b/firedrake/adjoint_utils/function.py
@@ -47,7 +47,6 @@ def wrapper(self, b, *args, **kwargs):
                 output = project(self, b, *args, **kwargs)
 
             if annotate:
-                tape._eagerly_checkpoint_outputs = True
                 block.add_output(output.create_block_variable())
 
             return output
diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py
index 42c5319430..e9556e06ea 100644
--- a/firedrake/adjoint_utils/solving.py
+++ b/firedrake/adjoint_utils/solving.py
@@ -61,7 +61,6 @@ def wrapper(*args, **kwargs):
                 block_variable = args[1].create_block_variable()
             else:
                 block_variable = args[1].function.create_block_variable()
-            tape._eagerly_checkpoint_outputs = True
             block.add_output(block_variable)
 
         return output
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index e41260ea1a..fddbd8157a 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -47,7 +47,7 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_kwargs = kwargs
             self._ad_nlvs = None
             self._ad_adj_cache = {}
-            self._ad_adj_varsolver = None
+            self._ad_adj_lvs = None
             self._ad_dJdu = None
 
         return wrapper
@@ -89,14 +89,14 @@ def wrapper(self, **kwargs):
 
                 # Adjoint variational solver.
                 with stop_annotating():
-                    if not self._ad_adj_varsolver:
-                        problem = self._ad_create_lvs_problem(
+                    if not self._ad_adj_lvs:
+                        problem = self._ad_adj_lvs_problem(
                             problem._ad_u.function_space(), block)
-                        self._ad_adj_varsolver = LinearVariationalSolver(
+                        self._ad_adj_lvs = LinearVariationalSolver(
                             problem, solver_parameters=self.parameters)
 
                 block._ad_dJdu = self._ad_dJdu
-                block._ad_adj_varsolver = self._ad_adj_varsolver
+                block._ad_adj_lvs = self._ad_adj_lvs
 
                 tape.add_block(block)
 
@@ -104,10 +104,6 @@ def wrapper(self, **kwargs):
                 out = solve(self, **kwargs)
 
             if annotate:
-                # Eagerly checkpoint the outputs. This is necessary for the
-                # for the adjoint solver to work correctly. The adjoint solver
-                # will always use the checkpoint forward coefficients.
-                tape._eagerly_checkpoint_outputs = True
                 block.add_output(self._ad_problem._ad_u.create_block_variable())
 
             return out
@@ -135,23 +131,28 @@ def _ad_problem_clone(self, problem, dependencies):
         return nlvp
 
     @no_annotations
-    def _ad_create_lvs_problem(self, fwd_func_space, block):
-        """Creates a LinearVariationalProblem instance for the adjoint variational problem."""
+    def _ad_adj_lvs_problem(self, fwd_func_space, block):
+        """Create the adjoint variational problem."""
         from firedrake import Cofunction, Function, LinearVariationalProblem
 
+        # Right-hand side of the adjoint equation.
         self._ad_dJdu = Cofunction(fwd_func_space.dual())
         adj_sol = Function(fwd_func_space)
         # Homogeneous boundary conditions for the adjoint problem
         # when Dirichlet boundary conditions are applied.
         bcs = block._homogenize_bcs()
+        # Create a temporary LinearVariationalProblem instance to
+        # using the adjoint form `block.adj_F`.
         tmp_problem = LinearVariationalProblem(
             block.adj_F, self._ad_dJdu, adj_sol, bcs=bcs)
+        # The `block.adj_F` coefficients hold the output references.
+        # We do not want to modify the user-defined values. Hence, the adjoint
+        # linear variational problem is created with a deep copy of the forward
+        # outputs.
         _ad_count_map, _, J_replace_map = self._build_count_map(
             tmp_problem, block._dependencies)
         lvp = LinearVariationalProblem(
-            replace(tmp_problem.J, J_replace_map),
-            self._ad_dJdu,
-            adj_sol,
+            replace(tmp_problem.J, J_replace_map), self._ad_dJdu, adj_sol,
             bcs=tmp_problem.bcs)
         lvp._ad_count_map_update(_ad_count_map)
         del tmp_problem

From b24232cd81acf4017de1dbb83fcdd751c805c8d0 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 5 Aug 2024 13:16:54 +0100
Subject: [PATCH 11/46] linting

---
 firedrake/adjoint_utils/blocks/solving.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 2921fafb66..4446d3bef8 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -649,7 +649,7 @@ def _adjoint_solve(self, dJdu):
                             firedrake.Function, firedrake.Constant,
                             firedrake.Cofunction)):
                     problem.J.coefficients()[index].assign(
-                                block_variable.checkpoint)
+                        block_variable.checkpoint)
                 elif isinstance(
                         block_variable.checkpoint, DelegatedFunctionCheckpoint
                 ):

From 4159f79600a0f758073748eeb4c0a51ad0d462d1 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 5 Aug 2024 15:16:25 +0100
Subject: [PATCH 12/46] keep dFdm cache

---
 firedrake/adjoint_utils/blocks/solving.py | 28 +++++++++++++++--------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 4446d3bef8..1a70e992be 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -633,7 +633,7 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         func.assign(self._ad_nlvs._problem.u)
         return func
 
-    def _adjoint_solve(self, dJdu):
+    def _adjoint_solve(self, dJdu, adj_sol):
         # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
         bcs = self._homogenize_bcs()
         for bc in bcs:
@@ -644,6 +644,8 @@ def _adjoint_solve(self, dJdu):
             # (user-defined variables) references.
             if block_variable.output in self.adj_F.coefficients():
                 index = self.adj_F.coefficients().index(block_variable.output)
+                # Update adjoint problem coefficients with the
+                # block_variable checkpoint values.
                 if isinstance(
                         block_variable.checkpoint, (
                             firedrake.Function, firedrake.Constant,
@@ -662,11 +664,10 @@ def _adjoint_solve(self, dJdu):
             problem.J.coefficients()[index].assign(
                 bv.checkpoint)
 
-        self._ad_adj_lvs.parameters.update(self.solver_params)
         # Update the right hand side of the adjoint equation.
         self._ad_dJdu.assign(dJdu)
         self._ad_adj_lvs.solve()
-        return self._ad_adj_lvs._problem.u
+        adj_sol.assign(self._ad_adj_lvs._problem.u)
 
     def _ad_assign_map(self, form):
         count_map = self._ad_nlvs._problem._ad_count_map
@@ -707,15 +708,19 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
+        # Forward form.
         F_form = self._create_F_form()
 
         dJdu = adj_inputs[0].copy()
-        self.adj_sol = self._adjoint_solve(dJdu)
+        adj_sol = firedrake.Function(self.function_space)
+        # Solve the adjoint equation and update the adjoint solution
+        # (`adj_sol`).
+        self._adjoint_solve(dJdu, adj_sol)
+        self.adj_sol = adj_sol
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self.compute_adj_bdy(
-                self.adj_sol, adj_sol_bdy, self._ad_adj_lvs._problem.F,
-                dJdu)
+                self.adj_sol, adj_sol_bdy, self._ad_adj_lvs._problem.F, dJdu)
         if self.adj_cb is not None:
             self.adj_cb(self.adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
@@ -723,7 +728,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
 
         r = {}
         r["form"] = F_form
-        r["adj_sol"] = self.adj_sol
+        r["adj_sol"] = adj_sol
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
 
@@ -762,8 +767,13 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
             dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs)
             return dFdm
 
-        dFdm = -firedrake.derivative(self.lhs, c, trial_function)
-        dFdm = firedrake.adjoint(dFdm)
+        # dFdm_cache works with original variables, not block saved outputs.
+        if c in self._dFdm_cache:
+            dFdm = self._dFdm_cache[c]
+        else:
+            dFdm = -firedrake.derivative(self.lhs, c, trial_function)
+            dFdm = firedrake.adjoint(dFdm)
+            self._dFdm_cache[c] = dFdm
 
         # Replace the form coefficients with checkpointed values.
         replace_map = self._replace_map(dFdm)

From 3fc41ddd5f3e61063da745589894520742ffcf76 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 5 Aug 2024 16:42:03 +0100
Subject: [PATCH 13/46] LinearSolver if jacobian is constant

---
 firedrake/adjoint_utils/blocks/solving.py     | 84 +++++++++++--------
 firedrake/adjoint_utils/variational_solver.py | 41 +++++----
 2 files changed, 74 insertions(+), 51 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 1a70e992be..8c13c9a488 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -634,40 +634,51 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         return func
 
     def _adjoint_solve(self, dJdu, adj_sol):
+
         # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
         bcs = self._homogenize_bcs()
         for bc in bcs:
             bc.apply(dJdu)
-        problem = self._ad_adj_lvs._problem
-        for block_variable in self.get_dependencies():
-            # The self.adj_F coefficients hold the forward output
-            # (user-defined variables) references.
-            if block_variable.output in self.adj_F.coefficients():
-                index = self.adj_F.coefficients().index(block_variable.output)
-                # Update adjoint problem coefficients with the
-                # block_variable checkpoint values.
-                if isinstance(
-                        block_variable.checkpoint, (
-                            firedrake.Function, firedrake.Constant,
-                            firedrake.Cofunction)):
-                    problem.J.coefficients()[index].assign(
-                        block_variable.checkpoint)
-                elif isinstance(
-                        block_variable.checkpoint, DelegatedFunctionCheckpoint
-                ):
-                    problem.J.coefficients()[index].assign(
-                        block_variable.checkpoint.restore())
-
-        bv = self.get_outputs()[0]
-        if bv.output in self.adj_F.coefficients():
-            index = self.adj_F.coefficients().index(bv.output)
-            problem.J.coefficients()[index].assign(
-                bv.checkpoint)
-
-        # Update the right hand side of the adjoint equation.
-        self._ad_dJdu.assign(dJdu)
-        self._ad_adj_lvs.solve()
-        adj_sol.assign(self._ad_adj_lvs._problem.u)
+        if isinstance(self._ad_adj_solver, firedrake.LinearVariationalSolver):
+            problem = self._ad_adj_solver._problem
+            for block_variable in self.get_dependencies():
+                # The self.adj_F coefficients hold the forward output
+                # (user-defined variables) references.
+                if block_variable.output in self.adj_F.coefficients():
+                    index = self.adj_F.coefficients().index(block_variable.output)
+                    # Update adjoint problem coefficients with the
+                    # block_variable checkpoint values.
+                    if isinstance(
+                            block_variable.checkpoint, (
+                                firedrake.Function, firedrake.Constant,
+                                firedrake.Cofunction)):
+                        problem.J.coefficients()[index].assign(
+                            block_variable.checkpoint)
+                    elif isinstance(
+                            block_variable.checkpoint, DelegatedFunctionCheckpoint
+                    ):
+                        problem.J.coefficients()[index].assign(
+                            block_variable.checkpoint.restore())
+
+            bv = self.get_outputs()[0]
+            if bv.output in self.adj_F.coefficients():
+                index = self.adj_F.coefficients().index(bv.output)
+                problem.J.coefficients()[index].assign(
+                    bv.checkpoint)
+
+            # Update the right hand side of the adjoint equation.
+            self._ad_dJdu.assign(dJdu)
+            # Solve the adjoint equation.
+            self._ad_adj_solver.solve()
+            adj_sol.assign(self._ad_adj_solver._problem.u)
+        elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
+            self._ad_dJdu.assign(dJdu)
+            self._ad_adj_solver.solve(adj_sol, self._ad_dJdu)
+        else:
+            raise NotImplementedError(
+                "Only LinearVariationalSolver and LinearSolver are supported."
+            )
+        return adj_sol
 
     def _ad_assign_map(self, form):
         count_map = self._ad_nlvs._problem._ad_count_map
@@ -715,12 +726,19 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         adj_sol = firedrake.Function(self.function_space)
         # Solve the adjoint equation and update the adjoint solution
         # (`adj_sol`).
-        self._adjoint_solve(dJdu, adj_sol)
+        adj_sol = self._adjoint_solve(dJdu, adj_sol)
         self.adj_sol = adj_sol
         adj_sol_bdy = None
         if compute_bdy:
-            adj_sol_bdy = self.compute_adj_bdy(
-                self.adj_sol, adj_sol_bdy, self._ad_adj_lvs._problem.F, dJdu)
+            if isinstance(self._ad_adj_solver, firedrake.LinearVariationalSolver):
+                adj_sol_bdy = self.compute_adj_bdy(
+                    adj_sol, adj_sol_bdy, self._ad_adj_lvs._problem.F, dJdu)
+            elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
+                print("trying to compute adjoint boundary")
+            else:
+                raise NotImplementedError(
+                    "Only LinearVariationalSolver and LinearSolver are supported."
+                )                                    
         if self.adj_cb is not None:
             self.adj_cb(self.adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index fddbd8157a..477456347c 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -47,7 +47,7 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_kwargs = kwargs
             self._ad_nlvs = None
             self._ad_adj_cache = {}
-            self._ad_adj_lvs = None
+            self._ad_adj_solver = None
             self._ad_dJdu = None
 
         return wrapper
@@ -60,7 +60,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver
+            from firedrake import LinearVariationalSolver, assemble, LinearSolver, Cofunction
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -86,17 +86,25 @@ def wrapper(self, **kwargs):
                     )
 
                 block._ad_nlvs = self._ad_nlvs
+                if not self._ad_dJdu:
+                    # Right-hand side of the adjoint equation.
+                    self._ad_dJdu = Cofunction(block.function_space.dual())
+                block._ad_dJdu = self._ad_dJdu
 
-                # Adjoint variational solver.
+                # Adjoint solver: LinearVariationalSolver if the forward
+                # problem in nonlinear, otherwise LinearSolver.
                 with stop_annotating():
-                    if not self._ad_adj_lvs:
-                        problem = self._ad_adj_lvs_problem(
-                            problem._ad_u.function_space(), block)
-                        self._ad_adj_lvs = LinearVariationalSolver(
-                            problem, solver_parameters=self.parameters)
-
-                block._ad_dJdu = self._ad_dJdu
-                block._ad_adj_lvs = self._ad_adj_lvs
+                    if not self._ad_adj_solver:
+                        if block._ad_nlvs._problem._constant_jacobian:
+                            self._ad_adj_solver = LinearSolver(
+                                assemble(block.adj_F),
+                                solver_parameters=self.parameters)
+                        else:
+                            problem = self._ad_adj_lvs_problem(block)
+                            self._ad_adj_solver = LinearVariationalSolver(
+                                    problem, solver_parameters=self.parameters)
+
+                block._ad_adj_solver = self._ad_adj_solver
 
                 tape.add_block(block)
 
@@ -131,20 +139,17 @@ def _ad_problem_clone(self, problem, dependencies):
         return nlvp
 
     @no_annotations
-    def _ad_adj_lvs_problem(self, fwd_func_space, block):
+    def _ad_adj_lvs_problem(self, block):
         """Create the adjoint variational problem."""
-        from firedrake import Cofunction, Function, LinearVariationalProblem
-
-        # Right-hand side of the adjoint equation.
-        self._ad_dJdu = Cofunction(fwd_func_space.dual())
-        adj_sol = Function(fwd_func_space)
+        from firedrake import Function, LinearVariationalProblem
+        adj_sol = Function(block.function_space)
         # Homogeneous boundary conditions for the adjoint problem
         # when Dirichlet boundary conditions are applied.
         bcs = block._homogenize_bcs()
         # Create a temporary LinearVariationalProblem instance to
         # using the adjoint form `block.adj_F`.
         tmp_problem = LinearVariationalProblem(
-            block.adj_F, self._ad_dJdu, adj_sol, bcs=bcs)
+            block.adj_F, block._ad_dJdu, adj_sol, bcs=bcs)
         # The `block.adj_F` coefficients hold the output references.
         # We do not want to modify the user-defined values. Hence, the adjoint
         # linear variational problem is created with a deep copy of the forward

From ad05a12bb5aaf5dc53efadc6e6fe77d160737cc5 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 5 Aug 2024 16:44:26 +0100
Subject: [PATCH 14/46] update test jacobian is constant

---
 firedrake/adjoint_utils/blocks/solving.py     | 3 +--
 firedrake/adjoint_utils/variational_solver.py | 2 +-
 tests/regression/test_adjoint_operators.py    | 6 ++++--
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 8c13c9a488..1600737ec3 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -737,8 +737,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
                 print("trying to compute adjoint boundary")
             else:
                 raise NotImplementedError(
-                    "Only LinearVariationalSolver and LinearSolver are supported."
-                )                                    
+                    "Only LinearVariationalSolver and LinearSolver are supported.")
         if self.adj_cb is not None:
             self.adj_cb(self.adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 477456347c..f79037d01c 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -102,7 +102,7 @@ def wrapper(self, **kwargs):
                         else:
                             problem = self._ad_adj_lvs_problem(block)
                             self._ad_adj_solver = LinearVariationalSolver(
-                                    problem, solver_parameters=self.parameters)
+                                problem, solver_parameters=self.parameters)
 
                 block._ad_adj_solver = self._ad_adj_solver
 
diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 332f2bc004..84dc7a3030 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -950,7 +950,8 @@ def test_riesz_representation_for_adjoints():
 
 
 @pytest.mark.skipcomplex
-def test_lvs_constant_jacobian():
+@pytest.mark.parametrize("constant_jacobian", [False, True])
+def test_lvs_constant_jacobian(constant_jacobian):
     mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
     space = FunctionSpace(mesh, "Lagrange", 1)
@@ -962,7 +963,8 @@ def test_lvs_constant_jacobian():
         u_ref = u.copy(deepcopy=True)
     v = Function(space, name="v")
     problem = LinearVariationalProblem(
-        inner(trial, test) * dx, inner(u, test) * dx, v)
+        inner(trial, test) * dx, inner(u, test) * dx, v,
+        constant_jacobian=constant_jacobian)
     solver = LinearVariationalSolver(problem)
     solver.solve()
     J = assemble(v * v * dx)

From 5a9d24661451b9f6fbcb0a64185cc4ce2c445a8e Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 5 Aug 2024 17:07:05 +0100
Subject: [PATCH 15/46] wip

---
 firedrake/adjoint_utils/blocks/solving.py     | 8 +++++---
 firedrake/adjoint_utils/variational_solver.py | 3 +--
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 1600737ec3..428f0e285d 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -643,7 +643,7 @@ def _adjoint_solve(self, dJdu, adj_sol):
             problem = self._ad_adj_solver._problem
             for block_variable in self.get_dependencies():
                 # The self.adj_F coefficients hold the forward output
-                # (user-defined variables) references.
+                # references.
                 if block_variable.output in self.adj_F.coefficients():
                     index = self.adj_F.coefficients().index(block_variable.output)
                     # Update adjoint problem coefficients with the
@@ -730,11 +730,13 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         self.adj_sol = adj_sol
         adj_sol_bdy = None
         if compute_bdy:
+            # TODO: I need to test this part of the code.
             if isinstance(self._ad_adj_solver, firedrake.LinearVariationalSolver):
                 adj_sol_bdy = self.compute_adj_bdy(
-                    adj_sol, adj_sol_bdy, self._ad_adj_lvs._problem.F, dJdu)
+                    adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.F, dJdu)
             elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
-                print("trying to compute adjoint boundary")
+                adj_sol_bdy = self.compute_adj_bdy(
+                    adj_sol, adj_sol_bdy, self._ad_adj_solver.A, dJdu)
             else:
                 raise NotImplementedError(
                     "Only LinearVariationalSolver and LinearSolver are supported.")
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index f79037d01c..d7dff5574d 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -91,8 +91,7 @@ def wrapper(self, **kwargs):
                     self._ad_dJdu = Cofunction(block.function_space.dual())
                 block._ad_dJdu = self._ad_dJdu
 
-                # Adjoint solver: LinearVariationalSolver if the forward
-                # problem in nonlinear, otherwise LinearSolver.
+                # Adjoint solver.
                 with stop_annotating():
                     if not self._ad_adj_solver:
                         if block._ad_nlvs._problem._constant_jacobian:

From d0046538c2b08c79fc69ce0ddf6068bdb4dcc435 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 6 Aug 2024 07:23:15 +0100
Subject: [PATCH 16/46] wip

---
 firedrake/adjoint_utils/blocks/solving.py     | 29 ++++++++-----------
 firedrake/adjoint_utils/variational_solver.py |  2 --
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 428f0e285d..ae1c058bd0 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -6,8 +6,7 @@
 from pyadjoint import Block, stop_annotating
 from pyadjoint.enlisting import Enlist
 import firedrake
-from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint, \
-    DelegatedFunctionCheckpoint
+from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
 from .block_utils import isconstant
 
 
@@ -634,11 +633,15 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         return func
 
     def _adjoint_solve(self, dJdu, adj_sol):
-
         # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
         bcs = self._homogenize_bcs()
         for bc in bcs:
             bc.apply(dJdu)
+
+        # Update the right hand side of the adjoint equation.
+        self._ad_dJdu.assign(dJdu)
+
+        # Update the left hand side coefficients of the adjoint equation.
         if isinstance(self._ad_adj_solver, firedrake.LinearVariationalSolver):
             problem = self._ad_adj_solver._problem
             for block_variable in self.get_dependencies():
@@ -646,33 +649,25 @@ def _adjoint_solve(self, dJdu, adj_sol):
                 # references.
                 if block_variable.output in self.adj_F.coefficients():
                     index = self.adj_F.coefficients().index(block_variable.output)
-                    # Update adjoint problem coefficients with the
-                    # block_variable checkpoint values.
                     if isinstance(
-                            block_variable.checkpoint, (
+                            block_variable.output, (
                                 firedrake.Function, firedrake.Constant,
                                 firedrake.Cofunction)):
+                        # `problem.J` is a deep copy of `self.adj_F`.
+                        # The indices of `self.adj_F` serve as a map for
+                        # updating the coefficients in the adjoint problem.
                         problem.J.coefficients()[index].assign(
-                            block_variable.checkpoint)
-                    elif isinstance(
-                            block_variable.checkpoint, DelegatedFunctionCheckpoint
-                    ):
-                        problem.J.coefficients()[index].assign(
-                            block_variable.checkpoint.restore())
-
+                            block_variable.saved_output)
             bv = self.get_outputs()[0]
             if bv.output in self.adj_F.coefficients():
                 index = self.adj_F.coefficients().index(bv.output)
                 problem.J.coefficients()[index].assign(
                     bv.checkpoint)
-
-            # Update the right hand side of the adjoint equation.
-            self._ad_dJdu.assign(dJdu)
             # Solve the adjoint equation.
             self._ad_adj_solver.solve()
             adj_sol.assign(self._ad_adj_solver._problem.u)
         elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
-            self._ad_dJdu.assign(dJdu)
+            # Solve the adjoint equation.
             self._ad_adj_solver.solve(adj_sol, self._ad_dJdu)
         else:
             raise NotImplementedError(
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index d7dff5574d..967104ecff 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -145,8 +145,6 @@ def _ad_adj_lvs_problem(self, block):
         # Homogeneous boundary conditions for the adjoint problem
         # when Dirichlet boundary conditions are applied.
         bcs = block._homogenize_bcs()
-        # Create a temporary LinearVariationalProblem instance to
-        # using the adjoint form `block.adj_F`.
         tmp_problem = LinearVariationalProblem(
             block.adj_F, block._ad_dJdu, adj_sol, bcs=bcs)
         # The `block.adj_F` coefficients hold the output references.

From 7c75e4341f2d77dbe2449b4fe67afc8ee0d00d7f Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 6 Aug 2024 14:29:22 +0100
Subject: [PATCH 17/46] compute adjoint bdy working

---
 .../adjoint_utils/blocks/dirichlet_bc.py      |  3 +-
 firedrake/adjoint_utils/blocks/solving.py     | 38 +++++++++----------
 firedrake/adjoint_utils/variational_solver.py | 12 +++---
 tests/regression/test_adjoint_operators.py    | 30 +++++++++++++++
 4 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/dirichlet_bc.py b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
index 7a4784812c..b06d367da1 100644
--- a/firedrake/adjoint_utils/blocks/dirichlet_bc.py
+++ b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
@@ -60,7 +60,8 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                     adj_value = firedrake.Function(self.collapsed_space, vec.dat)
 
                 if adj_value.ufl_shape == () or adj_value.ufl_shape[0] <= 1:
-                    r = adj_value.dat.data_ro.sum()
+                    R = firedrake.FunctionSpace(self.parent_space.mesh(), "R", 0)
+                    r = firedrake.Function(R.dual(), val=adj_value.dat.data_ro.sum())
                 else:
                     output = []
                     subindices = _extract_subindices(self.function_space)
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index ae1c058bd0..f378ea1c80 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -633,11 +633,6 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         return func
 
     def _adjoint_solve(self, dJdu, adj_sol):
-        # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
-        bcs = self._homogenize_bcs()
-        for bc in bcs:
-            bc.apply(dJdu)
-
         # Update the right hand side of the adjoint equation.
         self._ad_dJdu.assign(dJdu)
 
@@ -655,7 +650,7 @@ def _adjoint_solve(self, dJdu, adj_sol):
                                 firedrake.Cofunction)):
                         # `problem.J` is a deep copy of `self.adj_F`.
                         # The indices of `self.adj_F` serve as a map for
-                        # updating the coefficients in the adjoint problem.
+                        # updating the coefficients.
                         problem.J.coefficients()[index].assign(
                             block_variable.saved_output)
             bv = self.get_outputs()[0]
@@ -716,25 +711,31 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         )
         # Forward form.
         F_form = self._create_F_form()
-
-        dJdu = adj_inputs[0].copy()
+        dJdu = adj_inputs[0]
+        dJdu_copy = dJdu.copy()
         adj_sol = firedrake.Function(self.function_space)
         # Solve the adjoint equation and update the adjoint solution
         # (`adj_sol`).
+        # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
+        bcs = self._homogenize_bcs()
+        for bc in bcs:
+            bc.apply(dJdu)
         adj_sol = self._adjoint_solve(dJdu, adj_sol)
         self.adj_sol = adj_sol
         adj_sol_bdy = None
+
         if compute_bdy:
-            # TODO: I need to test this part of the code.
-            if isinstance(self._ad_adj_solver, firedrake.LinearVariationalSolver):
-                adj_sol_bdy = self.compute_adj_bdy(
-                    adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.F, dJdu)
+            if isinstance(
+                self._ad_adj_solver, firedrake.LinearVariationalSolver
+            ):
+                dFdu_adj_form = self._ad_adj_solver._problem.J
             elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
-                adj_sol_bdy = self.compute_adj_bdy(
-                    adj_sol, adj_sol_bdy, self._ad_adj_solver.A, dJdu)
-            else:
-                raise NotImplementedError(
-                    "Only LinearVariationalSolver and LinearSolver are supported.")
+                # Jacobian is constant in this case.
+                dFdu_adj_form = self.adj_F
+
+            adj_sol_bdy = self.compute_adj_bdy(
+                adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
+
         if self.adj_cb is not None:
             self.adj_cb(self.adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
@@ -793,8 +794,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         replace_map = self._replace_map(dFdm)
         replace_map[self.func] = self.get_outputs()[0].saved_output
         dFdm = replace(dFdm, replace_map)
-
-        dFdm = ufl.Action(dFdm, adj_sol)
+        dFdm = dFdm * adj_sol
         dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs)
 
         return dFdm
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 967104ecff..43ba311d49 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -94,12 +94,15 @@ def wrapper(self, **kwargs):
                 # Adjoint solver.
                 with stop_annotating():
                     if not self._ad_adj_solver:
+                        # Homogeneous boundary conditions for the adjoint problem
+                        # when Dirichlet boundary conditions are applied.
+                        bcs = block._homogenize_bcs()
                         if block._ad_nlvs._problem._constant_jacobian:
                             self._ad_adj_solver = LinearSolver(
-                                assemble(block.adj_F),
+                                assemble(block.adj_F, bcs=bcs),
                                 solver_parameters=self.parameters)
                         else:
-                            problem = self._ad_adj_lvs_problem(block)
+                            problem = self._ad_adj_lvs_problem(block, bcs)
                             self._ad_adj_solver = LinearVariationalSolver(
                                 problem, solver_parameters=self.parameters)
 
@@ -138,13 +141,10 @@ def _ad_problem_clone(self, problem, dependencies):
         return nlvp
 
     @no_annotations
-    def _ad_adj_lvs_problem(self, block):
+    def _ad_adj_lvs_problem(self, block, bcs):
         """Create the adjoint variational problem."""
         from firedrake import Function, LinearVariationalProblem
         adj_sol = Function(block.function_space)
-        # Homogeneous boundary conditions for the adjoint problem
-        # when Dirichlet boundary conditions are applied.
-        bcs = block._homogenize_bcs()
         tmp_problem = LinearVariationalProblem(
             block.adj_F, block._ad_dJdu, adj_sol, bcs=bcs)
         # The `block.adj_F` coefficients hold the output references.
diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 84dc7a3030..68c21370fc 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -970,3 +970,33 @@ def test_lvs_constant_jacobian(constant_jacobian):
     J = assemble(v * v * dx)
     dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
+
+@pytest.mark.skipcomplex
+@pytest.mark.parametrize("constant_jacobian", [False, True])
+def test_adjoint_solver_compute_bdy(constant_jacobian):
+    # Testing the case where is required to compute the adjoint
+    # boundary condition.
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+    trial = TrialFunction(space)
+    sol = Function(space, name="sol")
+    # Dirichlet boundary conditions
+    R = FunctionSpace(mesh, "R", 0)
+    a = Function(R, val=1.0)
+    b = Function(R, val=2.0)
+    bc_left = DirichletBC(space, a, 1)
+    bc_right = DirichletBC(space, b, 2)
+    bc = [bc_left, bc_right]
+    F = dot(grad(trial), grad(test)) * dx
+    problem = LinearVariationalProblem(lhs(F), rhs(F), sol, bcs=bc,
+                                       constant_jacobian=constant_jacobian)
+    solver = LinearVariationalSolver(problem)
+    solver.solve()
+    J = assemble(sol * sol * dx)
+    J_hat = ReducedFunctional(J, [Control(a), Control(b)])
+
+    assert taylor_test(
+        J_hat, [a, b], [Function(R, val=rand(1)), Function(R, val=rand(1))]
+    ) > 1.9

From bf17de84c526c5e28dab211793e41f463083162f Mon Sep 17 00:00:00 2001
From: Daiane Iglesia Dolci <63597005+Ig-dolci@users.noreply.github.com>
Date: Tue, 6 Aug 2024 14:36:54 +0100
Subject: [PATCH 18/46] small changes

---
 firedrake/adjoint_utils/blocks/solving.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index f378ea1c80..cd09194b1e 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -650,7 +650,7 @@ def _adjoint_solve(self, dJdu, adj_sol):
                                 firedrake.Cofunction)):
                         # `problem.J` is a deep copy of `self.adj_F`.
                         # The indices of `self.adj_F` serve as a map for
-                        # updating the coefficients.
+                        # updating the coefficients of the adjoint solver.
                         problem.J.coefficients()[index].assign(
                             block_variable.saved_output)
             bv = self.get_outputs()[0]
@@ -714,12 +714,12 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         dJdu = adj_inputs[0]
         dJdu_copy = dJdu.copy()
         adj_sol = firedrake.Function(self.function_space)
-        # Solve the adjoint equation and update the adjoint solution
-        # (`adj_sol`).
         # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
         bcs = self._homogenize_bcs()
         for bc in bcs:
             bc.apply(dJdu)
+        # Solve the adjoint equation and update the adjoint solution
+        # (`adj_sol`).
         adj_sol = self._adjoint_solve(dJdu, adj_sol)
         self.adj_sol = adj_sol
         adj_sol_bdy = None
@@ -737,7 +737,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
                 adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
 
         if self.adj_cb is not None:
-            self.adj_cb(self.adj_sol)
+            self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
             self.adj_bdy_cb(adj_sol_bdy)
 

From 022ef278a6710da4e205f51b539cf200e1c46966 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 6 Aug 2024 14:37:36 +0100
Subject: [PATCH 19/46] linting

---
 tests/regression/test_adjoint_operators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 68c21370fc..1e1dfdfba4 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -971,13 +971,13 @@ def test_lvs_constant_jacobian(constant_jacobian):
     dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
 
+
 @pytest.mark.skipcomplex
 @pytest.mark.parametrize("constant_jacobian", [False, True])
 def test_adjoint_solver_compute_bdy(constant_jacobian):
     # Testing the case where is required to compute the adjoint
     # boundary condition.
     mesh = UnitIntervalMesh(10)
-    X = SpatialCoordinate(mesh)
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
     trial = TrialFunction(space)

From bfcf4c7504fa3a99892c7b953661162306fc2f8f Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 6 Aug 2024 14:39:02 +0100
Subject: [PATCH 20/46] dd

---
 firedrake/adjoint_utils/blocks/solving.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index cd09194b1e..311ba14993 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -794,6 +794,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         replace_map = self._replace_map(dFdm)
         replace_map[self.func] = self.get_outputs()[0].saved_output
         dFdm = replace(dFdm, replace_map)
+
         dFdm = dFdm * adj_sol
         dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs)
 

From f5f705bf1efb57aec6a547da593aa01ebd156a27 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 6 Aug 2024 20:02:54 +0100
Subject: [PATCH 21/46] Adjoint variational solver when the jacobian is not
 constant

---
 firedrake/adjoint_utils/blocks/solving.py     | 108 +++++++++---------
 firedrake/adjoint_utils/variational_solver.py |  43 ++++---
 tests/regression/test_adjoint_operators.py    |  11 ++
 3 files changed, 84 insertions(+), 78 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 311ba14993..3232e1c542 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -632,43 +632,47 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         func.assign(self._ad_nlvs._problem.u)
         return func
 
-    def _adjoint_solve(self, dJdu, adj_sol):
+    def _adjoint_solve(self, dJdu, compute_bdy):
+        dJdu_copy = dJdu.copy()
+        # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
+        bcs = self._homogenize_bcs()
+        for bc in bcs:
+            bc.apply(dJdu)
         # Update the right hand side of the adjoint equation.
         self._ad_dJdu.assign(dJdu)
 
         # Update the left hand side coefficients of the adjoint equation.
-        if isinstance(self._ad_adj_solver, firedrake.LinearVariationalSolver):
-            problem = self._ad_adj_solver._problem
-            for block_variable in self.get_dependencies():
-                # The self.adj_F coefficients hold the forward output
-                # references.
-                if block_variable.output in self.adj_F.coefficients():
-                    index = self.adj_F.coefficients().index(block_variable.output)
-                    if isinstance(
-                            block_variable.output, (
-                                firedrake.Function, firedrake.Constant,
-                                firedrake.Cofunction)):
-                        # `problem.J` is a deep copy of `self.adj_F`.
-                        # The indices of `self.adj_F` serve as a map for
-                        # updating the coefficients of the adjoint solver.
-                        problem.J.coefficients()[index].assign(
-                            block_variable.saved_output)
-            bv = self.get_outputs()[0]
-            if bv.output in self.adj_F.coefficients():
-                index = self.adj_F.coefficients().index(bv.output)
-                problem.J.coefficients()[index].assign(
-                    bv.checkpoint)
-            # Solve the adjoint equation.
-            self._ad_adj_solver.solve()
-            adj_sol.assign(self._ad_adj_solver._problem.u)
-        elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
-            # Solve the adjoint equation.
-            self._ad_adj_solver.solve(adj_sol, self._ad_dJdu)
-        else:
-            raise NotImplementedError(
-                "Only LinearVariationalSolver and LinearSolver are supported."
-            )
-        return adj_sol
+        problem = self._ad_adj_solver._problem
+        for block_variable in self.get_dependencies():
+            # The self.adj_F coefficients hold the forward output
+            # references.
+            if block_variable.output in self.adj_F.coefficients():
+                index = self.adj_F.coefficients().index(block_variable.output)
+                if isinstance(
+                        block_variable.output, (
+                            firedrake.Function, firedrake.Constant,
+                            firedrake.Cofunction)):
+                    # `problem.J` is a deep copy of `self.adj_F`.
+                    # The indices of `self.adj_F` serve as a map for
+                    # updating the coefficients of the adjoint solver.
+                    problem.J.coefficients()[index].assign(
+                        block_variable.saved_output)
+        bv = self.get_outputs()[0]
+        if bv.output in self.adj_F.coefficients():
+            index = self.adj_F.coefficients().index(bv.output)
+            problem.J.coefficients()[index].assign(
+                bv.checkpoint)
+        # Solve the adjoint equation.
+        self._ad_adj_solver.solve()
+
+        adj_sol = firedrake.Function(self.function_space)
+        adj_sol.assign(self._ad_adj_solver._problem.u)
+        adj_sol_bdy = None
+        if compute_bdy:
+            adj_sol_bdy = self.compute_adj_bdy(
+                adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.J,
+                dJdu_copy)
+        return adj_sol, adj_sol_bdy
 
     def _ad_assign_map(self, form):
         count_map = self._ad_nlvs._problem._ad_count_map
@@ -706,36 +710,30 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
         return dFdu
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
+        dJdu = adj_inputs[0]
+        dJdu = dJdu.copy()
+
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
         # Forward form.
         F_form = self._create_F_form()
-        dJdu = adj_inputs[0]
-        dJdu_copy = dJdu.copy()
-        adj_sol = firedrake.Function(self.function_space)
-        # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
-        bcs = self._homogenize_bcs()
-        for bc in bcs:
-            bc.apply(dJdu)
-        # Solve the adjoint equation and update the adjoint solution
-        # (`adj_sol`).
-        adj_sol = self._adjoint_solve(dJdu, adj_sol)
-        self.adj_sol = adj_sol
-        adj_sol_bdy = None
 
-        if compute_bdy:
-            if isinstance(
-                self._ad_adj_solver, firedrake.LinearVariationalSolver
-            ):
-                dFdu_adj_form = self._ad_adj_solver._problem.J
-            elif isinstance(self._ad_adj_solver, firedrake.LinearSolver):
-                # Jacobian is constant in this case.
-                dFdu_adj_form = self.adj_F
+        if self._ad_nlvs._problem._constant_jacobian:
+            dFdu_form = self.adj_F
+            # Replace the form coefficients with checkpointed values.
+            replace_map = self._replace_map(dFdu_form)
+            replace_map[self.func] = self.get_outputs()[0].saved_output
+            dFdu_form = replace(dFdu_form, replace_map)
 
-            adj_sol_bdy = self.compute_adj_bdy(
-                adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
+            adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
+                dFdu_form, dJdu, compute_bdy
+            )
+        else:
+            # Adjoint solver using linear variational solver.
+            adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
 
+        self.adj_sol = 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:
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 43ba311d49..74c454bd94 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -60,7 +60,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver, assemble, LinearSolver, Cofunction
+            from firedrake import LinearVariationalSolver, Cofunction
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -86,27 +86,21 @@ def wrapper(self, **kwargs):
                     )
 
                 block._ad_nlvs = self._ad_nlvs
-                if not self._ad_dJdu:
-                    # Right-hand side of the adjoint equation.
-                    self._ad_dJdu = Cofunction(block.function_space.dual())
-                block._ad_dJdu = self._ad_dJdu
-
-                # Adjoint solver.
-                with stop_annotating():
-                    if not self._ad_adj_solver:
-                        # Homogeneous boundary conditions for the adjoint problem
-                        # when Dirichlet boundary conditions are applied.
-                        bcs = block._homogenize_bcs()
-                        if block._ad_nlvs._problem._constant_jacobian:
-                            self._ad_adj_solver = LinearSolver(
-                                assemble(block.adj_F, bcs=bcs),
-                                solver_parameters=self.parameters)
-                        else:
-                            problem = self._ad_adj_lvs_problem(block, bcs)
-                            self._ad_adj_solver = LinearVariationalSolver(
-                                problem, solver_parameters=self.parameters)
-
-                block._ad_adj_solver = self._ad_adj_solver
+                if not self._ad_problem._constant_jacobian:
+                    if not self._ad_dJdu:
+                        # Right-hand side of the adjoint equation.
+                        self._ad_dJdu = Cofunction(block.function_space.dual())
+                    block._ad_dJdu = self._ad_dJdu
+
+                    # Adjoint solver.
+                    if not self._ad_problem._constant_jacobian:
+                        with stop_annotating():
+                            if not self._ad_adj_solver:
+                                problem = self._ad_adj_lvs_problem(block)
+                                self._ad_adj_solver = LinearVariationalSolver(
+                                    problem, solver_parameters=self.parameters)
+
+                    block._ad_adj_solver = self._ad_adj_solver
 
                 tape.add_block(block)
 
@@ -141,9 +135,12 @@ def _ad_problem_clone(self, problem, dependencies):
         return nlvp
 
     @no_annotations
-    def _ad_adj_lvs_problem(self, block, bcs):
+    def _ad_adj_lvs_problem(self, block):
         """Create the adjoint variational problem."""
         from firedrake import Function, LinearVariationalProblem
+        # Homogeneous boundary conditions for the adjoint problem
+        # when Dirichlet boundary conditions are applied.
+        bcs = block._homogenize_bcs()
         adj_sol = Function(block.function_space)
         tmp_problem = LinearVariationalProblem(
             block.adj_F, block._ad_dJdu, adj_sol, bcs=bcs)
diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 1e1dfdfba4..9588a3bf5b 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -968,7 +968,18 @@ def test_lvs_constant_jacobian(constant_jacobian):
     solver = LinearVariationalSolver(problem)
     solver.solve()
     J = assemble(v * v * dx)
+
+    assert "dFdu_adj" not in solver._ad_adj_cache
+
     dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
+
+    cached_dFdu_adj = solver._ad_adj_cache.get("dFdu_adj", None)
+    assert (cached_dFdu_adj is None) == (not constant_jacobian)
+    assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
+
+    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
+
+    assert cached_dFdu_adj is solver._ad_adj_cache.get("dFdu_adj", None)
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
 
 

From bc205ab4a3c66c8230a50a3306e4a26aa42fbe86 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 6 Aug 2024 21:22:24 +0100
Subject: [PATCH 22/46] compute adjoint boundary as a private method

---
 firedrake/adjoint_utils/blocks/solving.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 3232e1c542..1bd97bc810 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -203,12 +203,12 @@ def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
 
         adj_sol_bdy = None
         if compute_bdy:
-            adj_sol_bdy = self.compute_adj_bdy(
+            adj_sol_bdy = self._compute_adj_bdy(
                 adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
 
         return adj_sol, adj_sol_bdy
 
-    def compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
+    def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
         adj_sol_bdy = firedrake.Function(
             self.function_space.dual(), dJdu.dat - firedrake.assemble(
                 firedrake.action(dFdu_adj_form, adj_sol)).dat)
@@ -669,7 +669,7 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         adj_sol.assign(self._ad_adj_solver._problem.u)
         adj_sol_bdy = None
         if compute_bdy:
-            adj_sol_bdy = self.compute_adj_bdy(
+            adj_sol_bdy = self._compute_adj_bdy(
                 adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.J,
                 dJdu_copy)
         return adj_sol, adj_sol_bdy

From 0c05743a60b22670ceb8c436fa127f0e8589282b Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Fri, 23 Aug 2024 14:07:23 +0100
Subject: [PATCH 23/46] Remove dJdu attribute; define form and boundary
 variables when necessary.

---
 firedrake/adjoint_utils/blocks/solving.py     | 25 +++++++++++--------
 firedrake/adjoint_utils/variational_solver.py | 14 ++++-------
 2 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 1bd97bc810..99a4838dc8 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -8,7 +8,7 @@
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
 from .block_utils import isconstant
-
+import weakref
 
 def extract_subfunction(u, V):
     """If V is a subspace of the function-space of u, return the component of u that is in that subspace."""
@@ -48,7 +48,7 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         # Equation RHS
         self.rhs = rhs
         # Solution function
-        self.func = func
+        self._func = weakref.ref(func)
         self.function_space = self.func.function_space()
         # Boundary conditions
         self.bcs = []
@@ -72,6 +72,10 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         self.add_dependency(mesh)
         self._init_solver_parameters(args, kwargs)
 
+    @property
+    def func(self):
+        return self._func()
+
     def _init_solver_parameters(self, args, kwargs):
         self.forward_args = kwargs.pop("forward_args", [])
         self.forward_kwargs = kwargs.pop("forward_kwargs", {})
@@ -638,8 +642,6 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         bcs = self._homogenize_bcs()
         for bc in bcs:
             bc.apply(dJdu)
-        # Update the right hand side of the adjoint equation.
-        self._ad_dJdu.assign(dJdu)
 
         # Update the left hand side coefficients of the adjoint equation.
         problem = self._ad_adj_solver._problem
@@ -655,8 +657,11 @@ def _adjoint_solve(self, dJdu, compute_bdy):
                     # `problem.J` is a deep copy of `self.adj_F`.
                     # The indices of `self.adj_F` serve as a map for
                     # updating the coefficients of the adjoint solver.
-                    problem.J.coefficients()[index].assign(
+                    problem.F.coefficients()[index].assign(
                         block_variable.saved_output)
+        # Update the right hand side of the adjoint equation.
+        # problem.F._component[1] is the right hand side of the adjoint.
+        problem.F._components[1].assign(dJdu)
         bv = self.get_outputs()[0]
         if bv.output in self.adj_F.coefficients():
             index = self.adj_F.coefficients().index(bv.output)
@@ -716,8 +721,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
-        # Forward form.
-        F_form = self._create_F_form()
 
         if self._ad_nlvs._problem._constant_jacobian:
             dFdu_form = self.adj_F
@@ -740,7 +743,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
-        r["form"] = F_form
         r["adj_sol"] = adj_sol
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
@@ -750,25 +752,26 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if not self.linear and self.func == block_variable.output:
             # We are not able to calculate derivatives wrt initial guess.
             return None
-        F_form = prepared["form"]
         adj_sol = prepared["adj_sol"]
-        adj_sol_bdy = prepared["adj_sol_bdy"]
         c = block_variable.output
-        c_rep = block_variable.saved_output
 
         if isinstance(c, firedrake.Function):
             trial_function = firedrake.TrialFunction(c.function_space())
         elif isinstance(c, firedrake.Constant):
+            F_form = self._create_F_form()
             mesh = F_form.ufl_domain()
             trial_function = firedrake.TrialFunction(
                 c._ad_function_space(mesh)
             )
         elif isinstance(c, firedrake.DirichletBC):
+            adj_sol_bdy = prepared["adj_sol_bdy"]
             tmp_bc = c.reconstruct(
                 g=extract_subfunction(adj_sol_bdy, c.function_space())
             )
             return [tmp_bc]
         elif isinstance(c, firedrake.MeshGeometry):
+            c_rep = block_variable.saved_output
+            F_form = self._create_F_form()
             # Using CoordianteDerivative requires us to do action before
             # differentiating, might change in the future.
             F_form_tmp = firedrake.action(F_form, adj_sol)
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 74c454bd94..7064d3be17 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -48,7 +48,6 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_nlvs = None
             self._ad_adj_cache = {}
             self._ad_adj_solver = None
-            self._ad_dJdu = None
 
         return wrapper
 
@@ -87,11 +86,6 @@ def wrapper(self, **kwargs):
 
                 block._ad_nlvs = self._ad_nlvs
                 if not self._ad_problem._constant_jacobian:
-                    if not self._ad_dJdu:
-                        # Right-hand side of the adjoint equation.
-                        self._ad_dJdu = Cofunction(block.function_space.dual())
-                    block._ad_dJdu = self._ad_dJdu
-
                     # Adjoint solver.
                     if not self._ad_problem._constant_jacobian:
                         with stop_annotating():
@@ -137,13 +131,15 @@ def _ad_problem_clone(self, problem, dependencies):
     @no_annotations
     def _ad_adj_lvs_problem(self, block):
         """Create the adjoint variational problem."""
-        from firedrake import Function, LinearVariationalProblem
+        from firedrake import Function, Cofunction, LinearVariationalProblem
         # Homogeneous boundary conditions for the adjoint problem
         # when Dirichlet boundary conditions are applied.
         bcs = block._homogenize_bcs()
         adj_sol = Function(block.function_space)
+        right_hand_side = Cofunction(block.function_space.dual())
         tmp_problem = LinearVariationalProblem(
-            block.adj_F, block._ad_dJdu, adj_sol, bcs=bcs)
+            block.adj_F, right_hand_side,
+            adj_sol, bcs=bcs)
         # The `block.adj_F` coefficients hold the output references.
         # We do not want to modify the user-defined values. Hence, the adjoint
         # linear variational problem is created with a deep copy of the forward
@@ -151,7 +147,7 @@ def _ad_adj_lvs_problem(self, block):
         _ad_count_map, _, J_replace_map = self._build_count_map(
             tmp_problem, block._dependencies)
         lvp = LinearVariationalProblem(
-            replace(tmp_problem.J, J_replace_map), self._ad_dJdu, adj_sol,
+            replace(tmp_problem.J, J_replace_map), right_hand_side, adj_sol,
             bcs=tmp_problem.bcs)
         lvp._ad_count_map_update(_ad_count_map)
         del tmp_problem

From aaac6fd530cdf8048a84546e4e669e364a6b25ce Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Fri, 23 Aug 2024 14:13:53 +0100
Subject: [PATCH 24/46] lint

---
 firedrake/adjoint_utils/blocks/solving.py     | 1 +
 firedrake/adjoint_utils/variational_solver.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 99a4838dc8..28becdd8cc 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -10,6 +10,7 @@
 from .block_utils import isconstant
 import weakref
 
+
 def extract_subfunction(u, V):
     """If V is a subspace of the function-space of u, return the component of u that is in that subspace."""
     if V.index is not None:
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 7064d3be17..e1f28c2a73 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver, Cofunction
+            from firedrake import LinearVariationalSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()

From 342f093104ec7ae97c9e17336a7ab27631daaf99 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 26 Aug 2024 09:20:34 +0100
Subject: [PATCH 25/46] Update the problem.J coefficients

---
 firedrake/adjoint_utils/blocks/solving.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 28becdd8cc..e7f9283b3b 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -655,10 +655,11 @@ def _adjoint_solve(self, dJdu, compute_bdy):
                         block_variable.output, (
                             firedrake.Function, firedrake.Constant,
                             firedrake.Cofunction)):
-                    # `problem.J` is a deep copy of `self.adj_F`.
+                    # `problem.J` (Jacobian operator) is a deep copy of
+                    # `self.adj_F`.
                     # The indices of `self.adj_F` serve as a map for
                     # updating the coefficients of the adjoint solver.
-                    problem.F.coefficients()[index].assign(
+                    problem.J.coefficients()[index].assign(
                         block_variable.saved_output)
         # Update the right hand side of the adjoint equation.
         # problem.F._component[1] is the right hand side of the adjoint.
@@ -753,26 +754,25 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if not self.linear and self.func == block_variable.output:
             # We are not able to calculate derivatives wrt initial guess.
             return None
+        F_form = self._create_F_form()
         adj_sol = prepared["adj_sol"]
+        adj_sol_bdy = prepared["adj_sol_bdy"]
         c = block_variable.output
+        c_rep = block_variable.saved_output
 
         if isinstance(c, firedrake.Function):
             trial_function = firedrake.TrialFunction(c.function_space())
         elif isinstance(c, firedrake.Constant):
-            F_form = self._create_F_form()
             mesh = F_form.ufl_domain()
             trial_function = firedrake.TrialFunction(
                 c._ad_function_space(mesh)
             )
         elif isinstance(c, firedrake.DirichletBC):
-            adj_sol_bdy = prepared["adj_sol_bdy"]
             tmp_bc = c.reconstruct(
                 g=extract_subfunction(adj_sol_bdy, c.function_space())
             )
             return [tmp_bc]
         elif isinstance(c, firedrake.MeshGeometry):
-            c_rep = block_variable.saved_output
-            F_form = self._create_F_form()
             # Using CoordianteDerivative requires us to do action before
             # differentiating, might change in the future.
             F_form_tmp = firedrake.action(F_form, adj_sol)

From bc3d7658347db59be44e6fd530e71f37b2ee25a0 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 26 Aug 2024 19:56:02 +0100
Subject: [PATCH 26/46] Delegated block solvers and forms.

---
 firedrake/adjoint_utils/blocks/__init__.py    |   3 +-
 firedrake/adjoint_utils/blocks/solving.py     | 177 ++++++++++++++----
 firedrake/adjoint_utils/function.py           |  21 +--
 firedrake/adjoint_utils/variational_solver.py |  39 ++--
 4 files changed, 183 insertions(+), 57 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py
index bf83b896cc..354c342b19 100644
--- a/firedrake/adjoint_utils/blocks/__init__.py
+++ b/firedrake/adjoint_utils/blocks/__init__.py
@@ -1,7 +1,8 @@
 from .assembly import AssembleBlock  # NOQA F401
 from .solving import GenericSolveBlock, SolveLinearSystemBlock, \
     ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock, \
-    NonlinearVariationalSolveBlock  # NOQA F401
+    NonlinearVariationalSolveBlock, SolverType, DelegatedBlockSolver, \
+    DelegatedBlockForm, FormType  # NOQA F401
 from .function import FunctionAssignBlock, FunctionMergeBlock, \
     SubfunctionBlock  # NOQA F401
 from .dirichlet_bc import DirichletBCBlock  # NOQA F401
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index e7f9283b3b..6a98962114 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -1,9 +1,11 @@
 import numpy
+from enum import Enum
+from abc import ABC, abstractmethod
 import ufl
 from ufl import replace
 from ufl.formatting.ufl2unicode import ufl2unicode
 
-from pyadjoint import Block, stop_annotating
+from pyadjoint import Block, stop_annotating, get_working_tape
 from pyadjoint.enlisting import Enlist
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
@@ -85,18 +87,33 @@ def _init_solver_parameters(self, args, kwargs):
         self.assemble_kwargs = {}
 
     def __str__(self):
-        return "solve({} = {})".format(ufl2unicode(self.lhs),
-                                       ufl2unicode(self.rhs))
+        lhs = self.lhs
+        rhs = self.rhs
+        if isinstance(lhs, DelegatedBlockForm):
+            lhs = lhs.restore_form()
+        if isinstance(rhs, DelegatedBlockForm):
+            rhs = rhs.restore_form()
+        return "solve({} = {})".format(ufl2unicode(lhs),
+                                       ufl2unicode(rhs))
 
     def _create_F_form(self):
         # Process the equation forms, replacing values with checkpoints,
         # and gathering lhs and rhs in one single form.
         if self.linear:
             tmp_u = firedrake.Function(self.function_space)
-            F_form = firedrake.action(self.lhs, tmp_u) - self.rhs
+            lhs = self.lhs
+            if isinstance(lhs, DelegatedBlockForm):
+                lhs = lhs.restore_form()
+            rhs = self.rhs
+            if isinstance(rhs, DelegatedBlockForm):
+                rhs = self.rhs.restore_form()
+            F_form = firedrake.action(lhs, tmp_u) - rhs
         else:
             tmp_u = self.func
-            F_form = self.lhs
+            lhs = self.lhs
+            if isinstance(lhs, DelegatedBlockForm):
+                lhs = lhs.restore_form()
+            F_form = lhs
 
         replace_map = self._replace_map(F_form)
         replace_map[tmp_u] = self.get_outputs()[0].saved_output
@@ -509,10 +526,16 @@ def _replace_recompute_form(self):
         func = self._create_initial_guess()
 
         bcs = self._recover_bcs()
-        lhs = self._replace_form(self.lhs, func=func)
+        lhs = self.lhs
+        if isinstance(lhs, DelegatedBlockForm):
+            lhs = lhs.restore_form()
+        lhs = self._replace_form(lhs, func=func)
         rhs = 0
         if self.linear:
-            rhs = self._replace_form(self.rhs)
+            rhs = self.rhs
+            if isinstance(rhs, DelegatedBlockForm):
+                rhs = rhs.restore_form()
+            rhs = self._replace_form(rhs)
 
         return lhs, rhs, func, bcs
 
@@ -616,14 +639,13 @@ def __init__(self, equation, func, bcs, adj_F, adj_cache, problem_J,
         self.adj_F = adj_F
         self._adj_cache = adj_cache
         self._dFdm_cache = adj_cache.setdefault("dFdm_cache", {})
-        self.problem_J = problem_J
         self.solver_params = solver_params.copy()
         self.solver_kwargs = solver_kwargs
 
         super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs})
 
-        if self.problem_J is not None:
-            for coeff in self.problem_J.coefficients():
+        if problem_J is not None:
+            for coeff in problem_J.coefficients():
                 self.add_dependency(coeff, no_duplicates=True)
 
     def _init_solver_parameters(self, args, kwargs):
@@ -632,9 +654,13 @@ def _init_solver_parameters(self, args, kwargs):
 
     def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         self._ad_nlvs_replace_forms()
-        self._ad_nlvs.parameters.update(self.solver_params)
-        self._ad_nlvs.solve()
-        func.assign(self._ad_nlvs._problem.u)
+        nlvs = self._ad_nlvs
+        if isinstance(nlvs, DelegatedBlockSolver):
+            nlvs = self._ad_nlvs.restore_solver()
+        nlvs.parameters.update(self.solver_params)
+        nlvs.solve()
+        func.assign(nlvs._problem.u)
+        del nlvs
         return func
 
     def _adjoint_solve(self, dJdu, compute_bdy):
@@ -645,44 +671,53 @@ def _adjoint_solve(self, dJdu, compute_bdy):
             bc.apply(dJdu)
 
         # Update the left hand side coefficients of the adjoint equation.
-        problem = self._ad_adj_solver._problem
+        adj_solver = self._ad_adj_solver
+        if isinstance(adj_solver, DelegatedBlockSolver):
+            adj_solver = adj_solver.restore_solver()
+        problem = adj_solver._problem
         for block_variable in self.get_dependencies():
             # The self.adj_F coefficients hold the forward output
             # references.
-            if block_variable.output in self.adj_F.coefficients():
-                index = self.adj_F.coefficients().index(block_variable.output)
+            adj_F = self.adj_F
+            if isinstance(adj_F, DelegatedBlockForm):
+                adj_F = adj_F.restore_form()
+            if block_variable.output in adj_F.coefficients():
+                index = adj_F.coefficients().index(
+                    block_variable.output)
                 if isinstance(
                         block_variable.output, (
                             firedrake.Function, firedrake.Constant,
                             firedrake.Cofunction)):
                     # `problem.J` (Jacobian operator) is a deep copy of
-                    # `self.adj_F`.
-                    # The indices of `self.adj_F` serve as a map for
-                    # updating the coefficients of the adjoint solver.
+                    # `self.adj_F`. Hence, the indices of `self.adj_F`
+                    # is serving as a map for updating the coefficients of
+                    # the adjoint solver.
                     problem.J.coefficients()[index].assign(
                         block_variable.saved_output)
         # Update the right hand side of the adjoint equation.
         # problem.F._component[1] is the right hand side of the adjoint.
         problem.F._components[1].assign(dJdu)
         bv = self.get_outputs()[0]
-        if bv.output in self.adj_F.coefficients():
-            index = self.adj_F.coefficients().index(bv.output)
+        if bv.output in adj_F.coefficients():
+            index = adj_F.coefficients().index(bv.output)
             problem.J.coefficients()[index].assign(
                 bv.checkpoint)
         # Solve the adjoint equation.
-        self._ad_adj_solver.solve()
+        adj_solver.solve()
 
         adj_sol = firedrake.Function(self.function_space)
-        adj_sol.assign(self._ad_adj_solver._problem.u)
+        adj_sol.assign(adj_solver._problem.u)
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self._compute_adj_bdy(
-                adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.J,
-                dJdu_copy)
+                adj_sol, adj_sol_bdy, adj_solver._problem.J, dJdu_copy)
         return adj_sol, adj_sol_bdy
 
     def _ad_assign_map(self, form):
-        count_map = self._ad_nlvs._problem._ad_count_map
+        nlvs = self._ad_nlvs
+        if isinstance(nlvs, DelegatedBlockSolver):
+            nlvs = nlvs.restore_solver()
+        count_map = nlvs._problem._ad_count_map
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -703,7 +738,10 @@ def _ad_assign_coefficients(self, form):
             coeff.assign(value)
 
     def _ad_nlvs_replace_forms(self):
-        problem = self._ad_nlvs._problem
+        nlvs = self._ad_nlvs
+        if isinstance(nlvs, DelegatedBlockSolver):
+            nlvs = nlvs.restore_solver()
+        problem = nlvs._problem
         self._ad_assign_coefficients(problem.F)
         self._ad_assign_coefficients(problem.J)
 
@@ -712,7 +750,10 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
             dFdu = self._adj_cache["dFdu_adj"]
         else:
             dFdu = super()._assemble_dFdu_adj(dFdu_adj_form, **kwargs)
-            if self._ad_nlvs._problem._constant_jacobian:
+            nlvs = self._ad_nlvs
+            if isinstance(nlvs, DelegatedBlockSolver):
+                nlvs = nlvs.restore_solver(SolverType.FORWARD)
+            if nlvs._constant_jacobian:
                 self._adj_cache["dFdu_adj"] = dFdu
         return dFdu
 
@@ -723,9 +764,14 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
-
-        if self._ad_nlvs._problem._constant_jacobian:
-            dFdu_form = self.adj_F
+        nlvs = self._ad_nlvs
+        if isinstance(nlvs, DelegatedBlockSolver):
+            nlvs = nlvs.restore_solver()
+        if nlvs._problem._constant_jacobian:
+            adj_F = self.adj_F
+            if isinstance(adj_F, DelegatedBlockForm):
+                adj_F = adj_F.restore_form()
+            dFdu_form = adj_F
             # Replace the form coefficients with checkpointed values.
             replace_map = self._replace_map(dFdu_form)
             replace_map[self.func] = self.get_outputs()[0].saved_output
@@ -788,7 +834,10 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if c in self._dFdm_cache:
             dFdm = self._dFdm_cache[c]
         else:
-            dFdm = -firedrake.derivative(self.lhs, c, trial_function)
+            lhs = self.lhs
+            if isinstance(lhs, DelegatedBlockForm):
+                lhs = lhs.restore_form()
+            dFdm = -firedrake.derivative(lhs, c, trial_function)
             dFdm = firedrake.adjoint(dFdm)
             self._dFdm_cache[c] = dFdm
 
@@ -959,3 +1008,67 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
     def __str__(self):
         target_string = f"〈{str(self.target_space.ufl_element().shortstr())}〉"
         return f"project({self.get_dependencies()[0]}, {target_string}))"
+
+
+class SolverType(Enum):
+    FORWARD = 1
+    ADJOINT = 2
+
+
+class FormType(Enum):
+    ADJOINT = 1
+    LHS = 2
+    RHS = 3
+
+
+class Solver(dict):
+    def __init__(self, forward, adjoint):
+        super().__init__()
+        self.solvers = {SolverType.FORWARD: forward, SolverType.ADJOINT: adjoint}
+
+    def __setitem__(self, key, value):
+        if key in self.solvers:
+            self.solvers[key] = value
+        super().__setitem__(key, value)
+
+    def __getitem__(self, key):
+        if key in self.solvers:
+            return self.solvers[key]
+        return super().__getitem__(key)
+
+    def finalize(self):
+        f = weakref.finalize(self, self.solvers)
+        f.atexit = False
+
+
+class DelegatedBlockSolver:
+    def __init__(self, delegated_block, solver_type):
+        # Block number that this solver is delegated to.
+        self.delegated_block = delegated_block
+        self.solver_type = solver_type
+
+    def restore_solver(self):
+        block = get_working_tape()._blocks[self.delegated_block]
+        if self.solver_type == SolverType.FORWARD:
+            return block._ad_nlvs
+        elif self.solver_type == SolverType.ADJOINT:
+            return block._ad_adj_solver
+
+
+class DelegatedBlockForm:
+    def __init__(self, delegated_block, form_type):
+        self.delegated_block = delegated_block
+        self.form_type = form_type
+
+    def restore_form(self):
+        if self.form_type == FormType.ADJOINT:
+            return get_working_tape()._blocks[self.delegated_block].adj_F
+        elif self.form_type == FormType.LHS:
+            return get_working_tape()._blocks[self.delegated_block].lhs
+        elif self.form_type == FormType.RHS:
+            return get_working_tape()._blocks[self.delegated_block].rhs
+        else:
+            raise NotImplementedError(
+                "Only adjoint form is supported for now. If you need to "
+                "restore another form, please create a new form type."
+            )
\ No newline at end of file
diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py
index 4cd7e40269..4bdc0666a6 100644
--- a/firedrake/adjoint_utils/function.py
+++ b/firedrake/adjoint_utils/function.py
@@ -219,6 +219,15 @@ def _ad_create_checkpoint(self):
             return CheckpointFunction(self)
         else:
             return self.copy(deepcopy=True)
+        
+    def _ad_clear_checkpoint(self, checkpoint):
+        # DelegatedFunctionCheckpoint does not hold data. Actually,
+        # it is a reference to another checkpoint function. Thus, clear
+        # can be unsafe once we will lose the reference to the original
+        # checkpoint.
+        if not isinstance(checkpoint, DelegatedFunctionCheckpoint):
+            checkpoint = None
+        return checkpoint
 
     def _ad_convert_riesz(self, value, options=None):
         from firedrake import Function, Cofunction
@@ -398,17 +407,5 @@ def _applyBinary(self, f, y):
             npdata[i] = f(npdata[i], npdatay[i])
         vec.set_local(npdata)
 
-    def _ad_from_petsc(self, vec):
-        with self.dat.vec_wo as self_v:
-            vec.copy(result=self_v)
-
-    def _ad_to_petsc(self, vec=None):
-        with self.dat.vec_ro as self_v:
-            if vec:
-                self_v.copy(result=vec)
-            else:
-                vec = self_v.copy()
-        return vec
-
     def __deepcopy__(self, memodict={}):
         return self.copy(deepcopy=True)
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index e1f28c2a73..d0efd46003 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -1,7 +1,8 @@
 import copy
 from functools import wraps
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
-from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
+from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock, \
+    DelegatedBlockSolver, DelegatedBlockForm, FormType, SolverType
 from ufl import replace
 
 
@@ -48,6 +49,7 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_nlvs = None
             self._ad_adj_cache = {}
             self._ad_adj_solver = None
+            self._ad_block_to_delegate = None
 
         return wrapper
 
@@ -83,18 +85,30 @@ def wrapper(self, **kwargs):
                         self._ad_problem_clone(self._ad_problem, block.get_dependencies()),
                         **self._ad_kwargs
                     )
-
-                block._ad_nlvs = self._ad_nlvs
+                    block._ad_nlvs = self._ad_nlvs
+                    self._ad_block_to_delegate = len(tape._blocks)
+                else:
+                    block._ad_nlvs = DelegatedBlockSolver(
+                        self._ad_block_to_delegate, SolverType.FORWARD)
+                    block.adj_F = DelegatedBlockForm(
+                        self._ad_block_to_delegate, FormType.ADJOINT)
+                    block.lhs = DelegatedBlockForm(
+                        self._ad_block_to_delegate, FormType.LHS)
+                    block.rhs = DelegatedBlockForm(
+                        self._ad_block_to_delegate, FormType.RHS)
+
+                # Adjoint solver.
+                self.delegated_block_ref = None
                 if not self._ad_problem._constant_jacobian:
-                    # Adjoint solver.
-                    if not self._ad_problem._constant_jacobian:
-                        with stop_annotating():
-                            if not self._ad_adj_solver:
-                                problem = self._ad_adj_lvs_problem(block)
-                                self._ad_adj_solver = LinearVariationalSolver(
-                                    problem, solver_parameters=self.parameters)
-
-                    block._ad_adj_solver = self._ad_adj_solver
+                    with stop_annotating():
+                        if not self._ad_adj_solver:
+                            problem = self._ad_adj_lvs_problem(block)
+                            self._ad_adj_solver = LinearVariationalSolver(
+                                problem, solver_parameters=self.parameters)
+                            block._ad_adj_solver = self._ad_adj_solver
+                        else:
+                            block._ad_adj_solver = DelegatedBlockSolver(
+                                self._ad_block_to_delegate, SolverType.ADJOINT)
 
                 tape.add_block(block)
 
@@ -181,3 +195,4 @@ def _build_count_map(self, problem, dependencies):
                     J_replace_map[coeff] = coeff.copy()
                 _ad_count_map[J_replace_map[coeff]] = coeff.count()
         return _ad_count_map, F_replace_map, J_replace_map
+

From 9c9e3763cdec7ba7928ec6cafac59abe7659aa91 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 27 Aug 2024 08:11:59 +0100
Subject: [PATCH 27/46] minor changes

---
 firedrake/adjoint_utils/blocks/__init__.py    |   3 +-
 firedrake/adjoint_utils/blocks/solving.py     | 186 ++++--------------
 firedrake/adjoint_utils/function.py           |  21 +-
 firedrake/adjoint_utils/variational_solver.py |  25 +--
 4 files changed, 56 insertions(+), 179 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py
index 354c342b19..bf83b896cc 100644
--- a/firedrake/adjoint_utils/blocks/__init__.py
+++ b/firedrake/adjoint_utils/blocks/__init__.py
@@ -1,8 +1,7 @@
 from .assembly import AssembleBlock  # NOQA F401
 from .solving import GenericSolveBlock, SolveLinearSystemBlock, \
     ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock, \
-    NonlinearVariationalSolveBlock, SolverType, DelegatedBlockSolver, \
-    DelegatedBlockForm, FormType  # NOQA F401
+    NonlinearVariationalSolveBlock  # NOQA F401
 from .function import FunctionAssignBlock, FunctionMergeBlock, \
     SubfunctionBlock  # NOQA F401
 from .dirichlet_bc import DirichletBCBlock  # NOQA F401
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 6a98962114..862234d8d2 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -1,11 +1,9 @@
 import numpy
-from enum import Enum
-from abc import ABC, abstractmethod
 import ufl
 from ufl import replace
 from ufl.formatting.ufl2unicode import ufl2unicode
 
-from pyadjoint import Block, stop_annotating, get_working_tape
+from pyadjoint import Block, stop_annotating
 from pyadjoint.enlisting import Enlist
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
@@ -87,33 +85,18 @@ def _init_solver_parameters(self, args, kwargs):
         self.assemble_kwargs = {}
 
     def __str__(self):
-        lhs = self.lhs
-        rhs = self.rhs
-        if isinstance(lhs, DelegatedBlockForm):
-            lhs = lhs.restore_form()
-        if isinstance(rhs, DelegatedBlockForm):
-            rhs = rhs.restore_form()
-        return "solve({} = {})".format(ufl2unicode(lhs),
-                                       ufl2unicode(rhs))
+        return "solve({} = {})".format(ufl2unicode(self.lhs),
+                                       ufl2unicode(self.rhs))
 
     def _create_F_form(self):
         # Process the equation forms, replacing values with checkpoints,
         # and gathering lhs and rhs in one single form.
         if self.linear:
             tmp_u = firedrake.Function(self.function_space)
-            lhs = self.lhs
-            if isinstance(lhs, DelegatedBlockForm):
-                lhs = lhs.restore_form()
-            rhs = self.rhs
-            if isinstance(rhs, DelegatedBlockForm):
-                rhs = self.rhs.restore_form()
-            F_form = firedrake.action(lhs, tmp_u) - rhs
+            F_form = firedrake.action(self.lhs, tmp_u) - self.rhs
         else:
             tmp_u = self.func
-            lhs = self.lhs
-            if isinstance(lhs, DelegatedBlockForm):
-                lhs = lhs.restore_form()
-            F_form = lhs
+            F_form = self.lhs
 
         replace_map = self._replace_map(F_form)
         replace_map[tmp_u] = self.get_outputs()[0].saved_output
@@ -526,16 +509,10 @@ def _replace_recompute_form(self):
         func = self._create_initial_guess()
 
         bcs = self._recover_bcs()
-        lhs = self.lhs
-        if isinstance(lhs, DelegatedBlockForm):
-            lhs = lhs.restore_form()
-        lhs = self._replace_form(lhs, func=func)
+        lhs = self._replace_form(self.lhs, func=func)
         rhs = 0
         if self.linear:
-            rhs = self.rhs
-            if isinstance(rhs, DelegatedBlockForm):
-                rhs = rhs.restore_form()
-            rhs = self._replace_form(rhs)
+            rhs = self._replace_form(self.rhs)
 
         return lhs, rhs, func, bcs
 
@@ -639,13 +616,14 @@ def __init__(self, equation, func, bcs, adj_F, adj_cache, problem_J,
         self.adj_F = adj_F
         self._adj_cache = adj_cache
         self._dFdm_cache = adj_cache.setdefault("dFdm_cache", {})
+        self.problem_J = problem_J
         self.solver_params = solver_params.copy()
         self.solver_kwargs = solver_kwargs
 
         super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs})
 
-        if problem_J is not None:
-            for coeff in problem_J.coefficients():
+        if self.problem_J is not None:
+            for coeff in self.problem_J.coefficients():
                 self.add_dependency(coeff, no_duplicates=True)
 
     def _init_solver_parameters(self, args, kwargs):
@@ -654,13 +632,9 @@ def _init_solver_parameters(self, args, kwargs):
 
     def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         self._ad_nlvs_replace_forms()
-        nlvs = self._ad_nlvs
-        if isinstance(nlvs, DelegatedBlockSolver):
-            nlvs = self._ad_nlvs.restore_solver()
-        nlvs.parameters.update(self.solver_params)
-        nlvs.solve()
-        func.assign(nlvs._problem.u)
-        del nlvs
+        self._ad_nlvs.parameters.update(self.solver_params)
+        self._ad_nlvs.solve()
+        func.assign(self._ad_nlvs._problem.u)
         return func
 
     def _adjoint_solve(self, dJdu, compute_bdy):
@@ -671,53 +645,47 @@ def _adjoint_solve(self, dJdu, compute_bdy):
             bc.apply(dJdu)
 
         # Update the left hand side coefficients of the adjoint equation.
-        adj_solver = self._ad_adj_solver
-        if isinstance(adj_solver, DelegatedBlockSolver):
-            adj_solver = adj_solver.restore_solver()
-        problem = adj_solver._problem
+        problem = self._ad_adj_solver._problem
         for block_variable in self.get_dependencies():
             # The self.adj_F coefficients hold the forward output
             # references.
-            adj_F = self.adj_F
-            if isinstance(adj_F, DelegatedBlockForm):
-                adj_F = adj_F.restore_form()
-            if block_variable.output in adj_F.coefficients():
-                index = adj_F.coefficients().index(
-                    block_variable.output)
+            if block_variable.output in self.adj_F.coefficients():
+                index = self.adj_F.coefficients().index(block_variable.output)
                 if isinstance(
                         block_variable.output, (
                             firedrake.Function, firedrake.Constant,
                             firedrake.Cofunction)):
                     # `problem.J` (Jacobian operator) is a deep copy of
-                    # `self.adj_F`. Hence, the indices of `self.adj_F`
-                    # is serving as a map for updating the coefficients of
-                    # the adjoint solver.
+                    # `self.adj_F`.
+                    # The indices of `self.adj_F` serve as a map for
+                    # updating the coefficients of the adjoint solver.
                     problem.J.coefficients()[index].assign(
                         block_variable.saved_output)
-        # Update the right hand side of the adjoint equation.
-        # problem.F._component[1] is the right hand side of the adjoint.
-        problem.F._components[1].assign(dJdu)
+
         bv = self.get_outputs()[0]
-        if bv.output in adj_F.coefficients():
-            index = adj_F.coefficients().index(bv.output)
+        if bv.output in self.adj_F.coefficients():
+            index = self.adj_F.coefficients().index(bv.output)
             problem.J.coefficients()[index].assign(
                 bv.checkpoint)
+
+        # Update the right hand side of the adjoint equation.
+        # problem.F._component[1] is the right hand side of the adjoint.
+        problem.F._components[1].assign(dJdu)
+
         # Solve the adjoint equation.
-        adj_solver.solve()
+        self._ad_adj_solver.solve()
 
         adj_sol = firedrake.Function(self.function_space)
-        adj_sol.assign(adj_solver._problem.u)
+        adj_sol.assign(self._ad_adj_solver._problem.u)
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self._compute_adj_bdy(
-                adj_sol, adj_sol_bdy, adj_solver._problem.J, dJdu_copy)
+                adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.J,
+                dJdu_copy)
         return adj_sol, adj_sol_bdy
 
     def _ad_assign_map(self, form):
-        nlvs = self._ad_nlvs
-        if isinstance(nlvs, DelegatedBlockSolver):
-            nlvs = nlvs.restore_solver()
-        count_map = nlvs._problem._ad_count_map
+        count_map = self._ad_nlvs._problem._ad_count_map
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -738,10 +706,7 @@ def _ad_assign_coefficients(self, form):
             coeff.assign(value)
 
     def _ad_nlvs_replace_forms(self):
-        nlvs = self._ad_nlvs
-        if isinstance(nlvs, DelegatedBlockSolver):
-            nlvs = nlvs.restore_solver()
-        problem = nlvs._problem
+        problem = self._ad_nlvs._problem
         self._ad_assign_coefficients(problem.F)
         self._ad_assign_coefficients(problem.J)
 
@@ -750,10 +715,7 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
             dFdu = self._adj_cache["dFdu_adj"]
         else:
             dFdu = super()._assemble_dFdu_adj(dFdu_adj_form, **kwargs)
-            nlvs = self._ad_nlvs
-            if isinstance(nlvs, DelegatedBlockSolver):
-                nlvs = nlvs.restore_solver(SolverType.FORWARD)
-            if nlvs._constant_jacobian:
+            if self._ad_nlvs._problem._constant_jacobian:
                 self._adj_cache["dFdu_adj"] = dFdu
         return dFdu
 
@@ -764,14 +726,9 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
-        nlvs = self._ad_nlvs
-        if isinstance(nlvs, DelegatedBlockSolver):
-            nlvs = nlvs.restore_solver()
-        if nlvs._problem._constant_jacobian:
-            adj_F = self.adj_F
-            if isinstance(adj_F, DelegatedBlockForm):
-                adj_F = adj_F.restore_form()
-            dFdu_form = adj_F
+
+        if self._ad_nlvs._problem._constant_jacobian:
+            dFdu_form = self.adj_F
             # Replace the form coefficients with checkpointed values.
             replace_map = self._replace_map(dFdu_form)
             replace_map[self.func] = self.get_outputs()[0].saved_output
@@ -834,10 +791,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if c in self._dFdm_cache:
             dFdm = self._dFdm_cache[c]
         else:
-            lhs = self.lhs
-            if isinstance(lhs, DelegatedBlockForm):
-                lhs = lhs.restore_form()
-            dFdm = -firedrake.derivative(lhs, c, trial_function)
+            dFdm = -firedrake.derivative(self.lhs, c, trial_function)
             dFdm = firedrake.adjoint(dFdm)
             self._dFdm_cache[c] = dFdm
 
@@ -1008,67 +962,3 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
     def __str__(self):
         target_string = f"〈{str(self.target_space.ufl_element().shortstr())}〉"
         return f"project({self.get_dependencies()[0]}, {target_string}))"
-
-
-class SolverType(Enum):
-    FORWARD = 1
-    ADJOINT = 2
-
-
-class FormType(Enum):
-    ADJOINT = 1
-    LHS = 2
-    RHS = 3
-
-
-class Solver(dict):
-    def __init__(self, forward, adjoint):
-        super().__init__()
-        self.solvers = {SolverType.FORWARD: forward, SolverType.ADJOINT: adjoint}
-
-    def __setitem__(self, key, value):
-        if key in self.solvers:
-            self.solvers[key] = value
-        super().__setitem__(key, value)
-
-    def __getitem__(self, key):
-        if key in self.solvers:
-            return self.solvers[key]
-        return super().__getitem__(key)
-
-    def finalize(self):
-        f = weakref.finalize(self, self.solvers)
-        f.atexit = False
-
-
-class DelegatedBlockSolver:
-    def __init__(self, delegated_block, solver_type):
-        # Block number that this solver is delegated to.
-        self.delegated_block = delegated_block
-        self.solver_type = solver_type
-
-    def restore_solver(self):
-        block = get_working_tape()._blocks[self.delegated_block]
-        if self.solver_type == SolverType.FORWARD:
-            return block._ad_nlvs
-        elif self.solver_type == SolverType.ADJOINT:
-            return block._ad_adj_solver
-
-
-class DelegatedBlockForm:
-    def __init__(self, delegated_block, form_type):
-        self.delegated_block = delegated_block
-        self.form_type = form_type
-
-    def restore_form(self):
-        if self.form_type == FormType.ADJOINT:
-            return get_working_tape()._blocks[self.delegated_block].adj_F
-        elif self.form_type == FormType.LHS:
-            return get_working_tape()._blocks[self.delegated_block].lhs
-        elif self.form_type == FormType.RHS:
-            return get_working_tape()._blocks[self.delegated_block].rhs
-        else:
-            raise NotImplementedError(
-                "Only adjoint form is supported for now. If you need to "
-                "restore another form, please create a new form type."
-            )
\ No newline at end of file
diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py
index 4bdc0666a6..4cd7e40269 100644
--- a/firedrake/adjoint_utils/function.py
+++ b/firedrake/adjoint_utils/function.py
@@ -219,15 +219,6 @@ def _ad_create_checkpoint(self):
             return CheckpointFunction(self)
         else:
             return self.copy(deepcopy=True)
-        
-    def _ad_clear_checkpoint(self, checkpoint):
-        # DelegatedFunctionCheckpoint does not hold data. Actually,
-        # it is a reference to another checkpoint function. Thus, clear
-        # can be unsafe once we will lose the reference to the original
-        # checkpoint.
-        if not isinstance(checkpoint, DelegatedFunctionCheckpoint):
-            checkpoint = None
-        return checkpoint
 
     def _ad_convert_riesz(self, value, options=None):
         from firedrake import Function, Cofunction
@@ -407,5 +398,17 @@ def _applyBinary(self, f, y):
             npdata[i] = f(npdata[i], npdatay[i])
         vec.set_local(npdata)
 
+    def _ad_from_petsc(self, vec):
+        with self.dat.vec_wo as self_v:
+            vec.copy(result=self_v)
+
+    def _ad_to_petsc(self, vec=None):
+        with self.dat.vec_ro as self_v:
+            if vec:
+                self_v.copy(result=vec)
+            else:
+                vec = self_v.copy()
+        return vec
+
     def __deepcopy__(self, memodict={}):
         return self.copy(deepcopy=True)
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index d0efd46003..4ff1ffeb58 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -1,8 +1,7 @@
 import copy
 from functools import wraps
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
-from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock, \
-    DelegatedBlockSolver, DelegatedBlockForm, FormType, SolverType
+from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
 from ufl import replace
 
 
@@ -49,7 +48,6 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_nlvs = None
             self._ad_adj_cache = {}
             self._ad_adj_solver = None
-            self._ad_block_to_delegate = None
 
         return wrapper
 
@@ -85,30 +83,18 @@ def wrapper(self, **kwargs):
                         self._ad_problem_clone(self._ad_problem, block.get_dependencies()),
                         **self._ad_kwargs
                     )
-                    block._ad_nlvs = self._ad_nlvs
-                    self._ad_block_to_delegate = len(tape._blocks)
-                else:
-                    block._ad_nlvs = DelegatedBlockSolver(
-                        self._ad_block_to_delegate, SolverType.FORWARD)
-                    block.adj_F = DelegatedBlockForm(
-                        self._ad_block_to_delegate, FormType.ADJOINT)
-                    block.lhs = DelegatedBlockForm(
-                        self._ad_block_to_delegate, FormType.LHS)
-                    block.rhs = DelegatedBlockForm(
-                        self._ad_block_to_delegate, FormType.RHS)
+
+                block._ad_nlvs = self._ad_nlvs
 
                 # Adjoint solver.
-                self.delegated_block_ref = None
                 if not self._ad_problem._constant_jacobian:
                     with stop_annotating():
                         if not self._ad_adj_solver:
                             problem = self._ad_adj_lvs_problem(block)
                             self._ad_adj_solver = LinearVariationalSolver(
                                 problem, solver_parameters=self.parameters)
-                            block._ad_adj_solver = self._ad_adj_solver
-                        else:
-                            block._ad_adj_solver = DelegatedBlockSolver(
-                                self._ad_block_to_delegate, SolverType.ADJOINT)
+
+                    block._ad_adj_solver = self._ad_adj_solver
 
                 tape.add_block(block)
 
@@ -195,4 +181,3 @@ def _build_count_map(self, problem, dependencies):
                     J_replace_map[coeff] = coeff.copy()
                 _ad_count_map[J_replace_map[coeff]] = coeff.count()
         return _ad_count_map, F_replace_map, J_replace_map
-

From 1357a17e919d954310d9d2c277337d5adf4e1213 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 10 Sep 2024 13:31:24 +0100
Subject: [PATCH 28/46] remove weak ref; minor changes; LinearSolver if
 Jacobian is constant

---
 firedrake/adjoint_utils/blocks/solving.py     | 28 ++++++-------------
 firedrake/adjoint_utils/variational_solver.py | 12 +++++---
 2 files changed, 16 insertions(+), 24 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 862234d8d2..42a95985f7 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -1,6 +1,7 @@
 import numpy
 import ufl
 from ufl import replace
+from numbers import Number
 from ufl.formatting.ufl2unicode import ufl2unicode
 
 from pyadjoint import Block, stop_annotating
@@ -8,7 +9,6 @@
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
 from .block_utils import isconstant
-import weakref
 
 
 def extract_subfunction(u, V):
@@ -49,7 +49,7 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         # Equation RHS
         self.rhs = rhs
         # Solution function
-        self._func = weakref.ref(func)
+        self.func = func
         self.function_space = self.func.function_space()
         # Boundary conditions
         self.bcs = []
@@ -73,10 +73,6 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         self.add_dependency(mesh)
         self._init_solver_parameters(args, kwargs)
 
-    @property
-    def func(self):
-        return self._func()
-
     def _init_solver_parameters(self, args, kwargs):
         self.forward_args = kwargs.pop("forward_args", [])
         self.forward_kwargs = kwargs.pop("forward_kwargs", {})
@@ -723,23 +719,14 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         dJdu = adj_inputs[0]
         dJdu = dJdu.copy()
 
+        F_form = self._create_F_form()
+
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
 
-        if self._ad_nlvs._problem._constant_jacobian:
-            dFdu_form = self.adj_F
-            # Replace the form coefficients with checkpointed values.
-            replace_map = self._replace_map(dFdu_form)
-            replace_map[self.func] = self.get_outputs()[0].saved_output
-            dFdu_form = replace(dFdu_form, replace_map)
-
-            adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
-                dFdu_form, dJdu, compute_bdy
-            )
-        else:
-            # Adjoint solver using linear variational solver.
-            adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
+        # Adjoint solver using linear variational solver.
+        adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
 
         self.adj_sol = adj_sol
         if self.adj_cb is not None:
@@ -748,6 +735,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
+        r["form"] = F_form
         r["adj_sol"] = adj_sol
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
@@ -757,7 +745,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if not self.linear and self.func == block_variable.output:
             # We are not able to calculate derivatives wrt initial guess.
             return None
-        F_form = self._create_F_form()
+        F_form = prepared["form"]
         adj_sol = prepared["adj_sol"]
         adj_sol_bdy = prepared["adj_sol_bdy"]
         c = block_variable.output
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 4ff1ffeb58..bc4bbcc85f 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver
+            from firedrake import LinearVariationalSolver, assemble, LinearSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -87,10 +87,14 @@ def wrapper(self, **kwargs):
                 block._ad_nlvs = self._ad_nlvs
 
                 # Adjoint solver.
-                if not self._ad_problem._constant_jacobian:
+                if not self._ad_adj_solver:
                     with stop_annotating():
-                        if not self._ad_adj_solver:
-                            problem = self._ad_adj_lvs_problem(block)
+                        problem = self._ad_adj_lvs_problem(block)
+                        if self._ad_problem._constant_jacobian:
+                            self._ad_adj_solver = LinearSolver(
+                                assemble(problem.J),
+                                solver_parameters=self.parameters)
+                        else:
                             self._ad_adj_solver = LinearVariationalSolver(
                                 problem, solver_parameters=self.parameters)
 

From d8345eeb1e34f9b9272023da1d691aeb7a369ad9 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Wed, 11 Sep 2024 07:36:29 +0100
Subject: [PATCH 29/46] lvs when not self._ad_problem._constant_jacobian.

---
 firedrake/adjoint_utils/blocks/solving.py     | 22 +++++++++++++------
 firedrake/adjoint_utils/variational_solver.py | 16 +++++---------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 42a95985f7..e0afed82c1 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -1,7 +1,6 @@
 import numpy
 import ufl
 from ufl import replace
-from numbers import Number
 from ufl.formatting.ufl2unicode import ufl2unicode
 
 from pyadjoint import Block, stop_annotating
@@ -9,6 +8,7 @@
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
 from .block_utils import isconstant
+import weakref
 
 
 def extract_subfunction(u, V):
@@ -719,14 +719,23 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         dJdu = adj_inputs[0]
         dJdu = dJdu.copy()
 
-        F_form = self._create_F_form()
-
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
 
-        # Adjoint solver using linear variational solver.
-        adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
+        if self._ad_nlvs._problem._constant_jacobian:
+            dFdu_form = self.adj_F
+            # Replace the form coefficients with checkpointed values.
+            replace_map = self._replace_map(dFdu_form)
+            replace_map[self.func] = self.get_outputs()[0].saved_output
+            dFdu_form = replace(dFdu_form, replace_map)
+
+            adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
+                dFdu_form, dJdu, compute_bdy
+            )
+        else:
+            # Adjoint solver using linear variational solver.
+            adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
 
         self.adj_sol = adj_sol
         if self.adj_cb is not None:
@@ -735,7 +744,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
-        r["form"] = F_form
         r["adj_sol"] = adj_sol
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
@@ -745,7 +753,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if not self.linear and self.func == block_variable.output:
             # We are not able to calculate derivatives wrt initial guess.
             return None
-        F_form = prepared["form"]
+        F_form = self._create_F_form()
         adj_sol = prepared["adj_sol"]
         adj_sol_bdy = prepared["adj_sol_bdy"]
         c = block_variable.output
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index bc4bbcc85f..cee482bff2 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver, assemble, LinearSolver
+            from firedrake import LinearVariationalSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -87,16 +87,12 @@ def wrapper(self, **kwargs):
                 block._ad_nlvs = self._ad_nlvs
 
                 # Adjoint solver.
-                if not self._ad_adj_solver:
-                    with stop_annotating():
-                        problem = self._ad_adj_lvs_problem(block)
-                        if self._ad_problem._constant_jacobian:
-                            self._ad_adj_solver = LinearSolver(
-                                assemble(problem.J),
-                                solver_parameters=self.parameters)
-                        else:
+                if not self._ad_problem._constant_jacobian:
+                    if not self._ad_adj_solver:
+                        with stop_annotating():
                             self._ad_adj_solver = LinearVariationalSolver(
-                                problem, solver_parameters=self.parameters)
+                                self._ad_adj_lvs_problem(block),
+                                solver_parameters=self.parameters)
 
                     block._ad_adj_solver = self._ad_adj_solver
 

From f0270e163b81a20ee5aba557925e02b4179605d7 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Wed, 11 Sep 2024 07:38:43 +0100
Subject: [PATCH 30/46] flake8

---
 firedrake/adjoint_utils/blocks/solving.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index e0afed82c1..47313f9615 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -8,7 +8,6 @@
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
 from .block_utils import isconstant
-import weakref
 
 
 def extract_subfunction(u, V):

From 58bec52ab84639c41105812acfd63047ed4f3db2 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sat, 14 Sep 2024 06:20:12 +0100
Subject: [PATCH 31/46] code for constant jacobian

---
 .../adjoint_utils/blocks/dirichlet_bc.py      |   3 +-
 firedrake/adjoint_utils/blocks/solving.py     | 107 ++++++++++--------
 firedrake/adjoint_utils/variational_solver.py |  24 ++--
 3 files changed, 76 insertions(+), 58 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/dirichlet_bc.py b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
index b06d367da1..7a4784812c 100644
--- a/firedrake/adjoint_utils/blocks/dirichlet_bc.py
+++ b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
@@ -60,8 +60,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                     adj_value = firedrake.Function(self.collapsed_space, vec.dat)
 
                 if adj_value.ufl_shape == () or adj_value.ufl_shape[0] <= 1:
-                    R = firedrake.FunctionSpace(self.parent_space.mesh(), "R", 0)
-                    r = firedrake.Function(R.dual(), val=adj_value.dat.data_ro.sum())
+                    r = adj_value.dat.data_ro.sum()
                 else:
                     output = []
                     subindices = _extract_subindices(self.function_space)
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 47313f9615..0d62a5060d 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -639,44 +639,41 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         for bc in bcs:
             bc.apply(dJdu)
 
-        # Update the left hand side coefficients of the adjoint equation.
-        problem = self._ad_adj_solver._problem
-        for block_variable in self.get_dependencies():
-            # The self.adj_F coefficients hold the forward output
-            # references.
-            if block_variable.output in self.adj_F.coefficients():
-                index = self.adj_F.coefficients().index(block_variable.output)
-                if isinstance(
-                        block_variable.output, (
-                            firedrake.Function, firedrake.Constant,
-                            firedrake.Cofunction)):
-                    # `problem.J` (Jacobian operator) is a deep copy of
-                    # `self.adj_F`.
-                    # The indices of `self.adj_F` serve as a map for
-                    # updating the coefficients of the adjoint solver.
-                    problem.J.coefficients()[index].assign(
-                        block_variable.saved_output)
-
-        bv = self.get_outputs()[0]
-        if bv.output in self.adj_F.coefficients():
-            index = self.adj_F.coefficients().index(bv.output)
-            problem.J.coefficients()[index].assign(
-                bv.checkpoint)
-
-        # Update the right hand side of the adjoint equation.
-        # problem.F._component[1] is the right hand side of the adjoint.
-        problem.F._components[1].assign(dJdu)
-
-        # Solve the adjoint equation.
-        self._ad_adj_solver.solve()
+        if self._ad_nlvs._problem._constant_jacobian:
+            if self._adj_cache["recomputed_block"] == self._recomputed_block - 1:
+                self._adj_cache["recomputed_block"] = self._recomputed_block
+                self._adj_cache["dFdu_adj"] = firedrake.assemble(
+                    self._ad_update_jacobian_adj(self._adj_cache["dFdu_adj_form"]))
+
+            # Update the solver matrix, which the assembled Jacobian.
+            # self._ad_adj_solver.A = self._adj_cache["dFdu_adj"]
+            # self._ad_adj_solver.solve(adj_sol, dJdu)
+            adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
+                self._adj_cache["dFdu_adj_form"], dJdu, compute_bdy)
+            # print(self._ad_adj_solver.A.M.values.max(), self._recomputed_block, adj_sol.dat.data_ro.max())
+            if compute_bdy:
+                jac_adj = self._adj_cache["dFdu_adj"]
+        else:
+            # Update the left hand side coefficients of the adjoint equation.
+            problem = self._ad_adj_solver._problem
+            self._ad_update_jacobian_adj(problem.J)
+            # Update the right hand side of the adjoint equation.
+            # problem.F._component[1] is the right hand side of the adjoint.
+            problem.F._components[1].assign(dJdu)
+
+            # Solve the adjoint equation.
+            self._ad_adj_solver.solve()
+            # M = firedrake.assemble(problem.J)
+            adj_sol = firedrake.Function(self.function_space)
+            adj_sol.assign(self._ad_adj_solver._problem.u)
+            # print(M.M.values.max(), self._recomputed_block, adj_sol.dat.data_ro.max())
+            if compute_bdy:
+                jac_adj = self._ad_adj_solver._problem.J
 
-        adj_sol = firedrake.Function(self.function_space)
-        adj_sol.assign(self._ad_adj_solver._problem.u)
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self._compute_adj_bdy(
-                adj_sol, adj_sol_bdy, self._ad_adj_solver._problem.J,
-                dJdu_copy)
+                adj_sol, adj_sol_bdy, jac_adj, dJdu_copy)
         return adj_sol, adj_sol_bdy
 
     def _ad_assign_map(self, form):
@@ -700,6 +697,28 @@ def _ad_assign_coefficients(self, form):
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
+    def _ad_update_jacobian_adj(self, jac_adj):
+        for block_variable in self.get_dependencies():
+            # The self.adj_F coefficients hold the forward output
+            # references.
+            if block_variable.output in self.adj_F.coefficients():
+                index = self.adj_F.coefficients().index(block_variable.output)
+                if isinstance(
+                        block_variable.output, (
+                            firedrake.Function, firedrake.Constant,
+                            firedrake.Cofunction)):
+                    # jac_adj is a deep copy of `self.adj_F`. Thus, the indices
+                    # of `self.adj_F` serve as a map for updating the
+                    # coefficients of the adjoint solver.
+                    jac_adj.coefficients()[index].assign(
+                        block_variable.saved_output)
+
+        bv = self.get_outputs()[0]
+        if bv.output in self.adj_F.coefficients():
+            index = self.adj_F.coefficients().index(bv.output)
+            jac_adj.coefficients()[index].assign(bv.checkpoint)
+        return jac_adj
+
     def _ad_nlvs_replace_forms(self):
         problem = self._ad_nlvs._problem
         self._ad_assign_coefficients(problem.F)
@@ -716,25 +735,16 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         dJdu = adj_inputs[0]
+
+        F_form = self._create_F_form()
+
         dJdu = dJdu.copy()
 
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
 
-        if self._ad_nlvs._problem._constant_jacobian:
-            dFdu_form = self.adj_F
-            # Replace the form coefficients with checkpointed values.
-            replace_map = self._replace_map(dFdu_form)
-            replace_map[self.func] = self.get_outputs()[0].saved_output
-            dFdu_form = replace(dFdu_form, replace_map)
-
-            adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
-                dFdu_form, dJdu, compute_bdy
-            )
-        else:
-            # Adjoint solver using linear variational solver.
-            adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
+        adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
 
         self.adj_sol = adj_sol
         if self.adj_cb is not None:
@@ -743,6 +753,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
+        r["form"] = F_form
         r["adj_sol"] = adj_sol
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
@@ -752,7 +763,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if not self.linear and self.func == block_variable.output:
             # We are not able to calculate derivatives wrt initial guess.
             return None
-        F_form = self._create_F_form()
+        F_form = prepared["form"]
         adj_sol = prepared["adj_sol"]
         adj_sol_bdy = prepared["adj_sol_bdy"]
         c = block_variable.output
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index cee482bff2..3362323f68 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver
+            from firedrake import LinearVariationalSolver, LinearSolver, assemble
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -87,14 +87,22 @@ def wrapper(self, **kwargs):
                 block._ad_nlvs = self._ad_nlvs
 
                 # Adjoint solver.
-                if not self._ad_problem._constant_jacobian:
-                    if not self._ad_adj_solver:
-                        with stop_annotating():
-                            self._ad_adj_solver = LinearVariationalSolver(
-                                self._ad_adj_lvs_problem(block),
+                # How many times the block has been recomputed.
+                self._ad_adj_cache.update({"recomputed_block": 0})
+                if not self._ad_adj_solver:
+                    with stop_annotating():
+                        problem = self._ad_adj_lvs_problem(block)
+                        if self._ad_problem._constant_jacobian:
+                            assembled_J = assemble(problem.J)
+                            self._ad_adj_cache.update({"dFdu_adj": assembled_J})
+                            self._ad_adj_cache.update({"dFdu_adj_form": problem.J})
+                            self._ad_adj_solver = LinearSolver(
+                                self._ad_adj_cache["dFdu_adj"],
                                 solver_parameters=self.parameters)
-
-                    block._ad_adj_solver = self._ad_adj_solver
+                        else:
+                            self._ad_adj_solver = LinearVariationalSolver(
+                                problem, solver_parameters=self.parameters)
+                block._ad_adj_solver = self._ad_adj_solver
 
                 tape.add_block(block)
 

From 6bb1d148095d75082b14bc5875ce2392548d1abc Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sat, 14 Sep 2024 06:20:51 +0100
Subject: [PATCH 32/46] Adapting test

---
 tests/regression/test_adjoint_operators.py | 72 +++++++---------------
 1 file changed, 22 insertions(+), 50 deletions(-)

diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index d4f6184484..380bc48205 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -952,35 +952,37 @@ def test_riesz_representation_for_adjoints():
 @pytest.mark.skipcomplex
 @pytest.mark.parametrize("constant_jacobian", [False, True])
 def test_lvs_constant_jacobian(constant_jacobian):
-    mesh = UnitIntervalMesh(10)
+    mesh = IntervalMesh(10, 0, 2)
     X = SpatialCoordinate(mesh)
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
     trial = TrialFunction(space)
-
     u = Function(space, name="u").interpolate(X[0] - 0.5)
     with stop_annotating():
         u_ref = u.copy(deepcopy=True)
-    v = Function(space, name="v")
-    problem = LinearVariationalProblem(
-        inner(trial, test) * dx, inner(u, test) * dx, v,
-        constant_jacobian=constant_jacobian)
-    solver = LinearVariationalSolver(problem)
-    solver.solve()
-    J = assemble(v * v * dx)
-
-    assert "dFdu_adj" not in solver._ad_adj_cache
-
-    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
-
-    cached_dFdu_adj = solver._ad_adj_cache.get("dFdu_adj", None)
-    assert (cached_dFdu_adj is None) == (not constant_jacobian)
+    def J(u):
+        v = Function(space, name="v")
+        problem = LinearVariationalProblem(
+            inner(trial, test) * dx, inner(u, test) * dx, v,
+            constant_jacobian=constant_jacobian)
+        solver = LinearVariationalSolver(problem)
+        solver.solve()
+        return assemble(v * v * dx), solver
+   
+    J_val, solver = J(u)
+    assert ("dFdu_adj" in solver._ad_adj_cache) == constant_jacobian
+    J_hat = ReducedFunctional(J_val, Control(u))
+    dJ = J_hat.derivative(options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
+    with stop_annotating():
+        u0 = Function(space).interpolate(X[0] - 0.6)
+        u_ref0 = u0.copy(deepcopy=True)
+        J_val0, _ = J(u0)
+    
+    assert np.allclose(J_hat(u0), J_val0)
+    dJ = J_hat.derivative(options={"riesz_representation": "l2"})
 
-    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
-
-    assert cached_dFdu_adj is solver._ad_adj_cache.get("dFdu_adj", None)
-    assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
+    assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref0, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex  # Taping for complex-valued 0-forms not yet done
@@ -1003,33 +1005,3 @@ def test_cofunction_assign_functional():
     assert np.allclose(float(Jhat.derivative()), 1.0)
     f.assign(2.0)
     assert np.allclose(Jhat(f), 2.0)
-
-
-@pytest.mark.skipcomplex
-@pytest.mark.parametrize("constant_jacobian", [False, True])
-def test_adjoint_solver_compute_bdy(constant_jacobian):
-    # Testing the case where is required to compute the adjoint
-    # boundary condition.
-    mesh = UnitIntervalMesh(10)
-    space = FunctionSpace(mesh, "Lagrange", 1)
-    test = TestFunction(space)
-    trial = TrialFunction(space)
-    sol = Function(space, name="sol")
-    # Dirichlet boundary conditions
-    R = FunctionSpace(mesh, "R", 0)
-    a = Function(R, val=1.0)
-    b = Function(R, val=2.0)
-    bc_left = DirichletBC(space, a, 1)
-    bc_right = DirichletBC(space, b, 2)
-    bc = [bc_left, bc_right]
-    F = dot(grad(trial), grad(test)) * dx
-    problem = LinearVariationalProblem(lhs(F), rhs(F), sol, bcs=bc,
-                                       constant_jacobian=constant_jacobian)
-    solver = LinearVariationalSolver(problem)
-    solver.solve()
-    J = assemble(sol * sol * dx)
-    J_hat = ReducedFunctional(J, [Control(a), Control(b)])
-
-    assert taylor_test(
-        J_hat, [a, b], [Function(R, val=rand(1)), Function(R, val=rand(1))]
-    ) > 1.9

From 1e67f5a87e6fdedf0f737a303d5c5fd6f58ef47b Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sat, 14 Sep 2024 06:49:16 +0100
Subject: [PATCH 33/46] flake8

---
 tests/regression/test_adjoint_operators.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 380bc48205..0b2e0865d0 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -960,6 +960,7 @@ def test_lvs_constant_jacobian(constant_jacobian):
     u = Function(space, name="u").interpolate(X[0] - 0.5)
     with stop_annotating():
         u_ref = u.copy(deepcopy=True)
+
     def J(u):
         v = Function(space, name="v")
         problem = LinearVariationalProblem(
@@ -968,7 +969,7 @@ def J(u):
         solver = LinearVariationalSolver(problem)
         solver.solve()
         return assemble(v * v * dx), solver
-   
+
     J_val, solver = J(u)
     assert ("dFdu_adj" in solver._ad_adj_cache) == constant_jacobian
     J_hat = ReducedFunctional(J_val, Control(u))
@@ -978,7 +979,7 @@ def J(u):
         u0 = Function(space).interpolate(X[0] - 0.6)
         u_ref0 = u0.copy(deepcopy=True)
         J_val0, _ = J(u0)
-    
+
     assert np.allclose(J_hat(u0), J_val0)
     dJ = J_hat.derivative(options={"riesz_representation": "l2"})
 

From 3f4306f0dec5a2747ce85cb0cb2bea566cf6b419 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 7 Oct 2024 20:04:41 -0300
Subject: [PATCH 34/46] solvers dictionary

---
 firedrake/adjoint_utils/blocks/solving.py     | 116 +++++++++---------
 firedrake/adjoint_utils/variational_solver.py |  39 +++---
 2 files changed, 72 insertions(+), 83 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 8529aabe73..9d1eb8df67 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -3,7 +3,7 @@
 from ufl import replace
 from ufl.formatting.ufl2unicode import ufl2unicode
 
-from pyadjoint import Block, stop_annotating
+from pyadjoint import Block, stop_annotating, get_working_tape
 from pyadjoint.enlisting import Enlist
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
@@ -617,6 +617,8 @@ def __init__(self, equation, func, bcs, adj_F, adj_cache, problem_J,
         self.problem_J = problem_J
         self.solver_params = solver_params.copy()
         self.solver_kwargs = solver_kwargs
+        # The number of times the block has been recomputed.
+        self._recomputed_block = 0
 
         super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs})
 
@@ -628,11 +630,22 @@ def _init_solver_parameters(self, args, kwargs):
         super()._init_solver_parameters(args, kwargs)
         solve_init_params(self, args, kwargs, varform=True)
 
+    def recompute_component(self, inputs, block_variable, idx, prepared):
+        tape = get_working_tape()
+        if self._ad_solvers["recompute"] == tape.recompute - 1:
+            # Update how many times the block has been recomputed.
+            self._ad_solvers["recompute"] = tape.recompute
+            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+                self._ad_solvers["forward_nlvs"].invalidate_jacobian()
+                self._ad_solvers["adjoint_lvs"].invalidate_jacobian()
+
+        return super().recompute_component(inputs, block_variable, idx, prepared)
+
     def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
         self._ad_nlvs_replace_forms()
-        self._ad_nlvs.parameters.update(self.solver_params)
-        self._ad_nlvs.solve()
-        func.assign(self._ad_nlvs._problem.u)
+        self._ad_solvers["forward_nlvs"].parameters.update(self.solver_params)
+        self._ad_solvers["forward_nlvs"].solve()
+        func.assign(self._ad_solvers["forward_nlvs"]._problem.u)
         return func
 
     def _adjoint_solve(self, dJdu, compute_bdy):
@@ -642,36 +655,43 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         for bc in bcs:
             bc.apply(dJdu)
 
-        if self._ad_nlvs._problem._constant_jacobian:
-            if self._adj_cache["recomputed_block"] == self._recomputed_block - 1:
-                self._adj_cache["recomputed_block"] = self._recomputed_block
-                self._adj_cache["dFdu_adj"] = firedrake.assemble(
-                    self._ad_update_jacobian_adj(self._adj_cache["dFdu_adj_form"]))
-
-            # Update the solver matrix, which the assembled Jacobian.
-            # self._ad_adj_solver.A = self._adj_cache["dFdu_adj"]
-            # self._ad_adj_solver.solve(adj_sol, dJdu)
-            adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
-                self._adj_cache["dFdu_adj_form"], dJdu, compute_bdy)
-            # print(self._ad_adj_solver.A.M.values.max(), self._recomputed_block, adj_sol.dat.data_ro.max())
-            if compute_bdy:
-                jac_adj = self._adj_cache["dFdu_adj"]
-        else:
-            # Update the left hand side coefficients of the adjoint equation.
-            problem = self._ad_adj_solver._problem
-            self._ad_update_jacobian_adj(problem.J)
-            # Update the right hand side of the adjoint equation.
-            # problem.F._component[1] is the right hand side of the adjoint.
-            problem.F._components[1].assign(dJdu)
-
-            # Solve the adjoint equation.
-            self._ad_adj_solver.solve()
-            # M = firedrake.assemble(problem.J)
-            adj_sol = firedrake.Function(self.function_space)
-            adj_sol.assign(self._ad_adj_solver._problem.u)
-            # print(M.M.values.max(), self._recomputed_block, adj_sol.dat.data_ro.max())
-            if compute_bdy:
-                jac_adj = self._ad_adj_solver._problem.J
+        # Update the left hand side coefficients of the adjoint equation.
+        problem = self._ad_solvers["adjoint_lvs"]._problem
+        if (
+            (not problem._constant_jacobian) or
+            (problem._constant_jacobian and
+             not self._ad_solvers["adjoint_lvs"]._ctx._jacobian_assembled)
+        ):
+            for block_variable in self.get_dependencies():
+                # The self.adj_F coefficients hold the forward output
+                # references.
+                if block_variable.output in self.adj_F.coefficients():
+                    index = self.adj_F.coefficients().index(block_variable.output)
+                    if isinstance(
+                            block_variable.output, (
+                                firedrake.Function, firedrake.Constant,
+                                firedrake.Cofunction)):
+                        # jproblem.J is a deep copy of `self.adj_F`. Thus, the
+                        # indices of `self.adj_F` serve as a map for updating
+                        # the coefficients of the adjoint solver.
+                        problem.J.coefficients()[index].assign(
+                            block_variable.saved_output)
+
+            bv = self.get_outputs()[0]
+            if bv.output in self.adj_F.coefficients():
+                index = self.adj_F.coefficients().index(bv.output)
+                problem.J.coefficients()[index].assign(bv.checkpoint)
+
+        # Update the right hand side of the adjoint equation.
+        # problem.F._component[1] is the right hand side of the adjoint.
+        problem.F._components[1].assign(dJdu)
+
+        # Solve the adjoint equation.
+        self._ad_solvers["adjoint_lvs"].solve()
+        adj_sol = firedrake.Function(self.function_space)
+        adj_sol.assign(self._ad_solvers["adjoint_lvs"]._problem.u)
+        if compute_bdy:
+            jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
 
         adj_sol_bdy = None
         if compute_bdy:
@@ -680,7 +700,7 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         return adj_sol, adj_sol_bdy
 
     def _ad_assign_map(self, form):
-        count_map = self._ad_nlvs._problem._ad_count_map
+        count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -700,30 +720,8 @@ def _ad_assign_coefficients(self, form):
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
-    def _ad_update_jacobian_adj(self, jac_adj):
-        for block_variable in self.get_dependencies():
-            # The self.adj_F coefficients hold the forward output
-            # references.
-            if block_variable.output in self.adj_F.coefficients():
-                index = self.adj_F.coefficients().index(block_variable.output)
-                if isinstance(
-                        block_variable.output, (
-                            firedrake.Function, firedrake.Constant,
-                            firedrake.Cofunction)):
-                    # jac_adj is a deep copy of `self.adj_F`. Thus, the indices
-                    # of `self.adj_F` serve as a map for updating the
-                    # coefficients of the adjoint solver.
-                    jac_adj.coefficients()[index].assign(
-                        block_variable.saved_output)
-
-        bv = self.get_outputs()[0]
-        if bv.output in self.adj_F.coefficients():
-            index = self.adj_F.coefficients().index(bv.output)
-            jac_adj.coefficients()[index].assign(bv.checkpoint)
-        return jac_adj
-
     def _ad_nlvs_replace_forms(self):
-        problem = self._ad_nlvs._problem
+        problem = self._ad_solvers["forward_nlvs"]._problem
         self._ad_assign_coefficients(problem.F)
         self._ad_assign_coefficients(problem.J)
 
@@ -732,7 +730,7 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
             dFdu = self._adj_cache["dFdu_adj"]
         else:
             dFdu = super()._assemble_dFdu_adj(dFdu_adj_form, **kwargs)
-            if self._ad_nlvs._problem._constant_jacobian:
+            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
                 self._adj_cache["dFdu_adj"] = dFdu
         return dFdu
 
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 3362323f68..90d8fc5f96 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -45,9 +45,9 @@ def wrapper(self, problem, *args, **kwargs):
             self._ad_problem = problem
             self._ad_args = args
             self._ad_kwargs = kwargs
-            self._ad_nlvs = None
+            self._ad_solvers = {"forward_nlvs": None, "adjoint_lvs": None,
+                                "recompute": 0}
             self._ad_adj_cache = {}
-            self._ad_adj_solver = None
 
         return wrapper
 
@@ -77,32 +77,23 @@ def wrapper(self, **kwargs):
                                                        solver_kwargs=self._ad_kwargs,
                                                        ad_block_tag=self.ad_block_tag,
                                                        **sb_kwargs)
+
                 # Forward variational solver.
-                if not self._ad_nlvs:
-                    self._ad_nlvs = type(self)(
+                if not self._ad_solvers["forward_nlvs"]:
+                    self._ad_solvers["forward_nlvs"] = type(self)(
                         self._ad_problem_clone(self._ad_problem, block.get_dependencies()),
                         **self._ad_kwargs
                     )
 
-                block._ad_nlvs = self._ad_nlvs
-
                 # Adjoint solver.
-                # How many times the block has been recomputed.
-                self._ad_adj_cache.update({"recomputed_block": 0})
-                if not self._ad_adj_solver:
+                if not self._ad_solvers["adjoint_lvs"]:
                     with stop_annotating():
-                        problem = self._ad_adj_lvs_problem(block)
-                        if self._ad_problem._constant_jacobian:
-                            assembled_J = assemble(problem.J)
-                            self._ad_adj_cache.update({"dFdu_adj": assembled_J})
-                            self._ad_adj_cache.update({"dFdu_adj_form": problem.J})
-                            self._ad_adj_solver = LinearSolver(
-                                self._ad_adj_cache["dFdu_adj"],
-                                solver_parameters=self.parameters)
-                        else:
-                            self._ad_adj_solver = LinearVariationalSolver(
-                                problem, solver_parameters=self.parameters)
-                block._ad_adj_solver = self._ad_adj_solver
+                        self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver(
+                            self._ad_adj_lvs_problem(
+                                block, self._ad_problem._constant_jacobian),
+                            solver_parameters=self.parameters)
+
+                block._ad_solvers = self._ad_solvers
 
                 tape.add_block(block)
 
@@ -137,7 +128,7 @@ def _ad_problem_clone(self, problem, dependencies):
         return nlvp
 
     @no_annotations
-    def _ad_adj_lvs_problem(self, block):
+    def _ad_adj_lvs_problem(self, block, constant_jacobian):
         """Create the adjoint variational problem."""
         from firedrake import Function, Cofunction, LinearVariationalProblem
         # Homogeneous boundary conditions for the adjoint problem
@@ -147,7 +138,7 @@ def _ad_adj_lvs_problem(self, block):
         right_hand_side = Cofunction(block.function_space.dual())
         tmp_problem = LinearVariationalProblem(
             block.adj_F, right_hand_side,
-            adj_sol, bcs=bcs)
+            adj_sol, bcs=bcs, constant_jacobian=constant_jacobian)
         # The `block.adj_F` coefficients hold the output references.
         # We do not want to modify the user-defined values. Hence, the adjoint
         # linear variational problem is created with a deep copy of the forward
@@ -156,7 +147,7 @@ def _ad_adj_lvs_problem(self, block):
             tmp_problem, block._dependencies)
         lvp = LinearVariationalProblem(
             replace(tmp_problem.J, J_replace_map), right_hand_side, adj_sol,
-            bcs=tmp_problem.bcs)
+            bcs=tmp_problem.bcs, constant_jacobian=constant_jacobian)
         lvp._ad_count_map_update(_ad_count_map)
         del tmp_problem
         return lvp

From 0294c9586129705e87ff883a97faaf3c1825bf7e Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 7 Oct 2024 20:06:01 -0300
Subject: [PATCH 35/46] Update test for the new solvers.

---
 tests/regression/test_adjoint_operators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 0b2e0865d0..7570746b46 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -840,7 +840,7 @@ def test_assign_cofunction(solve_type):
             solver.solve()
         J += assemble(((sol + Constant(1.0)) ** 2) * dx)
     rf = ReducedFunctional(J, Control(k))
-    assert rf(k) == J
+    assert np.isclose(rf(k), J, rtol=1e-10)
     assert taylor_test(rf, k, Function(V).assign(0.1)) > 1.9
 
 
@@ -971,7 +971,7 @@ def J(u):
         return assemble(v * v * dx), solver
 
     J_val, solver = J(u)
-    assert ("dFdu_adj" in solver._ad_adj_cache) == constant_jacobian
+    # assert ("dFdu_adj" in solver._ad_adj_cache) == constant_jacobian
     J_hat = ReducedFunctional(J_val, Control(u))
     dJ = J_hat.derivative(options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)

From f05e8b95a065c2c4b69fd08516df74e257600a1c Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 8 Oct 2024 16:20:28 -0300
Subject: [PATCH 36/46] Add adj_args and adj_kwargs; add a property in nlvs;
 remove _assemble_dFdu_adj

---
 firedrake/adjoint_utils/blocks/solving.py     | 11 +---------
 firedrake/adjoint_utils/variational_solver.py | 21 ++++++++++---------
 firedrake/variational_solver.py               |  5 +++++
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 9d1eb8df67..8f53334a39 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -660,7 +660,7 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         if (
             (not problem._constant_jacobian) or
             (problem._constant_jacobian and
-             not self._ad_solvers["adjoint_lvs"]._ctx._jacobian_assembled)
+             not self._ad_solvers["adjoint_lvs"].jacobian_assembled)
         ):
             for block_variable in self.get_dependencies():
                 # The self.adj_F coefficients hold the forward output
@@ -725,15 +725,6 @@ def _ad_nlvs_replace_forms(self):
         self._ad_assign_coefficients(problem.F)
         self._ad_assign_coefficients(problem.J)
 
-    def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
-        if "dFdu_adj" in self._adj_cache:
-            dFdu = self._adj_cache["dFdu_adj"]
-        else:
-            dFdu = super()._assemble_dFdu_adj(dFdu_adj_form, **kwargs)
-            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
-                self._adj_cache["dFdu_adj"] = dFdu
-        return dFdu
-
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         dJdu = adj_inputs[0]
 
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 90d8fc5f96..6a4ff5d989 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver, LinearSolver, assemble
+            from firedrake import LinearVariationalSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -85,13 +85,12 @@ def wrapper(self, **kwargs):
                         **self._ad_kwargs
                     )
 
-                # Adjoint solver.
+                # Adjoint variational solver.
                 if not self._ad_solvers["adjoint_lvs"]:
                     with stop_annotating():
                         self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver(
-                            self._ad_adj_lvs_problem(
-                                block, self._ad_problem._constant_jacobian),
-                            solver_parameters=self.parameters)
+                            self._ad_adj_lvs_problem(block),
+                            *block.adj_args, **block.adj_kwargs)
 
                 block._ad_solvers = self._ad_solvers
 
@@ -128,7 +127,7 @@ def _ad_problem_clone(self, problem, dependencies):
         return nlvp
 
     @no_annotations
-    def _ad_adj_lvs_problem(self, block, constant_jacobian):
+    def _ad_adj_lvs_problem(self, block):
         """Create the adjoint variational problem."""
         from firedrake import Function, Cofunction, LinearVariationalProblem
         # Homogeneous boundary conditions for the adjoint problem
@@ -138,16 +137,18 @@ def _ad_adj_lvs_problem(self, block, constant_jacobian):
         right_hand_side = Cofunction(block.function_space.dual())
         tmp_problem = LinearVariationalProblem(
             block.adj_F, right_hand_side,
-            adj_sol, bcs=bcs, constant_jacobian=constant_jacobian)
+            adj_sol, bcs=bcs,
+            constant_jacobian=self._ad_problem._constant_jacobian)
         # The `block.adj_F` coefficients hold the output references.
         # We do not want to modify the user-defined values. Hence, the adjoint
-        # linear variational problem is created with a deep copy of the forward
-        # outputs.
+        # 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(
             tmp_problem, block._dependencies)
         lvp = LinearVariationalProblem(
             replace(tmp_problem.J, J_replace_map), right_hand_side, adj_sol,
-            bcs=tmp_problem.bcs, constant_jacobian=constant_jacobian)
+            bcs=tmp_problem.bcs,
+            constant_jacobian=self._ad_problem._constant_jacobian)
         lvp._ad_count_map_update(_ad_count_map)
         del tmp_problem
         return lvp
diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py
index 609d599800..0e411d6d11 100644
--- a/firedrake/variational_solver.py
+++ b/firedrake/variational_solver.py
@@ -332,6 +332,11 @@ def solve(self, bounds=None):
         comm = self._problem.u_restrict.function_space().mesh()._comm
         PETSc.garbage_cleanup(comm=comm)
 
+    @property
+    def jacobian_assembled(self):
+        r"""Whether the Jacobian has been assembled."""
+        return self._ctx._jacobian_assembled
+
 
 class LinearVariationalProblem(NonlinearVariationalProblem):
     r"""Linear variational problem a(u, v) = L(v)."""

From 92f7f81c96b7b840fe62da581af30c1400c04c81 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 8 Oct 2024 16:30:20 -0300
Subject: [PATCH 37/46] flake8; minimal tests modifications

---
 firedrake/adjoint_utils/blocks/solving.py   |  8 ++---
 tests/regression/test_adjoint_operators.py  | 34 +++++++++------------
 tests/regression/test_poisson_strong_bcs.py |  4 +++
 3 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 8f53334a39..88113298f1 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -617,8 +617,6 @@ def __init__(self, equation, func, bcs, adj_F, adj_cache, problem_J,
         self.problem_J = problem_J
         self.solver_params = solver_params.copy()
         self.solver_kwargs = solver_kwargs
-        # The number of times the block has been recomputed.
-        self._recomputed_block = 0
 
         super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs})
 
@@ -658,9 +656,9 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         # Update the left hand side coefficients of the adjoint equation.
         problem = self._ad_solvers["adjoint_lvs"]._problem
         if (
-            (not problem._constant_jacobian) or
-            (problem._constant_jacobian and
-             not self._ad_solvers["adjoint_lvs"].jacobian_assembled)
+            (not problem._constant_jacobian)
+            or (problem._constant_jacobian
+                and not self._ad_solvers["adjoint_lvs"].jacobian_assembled)
         ):
             for block_variable in self.get_dependencies():
                 # The self.adj_F coefficients hold the forward output
diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 7570746b46..0ae0d3ea80 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -952,38 +952,32 @@ def test_riesz_representation_for_adjoints():
 @pytest.mark.skipcomplex
 @pytest.mark.parametrize("constant_jacobian", [False, True])
 def test_lvs_constant_jacobian(constant_jacobian):
-    mesh = IntervalMesh(10, 0, 2)
+    mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
     trial = TrialFunction(space)
+
     u = Function(space, name="u").interpolate(X[0] - 0.5)
     with stop_annotating():
         u_ref = u.copy(deepcopy=True)
+    v = Function(space, name="v")
+    problem = LinearVariationalProblem(
+        inner(trial, test) * dx, inner(u, test) * dx, v,
+        constant_jacobian=constant_jacobian)
+    solver = LinearVariationalSolver(problem)
+    solver.solve()
+    J = assemble(v * v * dx)
 
-    def J(u):
-        v = Function(space, name="v")
-        problem = LinearVariationalProblem(
-            inner(trial, test) * dx, inner(u, test) * dx, v,
-            constant_jacobian=constant_jacobian)
-        solver = LinearVariationalSolver(problem)
-        solver.solve()
-        return assemble(v * v * dx), solver
+    assert "dFdu_adj" not in solver._ad_adj_cache
+
+    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
 
-    J_val, solver = J(u)
-    # assert ("dFdu_adj" in solver._ad_adj_cache) == constant_jacobian
-    J_hat = ReducedFunctional(J_val, Control(u))
-    dJ = J_hat.derivative(options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
-    with stop_annotating():
-        u0 = Function(space).interpolate(X[0] - 0.6)
-        u_ref0 = u0.copy(deepcopy=True)
-        J_val0, _ = J(u0)
 
-    assert np.allclose(J_hat(u0), J_val0)
-    dJ = J_hat.derivative(options={"riesz_representation": "l2"})
+    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
 
-    assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref0, test) * dx).dat.data_ro)
+    assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex  # Taping for complex-valued 0-forms not yet done
diff --git a/tests/regression/test_poisson_strong_bcs.py b/tests/regression/test_poisson_strong_bcs.py
index 2655f98b91..23d3fe634a 100644
--- a/tests/regression/test_poisson_strong_bcs.py
+++ b/tests/regression/test_poisson_strong_bcs.py
@@ -90,3 +90,7 @@ def test_poisson_analytic_linear_parallel():
     error = run_test_linear(1, 1)
     print('[%d]' % MPI.COMM_WORLD.rank, 'error:', error)
     assert error < 5e-6
+
+
+if __name__ == '__main__':
+    test_poisson_analytic_linear_parallel()
\ No newline at end of file

From 2f8c84b68674171b4eb0d5f9a141849a83f66b58 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 8 Oct 2024 16:32:42 -0300
Subject: [PATCH 38/46] flake8

---
 tests/regression/test_poisson_strong_bcs.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tests/regression/test_poisson_strong_bcs.py b/tests/regression/test_poisson_strong_bcs.py
index 23d3fe634a..2655f98b91 100644
--- a/tests/regression/test_poisson_strong_bcs.py
+++ b/tests/regression/test_poisson_strong_bcs.py
@@ -90,7 +90,3 @@ def test_poisson_analytic_linear_parallel():
     error = run_test_linear(1, 1)
     print('[%d]' % MPI.COMM_WORLD.rank, 'error:', error)
     assert error < 5e-6
-
-
-if __name__ == '__main__':
-    test_poisson_analytic_linear_parallel()
\ No newline at end of file

From de687bd4f69d3838a2ae42c397783c86c0318464 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Tue, 8 Oct 2024 21:26:40 -0300
Subject: [PATCH 39/46] minor changes

---
 firedrake/adjoint_utils/blocks/solving.py | 25 +++++++++--------------
 1 file changed, 10 insertions(+), 15 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 88113298f1..bab05a673d 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -686,16 +686,15 @@ def _adjoint_solve(self, dJdu, compute_bdy):
 
         # Solve the adjoint equation.
         self._ad_solvers["adjoint_lvs"].solve()
-        adj_sol = firedrake.Function(self.function_space)
-        adj_sol.assign(self._ad_solvers["adjoint_lvs"]._problem.u)
         if compute_bdy:
             jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
 
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self._compute_adj_bdy(
-                adj_sol, adj_sol_bdy, jac_adj, dJdu_copy)
-        return adj_sol, adj_sol_bdy
+                self._ad_solvers["adjoint_lvs"]._problem.u, adj_sol_bdy,
+                jac_adj, dJdu_copy)
+        return self._ad_solvers["adjoint_lvs"]._problem.u, adj_sol_bdy
 
     def _ad_assign_map(self, form):
         count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
@@ -724,26 +723,22 @@ def _ad_nlvs_replace_forms(self):
         self._ad_assign_coefficients(problem.J)
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
-        dJdu = adj_inputs[0]
-
-        F_form = self._create_F_form()
-
-        dJdu = dJdu.copy()
-
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
-
-        adj_sol, adj_sol_bdy = self._adjoint_solve(dJdu, compute_bdy)
-        self.adj_state = adj_sol
+        adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy)
+        if not self.adj_state:
+            self.adj_state = adj_sol
+        else:
+            self.adj_state.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:
             self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
-        r["form"] = F_form
-        r["adj_sol"] = adj_sol
+        r["form"] = self._create_F_form()
+        r["adj_sol"] = self.adj_state
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
 

From 23c82b62c3099db93bd8100d415a2ff894d4a1f1 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Wed, 9 Oct 2024 09:37:57 -0300
Subject: [PATCH 40/46] minor changes

---
 firedrake/adjoint_utils/blocks/solving.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index bab05a673d..4f5a444100 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -728,9 +728,8 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         )
         adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy)
         if not self.adj_state:
-            self.adj_state = adj_sol
-        else:
-            self.adj_state.assign(adj_sol)
+            self.adj_state = firedrake.Function(adj_sol.function_space())
+        self.adj_state.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:
@@ -738,7 +737,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
 
         r = {}
         r["form"] = self._create_F_form()
-        r["adj_sol"] = self.adj_state
+        r["adj_sol"] = adj_sol
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
 

From fdcb2b31b01914ff5b49b7c46223f02efbe27a77 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Sun, 13 Oct 2024 23:48:50 -0400
Subject: [PATCH 41/46] LinearSolver for constant jacobian

---
 firedrake/adjoint_utils/blocks/solving.py     | 104 ++++++++++--------
 firedrake/adjoint_utils/variational_solver.py |  47 ++++----
 2 files changed, 87 insertions(+), 64 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 4f5a444100..fba392f9f4 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -635,12 +635,11 @@ def recompute_component(self, inputs, block_variable, idx, prepared):
             self._ad_solvers["recompute"] = tape.recompute
             if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
                 self._ad_solvers["forward_nlvs"].invalidate_jacobian()
-                self._ad_solvers["adjoint_lvs"].invalidate_jacobian()
-
+                self._ad_solvers["update_adjoint"] = True
         return super().recompute_component(inputs, block_variable, idx, prepared)
 
     def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
-        self._ad_nlvs_replace_forms()
+        self._ad_solver_replace_forms(forward=True, adjoint=False)
         self._ad_solvers["forward_nlvs"].parameters.update(self.solver_params)
         self._ad_solvers["forward_nlvs"].solve()
         func.assign(self._ad_solvers["forward_nlvs"]._problem.u)
@@ -653,51 +652,50 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         for bc in bcs:
             bc.apply(dJdu)
 
-        # Update the left hand side coefficients of the adjoint equation.
-        problem = self._ad_solvers["adjoint_lvs"]._problem
         if (
-            (not problem._constant_jacobian)
-            or (problem._constant_jacobian
-                and not self._ad_solvers["adjoint_lvs"].jacobian_assembled)
+            self._ad_solvers["forward_nlvs"]._problem._constant_jacobian
+            and self._ad_solvers["update_adjoint"]
         ):
-            for block_variable in self.get_dependencies():
-                # The self.adj_F coefficients hold the forward output
-                # references.
-                if block_variable.output in self.adj_F.coefficients():
-                    index = self.adj_F.coefficients().index(block_variable.output)
-                    if isinstance(
-                            block_variable.output, (
-                                firedrake.Function, firedrake.Constant,
-                                firedrake.Cofunction)):
-                        # jproblem.J is a deep copy of `self.adj_F`. Thus, the
-                        # indices of `self.adj_F` serve as a map for updating
-                        # the coefficients of the adjoint solver.
-                        problem.J.coefficients()[index].assign(
-                            block_variable.saved_output)
-
-            bv = self.get_outputs()[0]
-            if bv.output in self.adj_F.coefficients():
-                index = self.adj_F.coefficients().index(bv.output)
-                problem.J.coefficients()[index].assign(bv.checkpoint)
-
-        # Update the right hand side of the adjoint equation.
-        # problem.F._component[1] is the right hand side of the adjoint.
-        problem.F._components[1].assign(dJdu)
-
-        # Solve the adjoint equation.
-        self._ad_solvers["adjoint_lvs"].solve()
+            self._ad_solver_replace_forms(forward=False, adjoint=True)
+            self._ad_solvers["adjoint_lvs"] = type(self._ad_solvers["adjoint_lvs"])(
+                self._assemble_dFdu_adj(
+                    self._ad_solvers["adjoint_lvs"].A.form, **self.assemble_kwargs.copy()),
+                *self.adj_args, **self.adj_kwargs)
+            self._ad_solvers["update_adjoint"] = False
+        elif not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+            # Update left hand side of the adjoint equation.
+            self._ad_solver_replace_forms(forward=False, adjoint=True)
+
+            # Update the right hand side of the adjoint equation.
+            # problem.F._component[1] is the right hand side of the adjoint.
+            self._ad_solvers["adjoint_lvs"]._problem.F._components[1].assign(dJdu)
+
+            # Solve the adjoint linear variational solver.
+            self._ad_solvers["adjoint_lvs"].solve()
+            u_sol = self._ad_solvers["adjoint_lvs"]._problem.u
+
+        if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+            u_sol = firedrake.Function(
+                self._ad_solvers["forward_nlvs"]._problem.u.function_space())
+            self._ad_solvers["adjoint_lvs"].solve(u_sol, dJdu)
+
         if compute_bdy:
             jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
 
         adj_sol_bdy = None
         if compute_bdy:
             adj_sol_bdy = self._compute_adj_bdy(
-                self._ad_solvers["adjoint_lvs"]._problem.u, adj_sol_bdy,
-                jac_adj, dJdu_copy)
-        return self._ad_solvers["adjoint_lvs"]._problem.u, adj_sol_bdy
+                u_sol, adj_sol_bdy, jac_adj, dJdu_copy)
+        return u_sol, adj_sol_bdy
 
-    def _ad_assign_map(self, form):
-        count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
+    def _ad_assign_map(self, form, forward, adjoint):
+        if forward:
+            count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
+        elif adjoint:
+            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+                count_map = self._ad_solvers["adj_ad_count_map"]
+            else:
+                count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -710,17 +708,33 @@ def _ad_assign_map(self, form):
                 if coeff_count in form_ad_count_map:
                     assign_map[form_ad_count_map[coeff_count]] = \
                         block_variable.saved_output
+                    
+        if adjoint and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+            block_variable = self.get_outputs()[0]
+            coeff_count = block_variable.output.count()
+            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):
-        assign_map = self._ad_assign_map(form)
+    def _ad_assign_coefficients(self, form, forward, adjoint):
+        assign_map = self._ad_assign_map(
+            form, forward, adjoint)
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
-    def _ad_nlvs_replace_forms(self):
-        problem = self._ad_solvers["forward_nlvs"]._problem
-        self._ad_assign_coefficients(problem.F)
-        self._ad_assign_coefficients(problem.J)
+    def _ad_solver_replace_forms(self, forward, adjoint):
+        if forward:
+            problem = self._ad_solvers["forward_nlvs"]._problem
+            self._ad_assign_coefficients(problem.F, forward, adjoint)
+            self._ad_assign_coefficients(problem.J, forward, adjoint)
+        else:
+            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+                form = self._ad_solvers["adjoint_lvs"].A.form
+            else:
+                form = self._ad_solvers["adjoint_lvs"]._problem.J
+
+            self._ad_assign_coefficients(form, forward, adjoint)
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 6a4ff5d989..db98b40f90 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver
+            from firedrake import LinearVariationalSolver, LinearSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -88,9 +88,16 @@ def wrapper(self, **kwargs):
                 # Adjoint variational solver.
                 if not self._ad_solvers["adjoint_lvs"]:
                     with stop_annotating():
-                        self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver(
-                            self._ad_adj_lvs_problem(block),
-                            *block.adj_args, **block.adj_kwargs)
+                        problem = self._ad_adj_lvs_problem(block)
+                        if self._ad_problem._constant_jacobian:
+                            self._ad_solvers["adjoint_lvs"] = LinearSolver(
+                                block._assemble_dFdu_adj(problem.J, **block.assemble_kwargs.copy()),
+                                *block.adj_args, **block.adj_kwargs)
+                            self._ad_solvers["adj_ad_count_map"] = problem._ad_count_map
+                            self._ad_solvers["update_adjoint"] = False
+                        else:
+                            self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver(
+                                self._ad_adj_lvs_problem(block), *block.adj_args, **block.adj_kwargs)
 
                 block._ad_solvers = self._ad_solvers
 
@@ -115,8 +122,8 @@ def _ad_problem_clone(self, problem, dependencies):
         expressions, we'll instead create clones of them.
         """
         from firedrake import NonlinearVariationalProblem
-        _ad_count_map, F_replace_map, J_replace_map = self._build_count_map(
-            problem, dependencies)
+        _ad_count_map, J_replace_map, F_replace_map = self._build_count_map(
+            problem.J, dependencies, Form=problem.F)
         nlvp = NonlinearVariationalProblem(replace(problem.F, F_replace_map),
                                            F_replace_map[problem.u_restrict],
                                            bcs=problem.bcs,
@@ -143,8 +150,9 @@ def _ad_adj_lvs_problem(self, block):
         # We do not want to modify the user-defined values. Hence, the adjoint
         # 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(
-            tmp_problem, block._dependencies)
+        _ad_count_map, J_replace_map, _ = self._build_count_map(
+            block.adj_F, block._dependencies
+        )
         lvp = LinearVariationalProblem(
             replace(tmp_problem.J, J_replace_map), right_hand_side, adj_sol,
             bcs=tmp_problem.bcs,
@@ -153,24 +161,25 @@ def _ad_adj_lvs_problem(self, block):
         del tmp_problem
         return lvp
 
-    def _build_count_map(self, problem, dependencies):
+    def _build_count_map(self, J, dependencies, Form=None):
         from firedrake import Function
 
         F_replace_map = {}
         J_replace_map = {}
-
-        F_coefficients = problem.F.coefficients()
-        J_coefficients = problem.J.coefficients()
+        if Form:
+            F_coefficients = Form.coefficients()
+        J_coefficients = J.coefficients()
 
         _ad_count_map = {}
         for block_variable in dependencies:
             coeff = block_variable.output
-            if coeff in F_coefficients and coeff not in F_replace_map:
-                if isinstance(coeff, Function) and coeff.ufl_element().family() == "Real":
-                    F_replace_map[coeff] = copy.deepcopy(coeff)
-                else:
-                    F_replace_map[coeff] = coeff.copy(deepcopy=True)
-                _ad_count_map[F_replace_map[coeff]] = coeff.count()
+            if Form:
+                if coeff in F_coefficients and coeff not in F_replace_map:
+                    if isinstance(coeff, Function) and coeff.ufl_element().family() == "Real":
+                        F_replace_map[coeff] = copy.deepcopy(coeff)
+                    else:
+                        F_replace_map[coeff] = coeff.copy(deepcopy=True)
+                    _ad_count_map[F_replace_map[coeff]] = coeff.count()
 
             if coeff in J_coefficients and coeff not in J_replace_map:
                 if coeff in F_replace_map:
@@ -180,4 +189,4 @@ def _build_count_map(self, problem, dependencies):
                 else:
                     J_replace_map[coeff] = coeff.copy()
                 _ad_count_map[J_replace_map[coeff]] = coeff.count()
-        return _ad_count_map, F_replace_map, J_replace_map
+        return _ad_count_map, J_replace_map, F_replace_map

From 5fee9ee39ef4b8047f904b62adb24812e76856b8 Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 14 Oct 2024 00:22:47 -0400
Subject: [PATCH 42/46] minor changes

---
 firedrake/adjoint_utils/blocks/solving.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index fba392f9f4..4d64a4863e 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -679,11 +679,13 @@ def _adjoint_solve(self, dJdu, compute_bdy):
                 self._ad_solvers["forward_nlvs"]._problem.u.function_space())
             self._ad_solvers["adjoint_lvs"].solve(u_sol, dJdu)
 
-        if compute_bdy:
-            jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
-
         adj_sol_bdy = None
         if compute_bdy:
+            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+                jac_adj = self._ad_solvers["adjoint_lvs"].A.form
+            else:
+                jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
+
             adj_sol_bdy = self._compute_adj_bdy(
                 u_sol, adj_sol_bdy, jac_adj, dJdu_copy)
         return u_sol, adj_sol_bdy

From 01bbe5ce030cba07edcc695ec3ef421e5a67100d Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 14 Oct 2024 10:12:56 -0400
Subject: [PATCH 43/46] Keep only variational solver.

---
 firedrake/adjoint_utils/blocks/solving.py     | 85 +++++++++----------
 firedrake/adjoint_utils/variational_solver.py | 14 +--
 2 files changed, 42 insertions(+), 57 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 4d64a4863e..994f4fc253 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -2,6 +2,7 @@
 import ufl
 from ufl import replace
 from ufl.formatting.ufl2unicode import ufl2unicode
+from enum import Enum
 
 from pyadjoint import Block, stop_annotating, get_working_tape
 from pyadjoint.enlisting import Enlist
@@ -24,6 +25,12 @@ def extract_subfunction(u, V):
         return u
 
 
+class Solver(Enum):
+    """Enum for solver types."""
+    FORWARD = 0
+    ADJOINT = 1
+
+
 class GenericSolveBlock(Block):
     pop_kwargs_keys = ["adj_cb", "adj_bdy_cb", "adj2_cb", "adj2_bdy_cb",
                        "forward_args", "forward_kwargs", "adj_args",
@@ -630,16 +637,16 @@ def _init_solver_parameters(self, args, kwargs):
 
     def recompute_component(self, inputs, block_variable, idx, prepared):
         tape = get_working_tape()
-        if self._ad_solvers["recompute"] == tape.recompute - 1:
+        if self._ad_solvers["recompute_count"] == tape.recompute_count - 1:
             # Update how many times the block has been recomputed.
-            self._ad_solvers["recompute"] = tape.recompute
+            self._ad_solvers["recompute_count"] = tape.recompute_count
             if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
                 self._ad_solvers["forward_nlvs"].invalidate_jacobian()
                 self._ad_solvers["update_adjoint"] = True
         return super().recompute_component(inputs, block_variable, idx, prepared)
 
     def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
-        self._ad_solver_replace_forms(forward=True, adjoint=False)
+        self._ad_solver_replace_forms()
         self._ad_solvers["forward_nlvs"].parameters.update(self.solver_params)
         self._ad_solvers["forward_nlvs"].solve()
         func.assign(self._ad_solvers["forward_nlvs"]._problem.u)
@@ -656,48 +663,34 @@ def _adjoint_solve(self, dJdu, compute_bdy):
             self._ad_solvers["forward_nlvs"]._problem._constant_jacobian
             and self._ad_solvers["update_adjoint"]
         ):
-            self._ad_solver_replace_forms(forward=False, adjoint=True)
-            self._ad_solvers["adjoint_lvs"] = type(self._ad_solvers["adjoint_lvs"])(
-                self._assemble_dFdu_adj(
-                    self._ad_solvers["adjoint_lvs"].A.form, **self.assemble_kwargs.copy()),
-                *self.adj_args, **self.adj_kwargs)
+            # Update left hand side of the adjoint equation.
+            self._ad_solver_replace_forms(Solver.ADJOINT)
+            self._ad_solvers["adjoint_lvs"].invalidate_jacobian()
             self._ad_solvers["update_adjoint"] = False
         elif not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
             # Update left hand side of the adjoint equation.
-            self._ad_solver_replace_forms(forward=False, adjoint=True)
-
-            # Update the right hand side of the adjoint equation.
-            # problem.F._component[1] is the right hand side of the adjoint.
-            self._ad_solvers["adjoint_lvs"]._problem.F._components[1].assign(dJdu)
+            self._ad_solver_replace_forms(Solver.ADJOINT)
 
-            # Solve the adjoint linear variational solver.
-            self._ad_solvers["adjoint_lvs"].solve()
-            u_sol = self._ad_solvers["adjoint_lvs"]._problem.u
+        # Update the right hand side of the adjoint equation.
+        # problem.F._component[1] is the right hand side of the adjoint.
+        self._ad_solvers["adjoint_lvs"]._problem.F._components[1].assign(dJdu)
 
-        if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
-            u_sol = firedrake.Function(
-                self._ad_solvers["forward_nlvs"]._problem.u.function_space())
-            self._ad_solvers["adjoint_lvs"].solve(u_sol, dJdu)
+        # Solve the adjoint linear variational solver.
+        self._ad_solvers["adjoint_lvs"].solve()
+        u_sol = self._ad_solvers["adjoint_lvs"]._problem.u
 
         adj_sol_bdy = None
         if compute_bdy:
-            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
-                jac_adj = self._ad_solvers["adjoint_lvs"].A.form
-            else:
-                jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
-
+            jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
             adj_sol_bdy = self._compute_adj_bdy(
                 u_sol, adj_sol_bdy, jac_adj, dJdu_copy)
         return u_sol, adj_sol_bdy
 
-    def _ad_assign_map(self, form, forward, adjoint):
-        if forward:
+    def _ad_assign_map(self, form, solver):
+        if solver == Solver.FORWARD:
             count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
-        elif adjoint:
-            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
-                count_map = self._ad_solvers["adj_ad_count_map"]
-            else:
-                count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map
+        else:
+            count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -710,8 +703,11 @@ def _ad_assign_map(self, form, forward, adjoint):
                 if coeff_count in form_ad_count_map:
                     assign_map[form_ad_count_map[coeff_count]] = \
                         block_variable.saved_output
-                    
-        if adjoint and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+
+        if (
+            solver == Solver.ADJOINT
+            and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian
+        ):
             block_variable = self.get_outputs()[0]
             coeff_count = block_variable.output.count()
             if coeff_count in form_ad_count_map:
@@ -719,24 +715,19 @@ def _ad_assign_map(self, form, forward, adjoint):
                     block_variable.saved_output
         return assign_map
 
-    def _ad_assign_coefficients(self, form, forward, adjoint):
-        assign_map = self._ad_assign_map(
-            form, forward, adjoint)
+    def _ad_assign_coefficients(self, form, solver):
+        assign_map = self._ad_assign_map(form, solver)
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
-    def _ad_solver_replace_forms(self, forward, adjoint):
-        if forward:
+    def _ad_solver_replace_forms(self, solver=Solver.FORWARD):
+        if solver == Solver.FORWARD:
             problem = self._ad_solvers["forward_nlvs"]._problem
-            self._ad_assign_coefficients(problem.F, forward, adjoint)
-            self._ad_assign_coefficients(problem.J, forward, adjoint)
+            self._ad_assign_coefficients(problem.F, solver)
+            self._ad_assign_coefficients(problem.J, solver)
         else:
-            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
-                form = self._ad_solvers["adjoint_lvs"].A.form
-            else:
-                form = self._ad_solvers["adjoint_lvs"]._problem.J
-
-            self._ad_assign_coefficients(form, forward, adjoint)
+            self._ad_assign_coefficients(
+                self._ad_solvers["adjoint_lvs"]._problem.J, solver)
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         compute_bdy = self._should_compute_boundary_adjoint(
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index db98b40f90..a61555f1a9 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -46,7 +46,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": 0}
+                                "recompute_count": 0}
             self._ad_adj_cache = {}
 
         return wrapper
@@ -59,7 +59,7 @@ def wrapper(self, **kwargs):
             Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
             for the purposes of the adjoint computation (such as projecting fields to other function spaces
             for the purposes of visualisation)."""
-            from firedrake import LinearVariationalSolver, LinearSolver
+            from firedrake import LinearVariationalSolver
             annotate = annotate_tape(kwargs)
             if annotate:
                 tape = get_working_tape()
@@ -88,16 +88,10 @@ def wrapper(self, **kwargs):
                 # Adjoint variational solver.
                 if not self._ad_solvers["adjoint_lvs"]:
                     with stop_annotating():
-                        problem = self._ad_adj_lvs_problem(block)
+                        self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver(
+                            self._ad_adj_lvs_problem(block), *block.adj_args, **block.adj_kwargs)
                         if self._ad_problem._constant_jacobian:
-                            self._ad_solvers["adjoint_lvs"] = LinearSolver(
-                                block._assemble_dFdu_adj(problem.J, **block.assemble_kwargs.copy()),
-                                *block.adj_args, **block.adj_kwargs)
-                            self._ad_solvers["adj_ad_count_map"] = problem._ad_count_map
                             self._ad_solvers["update_adjoint"] = False
-                        else:
-                            self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver(
-                                self._ad_adj_lvs_problem(block), *block.adj_args, **block.adj_kwargs)
 
                 block._ad_solvers = self._ad_solvers
 

From 74de2acf9086568a3adc9feb7c98e0d7409ed4bb Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 14 Oct 2024 10:32:02 -0400
Subject: [PATCH 44/46] Update constant Jacobian test according new adjoint
 solver.

---
 tests/regression/test_adjoint_operators.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/regression/test_adjoint_operators.py b/tests/regression/test_adjoint_operators.py
index 0ae0d3ea80..962f2a0194 100644
--- a/tests/regression/test_adjoint_operators.py
+++ b/tests/regression/test_adjoint_operators.py
@@ -969,14 +969,15 @@ def test_lvs_constant_jacobian(constant_jacobian):
     solver.solve()
     J = assemble(v * v * dx)
 
-    assert "dFdu_adj" not in solver._ad_adj_cache
-
-    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
+    J_hat = ReducedFunctional(J, Control(u))
 
+    dJ = J_hat.derivative(options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
 
-    dJ = compute_gradient(J, Control(u), options={"riesz_representation": "l2"})
+    u_ref = Function(space, name="u").interpolate(X[0] - 0.1)
+    J_hat(u_ref)
 
+    dJ = J_hat.derivative(options={"riesz_representation": "l2"})
     assert np.allclose(dJ.dat.data_ro, 2 * assemble(inner(u_ref, test) * dx).dat.data_ro)
 
 

From 850e210450dfcdc6bf9479bf36bc96a1b133ce4b Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Mon, 14 Oct 2024 14:52:53 -0400
Subject: [PATCH 45/46] Minor chancge

---
 firedrake/variational_solver.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py
index 0e411d6d11..609d599800 100644
--- a/firedrake/variational_solver.py
+++ b/firedrake/variational_solver.py
@@ -332,11 +332,6 @@ def solve(self, bounds=None):
         comm = self._problem.u_restrict.function_space().mesh()._comm
         PETSc.garbage_cleanup(comm=comm)
 
-    @property
-    def jacobian_assembled(self):
-        r"""Whether the Jacobian has been assembled."""
-        return self._ctx._jacobian_assembled
-
 
 class LinearVariationalProblem(NonlinearVariationalProblem):
     r"""Linear variational problem a(u, v) = L(v)."""

From 1dc5892c24e3bfab3d8cc2ba037ff79f3f6aca0b Mon Sep 17 00:00:00 2001
From: Iglesia Dolci <daia.dolci@gmail.com>
Date: Thu, 24 Oct 2024 12:07:57 -0300
Subject: [PATCH 46/46] Test with right branch

---
 .github/workflows/build.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 675946f845..218ac9c379 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -83,6 +83,7 @@ jobs:
             --install defcon \
             --install gadopt \
             --install asQ \
+            --package-branch pyadjoint dolci/tape_recompute_count \
             || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |