def create_loader(generator, graphs, batch_size, shuffle, pin_memory=False):
train_batch = GraphBatch.from_networkx_list(graphs,
n_edge_feat=1,
n_node_feat=generator.n_parts,
n_glob_feat=1)
target_batch = GraphBatch.from_networkx_list(graphs,
n_edge_feat=16,
n_node_feat=1,
n_glob_feat=1,
feature_key="target")
train_list = train_batch.to_data_list()
target_list = target_batch.to_data_list()
if batch_size is None:
batch_size = len(train_list)
return GraphDataLoader(
list(zip(train_list, target_list)),
batch_size=batch_size,
shuffle=shuffle,
pin_memory=pin_memory,
)
def grid_data(request):
def newg(g):
return nx_random_features(g, 5, 4, 3)
if request.param is GraphData:
g = nx_to_directed(newg(nx.grid_graph([2, 4, 3])))
return GraphData.from_networkx(g)
elif request.param is GraphBatch:
graphs = [
nx_to_directed(newg(nx.grid_graph([2, 4, 3]))) for _ in range(10)
]
return GraphBatch.from_networkx_list(graphs)
else:
raise ValueError()
def test_k_hop_random_graph(k):
g1 = nx.grid_graph(dim=[2, 3, 4])
g2 = nx.grid_graph(dim=[2, 3, 4])
g1 = nx_to_directed(g1)
g2 = nx_to_directed(g2)
nx_random_features(g1, 5, 4, 3)
nx_random_features(g2, 5, 4, 3)
batch = GraphBatch.from_networkx_list([g1, g2])
nodes = torch.BoolTensor([False] * batch.num_nodes)
nodes[0] = True
node_mask = induce(batch, nodes, k)
subgraph = batch.apply_node_mask(node_mask)
print(subgraph.info())
