Implement linear combination assignment

[jorgensd]

Implementation of Linear Combination Assignments in FEniCSx

The code was written by me and then refined by using Gemini (August 2026).
The following summary was generated with the assistance of Gemini (August 2026).

This document outlines the practical usage and mathematical foundation for evaluating linear combinations of finite element functions, parsing their components, and deriving their corresponding adjoints within the dolfinx-adjoint framework.

Practical Usage in dolfinx-adjoint

In dolfinx-adjoint, this mathematical optimization is exposed through the assign(linear_expression, z) function. Instead of treating a linear combination as a generic UFL expression that requires a costly $L^2$-projection or spatial interpolation, the framework intercepts the assignment. By recognizing that the operation is purely linear and operates within the same function space, it translates the continuous UFL expression into direct, highly efficient vector operations on the underlying degrees of freedom (DOFs).

Parsing the Linear Expression

Before the assignment can be executed on the DOF arrays, the continuous UFL expression must be decomposed. The framework recursively traverses the UFL abstract syntax tree (using extract_linear_combination) to break the expression down into a standard sum of weighted functions:

$$\text{expr}(x) = \sum_{k} c_k v_k(x)$$

During this traversal, the algorithm enforces a strict separation for each term:

1. The Spatial Function ($v_k$): It isolates the single standard finite element function in the term. Crucially, every $v_k$ discovered must share the exact same function space $V$ as the target variable $z$. If multiple spatial functions are multiplied together (a non-linear operation like $v_1 \times v_2$), the parser raises an error, as this cannot be evaluated via direct DOF addition.
2. The Scalar Weight ($c_k$): All other components in the term are folded into the coefficient $c_k$. This weight is not limited to a single float; it can be a FEniCSx Constant, a Real space function (a globally uniform constant stored as a finite element function), or a complex product and division of multiple scalar values.

Once parsed, the expression is reduced to a list of (weight, function) pairs, which can then be directly applied to the arrays.

---

The Forward Pass: Interpolation as an $\mathbb{R}^N$ Identity

In the finite element method, a function $u$ belonging to a function space $V$ is represented as a linear combination of basis functions $\phi_i(x)$:

$$u(x) = \sum_{i=1}^N u_i \phi_i(x)$$

where $u_i$ are the discrete degrees of freedom. The interpolation operator $\Pi$ maps a continuous function $f(x)$ into $V$ using the dual basis of nodal functionals $l_i$:

$$\Pi(f) = \sum_{i=1}^N l_i(f) \phi_i(x)$$

By definition, the basis functions and nodal functionals satisfy the biorthogonality property:

$$l_i(\phi_j) = \delta_{ij}$$

When $f(x)$ is strictly a linear combination of other finite element functions defined on the exact same function space $V$ (as guaranteed by the parsing step), it can be written as:

$$f(x) = \sum_{j=1}^N f_j \phi_j(x)$$

Applying the nodal functional $l_i$ to this $f(x)$ yields:

$$l_i(f) = l_i\left( \sum_{j=1}^N f_j \phi_j \right) = \sum_{j=1}^N f_j l_i(\phi_j) = \sum_{j=1}^N f_j \delta_{ij} = f_i$$

Because $l_i(f) = f_i$, the continuous interpolation operator reduces directly to an identity mapping in $\mathbb{R}^N$. This mathematically proves that assigning a linear combination of functions from the same space does not require solving a projection system; it is perfectly exact to apply the weights $c_k$ and functions $v_k$ directly to the underlying DOF arrays.

---

The Adjoint Pass: Deriving Sensitivities

In the algorithmic differentiation tape, let $J$ be the objective functional. For a forward assignment $u = \text{expr}(v, r)$ where $v$ is a spatial function and $r$ is a global constant weight, we receive the incoming adjoint sensitivities $\bar{u}_i = \frac{\partial J}{\partial u_i}$.

Adjoint of a Spatial Function Term ($v$)

Because the assignment maps DOFs pointwise (i.e., $u_i$ depends only on $v_i$), the chain rule for the derivative with respect to the spatial function $v$ is strictly local to each DOF:

$$\bar{v}_i = \frac{\partial J}{\partial v_i} = \frac{\partial J}{\partial u_i} \frac{\partial u_i}{\partial v_i} = \bar{u}_i \frac{\partial \text{expr}_i}{\partial v_i}$$

In the UFL-based implementation, this is evaluated by taking the symbolic derivative of the expression with respect to $v$, in the direction of the adjoint vector $\bar{u}$. Because the underlying operation is a linear combination, applying the directional derivative automatically multiplies the local coefficient by the incoming adjoint sensitivity.

Adjoint of a Constant Term ($r$)

When assigning a constant $r$ (such as a function in the Real space) to a spatial function $u$, the scalar $r$ is broadcasted to all $N$ degrees of freedom in the mesh. Consequently, a change in $r$ affects every $u_i$. The chain rule dictates that the adjoint derivative must sum the contributions from all DOFs:

$$\bar{r} = \frac{\partial J}{\partial r} = \sum_{i=1}^N \frac{\partial J}{\partial u_i} \frac{\partial u_i}{\partial r} = \sum_{i=1}^N \bar{u}_i \frac{\partial \text{expr}_i}{\partial r}$$

To implement this, the code performs the following steps:

1. Computes the UFL derivative of the expression with respect to $r$ in the direction of $1.0$ (a unit perturbation in the Real space).
2. Evaluates this derivative vector $\frac{\partial \text{expr}}{\partial r}$ at the degrees of freedom of the target space.
3. Computes the standard $\mathbb{R}^N$ inner product between the derivative vector and the adjoint vector $\bar{u}$.

This inner product exactly executes the required summation over $i$, collapsing the $N$-dimensional spatial sensitivities down to a single global scalar adjoint $\bar{r}$.

[finsberg]

Very nice feature! Added a few comments.
Should probably also add some tests that checks error handling when functions for different function spaces are used and in general a bit more exhausting unit testing of the utility functions in utils.py