Source code for dolfinx_adjoint.assembly

import typing

from mpi4py import MPI

import dolfinx
import numpy
import numpy.typing as npt
import ufl
from pyadjoint.overloaded_type import create_overloaded_object
from pyadjoint.tape import annotate_tape, get_working_tape, stop_annotating

from .blocks.assembly import AssembleBlock


[docs] def assemble_scalar(form: ufl.Form, **kwargs): """Assemble as scalar value from a form. Args: form: Symbolic form (UFL) to assemble. kwargs: Keyword arguments to pass to the assembly routine. Includes ``"ad_block_tag"`` to tag the block in the adjoint tape, ``"annotate"`` to control whether the assembly is annotated in the adjoint tape, ``"jit_options"`` for JIT compilation options, and ``"form_compiler_options"`` for form compiler options and ``"entity_map"`` for assembling with Arguments and coefficients form meshes that has some relation. """ ad_block_tag = kwargs.pop("ad_block_tag", None) annotate = annotate_tape(kwargs) with stop_annotating(): compiled_form = dolfinx.fem.form( form, jit_options=kwargs.pop("jit_options", None), form_compiler_options=kwargs.pop("form_compiler_options", None), entity_maps=kwargs.pop("entity_maps", None), ) local_output = dolfinx.fem.assemble_scalar(compiled_form) comm = compiled_form.mesh.comm output = comm.allreduce(local_output, op=MPI.SUM) assert isinstance(output, float) output = create_overloaded_object(output) if annotate: block = AssembleBlock(form, ad_block_tag=ad_block_tag) tape = get_working_tape() tape.add_block(block) block.add_output(output.block_variable) return output
[docs] def error_norm( u_ex: ufl.core.expr.Expr, u: ufl.core.expr.Expr, norm_type=typing.Literal["L2", "H1"], jit_options: typing.Optional[dict] = None, form_compiler_options: typing.Optional[dict] = None, entity_map: typing.Optional[dict[dolfinx.mesh.Mesh, npt.NDArray[numpy.int32]]] = None, ad_block_tag: typing.Optional[str] = None, annotate: bool = True, ) -> float: """Compute the error norm between the exact solution and the computed solution. Args: u_ex: The exact solution as a UFL expression. u: The computed solution as a UFL expression. norm_type: The type of norm to compute, either "L2" or "H1". jit_options: Optional JIT compilation options. form_compiler_options: Optional form compiler options. entity_map: Optional mapping from mesh entities to submesh entities. ad_block_tag: Optional tag for the block in the adjoint tape. annotate: Whether to annotate the assignment in the adjoint tape. Returns: The computed error norm as a float. """ diff = u_ex - u norm = ufl.inner(diff, diff) * ufl.dx if norm_type == "H1": norm += ufl.inner(ufl.grad(diff), ufl.grad(diff)) * ufl.dx return numpy.sqrt( assemble_scalar( norm, jit_options=jit_options, form_compiler_options=form_compiler_options, entity_maps=entity_map, ad_block_tag=ad_block_tag, annotate=annotate, ) )