Support for `ufl.Argument`

[jorgensd]

It would be great to support ufl.Argument in external operators.

For a single input argument, you would get "one" coefficient per degree of freedom, which has to be assembled appropriately.

[jorgensd]

Minimal failing example:

from mpi4py import MPI

import basix.ufl
import dolfinx
import ufl
from dolfinx_external_operator import (
    FEMExternalOperator,
    evaluate_external_operators,
    evaluate_operands,
    replace_external_operators,
)


def g_impl(uuh):
    """Identity operator"""
    output = uuh.copy()
    return output.reshape(-1)


def g_external(derivatives):
    if derivatives == (0,):  # no derivation, the function itself
        return g_impl
    else:
        return NotImplementedError


quadrature_degree = 2
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 5, 5)
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)

V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
u = ufl.TrialFunction(V)

Qe = basix.ufl.quadrature_element(
    mesh.basix_cell(), degree=quadrature_degree, value_shape=())
Q = dolfinx.fem.functionspace(mesh, Qe)

f = FEMExternalOperator(
    u, function_space=Q, external_function=g_external
)

dx = ufl.Measure("dx", domain=mesh)
v = ufl.TestFunction(V)
a = f*v*dx

a_replaced, a_external_operators = replace_external_operators(a)
# Pack coefficients for g
evaluated_operands = evaluate_operands(a_external_operators)
evaluate_external_operators(
    a_external_operators, evaluated_operands)

a_compiled = dolfinx.fem.form(a_replaced)
A = dolfinx.fem.assemble_matrix(a_compiled)
A.assemble()

returning

Traceback (most recent call last):
  File "/root/shared/test/example.py", line 54, in <module>
    evaluate_external_operators(
  File "/root/shared/src/dolfinx_external_operator/external_operator.py", line 178, in evaluate_external_operators
    np.copyto(external_operator.ref_coefficient.x.array, external_operator_eval)
ValueError: could not broadcast input array from shape (450,) into shape (150,)

as there are three basis functions per "coefficient" entry in the operator.