def test_loader_zipped():
datalist1 = [GraphData.random(5, 4, 3) for _ in range(32 * 5)]
datalist2 = [GraphData.random(5, 4, 3) for _ in range(32 * 5)]
loader = GraphDataLoader(datalist1, datalist2, batch_size=32, shuffle=True)
for a, b in loader:
assert isinstance(a, GraphBatch)
assert isinstance(b, GraphBatch)
assert a is not b
def data(request):
args = (5, 4, 3)
kwargs = dict(min_nodes=10, max_nodes=10, min_edges=5, max_edges=5)
if request.param is GraphData:
return GraphData.random(*args, **kwargs)
else:
return GraphBatch.random_batch(100, *args, **kwargs)
def test_fully_connected_singe_graph_batch():
deterministic_seed(0)
data = GraphData.random(5, 4, 3)
batch = GraphBatch.from_data_list([data])
t = FullyConnected()
batch2 = t(batch)
assert batch2.edges.shape[1] > batch.edges.shape[1]
def run():
data = GraphData.random(5, 4, 3, min_edges=1000, min_nodes=1000)
ij = torch.cat([data.edges])
sparse_mask = scatter_coo(ij,
1,
expand=True,
size=(data.num_nodes, data.num_nodes))
def data(request):
deterministic_seed(0)
data_cls = request.param
if data_cls is GraphData:
return GraphData.random(5, 4, 3)
else:
return GraphBatch.random_batch(10, 5, 4, 3)
def test_view_graph_data(self, slices):
data = GraphData.random(5, 5, 5, min_nodes=10, max_nodes=10)
assert data.shape == (5, 5, 5)
data_view = data.view(*slices)
print(data_view.shape)
# assert data_view.shape == (5, 3, 5)
assert data_view.share_storage(data)
def test_loader_first():
datalist = [GraphData.random(5, 4, 3) for _ in range(32 * 5)]
loader = GraphDataLoader(datalist, batch_size=32, shuffle=True)
batch = loader.first()
assert isinstance(batch, GraphBatch)
assert batch.shape == (5, 4, 3)
assert batch.num_graphs == 32
def test_floyd_warshall_neighbors(self, nodes, depth, return_matrix):
data = GraphData.random(5, 4, 3, min_edges=300, min_nodes=100)
print(data.density())
x = floyd_warshall_neighbors(data,
nodes,
depth,
return_matrix=return_matrix)
print(x)
def test_loader_dataset():
datalist = [GraphData.random(5, 4, 3) for _ in range(32 * 4)]
dataset = GraphDataset(datalist)
for batch in GraphDataLoader(dataset, shuffle=True, batch_size=32):
print(batch.size)
assert isinstance(batch, GraphBatch)
assert batch.size[-1] == 32
def test_neighbors(edges, source, kwargs, expected):
data = GraphData.random(
5,
4,
3,
min_nodes=10,
max_nodes=10,
min_edges=edges.shape[1],
max_edges=edges.shape[1],
)
data.edges = edges
data.debug()
res = neighbors(data, source, **kwargs)
print(res)
assert torch.all(res == expected)
def test_k_hop(edges, k, source, expected):
deterministic_seed(0)
data = GraphData.random(
5,
4,
3,
min_nodes=10,
max_nodes=10,
min_edges=edges.shape[1],
max_edges=edges.shape[1],
)
data.edges = edges
data.debug()
res = induce(data, source, k)
print(res)
assert torch.all(res == expected)
def test_floyd_warshall_neighbors_bool(self, nodes, depth, return_matrix):
data = GraphData.random(5, 4, 3, min_edges=300, min_nodes=100)
if isinstance(nodes, tuple):
bool_nodes = []
for n in nodes:
_bool_nodes = torch.BoolTensor([False] * data.num_nodes)
_bool_nodes[n] = True
bool_nodes.append(_bool_nodes)
bool_nodes = tuple(bool_nodes)
else:
bool_nodes = torch.BoolTensor([False] * data.num_nodes)
bool_nodes[nodes] = True
x = floyd_warshall_neighbors(data,
bool_nodes,
return_matrix=return_matrix,
depth=depth)
print(x)
def test_find_neighbors():
data = GraphData.random(5, 4, 3, min_nodes=1000, min_edges=1000)
W = floyd_warshall(data)
print(W.__class__)
nodes = torch.LongTensor([[0], [1], [2], [3]])
x = W[nodes]
noninf = x != float("inf")
reachable = x <= 3
print(x.__class__)
print(noninf.__class__)
c = torch.logical_and(noninf, reachable)
d = torch.where(c)
print(d)
neighbors = d[1]
print(neighbors)
from caldera.utils import torch_scatter_group
neighbors = torch_scatter_group(d[-1], d[0]) # d[0], d[-1])
print(neighbors)
def data(self, request):
n, e = request.param
return GraphData.random(5, 3, 2, min_nodes=n, min_edges=e)
def test_to_coo_matrix():
data = GraphData.random(5, 4, 3, min_edges=1000, min_nodes=1000)
W = to_sparse_coo_matrix(data)
assert W.size() == torch.Size([data.num_edges, data.num_edges, 4])
def test_fully_connected_singe_graph_batch():
deterministic_seed(0)
data = GraphData.random(5, 4, 3)
t = FullyConnected()
data2 = t(data)
assert data2.edges.shape[1] > data.edges.shape[1]
def test_to_coo_matrix_with_fill():
data = GraphData.random(5, 4, 3, min_edges=1000, min_nodes=1000)
W = to_sparse_coo_matrix(data, fill_value=1)
print(W)
assert W.size() == torch.Size([data.num_edges, data.num_edges])
def test_serialize_graph_data():
data = GraphData.random(5, 4, 3)
pickle.loads(pickle.dumps(data))
def test_loader_mem_sizes():
datalist = [GraphData.random(5, 4, 3) for _ in range(32 * 5)]
loader = GraphDataLoader(datalist, batch_size=1, shuffle=True)
print(loader.mem_sizes())
print(loader.mem_sizes().to(torch.float).std())
def test_loader_limit_mem_sizes():
datalist = [GraphData.random(5, 4, 3) for _ in range(32 * 5)]
loader = GraphDataLoader(datalist, batch_size=1, shuffle=True)
assert not list(loader(limit_mem_size=10))
assert list(loader(limit_mem_size=1000))
def test_loader():
datalist = [GraphData.random(5, 4, 3) for _ in range(32 * 5)]
loader = GraphDataLoader(datalist, batch_size=32, shuffle=True)
for batch in loader:
assert batch.size[2] == 32
def test_graph_data_random():
assert GraphData.random(5, 5, 5)
def test_floyd_warshall():
data = GraphData.random(5, 4, 3, min_nodes=1000, min_edges=1000)
W = floyd_warshall(data)
assert torch.is_tensor(W)
def random_loader(data_size, batch_size):
datalist = [GraphData.random(5, 5, 5) for _ in range(data_size)]
return GraphDataLoader(datalist, batch_size)
