# External operators¶

Firedrake provides a seamless and easy-to-use interface to incorporate arbitrary operators within PDE
systems with automated assembly and differentiation capabilities. This allows for embedding complex
operators, such as machine learning models, into PDE or PDE-constrained optimisation problems.
We refer to these operators as *external operators*, as their implementation is left to be specified by
the user, i.e “external” with respect to Firedrake.

External operators are represented by the `AbstractExternalOperator`

base class,
which can be subclassed for defining specific operators. While the base class documentation
provides the syntax, this document introduces the external operator abstraction in more detail and
explains what you need to do if you want to define your own external operator. The external operator
framework is discussed in more detail in [BH21] and [BHF24].

## Background¶

An external operator is an operator that maps operands defined on finite element spaces to an output in another finite element space. Let’s consider the case of an external operator \(N\) with a single operand, i.e. we have:

where \(V\) is a finite element space. Because the spaces we are interested in are reflexive, we can express the operator \(N\) as a linear form \(L_{N}\), defined as:

where \(V^{*}\) is the dual space of \(V\), and where the following equivalence holds:

for \(u \in V\). For sake of simplicity, in what follows, we use \(N\) to denote both the linear form \(L_{N}\)
and the operator \(N\). Expressing external operators as linear forms facilitates their
composition with variational forms of PDEs. Note that while \(N\) is linear with respect to \(v^{*}\),
its *argument*, it can be nonlinear with respect to \(u\), its *operand*. In Firedrake, \(u\) would
be a `Function`

object, and \(v^{*}\) a `Coargument`

object.

Arbitrary UFL expressions defined on arbitrary finite element spaces, i.e. other than \(V\), can also
be passed as operands to the external operator using the symbolic `Interpolate`

operator.

### Differentiation¶

An essential aspect of the external operator abstraction is differentiation, not least because solving a PDE or a PDE-constrained optimisation problem involving an external operator may require differentiating through it. In order to illustrate how differentiation works, let’s consider the case of a variational form containing the previously defined external operator \(N(u; v^{*})\). We want to find \(u \in V\) such that:

In order to solve this variational problem, we need to compute the Jacobian of \(F\) with respect to \(u\) , see solving variational problems, which involves differentiating \(N\) with respect to \(u\), as a consequence of the chain rule. More specifically, the Jacobian of \(F\) can be written as:

where \(\hat{u} \in V\) is the Gâteaux direction, i.e. a trial function on \(V\). The Jacobian of \(N\), \(\frac{dN(u; \hat{u}, v^{*})}{du}\), is also an external operator with operand \(u\) and whose arguments are \(\hat{u} \in V\) and \(v^{*} \in V^{*}\). Note that the Jacobian of \(N\) is a 2-form, i.e. a linear form with 2 (unknown) arguments.

### Action/Adjoint¶

The external operator interface doesn’t only provide differentiation capabilities, but also allows for performing symbolic operations on external operators. There are two main symbolic operations that often arise when solving PDE systems, namely \((i)\) taking the action of a linear form on a given input, e.g. to represent the Jacobian-vector product of \(N\), and \((ii)\) taking the adjoint of a linear 2-form, e.g. to represent the Hermitian transpose of the Jacobian of \(N\). In Firedrake, these operations are represented by the \(\operatorname{action}\) and \(\operatorname{adjoint}\) functions.

It turns out that those two symbolic operations simply translate into updating the arguments of the linear form on which they are applied. For example, let’s consider the Jacobian of the external operator \(N\), we have:

where the external operator \(\frac{dN(u; w, v^{*})}{du}\) is the Jacobian action of \(N\) on \(w\).
Hence, for a given `Function`

\(w\) in \(V\), the external operator \(\frac{dN(u; w, v^{*})}{du}\)
is linear with respect to \(w\), which is a known function in \(V\), and \(v^{*}\), which is an unknown
function in \(V^{*}\), i.e. a `Coargument`

in \(V^{*}\). Therefore, \(\frac{dN(u; w, v^{*})}{du}\)
is a 1-form, i.e. it is a linear form with respect to one unknown argument.

Similarly, the adjoint of the Jacobian of \(N\) can be written as:

where the external operator \(\frac{dN(u; v^{*}, \hat{u})}{du}\) is the adjoint of the Jacobian of \(N\), which is still a 2-form as it is linear with respect to two unknown arguments, namely the coargument \(v^{*} \in V^{*}\) and the trial function \(\hat{u} \in V\).

Put simply, taking the action of an external operator consumes its last argument, while taking the adjoint of a 2-form external operator swaps its two arguments. Given that linear form’s arguments in Firedrake are numbered, we can also say that the action results in consuming the highest-numbered argument, while the adjoint simply swaps the arguments’ numbers.

### Assembly¶

The last ingredient needed for external operators to work is their coupling with the finite element assembly. The external operator interface composes seamlessly with the Firedrake assembly system, and one can assemble external operators as well as variational forms and expressions containing external operators. Firedrake assembles expressions or variational forms containing external operators by representing the assembled object as a directed acyclic graph (DAG). For example, assembling the Jacobian of \(F\), see differentiation section, would result in the following DAG:

The assembly is achieved by traversing this DAG and evaluating each node on the fly.
In this DAG representation, external operators are considered as distinct nodes, whose
implementation is, by definition, left to be specified by the user. Consequently,
in order to evaluate the expression of interest, we need a mechanism to map the assembly of the
external operator nodes in the DAG to the corresponding implementation specified by the user. This stage is
referred to as the *external operator assembly*. This is achieved by the `AbstractExternalOperator`

base class, whose main role is to orchestrate the external operator assembly.

Depending on the specific use case considered, one may need to specify one or more evaluation implementations for a particular external operator. For instance, if the PDE system considered only requires to evaluate the external operator, the user-defined implementation only needs to specify the implementation for evaluating that operator \(N(u; v^{*})\). On the other hand, if the PDE system considered also requires evaluating the Jacobian of \(N\), e.g. when solving a PDE whose residual contains an external operator that has the solution \(u\) as operand, the user implementation needs to specify the implementation for evaluating the Jacobian \(\frac{dN(u; \hat{u}, v^{*})}{du}\) as well. The external operator might also be further equipped with the implementation for evaluating the action of the Jacobian adjoint of \(N\), i.e. \(\frac{dN(u; w^{*}, \hat{u})}{du}\) with \(w^{*} \in V^{*}\), e.g. when solving a PDE-constrained optimisation problem, where the cost function derivative is computed using the adjoint method.

The external operator assembly is facilitated by the fact that for a given operator \(N(u; v^{*})\), the type of external operators being assembled (e.g. \(N\), its Jacobian, etc.) is fully determined by: \((i)\) the derivative multi-index and \((ii)\) and the external operator arguments. More specifically, the derivatives multi-index of an external operator indicates the derivative(s) taken with respect to the operand(s). On the other hand, the arguments indicate whether the highest-numbered argument was replaced, i.e. if the action was taken, or if the arguments’ numbers were swapped, i.e. if the adjoint was taken, see previous section.

Finally, the number of arguments of an external operator also determines the type of its output after
being assembled. For instance, given that an external operator with one argument is a 1-form, its assembly
would result in a `Function`

or a `Cofunction`

. Similarly, an external operator
with two arguments, i.e. a 2-form, would produce a `MatrixBase`

object. The following table
illustrates an external operator \(N\) with two operands,
\(N \colon V \times V \times V^{*} \rightarrow \mathbb{R}\), along with the external operators
resulting from different symbolic operations on it, such as differentiation or action/adjoint.
Additionally, the table includes the derivative multi-index, argument slots, and assembly type
of each external operator:

## Build your own external operator¶

External operators can be used to incorporate arbitrary operators within Firedrake. Some external
operators are already implemented, such as the `MLOperator`

for embedding machine learning
models within Firedrake, whose implementation is explained in more detail in [BH21] and [BHF24].
However, you may want to build your own external operator for your specific problem.
In this section, we discuss how new external operators can be defined.

To define a new external operator, one first needs to subclass the `AbstractExternalOperator`

class. Then, the external operator subclass needs to be equipped with methods specifying how the
different types of external operator arising in the PDE system considered can be evaluated, such as the
Jacobian of the operator. Note that you only need to specify the evaluation methods required for your
problem of interest.

We have previously seen the different external operators you can get as a result of symbolic operations
such as differentiation, action or adjoint, can all be understood as changes in the derivative multi-index
and/or the arguments of the external operator. The external operator interface uses this logic to allow
users specifying which external operator implementation each method correspond to. More specifically,
each evaluation method of the subclass needs to be decorated with the *assemble_method* decorator. This
decorator takes in two arguments: \((i)\) the derivative multi-index, and \((ii)\) and a tuple containing the
arguments’ numbers, wherein arguments that are not of type `Argument`

or `Coargument`

are denoted with \(None\).

For instance, the Jacobian of the previously introduced external operator \(N(u, m; v^{*})\) with
respect to \(u \in V\), i.e. \(\frac{\partial N(u, m; \hat{u}, v^{*})}{\partial u}\), has two arguments:
\(v^{*}\) and \(\hat{u}\). Linear forms’ arguments are numbered incrementally, starting from 0, as new
arguments are added to the linear form as a result of differentiation. Hence, the second entry of the
*assemble_method* decorator for this Jacobian would be *(0, 1)*. As for the derivative multi-index, it
is *(1, 0)* as we differentiated with respect to the first operand, i.e. \(u\). Therefore, the
decorator for the evaluation of the Jacobian would be:

```
@assemble_method((1, 0), (0, 1))
```

If we now take the adjoint of that Jacobian, i.e. its Hermitian transpose, we would still have the same derivative multi-index, but the arguments would be swapped. Hence, the specified decorator should be:

```
@assemble_method((1, 0), (1, 0))
```

Likewise, if we take the action of the Hermitian transpose of the Jacobian matrix on a given
`Cofunction`

, the highest-numbered argument will be replaced by this cofunction. This implies that
the highest number in the second tuple of the decorator will be replaced by *None* as cofunctions are
not `Argument`

or `Coargument`

, which results in:

```
@assemble_method((1, 0), (None, 0))
```

The following example illustrates how to define a new external operator, named *MyExternalOperator*,
and sketch out a few evaluation methods:

```
class MyExternalOperator(AbstractExternalOperator):
def __init__(self, *args, **kwargs):
...
@assemble_method((0, 0), (0,))
# or @assemble_method(0, (0,))
def N(self, *args, *kwargs):
"""Evaluate the external operator N"""
...
@assemble_method((1, 0), (0, 1))
def dNdu(self, *args, **kwargs):
"""Evaluate the Jacobian dNdu"""
...
@assemble_method((1, 0), (0, None))
def dNdu_action(self, *args, **kwargs):
"""Evaluate the action of the Jacobian dNdu"""
...
@assemble_method((0, 1), (1, 0))
def dNdm_adjoint(self, *args, **kwargs):
"""Evaluate the Hermitian transpose of the Jacobian dNdm"""
...
@assemble_method((0, 1), (None, 0))
def dNdm_adjoint_action(self, *args, **kwargs):
"""Evaluate the action of the Hermitian transpose of the Jacobian dNdm"""
...
...
```

While the above template addresses the case of an external operator with two operands, there are no restrictions on the number of operands an external operator can have. Additionally, not all the above methods need to be implemented; only those required for the problem of interest should be included.

## A simple example: the translation operator¶

In this section, we showcase the external operator interface on a very simple example, namely the translation operator \(N \colon V \times V \times V^{*} \rightarrow \mathbb{R}\), or equivalently \(N \colon V \times V \rightarrow V\), defined as:

Note that building an external operator for the above operation is, in practice, not necessary as this
can already be readily implemented using Firedrake’s built-in functionalities. Also, because this
translation operation is fully defined in Firedrake, the evaluation methods of the external operator we will
build rely on Firedrake code. However, the external operator evaluation methods can contain any Python
code as long as they return compatible objects, e.g. `Function`

or `MatrixBase`

objects.

\(N\) takes in two operands \(f, u \in V\) and one argument \(v^{*} \in V^{*}\). When assembled,
this external operator returns a `Function`

in \(V\) since the linear form \(N\) can also
be seen as an operator mapping to \(V\), as previously discussed. To construct \(N\),
we need to subclass the `AbstractExternalOperator`

class and specify how \(N\) can be assembled.
Given that \(N\) has \((0,)\) as derivative multi-index and that it only has one argument,
the translation operator subclass can be defined as:

```
class TranslationOperator(AbstractExternalOperator):
def __init__(self, *operands, function_space, operator_data, **kwargs):
AbstractExternalOperator.__init__(self, *operands,
function_space=function_space,
operator_data=operator_data,
**kwargs)
@assemble_method(0, (0,))
def assemble_N(self, *args, **kwargs):
"""Evaluate the translation operator N"""
u, f = self.ufl_operands
N = assemble(u - f)
return N
N = TranslationOperator(u, f, function_space=V)
```

Note that the above subclass takes in an *operator_data* argument. This keyword argument allows
users to attach data specfic to their operator. The data stashed in *operator_data* are accessible to all
the external operators induced by the original operator. For instance, the Jacobian of an external
operator \(N\), which is an external operator, will also have access to the content of *operator_data*.

Now that we have specified the implementation for evaluating \(N\), we can assemble it:

```
assembled_N = assemble(N)
assert np.allclose(assembled_N.dat.data_ro, u.dat.data_ro[:] - f.dat.data_ro[:])
```

Assembling an external operator is often not enough, in particular as the external operator of interest can be used in a variational form, which may require providing an implementation for its Jacobian as well. For example, let’s consider the following variational problem for \(u \in H^{1}(\Omega)\)

Using \(N\), we can derive the following variational form

Solving this variational problem necessitates calculating the Jacobian of the above residual form, which in turn requires computing the Jacobian \(\frac{\partial N(u, f; \hat{u}, v^{*})}{\partial u}\), which in this case is the identity matrix. Hence, we now need to add an implementation specifying how the Jacobian of \(N\) can be assembled:

```
class TranslationOperator(AbstractExternalOperator):
def __init__(self, *operands, function_space, **kwargs):
AbstractExternalOperator.__init__(self, *operands, function_space=function_space, **kwargs)
@assemble_method(0, (0,))
def assemble_N(self, *args, **kwargs):
"""Evaluate N"""
u, f = self.ufl_operands
N = assemble(u - f)
return N
@assemble_method((1, 0), (0, 1))
def assemble_Jacobian(self, *args, **kwargs):
"""Evaluate the Jacobian of N"""
dNdu = Function(self.function_space()).assign(1)
# Construct the Jacobian matrix
integral_types = set(['cell'])
assembly_opts = kwargs.get('assembly_opts')
J = self._matrix_builder((), assembly_opts, integral_types)
with dNdu.dat.vec as vec:
J.petscmat.setDiagonal(vec)
return J
```

Note that the above implementation first constructs the Jacobian matrix \(J\) before populating its diagonal.
This can be achieved using the *_matrix_builder* external operator’s helper function. The variational
problem can now be solved

```
u = Function(V)
v = TestFunction(V)
bcs = DirichletBC(V, 0, "on_boundary")
N = TranslationOperator(u, f, function_space=V)
F = (inner(grad(u), grad(v)) + inner(N, v)) * dx
solve(F == 0, u, bcs=bcs)
```

### Matrix-free¶

In many cases, computing the Jacobian of the residual form is not appropriate, or even not possible.
Instead, one may want to use matrix-free methods to solve the PDE problem of interest. In that case,
the Jacobian of \(F\) won’t be assembled. Instead, only the action of the Jacobian will be used. As a
consequence, our external operator subclass will need to be equipped with an implementation stating how
the action of the Jacobian of \(N\) on a given `Function`

\(w\) can be assembled, i.e. how to
compute \(\frac{\partial N(u, f; w, v^{*})}{\partial u}\). In this case, this implementation should simply
return \(w\) as the Jacobian is the identity matrix.

```
class TranslationOperator(AbstractExternalOperator):
def __init__(self, *operands, function_space, **kwargs):
AbstractExternalOperator.__init__(self, *operands, function_space=function_space, **kwargs)
@assemble_method(0, (0,))
def assemble_N(self, *args, **kwargs):
"""Evaluate N"""
u, f = self.ufl_operands
N = assemble(u - f)
return N
@assemble_method((1, 0), (0, None))
def assemble_Jacobian_action(self, *args, **kwargs):
"""Evaluate the action of the Jacobian"""
w = self.argument_slots()[-1]
return w
```

The arguments of an external operator can be obtained via *.argument_slots()*. This will return
all the arguments of the external operator, independently of whether they are
`Argument`

/ `Coargument`

or `Function`

/ `Cofunction`

. If you only want
the unknown arguments, for example to determine the arity of the external operator,
you can instead use *.arguments()*. We can now solve the variational problem using any matrix-free method:

```
u = Function(V)
N = TranslationOperator(u, f, function_space=V)
F = (inner(grad(u), grad(v)) + inner(N, v)) * dx
solve(F == 0, u, bcs=bcs, solver_parameters={"mat_type": "matfree",
"ksp_type": "cg",
"pc_type": "none"})
```

### Inverse problems¶

External operators can also be embedded in PDE-constrained optimisation problems. For instance, let’s consider the following inverse problem driven by the elliptic PDE previously introduced:

subject to

where \(u(\cdot)\) is the solution operator, \(u^{obs}\) refers to some observables, \(\mathcal{R}\) is a regularisation term, and \(f_{0}\) is a guess. In our case, we consider a general Tikhonov regularization, that is:

The above regulariser can be used to incorporate prior knowledge into the problem via some guess \(f_{0}\).
The regulariser \(\mathcal{R}\) can also be defined with the *TranslationOperator* we introduced in the
previous section, but this time with the operands \(f\) and \(f_{0}\).

We use the *firedrake.adjoint* package to automatically compute the gradient of the cost function \(J\)
for the optimisation. Evaluating the functional \(J\) requires evaluating the external operator
\(\mathcal{R}(f, f_{0}; v^{*})\). On the other hand, computing the gradient of \(J\) using the adjoint
method involves evaluating the action of the Jacobian adjoint of the external operator, i.e.
\(\frac{\partial \mathcal{R}(f, f_{0}; y, \hat{f})}{\partial f}\) for a given \(y \in V^{*}\) and
\(\forall \hat{f} \in V\). We already implemented the evaluation method for the *TranslationOperator*.
We now need to add the method for \(\frac{\partial \mathcal{R}(f, f_{0}; y, \hat{f})}{\partial f}\), which in
this case is trivial as the Jacobian is the identity matrix

```
class TranslationOperator(AbstractExternalOperator):
def __init__(self, *operands, function_space, **kwargs):
AbstractExternalOperator.__init__(self, *operands, function_space=function_space, **kwargs)
@assemble_method(0, (0,))
def assemble_R(self, *args, **kwargs):
"""Evaluate the regulariser R"""
f, f0 = self.ufl_operands
N = assemble(f - f0)
return N
@assemble_method((1, 0), (None, 0))
def assemble_Jacobian_adjoint_action(self, *args, **kwargs):
"""Evaluate the action of the Hermitian transpose of the Jacobian of R"""
y, _ = self.argument_slots()
return y
```

We define the observables by adding noise to the exact solution of the PDE associated with a given rhs \(f_{exact}\):

We can now solve the PDE-constrained optimisation problem using the *firedrake.adjoint* package.
For this, we employ the *BFGS* algorithm:

```
R = partial(TranslationOperator, function_space=V)
def J(f):
F = (inner(grad(u), grad(v)) + inner(u, v) - inner(f, v)) * dx
solve(F == 0, u, bcs=bcs)
return assemble(0.5 * (u - u_obs) ** 2 * dx + 0.5 * alpha * R(f, f_0) ** 2 * dx)
c = Control(f)
Jhat = ReducedFunctional(J(f), c)
f_opt = minimize(Jhat, method= "BFGS")
```

The above code will execute the optimisation and call the external operator subclass every time the functional \(J\) or its gradient is evaluated.

PDE systems implemented in Firedrake can be specified with one or more external operators.
External operators can also be embedded with each other as long as the function spaces match.
In the above example, we only used the *TranslationOperator* to define the regulariser. However, we
could also have used it inside the PDE as we did it in the previous section. In that case, we would
end up with two *TranslationOperator* s, one in the cost function with operands \(f\) and \(f_{0}\), and
one in the PDE with operands \(u\) and \(f\). The external operator subclass would then need to be equipped with
the methods to evaluate \(J\), which implies solving the PDE, and computing its derivative, which involves
solving the adjoint equation.