fem.forms

class dolfinx_graph_ad.fem.forms.FormNode(object: Any, ufl_form: Form, version=0, **kwargs)

Bases: AbstractNode

Node for the operation dolfinx.fem.form().

In order to compile the form from the ufl form in the forward pass, the ufl form needs to be saved in the node, together with the arguments it was compiled with.

Variables:

• object (Any) – The object that is being compiled.

• ufl_form (ufl.form.Form) – The ufl form that is being compiled.

• kwargs (dict) – The arguments of dolfinx.fem.form() the form was compiled with, keyed by parameter name and without the form itself.

• values (dict) – Copies of the coefficients and constants of the form, keyed by the originals, taken when the form was assembled. None before.

bind(_graph) → FormNode

Record a fresh form version with the inputs this assembly reads.

Parameters:

_graph (Graph) – The graph the assembly is recorded in.

Returns:

The node of the form for this assembly.

Return type:

FormNode

recorded()

The form at the values it was assembled with, and these values.

Returns:

The form and the values it is evaluated at, keyed by the originals.

Return type:

tuple

Raises:

RuntimeError – If the form has not been assembled, so that no values were copied.

release()

Releases the compiled form, the ufl form, the arguments it was compiled from and the values it was assembled with.

class dolfinx_graph_ad.fem.forms.Form_Coefficient_Edge(predecessor: Node, successor: Node, ctx=None)

Bases: Edge

The scalar form’s reverse derivative at its recorded input values.

Variables:

ctx – The original coefficient differentiated by this edge.

calculate_adjoint(value)

The method provides the adjoint equation for the derivative of the form with respect to a coefficient.

Returns:

The accumulated gradient up to this point in the computational graph. Only its entries owned by the calling rank are valid.

Return type:

(PETSc.Vec)

class dolfinx_graph_ad.fem.forms.Form_Constant_Edge(predecessor: Node, successor: Node, ctx=None)

Bases: Edge

Edge providing the adjoint equation for the derivative of the form with respect to a constant.

Variables:

ctx – The original Constant differentiated by this edge.

calculate_adjoint(value)

The method provides the adjoint equation for the derivative of the form with respect to a constant.

Returns:

The accumulated gradient up to this point in the computational graph, a number for a scalar constant.

Return type:

float or complex or PETSc.Vec

dolfinx_graph_ad.fem.forms.form(*args, **kwargs)

OVERLOADS: dolfinx.fem.form(). Create a Form or an array of Forms.

The overloaded function adds the functionality to keep track of the dependencies in the computational graph and the adjoint equations. The original functionality is kept.

Parameters:

• args – Arguments to dolfinx.fem.form().

• kwargs – Keyword arguments to dolfinx.fem.form().

• graph (graph, optional) – An additional keyword argument to specifier whether the assemble operation should be added to the graph. If not present, the original functionality of dolfinx is used without any additional functionalities.

Returns:

Compiled finite element Form.

Note

When a form is compiled from a UFL form, the dependencies in the symbolic equation is lost. The resulting compiled form does not support symbolic differentiation. The graph and the custom edges are used to keep track of the dependencies and the adjoint equations.

Note

The arguments the form is compiled with are recorded together with it, so that the compilations the node and its edges perform later are the compilation the caller asked for.