graph.dag

class dolfinx_graph_ad.graph.dag.Graph

Bases: object

Class for the computational graph

The computational graph is a directed acyclic graph (DAG) that represents the operations, objects and dependencies in a forward simulation in DOLFINx.

Example

The graph object can be initialised and the default nodes and edges can be added as follows: >>> graph = Graph() >>> node1 = Node(object_1, name = “Node 1”) >>> node2 = Node(object_2, name = “Node 2”) >>> edge = Edge(node1, node2) >>> graph.add_node(node1) >>> graph.add_node(node2) >>> graph.add_edge(edge)

The derivatives can then be caculated using the backpropagation method: >>> graph.backprop(object_2, object_1)

add_edge(edge: Edge)

Add an edge to the graph

Parameters:: edge (Edge) – The edge to be added to the graph
Raises:: ValueError – If an edge already connects the predecessor to the successor

Note

Both nodes of the edge have to be part of the graph, since the graph is what keeps them alive: an edge references its successor weakly, so that the graph stays free of reference cycles.

add_node(node: Node)

Add a node to the graph

Parameters:: node (Node) – The node to be added to the graph

Perform backpropagation in the graph

The gradients of all previous calls are reset before the propagation is started.

Multiple outputs produce one gradient of their weighted sum per control. For scalar objectives J_j and seeds s_j, the gradient for control m_i is sum_j s_j * dJ_j/dm_i. To obtain separate gradients for each objective, call backprop once per objective with all requested controls.

If controls are given, only the edges on the paths from the controls to the outputs are executed, in a single propagation. The propagation therefore ends at the controls, whose gradients are stored and returned, even if a control is an intermediate node of the graph. A control that depends on another control stores its gradient and also passes it on, so that each control receives its total derivative.

If no control is given, all edges the outputs depend on are executed and the propagation continues until it reaches the leaves of this dependency subgraph. Gradients are then stored only in these leaves; intermediate nodes are passed through without storing their gradients. The gradients can be read afterwards from Node.grad of the respective nodes.

Parameters:

outputs – The recorded object to differentiate, or its exact Node, or a sequence of them. An object selects its latest version; a Node selects that version.
inputs – A control object or Node, or a sequence of control objects and Nodes in the desired return order. Defaults to None. If None, the propagation is carried out down to the dependency leaves of the outputs.
grad_outputs (float or PETSc.Vec, or a sequence of them, optional) – The seed, i.e. the adjoint value the propagation is started with, or one per output.

Returns:

One gradient of the weighted output sum per requested control, in control order. An entry is None for a control no output depends on. Without controls, returns None and stores leaf gradients. A PETSc.Vec gradient has the control’s vector layout, including ghost entries, but only its entries owned by the calling rank are valid.

Return type:

tuple or None

Raises:

TypeError – If a control does not represent a numerical value and can therefore not store a gradient.

Example

Given two recorded scalar objectives and three recorded controls:

>>> controls = (m1, m2, m3)
>>> summed = graph_.backprop((J1, J2), controls)
>>> weighted = graph_.backprop((J1, J2), controls, grad_outputs=(1.0, 0.5))

summed and weighted each contain three gradients.

capture(value)

Reference the exact version an operation reads, see save().

Parameters:: value – A Function or Constant, or an exact node owned by this graph. Objects resolve to their latest recorded version; exact nodes retain their version. Untracked values remain constant for this graph and are neither registered nor copied.
Returns:: A weak reference to the node, whose snapshot holds its saved value, or the value itself when it is untracked.

property edges: ValuesView[Edge]

The recorded edges, in insertion order.

Returns:: A live view of all edges. Add edges with add_edge().
Return type:: ValuesView

get_dependencies(end_nodes: Sequence[AbstractNode]) → set[Edge]

Get the edges of all operations the end nodes are the result of

Parameters:: end_nodes (Sequence[AbstractNode]) – The end nodes, whose dependencies are collected together.
Returns:: The edges into the end nodes and into the nodes they depend on, see get_path().
Return type:: set[Edge]

get_edge(predecessor: Node, successor: Node)

Get an edge from the graph

Parameters:

predecessor (Node) – The predecessor node of the edge
successor (Node) – The successor node of the edge

Returns:

The edge with the given predecessor and successor

Return type:

Edge

get_node(value, version=None) → AbstractNode | None

Get the node of a recorded object or an exact node owned by this graph.

Parameters:

value – The recorded object, mesh or exact node.
version (int, optional) – The recorded version of an object. Defaults to its latest version. Must be None when value is an exact node, which already identifies its version.

Returns:

The matching node, or None if it is not recorded in this graph. An exact node must match by identity.

Return type:

AbstractNode or None

Raises:

TypeError – If a non-None version is supplied with an exact node.

get_path(start_nodes: Sequence[AbstractNode], end_nodes: Sequence[AbstractNode]) → set[Edge]

Get the edges on the paths from the start nodes to the end nodes

Parameters:

start_nodes (Sequence[AbstractNode]) – The start nodes, whose paths are collected together.
end_nodes (Sequence[AbstractNode]) – The end nodes, whose paths are collected together.

Returns:

The edges along the paths.

Return type:

set[Edge]

property nodes: ValuesView[AbstractNode]

The recorded nodes, in insertion order.

Returns:: A live view of all versions. Add nodes with add_node().
Return type:: ValuesView

property nx_graph: DiGraph

The cached NetworkX representation of the graph.

Nodes themselves are the keys. The representation is cached and cleared when the graph is deleted, so it is empty if kept past the graph.

Returns:: The networkx graph representation of the graph
Return type:: nx_graph (nx.DiGraph)

print(detailed=False) → None

Print the graph structure

Parameters:: detailed (bool, optional) – Whether to print the graph in detailed mode. Defaults to False.

release(): Release the values saved in the graph

reset_grads(): Reset the gradients in the graph

save(node: Node)

Save a new version before any later recorded call overwrites it.

Parameters:: node – The Function or Constant version to save.

Note

The saved copy includes owned and ghost entries. Every rank saves and releases corresponding versions in the same order.

track(value) → Node

Register and save an existing Function or Constant at its current value.

Parameters:: value – A DOLFINx Function or Constant owned by the caller.
Returns:: The control node. A value that is already recorded in this graph returns its latest node unchanged.
Return type:: Node
Raises:: TypeError – If the value is not a DOLFINx Function or Constant.

Note

Saving preserves owned and ghost entries locally.

visualise(ax=None, style: str | Callable = 'planar', print_edge_labels=True, path=None)

Visualise the graph

Parameters:

ax (matplotlib.axes.Axes, optional) – The axes to draw the graph in. If None, a new figure and axes are created. Defaults to None.
style (str, optional) – The style of the visualisation. Defaults to “planar”.
print_edge_labels (bool, optional) – Whether to print the edge labels. Defaults to True.
path (set, optional) – The set of edges to be highlighted in the visualisation. Defaults to None.

Returns:

The axes containing the visualisation of the graph.

Return type:

matplotlib.axes.Axes