Extract linear combination from ufl-expression

[jorgensd]

It is often common to extract linear combinations $\sum_i c_i u_i$ or a set of scalar weights $c_i$ and function $u_i\in V$ for quick assignment without having to call interpolation or projection, ref #486.

AI assistance

The code and documentation was generated by prompting Gemini (August 2026) based on my own initial implementations without a DAGTraverser.
I reviewed, tested, and take responsibility for the final contribution.

Deprecated example after discussions in this PR

Initially to catch scalar valued real elements, one would have to overload the dag-traverser. This is no longer true, as we use ufl.checks.is_scalar_constant_expression to have a common base for both Firedrake and FEniCS.

The old example can be found below:

import ufl
from ufl.algorithms.extract_linear_combination import LinearCombinationExtractor
from functools import singledispatchmethod

class CustomLinearCombinationExtractor(LinearCombinationExtractor):
    """Custom extractor to handle linear combinations of UFL expressions."""

    @singledispatchmethod
    def process(self, o: ufl.classes.Expr, **kwargs):
        """Fallback for any unsupported node types."""
        return super().process(o, **kwargs)

    @process.register(dolfinx.fem.Function)
    def _(self, o, **kwargs):
        if o.function_space.ufl_element().is_real and o.ufl_shape == ():
            return o
        return [(1.0, o)]


def extract_linear_combination(
    expr: ufl.core.expr.Expr,
) -> list[tuple[ufl.core.expr.Expr, ufl.classes.Coefficient]]:
    """Wrapper to initialize traverser and extract linear combinations.

    Returns:
        A list of tuples where the first element is the UFL expression of the
        weight, and the second element is the base UFL Coefficient (spatial function).
    """
    extractor = CustomLinearCombinationExtractor()
    final_result = extractor(expr)

    if not isinstance(final_result, list):
        raise ValueError("Expression evaluated to a pure scalar, no spatial functions found.")

    return final_result

[pbrubeck]

Wouldn't dolfinx.fem.Function inherit from ufl.Coefficient? Why do you need to subclass the DAGTraverser in the example from the description?

[jorgensd]

Wouldn't dolfinx.fem.Function inherit from ufl.Coefficient? Why do you need to subclass the DAGTraverser in the example from the description?

I would need to sub-class it to detect a Real-valued finite element, which acts as a differentiable constant, and only has DOLFINx specific data, such as ufl_element().is_real which is not part of the UFL standard.

[pbrubeck]

I would need to sub-class it to detect a Real-valued finite element, which acts as a differentiable constant, and only has DOLFINx specific data, such as ufl_element().is_real which is not part of the UFL standard.

But the UFL standard can identify piecewise and global constants. I think we should try to do that here, since Firedrake also has a Real element

[jorgensd]

I would need to sub-class it to detect a Real-valued finite element, which acts as a differentiable constant, and only has DOLFINx specific data, such as ufl_element().is_real which is not part of the UFL standard.

But the UFL standard can identify piecewise and global constants. I think we should try to do that here, since Firedrake also has a Real element

Great point! I've adjusted the code, and it passes my test suite in DOLFINx with real element coefficients.

[pbrubeck]

Why aren't we supporting this case?

[jorgensd]

In the same way that we are not supporting ufl.core.expr.Expr as a general term.
The only two classes that for now uses BaseCoefficient is Cofunction and Function, which have specializations for in the code. This would be a good check if someone implements a new class based on BaseCoefficient, which they would have to check if the current implementation is correct.

[pbrubeck]

Wouldn't we want linear combinations to just work on newly implemented subclasses of BaseCoefficient?

[jorgensd]

I guess it depends on how it should be handled.
Currently there is a split between how Coefficient and CoFunction is handled, as coefficient checks for
ufl.checks.is_scalar_constant_expression(o) to factor it out of a linear combination, while this is not the case for the CoFunction, as it is always a scalar_constant_expression.

[pbrubeck]

I don't see why would one want distinct treatment for Coefficient, Cofunction and Matrix.

[pbrubeck]

What are the assumptions here? Are we allowing affine expressions, i.e. sum_i a_i v_i + Constant?

[jorgensd]

The assumption here is that if we have (sum a_i v_i), then in some cases a_i might be a scalar valued real space coefficient, and should therefore be returned as a float rather than the tuple (1, a_i).

We are not allowing for a_i v_i + constant in this PR.

[pbrubeck]

Why is this unique to Coefficient? We should give no special treatment to Coefficient vs Cofunction