def test_globals():
graph = Graph(global_features=torch.rand(3))
validate_graph(graph)
graph = Graph(node_features=torch.rand(6, 2),
edge_features=torch.rand(5, 2),
global_features=torch.rand(3),
senders=torch.tensor([0, 0, 1, 1, 2]),
receivers=torch.tensor([0, 0, 3, 4, 5]))
validate_graph(graph)
def test_collate_dicts(graphs_nx, features_shapes, device):
graphs_in = [
add_random_features(Graph.from_networkx(g),
**features_shapes).to(device) for g in graphs_nx
]
graphs_out = list(reversed(graphs_in))
xs = torch.rand(len(graphs_in), 10, 32)
ys = torch.rand(len(graphs_in), 7)
samples = [{
'in': gi,
'x': x,
'y': y,
'out': go
} for gi, x, y, go in zip(graphs_in, xs, ys, graphs_out)]
batch = GraphBatch.collate(samples)
for g1, g2 in zip(graphs_in, batch['in']):
assert_graphs_equal(g1, g2)
torch.testing.assert_allclose(xs, batch['x'])
torch.testing.assert_allclose(ys, batch['y'])
for g1, g2 in zip(graphs_out, batch['out']):
assert_graphs_equal(g1, g2)
def test_graph_properties(graph_nx):
graph_nx = add_dummy_features(graph_nx)
graph = Graph.from_networkx(graph_nx)
assert list(graph.degree) == [d for _, d in graph_nx.degree]
assert list(graph.in_degree) == [d for _, d in graph_nx.in_degree]
assert list(graph.out_degree) == [d for _, d in graph_nx.out_degree]
def test_global_functions(graph_nx):
graph_nx = add_dummy_features(graph_nx)
graph = Graph.from_networkx(graph_nx)
assert graph.global_features.shape == graph.global_features_shape
assert graph.global_features_as_nodes.shape == (
graph.num_nodes, *graph.global_features_shape)
assert graph.global_features_as_edges.shape == (
graph.num_edges, *graph.global_features_shape)
def test_nodes():
graph = Graph(num_nodes=0)
validate_graph(graph)
assert graph.num_nodes == 0
graph = Graph(num_nodes=10)
validate_graph(graph)
assert graph.num_nodes == 10
graph = Graph(node_features=torch.rand(15, 2))
validate_graph(graph)
assert graph.num_nodes == 15
assert graph.node_features_shape == (2, )
with pytest.raises(ValueError):
Graph(num_nodes=-1)
with pytest.raises(ValueError):
Graph(num_nodes=0, node_features=torch.rand(15, 2))
def test_node_functions(graph_nx):
graph_nx = add_dummy_features(graph_nx)
graph = Graph.from_networkx(graph_nx)
# Features of the outgoing edges
# By node index
for node_index in range(graph.num_nodes):
assert graph.out_edge_features[node_index].shape[
1:] == graph.edge_features_shape
# Iterator
for out_edges in iter(graph.out_edge_features):
assert out_edges.shape[1:] == graph.edge_features_shape
# As tensor
assert graph.out_edge_features(
aggregation='sum').shape == (graph.num_nodes,
*graph.edge_features_shape)
# Features of the incoming edges
# By node index
for node_index in range(graph.num_nodes):
assert graph.in_edge_features[node_index].shape[
1:] == graph.edge_features_shape
# Iterator
for in_edges in iter(graph.in_edge_features):
assert in_edges.shape[1:] == graph.edge_features_shape
# As tensor
assert graph.in_edge_features(
aggregation='sum').shape == (graph.num_nodes,
*graph.edge_features_shape)
# Features of the successor nodes
# By node index
for node_index in range(graph.num_nodes):
assert graph.successor_features[node_index].shape[
1:] == graph.node_features_shape
# Iterator
for in_edges in iter(graph.successor_features):
assert in_edges.shape[1:] == graph.node_features_shape
# As tensor
assert graph.successor_features(
aggregation='sum').shape == (graph.num_nodes,
*graph.node_features_shape)
# Features of the predecessor nodes
# By node index
for node_index in range(graph.num_nodes):
assert graph.predecessor_features[node_index].shape[
1:] == graph.node_features_shape
# Iterator
for in_edges in iter(graph.predecessor_features):
assert in_edges.shape[1:] == graph.node_features_shape
# As tensor
assert graph.predecessor_features(
aggregation='sum').shape == (graph.num_nodes,
*graph.node_features_shape)
def test_empty():
graph = Graph()
validate_graph(graph)
assert graph.num_nodes == 0
assert graph.node_features is None
assert graph.node_features_shape is None
assert graph.num_edges == len(graph.senders) == len(graph.receivers) == 0
assert graph.edge_features is None
assert graph.edge_features_shape is None
assert graph.global_features is None
assert graph.global_features_shape is None
def test_corner_cases(features_shapes, device):
# Only some graphs have node/edge features, global features are either present on all of them or absent from all
gfs = features_shapes['global_features_shape']
graphs = [
add_random_features(Graph(num_nodes=0, num_edges=0),
global_features_shape=gfs),
add_random_features(Graph(num_nodes=0, num_edges=0),
global_features_shape=gfs),
add_random_features(Graph(num_nodes=3, num_edges=0),
**features_shapes),
add_random_features(Graph(num_nodes=0, num_edges=0),
**features_shapes),
add_random_features(
Graph(num_nodes=2,
senders=torch.tensor([0, 1]),
receivers=torch.tensor([1, 0])), **features_shapes)
]
graphbatch = GraphBatch.from_graphs(graphs).to(device)
validate_batch(graphbatch)
for g_orig, g_batch in zip(graphs, graphbatch):
assert_graphs_equal(g_orig, g_batch.cpu())
# Global features should be either present on all graphs or absent from all graphs
with pytest.raises(ValueError):
GraphBatch.from_graphs([
Graph(num_nodes=0, num_edges=0),
add_random_features(Graph(num_nodes=0, num_edges=0),
global_features_shape=10)
])
with pytest.raises(ValueError):
GraphBatch.from_graphs([
add_random_features(Graph(num_nodes=0, num_edges=0),
global_features_shape=10),
Graph(num_nodes=0, num_edges=0)
])
def test_collate_tuples(graphs_nx, features_shapes, device):
graphs_in = [
add_random_features(Graph.from_networkx(g),
**features_shapes).to(device) for g in graphs_nx
]
graphs_out = list(reversed(graphs_in))
xs = torch.rand(len(graphs_in), 10, 32)
ys = torch.rand(len(graphs_in), 7)
samples = list(zip(graphs_in, xs, ys, graphs_out))
batch = GraphBatch.collate(samples)
for g1, g2 in zip(graphs_in, batch[0]):
assert_graphs_equal(g1, g2)
torch.testing.assert_allclose(xs, batch[1])
torch.testing.assert_allclose(ys, batch[2])
for g1, g2 in zip(graphs_out, batch[3]):
assert_graphs_equal(g1, g2)
def test_edge_functions(graph_nx):
graph_nx = add_dummy_features(graph_nx)
graph = Graph.from_networkx(graph_nx)
# Edge features
# By edge index
for edge_index in range(graph.num_edges):
assert graph.edge_features[
edge_index].shape == graph.edge_features_shape
# Iterator
for edge_features in iter(graph.edge_features):
assert edge_features.shape == graph.edge_features_shape
# As tensor
assert graph.edge_features.shape == (graph.num_edges,
*graph.edge_features_shape)
# Features of the sender nodes
# By edge index
for edge_index in range(graph.num_edges):
assert graph.sender_features[
edge_index].shape == graph.node_features_shape
# Iterator
for edge_features in graph.sender_features:
assert edge_features.shape == graph.node_features_shape
# As tensor
assert graph.sender_features().shape == (graph.num_edges,
*graph.node_features_shape)
# Features of the receiver nodes
# By edge index
for edge_index in range(graph.num_edges):
assert graph.receiver_features[
edge_index].shape == graph.node_features_shape
# Iterator
for edge_features in graph.receiver_features:
assert edge_features.shape == graph.node_features_shape
# As tensor
assert graph.receiver_features().shape == (graph.num_edges,
*graph.node_features_shape)
def test_edges():
graph = Graph(num_nodes=6,
senders=torch.tensor([0, 1, 2, 5, 5]),
receivers=torch.tensor([3, 4, 5, 5, 5]))
validate_graph(graph)
assert graph.num_edges == len(graph.senders) == len(graph.receivers) == 5
graph = Graph(num_nodes=6,
edge_features=torch.rand(5, 2),
senders=torch.tensor([0, 1, 2, 5, 5]),
receivers=torch.tensor([3, 4, 5, 5, 5]))
validate_graph(graph)
assert graph.num_edges == len(graph.senders) == len(
graph.receivers) == len(graph.edge_features) == 5
assert graph.edge_features_shape == (2, )
# Negative number of edges
with pytest.raises(ValueError):
Graph(num_edges=-1)
# Senders and receivers not given
with pytest.raises(ValueError):
Graph(num_edges=3)
# Senders not given
with pytest.raises(ValueError):
Graph(num_edges=3, receivers=torch.arange(10))
# Receivers not given
with pytest.raises(ValueError):
Graph(num_edges=3, senders=torch.arange(10))
# Senders and receivers given, but not matching number of edges
with pytest.raises(ValueError):
Graph(num_edges=3,
senders=torch.arange(10),
receivers=torch.arange(10))
# Edges on a graph with no nodes
with pytest.raises(ValueError):
Graph(senders=torch.tensor([0, 1, 2]),
receivers=torch.tensor([3, 4, 5]))
# Different number of senders and receivers
with pytest.raises(ValueError):
Graph(num_nodes=6,
senders=torch.tensor([0]),
receivers=torch.tensor([3, 4, 5]))
# Indexes out-of-bounds
with pytest.raises(ValueError):
Graph(num_nodes=6,
senders=torch.tensor([0, 1, 1000]),
receivers=torch.tensor([3, 4, 5]))
# Indexes out-of-bounds
with pytest.raises(ValueError):
Graph(num_nodes=6,
senders=torch.tensor([0, 1, 2]),
receivers=torch.tensor([3, 4, 1000]))
# Indexes out-of-bounds
with pytest.raises(ValueError):
Graph(num_nodes=6,
senders=torch.tensor([-1000, 1, 2]),
receivers=torch.tensor([3, 4, 5]))
# Indexes out-of-bounds
with pytest.raises(ValueError):
Graph(num_nodes=6,
senders=torch.tensor([0, 1, 2]),
receivers=torch.tensor([-1000, 4, 5]))
# Senders, receivers and number of edges given, but not matching features
with pytest.raises(ValueError):
Graph(num_nodes=6,
senders=torch.tensor([0, 1]),
receivers=torch.tensor([3, 4]),
edge_features=torch.rand(9, 2))
def test_from_networkx(graph_nx, features_shapes):
graph_nx = add_random_features(graph_nx, **features_shapes)
graph = Graph.from_networkx(graph_nx)
assert_graphs_equal(graph_nx, graph)
def graphs() -> Sequence[Graph]:
return [Graph.from_networkx(g) for g in graphs_for_test().values()]
def graph(request) -> Graph:
return Graph.from_networkx(request.param)