From ee4e99113e40423b15c8c8f1b7d7449c47a5e06c Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Thu, 20 Aug 2026 12:17:13 +0000
Subject: [PATCH 1/9] Add linear combination extractor. Addresses
 https://github.com/FEniCS/ufl/issues/486 Written with the help of Gemini
 (August 2026). I take full responsibility for the correctness and testing of
 the code.

---
 test/test_extract_linear_combination.py      | 141 +++++++++++++++++
 ufl/algorithms/extract_linear_combination.py | 151 +++++++++++++++++++
 2 files changed, 292 insertions(+)
 create mode 100644 test/test_extract_linear_combination.py
 create mode 100644 ufl/algorithms/extract_linear_combination.py

diff --git a/test/test_extract_linear_combination.py b/test/test_extract_linear_combination.py
new file mode 100644
index 000000000..f500cacbe
--- /dev/null
+++ b/test/test_extract_linear_combination.py
@@ -0,0 +1,141 @@
+import pytest
+from utils import LagrangeElement
+
+import ufl
+from ufl.algorithms.extract_linear_combination import extract_linear_combination
+
+
+@pytest.fixture
+def domain():
+    return ufl.Mesh(LagrangeElement(ufl.triangle, 1, (2,)))
+
+
+@pytest.fixture
+def V(domain):
+    """Standard spatial function space."""
+    el = LagrangeElement(domain.ufl_cell(), 3)
+    return ufl.FunctionSpace(domain, el)
+
+
+@pytest.fixture
+def V_vec(domain):
+    """Vector function space for indexed tests."""
+    return ufl.FunctionSpace(domain, LagrangeElement(domain.ufl_cell(), 2, (3,)))
+
+
+def test_valid_linear_combinations(V):
+    """Test standard valid combinations.
+
+    Sum, Product, Division, Negative, and float/int terminals.
+    """
+    u = ufl.Coefficient(V)
+    v = ufl.Coefficient(V)
+
+    # Includes: IntValue (3), FloatValue (1.5), Sum, Negative (-v/4.0), Product, Division
+    expr = 2.0 * u - v / 4.0 + ufl.as_ufl(3) * u - ufl.classes.FloatValue(1.5) * v
+
+    res = extract_linear_combination(expr)
+
+    assert len(res) == 4
+    # The DAG traversal evaluates leaves to roots, maintaining term extraction
+    # Order might depend slightly on UFL's internal DAG sorting, but usually
+    # follows algebraic definition
+    assert (2.0, u) in res
+    assert (-0.25, v) in res
+    assert (3.0, u) in res
+    assert (-1.5, v) in res
+
+
+def test_scalars_and_powers(V, domain):
+    """Test evaluating constants, real functions, and powers as scalar weights."""
+    u = ufl.Coefficient(V)
+
+    c = ufl.Constant(domain)
+    d = ufl.Constant(domain)
+    expr = (d**2) * u + (c + d**2) * u
+
+    res = extract_linear_combination(expr)
+
+    assert len(res) == 2
+    assert (d**2, u) in res
+    assert (c + d**2, u) in res
+
+
+def test_ufl_zero_node(V):
+    """Test that UFL Zero nodes are handled correctly (usually evaluates to an empty list)."""
+    u = ufl.Coefficient(V)
+
+    # Create an explicit Zero node with the shape of u
+    z = ufl.classes.Zero(u.ufl_shape, u.ufl_free_indices, u.ufl_index_dimensions)
+
+    expr = u + z
+    res = extract_linear_combination(expr)
+
+    assert len(res) == 1
+    assert res[0] == (1.0, u)
+
+
+def test_nonlinear_errors(V):
+    """Test that non-linear operations strictly raise ValueErrors."""
+    u = ufl.Coefficient(V)
+    v = ufl.Coefficient(V)
+
+    with pytest.raises(ValueError, match="product of two spatial functions"):
+        extract_linear_combination(u * v)
+
+    with pytest.raises(ValueError, match="division by a spatial function"):
+        extract_linear_combination(u / v)
+
+    with pytest.raises(ValueError, match="power involving a spatial function"):
+        extract_linear_combination(u**2)
+
+
+def test_invalid_additions(V, domain):
+    """Test that adding a scalar to a spatial function is caught."""
+    u = ufl.Coefficient(V)
+    c = ufl.Constant(domain)
+
+    with pytest.raises(
+        ValueError, match="Cannot directly add a raw scalar expression to a spatial function"
+    ):
+        extract_linear_combination(u + c)
+
+    with pytest.raises(
+        ValueError, match="Cannot directly add a raw scalar expression to a spatial function"
+    ):
+        extract_linear_combination(u + 5.0)
+
+
+def test_pure_scalar_error(domain):
+    """Test that evaluating an expression with NO spatial functions raises an error."""
+    c = ufl.Constant(domain)
+    r_func = ufl.Constant(domain)
+
+    expr = c * 5.0 + (r_func**2)
+
+    with pytest.raises(ValueError, match="Expression evaluated to a pure scalar"):
+        extract_linear_combination(expr)
+
+
+def test_forbidden_operations(V_vec):
+    """Test that indexed vectors and component tensors trigger a NotImplementedError."""
+    u = ufl.Coefficient(V_vec)
+
+    with pytest.raises(
+        NotImplementedError, match="Direct array assignment of indexed vector components"
+    ):
+        extract_linear_combination(u[0])
+
+
+def test_negative(V, domain):
+    """Test that negative operations are handled correctly."""
+    r_func = ufl.Constant(domain)
+    u = ufl.Coefficient(V)
+    expr = -r_func * u + 2 * u
+
+    res = extract_linear_combination(expr)
+
+    assert len(res) == 2
+
+    assert (-r_func, u) in res
+    assert (2.0, u) in res
diff --git a/ufl/algorithms/extract_linear_combination.py b/ufl/algorithms/extract_linear_combination.py
new file mode 100644
index 000000000..3e2a2d5c3
--- /dev/null
+++ b/ufl/algorithms/extract_linear_combination.py
@@ -0,0 +1,151 @@
+import ufl
+from ufl.corealg.dag_traverser import DAGTraverser
+from functools import singledispatchmethod
+
+
+class LinearCombinationExtractor(DAGTraverser):
+    """Bottom-up DAG traverser for extracting linear combinations.
+
+    To process an arbitrary mathematical expression, this traverser categorizes
+    every node in the DAG into one of two states, returning different types for each:
+
+    1. Scalar Weights (Returns: {py:class}`ufl.core.expr.Expr`)
+       If a node and all its children represent a global scalar value (e.g.,
+       {py:class}`ufl.FloatValue`, {py:class}`ufl.Constant`), the traverser
+       propagates the actual UFL expression upwards. It does not evaluate them
+       to Python floats, preserving the full UFL AST of the constants.
+
+    2. Spatial Fields (Returns: `list[tuple[ufl.core.expr.Expr, ufl.Coefficient]]`)
+       If a node contains spatial functions (standard Coefficients), it must
+       maintain the strict algebraic structure of a linear combination. Therefore,
+       it returns a list of `(weight, function)` tuples, where `weight` is the
+       accumulated UFL expression and `function` is the base spatial field.
+
+    By strictly distinguishing between the two (checking `isinstance(..., list)`),
+    the traverser can safely apply algebraic rules (e.g., multiplying a list by a
+    scalar weight expression distributes the weight) and instantly catch illegal
+    non-linear operations (e.g., attempting to multiply two lists together).
+    """
+
+    def __init__(self, **kwargs):
+        # Disable compression to avoid hashing unhashable return types (like lists)
+        # while preserving the `_visited_cache` memoization.
+        kwargs["compress"] = False
+        super().__init__(**kwargs)
+
+    @singledispatchmethod
+    def process(self, o: ufl.classes.Expr, **kwargs):
+        """Fallback for any unsupported node types."""
+        raise ValueError(f"Unsupported UFL node type for linear combinations: {type(o)}")
+
+    # ---------------------------------------------------------
+    # 1. Terminals (Leaves) - No children to evaluate
+    # ---------------------------------------------------------
+    @process.register(ufl.classes.IntValue)
+    @process.register(ufl.classes.FloatValue)
+    @process.register(ufl.classes.ScalarValue)
+    def _(self, o, **kwargs):
+        # Return the UFL expression itself
+        return o
+
+    @process.register(ufl.classes.Zero)
+    def _(self, o, **kwargs):
+        return o if o.ufl_shape == () else []
+
+    @process.register(ufl.Constant)
+    def _(self, o, **kwargs):
+        if o.ufl_shape == ():
+            return o
+        raise ValueError(f"Only scalar constants are supported, got shape {o.ufl_shape}")
+
+    @process.register(ufl.classes.Coefficient)
+    def _(self, o, **kwargs):
+        return [(ufl.as_ufl(1.0), o)]
+
+    # ---------------------------------------------------------
+    # 2. Operators - Use @postorder to evaluate operands first
+    # ---------------------------------------------------------
+    @process.register(ufl.classes.Sum)
+    @DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        # If no operands are lists, this is a pure scalar addition.
+        # We construct a new UFL expression by safely summing them.
+        if all(not isinstance(op, list) for op in operands):
+            res = operands[0]
+            for op in operands[1:]:
+                res = res + op
+            return res
+
+        # Otherwise, accumulate the spatial functions
+        res = []
+        for op_res in operands:
+            if isinstance(op_res, list):
+                res.extend(op_res)
+            else:
+                raise ValueError(
+                    "Cannot directly add a raw scalar expression to a spatial function."
+                )
+        return res
+
+    @process.register(ufl.classes.Product)
+    @DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        op1_res, op2_res = operands
+        is_list1 = isinstance(op1_res, list)
+        is_list2 = isinstance(op2_res, list)
+
+        if not is_list1 and not is_list2:
+            return op1_res * op2_res  # UFL operator overloading takes over
+        elif not is_list1 and is_list2:
+            return [(op1_res * w, f) for w, f in op2_res]
+        elif not is_list2 and is_list1:
+            return [(op2_res * w, f) for w, f in op1_res]
+
+        raise ValueError("Non-linear expression detected: product of two spatial functions.")
+
+    @process.register(ufl.classes.Division)
+    @DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        num_res, den_res = operands
+        if isinstance(den_res, list):
+            raise ValueError("Non-linear expression detected: division by a spatial function.")
+
+        if not isinstance(num_res, list):
+            return num_res / den_res
+        return [(w / den_res, f) for w, f in num_res]
+
+    @process.register(ufl.classes.Power)
+    @DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        base_res, exp_res = operands
+        if isinstance(base_res, list) or isinstance(exp_res, list):
+            raise ValueError("Non-linear expression detected: power involving a spatial function.")
+        return base_res**exp_res
+
+    # ---------------------------------------------------------
+    # 3. Forbidden Operations
+    # ---------------------------------------------------------
+    @process.register(ufl.classes.Indexed)
+    @process.register(ufl.classes.ComponentTensor)
+    def _(self, o, **kwargs):
+        raise NotImplementedError(
+            "Direct array assignment of indexed vector components is not supported."
+        )
+
+
+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 = LinearCombinationExtractor()
+    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

From 18c17055d2ed2a4ee8a9f291e1e05315b28086cf Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Thu, 20 Aug 2026 12:25:09 +0000
Subject: [PATCH 2/9] Add some more documentation

---
 ufl/algorithms/extract_linear_combination.py | 29 ++++++++++++++++----
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/ufl/algorithms/extract_linear_combination.py b/ufl/algorithms/extract_linear_combination.py
index 3e2a2d5c3..eef463c49 100644
--- a/ufl/algorithms/extract_linear_combination.py
+++ b/ufl/algorithms/extract_linear_combination.py
@@ -1,6 +1,15 @@
+"""Tools to extract linear combinations from UFL expressions."""
+
+# Copyright (C) 2026 Jørgen S. Dokken
+#
+# This file is part of UFL (https://www.fenicsproject.org)
+#
+# SPDX-License-Identifier:    LGPL-3.0-or-later
+
+from functools import singledispatchmethod
+
 import ufl
 from ufl.corealg.dag_traverser import DAGTraverser
-from functools import singledispatchmethod
 
 
 class LinearCombinationExtractor(DAGTraverser):
@@ -28,8 +37,11 @@ class LinearCombinationExtractor(DAGTraverser):
     """
 
     def __init__(self, **kwargs):
-        # Disable compression to avoid hashing unhashable return types (like lists)
-        # while preserving the `_visited_cache` memoization.
+        """Initialize LinearCombinationExtractor with memoization and no compression.
+
+        Compression is disabled to avoid hashing unhashable return types (like lists)
+        while preserving the `_visited_cache` memoization.
+        """
         kwargs["compress"] = False
         super().__init__(**kwargs)
 
@@ -91,17 +103,22 @@ def _(self, o, *operands, **kwargs):
     @DAGTraverser.postorder
     def _(self, o, *operands, **kwargs):
         op1_res, op2_res = operands
+        # Each of the operands are either a scalar UFL expression (float, Constant, etc.)
+        # or a list of (weight, function) tuples.
+        # The following cases are possible:
+        # 1. Both operands are scalars: return the product of the two UFL expressions.
+        # 2. One operand is a scalar, the other is a list: distribute the scalar across the list.
+        # 3. Both operands are lists: this is a non-linear operation and should raise an error.
         is_list1 = isinstance(op1_res, list)
         is_list2 = isinstance(op2_res, list)
-
         if not is_list1 and not is_list2:
             return op1_res * op2_res  # UFL operator overloading takes over
         elif not is_list1 and is_list2:
             return [(op1_res * w, f) for w, f in op2_res]
         elif not is_list2 and is_list1:
             return [(op2_res * w, f) for w, f in op1_res]
-
-        raise ValueError("Non-linear expression detected: product of two spatial functions.")
+        else:
+            raise ValueError("Non-linear expression detected: product of two spatial functions.")
 
     @process.register(ufl.classes.Division)
     @DAGTraverser.postorder

From c444ee0d3f8c936fbf11c890b915080a2a53255a Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Thu, 20 Aug 2026 13:59:15 +0000
Subject: [PATCH 3/9] Extend to some more operators

---
 test/test_extract_linear_combination.py      | 43 ++++++++++++++++++++
 ufl/algorithms/extract_linear_combination.py | 36 ++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/test/test_extract_linear_combination.py b/test/test_extract_linear_combination.py
index f500cacbe..101fffc41 100644
--- a/test/test_extract_linear_combination.py
+++ b/test/test_extract_linear_combination.py
@@ -139,3 +139,46 @@ def test_negative(V, domain):
 
     assert (-r_func, u) in res
     assert (2.0, u) in res
+
+
+def test_matrix_linear_combination(V, domain):
+    """Test linear combinations involving ufl.Matrix."""
+    # Matrix requires a row space and column space
+    A = ufl.Matrix(V, V)
+    c = ufl.Constant(domain)
+    expr = 2.0 * A + (0.3 + c**2) * A
+
+    res = extract_linear_combination(expr)
+
+    assert len(res) == 2
+    # NOTE: Matrices store scalar weights under `weights`
+    assert (2.0, A) in res
+    assert (0.3 + c**2, A) in res
+
+
+def test_cofunction_linear_combination(V):
+    """Test linear combinations involving ufl.Cofunction."""
+    # Cofunction requires a dual space (it will raise an error if given a primal space)
+    V_dual = V.dual()
+    c = ufl.Cofunction(V_dual)
+
+    # NOTE: Cofunctions can only be multiplied by integers
+    expr = -2 * c + 5 * c
+
+    res = extract_linear_combination(expr)
+
+    assert len(res) == 2
+    assert (-2.0, c) in res
+    assert (5.0, c) in res
+
+
+def test_matrix_nonlinear_error(V):
+    """Test that matrices are protected by the same non-linear guardrails."""
+    A = ufl.Matrix(V, V)
+    u = ufl.Coefficient(V)
+
+    with pytest.raises(
+        ValueError, match=r"Non-linear expression detected: product of two spatial functions."
+    ):
+        # Cannot multiply a matrix by a coefficient algebraically in this block
+        extract_linear_combination(A * u)
diff --git a/ufl/algorithms/extract_linear_combination.py b/ufl/algorithms/extract_linear_combination.py
index eef463c49..edd57a31e 100644
--- a/ufl/algorithms/extract_linear_combination.py
+++ b/ufl/algorithms/extract_linear_combination.py
@@ -99,6 +99,35 @@ def _(self, o, *operands, **kwargs):
                 )
         return res
 
+    @process.register(ufl.Action)
+    def _(self, o, **kwargs):
+        # An Action node represents a matrix-vector product (e.g., A * u).
+        # This cannot be reduced to a simple algebraic linear combination of arrays.
+        raise ValueError("Non-linear expression detected: product of two spatial functions.")
+
+    @process.register(ufl.classes.FormSum)
+    @process.register(ufl.form.FormSum)
+    def _(self, o, **kwargs):
+        res = []
+        components = o.components()
+        weights = o.weights()
+        for weight, comp in zip(weights, components):
+            # Evaluate the base component (e.g., Matrix or Cofunction)
+            comp_res = self(comp, **kwargs)
+
+            # Evaluate the weight (in case it contains sub-expressions)
+            w_res = self(weight, **kwargs) if isinstance(weight, ufl.classes.Expr) else weight
+
+            if isinstance(comp_res, list):
+                # Distribute this FormSum weight into the component's linear combination
+                res.extend([(w_res * w, f) for w, f in comp_res])
+            else:
+                raise ValueError(
+                    "Cannot directly add a raw scalar expression to a spatial function."
+                )
+
+        return res
+
     @process.register(ufl.classes.Product)
     @DAGTraverser.postorder
     def _(self, o, *operands, **kwargs):
@@ -149,6 +178,13 @@ def _(self, o, **kwargs):
             "Direct array assignment of indexed vector components is not supported."
         )
 
+    @process.register(ufl.Cofunction)
+    @process.register(ufl.Matrix)
+    def _(self, o, **kwargs):
+        # Cofunctions and Matrices are symbolic algebraic terminals.
+        # They act purely as spatial fields in linear combinations.
+        return [(ufl.as_ufl(1.0), o)]
+
 
 def extract_linear_combination(
     expr: ufl.core.expr.Expr,

From 50b8ccd2c6aac968ba17affdccac4700117e24fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= <dokken92@gmail.com>
Date: Thu, 20 Aug 2026 16:14:22 +0200
Subject: [PATCH 4/9] Add copyright header
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>
---
 test/test_extract_linear_combination.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/test/test_extract_linear_combination.py b/test/test_extract_linear_combination.py
index 101fffc41..8ae529fd7 100644
--- a/test/test_extract_linear_combination.py
+++ b/test/test_extract_linear_combination.py
@@ -1,3 +1,9 @@
+# Copyright (C) 2026 Jørgen S. Dokken
+#
+# This file is part of UFL (https://www.fenicsproject.org)
+#
+# SPDX-License-Identifier:    LGPL-3.0-or-later
+
 import pytest
 from utils import LagrangeElement
 

From 3d509cc68aaea30743ebaefd734c7908c8621a37 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 06:42:50 +0000
Subject: [PATCH 5/9] Generalize cofunction test and remove wrong comment

---
 test/test_extract_linear_combination.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/test/test_extract_linear_combination.py b/test/test_extract_linear_combination.py
index 8ae529fd7..4170616b6 100644
--- a/test/test_extract_linear_combination.py
+++ b/test/test_extract_linear_combination.py
@@ -167,14 +167,12 @@ def test_cofunction_linear_combination(V):
     # Cofunction requires a dual space (it will raise an error if given a primal space)
     V_dual = V.dual()
     c = ufl.Cofunction(V_dual)
-
-    # NOTE: Cofunctions can only be multiplied by integers
-    expr = -2 * c + 5 * c
-
+    d = ufl.Constant(V.ufl_domain())
+    expr = -4 * d * c + 5.0 * c
     res = extract_linear_combination(expr)
 
     assert len(res) == 2
-    assert (-2.0, c) in res
+    assert (-4 * d, c) in res
     assert (5.0, c) in res
 
 

From 2f489ef1d5ac3c7ec9bd5209428579d3cf8a9bff Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 06:57:10 +0000
Subject: [PATCH 6/9] Propagate processing of baseform through dagtraverser

---
 ufl/algorithms/apply_coefficient_split.py    |  5 +-
 ufl/algorithms/apply_derivatives.py          | 78 ++++++++++----------
 ufl/algorithms/cancel_jacobian_products.py   |  9 ++-
 ufl/algorithms/extract_linear_combination.py | 17 ++---
 ufl/corealg/dag_traverser.py                 |  8 +-
 5 files changed, 61 insertions(+), 56 deletions(-)

diff --git a/ufl/algorithms/apply_coefficient_split.py b/ufl/algorithms/apply_coefficient_split.py
index 0a221fb75..0bc59cd43 100644
--- a/ufl/algorithms/apply_coefficient_split.py
+++ b/ufl/algorithms/apply_coefficient_split.py
@@ -26,6 +26,7 @@
 )
 from ufl.core.multiindex import indices
 from ufl.corealg.dag_traverser import DAGTraverser
+from ufl.form import BaseForm
 from ufl.tensors import as_tensor
 
 
@@ -36,8 +37,8 @@ def __init__(
         self,
         coefficient_split: dict,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise.
 
diff --git a/ufl/algorithms/apply_derivatives.py b/ufl/algorithms/apply_derivatives.py
index 19516e2a0..e1de2ee06 100644
--- a/ufl/algorithms/apply_derivatives.py
+++ b/ufl/algorithms/apply_derivatives.py
@@ -158,8 +158,8 @@ def __init__(
         self,
         var_shape: tuple,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
@@ -199,11 +199,11 @@ def independent_operator(self, o):
     # --- Error checking for missing handlers and unexpected types
 
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -723,8 +723,8 @@ def __init__(
         self,
         geometric_dimension: int,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(
@@ -738,11 +738,11 @@ def __init__(
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -1001,8 +1001,8 @@ def __init__(
         self,
         topological_dimension: int,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(
@@ -1016,7 +1016,7 @@ def __init__(
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
@@ -1115,8 +1115,8 @@ def __init__(
         self,
         var: Expr,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(
@@ -1160,11 +1160,11 @@ def _make_identity(self, sh):
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -1272,8 +1272,8 @@ def __init__(
         arguments: ExprList,
         coefficient_derivatives: ExprMapping,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(
@@ -1306,11 +1306,11 @@ def __init__(
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -1696,8 +1696,8 @@ def __init__(
         coefficient_derivatives: ExprMapping,
         outer_base_form_op: Expr,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(
@@ -1713,11 +1713,11 @@ def __init__(
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -1778,8 +1778,8 @@ class DerivativeRuleDispatcher(DAGTraverser):
     def __init__(
         self,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
@@ -1788,16 +1788,18 @@ def __init__(
         self.pending_operations = ()
         # Create DAGTraverser caches.
         self._dag_traverser_cache: dict[
-            tuple[type, Expr] | tuple[type, Expr, Expr, Expr] | tuple[type, Expr, Expr, Expr, Expr],
+            tuple[type, Expr | BaseForm]
+            | tuple[type, Expr | BaseForm, Expr | BaseForm, Expr | BaseForm]
+            | tuple[type, Expr | BaseForm, Expr | BaseForm, Expr | BaseForm, Expr | BaseForm],
             DAGTraverser,
         ] = {}
 
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -2095,8 +2097,8 @@ def __init__(
         arguments: ExprList,
         coefficient_derivatives: ExprMapping,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(
@@ -2122,11 +2124,11 @@ def __init__(
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
@@ -2229,19 +2231,21 @@ class CoordinateDerivativeRuleDispatcher(DAGTraverser):
     def __init__(
         self,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
-        self._dag_traverser_cache: dict[tuple[type, Expr, Expr, Expr], DAGTraverser] = {}
+        self._dag_traverser_cache: dict[
+            tuple[type, Expr | BaseForm, Expr | BaseForm, Expr | BaseForm], DAGTraverser
+        ] = {}
 
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``.
 
         Args:
-            o: `Expr` to be processed.
+            o: `Expr` or `BaseForm` to be processed.
 
         Returns:
             Processed object.
diff --git a/ufl/algorithms/cancel_jacobian_products.py b/ufl/algorithms/cancel_jacobian_products.py
index d6c3d103f..d1e72ef54 100644
--- a/ufl/algorithms/cancel_jacobian_products.py
+++ b/ufl/algorithms/cancel_jacobian_products.py
@@ -61,6 +61,7 @@
 from ufl.corealg.dag_traverser import DAGTraverser
 from ufl.corealg.map_dag import map_expr_dag
 from ufl.domain import extract_unique_domain
+from ufl.form import BaseForm
 
 
 def _flatten_product(expr, factors):
@@ -93,15 +94,15 @@ class IndexSumSimplifier(DAGTraverser):
     def __init__(
         self,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
         self._rules: dict[tuple, IndexReplacer] = {}
 
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``."""
         return super().process(o)
 
@@ -260,7 +261,7 @@ def match(self, with_k, rest, k):
     # Work around singledispatchmethod inheritance issue;
     # see https://bugs.python.org/issue36457.
     @singledispatchmethod
-    def process(self, o: Expr) -> Expr:
+    def process(self, o: Expr | BaseForm) -> Expr | BaseForm:
         """Process ``o``."""
         return super().process(o)
 
diff --git a/ufl/algorithms/extract_linear_combination.py b/ufl/algorithms/extract_linear_combination.py
index edd57a31e..433ad0634 100644
--- a/ufl/algorithms/extract_linear_combination.py
+++ b/ufl/algorithms/extract_linear_combination.py
@@ -50,6 +50,10 @@ def process(self, o: ufl.classes.Expr, **kwargs):
         """Fallback for any unsupported node types."""
         raise ValueError(f"Unsupported UFL node type for linear combinations: {type(o)}")
 
+    @process.register(ufl.coefficient.BaseCoefficient)
+    def _(self, o, **kwargs):
+        raise NotImplementedError(f"Unsupported UFL node type for linear combinations: {type(o)}")
+
     # ---------------------------------------------------------
     # 1. Terminals (Leaves) - No children to evaluate
     # ---------------------------------------------------------
@@ -70,6 +74,8 @@ def _(self, o, **kwargs):
             return o
         raise ValueError(f"Only scalar constants are supported, got shape {o.ufl_shape}")
 
+    @process.register(ufl.Cofunction)
+    @process.register(ufl.Matrix)
     @process.register(ufl.classes.Coefficient)
     def _(self, o, **kwargs):
         return [(ufl.as_ufl(1.0), o)]
@@ -178,17 +184,10 @@ def _(self, o, **kwargs):
             "Direct array assignment of indexed vector components is not supported."
         )
 
-    @process.register(ufl.Cofunction)
-    @process.register(ufl.Matrix)
-    def _(self, o, **kwargs):
-        # Cofunctions and Matrices are symbolic algebraic terminals.
-        # They act purely as spatial fields in linear combinations.
-        return [(ufl.as_ufl(1.0), o)]
-
 
 def extract_linear_combination(
-    expr: ufl.core.expr.Expr,
-) -> list[tuple[ufl.core.expr.Expr, ufl.classes.Coefficient]]:
+    expr: ufl.core.expr.Expr | ufl.form.BaseForm,
+) -> list[tuple[ufl.core.expr.Expr, ufl.coefficient.BaseCoefficient]]:
     """Wrapper to initialize traverser and extract linear combinations.
 
     Returns:
diff --git a/ufl/corealg/dag_traverser.py b/ufl/corealg/dag_traverser.py
index 2d86b1ce3..643fe99c2 100644
--- a/ufl/corealg/dag_traverser.py
+++ b/ufl/corealg/dag_traverser.py
@@ -24,15 +24,15 @@ class DAGTraverser:
     def __init__(
         self,
         compress: bool | None = True,
-        visited_cache: dict[tuple, Expr] | None = None,
-        result_cache: dict[Expr, Expr] | None = None,
+        visited_cache: dict[tuple, Expr | BaseForm] | None = None,
+        result_cache: dict[Expr | BaseForm, Expr | BaseForm] | None = None,
     ) -> None:
         """Initialise."""
         self._compress = compress
         self._visited_cache = {} if visited_cache is None else visited_cache
         self._result_cache = {} if result_cache is None else result_cache
 
-    def __call__(self, node: Expr, **kwargs) -> Any:
+    def __call__(self, node: Expr | BaseForm, **kwargs) -> Any:
         """Perform memoised DAG traversal with ``process`` singledispatch method.
 
         Args:
@@ -65,7 +65,7 @@ def __call__(self, node: Expr, **kwargs) -> Any:
             return result
 
     @singledispatchmethod
-    def process(self, o: Expr, **kwargs) -> Any:
+    def process(self, o: Expr | BaseForm, **kwargs) -> Any:
         """Process node by type.
 
         Args:

From 3c7421b027601cfb934f4d96e1b793d11a6aaaf2 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 07:11:26 +0000
Subject: [PATCH 7/9] Remove test that didn't check what it claimed as UFL
 eagerly simplifies.

---
 test/test_extract_linear_combination.py | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/test/test_extract_linear_combination.py b/test/test_extract_linear_combination.py
index 4170616b6..2049dd6b5 100644
--- a/test/test_extract_linear_combination.py
+++ b/test/test_extract_linear_combination.py
@@ -67,20 +67,6 @@ def test_scalars_and_powers(V, domain):
     assert (c + d**2, u) in res
 
 
-def test_ufl_zero_node(V):
-    """Test that UFL Zero nodes are handled correctly (usually evaluates to an empty list)."""
-    u = ufl.Coefficient(V)
-
-    # Create an explicit Zero node with the shape of u
-    z = ufl.classes.Zero(u.ufl_shape, u.ufl_free_indices, u.ufl_index_dimensions)
-
-    expr = u + z
-    res = extract_linear_combination(expr)
-
-    assert len(res) == 1
-    assert res[0] == (1.0, u)
-
-
 def test_nonlinear_errors(V):
     """Test that non-linear operations strictly raise ValueErrors."""
     u = ufl.Coefficient(V)

From a515953a841c70487cfbc0ffe646664b6d2c8d35 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 07:44:10 +0000
Subject: [PATCH 8/9] Add fix for real-valued elements

---
 ufl/algorithms/extract_linear_combination.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ufl/algorithms/extract_linear_combination.py b/ufl/algorithms/extract_linear_combination.py
index 433ad0634..50701eb28 100644
--- a/ufl/algorithms/extract_linear_combination.py
+++ b/ufl/algorithms/extract_linear_combination.py
@@ -78,6 +78,9 @@ def _(self, o, **kwargs):
     @process.register(ufl.Matrix)
     @process.register(ufl.classes.Coefficient)
     def _(self, o, **kwargs):
+        # Check for real-valued elements
+        if ufl.checks.is_scalar_constant_expression(o):
+            return o
         return [(ufl.as_ufl(1.0), o)]
 
     # ---------------------------------------------------------

From 3a6cc356c4ec1979188cad555ca962fc4c156aea Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 07:52:51 +0000
Subject: [PATCH 9/9] Separate cofunction and matrix from coefficient as you
 can't check if they are scalar constant expressions

---
 ufl/algorithms/extract_linear_combination.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ufl/algorithms/extract_linear_combination.py b/ufl/algorithms/extract_linear_combination.py
index 50701eb28..7bfcfe6d0 100644
--- a/ufl/algorithms/extract_linear_combination.py
+++ b/ufl/algorithms/extract_linear_combination.py
@@ -76,6 +76,9 @@ def _(self, o, **kwargs):
 
     @process.register(ufl.Cofunction)
     @process.register(ufl.Matrix)
+    def _(self, o, **kwargs):
+        return [(ufl.as_ufl(1.0), o)]
+
     @process.register(ufl.classes.Coefficient)
     def _(self, o, **kwargs):
         # Check for real-valued elements