1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
|
import torch
from torch_geometric import seed_everything
from torch_geometric.data import Data
from torch_geometric.datasets import InfectionDataset
from torch_geometric.datasets.graph_generator import ERGraph, GraphGenerator
class DummyGraph(GraphGenerator):
def __call__(self) -> Data:
edge_index = torch.tensor([
[0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9],
[1, 0, 2, 1, 3, 2, 4, 3, 5, 4, 6, 5, 7, 6, 8, 7, 9, 8],
])
return Data(num_nodes=10, edge_index=edge_index)
def test_infection_dataset():
seed_everything(12345)
graph_generator = DummyGraph()
dataset = InfectionDataset(graph_generator, num_infected_nodes=2,
max_path_length=2)
assert str(dataset) == ('InfectionDataset(1, '
'graph_generator=DummyGraph(), '
'num_infected_nodes=2, '
'max_path_length=2)')
assert len(dataset) == 1
data = dataset[0]
assert len(data) == 4
assert data.x.size() == (10, 2)
assert data.x[:, 0].sum() == 8 and data.x[:, 1].sum() == 2
assert torch.equal(data.edge_index, graph_generator().edge_index)
assert data.y.size() == (10, )
# With `seed=12345`, node 0 and node 7 will be infected:
assert data.x[0].tolist() == [0, 1] # First infected node.
assert data.x[7].tolist() == [0, 1] # Second infected node.
assert data.y.tolist() == [0, 1, 2, 3, 3, 2, 1, 0, 1, 2]
assert data.edge_mask.tolist() == [
1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0
]
def test_infection_dataset_reproducibility():
graph_generator = ERGraph(num_nodes=500, edge_prob=0.004)
seed_everything(12345)
dataset1 = InfectionDataset(graph_generator, num_infected_nodes=50,
max_path_length=5)
seed_everything(12345)
dataset2 = InfectionDataset(graph_generator, num_infected_nodes=50,
max_path_length=5)
assert torch.equal(dataset1[0].edge_mask, dataset2[0].edge_mask)
|
