DOLFINx-adjoint Blocks API

The following is the API of the blocks implemented in the DOLFINx-adjoint package, which is used within the pyadjoint framework.

class dolfinx_adjoint.blocks.AssembleBlock(form: Form, ad_block_tag: str | None = None, jit_options: dict | None = None, form_compiler_options: dict | None = None, entity_maps: dict[Mesh, ndarray[Any, dtype[int32]]] | None = None)

Block for assembling a symbolic UFL form into a tensor.

Parameters:

• form – The UFL form to assemble.

• ad_block_tag – Tag for the block in the adjoint tape.

• jit_options – Dictionary of options for JIT compilation.

• form_compiler_options – Dictionary of options for the form compiler.

• entity_maps – Dictionary mapping meshes to entity maps for assembly.

compute_action_adjoint(adj_input: float | Vector, arity_form: int, form: Form | None = None, c_rep: Coefficient | Constant | None = None, space: FunctionSpace | None = None, dform: Form | None = None)

This computes the action of the adjoint of the derivative of form wrt c_rep on adj_input.

In other words, it returns:

\[\left\langle\left(\frac{\partial form}{\partial c_{rep}}\right)^*, adj_{input} \right\rangle\]

• If form has arity 0, then \(\frac{\partial form}{\partial c_{rep}}\) is a 1-form and adj_input a float, we can simply use the * operator.

• If form has arity 1 then \(\frac{\partial form}{\partial c_{rep}}\) is a 2-form and we can symbolically take its adjoint and then apply the action on adj_input, to finally assemble the result.

Parameters:

• adj_input – The input to the adjoint operation, typically a scalar or vector.

• arity_form – The arity of the form, i.e., 0 for scalar, 1 for vector, 2 for matrix etc.

• form – The UFL form to differentiate if dform is not provided.

• c_rep – The coefficient or constant with respect to which the derivative is taken.

• space – The function space associated with the c_rep to form an ufl.Argument in.

• dform – Pre-computed derivative form, \(\frac{\partial form}{\partial c_{rep}}\).

evaluate_adj_component(inputs, adj_inputs, block_variable, idx, prepared=None)

This method should be overridden.

The method should implement a routine for evaluating the adjoint of the block that corresponds to one dependency. If one considers the adjoint action a vector right multiplied with the Jacobian matrix, then this method should return one entry in the resulting product, where the entry returned is decided by the argument idx.

Parameters:

• inputs (list) – A list of the saved input values, determined by the dependencies list.

• adj_inputs (list) – A list of the adjoint input values, determined by the outputs list.

• block_variable (BlockVariable) – The block variable of the dependency corresponding to index idx.

• idx (int) – The index of the component to compute.

• prepared (object) – Anything returned by the prepare_evaluate_adj method. Default is None.

Returns:

The resulting product.

Return type:

An object of a type consistent with the adj_value type of block_variable

evaluate_hessian_component(inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None)

This method must be overridden.

The method should implement a routine for evaluating the hessian of the block. It is preferable that a “Forward-over-Reverse” scheme is used. Thus the hessians are evaluated in reverse (starting with the last block on the tape).

evaluate_tlm_component(inputs, tlm_inputs, block_variable, idx, prepared=None)

This method should be overridden.

The method should implement a routine for computing the tangent linear model of the block that corresponds to one output. If one considers the tangent linear action as a Jacobian matrix multiplied with a vector, then this method should return one entry in the resulting product, where the entry returned is decided by the argument idx.

Parameters:

• inputs (list) – A list of the saved input values, determined by the dependencies list.

• tlm_inputs (list) – A list of the tlm input values, determined by the dependencies list.

• block_variable (BlockVariable) – The block variable of the output corresponding to index idx.

• idx (int) – The index of the component to compute.

• prepared (object) – Anything returned by the prepare_evaluate_tlm method. Default is None.

Returns:

The resulting product.

Return type:

An object of the same type as block_variable.saved_output

prepare_evaluate_adj(inputs, adj_inputs, relevant_dependencies)

Runs preparations before evalute_adj_component is ran.

The return value is supplied to each of the subsequent evaluate_adj_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• adj_inputs – The adjoint inputs

• relevant_dependencies – A list of the relevant block variables for evaluate_adj_component.

Returns:

Anything. The returned value is supplied to evaluate_adj_component

prepare_evaluate_hessian(inputs, hessian_inputs, adj_inputs, relevant_dependencies)

Runs preparations before evalute_hessian_component is ran for each relevant dependency.

The return value is supplied to each of the subsequent evaluate_hessian_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• hessian_inputs – The hessian inputs

• adj_inputs – The adjoint inputs

• relevant_dependencies – A list of the relevant block variables for evaluate_hessian_component.

Returns:

Anything. The returned value is supplied to evaluate_hessian_component

prepare_evaluate_tlm(inputs, tlm_inputs, relevant_outputs)

Runs preparations before evalute_tlm_component is ran.

The return value is supplied to each of the subsequent evaluate_tlm_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• tlm_inputs – The tlm inputs

• relevant_outputs – A list of the relevant block variables for evaluate_tlm_component.

Returns:

Anything. The returned value is supplied to evaluate_tlm_component

prepare_recompute_component(inputs, relevant_outputs)

Runs preparations before recompute_component is ran.

The return value is supplied to each of the subsequent recompute_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• relevant_outputs – A list of the relevant block variables for recompute_component.

Returns:

Anything. The returned value is supplied to recompute_component

recompute_component(inputs, block_variable, idx, prepared)

This method must be overridden.

The method should implement a routine for recomputing one output of the block in the forward computations. The output to recompute is determined by the idx argument, which corresponds to the index of the output in the outputs list. If the block only has a single output, then idx will always be 0.

Parameters:

• inputs (list) – A list of the saved input values, determined by the dependencies list.

• block_variable (BlockVariable) – The block variable of the output corresponding to index idx.

• idx (int) – The index of the output to compute.

• prepared (object) – Anything returned by the prepare_recompute_component method. Default is None.

Returns:

An object of the same type as block_variable.checkpoint which is determined by OverloadedType._ad_create_checkpoint (often the same as block_variable.saved_output): The new output.

class dolfinx_adjoint.blocks.FunctionAssignBlock(other: inexact | int | float, func: Function, ad_block_tag: str | None = None)

evaluate_adj_component(inputs, adj_inputs, block_variable, idx, prepared=None)

This method should be overridden.

The method should implement a routine for evaluating the adjoint of the block that corresponds to one dependency. If one considers the adjoint action a vector right multiplied with the Jacobian matrix, then this method should return one entry in the resulting product, where the entry returned is decided by the argument idx.

Parameters:

• inputs (list) – A list of the saved input values, determined by the dependencies list.

• adj_inputs (list) – A list of the adjoint input values, determined by the outputs list.

• block_variable (BlockVariable) – The block variable of the dependency corresponding to index idx.

• idx (int) – The index of the component to compute.

• prepared (object) – Anything returned by the prepare_evaluate_adj method. Default is None.

Returns:

The resulting product.

Return type:

An object of a type consistent with the adj_value type of block_variable

evaluate_hessian_component(inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None)

This method must be overridden.

The method should implement a routine for evaluating the hessian of the block. It is preferable that a “Forward-over-Reverse” scheme is used. Thus the hessians are evaluated in reverse (starting with the last block on the tape).

evaluate_tlm_component(inputs, tlm_inputs, block_variable, idx, prepared=None)

This method should be overridden.

The method should implement a routine for computing the tangent linear model of the block that corresponds to one output. If one considers the tangent linear action as a Jacobian matrix multiplied with a vector, then this method should return one entry in the resulting product, where the entry returned is decided by the argument idx.

Parameters:

• inputs (list) – A list of the saved input values, determined by the dependencies list.

• tlm_inputs (list) – A list of the tlm input values, determined by the dependencies list.

• block_variable (BlockVariable) – The block variable of the output corresponding to index idx.

• idx (int) – The index of the component to compute.

• prepared (object) – Anything returned by the prepare_evaluate_tlm method. Default is None.

Returns:

The resulting product.

Return type:

An object of the same type as block_variable.saved_output

prepare_evaluate_adj(inputs, adj_inputs, relevant_dependencies)

Runs preparations before evalute_adj_component is ran.

The return value is supplied to each of the subsequent evaluate_adj_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• adj_inputs – The adjoint inputs

• relevant_dependencies – A list of the relevant block variables for evaluate_adj_component.

Returns:

Anything. The returned value is supplied to evaluate_adj_component

prepare_evaluate_hessian(inputs, hessian_inputs, adj_inputs, relevant_dependencies)

Runs preparations before evalute_hessian_component is ran for each relevant dependency.

The return value is supplied to each of the subsequent evaluate_hessian_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• hessian_inputs – The hessian inputs

• adj_inputs – The adjoint inputs

• relevant_dependencies – A list of the relevant block variables for evaluate_hessian_component.

Returns:

Anything. The returned value is supplied to evaluate_hessian_component

prepare_evaluate_tlm(inputs, tlm_inputs, relevant_outputs)

Runs preparations before evalute_tlm_component is ran.

The return value is supplied to each of the subsequent evaluate_tlm_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• tlm_inputs – The tlm inputs

• relevant_outputs – A list of the relevant block variables for evaluate_tlm_component.

Returns:

Anything. The returned value is supplied to evaluate_tlm_component

prepare_recompute_component(inputs, relevant_outputs)

Runs preparations before recompute_component is ran.

The return value is supplied to each of the subsequent recompute_component calls. This method is intended to be overridden for blocks that require such preparations, by default there is none.

Parameters:

• inputs – The values of the inputs

• relevant_outputs – A list of the relevant block variables for recompute_component.

Returns:

Anything. The returned value is supplied to recompute_component

recompute_component(inputs, block_variable, idx, prepared)

This method must be overridden.

The method should implement a routine for recomputing one output of the block in the forward computations. The output to recompute is determined by the idx argument, which corresponds to the index of the output in the outputs list. If the block only has a single output, then idx will always be 0.

Parameters:

• inputs (list) – A list of the saved input values, determined by the dependencies list.

• block_variable (BlockVariable) – The block variable of the output corresponding to index idx.

• idx (int) – The index of the output to compute.

• prepared (object) – Anything returned by the prepare_recompute_component method. Default is None.

Returns:

An object of the same type as block_variable.checkpoint which is determined by OverloadedType._ad_create_checkpoint (often the same as block_variable.saved_output): The new output.