fem.petsc

Record each linear or nonlinear solve at its own dependencies and values.

class dolfinx_graph_ad.fem.petsc.AdjointProblem(a: Form, b: Vec, u: Function, bcs=None, petsc_options: dict | None = None, petsc_options_prefix: str | None = None, **kwargs)

Bases: LinearProblem

OVERLOADS: dolfinx.fem.petsc.LinearProblem.

Solve an adjoint bilinear form against a vector with homogeneous Dirichlet conditions. DOLFINx owns and destroys the matrix, vectors, and KSP.

solve() → Function

Solve using the supplied vector, without RHS assembly or lifting.

Returns:

The solution function with updated ghost entries.

class dolfinx_graph_ad.fem.petsc.LinearProblem(*args, **kwargs)

Bases: LinearProblem

OVERLOADS: dolfinx.fem.petsc.LinearProblem. Linear problem class for solving the linear problem

The overloaded class modifies the initialization of the LinearProblem to keep track of the dependencies in the computational graph and the adjoint equations. The original functionality is kept.

solve(*args, **kwargs)

OVERLOADS: dolfinx.fem.petsc.LinearProblem.solve() Solve the linear problem into the function u, which is returned.

Parameters:

  • args – Arguments to dolfinx.fem.petsc.LinearProblem.solve()

  • kwargs – Keyword arguments to dolfinx.fem.petsc.LinearProblem.solve()

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

  • version (int, optional) – The version of the solution in the graph. Defaults to the version following its latest recorded one, or 0 if unrecorded.

Returns:

The solution function u after solving the linear problem.

Return type:

fem.Function

Raises:

ValueError – If the problem was not recorded in the graph by its constructor, since the solve then has no residual to differentiate.

class dolfinx_graph_ad.fem.petsc.LinearProblemNode(object, arguments, forms, options, **kwargs)

Bases: ProblemNode

A recorded version of dolfinx.fem.petsc.LinearProblem.

property residual

The residual F(u) = a(u) - L.

class dolfinx_graph_ad.fem.petsc.NonlinearProblem(*args, **kwargs)

Bases: NonlinearProblem

OVERLOADS: dolfinx.fem.petsc.NonlinearProblem. Nonlinear problem class for solving the non-linear problem

The overloaded class modifies the initialization of the NonlinearProblem to keep track of the dependencies in the computational graph and the adjoint equations. The original functionality is kept.

solve(*args, **kwargs)

OVERLOADS: dolfinx.fem.petsc.NonlinearProblem.solve()

Solve the non-linear problem into the function u, which is returned.

Parameters:

  • args – Arguments to dolfinx.fem.petsc.NonlinearProblem.solve()

  • kwargs – Keyword arguments to dolfinx.fem.petsc.NonlinearProblem.solve()

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

  • version (int, optional) – The version of the solution in the graph. Defaults to the version following its latest recorded one, or 0 if unrecorded.

Returns:

The solution function u after solving the nonlinear problem.

Return type:

fem.Function

Raises:

ValueError – If the problem was not recorded in the graph by its constructor, since the solve then has no residual to differentiate.

class dolfinx_graph_ad.fem.petsc.NonlinearProblemNode(object, arguments, forms, options, **kwargs)

Bases: ProblemNode

A recorded nonlinear problem version.

property residual

The supplied residual F(u) = 0, with no previous-state stand-ins.

class dolfinx_graph_ad.fem.petsc.ProblemNode(object, arguments, forms, options, **kwargs)

Bases: AbstractNode

A recorded problem version with its own dependencies and captured values.

Variables:

  • residual (ufl.Form) – The residual, built on first access during preparation.

  • values (dict) – Captured input versions, or None before preparation.

  • adjoint_form_kwargs (dict) – Arguments used to compile derivative forms.

capture_inputs(_graph, previous_values)

Capture controls and boundary data before the native solve.

Parameters:

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

  • previous_values – Captured pre-solve values keyed by stand-in coefficients.

release()

Release the problem, its recording metadata and its adjoint.

solve_adjoint(seed)

Solve (∂F/∂u)ᵀ λ = −seed at this forward solve’s saved values.

Each control edge creates its own adjoint problem. Construction and release are collective in the same edge order on every mesh rank; solution ghosts are updated before the edge assembles its sensitivity.

Parameters:

seed (PETSc.Vec) – The adjoint value of the solution. It is not modified.

Returns:

The residual at saved values, adjoint solution and saved value mapping.

Return type:

tuple

Raises:

RuntimeError – If the solve has not been performed, so that no values were copied.

class dolfinx_graph_ad.fem.petsc.Problem_Boundary_Edge(predecessor: Node, successor: Node, ctx=None)

Bases: Edge

Sensitivity on the owned dofs where this boundary condition survives.

calculate_adjoint(value)

Return the lifting sensitivity plus the direct constrained-state seed.

Differentiate λᵀF with respect to the solved state, then add the seed because u = g on constrained dofs. Restrict to owned dofs surviving later conditions; the boundary-condition edge maps these entries back to its control.

Parameters:

value – The accumulated solution seed, left unchanged.

Returns:

The state-layout gradient restricted to surviving owned dofs.

Return type:

PETSc.Vec

class dolfinx_graph_ad.fem.petsc.Problem_Coefficient_Edge(predecessor: Node, successor: Node, ctx=None)

Bases: Edge

Sensitivity of a solve to a recorded coefficient.

calculate_adjoint(value)

Return the coefficient derivative of the pairing λᵀF.

Parameters:

value – The accumulated solution seed.

Returns:

The gradient, with valid owned entries.

Return type:

PETSc.Vec

class dolfinx_graph_ad.fem.petsc.Problem_Constant_Edge(predecessor: Node, successor: Node, ctx=None)

Bases: Edge

Sensitivity of a solve to a recorded Constant.

calculate_adjoint(value)

Differentiate λᵀF through a real-element stand-in for the Constant.

Parameters:

value – The accumulated solution seed.

Returns:

The scalar or component gradient.

Return type:

float or complex or PETSc.Vec

class dolfinx_graph_ad.fem.petsc.SolveNode(object: Any, **kwargs)

Bases: FunctionNode

The saved solution of one recorded solve.

Its incoming edge identifies the problem version that produced the solution.