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,
)
)