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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
|
import pytest
import torch
from torch_geometric.data.graph_store import EdgeAttr, EdgeLayout
from torch_geometric.testing import MyGraphStore, get_random_edge_index
from torch_geometric.utils import (
to_torch_coo_tensor,
to_torch_csc_tensor,
to_torch_csr_tensor,
)
def test_graph_store():
graph_store = MyGraphStore()
assert str(graph_store) == 'MyGraphStore()'
coo = torch.tensor([0, 1]), torch.tensor([1, 2])
csr = torch.tensor([0, 1, 2]), torch.tensor([1, 2])
csc = torch.tensor([0, 1]), torch.tensor([0, 0, 1, 2])
graph_store['edge_type', 'coo'] = coo
graph_store['edge_type', 'csr'] = csr
graph_store['edge_type', 'csc'] = csc
assert torch.equal(graph_store['edge_type', 'coo'][0], coo[0])
assert torch.equal(graph_store['edge_type', 'coo'][1], coo[1])
assert torch.equal(graph_store['edge_type', 'csr'][0], csr[0])
assert torch.equal(graph_store['edge_type', 'csr'][1], csr[1])
assert torch.equal(graph_store['edge_type', 'csc'][0], csc[0])
assert torch.equal(graph_store['edge_type', 'csc'][1], csc[1])
assert len(graph_store.get_all_edge_attrs()) == 3
del graph_store['edge_type', 'coo']
with pytest.raises(KeyError):
graph_store['edge_type', 'coo']
with pytest.raises(KeyError):
graph_store['edge_type_2', 'coo']
def test_graph_store_conversion():
graph_store = MyGraphStore()
edge_index = get_random_edge_index(100, 100, 300)
adj = to_torch_coo_tensor(edge_index, size=(100, 100))
coo = (adj.indices()[0], adj.indices()[1])
adj = to_torch_csr_tensor(edge_index, size=(100, 100))
csr = (adj.crow_indices(), adj.col_indices())
adj = to_torch_csc_tensor(edge_index, size=(100, 100))
csc = (adj.row_indices(), adj.ccol_indices())
graph_store.put_edge_index(coo, ('v', '1', 'v'), 'coo', size=(100, 100))
graph_store.put_edge_index(csr, ('v', '2', 'v'), 'csr', size=(100, 100))
graph_store.put_edge_index(csc, ('v', '3', 'v'), 'csc', size=(100, 100))
# Convert to COO:
row_dict, col_dict, perm_dict = graph_store.coo()
assert len(row_dict) == len(col_dict) == len(perm_dict) == 3
for row, col, perm in zip(row_dict.values(), col_dict.values(),
perm_dict.values()):
assert torch.equal(row.sort()[0], coo[0].sort()[0])
assert torch.equal(col.sort()[0], coo[1].sort()[0])
assert perm is None
# Convert to CSR:
row_dict, col_dict, perm_dict = graph_store.csr()
assert len(row_dict) == len(col_dict) == len(perm_dict) == 3
for row, col in zip(row_dict.values(), col_dict.values()):
assert torch.equal(row, csr[0])
assert torch.equal(col.sort()[0], csr[1].sort()[0])
# Convert to CSC:
row_dict, col_dict, perm_dict = graph_store.csc()
assert len(row_dict) == len(col_dict) == len(perm_dict) == 3
for row, col in zip(row_dict.values(), col_dict.values()):
assert torch.equal(row.sort()[0], csc[0].sort()[0])
assert torch.equal(col, csc[1])
# Ensure that 'edge_types' parameters work as intended:
out = graph_store.coo([('v', '1', 'v')])
assert torch.equal(list(out[0].values())[0], coo[0])
assert torch.equal(list(out[1].values())[0], coo[1])
# Ensure that 'store' parameter works as intended:
key = EdgeAttr(edge_type=('v', '1', 'v'), layout=EdgeLayout.CSR,
is_sorted=False, size=(100, 100))
with pytest.raises(KeyError):
graph_store[key]
out = graph_store.csr([('v', '1', 'v')], store=True)
assert torch.equal(list(out[0].values())[0], csr[0])
assert torch.equal(list(out[1].values())[0].sort()[0], csr[1].sort()[0])
out = graph_store[key]
assert torch.equal(out[0], csr[0])
assert torch.equal(out[1].sort()[0], csr[1].sort()[0])
|
