BUG: Fieldsplit doesn't work with a `Cofunction` right hand side.

[JHopeCollins]

Describe the bug

Using a fieldsplit preconditioner with a Cofunction as the right hand side of a LinearVariationalProblem crashes because Cofunction doesn't have a _ufl_class_ attribute.

Steps to Reproduce

from firedrake import *

mesh = UnitIntervalMesh(4)
V = FunctionSpace(mesh, "CG", 1)
W = V*V

u, p = TrialFunctions(W)
v, q = TestFunctions(W)

a = (u*v + p*q)*dx

L = Cofunction(W.dual())

params = {
    'pc_type': 'fieldsplit',
    'pc_fieldsplit_type': 'additive',
    'fieldsplit': {
        'ksp_type': 'preonly',
        'pc_type': 'lu'
    }
}

w = Function(W)
problem = LinearVariationalProblem(a, L, w)
solver = LinearVariationalSolver(problem, solver_parameters=params)

solver.solve()

Expected behavior
I would expect that this would solve the problem without error in one iteration.

Error message

Traceback (most recent call last):
  File "/home/jhc/codes/firedrake_examples/mfe/cofunction_fieldsplit.py", line 28, 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 283, in solve
    self.snes.solve(None, work)
  File "petsc4py/PETSc/SNES.pyx", line 1555, in petsc4py.PETSc.SNES.solve
  File "petsc4py/PETSc/PETSc.pyx", line 323, in petsc4py.PETSc.PetscPythonErrorHandler
  File "petsc4py/PETSc/PETSc.pyx", line 323, in petsc4py.PETSc.PetscPythonErrorHandler
  File "petsc4py/PETSc/PETSc.pyx", line 323, in petsc4py.PETSc.PetscPythonErrorHandler
  [Previous line repeated 6 more times]
  File "petsc4py/PETSc/petscdmshell.pxi", line 341, in petsc4py.PETSc.DMSHELL_CreateFieldDecomposition
  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 347, in create_field_decomposition
    ctxs = ctx.split([(i, ) for i in range(len(W))])
  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/solving_utils.py", line 328, in split
    F = splitter.split(problem.F, argument_indices=(field, ))
  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/formmanipulation.py", line 59, in split
    f = map_integrand_dags(self, form)
  File "/home/jhc/codes/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 73, in map_integrand_dags
    return map_integrands(lambda expr: map_expr_dag(function, expr, compress),
  File "/home/jhc/codes/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 37, in map_integrands
    mapped_components = [map_integrands(function, component, only_integral_type)
  File "/home/jhc/codes/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 37, in <listcomp>
    mapped_components = [map_integrands(function, component, only_integral_type)
  File "/home/jhc/codes/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 66, in map_integrands
    return function(integrand)
  File "/home/jhc/codes/firedrake/src/ufl/ufl/algorithms/map_integrands.py", line 73, in <lambda>
    return map_integrands(lambda expr: map_expr_dag(function, expr, compress),
  File "/home/jhc/codes/firedrake/src/ufl/ufl/corealg/map_dag.py", line 34, in map_expr_dag
    result, = map_expr_dags(function, [expression], compress=compress,
  File "/home/jhc/codes/firedrake/src/ufl/ufl/corealg/map_dag.py", line 100, in map_expr_dags
    r = handlers[v._ufl_typecode_](v, *[vcache[u] for u in v.ufl_operands])
  File "/home/jhc/codes/firedrake/src/ufl/ufl/corealg/multifunction.py", line 99, in undefined
    raise ValueError(f"No handler defined for {o._ufl_class_.__name__}.")
AttributeError: 'Cofunction' object has no attribute '_ufl_class_'

Environment:

• OS: PopOS 22.04
• Python version: 3.10.12
• Output of 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: False
    slepc: False
    packages: []
    honour_pythonpath: False
    opencascade: False
    tinyasm: False
    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                        |e2805c4   |False     |
|PyOP2               |master                        |fbde61f9  |False     |
|fiat                |master                        |e7b2909   |False     |
|firedrake           |master                        |ebd190f27 |False     |
|h5py                |firedrake                     |4c01efa9  |False     |
|libspatialindex     |master                        |4768bf3   |True      |
|libsupermesh        |master                        |dbe226b   |False     |
|loopy               |main                          |8158afdb  |False     |
|petsc               |firedrake                     |09f36907a6|False     |
|pyadjoint           |master                        |f194553   |False     |
|pytest-mpi          |main                          |a478bc8   |False     |
|tsfc                |master                        |799191d   |False     |
|ufl                 |master                        |054b0617  |False     |
---------------------------------------------------------------------------

[JHopeCollins]

@nbouziani @pbrubeck any idea why this is breaking?

[connorjward]

The backtrace makes it look like the issue is with ufl.split. Perhaps a more minimal example could be constructed?

It also looks like _ufl_class_ is only used in undefined. So this issue is really two issues:

1. split doesn't know what to do with Cofunction
2. Cofunction doesn't define _ufl_class_ so the error results in a less-than-informative error message.

[JHopeCollins]

The split call is the method of firedrake.ExtractSubBlock, rather than ufl.split, which splits a form rather than a Function.
ufl.split(Cofunction) also doesn't work (TypeError: 'Cofunction' object is not subscriptable) but I don't think that's the error here.

Seems the error is coming from somewhere in the ufl.map_expr_dags code but I'm having trouble untangling what is actually getting called there.

[connorjward]

I don't understand what is happening inside split in solving_utils.py so I'm not sure how to proceed.

The _ufl_class_ error can be avoided by adding the line

class ExtractSubBlock(MultiFunction):
    ...
    cofunction = expr  # == MultiFunction.reuse_if_untouched

but this just leads to a shape error further on so it might be wrong.

[Ig-dolci]

This bug is affecting the recently merged PR #3723 (Variational adjoint solver) for problems involving a mixed function space.

[JHopeCollins]

I was looking at this again the other day, I'll have a shot at trying to fix it probably next week.