File size: 1,375 Bytes
916755e | 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 | import pytest
import torch
from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
from gnn4colliders.graphs import build_dgl_graph
dgl = pytest.importorskip("dgl")
def _sample(index: int, node_count: int) -> GraphSample:
nodes = torch.arange(node_count * 3, dtype=torch.float32).reshape(node_count, 3)
return GraphSample(
graph=build_dgl_graph(nodes),
label=torch.tensor(index),
global_features=torch.tensor([float(index), -float(index)]),
metadata=EventMetadata(index, index + 0.5, f"sample_{index}"),
)
@pytest.mark.parametrize("counts", [(1,), (1, 2), (1, 3, 2)])
def test_graph_batch_preserves_heterogeneous_order_and_metadata(counts):
samples = [_sample(index, count) for index, count in enumerate(counts)]
batch = batch_graph_samples(samples)
assert batch.labels.tolist() == list(range(len(counts)))
assert batch.metadata.sample_id == tuple(f"sample_{i}" for i in range(len(counts)))
assert batch.metadata.weight.tolist() == pytest.approx(
[i + 0.5 for i in range(len(counts))]
)
assert batch.global_features[:, 0].tolist() == list(map(float, range(len(counts))))
assert batch.graph.batch_num_nodes().tolist() == list(counts)
assert batch.graph.batch_num_edges().tolist() == [
count * (count - 1) if count > 1 else 1 for count in counts
]
|