From 13e1fe1413f2ca089d26f0a7c8e54350d94bdff9 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 19 Aug 2026 12:16:22 +0000
Subject: [PATCH 01/13] Implement linear combination assignment

---
 .../blocks/function_assigner.py               | 166 +++++++++---------
 src/dolfinx_adjoint/types/function.py         |   7 +-
 src/dolfinx_adjoint/utils.py                  | 126 +++++++++++++
 tests/test_assign.py                          |  58 ++++++
 4 files changed, 268 insertions(+), 89 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index 5cb0c18..9128a33 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -2,11 +2,13 @@
 
 import dolfinx
 import numpy as np
+import numpy.typing as npt
 import ufl
 from pyadjoint import AdjFloat, Block, OverloadedType
+from ufl.corealg.traversal import traverse_unique_terminals
 from ufl.formatting.ufl2unicode import ufl2unicode
 
-from ..utils import function_from_vector
+from ..utils import assign_linear_combination, function_from_vector, extract_linear_combination
 from ._vector import _vector
 
 
@@ -14,7 +16,6 @@ class FunctionAssignBlock(Block):
     def __init__(
         self,
         other: typing.Union[np.inexact, int, float],
-        func: dolfinx.fem.Function,
         ad_block_tag: typing.Optional[str] = None,
     ):
         super().__init__(ad_block_tag=ad_block_tag)
@@ -25,15 +26,15 @@ def __init__(
         elif isinstance(other, float) or isinstance(other, int):
             other = AdjFloat(other)
             self.add_dependency(other, no_duplicates=True)
-        elif not (isinstance(other, float) or isinstance(other, int)):
-            raise NotImplementedError("This should eventually be supported")
-            # # Assume that this is a point-wise evaluated UFL expression (firedrake only)
-            # for op in traverse_unique_terminals(other):
-            #     if isinstance(op, OverloadedType):
-            #         self.add_dependency(op, no_duplicates=True)
-            # self.expr = other
         else:
-            raise NotImplementedError("We should not get here!")
+            # Extract linear combination
+            lin_comb = extract_linear_combination(other)
+            if len(lin_comb) == 0:
+                raise ValueError("No linear combination found in the expression.")
+            for op in traverse_unique_terminals(other):
+                if isinstance(op, OverloadedType):
+                    self.add_dependency(op, no_duplicates=True)
+            self.expr = other
 
     def _replace_with_saved_output(self):
         if self.expr is None:
@@ -53,85 +54,76 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         expr = self._replace_with_saved_output()
         return expr, adj_input_func
 
+    def _compute_adjoint_of_broadcast(self, input: dolfinx.la.Vector | npt.NDArray | float | int) -> float | int:
+        # Adjoint of a broadcast is just a sum
+        if isinstance(input, dolfinx.la.Vector):
+            one = dolfinx.la.vector(input.index_map, input.block_size, input.array.dtype)
+            one.array[:] = 1
+            return dolfinx.cpp.la.inner_product(input._cpp_object, one._cpp_object)
+        else:
+            try:
+                return input.sum()
+            except AttributeError:
+                # Catch the case where input[0] is just a float
+                return input
+
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
+        bo = block_variable.output
         if self.expr is None:
-            if isinstance(block_variable.output, AdjFloat):
-                # Adjoint of a broadcast is just a sum
-                if isinstance(adj_inputs[0], dolfinx.la.Vector):
-                    vec = adj_inputs[0]
-                    one = dolfinx.la.vector(
-                        adj_inputs[0].index_map, adj_inputs[0].block_size, adj_inputs[0].array.dtype
-                    )
-                    one.array[:] = 1
-                    return dolfinx.cpp.la.inner_product(vec._cpp_object, one._cpp_object)
-                else:
-                    try:
-                        return adj_inputs[0].sum()
-                    except AttributeError:
-                        # Catch the case where adj_inputs[0] is just a float
-                        return adj_inputs[0]
-            elif isinstance(func := block_variable.output, dolfinx.fem.Function):
+            assert len(adj_inputs) == 1
+            if isinstance(func := bo, AdjFloat):
+                return self._compute_adjoint_of_broadcast(adj_inputs[0])
+            elif isinstance(func, dolfinx.fem.Function):
                 assert func.function_space == prepared.function_space
                 vec = _vector(
                     prepared.x.index_map, prepared.x.block_size, func.function_space, dtype=prepared.x.array.dtype
                 )
                 vec.array[:] = prepared.x.array[:]
                 return vec
+            elif isinstance(bo, dolfinx.fem.Constant):
+                raise NotImplementedError(
+                    "Adjoint for Constant assignment not implemented, use dolfinx_adjoint.Constant instead."
+                )
+            else:
+                raise NotImplementedError(f"Adjoint for {block_variable=} not implemented.")
+        else:
+            # Linear combination
+            expr, adj_input_func = prepared
+            vec = _vector(
+                bo.x.index_map,
+                bo.x.block_size,
+                bo.function_space,
+                dtype=bo.x.array.dtype,
+            )
+            if isinstance(bo, dolfinx.fem.Function) and bo.function_space == adj_input_func.function_space:
+                # Differentiate with respect to one of the input functions
+                diff_expr = ufl.algorithms.expand_derivatives(
+                    ufl.derivative(expr, block_variable.saved_output, adj_input_func)
+                )
+                temp_func = dolfinx.fem.Function(bo.function_space)
+                assign_linear_combination(diff_expr, temp_func)
+                vec.array[:] = temp_func.x.array[:]
+                return vec
+            elif isinstance(bo, dolfinx.fem.Function) and bo.ufl_element().is_real:
+                # Differentiate with respect to a real function (constant stored as Function)
+                # Create a perturbation direction in the Real space (value = 1.0)
+                direction = dolfinx.fem.Function(bo.function_space)
+                direction.x.array[0] = 1.0
+
+                # Differentiate expr w.r.t 'bo' in that direction
+                diff_expr = ufl.algorithms.expand_derivatives(
+                    ufl.derivative(expr, block_variable.saved_output, direction)
+                )
 
+                # Evaluate the derivative at the DOFs of the target space V
+                diff_eval = dolfinx.fem.Function(adj_input_func.function_space)
+                assign_linear_combination(diff_expr, diff_eval)
+
+                # Chain rule: dot product of (dz/dr) and adjoint inputs (bar_u)
+                vec.array[0] = dolfinx.cpp.la.inner_product(diff_eval.x._cpp_object, adj_input_func.x._cpp_object)
+                return vec
             else:
                 raise NotImplementedError(f"Adjoint for {block_variable=} not implemented.")
-            # elif isinstance(block_variable.output, dolfinx.fem.Constant):
-        #         R = block_variable.output._ad_function_space(prepared.function_space.mesh)
-        #         return self._adj_assign_constant(prepared, R)
-        #     else:
-        #         adj_output = dolfinx.fem.Function(
-        #             block_variable.output.function_space())
-        #         adj_output.assign(prepared)
-        #         return adj_output.vector()
-        # else:
-        #     # Linear combination
-        #     expr, adj_input_func = prepared
-        #     adj_output = dolfinx.fem.Function(adj_input_func.function_space)
-        #     if not isinstance(block_variable.output, dolfinx.fem.Constant):
-        #         diff_expr = ufl.algorithms.expand_derivatives(
-        #             ufl.derivative(expr, block_variable.saved_output, adj_input_func)
-        #         )
-        #         adj_output.assign(diff_expr)
-        #     else:
-        #         mesh = adj_output.function_space().mesh()
-        #         diff_expr = ufl.algorithms.expand_derivatives(
-        #             ufl.derivative(
-        #                 expr,
-        #                 block_variable.saved_output,
-        #                 create_constant(1., domain=mesh)
-        #             )
-        #         )
-        #         adj_output.assign(diff_expr)
-        #         return adj_output.vector().inner(adj_input_func.vector())
-
-        #     if isinstance(block_variable.output, dolfin.Constant):
-        #         R = block_variable.output._ad_function_space(adj_output.function_space().mesh())
-        #         return self._adj_assign_constant(adj_output, R)
-        #     else:
-        #         return adj_output.vector()
-
-    def _adj_assign_constant(self, adj_output, constant_fs):
-        r = dolfinx.fem.Function(constant_fs)
-        shape = r.ufl_shape
-        raise NotImplementedError("Not implemented for constants.")
-
-        if shape == () or shape[0] == 1:
-            # Scalar Constant
-            raise NotImplementedError("Not implemented for scalar constants yet.")
-            # r.vector()[:] = adj_output.vector().sum()
-        # else:
-        #     # We assume the shape of the constant == shape of the output function if not scalar.
-        #     # This assumption is due to FEniCS not supporting products with non-scalar constants in assign.
-        #     values = []
-        #     for i in range(shape[0]):
-        #         values.append(adj_output.sub(i, deepcopy=True).vector().sum())
-        #     r.assign(dolfin.Constant(values))
-        return r.vector()
 
     def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
         if self.expr is None:
@@ -144,11 +136,13 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
             return tlm_inputs[0]
         expr = prepared
         dudm = dolfinx.fem.Function(block_variable.output.function_space)
+        dudm.x.array[:] = 0.0
         dudmi = dolfinx.fem.Function(block_variable.output.function_space)
         for dep in self.get_dependencies():
             if dep.tlm_value:
-                dudmi.assign(ufl.algorithms.expand_derivatives(ufl.derivative(expr, dep.saved_output, dep.tlm_value)))
-                dudm.vector().axpy(1.0, dudmi.vector())
+                diff_expr = ufl.algorithms.expand_derivatives(ufl.derivative(expr, dep.saved_output, dep.tlm_value))
+                assign_linear_combination(diff_expr, dudmi)
+                dudm.x.array[:] += dudmi.x.array[:]
 
         return dudm
 
@@ -181,14 +175,12 @@ def recompute_component(self, inputs, block_variable, idx, prepared):
         # We should return the exact object instance to maintain C++ memory bindings
         # (especially for DirichletBCs), updating it in-place.
         output = block_variable.saved_output
-
-        try:
-            if output.function_space == prepared.function_space:
-                output.x.array[:] = prepared.x.array[:]
-        except AttributeError:
-            # Handling float value
+        if isinstance(prepared, dolfinx.fem.Function):
+            output.x.array[:] = prepared.x.array[:]
+        elif isinstance(prepared, (float, int)):
             output.x.array[:] = prepared
-
+        else:
+            assign_linear_combination(prepared, output)
         return output
 
     def __str__(self):
diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py
index 165c36c..13eb5bc 100644
--- a/src/dolfinx_adjoint/types/function.py
+++ b/src/dolfinx_adjoint/types/function.py
@@ -17,7 +17,7 @@
 
 from ..blocks.assembly import assemble_compiled_form
 from ..blocks.function_assigner import FunctionAssignBlock
-from ..utils import ad_kwargs, function_from_vector, gather
+from ..utils import ad_kwargs, assign_linear_combination, function_from_vector, gather
 
 
 class Function(dolfinx.fem.Function, FloatingType):
@@ -273,7 +273,7 @@ def assign(value: typing.Union[numpy.inexact, float, int], function: Function, *
     if annotate:
         if not isinstance(value, ufl.core.operator.Operator):
             value = create_overloaded_object(value)
-        block = FunctionAssignBlock(value, function, ad_block_tag=ad_block_tag)
+        block = FunctionAssignBlock(value, ad_block_tag=ad_block_tag)
         tape = get_working_tape()
         tape.add_block(block)
 
@@ -285,6 +285,9 @@ def assign(value: typing.Union[numpy.inexact, float, int], function: Function, *
                 "Function spaces of the value and function must match for assignment."
             )
             function.x.array[:] = value.x.array[:]
+        elif isinstance(value, ufl.core.expr.Expr):
+            # Linear combination of functions, e.g., 2*u + 3*v
+            assign_linear_combination(value, function)
         else:
             raise ValueError(f"Unsupported value type for assignment: {type(value)})")
     if annotate:
diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index 5f0b5d1..b4a7e57 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -3,6 +3,7 @@
 import dolfinx
 import numpy
 import numpy.typing as npt
+import ufl
 
 
 def function_from_vector(
@@ -41,3 +42,128 @@ class ad_kwargs(typing.TypedDict):
     """Tag for the block in the adjoint tape."""
     annotate: typing.NotRequired[bool]
     """Whether to annotate the assignment in the adjoint tape."""
+
+
+def extract_scalar_value(scalar_expr):
+    """Extract float from a scalar UFL expression."""
+    if isinstance(scalar_expr, (ufl.classes.IntValue, ufl.classes.FloatValue)):
+        return float(scalar_expr)
+    elif isinstance(scalar_expr, dolfinx.fem.Function):
+        # Check if it's a RealElement (constant stored as Function)
+        if scalar_expr.function_space.ufl_element().is_real and scalar_expr.ufl_shape == ():
+            return float(scalar_expr.x.array[0])
+        else:
+            raise ValueError(f"Cannot extract scalar from spatial Function: {scalar_expr}")
+    elif isinstance(scalar_expr, dolfinx.fem.Constant) and scalar_expr.ufl_shape == ():
+        val = scalar_expr.value
+        return float(val) if hasattr(val, "__float__") else float(val.item())
+    elif isinstance(scalar_expr, ufl.classes.ScalarValue):
+        return float(scalar_expr._value)
+    elif isinstance(scalar_expr, ufl.classes.Product):
+        result = 1.0
+        for op in scalar_expr.ufl_operands:
+            result *= extract_scalar_value(op)
+        return result
+    elif isinstance(scalar_expr, ufl.classes.Division):
+        num, den = scalar_expr.ufl_operands
+        return extract_scalar_value(num) / extract_scalar_value(den)
+    else:
+        raise ValueError(f"Cannot extract scalar from {type(scalar_expr)}: {scalar_expr}")
+
+
+def extract_function(expr) -> tuple[bool, dolfinx.fem.Function | None]:
+    """Recursively extract a Function from nested UFL expressions."""
+    if isinstance(expr, dolfinx.fem.Function):
+        is_real = expr.function_space.ufl_element().is_real
+        if is_real:
+            return (False, None)
+        return (False, expr)
+    elif isinstance(expr, (ufl.classes.Indexed, ufl.classes.ComponentTensor)):
+        return extract_function(expr.ufl_operands[0])
+    elif hasattr(expr, "ufl_operands"):
+        found_func = None
+        for op in expr.ufl_operands:
+            is_real, func = extract_function(op)
+            if func is not None:
+                if found_func is not None:
+                    raise ValueError(f"Non-linear expression detected: multiple spatial functions in {expr}")
+                found_func = func
+        return (False, found_func)
+    return (False, None)
+
+
+def extract_term(term):
+    """Extract (weight, function) from a single term."""
+    if isinstance(term, dolfinx.fem.Function):
+        is_real = term.function_space.ufl_element().is_real
+        if is_real:
+            return None
+        return (1.0, term)
+    elif isinstance(term, ufl.classes.ComponentTensor):
+        return extract_term(term.ufl_operands[0])
+    elif isinstance(term, ufl.classes.Indexed):
+        is_real, func = extract_function(term)
+        if func is None:
+            return None
+        return (1.0, func)
+    elif isinstance(term, ufl.classes.Product):
+        weight = 1.0
+        func = None
+        for op in term.ufl_operands:
+            is_real, extracted_func = extract_function(op)
+            if extracted_func is not None:
+                if func is not None:
+                    raise ValueError(f"Non-linear term detected: multiple spatial functions in {term}")
+                func = extracted_func
+            else:
+                weight *= extract_scalar_value(op)
+        return (weight, func) if func is not None else None
+    elif isinstance(term, ufl.classes.Division):
+        num, den = term.ufl_operands
+        denom_val = extract_scalar_value(den)
+        if isinstance(num, dolfinx.fem.Function):
+            is_real = num.function_space.ufl_element().is_real
+            if is_real:
+                return None
+            return (1.0 / denom_val, num)
+        elif isinstance(num, ufl.classes.Product):
+            result = extract_term(num)
+            return (result[0] / denom_val, result[1]) if result else None
+    return None
+
+
+def extract_linear_combination(expr: ufl.core.expr.Expr) -> list[tuple[float, dolfinx.fem.Function]]:
+    """Extract (weight, function) pairs from a UFL linear combination.
+
+    Analyzes expressions like: 0.5*u + 0.3*v + 0.2*w
+    Returns: [(0.5, u), (0.3, v), (0.2, w)]
+
+    :param expr: UFL expression (Sum, Product, or single Function)
+    :returns: List of (weight, function) tuples
+    """
+
+    # Parse the expression, flattening nested Sums recursively
+    if isinstance(expr, ufl.classes.Sum):
+        summands = expr.ufl_operands
+    else:
+        summands = [expr]
+    terms = []
+    for summand in summands:
+        if isinstance(summand, ufl.classes.Sum):
+            # Recursively flatten nested Sum structures
+            terms.extend(extract_linear_combination(summand))
+        else:
+            result = extract_term(summand)
+            if result is not None:
+                terms.append(result)
+    return terms
+
+
+def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.Function):
+    pairs = extract_linear_combination(value)
+    function.x.array[:] = 0.0
+    for weight, func in pairs:
+        if not func.function_space == function.function_space:
+            raise ValueError("Function spaces of all functions in the linear combination must match for assignment.")
+        function.x.array[:] += weight * func.x.array[:]
+    function.x.scatter_forward()
diff --git a/tests/test_assign.py b/tests/test_assign.py
index d72d177..66db911 100644
--- a/tests/test_assign.py
+++ b/tests/test_assign.py
@@ -2,6 +2,7 @@
 
 from mpi4py import MPI
 
+import basix
 import dolfinx
 import numpy
 import numpy as np
@@ -27,6 +28,63 @@ def mesh_3D():
     return dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 50, 50, 50)
 
 
+def test_assign_linear_combination(mesh_1D):
+    V = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+    f = Function(V, name="f")
+    f.interpolate(lambda x: 2.0 * x[0])
+    g = Function(V, name="g")
+    g.interpolate(lambda x: 3.0 * x[0] ** 2)
+    u = Function(V, name="u")
+
+    assign(3 * f - g, u)
+
+    J = assemble_scalar(u**2 * ufl.dx)
+    rf = pyadjoint.ReducedFunctional(J, pyadjoint.Control(f))
+    h = Function(V)
+    rng = np.random.default_rng(seed=42)
+    num_dofs_local = (V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts) * V.dofmap.index_map_bs
+    rand = rng.random(size=num_dofs_local, dtype=h.dtype)
+    h.x.array[:] = rand
+    h.x.scatter_forward()
+    assert pyadjoint.taylor_test(rf, f, h) > 1.9
+
+    rf2 = pyadjoint.ReducedFunctional(J, pyadjoint.Control(g))
+    assert pyadjoint.taylor_test(rf2, g, h) > 1.9
+
+
+def test_assign_lincomb_real_space(mesh_1D):
+    r_el = basix.ufl.real_element(mesh_1D.basix_cell(), value_shape=())
+    R = dolfinx.fem.functionspace(mesh_1D, r_el)
+    r = Function(R, name="r")
+    r.x.array[0] = 0.2
+
+    V = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+    v = Function(V, name="u")
+    v.interpolate(lambda x: 3.0 * x[0] ** 2)
+
+    z = -2 * v + 4 * r * v
+    u = Function(V, name="u_output")
+    assign(z, u)
+
+    J = assemble_scalar(u**2 * ufl.dx)
+    rf = pyadjoint.ReducedFunctional(J, pyadjoint.Control(r))
+    h = Function(R)
+    rng = np.random.default_rng(seed=42)
+    num_dofs_local = (R.dofmap.index_map.size_local + R.dofmap.index_map.num_ghosts) * R.dofmap.index_map_bs
+    rand = rng.random(size=num_dofs_local, dtype=h.dtype)
+    h.x.array[:] = rand
+    h.x.scatter_forward()
+    assert pyadjoint.taylor_test(rf, r, h) > 1.9
+
+    rf2 = pyadjoint.ReducedFunctional(J, pyadjoint.Control(v))
+    hv = Function(V)
+    num_dofs_local = (V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts) * V.dofmap.index_map_bs
+    rand = rng.random(size=num_dofs_local, dtype=hv.dtype)
+    hv.x.array[:] = rand
+    hv.x.scatter_forward()
+    assert pyadjoint.taylor_test(rf2, v, hv) > 1.9
+
+
 def test_multiple_assign_adjoint_accumulation(mesh_1D):
     """
     Test that assigning a Function multiple times and computing the adjoint

From 5773608b360fca32e13d9890c02f2e03e5b6ab95 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 19 Aug 2026 12:28:50 +0000
Subject: [PATCH 02/13] Sort imports

---
 src/dolfinx_adjoint/blocks/function_assigner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index 9128a33..b83bcf9 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -8,7 +8,7 @@
 from ufl.corealg.traversal import traverse_unique_terminals
 from ufl.formatting.ufl2unicode import ufl2unicode
 
-from ..utils import assign_linear_combination, function_from_vector, extract_linear_combination
+from ..utils import assign_linear_combination, extract_linear_combination, function_from_vector
 from ._vector import _vector
 
 

From 3c7b2e3fca1cbe4f8bd006ca7b458c97e6120cf2 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 19 Aug 2026 13:22:11 +0000
Subject: [PATCH 03/13] Add docstring

---
 src/dolfinx_adjoint/blocks/function_assigner.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index b83bcf9..d1e4128 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -13,6 +13,12 @@
 
 
 class FunctionAssignBlock(Block):
+    """Block for assigning data directly to a `Function` on the tape.
+
+    This block handles the assignment of a linear combination of `Function`s
+    or constants to a target `Function`.
+    """
+
     def __init__(
         self,
         other: typing.Union[np.inexact, int, float],
@@ -59,12 +65,12 @@ def _compute_adjoint_of_broadcast(self, input: dolfinx.la.Vector | npt.NDArray |
         if isinstance(input, dolfinx.la.Vector):
             one = dolfinx.la.vector(input.index_map, input.block_size, input.array.dtype)
             one.array[:] = 1
-            return dolfinx.cpp.la.inner_product(input._cpp_object, one._cpp_object)
+            return dolfinx.cpp.la.inner_product(input._cpp_object, one._cpp_object)  # type: ignore[arg-type]
         else:
-            try:
+            if hasattr(input, "sum"):
                 return input.sum()
-            except AttributeError:
-                # Catch the case where input[0] is just a float
+            else:
+                # Catch the case where input is just a float
                 return input
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):

From 1cb29908e8c315a8422d7edf9defe50dd2fbc0d0 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 19 Aug 2026 13:36:51 +0000
Subject: [PATCH 04/13] Slight reformatting and restructuring to avoid circular
 depenendcies and simplify copy logic

---
 src/dolfinx_adjoint/__init__.py               |  2 +-
 .../blocks/function_assigner.py               |  4 +-
 src/dolfinx_adjoint/function.py               | 51 +++++++++++++++++++
 src/dolfinx_adjoint/types/function.py         | 43 ++--------------
 src/dolfinx_adjoint/utils.py                  |  1 +
 5 files changed, 59 insertions(+), 42 deletions(-)
 create mode 100644 src/dolfinx_adjoint/function.py

diff --git a/src/dolfinx_adjoint/__init__.py b/src/dolfinx_adjoint/__init__.py
index 7a0a688..cef54ce 100644
--- a/src/dolfinx_adjoint/__init__.py
+++ b/src/dolfinx_adjoint/__init__.py
@@ -6,9 +6,9 @@
 import pyadjoint as _pyad
 
 from .assembly import assemble_scalar, error_norm
+from .function import assign
 from .solvers import LinearProblem, NonlinearProblem
 from .types import Constant, Function, dirichletbc
-from .types.function import assign
 
 meta = metadata("dolfinx_adjoint")
 __version__ = meta.get("Version")
diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index d1e4128..2e363d0 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -8,6 +8,7 @@
 from ufl.corealg.traversal import traverse_unique_terminals
 from ufl.formatting.ufl2unicode import ufl2unicode
 
+from ..types.function import Function as _Function
 from ..utils import assign_linear_combination, extract_linear_combination, function_from_vector
 from ._vector import _vector
 
@@ -21,7 +22,7 @@ class FunctionAssignBlock(Block):
 
     def __init__(
         self,
-        other: typing.Union[np.inexact, int, float],
+        other: np.inexact | int | float | _Function | ufl.core.expr.Expr,
         ad_block_tag: typing.Optional[str] = None,
     ):
         super().__init__(ad_block_tag=ad_block_tag)
@@ -34,6 +35,7 @@ def __init__(
             self.add_dependency(other, no_duplicates=True)
         else:
             # Extract linear combination
+            assert isinstance(other, ufl.core.expr.Expr), f"Expected UFL expression, got {type(other)}"
             lin_comb = extract_linear_combination(other)
             if len(lin_comb) == 0:
                 raise ValueError("No linear combination found in the expression.")
diff --git a/src/dolfinx_adjoint/function.py b/src/dolfinx_adjoint/function.py
new file mode 100644
index 0000000..9ce8e1d
--- /dev/null
+++ b/src/dolfinx_adjoint/function.py
@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import typing
+
+import dolfinx
+import numpy
+import ufl
+from pyadjoint.overloaded_type import (
+    create_overloaded_object,
+)
+from pyadjoint.tape import annotate_tape, get_working_tape, stop_annotating
+
+from .blocks.function_assigner import FunctionAssignBlock
+from .types.function import Function as _Function
+from .utils import ad_kwargs, assign_linear_combination
+
+
+def assign(value: typing.Union[numpy.inexact, float, int], function: _Function, **kwargs: typing.Unpack[ad_kwargs]):
+    """Assign a `value` to a :py:func:`dolfinx_adjoint.Function`.
+
+    Args:
+        value: The value to assign to the function.
+        function: The function to assign the value to.
+        *args: Additional positional arguments to pass to the assign method.
+        **kwargs: Additional keyword arguments to pass to the assign method.
+    """
+    # do not annotate in case of self assignment
+    ad_block_tag = kwargs.pop("ad_block_tag", None)
+    annotate = annotate_tape(kwargs) and value != function
+    if annotate:
+        if not isinstance(value, ufl.core.operator.Operator):
+            value = create_overloaded_object(value)
+        block = FunctionAssignBlock(value, ad_block_tag=ad_block_tag)
+        tape = get_working_tape()
+        tape.add_block(block)
+
+    with stop_annotating():
+        if isinstance(value, (numpy.inexact, float, int)):
+            function.x.array[:] = value
+        elif isinstance(value, dolfinx.fem.Function):
+            assert value.function_space == function.function_space, (
+                "Function spaces of the value and function must match for assignment."
+            )
+            function.x.array[:] = value.x.array[:]
+        elif isinstance(value, ufl.core.expr.Expr):
+            # Linear combination of functions, e.g., 2*u + 3*v
+            assign_linear_combination(value, function)
+        else:
+            raise ValueError(f"Unsupported value type for assignment: {type(value)})")
+    if annotate:
+        block.add_output(function.create_block_variable())
diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py
index 13eb5bc..dfe5808 100644
--- a/src/dolfinx_adjoint/types/function.py
+++ b/src/dolfinx_adjoint/types/function.py
@@ -13,11 +13,10 @@
     get_overloaded_class,
     register_overloaded_type,
 )
-from pyadjoint.tape import annotate_tape, get_working_tape, no_annotations, stop_annotating
+from pyadjoint.tape import no_annotations
 
 from ..blocks.assembly import assemble_compiled_form
-from ..blocks.function_assigner import FunctionAssignBlock
-from ..utils import ad_kwargs, assign_linear_combination, function_from_vector, gather
+from ..utils import function_from_vector, gather
 
 
 class Function(dolfinx.fem.Function, FloatingType):
@@ -189,7 +188,7 @@ def _ad_to_list(m):
     def _ad_copy(self):
         """Create a (deep) copy of the function."""
         r = get_overloaded_class(dolfinx.fem.Function)(self.function_space)
-        assign(self, r)
+        r.x.array[:] = self.x.array[:].copy()
         return r
 
     @staticmethod
@@ -256,39 +255,3 @@ def _ad_init_object(cls, obj):
 
 register_overloaded_type(Function, (dolfinx.fem.Function, Function))
 register_overloaded_type(Constant, (dolfinx.fem.Constant, Constant))
-
-
-def assign(value: typing.Union[numpy.inexact, float, int], function: Function, **kwargs: typing.Unpack[ad_kwargs]):
-    """Assign a `value` to a :py:func:`dolfinx_adjoint.Function`.
-
-    Args:
-        value: The value to assign to the function.
-        function: The function to assign the value to.
-        *args: Additional positional arguments to pass to the assign method.
-        **kwargs: Additional keyword arguments to pass to the assign method.
-    """
-    # do not annotate in case of self assignment
-    ad_block_tag = kwargs.pop("ad_block_tag", None)
-    annotate = annotate_tape(kwargs) and value != function
-    if annotate:
-        if not isinstance(value, ufl.core.operator.Operator):
-            value = create_overloaded_object(value)
-        block = FunctionAssignBlock(value, ad_block_tag=ad_block_tag)
-        tape = get_working_tape()
-        tape.add_block(block)
-
-    with stop_annotating():
-        if isinstance(value, (numpy.inexact, float, int)):
-            function.x.array[:] = value
-        elif isinstance(value, dolfinx.fem.Function):
-            assert value.function_space == function.function_space, (
-                "Function spaces of the value and function must match for assignment."
-            )
-            function.x.array[:] = value.x.array[:]
-        elif isinstance(value, ufl.core.expr.Expr):
-            # Linear combination of functions, e.g., 2*u + 3*v
-            assign_linear_combination(value, function)
-        else:
-            raise ValueError(f"Unsupported value type for assignment: {type(value)})")
-    if annotate:
-        block.add_output(function.create_block_variable())
diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index b4a7e57..14f0d7f 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -143,6 +143,7 @@ def extract_linear_combination(expr: ufl.core.expr.Expr) -> list[tuple[float, do
     """
 
     # Parse the expression, flattening nested Sums recursively
+    summands: list[ufl.core.expr.Expr]
     if isinstance(expr, ufl.classes.Sum):
         summands = expr.ufl_operands
     else:

From c11c743afd7b38b3991d78771cfbe56e94e1c568 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 19 Aug 2026 14:04:09 +0000
Subject: [PATCH 05/13] Use type hint from ufl.

---
 src/dolfinx_adjoint/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index 14f0d7f..9a0f829 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -143,7 +143,7 @@ def extract_linear_combination(expr: ufl.core.expr.Expr) -> list[tuple[float, do
     """
 
     # Parse the expression, flattening nested Sums recursively
-    summands: list[ufl.core.expr.Expr]
+    summands: list[ufl.core.expr.Expr] | tuple[ufl.core.terminal.FormArgument, ...]
     if isinstance(expr, ufl.classes.Sum):
         summands = expr.ufl_operands
     else:

From 3cb2922eaa3d0c7d9fe47a20dc37b0ec2aaac3e9 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Thu, 20 Aug 2026 08:09:07 +0000
Subject: [PATCH 06/13] Refactor code to avoid re-allocations of functions at
 eval

---
 .../blocks/function_assigner.py               | 105 ++++++++++++------
 src/dolfinx_adjoint/function.py               |   2 +-
 src/dolfinx_adjoint/types/function.py         |   7 ++
 3 files changed, 77 insertions(+), 37 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index 2e363d0..986a9ab 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -14,25 +14,52 @@
 
 
 class FunctionAssignBlock(Block):
-    """Block for assigning data directly to a `Function` on the tape.
-
-    This block handles the assignment of a linear combination of `Function`s
-    or constants to a target `Function`.
+    """Block for assigning data directly to a :py:class:`dolfinx_adjoint.Function` on the tape.
+
+    This block handles the assignment of a linear combination of ":py:class:`dolfinx_adjoint.Function` objects
+    or constants to a target :py:class:`dolfinx_adjoint.Function`.
+
+    Args:
+        other: The right-hand side of the assignment, which can be a :py:class:`dolfinx_adjoint.Function`,
+            a :py:class:`dolfinx_adjoint.Constant`, or a linear combination of :py:class:`dolfinx_adjoint.Function`
+            objects.
+        func: The target :py:class:`dolfinx_adjoint.Function` to which the assignment is made.
+        ad_block_tag: Optional tag for identifying the block in the adjoint tape.
+            If not provided, a default tag will be generated.
     """
 
+    _working_memory: list[_Function]
+    _one: _Function  # Array for storing the value 1.0 for adjoint of broadcast operations
+
     def __init__(
         self,
         other: np.inexact | int | float | _Function | ufl.core.expr.Expr,
+        func: _Function,
         ad_block_tag: typing.Optional[str] = None,
     ):
         super().__init__(ad_block_tag=ad_block_tag)
+
+        # Allocate working memory for adjoint computations
+        self._working_memory = []
+        for i in range(2):
+            vec = _vector(func.x.index_map, func.x.block_size, func.function_space, dtype=func.x.array.dtype)
+            self._working_memory.append(
+                _Function(func.function_space, x=vec, annotate=False, name=f"working_memory_{i}")
+            )
+
+        # Extract dependencies
         self.other = None
         self.expr = None
-        if isinstance(other, OverloadedType):
-            self.add_dependency(other, no_duplicates=True)
-        elif isinstance(other, float) or isinstance(other, int):
+
+        if isinstance(other, (float, int)) and not isinstance(other, OverloadedType):
             other = AdjFloat(other)
+
+        if isinstance(other, OverloadedType):
             self.add_dependency(other, no_duplicates=True)
+            # If the dependency is a scalar broadcast, allocate the ones vector
+            if isinstance(other, AdjFloat):
+                self._one = _Function(func.function_space, name="one", annotate=False)
+                self._one.x.array[:] = 1.0
         else:
             # Extract linear combination
             assert isinstance(other, ufl.core.expr.Expr), f"Expected UFL expression, got {type(other)}"
@@ -42,6 +69,16 @@ def __init__(
             for op in traverse_unique_terminals(other):
                 if isinstance(op, OverloadedType):
                     self.add_dependency(op, no_duplicates=True)
+
+            # Allocate extra memory for adjoint computations if any of the operands are real functions
+            for op in traverse_unique_terminals(other):
+                if isinstance(op, _Function) and op.function_space.ufl_element().is_real:
+                    vec_2 = _vector(op.x.index_map, op.x.block_size, op.function_space, dtype=op.x.array.dtype)
+                    self._working_memory.append(
+                        _Function(op.function_space, x=vec_2, annotate=False, name="working_memory_2")
+                    )
+                    break
+
             self.expr = other
 
     def _replace_with_saved_output(self):
@@ -62,12 +99,16 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         expr = self._replace_with_saved_output()
         return expr, adj_input_func
 
-    def _compute_adjoint_of_broadcast(self, input: dolfinx.la.Vector | npt.NDArray | float | int) -> float | int:
+    @classmethod
+    def _compute_adjoint_of_broadcast(
+        cls, input: dolfinx.la.Vector | npt.NDArray | float | int, one: _Function
+    ) -> float | int:
+        """
+        Computes the adjoint of a broadcast operation into an R^N vector, which is simply the sum of the input values.
+        """
         # Adjoint of a broadcast is just a sum
         if isinstance(input, dolfinx.la.Vector):
-            one = dolfinx.la.vector(input.index_map, input.block_size, input.array.dtype)
-            one.array[:] = 1
-            return dolfinx.cpp.la.inner_product(input._cpp_object, one._cpp_object)  # type: ignore[arg-type]
+            return dolfinx.cpp.la.inner_product(input._cpp_object, one.x._cpp_object)  # type: ignore[arg-type]
         else:
             if hasattr(input, "sum"):
                 return input.sum()
@@ -79,15 +120,12 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
         bo = block_variable.output
         if self.expr is None:
             assert len(adj_inputs) == 1
-            if isinstance(func := bo, AdjFloat):
-                return self._compute_adjoint_of_broadcast(adj_inputs[0])
-            elif isinstance(func, dolfinx.fem.Function):
-                assert func.function_space == prepared.function_space
-                vec = _vector(
-                    prepared.x.index_map, prepared.x.block_size, func.function_space, dtype=prepared.x.array.dtype
-                )
-                vec.array[:] = prepared.x.array[:]
-                return vec
+            if isinstance(bo, AdjFloat):
+                return self._compute_adjoint_of_broadcast(adj_inputs[0], self._one)
+            elif isinstance(bo, dolfinx.fem.Function):
+                assert bo.function_space == prepared.function_space
+                self._working_memory[0].x.array[:] = adj_inputs[0].array[:]
+                return self._working_memory[0].x
             elif isinstance(bo, dolfinx.fem.Constant):
                 raise NotImplementedError(
                     "Adjoint for Constant assignment not implemented, use dolfinx_adjoint.Constant instead."
@@ -97,25 +135,18 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
         else:
             # Linear combination
             expr, adj_input_func = prepared
-            vec = _vector(
-                bo.x.index_map,
-                bo.x.block_size,
-                bo.function_space,
-                dtype=bo.x.array.dtype,
-            )
             if isinstance(bo, dolfinx.fem.Function) and bo.function_space == adj_input_func.function_space:
                 # Differentiate with respect to one of the input functions
                 diff_expr = ufl.algorithms.expand_derivatives(
                     ufl.derivative(expr, block_variable.saved_output, adj_input_func)
                 )
-                temp_func = dolfinx.fem.Function(bo.function_space)
-                assign_linear_combination(diff_expr, temp_func)
-                vec.array[:] = temp_func.x.array[:]
-                return vec
+                assign_linear_combination(diff_expr, self._working_memory[0])
+                return self._working_memory[0].x
             elif isinstance(bo, dolfinx.fem.Function) and bo.ufl_element().is_real:
                 # Differentiate with respect to a real function (constant stored as Function)
                 # Create a perturbation direction in the Real space (value = 1.0)
-                direction = dolfinx.fem.Function(bo.function_space)
+                assert len(self._working_memory) == 3, "Working memory not allocated for real function adjoint."
+                direction = self._working_memory[2]
                 direction.x.array[0] = 1.0
 
                 # Differentiate expr w.r.t 'bo' in that direction
@@ -124,12 +155,14 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
                 )
 
                 # Evaluate the derivative at the DOFs of the target space V
-                diff_eval = dolfinx.fem.Function(adj_input_func.function_space)
+                diff_eval = self._working_memory[1]
                 assign_linear_combination(diff_expr, diff_eval)
 
                 # Chain rule: dot product of (dz/dr) and adjoint inputs (bar_u)
-                vec.array[0] = dolfinx.cpp.la.inner_product(diff_eval.x._cpp_object, adj_input_func.x._cpp_object)
-                return vec
+                self._working_memory[2].x.array[0] = dolfinx.cpp.la.inner_product(
+                    diff_eval.x._cpp_object, adj_input_func.x._cpp_object
+                )
+                return self._working_memory[2].x
             else:
                 raise NotImplementedError(f"Adjoint for {block_variable=} not implemented.")
 
@@ -143,9 +176,9 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         if self.expr is None:
             return tlm_inputs[0]
         expr = prepared
-        dudm = dolfinx.fem.Function(block_variable.output.function_space)
+        dudm = self._working_memory[0]
         dudm.x.array[:] = 0.0
-        dudmi = dolfinx.fem.Function(block_variable.output.function_space)
+        dudmi = self._working_memory[1]
         for dep in self.get_dependencies():
             if dep.tlm_value:
                 diff_expr = ufl.algorithms.expand_derivatives(ufl.derivative(expr, dep.saved_output, dep.tlm_value))
diff --git a/src/dolfinx_adjoint/function.py b/src/dolfinx_adjoint/function.py
index 9ce8e1d..7a1df8b 100644
--- a/src/dolfinx_adjoint/function.py
+++ b/src/dolfinx_adjoint/function.py
@@ -30,7 +30,7 @@ def assign(value: typing.Union[numpy.inexact, float, int], function: _Function,
     if annotate:
         if not isinstance(value, ufl.core.operator.Operator):
             value = create_overloaded_object(value)
-        block = FunctionAssignBlock(value, ad_block_tag=ad_block_tag)
+        block = FunctionAssignBlock(value, function, ad_block_tag=ad_block_tag)
         tape = get_working_tape()
         tape.add_block(block)
 
diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py
index dfe5808..d5f82ce 100644
--- a/src/dolfinx_adjoint/types/function.py
+++ b/src/dolfinx_adjoint/types/function.py
@@ -61,6 +61,8 @@ def __init__(
             annotate=kwargs.pop("annotate", True),
             **kwargs,
         )
+        if x is not None:
+            self._x = x  # Ensure that the input `x` is stored in case it is a _SpecialVector
 
     @classmethod
     def _ad_init_object(cls, obj):
@@ -202,6 +204,11 @@ def _ad_assign_numpy(dst: dolfinx.fem.Function, src: numpy.ndarray, offset: int)
         dst.x.scatter_forward()
         return dst, offset
 
+    @property
+    def x(self) -> dolfinx.la.Vector:
+        """Return the underlying vector of the function."""
+        return self._x
+
 
 class Constant(Function):
     """A class overloading {py:class}`dolfinx.fem.Constant`

From bc4f14c850f2af86662b1e7dd86cd5d34182299c Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Thu, 20 Aug 2026 08:14:43 +0000
Subject: [PATCH 07/13] added dtype

---
 src/dolfinx_adjoint/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index 9a0f829..ad5a474 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -44,7 +44,7 @@ class ad_kwargs(typing.TypedDict):
     """Whether to annotate the assignment in the adjoint tape."""
 
 
-def extract_scalar_value(scalar_expr):
+def extract_scalar_value(scalar_expr: ufl.core.expr.Expr) -> float:
     """Extract float from a scalar UFL expression."""
     if isinstance(scalar_expr, (ufl.classes.IntValue, ufl.classes.FloatValue)):
         return float(scalar_expr)
@@ -92,7 +92,7 @@ def extract_function(expr) -> tuple[bool, dolfinx.fem.Function | None]:
     return (False, None)
 
 
-def extract_term(term):
+def extract_term(term: ufl.core.expr.Expr) -> tuple[float, dolfinx.fem.Function] | None:
     """Extract (weight, function) from a single term."""
     if isinstance(term, dolfinx.fem.Function):
         is_real = term.function_space.ufl_element().is_real

From 96d0cb57bf2d526d9aecec8efd12e273f7424813 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= <dokken92@gmail.com>
Date: Thu, 20 Aug 2026 10:15:04 +0200
Subject: [PATCH 08/13] Apply suggestions from code review

Co-authored-by: Henrik Finsberg <henrikfinsberg@hotmail.com>
---
 src/dolfinx_adjoint/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index ad5a474..d2e33d1 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -160,7 +160,7 @@ def extract_linear_combination(expr: ufl.core.expr.Expr) -> list[tuple[float, do
     return terms
 
 
-def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.Function):
+def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.Function) -> None:
     pairs = extract_linear_combination(value)
     function.x.array[:] = 0.0
     for weight, func in pairs:

From 6eed64f5c96d7a93a741810fe029af84855638f6 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Thu, 20 Aug 2026 08:44:14 +0000
Subject: [PATCH 09/13] Add tests and some more special cases (assign real
 scalar function to other function, equivalent of assigning a constant)

---
 .../blocks/function_assigner.py               | 54 +++++++++---
 src/dolfinx_adjoint/function.py               | 10 ++-
 tests/test_assign.py                          | 87 +++++++++++++++++++
 3 files changed, 134 insertions(+), 17 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index 986a9ab..3bd881d 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -13,6 +13,29 @@
 from ._vector import _vector
 
 
+def is_real_scalar_func(func: dolfinx.fem.Function) -> bool:
+    """Check if a function is a real scalar function.
+
+    Args:
+        func: The function to check.
+    """
+    return func.ufl_element().is_real and func.ufl_shape == ()
+
+
+def create_function_with_special_vector(func: _Function, name: str | None = None) -> _Function:
+    """Create a new function with the same function space as `func` but with a special vector for adjoint computations.
+
+    Args:
+        func: The original function from which to derive the new function.
+
+    Returns:
+        A new function with the same function space as `func` but with a special vector for adjoint computations.
+    """
+    name = name or f"{func.name}_special_vector"
+    vec = _vector(func.x.index_map, func.x.block_size, func.function_space, dtype=func.x.array.dtype)
+    return _Function(func.function_space, x=vec, annotate=False, name=name)
+
+
 class FunctionAssignBlock(Block):
     """Block for assigning data directly to a :py:class:`dolfinx_adjoint.Function` on the tape.
 
@@ -42,10 +65,7 @@ def __init__(
         # Allocate working memory for adjoint computations
         self._working_memory = []
         for i in range(2):
-            vec = _vector(func.x.index_map, func.x.block_size, func.function_space, dtype=func.x.array.dtype)
-            self._working_memory.append(
-                _Function(func.function_space, x=vec, annotate=False, name=f"working_memory_{i}")
-            )
+            self._working_memory.append(create_function_with_special_vector(func, name=f"working_memory_{i}"))
 
         # Extract dependencies
         self.other = None
@@ -60,7 +80,13 @@ def __init__(
             if isinstance(other, AdjFloat):
                 self._one = _Function(func.function_space, name="one", annotate=False)
                 self._one.x.array[:] = 1.0
+            elif isinstance(other, _Function) and is_real_scalar_func(other):
+                self._working_memory.append(create_function_with_special_vector(other, name="working_memory_2"))
+                self._one = _Function(func.function_space, name="one", annotate=False)
+                self._one.x.array[:] = 1.0
         else:
+            self.expr = other
+
             # Extract linear combination
             assert isinstance(other, ufl.core.expr.Expr), f"Expected UFL expression, got {type(other)}"
             lin_comb = extract_linear_combination(other)
@@ -72,15 +98,10 @@ def __init__(
 
             # Allocate extra memory for adjoint computations if any of the operands are real functions
             for op in traverse_unique_terminals(other):
-                if isinstance(op, _Function) and op.function_space.ufl_element().is_real:
-                    vec_2 = _vector(op.x.index_map, op.x.block_size, op.function_space, dtype=op.x.array.dtype)
-                    self._working_memory.append(
-                        _Function(op.function_space, x=vec_2, annotate=False, name="working_memory_2")
-                    )
+                if isinstance(op, _Function) and is_real_scalar_func(op):
+                    self._working_memory.append(create_function_with_special_vector(op, name="working_memory_2"))
                     break
 
-            self.expr = other
-
     def _replace_with_saved_output(self):
         if self.expr is None:
             return None
@@ -123,7 +144,14 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
             if isinstance(bo, AdjFloat):
                 return self._compute_adjoint_of_broadcast(adj_inputs[0], self._one)
             elif isinstance(bo, dolfinx.fem.Function):
-                assert bo.function_space == prepared.function_space
+                if is_real_scalar_func(bo):
+                    # Adjoint of a broadcast into a real function (constant stored as Function)
+                    self._working_memory[2].x.array[0] = self._compute_adjoint_of_broadcast(adj_inputs[0], self._one)
+                    return self._working_memory[2].x
+                if bo.function_space != prepared.function_space:
+                    raise ValueError(
+                        "Function spaces of the block variable and prepared function must match for adjoint evaluation."
+                    )
                 self._working_memory[0].x.array[:] = adj_inputs[0].array[:]
                 return self._working_memory[0].x
             elif isinstance(bo, dolfinx.fem.Constant):
@@ -142,7 +170,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
                 )
                 assign_linear_combination(diff_expr, self._working_memory[0])
                 return self._working_memory[0].x
-            elif isinstance(bo, dolfinx.fem.Function) and bo.ufl_element().is_real:
+            elif isinstance(bo, dolfinx.fem.Function) and is_real_scalar_func(bo):
                 # Differentiate with respect to a real function (constant stored as Function)
                 # Create a perturbation direction in the Real space (value = 1.0)
                 assert len(self._working_memory) == 3, "Working memory not allocated for real function adjoint."
diff --git a/src/dolfinx_adjoint/function.py b/src/dolfinx_adjoint/function.py
index 7a1df8b..95b73af 100644
--- a/src/dolfinx_adjoint/function.py
+++ b/src/dolfinx_adjoint/function.py
@@ -38,10 +38,12 @@ def assign(value: typing.Union[numpy.inexact, float, int], function: _Function,
         if isinstance(value, (numpy.inexact, float, int)):
             function.x.array[:] = value
         elif isinstance(value, dolfinx.fem.Function):
-            assert value.function_space == function.function_space, (
-                "Function spaces of the value and function must match for assignment."
-            )
-            function.x.array[:] = value.x.array[:]
+            if value.function_space == function.function_space:
+                function.x.array[:] = value.x.array[:]
+            elif value.ufl_element().is_real and value.ufl_shape == ():
+                function.x.array[:] = value.x.array[0]
+            else:
+                raise ValueError("Function spaces of the value and function must match for assignment.")
         elif isinstance(value, ufl.core.expr.Expr):
             # Linear combination of functions, e.g., 2*u + 3*v
             assign_linear_combination(value, function)
diff --git a/tests/test_assign.py b/tests/test_assign.py
index 66db911..b1fce07 100644
--- a/tests/test_assign.py
+++ b/tests/test_assign.py
@@ -217,3 +217,90 @@ def test_assign_constant_derivative():
     dd = pyadjoint.AdjFloat(0.1)
     min_rate = pyadjoint.taylor_test(Jhat, d, dd)
     assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2)
+
+
+def test_assign_wrong_function_space(mesh_1D):
+    """Test that assigning a single function from a different space fails."""
+    # Create two different function spaces (different DoF counts)
+    V1 = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+    V2 = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+
+    f_wrong = Function(V2, name="f_wrong")
+    u = Function(V1, name="u")
+
+    # The assignment should fail due to mismatched array lengths/spaces
+    with pytest.raises(ValueError):
+        assign(f_wrong, u)
+
+
+def test_assign_linear_combination_wrong_function_space(mesh_1D):
+    """Test that assigning a linear combination with an incompatible term fails."""
+    V1 = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+    V2 = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+
+    f = Function(V1, name="f")
+    g_wrong = Function(V2, name="g_wrong")
+    u = Function(V1, name="u")
+
+    # 3 * f is valid for u, but subtracting g_wrong should trigger an error
+    with pytest.raises(ValueError):
+        assign(3 * f - g_wrong, u)
+
+
+@pytest.mark.xfail(strict=True, reason="Non-linear UFL expressions cannot be assigned directly via array operations.")
+def test_assign_non_linear_expression(mesh_1D):
+    """Test that assigning a non-linear expression fails."""
+    V = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+
+    f = Function(V, name="f")
+    u = Function(V, name="u")
+
+    # Create a non-linear expression (e.g., f squared)
+    non_linear_expr = f**2
+
+    # This should raise an error during the `extract_linear_combination` phase
+    # since `f**2` is a ufl.Power or ufl.Product, not a linear combination.
+    assign(non_linear_expr, u)
+
+
+def test_assign_real_function_equals_constant(mesh_1D):
+    """Test that assigning a Real space function behaves identically to a scalar constant,
+    both in the forward pass and the adjoint pass."""
+
+    # Standard spatial space and Real (global scalar) space
+    V = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+
+    # Note: Depending on your exact Basix/DOLFINx version, this might be ("R", 0)
+    R = dolfinx.fem.functionspace(mesh_1D, basix.ufl.real_element(mesh_1D.basix_cell(), value_shape=()))
+
+    target_real = Function(V, name="target_real")
+    target_const = Function(V, name="target_const")
+
+    val = 4.2
+    # Assign using a Function from a Real space
+    r_func = Function(R, name="r_func")
+    r_func.x.array[:] = val
+    assign(r_func, target_real)
+
+    # Assign using a raw float (which your block converts to AdjFloat)
+    assign(val, target_const)
+
+    # The resulting degrees of freedom should be exactly identical
+    np.testing.assert_allclose(
+        target_real.x.array,
+        target_const.x.array,
+        err_msg="Forward assignment of Real function and constant do not match.",
+    )
+
+    J = assemble_scalar(target_real**2 * ufl.dx)
+
+    # Test the sensitivity with respect to the Real function
+    rf = pyadjoint.ReducedFunctional(J, pyadjoint.Control(r_func))
+
+    # Create a perturbation direction in the Real space
+    h = Function(R, name="h")
+    h.x.array[:] = 0.75
+
+    # Verify the adjoint derivative is correct (should converge at rate ~ 2.0)
+    convergence_rate = pyadjoint.taylor_test(rf, r_func, h)
+    assert convergence_rate > 1.9, f"Taylor test failed with rate {convergence_rate}"

From 95f5d1bfebaac61fe0f7dd57c7c0ce39f13e5bfd Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 08:26:04 +0000
Subject: [PATCH 10/13] Use ufl checks to determine if scalar constant. Use new
 dagtraverser to extract linear combination. Create new floatifier traverser
 to convert the UFL expression for the weight to a float.

---
 .../blocks/function_assigner.py               |  17 +-
 src/dolfinx_adjoint/compat.py                 | 200 ++++++++++++++++++
 src/dolfinx_adjoint/utils.py                  | 170 +++++----------
 3 files changed, 257 insertions(+), 130 deletions(-)
 create mode 100644 src/dolfinx_adjoint/compat.py

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index 3bd881d..95583f5 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -13,15 +13,6 @@
 from ._vector import _vector
 
 
-def is_real_scalar_func(func: dolfinx.fem.Function) -> bool:
-    """Check if a function is a real scalar function.
-
-    Args:
-        func: The function to check.
-    """
-    return func.ufl_element().is_real and func.ufl_shape == ()
-
-
 def create_function_with_special_vector(func: _Function, name: str | None = None) -> _Function:
     """Create a new function with the same function space as `func` but with a special vector for adjoint computations.
 
@@ -80,7 +71,7 @@ def __init__(
             if isinstance(other, AdjFloat):
                 self._one = _Function(func.function_space, name="one", annotate=False)
                 self._one.x.array[:] = 1.0
-            elif isinstance(other, _Function) and is_real_scalar_func(other):
+            elif isinstance(other, _Function) and ufl.checks.is_scalar_constant_expression(other):
                 self._working_memory.append(create_function_with_special_vector(other, name="working_memory_2"))
                 self._one = _Function(func.function_space, name="one", annotate=False)
                 self._one.x.array[:] = 1.0
@@ -98,7 +89,7 @@ def __init__(
 
             # Allocate extra memory for adjoint computations if any of the operands are real functions
             for op in traverse_unique_terminals(other):
-                if isinstance(op, _Function) and is_real_scalar_func(op):
+                if isinstance(op, _Function) and ufl.checks.is_scalar_constant_expression(op):
                     self._working_memory.append(create_function_with_special_vector(op, name="working_memory_2"))
                     break
 
@@ -144,7 +135,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
             if isinstance(bo, AdjFloat):
                 return self._compute_adjoint_of_broadcast(adj_inputs[0], self._one)
             elif isinstance(bo, dolfinx.fem.Function):
-                if is_real_scalar_func(bo):
+                if ufl.checks.is_scalar_constant_expression(bo):
                     # Adjoint of a broadcast into a real function (constant stored as Function)
                     self._working_memory[2].x.array[0] = self._compute_adjoint_of_broadcast(adj_inputs[0], self._one)
                     return self._working_memory[2].x
@@ -170,7 +161,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
                 )
                 assign_linear_combination(diff_expr, self._working_memory[0])
                 return self._working_memory[0].x
-            elif isinstance(bo, dolfinx.fem.Function) and is_real_scalar_func(bo):
+            elif isinstance(bo, dolfinx.fem.Function) and ufl.checks.is_scalar_constant_expression(bo):
                 # Differentiate with respect to a real function (constant stored as Function)
                 # Create a perturbation direction in the Real space (value = 1.0)
                 assert len(self._working_memory) == 3, "Working memory not allocated for real function adjoint."
diff --git a/src/dolfinx_adjoint/compat.py b/src/dolfinx_adjoint/compat.py
new file mode 100644
index 0000000..b7fb53e
--- /dev/null
+++ b/src/dolfinx_adjoint/compat.py
@@ -0,0 +1,200 @@
+try:
+    from ufl.algorithms.extract_linear_combination import extract_linear_combination
+
+except ImportError:
+    # This is a workaround until dolfinx-adjoint only supports the version of UFL that has
+    # this feature
+
+    from functools import singledispatchmethod
+
+    import ufl
+    from ufl.corealg.dag_traverser import DAGTraverser
+
+    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):
+            """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)
+
+        @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)}")
+
+        @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
+        # ---------------------------------------------------------
+        @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.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
+            if ufl.checks.is_scalar_constant_expression(o):
+                return o
+            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.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):
+            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]
+            else:
+                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 | ufl.form.BaseForm,
+    ) -> list[tuple[ufl.core.expr.Expr, ufl.coefficient.BaseCoefficient]]:
+        """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
diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index d2e33d1..df09714 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -1,9 +1,12 @@
+from multiprocessing import process
 import typing
 
 import dolfinx
 import numpy
 import numpy.typing as npt
 import ufl
+from .compat import extract_linear_combination
+from functools import singledispatchmethod
 
 
 def function_from_vector(
@@ -44,127 +47,60 @@ class ad_kwargs(typing.TypedDict):
     """Whether to annotate the assignment in the adjoint tape."""
 
 
-def extract_scalar_value(scalar_expr: ufl.core.expr.Expr) -> float:
-    """Extract float from a scalar UFL expression."""
-    if isinstance(scalar_expr, (ufl.classes.IntValue, ufl.classes.FloatValue)):
-        return float(scalar_expr)
-    elif isinstance(scalar_expr, dolfinx.fem.Function):
-        # Check if it's a RealElement (constant stored as Function)
-        if scalar_expr.function_space.ufl_element().is_real and scalar_expr.ufl_shape == ():
-            return float(scalar_expr.x.array[0])
-        else:
-            raise ValueError(f"Cannot extract scalar from spatial Function: {scalar_expr}")
-    elif isinstance(scalar_expr, dolfinx.fem.Constant) and scalar_expr.ufl_shape == ():
-        val = scalar_expr.value
-        return float(val) if hasattr(val, "__float__") else float(val.item())
-    elif isinstance(scalar_expr, ufl.classes.ScalarValue):
-        return float(scalar_expr._value)
-    elif isinstance(scalar_expr, ufl.classes.Product):
-        result = 1.0
-        for op in scalar_expr.ufl_operands:
-            result *= extract_scalar_value(op)
-        return result
-    elif isinstance(scalar_expr, ufl.classes.Division):
-        num, den = scalar_expr.ufl_operands
-        return extract_scalar_value(num) / extract_scalar_value(den)
-    else:
-        raise ValueError(f"Cannot extract scalar from {type(scalar_expr)}: {scalar_expr}")
-
-
-def extract_function(expr) -> tuple[bool, dolfinx.fem.Function | None]:
-    """Recursively extract a Function from nested UFL expressions."""
-    if isinstance(expr, dolfinx.fem.Function):
-        is_real = expr.function_space.ufl_element().is_real
-        if is_real:
-            return (False, None)
-        return (False, expr)
-    elif isinstance(expr, (ufl.classes.Indexed, ufl.classes.ComponentTensor)):
-        return extract_function(expr.ufl_operands[0])
-    elif hasattr(expr, "ufl_operands"):
-        found_func = None
-        for op in expr.ufl_operands:
-            is_real, func = extract_function(op)
-            if func is not None:
-                if found_func is not None:
-                    raise ValueError(f"Non-linear expression detected: multiple spatial functions in {expr}")
-                found_func = func
-        return (False, found_func)
-    return (False, None)
-
-
-def extract_term(term: ufl.core.expr.Expr) -> tuple[float, dolfinx.fem.Function] | None:
-    """Extract (weight, function) from a single term."""
-    if isinstance(term, dolfinx.fem.Function):
-        is_real = term.function_space.ufl_element().is_real
-        if is_real:
-            return None
-        return (1.0, term)
-    elif isinstance(term, ufl.classes.ComponentTensor):
-        return extract_term(term.ufl_operands[0])
-    elif isinstance(term, ufl.classes.Indexed):
-        is_real, func = extract_function(term)
-        if func is None:
-            return None
-        return (1.0, func)
-    elif isinstance(term, ufl.classes.Product):
-        weight = 1.0
-        func = None
-        for op in term.ufl_operands:
-            is_real, extracted_func = extract_function(op)
-            if extracted_func is not None:
-                if func is not None:
-                    raise ValueError(f"Non-linear term detected: multiple spatial functions in {term}")
-                func = extracted_func
-            else:
-                weight *= extract_scalar_value(op)
-        return (weight, func) if func is not None else None
-    elif isinstance(term, ufl.classes.Division):
-        num, den = term.ufl_operands
-        denom_val = extract_scalar_value(den)
-        if isinstance(num, dolfinx.fem.Function):
-            is_real = num.function_space.ufl_element().is_real
-            if is_real:
-                return None
-            return (1.0 / denom_val, num)
-        elif isinstance(num, ufl.classes.Product):
-            result = extract_term(num)
-            return (result[0] / denom_val, result[1]) if result else None
-    return None
-
-
-def extract_linear_combination(expr: ufl.core.expr.Expr) -> list[tuple[float, dolfinx.fem.Function]]:
-    """Extract (weight, function) pairs from a UFL linear combination.
-
-    Analyzes expressions like: 0.5*u + 0.3*v + 0.2*w
-    Returns: [(0.5, u), (0.3, v), (0.2, w)]
-
-    :param expr: UFL expression (Sum, Product, or single Function)
-    :returns: List of (weight, function) tuples
-    """
-
-    # Parse the expression, flattening nested Sums recursively
-    summands: list[ufl.core.expr.Expr] | tuple[ufl.core.terminal.FormArgument, ...]
-    if isinstance(expr, ufl.classes.Sum):
-        summands = expr.ufl_operands
-    else:
-        summands = [expr]
-    terms = []
-    for summand in summands:
-        if isinstance(summand, ufl.classes.Sum):
-            # Recursively flatten nested Sum structures
-            terms.extend(extract_linear_combination(summand))
-        else:
-            result = extract_term(summand)
-            if result is not None:
-                terms.append(result)
-    return terms
-
-
 def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.Function) -> None:
+    """Assign a linear combination of functions to a function.
+
+    Arguments:
+        value: A linear combination of functions, e.g. `2*u + 3*v`.
+        function: The function to assign the linear combination to.
+    """
     pairs = extract_linear_combination(value)
     function.x.array[:] = 0.0
+    floatifier = Floatify()
     for weight, func in pairs:
         if not func.function_space == function.function_space:
             raise ValueError("Function spaces of all functions in the linear combination must match for assignment.")
-        function.x.array[:] += weight * func.x.array[:]
+        function.x.array[:] += floatifier.process(weight) * func.x.array[:]
     function.x.scatter_forward()
+
+
+class Floatify(ufl.corealg.dag_traverser.DAGTraverser):
+    """Traverser to convert a UFL expression into a float."""
+
+    def __init__(self, **kwargs):
+        """Convert a ufl expression into a float"""
+        super().__init__(**kwargs)
+
+    @singledispatchmethod
+    def process(self, o: ufl.classes.Expr, **kwargs):
+        return float(o)
+
+    @process.register(dolfinx.fem.Function)
+    def _(self, o, **kwargs):
+        if ufl.checks.is_scalar_constant_expression(o):
+            return o.x.array[0]
+        raise NotImplementedError(f"Unsupported UFL node type for floatification: {type(o)}")
+
+    @process.register(ufl.classes.Sum)
+    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        # operands is a tuple of the already-floatified children
+        return sum(operands)
+
+    @process.register(ufl.classes.Division)
+    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        # Division always has exactly two operands: numerator and denominator
+        return operands[0] / operands[1]
+
+    @process.register(ufl.classes.Power)
+    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        # Power has exactly two operands: base and exponent
+        return operands[0] ** operands[1]
+
+    @process.register(ufl.classes.Product)
+    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    def _(self, o, *operands, **kwargs):
+        # Product has exactly two operands: left and right
+        return operands[0] * operands[1]

From b09e59dda006806db17984df30b60333c7778479 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 08:26:24 +0000
Subject: [PATCH 11/13] Ruff formatting

---
 src/dolfinx_adjoint/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index df09714..03e392f 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -1,12 +1,12 @@
-from multiprocessing import process
 import typing
+from functools import singledispatchmethod
 
 import dolfinx
 import numpy
 import numpy.typing as npt
 import ufl
+
 from .compat import extract_linear_combination
-from functools import singledispatchmethod
 
 
 def function_from_vector(

From b579643d2829cbdc3a8805b930682290aa43e6de Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 08:29:10 +0000
Subject: [PATCH 12/13] Simplify check

---
 src/dolfinx_adjoint/function.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dolfinx_adjoint/function.py b/src/dolfinx_adjoint/function.py
index 95b73af..9927d10 100644
--- a/src/dolfinx_adjoint/function.py
+++ b/src/dolfinx_adjoint/function.py
@@ -40,7 +40,7 @@ def assign(value: typing.Union[numpy.inexact, float, int], function: _Function,
         elif isinstance(value, dolfinx.fem.Function):
             if value.function_space == function.function_space:
                 function.x.array[:] = value.x.array[:]
-            elif value.ufl_element().is_real and value.ufl_shape == ():
+            elif ufl.checks.is_scalar_constant_expression(value):
                 function.x.array[:] = value.x.array[0]
             else:
                 raise ValueError("Function spaces of the value and function must match for assignment.")

From 57e3b229bc0ca6ff4709290ecb7979c95773a38c Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Fri, 21 Aug 2026 08:35:24 +0000
Subject: [PATCH 13/13] Type fixing

---
 src/dolfinx_adjoint/compat.py |  2 +-
 src/dolfinx_adjoint/utils.py  | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/dolfinx_adjoint/compat.py b/src/dolfinx_adjoint/compat.py
index b7fb53e..e218266 100644
--- a/src/dolfinx_adjoint/compat.py
+++ b/src/dolfinx_adjoint/compat.py
@@ -192,7 +192,7 @@ def extract_linear_combination(
             weight, and the second element is the base UFL Coefficient (spatial function).
         """
         extractor = LinearCombinationExtractor()
-        final_result = extractor(expr)
+        final_result = extractor(expr)  # type: ignore
 
         if not isinstance(final_result, list):
             raise ValueError("Expression evaluated to a pure scalar, no spatial functions found.")
diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index 03e392f..9688661 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -5,6 +5,7 @@
 import numpy
 import numpy.typing as npt
 import ufl
+from ufl.corealg.dag_traverser import DAGTraverser
 
 from .compat import extract_linear_combination
 
@@ -64,7 +65,7 @@ def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.F
     function.x.scatter_forward()
 
 
-class Floatify(ufl.corealg.dag_traverser.DAGTraverser):
+class Floatify(DAGTraverser):
     """Traverser to convert a UFL expression into a float."""
 
     def __init__(self, **kwargs):
@@ -82,25 +83,25 @@ def _(self, o, **kwargs):
         raise NotImplementedError(f"Unsupported UFL node type for floatification: {type(o)}")
 
     @process.register(ufl.classes.Sum)
-    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    @DAGTraverser.postorder
     def _(self, o, *operands, **kwargs):
         # operands is a tuple of the already-floatified children
         return sum(operands)
 
     @process.register(ufl.classes.Division)
-    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    @DAGTraverser.postorder
     def _(self, o, *operands, **kwargs):
         # Division always has exactly two operands: numerator and denominator
         return operands[0] / operands[1]
 
     @process.register(ufl.classes.Power)
-    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    @DAGTraverser.postorder
     def _(self, o, *operands, **kwargs):
         # Power has exactly two operands: base and exponent
         return operands[0] ** operands[1]
 
     @process.register(ufl.classes.Product)
-    @ufl.corealg.dag_traverser.DAGTraverser.postorder
+    @DAGTraverser.postorder
     def _(self, o, *operands, **kwargs):
         # Product has exactly two operands: left and right
         return operands[0] * operands[1]