def test_data():
x = torch.tensor([[1, 3, 5], [2, 4, 6]], dtype=torch.float).t()
edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2, 1]])
data = Data(x=x, edge_index=edge_index).to(torch.device('cpu'))
N = data.num_nodes
assert data.x.tolist() == x.tolist()
assert data['x'].tolist() == x.tolist()
assert sorted(data.keys) == ['edge_index', 'x']
assert len(data) == 2
assert 'x' in data and 'edge_index' in data and 'pos' not in data
assert data.__cat_dim__('x', data.x) == 0
assert data.__cat_dim__('edge_index', data.edge_index) == -1
assert data.__cumsum__('x', data.x) is False
assert data.__cumsum__('edge_index', data.edge_index) is True
assert not data.x.is_contiguous()
data.contiguous()
assert data.x.is_contiguous()
assert not data.is_coalesced()
data.edge_index, _ = coalesce(data.edge_index, None, N, N)
assert data.is_coalesced()
clone = data.clone()
assert clone != data
assert len(clone) == len(data)
assert clone.x.tolist() == data.x.tolist()
assert clone.edge_index.tolist() == data.edge_index.tolist()
data['x'] = x + 1
assert data.x.tolist() == (x + 1).tolist()
assert data.__repr__() == 'Data(edge_index=[2, 4], x=[3, 2])'
dictionary = {'x': data.x, 'edge_index': data.edge_index}
data = Data.from_dict(dictionary)
assert sorted(data.keys) == ['edge_index', 'x']
assert not data.contains_isolated_nodes()
assert not data.contains_self_loops()
assert data.is_undirected()
assert not data.is_directed()
assert data.num_nodes == 3
assert data.num_edges == 4
assert data.num_features == 2
data.x = None
assert data.num_nodes == 3
data.edge_index = None
assert data.num_nodes is None
assert data.num_edges is None
def test_geometric(self):
x = torch.tensor([[1, 3, 5], [2, 4, 6]],
dtype=torch.float).t().to(device)
edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2,
1]]).to(device)
data = Data(x=x, edge_index=edge_index).to(torch.device('cpu'))
N = data.num_nodes
self.assertCountEqual(data.x.tolist(), x.tolist())
self.assertCountEqual(data['x'].tolist(), x.tolist())
self.assertEqual(['edge_index', 'x'], sorted(data.keys))
self.assertEqual(2, len(data))
self.assertIn('x', data)
self.assertIn('edge_index', data)
self.assertNotIn('pos', data)
self.assertEqual(0, data.__cat_dim__('x', data.x))
self.assertEqual(-1, data.__cat_dim__('edge_index', data.edge_index))
self.assertEqual(0, data.__inc__('x', data.x))
self.assertEqual(data.num_nodes,
data.__inc__('edge_index', data.edge_index))
data.contiguous()
self.assertTrue(data.x.is_contiguous())
data.edge_index, _ = coalesce(data.edge_index, None, N, N)
data = data.coalesce()
self.assertTrue(data.is_coalesced())
clone = data.clone()
self.assertNotEqual(clone, data)
self.assertEqual(len(clone), len(data))
self.assertCountEqual(clone.x.tolist(), data.x.tolist())
self.assertCountEqual(clone.edge_index.tolist(),
data.edge_index.tolist())
data['x'] = x + 1
self.assertCountEqual(data.x.tolist(), (x + 1).tolist())
self.assertEqual('Data(edge_index=[2, 4], x=[3, 2])', data.__repr__())
dictionary = {'x': data.x, 'edge_index': data.edge_index}
data = Data.from_dict(dictionary)
self.assertEqual(['edge_index', 'x'], sorted(data.keys))
self.assertTrue(data.is_undirected())
self.assertEqual(3, data.num_nodes)
self.assertEqual(4, data.num_edges)
self.assertIsNone(data.num_faces)
self.assertEqual(2, data.num_node_features)
self.assertEqual(2, data.num_features)
def test_geometric(self):
x = torch.tensor([[1, 3, 5], [2, 4, 6]], dtype=torch.float).t()
edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2, 1]])
data = Data(x=x, edge_index=edge_index).to(torch.device('cpu'))
N = data.num_nodes
self.assertTrue(data.x.tolist() == x.tolist())
self.assertTrue(data['x'].tolist() == x.tolist())
self.assertTrue(sorted(data.keys) == ['edge_index', 'x'])
self.assertTrue(len(data) == 2)
self.assertTrue('x' in data and 'edge_index' in data and 'pos' not in data)
self.assertTrue(data.__cat_dim__('x', data.x) == 0)
self.assertTrue(data.__cat_dim__('edge_index', data.edge_index) == -1)
self.assertTrue(data.__inc__('x', data.x) == 0)
self.assertTrue(data.__inc__('edge_index', data.edge_index) == data.num_nodes)
data.contiguous()
self.assertTrue(data.x.is_contiguous())
data.edge_index, _ = coalesce(data.edge_index, None, N, N)
data = data.coalesce()
self.assertTrue(data.is_coalesced())
clone = data.clone()
self.assertTrue(clone != data)
self.assertTrue(len(clone) == len(data))
self.assertTrue(clone.x.tolist() == data.x.tolist())
self.assertTrue(clone.edge_index.tolist() == data.edge_index.tolist())
data['x'] = x + 1
self.assertTrue(data.x.tolist() == (x + 1).tolist())
self.assertTrue(data.__repr__() == 'Data(edge_index=[2, 4], x=[3, 2])')
dictionary = {'x': data.x, 'edge_index': data.edge_index}
data = Data.from_dict(dictionary)
self.assertTrue(sorted(data.keys) == ['edge_index', 'x'])
self.assertTrue(data.is_undirected())
self.assertTrue(data.num_nodes == 3)
self.assertTrue(data.num_edges == 4)
self.assertTrue(data.num_faces is None)
self.assertTrue(data.num_node_features == 2)
self.assertTrue(data.num_features == 2)
def get(self, idx):
data = Data()
for key in self.data.keys:
item, slices = self.data[key], self.slices[key]
s = list(repeat(slice(None), item.dim()))
s[data.__cat_dim__(key, item)] = slice(slices[idx], slices[idx + 1])
data[key] = item[s]
return data
def __call__(self, data: Data) -> Data:
num_nodes, (row, col) = data.num_nodes, data.edge_index
edge_type = data.get('edge_type', torch.zeros_like(row))
arange = torch.arange(num_nodes, device=row.device)
full = row.new_full((num_nodes, ), num_nodes)
row = torch.cat([row, arange, full], dim=0)
col = torch.cat([col, full, arange], dim=0)
edge_index = torch.stack([row, col], dim=0)
new_type = edge_type.new_full((num_nodes, ), int(edge_type.max()) + 1)
edge_type = torch.cat([edge_type, new_type, new_type + 1], dim=0)
for key, value in data.items():
if key == 'edge_index' or key == 'edge_type':
continue
if isinstance(value, Tensor):
dim = data.__cat_dim__(key, value)
size = list(value.size())
fill_value = None
if key == 'edge_weight':
size[dim] = 2 * num_nodes
fill_value = 1.
elif data.is_edge_attr(key):
size[dim] = 2 * num_nodes
fill_value = 0.
elif data.is_node_attr(key):
size[dim] = 1
fill_value = 0.
if fill_value is not None:
new_value = value.new_full(size, fill_value)
data[key] = torch.cat([value, new_value], dim=dim)
data.edge_index = edge_index
data.edge_type = edge_type
if 'num_nodes' in data:
data.num_nodes = data.num_nodes + 1
return data
def test_data():
torch_geometric.set_debug(True)
x = torch.tensor([[1, 3, 5], [2, 4, 6]], dtype=torch.float).t()
edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2, 1]])
data = Data(x=x, edge_index=edge_index).to(torch.device('cpu'))
N = data.num_nodes
assert N == 3
assert data.x.tolist() == x.tolist()
assert data['x'].tolist() == x.tolist()
assert sorted(data.keys) == ['edge_index', 'x']
assert len(data) == 2
assert 'x' in data and 'edge_index' in data and 'pos' not in data
D = data.to_dict()
assert len(D) == 2
assert 'x' in D and 'edge_index' in D
D = data.to_namedtuple()
assert len(D) == 2
assert D.x is not None and D.edge_index is not None
assert data.__cat_dim__('x', data.x) == 0
assert data.__cat_dim__('edge_index', data.edge_index) == -1
assert data.__inc__('x', data.x) == 0
assert data.__inc__('edge_index', data.edge_index) == data.num_nodes
assert not data.x.is_contiguous()
data.contiguous()
assert data.x.is_contiguous()
assert not data.is_coalesced()
data = data.coalesce()
assert data.is_coalesced()
clone = data.clone()
assert clone != data
assert len(clone) == len(data)
assert clone.x.tolist() == data.x.tolist()
assert clone.edge_index.tolist() == data.edge_index.tolist()
# Test `data.to_heterogenous()`:
out = data.to_heterogeneous()
assert torch.allclose(data.x, out['0'].x)
assert torch.allclose(data.edge_index, out['0', '0'].edge_index)
data.edge_type = torch.tensor([0, 0, 1, 0])
out = data.to_heterogeneous()
assert torch.allclose(data.x, out['0'].x)
assert [store.num_edges for store in out.edge_stores] == [3, 1]
data.edge_type = None
data['x'] = x + 1
assert data.x.tolist() == (x + 1).tolist()
assert str(data) == 'Data(x=[3, 2], edge_index=[2, 4])'
dictionary = {'x': data.x, 'edge_index': data.edge_index}
data = Data.from_dict(dictionary)
assert sorted(data.keys) == ['edge_index', 'x']
assert not data.has_isolated_nodes()
assert not data.has_self_loops()
assert data.is_undirected()
assert not data.is_directed()
assert data.num_nodes == 3
assert data.num_edges == 4
assert data.num_faces is None
assert data.num_node_features == 2
assert data.num_features == 2
data.edge_attr = torch.randn(data.num_edges, 2)
assert data.num_edge_features == 2
data.edge_attr = None
data.x = None
assert data.num_nodes == 3
data.edge_index = None
assert data.num_nodes is None
assert data.num_edges == 0
data.num_nodes = 4
assert data.num_nodes == 4
data = Data(x=x, attribute=x)
assert len(data) == 2
assert data.x.tolist() == x.tolist()
assert data.attribute.tolist() == x.tolist()
face = torch.tensor([[0, 1], [1, 2], [2, 3]])
data = Data(num_nodes=4, face=face)
assert data.num_faces == 2
assert data.num_nodes == 4
data = Data(title='test')
assert str(data) == "Data(title='test')"
assert data.num_node_features == 0
assert data.num_edge_features == 0
key = value = 'test_value'
data[key] = value
assert data[key] == value
del data[value]
del data[value] # Deleting unset attributes should work as well.
assert data.get(key) is None
assert data.get('title') == 'test'
torch_geometric.set_debug(False)
def test_data():
torch_geometric.set_debug(True)
x = torch.tensor([[1, 3, 5], [2, 4, 6]], dtype=torch.float).t()
edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2, 1]])
data = Data(x=x, edge_index=edge_index).to(torch.device('cpu'))
N = data.num_nodes
assert N == 3
assert data.x.tolist() == x.tolist()
assert data['x'].tolist() == x.tolist()
assert sorted(data.keys) == ['edge_index', 'x']
assert len(data) == 2
assert 'x' in data and 'edge_index' in data and 'pos' not in data
D = data.to_dict()
assert len(D) == 2
assert 'x' in D and 'edge_index' in D
D = data.to_namedtuple()
assert len(D) == 2
assert D.x is not None and D.edge_index is not None
assert data.__cat_dim__('x', data.x) == 0
assert data.__cat_dim__('edge_index', data.edge_index) == -1
assert data.__inc__('x', data.x) == 0
assert data.__inc__('edge_index', data.edge_index) == data.num_nodes
assert not data.x.is_contiguous()
data.contiguous()
assert data.x.is_contiguous()
assert not data.is_coalesced()
data = data.coalesce()
assert data.is_coalesced()
clone = data.clone()
assert clone != data
assert len(clone) == len(data)
assert clone.x.tolist() == data.x.tolist()
assert clone.edge_index.tolist() == data.edge_index.tolist()
data['x'] = x + 1
assert data.x.tolist() == (x + 1).tolist()
assert str(data) == 'Data(x=[3, 2], edge_index=[2, 4])'
dictionary = {'x': data.x, 'edge_index': data.edge_index}
data = Data.from_dict(dictionary)
assert sorted(data.keys) == ['edge_index', 'x']
assert not data.has_isolated_nodes()
assert not data.has_self_loops()
assert data.is_undirected()
assert not data.is_directed()
assert data.num_nodes == 3
assert data.num_edges == 4
assert data.num_faces is None
assert data.num_node_features == 2
assert data.num_features == 2
data.edge_attr = torch.randn(data.num_edges, 2)
assert data.num_edge_features == 2
data.edge_attr = None
data.x = None
assert data.num_nodes == 3
data.edge_index = None
assert data.num_nodes is None
assert data.num_edges == 0
data.num_nodes = 4
assert data.num_nodes == 4
data = Data(x=x, attribute=x)
assert len(data) == 2
assert data.x.tolist() == x.tolist()
assert data.attribute.tolist() == x.tolist()
face = torch.tensor([[0, 1], [1, 2], [2, 3]])
data = Data(num_nodes=4, face=face)
assert data.num_faces == 2
assert data.num_nodes == 4
data = Data(title='test')
assert str(data) == "Data(title='test')"
assert data.num_node_features == 0
assert data.num_edge_features == 0
torch_geometric.set_debug(False)
