BUG: `matfree` does not work with NLVS and Cofunction.

[JHopeCollins]

Describe the bug
Assembling a matrix free Jacobian crashes for an NLVS when the form has a Cofunction.

Steps to Reproduce

from firedrake import *
from ufl.algorithms import expand_derivatives

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

u = Function(V)
v = TestFunction(V)
b = Cofunction(V.dual())

F = inner(u, v)*dx - b

J = derivative(F, u)

print(f"{type(J) = }")
print(f"{str(J) = }\n")

print(f"{type(expand_derivatives(J)) = }")
print(f"{str(expand_derivatives(J)) = }\n")

print("Building matrix-free solver")
nlvs = NonlinearVariationalSolver(
    NonlinearVariationalProblem(F, u),
    solver_parameters={'mat_type': 'matfree'})

Error message

type(J) = <class 'ufl.form.FormSum'>
str(J) = '1*d/dfj { w₂ * (conj((v_0))) }, with fh=ExprList(*(w₂,)), dfh/dfj = ExprList(*(v_1,)), and coefficient derivatives ExprMapping(*()) * dx(<Mesh #0>[everywhere], {}, {})\n  +  -1*d/dfj { cofunction_3 }, with fh=ExprList(*(w₂,)), dfh/dfj = ExprList(*(v_1,)), and coefficient derivatives ExprMapping(*())'

type(expand_derivatives(J)) = <class 'ufl.form.Form'>
str(expand_derivatives(J)) = 'v_1 * (conj((v_0))) * dx(<Mesh #0>[everywhere], {}, {})'

Building matrix-free solver
Traceback (most recent call last):
  File "/home/jhc/codes/fd/fd-dev/wrk/auxsnes/formsum_mat_mfe/formsum_mat_mfe.py", line 26, in <module>
    nlvs = NonlinearVariationalSolver(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 250, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 251, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/usr/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/adjoint_utils/variational_solver.py", line 47, in wrapper
    init(self, problem, *args, **kwargs)
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/variational_solver.py", line 307, in __init__
    ctx.set_jacobian(self.snes)
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/solving_utils.py", line 331, in set_jacobian
    snes.setJacobian(self.form_jacobian, J=self._jac.petscmat,
                                           ^^^^^^^^^
  File "/usr/lib/python3.12/functools.py", line 995, in __get__
    val = self.func(instance)
          ^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/solving_utils.py", line 539, in _jac
    return self._assembler_jac.allocate()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/assemble.py", line 362, in allocate
    return MatrixFreeAssembler(self._form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/assemble.py", line 954, in __new__
    raise TypeError(f"The first positional argument must be of ufl.Form or slate.TensorBase: got {type(form)} ({form})")
TypeError: The first positional argument must be of ufl.Form or slate.TensorBase: got <class 'ufl.form.FormSum'> (1*d/dfj { w₂ * (conj((v_0))) }, with fh=ExprList(*(w₂,)), dfh/dfj = ExprList(*(v_1,)), and coefficient derivatives ExprMapping(*()) * dx(<Mesh #0>[everywhere], {}, {})
  +  -1*d/dfj { cofunction_3 }, with fh=ExprList(*(w₂,)), dfh/dfj = ExprList(*(v_1,)), and coefficient derivatives ExprMapping(*()))

Additional Info
It appear that derivative(F, u) is returning a FormSum because it isn't removing the derivative(Cofunction, u) term. The matrix free assembly then complains that needs a Form not a FormSum. Using ufl.algorithms.expand_derivatives after calling derivative(F, u) returns a Form without the Cofunction term, so maybe we need a call to that somewhere in the NLVS setup.

[JHopeCollins]

@pbrubeck any ideas? I'm sure this used to work! I get the same error message with firedrake-2025.10.2.

[pbrubeck]

MatrixFreeAssembler should happily take any BaseForm and just call assemble on it. What happens if you remove the type check?

[JHopeCollins]

Should there be at least some check that it is a bilinear form or equivalent?

If I remove the type check then it hits an error further down because a FormSum doesn't have an assembler cache like we stash on a Form.

type(J) = <class 'ufl.form.FormSum'>
str(J) = '1*d/dfj { w₂ * (conj((v_0))) }, with fh=ExprList(*(w₂,)), dfh/dfj = ExprList(*(v_1,)), and coefficient derivatives ExprMapping(*()) * dx(<Mesh #0>[everywhere], {}, {})\n  +  -1*d/dfj { cofunction_3 }, with fh=ExprList(*(w₂,)), dfh/dfj = ExprList(*(v_1,)), and coefficient derivatives ExprMapping(*())'

type(expand_derivatives(J)) = <class 'ufl.form.Form'>
str(expand_derivatives(J)) = 'v_1 * (conj((v_0))) * dx(<Mesh #0>[everywhere], {}, {})'

Building matrix-free solver
Traceback (most recent call last):
  File "/home/jhc/codes/fd/fd-dev/wrk/auxsnes/formsum_mat_mfe/formsum_mat_mfe.py", line 26, in <module>
    nlvs = NonlinearVariationalSolver(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 250, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 251, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/usr/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/adjoint_utils/variational_solver.py", line 47, in wrapper
    init(self, problem, *args, **kwargs)
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/variational_solver.py", line 307, in __init__
    ctx.set_jacobian(self.snes)
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/solving_utils.py", line 331, in set_jacobian
    snes.setJacobian(self.form_jacobian, J=self._jac.petscmat,
                                           ^^^^^^^^^
  File "/usr/lib/python3.12/functools.py", line 995, in __get__
    val = self.func(instance)
          ^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/solving_utils.py", line 539, in _jac
    return self._assembler_jac.allocate()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/assemble.py", line 362, in allocate
    return MatrixFreeAssembler(self._form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jhc/codes/fd/fd-dev/src/firedrake/firedrake/assemble.py", line 975, in __new__
    if _FORM_CACHE_KEY not in form._cache:
                              ^^^^^^^^^^^
AttributeError: 'FormSum' object has no attribute '_cache'

[pbrubeck]

Should there be at least some check that it is a bilinear form or equivalent?

Ideally yes.

If I remove the type check then it hits an error further down because a FormSum doesn't have an assembler cache like we stash on a Form.

This means it's going through a codepath more specific to OneFormAssembler. We must have assumed somewhere that we were dealing with a Form and not a generic BaseForm.

I am surprised we have only noticed this until now.

[JHopeCollins]

I am surprised we have only noticed this until now.

Me too! I am sure that I have used a cofunction with NLVS and a matrix free pc before without issue.

This means it's going through a codepath more specific to OneFormAssembler. We must have assumed somewhere that we were dealing with a Form and not a generic BaseForm.

Yes, It looks like the code path is as follows:

1. SNESContext._assembler_jac is a BaseFormAssembler, assumedly because derivative(F, u) isn't eliminating the differentiated Cofunction (maybe erroneously?).
2. BaseFormAssembler.allocate internally uses MatrixFreeAssembler because mat_type='matfree'.
3. MatrixFreeAssembler inherits from FormAssembler, so we end up in FormAssembler.__new__. So now we are passing a BaseForm to something expecting a Form.
4. FormAssembler.__new__ checks to see if the "Form" already has an assembler in its ufl.Form._cache. However now we're in trouble because BaseForm doesn't have a cache.

I think there's three (non-exclusive) possibilities here:

1. MatrixFreeAssembler should inherit straight from AbstractFormAssembler and then handle the two cases of Form or BaseForm separately.
2. The _cache attribute should be pushed up the inheritance tree to ufl.BaseForm.
3. Somewhere along the road we should call expand_derivatives on the Jacobian form to try and simplify it before hitting any of this.

My gut feeling is that 2 is certainly the right thing to do, 3 is probably good to do (but I don't know if there are other implications of calling expand_derivatives), and 1 would be useful but could end up being a substantial refactor.

What do you think?

[pbrubeck]

I think there's three (non-exclusive) possibilities here:

1. MatrixFreeAssembler should inherit straight from AbstractFormAssembler and then handle the two cases of Form or BaseForm separately.
2. The _cache attribute should be pushed up the inheritance tree to ufl.BaseForm.
3. Somewhere along the road we should call expand_derivatives on the Jacobian form to try and simplify it before hitting any of this.

What do you think?

The real issue here is that a true BaseForm Jacobian is not supported if mat_type="matfree". Ideally we should accept anything and let BaseFormAssembler deal with it, as it knows how to dispatch an appropriate OneFormAssembler without calling expand_derivatives.

1. I'm not familiar with the inheritance model, but this is a very reasonable alternative to my proposal of treating everything as BaseForm.

2. Is a good thing to do, but this won't help unless we implement caching for BaseFormAssembler.

3. Should happen at some point, we should not be adding extra calls to expand_derivatives, it should happen at the right time, and only called by the form compiler.

[JHopeCollins]

The real issue here is that a true BaseForm Jacobian is not supported if mat_type="matfree". Ideally we should accept anything and let BaseFormAssembler deal with it, as it knows how to dispatch an appropriate OneFormAssembler without calling expand_derivatives.

Isn't this a 2-form? It's a Jacobian, and the original type check failure came after an if rank == 2 check.

1. I'm not familiar with the inheritance model, but this is a very reasonable alternative to my proposal of treating everything as BaseForm.

I've just had a chat with @connorjward and we think that a possible fix would be bumping the caching logic in FormAssembler.__new__ up to AbstractFormAssembler.__new__ so it also applies to BaseFormAssembler. We'd then just need to implement BaseFormAssembler._cache_key() -> None, and add the ufl.BaseForm._cache attribute.

Sound reasonable? It should be a faster route to fixing this bug than either a bigger refactor of the inheritance tree, or aa refactor to treat everything as a BaseForm. Although if you think that would be preferable long term then we can open an issue with some detail about it.

[pbrubeck]

Isn't this a 2-form? It's a Jacobian, and the original type check failure came after an if rank == 2 check.

Yes it it a 2-form, but MatrixFreeAssembler will assemble the action of the 2-Form with a function, so it becomes a 1-Form.

I've just had a chat with @connorjward and we think that a possible fix would be bumping the caching logic in FormAssembler.__new__ up to AbstractFormAssembler.__new__ so it also applies to BaseFormAssembler. We'd then just need to implement BaseFormAssembler._cache_key() -> None, and add the ufl.BaseForm._cache attribute.

Sound reasonable?

I do not think that implementing a cache key will immediately solve this particular issue, to fix the bug we just need to dispatch the correct Assembler class.

[JHopeCollins]

I do not think that implementing a cache key will immediately solve this particular issue, to fix the bug we just need to dispatch the correct Assembler class.

The BaseForm._cache_key would just be None, which mean "we aren't doing caching for this type of Assembler, just create a new one every time". It's actually what most of the assembler classes do. In essence it'd just be saying "BaseFormAssembler" is something we could do caching for in the future, but we aren't now but it's fine if we end up in the cache-checking code paths with a BaseForm.

[pbrubeck]

I do not think that implementing a cache key will immediately solve this particular issue, to fix the bug we just need to dispatch the correct Assembler class.

The BaseForm._cache_key would just be None, which mean "we aren't doing caching for this type of Assembler, just create a new one every time". It's actually what most of the assembler classes do. In essence it'd just be saying "BaseFormAssembler" is something we could do caching for in the future, but we aren't now but it's fine if we end up in the cache-checking code paths with a BaseForm.

I'm fine with this a first step towards caching, but I still view this a separate issue.

[JHopeCollins]

I'd agree it is a tangential issue in terms of what the code should be capable of, but it is a more immediate fix for this bug...

[pbrubeck]

I'd agree it is a tangential issue in terms of what the code should be capable of, but it is a more immediate fix for this bug...

Does it actually fix the bug?

[JHopeCollins]

We shall see...

[pbrubeck]

A workaround is to set an empty BaseForm._cache, or simply skip the caching if the attribute is not there. This will fix

F = inner(u, v)*dx - Cofunction(V.dual())
J = derivative(F, u)    
assemble(J, mat_type="matfree")

and solve(F == 0, u, solver_parameters={"mat_type": "matfree"})

But not

get_assembler(J, mat_type="matfreee").assemble()

The real issue is that mat_type="matfree" causes delayed derivative expansion, and an non-expanded CoefficientDerivative that should be zero (derivative(Cofunction(V.dual()), u)) fails to be evaluated.