def test_pemsbay_task_generator():
loader = PemsBayDatasetLoader(raw_data_dir="/tmp/")
dataset = loader.get_dataset(num_timesteps_in=6, num_timesteps_out=5)
for epoch in range(2):
for snapshot in dataset:
assert snapshot.edge_index.shape == (2, 2694)
assert snapshot.edge_attr.shape == (2694, )
assert snapshot.x.shape == (325, 2, 6)
assert snapshot.y.shape == (325, 2, 5)
def test_pemsbay():
loader = PemsBayDatasetLoader(raw_data_dir="/tmp/")
dataset = loader.get_dataset()
for epoch in range(2):
for snapshot in dataset:
assert snapshot.edge_index.shape == (2, 2694)
assert snapshot.edge_attr.shape == (2694, )
assert snapshot.x.shape == (325, 2, 12)
assert snapshot.y.shape == (325, 2, 12)