BUG: 'incompatible function spaces' error in adjoint interpolation via assemble

[jrmaddison]

Describe the bug

(For appropriate spaces, $V_1$, $V_2$, e.g. each $P_n$) Interpolation from $V_2$ to $V_1$ is a linear operator, say $A : V_2 \rightarrow V_1$. Adjoint interpolation is the linear operator $A^* : V_1^* \rightarrow V_2^*$ with

$$\forall v \in V_1^*, u \in V_2 \qquad v(A(u)) = A^{*}(v)(u).$$

Using Firedrake to compute $A^{*}(v)$ via assemble

assemble(action(adjoint(Interpolate(TestFunction(V_2), V_1)), v))

leads to an incompatible function spaces error.

Steps to Reproduce

import firedrake as fd

mesh = fd.UnitIntervalMesh(1)
V_1 = fd.FunctionSpace(mesh, "CG", 1)
V_2 = fd.FunctionSpace(mesh, "CG", 2)

# v \in V_1^*
# v(u) = \int_\Omega u^*
v = fd.Cofunction(V_1.dual(), name="v")
fd.assemble(fd.conj(fd.TestFunction(V_1)) * fd.dx, tensor=v)

# A : V_2 -> V_1
interp = fd.Interpolate(fd.TestFunction(V_2), V_1)

# A^* : V_1^* \rightarrow V_2^*
interp_adj = fd.adjoint(interp)

# A^*(v) \in V_2^*
# A^*(v)(u) = \int_\Omega [ A(u) ]^*
form = fd.action(interp_adj, v)

w = fd.assemble(form)

leads to

Traceback (most recent call last):
  File "/home/jmaddis2/build/firedrake/test/test.py", line 22, in <module>
    w = fd.assemble(form)
        ^^^^^^^^^^^^^^^^^
  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/jmaddis2/build/firedrake/firedrake/src/firedrake/firedrake/adjoint_utils/assembly.py", line 28, in wrapper
    form = BaseFormAssembler.preprocess_base_form(form)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jmaddis2/build/firedrake/firedrake/src/firedrake/firedrake/assemble.py", line 840, in preprocess_base_form
    expr = BaseFormAssembler.restructure_base_form_preorder(expr)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jmaddis2/build/firedrake/firedrake/src/firedrake/firedrake/assemble.py", line 691, in restructure_base_form_preorder
    return BaseFormAssembler.reconstruct_node_from_operands(expression, operands)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jmaddis2/build/firedrake/firedrake/src/firedrake/firedrake/assemble.py", line 648, in reconstruct_node_from_operands
    return expr._ufl_expr_reconstruct_(*operands)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jmaddis2/build/firedrake/firedrake/src/ufl/ufl/form.py", line 224, in _ufl_expr_reconstruct_
    return type(self)(*operands)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/jmaddis2/build/firedrake/firedrake/src/ufl/ufl/action.py", line 100, in __init__
    _check_function_spaces(left, right)
  File "/home/jmaddis2/build/firedrake/firedrake/src/ufl/ufl/action.py", line 203, in _check_function_spaces
    raise TypeError("Incompatible function spaces in Action")
TypeError: Incompatible function spaces in Action

Additional Info
Blocking #3939 and #3965. Seems to have been introduced recently.

[connorjward]

@pbrubeck any ideas?

[pbrubeck]

I'm familiar with the brokenness of action(adjoint(interpolate), Cofunction), but one could usually go around this with something like action(Cofunction, interpolate). If I replace form = fd.action(v, interp) I get the following backtrace:

ValueError                                Traceback (most recent call last)
Cell In[1], line 22
     18 # A^*(v) \in V_2^*
     19 # A^*(v)(u) = \int_\Omega [ A(u) ]^*
     20 form = fd.action(v, interp)
---> 22 w = fd.assemble(form)

File petsc4py/PETSc/Log.pyx:188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File petsc4py/PETSc/Log.pyx:189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/adjoint_utils/assembly.py:30, in annotate_assemble.<locals>.wrapper(form, *args, **kwargs)
     28         form = BaseFormAssembler.preprocess_base_form(form)
     29         kwargs['is_base_form_preprocessed'] = True
---> 30     output = assemble(form, *args, **kwargs)
     32 from firedrake.function import Function
     33 from firedrake.cofunction import Cofunction

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:134, in assemble(expr, *args, **kwargs)
    132     raise RuntimeError(f"Got unexpected args: {args}")
    133 tensor = kwargs.pop("tensor", None)
--> 134 return get_assembler(expr, *args, **kwargs).assemble(tensor=tensor)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:386, in BaseFormAssembler.assemble(self, tensor)
    384 # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly.
    385 visited = {}
--> 386 result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited)
    388 # Apply BCs after assembly
    389 rank = len(self._form.arguments())

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:620, in BaseFormAssembler.base_form_postorder_traversal(expr, visitor, visited)
    618         stack.extend(unvisited_children)
    619     else:
--> 620         visited[e] = visitor(e, *(visited[arg] for arg in operands))
    622 return visited[expr]

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:382, in BaseFormAssembler.assemble.<locals>.visitor(e, *operands)
    380 def visitor(e, *operands):
    381     t = tensor if e is self._form else None
--> 382     return self.base_form_assembly_visitor(e, t, *operands)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:558, in BaseFormAssembler.base_form_assembly_visitor(self, expr, tensor, *args)
    556 # Assembling the Jacobian action.
    557 if interpolator.nargs:
--> 558     return interpolator._interpolate(expression, output=tensor, default_missing_val=default_missing_val)
    559 # Assembling the operator
    560 if tensor is None:

File petsc4py/PETSc/Log.pyx:188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File petsc4py/PETSc/Log.pyx:189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/interpolation.py:853, in SameMeshInterpolator._interpolate(self, output, transpose, *function, **kwargs)
    851 function, = function
    852 if not hasattr(function, "dat"):
--> 853     raise ValueError("The expression had arguments: we therefore need to be given a Function (not an expression) to interpolate!")
    854 if transpose:
    855     mul = assembled_interpolator.handle.multTranspose

ValueError: The expression had arguments: we therefore need to be given a Function (not an expression) to interpolate!

[pbrubeck]

Also there is a difference between fd.interpolate and fd.Interpolate, if I use fd.interpolate I get
TypeError: Incompatible RHS for Action.

[pbrubeck]

Ah of course, from firedrake.__future__ import interpolate. This fixes form = fd.action(v, interp), but the issue with interp_adj still remains.

import firedrake as fd
from firedrake.__future__ import interpolate

mesh = fd.UnitIntervalMesh(1)
V_1 = fd.FunctionSpace(mesh, "CG", 1)
V_2 = fd.FunctionSpace(mesh, "CG", 2)

# v \in V_1^*
# v(u) = \int_\Omega u^*
v = fd.Cofunction(V_1.dual(), name="v")
fd.assemble(fd.conj(fd.TestFunction(V_1)) * fd.dx, tensor=v)

# A : V_2 -> V_1
interp = interpolate(fd.TestFunction(V_2), V_1)

form = fd.action(v, interp)

w = fd.assemble(form)

[jrmaddison]

For the case I needed the new Interpolator can be used as a workaround

import firedrake as fd
from firedrake.__future__ import Interpolator

mesh = fd.UnitIntervalMesh(1)
V_1 = fd.FunctionSpace(mesh, "CG", 1)
V_2 = fd.FunctionSpace(mesh, "CG", 2)

# v \in V_1^*
# v(u) = \int_\Omega u^*
v = fd.Cofunction(V_1.dual(), name="v")
fd.assemble(fd.conj(fd.TestFunction(V_1)) * fd.dx, tensor=v)

# A^*(v) \in V_2^*
# A^*(v)(u) = \int_\Omega [ A(u) ]^*
form = Interpolator(fd.TestFunction(V_2), V_1).interpolate(v, adjoint=True)

w = fd.assemble(form)

[pbrubeck]

The original MFE runs fine if you just replace TestFunction -> TrialFunction:

# A : V_2 -> V_1
interp = fd.Interpolate(fd.TrialFunction(V_2), V_1)

[pbrubeck]

Equivalently, form = fd.action(v, interp)

[jrmaddison]

Hmm, interesting, but is that correct? The interpolated expression, as passed to the Interpolator, should only take one argument, so I though TestFunction (number 0) should be used.

For action I'm not sure if that's equivalent in the complex case.

[pbrubeck]

I think the logic goes like this: Interpolate(Function(V_src), V_dst) should return a 1-form in V_dst.dual() a.k.a. a Function(V_dst), so it is natural for the 0th argument to live in V_dst.dual(). If you want to add an additional argument in V_src, then it needs to be Argument(V_src, number=1), therefore a TrialFunction(V_src).

I see that fd.__future__.interpolate (wrapping the old interpolate) silently promotes a TestFunction into TrialFunction.
https://github.com/firedrakeproject/firedrake/blame/58f3d6f26d2a564f75dcf8774c57a7ebe2d077c8/firedrake/interpolation.py#L295-299

Is there a reson to prefer Interpolator over fd.__future__.interpolate?

[jrmaddison]

I see, viewing it as a sesquilinear form on $V_1^* \times V_2$. I think I followed test code in the use of TestFunction.

Is there a reson to prefer Interpolator over fd.future.interpolate

No, I don't think so.

[pbrubeck]

Should we then throw an error for Interpolator(TestFunction(V_1), V_2)?

[jrmaddison]

Maybe the exception should be raised in Interpolate.__init__, possibly in UFL. In any case I hadn't realized Interpolate needs to use TrialFunction, e.g.

interp = fd.Interpolate(fd.TrialFunction(V_2), V_1)

is fine. I also think TestFunction worked until quite recently, but perhaps not for a good reason.