Graph data structures

Data structures for representing graphs and graph batches.

This module provides object-based graph classes (Node, Edge, Graph) and a flat batched container (GraphTuple) used by high-level GraphNet APIs.

class tf_gnns.tfgnns_datastructures.Node(node_attr_tensor)

Bases: object

Graph node with a feature tensor.

Parameters:

node_attr_tensor – Tensor-like node attributes with at least rank 2. The first dimension is usually batch-like in object graph mode.

get_state()

set_tensor(tensor)

copy()

class tf_gnns.tfgnns_datastructures.Edge(edge_attr_tensor, node_from, node_to)

Bases: object

Directed edge connecting two Node objects.

Parameters:

• edge_attr_tensor – Tensor-like edge attributes.

• node_from – Source node.

• node_to – Destination node.

set_tensor(edge_tensor)

copy(nodes_correspondence)

class tf_gnns.tfgnns_datastructures.Graph(nodes, edges, global_attr=None, NO_VALIDATION=True)

Bases: object

Object graph made of node and edge instances.

Parameters:

• nodes – List of Node instances.

• edges – List of Edge instances.

• global_attr – Optional graph-level attributes.

• NO_VALIDATION – If False, run connectivity validation checks.

is_equal_by_value(g2)

Checks if the graphs have the same values for node and edge attributes

compare_connectivity(g2)

Checks if the connectivity of two graphs is the same.

static validate_graph(self)

copy()

get_subgraph_from_nodes(nodes, edge_trimming_mode='+from+to')

Create a subgraph by filtering nodes and incident edges.

Parameters:

• nodes – Node subset to keep.

• edge_trimming_mode – Edge filter mode. Supported values are "+from+to" (keep edges where both endpoints are in nodes) and "-from+to" (keep edges where both endpoints are not in nodes).

Returns:

A new Graph with copied nodes and matching copied edges.

tf_gnns.tfgnns_datastructures.make_graph_tuple_from_graph_list(list_of_graphs)

Create a GraphTuple from a list of object graphs.

Parameters:

list_of_graphs – List of Graph objects with consistent feature dimensionality.

Returns:

A GraphTuple with flattened node/edge tensors and bookkeeping vectors (senders, receivers, n_nodes, n_edges).

Notes

This helper currently expects node and edge attributes in each input graph to have first dimension equal to 1.

class tf_gnns.tfgnns_datastructures.GraphTuple(nodes, edges, senders, receivers, n_nodes, n_edges, global_attr=None, global_reps_for_nodes=None, global_reps_for_edges=None, n_graphs=None)

Bases: object

Batched graph representation used by GraphNet tensor-dict paths.

A GraphTuple stores all node and edge features in contiguous tensors and keeps graph boundaries via n_nodes and n_edges vectors.

The GraphTuple makes multiple smaller graphs appear as a single large graph, with contiguous indexing for nodes and edges. This allows fast batched computation and takes advantage of default performance optimizations in deep learning frameworks.

__init__(nodes, edges, senders, receivers, n_nodes, n_edges, global_attr=None, global_reps_for_nodes=None, global_reps_for_edges=None, n_graphs=None)

Initialize a graph batch.

Parameters:

• nodes – Tensor-like node feature array with shape [sum(n_nodes), d_n].

• edges – Tensor-like edge feature array with shape [sum(n_edges), d_e].

• senders – Sender node indices for each edge. Indices are unique across graphs in the flattened representation.

• receivers – Receiver node indices for each edge. Indices are unique across graphs in the flattened representation.

• n_nodes – Per-graph node counts.

• n_edges – Per-graph edge counts.

• global_attr – Optional graph-level features of shape [n_graphs, d_g].

• global_reps_for_nodes – Optional precomputed mapping from node rows to graph ids.

• global_reps_for_edges – Optional precomputed mapping from edge rows to graph ids.

• n_graphs – Optional number of graphs.

update_reps_for_globals()

Build helper vectors mapping nodes/edges to graph indices.

assign_global(global_attr, check_shape=False)

Assign graph-level features.

Parameters:

• global_attr – Tensor-like global features.

• check_shape – If True, assert first dimension equals n_graphs.

is_equal_by_value(other_graph_tuple)

copy()

get_graph(graph_index)

Extract a single Graph from this batch.

Parameters:

graph_index – Zero-based index of the graph to extract.

Returns:

A new Graph object containing copied node/edge features.

to_tensor_dict()

Convert this graph batch to a GraphNet tensor dictionary.