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 5cb0c18..95583f5 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/src/dolfinx_adjoint/blocks/function_assigner.py @@ -2,38 +2,96 @@ 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 ..types.function import Function as _Function +from ..utils import assign_linear_combination, extract_linear_combination, function_from_vector from ._vector import _vector +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. + + 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: typing.Union[np.inexact, int, float], - func: dolfinx.fem.Function, + 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): + self._working_memory.append(create_function_with_special_vector(func, 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) - 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 + # 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 + 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 else: - raise NotImplementedError("We should not get here!") + 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) + 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) + + # 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 ufl.checks.is_scalar_constant_expression(op): + self._working_memory.append(create_function_with_special_vector(op, name="working_memory_2")) + break def _replace_with_saved_output(self): if self.expr is None: @@ -53,85 +111,79 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): expr = self._replace_with_saved_output() return expr, adj_input_func + @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): + 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() + 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): + 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 + assert len(adj_inputs) == 1 + if isinstance(bo, AdjFloat): + return self._compute_adjoint_of_broadcast(adj_inputs[0], self._one) + elif isinstance(bo, dolfinx.fem.Function): + 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 + if bo.function_space != prepared.function_space: + raise ValueError( + "Function spaces of the block variable and prepared function must match for adjoint evaluation." ) - 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 func.function_space == prepared.function_space - vec = _vector( - prepared.x.index_map, prepared.x.block_size, func.function_space, dtype=prepared.x.array.dtype + 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." + ) + else: + raise NotImplementedError(f"Adjoint for {block_variable=} not implemented.") + else: + # Linear combination + expr, adj_input_func = prepared + 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) + ) + assign_linear_combination(diff_expr, self._working_memory[0]) + return self._working_memory[0].x + 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." + direction = self._working_memory[2] + 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) ) - vec.array[:] = prepared.x.array[:] - return vec + # Evaluate the derivative at the DOFs of the target space V + 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) + 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.") - # 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: @@ -143,12 +195,14 @@ 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) - dudmi = dolfinx.fem.Function(block_variable.output.function_space) + dudm = self._working_memory[0] + dudm.x.array[:] = 0.0 + dudmi = self._working_memory[1] 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 +235,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/compat.py b/src/dolfinx_adjoint/compat.py new file mode 100644 index 0000000..e218266 --- /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) # type: ignore + + 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/function.py b/src/dolfinx_adjoint/function.py new file mode 100644 index 0000000..9927d10 --- /dev/null +++ b/src/dolfinx_adjoint/function.py @@ -0,0 +1,53 @@ +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, function, 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): + if value.function_space == function.function_space: + function.x.array[:] = value.x.array[:] + 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.") + 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 165c36c..d5f82ce 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, function_from_vector, gather +from ..utils import function_from_vector, gather class Function(dolfinx.fem.Function, FloatingType): @@ -62,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): @@ -189,7 +190,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 @@ -203,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` @@ -256,36 +262,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, function, 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[:] - 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 5f0b5d1..9688661 100644 --- a/src/dolfinx_adjoint/utils.py +++ b/src/dolfinx_adjoint/utils.py @@ -1,8 +1,13 @@ import typing +from functools import singledispatchmethod import dolfinx import numpy import numpy.typing as npt +import ufl +from ufl.corealg.dag_traverser import DAGTraverser + +from .compat import extract_linear_combination def function_from_vector( @@ -41,3 +46,62 @@ 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 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[:] += floatifier.process(weight) * func.x.array[:] + function.x.scatter_forward() + + +class Floatify(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) + @DAGTraverser.postorder + def _(self, o, *operands, **kwargs): + # operands is a tuple of the already-floatified children + return sum(operands) + + @process.register(ufl.classes.Division) + @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) + @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) + @DAGTraverser.postorder + def _(self, o, *operands, **kwargs): + # Product has exactly two operands: left and right + return operands[0] * operands[1] diff --git a/tests/test_assign.py b/tests/test_assign.py index d72d177..b1fce07 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 @@ -159,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}"