BUG: Multigrid doesn't work with a Cofunction right-hand-side

[JHopeCollins]

Describe the bug
A short description of what the bug is.

Multigrid doesn't appear to work with an assembled or Cofunction right hand side.

Steps to Reproduce
Steps to reproduce the behavior:

1. Code/test/example that was run '...' (as minimal as possible)
2. Command executed '...'

MFE:

from firedrake import *
from firedrake.petsc import PETSc

PETSc.Sys.popErrorHandler()

basemesh = UnitSquareMesh(4, 4)
hierarchy = MeshHierarchy(basemesh, refinement_levels=2)
mesh = hierarchy[-1]
x, y = SpatialCoordinate(mesh)

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

u = TrialFunction(V)
v = TestFunction(V)

A = inner(grad(u), grad(v))*dx
f = Function(V).interpolate(sin(x)*cos(y))

# L = f.riesz_representation()
# L = assemble(inner(f, v)*dx)
L = inner(f, v)*dx

bcs = DirichletBC(V, zero(), 1)

lu_params = { 
    'ksp_type': 'preonly',
    'pc_type': 'lu'
}

mg_params = { 
    'ksp_type': 'richardson',
    'ksp_rtol': 1e-5,
    'pc_type': 'mg',
    'mg': {
        'levels': {
            'pc_type': 'jacobi',
        },
        'coarse': lu_params
    }   
}

params = { 
    'ksp_monitor': None,
    'ksp_converged_reason': None,
}

# params.update(mg_params)
params.update(lu_params)

u = Function(V).zero()

problem = LinearVariationalProblem(A, L, u, bcs=bcs)
solver = LinearVariationalSolver(problem, solver_parameters=params)

solver.solve()

If L = inner(f, v)*dx then the solve succeeds with either lu_params or mg_params.
If L is the assembled form or Cofunction then the solve succeeds with lu_params but breaks with mg_params

Expected behavior
Describe what you expected to happen, in the case of a regression,
include details of a setup where this behaviour was last seen.

I would expect to be able to use multigrid with a Cofunction right hand side

Error message
Add error message with full backtrace, or log.
Please add these as text using three backticks (`) for highlighting.
Please do not add screenshots unless the bug is purely graphical.

Error message implies that coarsening a form involving a Cofunction isn't handled:

[0]PETSC ERROR:                                                                                                                                                                               
[0]PETSC ERROR: WARNING! There are unused option(s) set! Could be the program crashed before usage or a spelling mistake, etc!                                                                
[0]PETSC ERROR:   Option left: name:-firedrake_0_ksp_converged_reason (no value) source: code                                                                                                 
[0]PETSC ERROR:   Option left: name:-firedrake_0_ksp_monitor (no value) source: code                                                                                                          
[0]PETSC ERROR:   Option left: name:-firedrake_0_ksp_rtol value: 1e-05 source: code                                                                                                           
[0]PETSC ERROR:   Option left: name:-firedrake_0_ksp_type value: richardson source: code                                                                                                      
[0]PETSC ERROR:   Option left: name:-firedrake_0_mat_mumps_icntl_14 value: 200 source: code                                                                                                   
[0]PETSC ERROR:   Option left: name:-firedrake_0_mat_type value: aij source: code                                                                                                             
[0]PETSC ERROR:   Option left: name:-firedrake_0_mg_coarse_ksp_type value: preonly source: code                                                                                               
[0]PETSC ERROR:   Option left: name:-firedrake_0_mg_coarse_pc_type value: lu source: code                                                                                                     
[0]PETSC ERROR:   Option left: name:-firedrake_0_mg_levels_pc_type value: jacobi source: code                                                                                                 
[0]PETSC ERROR:   Option left: name:-firedrake_0_pc_type value: mg source: code                                                                                                               
[0]PETSC ERROR:   Option left: name:-firedrake_0_snes_type value: ksponly source: code                                                                                                        
[0]PETSC ERROR: See https://petsc.org/release/faq/ for trouble shooting.                                                              
[0]PETSC ERROR: Petsc Development GIT revision: v3.4.2-43105-gb9546569e7  GIT Date: 2023-10-04 16:14:07 +0100                         
[0]PETSC ERROR: mg_cofunction.py on a default named pop-os by jhc Tue Oct 17 16:36:18 2023                                            
[0]PETSC ERROR: Configure options PETSC_DIR=/home/jhc/codes/firedrake/src/petsc PETSC_ARCH=default --download-hypre --download-hwloc --download-scalapack --download-pnetcdf --download-ml --d
ownload-cmake --download-mpich --download-mumps --with-zlib --download-metis --download-suitesparse --download-superlu_dist --with-fortran-bindings=0 --with-debugging=0 --download-bison --wi
th-c2html=0 --download-ptscotch --download-parmetis --download-hdf5 --download-pastix --download-chaco --with-shared-libraries=1 --download-netcdf
[0]PETSC ERROR: #1 DMCoarsen() at /home/jhc/codes/firedrake/src/petsc/src/dm/interface/dm.c:3284                                      
[0]PETSC ERROR: #2 PCSetUp_MG() at /home/jhc/codes/firedrake/src/petsc/src/ksp/pc/impls/mg/mg.c:962                                   
[0]PETSC ERROR: #3 PCSetUp() at /home/jhc/codes/firedrake/src/petsc/src/ksp/pc/interface/precon.c:1069                                
[0]PETSC ERROR: #4 KSPSetUp() at /home/jhc/codes/firedrake/src/petsc/src/ksp/ksp/interface/itfunc.c:415                               
[0]PETSC ERROR: #5 KSPSolve_Private() at /home/jhc/codes/firedrake/src/petsc/src/ksp/ksp/interface/itfunc.c:836                       
[0]PETSC ERROR: #6 KSPSolve() at /home/jhc/codes/firedrake/src/petsc/src/ksp/ksp/interface/itfunc.c:1082                              
[0]PETSC ERROR: #7 SNESSolve_KSPONLY() at /home/jhc/codes/firedrake/src/petsc/src/snes/impls/ksponly/ksponly.c:49                     
[0]PETSC ERROR: #8 SNESSolve() at /home/jhc/codes/firedrake/src/petsc/src/snes/interface/snes.c:4635                                  
Traceback (most recent call last):                                                                                                    
  File "/home/jhc/codes/mfe_firedrake/mg_cofunction.py", line 53, in <module>                                                         
    solver.solve()                                                                                                                    
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/adjoint_utils/variational_solver.py", line 89, in wrapper                   
    out = solve(self, **kwargs)                                                                                                       
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/variational_solver.py", line 279, in solve                                  
    self.snes.solve(None, work)                                                                                                       
  File "petsc4py/PETSc/SNES.pyx", line 1541, in petsc4py.PETSc.SNES.solve                                                             
  File "petsc4py/PETSc/petscdmshell.pxi", line 239, in petsc4py.PETSc.DMSHELL_Coarsen                                                 
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/dmhooks.py", line 432, in coarsen                                           
    cctx = coarsen(ctx, coarsen)                                                                                                      
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper                                
    return dispatch(args[0].__class__)(*args, **kw)                                                                                   
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/mg/ufl_utils.py", line 254, in coarsen_snescontext                                                                                  
    problem = self(context._problem, self, coefficient_mapping=coefficient_mapping)                              
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper                                                                                                                               
    return dispatch(args[0].__class__)(*args, **kw)                                                                                                                                           
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/mg/ufl_utils.py", line 215, in coarsen_nlvp                                                                                         
    problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, J=J, Jp=Jp,                                                                                                                
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                                                                        
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                                                                        
  File "/home/jhc/codes/firedrake/src/pyadjoint/pyadjoint/tape.py", line 109, in wrapper                                                                                                      
    return function(*args, **kwargs)                                                                                                                                                          
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/adjoint_utils/variational_solver.py", line 26, in wrapper                                                                           
    self._ad_adj_F = adjoint(dFdu)                                                                                                                                                            
  File "petsc4py/PETSc/Log.pyx", line 188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                                                                        
  File "petsc4py/PETSc/Log.pyx", line 189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func                                                                                        
  File "/home/jhc/codes/firedrake/src/firedrake/firedrake/ufl_expr.py", line 336, in adjoint                                                                                                  
    return ufl.adjoint(form, reordered_arguments)                                                                                                                                             
  File "/home/jhc/codes/firedrake/src/ufl/ufl/formoperators.py", line 168, in adjoint                                                                                                         
    return compute_form_adjoint(form, reordered_arguments)                                                                                                                                    
  File "/home/jhc/codes/firedrake/src/ufl/ufl/algorithms/formtransformations.py", line 459, in compute_form_adjoint                                                                           
    raise ValueError("Expecting bilinear form.")                                                                                                                                              
ValueError: Expecting bilinear form.                

Environment:

• OS: [eg: Linux, MacOS, WSL (Windows Subsystem for Linux)] add this as a label too!
• Python version: [eg: 3.9.7]
• Output of firedrake-status
• Any relevant environment variables or modifications [eg: PYOP2_DEBUG=1]

Python version: 3.10.6
OS: PopOS 22.04 LTS

firedrake-status:

Firedrake Configuration:
    package_manager: True
    minimal_petsc: False
    mpicc: None
    mpicxx: None
    mpif90: None
    mpiexec: None
    disable_ssh: False
    honour_petsc_dir: False
    with_parmetis: True
    slepc: True
    packages: ['git+ssh://github.com/firedrakeproject/gusto.git@main#egg=gusto', 'git+ssh://github.com/firedrakeproject/Irksome.git#egg=Irksome']
    honour_pythonpath: False
    opencascade: False
    tinyasm: True
    torch: False
    petsc_int_type: int32
    cache_dir: /home/jhc/codes/firedrake/.cache
    complex: False
    remove_build_files: False
    with_blas: None
    netgen: False
Additions:
    None
Environment:
    PYTHONPATH: None
    PETSC_ARCH: None
    PETSC_DIR: None
Status of components:
---------------------------------------------------------------------------
|Package             |Branch                        |Revision  |Modified  |
---------------------------------------------------------------------------
|FInAT               |master                        |f493084   |False     |
|Irksome             |master                        |bd4a77d   |False     |
|PyOP2               |master                        |da14715c  |False     |
|TinyASM             |master                        |015a89a   |False     |
|fiat                |master                        |e440a06   |False     |
|firedrake           |master                        |88a095f09 |False     |
|gusto               |main                          |b51dc22f  |False     |
|h5py                |firedrake                     |6cc4c912  |False     |
|libspatialindex     |master                        |4768bf3   |True      |
|libsupermesh        |master                        |dbe226b   |False     |
|loopy               |main                          |8158afdb  |False     |
|petsc               |firedrake                     |b9546569e7|False     |
|pyadjoint           |master                        |cefba3f   |False     |
|pytest-mpi          |main                          |a478bc8   |False     |
|slepc               |firedrake                     |6b65b3a8d |False     |
|tsfc                |master                        |47c1324   |False     |
|ufl                 |master                        |cf66c4b3  |False     |
---------------------------------------------------------------------------

Additional Info
Add any other context about the problem here.

[pbrubeck]

We have not implemented coarsening of Cofunctions. Coarsening the original residual is not really required for a linear problem, although firedrake would mostly treat any problem as if it were nonlinear.

[nbouziani]

I have just implemented it. The PR will be there soon, I just need to add a test.

[wence-]

We have not implemented coarsening of Cofunctions. Coarsening the original residual is not really required for a linear problem, although firedrake would mostly treat any problem as if it were nonlinear.

The restriction operator definitely moves cofunctions to the coarse grid (but operates with functions because when we wrote it cofunctions didn't exist)

[pbrubeck]

The code that is currently missing has to do with the symbolic coarsening of the residual one-form when we coarsen a NLVP.

[nbouziani]

I have made a PR (#3181) where I indeed use restrict.