Use CachedSolverBlock for project (-> #4638)

[angus-g]

Description

The ProjectBlock was essentially an extension of SolveVarFormBlock that sets up the mass matrix system for the adjoint machinery. For the original (forward) project() call, the optimised solver-free path is used. However, any of the adjoint operations including a forward recomputation would solve the system.

The approach here replicates that behaviour, but it's routed through the new CachedSolverBlock. As a consequence however, we have annotation on the Projector factory, which means that adjoint-friendly projection can occur through either the bare project() call, or by creating a Projector object.

The implementation is a bit ugly: because a lot of the caching machinery belongs to the NonlinearVariationalSolverMixin, we essentially dispatch to a dummy initialiser when creating a BasicProjector to ensure the solver caches are created correctly. A call to project() annotates as though a solve occurred, but uses the underlying projection implementation.

The SupermeshSolveBlock is actually a standalone object, so we just create that when SupermeshProjector.project() is called. Again, this allows for both bare and object-based projection to be taped.

The final case is the "null" project from one function to another on the same space. This creates an Assigner, and I assume it will "just work" through the annotation on assign().

Questions

• Does this break the API (also the solve() change) by removing the annotate kwarg? Now we have to wrap things in with stop_annotating() or similar
• There's probably a better way to manage ad_block_tag
• There's probably a better way to manage the solver parameters (indeed, if you just solve inner(u,w)*dx == inner(u+c,w)*dx for a Real c and try to evaluate the tape, there's a divergence... guess there aren't any adjoint tests for dead simple solves)

[angus-g]

Not really sure why the test_adjoint_cleanup case is failing, it's such a simple solve, but it's failing on the adjoint and that should be fine? Unless something about the form manipulation or block dependencies has changed?

from firedrake import *
from firedrake.adjoint import *

continue_annotation()
mesh = SquareMesh(1, 1, 1, quadrilateral=True)

V = FunctionSpace(mesh, "CG", 1)
R = FunctionSpace(mesh, "R", 0)

u_0 = Function(V).assign(1.0)
u = Function(V).assign(u_0)
r = Function(R)

r.assign(2.0)
u.project(r * u)
r.assign(1.0)
u.project(r * u)

J = assemble((u - Function(V).assign(1.0)) ** 2 * dx)

pause_annotation()
reduced_functional = ReducedFunctional(J, Control(u_0))

# Deliberate additional derivative computation.
reduced_functional.derivative()

[angus-g]

Okay, I've got a slightly more minimal reproducer that avoids the project path. This actually doesn't rely on the CachedSolveBlock and highlights a difference between the bare solve() annotation (through SolveVarFormBlock) and the NonlinearVariationalSolver annotation. Because ProjectBlock uses SolveVarFormBlock, tihs is the issue we're hitting.

from firedrake import *                                                                                                                                                 
from firedrake.adjoint import *                                                                                                                                         
                                                                                                                                                                        
continue_annotation()                                                                                                                                                   
                                                                                                                                                                        
mesh = UnitSquareMesh(1, 1)                                                                                                                                             
V = FunctionSpace(mesh, "CG", 1)                                                                                                                                        
R = FunctionSpace(mesh, "R", 0)                                                                                                                                         
u = Function(V)                                                                                                                                                         
w = TestFunction(V)                                                                                                                                                     
c = Function(R, val=1.)                                                                                                                                                 
                                                                                                                                                                        
a = inner(TrialFunction(V), w) * dx                                                                                                                                     
L = inner(c*u, w) * dx                                                                                                                                                  

if reproduce:
    solve(a == L, u)  # diverges for this PR, *not* main
else:
    # always diverges
    problem = LinearVariationalProblem(a, L, u)
    NonlinearVariationalSolver(problem).solve()

J = assemble(u**2 * dx)                                                                                                                                                 
rf = ReducedFunctional(J, Control(c))                                                                                                                                   
                                                                                                                                                                        
pause_annotation()                                                                                                                                                      
                                                                                                                                                                        
rf.derivative()

Here, the bilinear form is essentially the projection form that was causing the issues above. When we go through the NLVP, the bilinear form becomes linear during the static method NonlinearVariationalProblem.compute_bc_lifting, which sets F = action(F, u). Here, we get a form like:

w₃ * (conj((v_0))) * dx(<Mesh #1>[everywhere], {}, {})
  +  -1 * w₃ * w₆ * (conj((v_0))) * dx(<Mesh #1>[everywhere], {}, {})

Note that u = w₃ is present in both terms. I guess as an adjoint, this becomes degenerate? Anyway, the bare solve() annotation constructed a SolveVarFormBlock directly from the bilinear form: it extracts the LHS and RHS, and uses a different trial function:

class GenericSolveBlock(Block):
    ...
    
    def _create_F_form(self):
        if self.linear:
            tmp_u = Function(self.function_space)
            F_form = action(self.lhs, tmp_u) - self.rhs
            
        ...

This means the form is something like

w₅₇ * (conj((v_0))) * dx(<Mesh #1>[everywhere], {}, {})
  +  -1 * w₃₉ * w₄₁ * (conj((v_0))) * dx(<Mesh #1>[everywhere], {}, {})

and so the derivative and adjoint thereof only has one term. This is probably the "correct" behaviour, because even if we work around the degenerate solve in the test case:

diff --git a/tests/firedrake/adjoint/test_assignment.py 
b/tests/firedrake/adjoint/test_assignment.py
index 342355076..d1884fc38 100644
--- a/tests/firedrake/adjoint/test_assignment.py
+++ b/tests/firedrake/adjoint/test_assignment.py
@@ -254,7 +254,7 @@ def test_adjoint_cleanup(scheduler, rg):

     r.assign(2.0)
     u.project(r * u)
-    r.assign(1.0)
+    r.assign(1.5)
     u.project(r * u)

     J = assemble((u - Function(V).assign(1.0)) ** 2 * dx)

The Taylor test fails:

>       assert taylor_test(reduced_functional, u_0, dtemp) > 1.99999999
E       AssertionError: assert 1.0009069989745567 > 1.99999999
E        +  where 1.0009069989745567 = taylor_test(<pyadjoint.reduced_functional.ReducedFunctional object at 0x7ff8097d3c80>, Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7ff80be92900>, FiniteElement('Q', quadrilateral, 1), name=None), Mesh(VectorElement(FiniteElement('Q', quadrilateral, 1), dim=2), 1)), 3), Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7ff80be92900>, FiniteElement('Q', quadrilateral, 1), name=None), Mesh(VectorElement(FiniteElement('Q', quadrilateral, 1), dim=2), 1)), 166))