From 0bd14ce3707c8ca6a3a33e5ff5e199207c84f695 Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Wed, 8 Jul 2026 15:42:26 -0400 Subject: [PATCH 1/2] an EventGraph with tests --- src/pathpyG/core/event_graph.py | 122 ++++++++++++++ tests/core/test_event_graph.py | 280 ++++++++++++++++++++++++++++++++ tests/core/test_graph.py | 12 +- tests/core/test_index_map.py | 2 +- 4 files changed, 411 insertions(+), 5 deletions(-) create mode 100644 src/pathpyG/core/event_graph.py create mode 100644 tests/core/test_event_graph.py diff --git a/src/pathpyG/core/event_graph.py b/src/pathpyG/core/event_graph.py new file mode 100644 index 00000000..97402aa9 --- /dev/null +++ b/src/pathpyG/core/event_graph.py @@ -0,0 +1,122 @@ +from __future__ import annotations +from typing import Tuple, Union +import numpy as np +import torch +from torch_geometric.data import Data +from pathpyG.algorithms.temporal import lift_order_temporal, temporal_shortest_paths +from pathpyG.core.graph import Graph +from pathpyG.core.index_map import IndexMap +from pathpyG.core.temporal_graph import TemporalGraph + + +class EventGraph(Graph): + def __init__( + self, + data: Data, + delta: Union[int, float], + fo_mapping: IndexMap | None = None, + num_fo_nodes: int | None = None, + mapping: IndexMap | None = None, + ) -> None: + + if "node_time" not in data: + raise ValueError("EventGraph requires a per-event `node_time` node attribute.") + + super().__init__(data, mapping=mapping) + + self.delta = delta + self.fo_mapping = fo_mapping if fo_mapping is not None else IndexMap() + if num_fo_nodes is not None: + self._num_fo_nodes = int(num_fo_nodes) + else: + self._num_fo_nodes = int(self.data.node_sequence.max().item()) + 1 + + ei = self.data.edge_index + self.data.edge_delta = self.data.node_time[ei[1]] - self.data.node_time[ei[0]] + + self._temporal_graph: TemporalGraph | None = None + + @classmethod + def from_temporal_graph(cls, g: TemporalGraph, delta: Union[int, float] = 1) -> "EventGraph": + ho_index = lift_order_temporal(g, delta) + m = g.data.time.size(0) # number of events (== number of first-order edges) + node_sequence = g.data.edge_index.as_tensor().t().contiguous() # [m, 2] + node_time = g.data.time.clone() # [m] + + data = Data( + edge_index=ho_index, + num_nodes=m, + node_sequence=node_sequence, + node_time=node_time, + ) + eg = cls(data, delta=delta, fo_mapping=g.mapping, num_fo_nodes=g.n) + + # Attach a clone of the temporal graph since we already have it + eg._temporal_graph = TemporalGraph(g.data.clone(), mapping=g.mapping) + + return eg + + def __str__(self) -> str: + events_str = "" + for i in range(self.n): + u_id, v_id, t = self.event_endpoints(i) + events_str += f"\n{u_id}->{v_id}@{t}" + return ( + f"EventGraph (delta={self.delta})" + f"{events_str}" + ) + + def __len__(self): + return self.n + + def __getitem__(self, key): + if isinstance(key, (int, np.integer)) and not isinstance(key, bool): + return self.event_endpoints(int(key)) + return super().__getitem__(key) + + def to(self, device: torch.device) -> "EventGraph": + super().to(device) + if self._temporal_graph is not None: + self._temporal_graph.to(device) + return self + + def to_temporal_graph(self) -> TemporalGraph: + if self._temporal_graph is None: + edge_index = self.data.node_sequence.t().contiguous() # [2, num_events] + self._temporal_graph = TemporalGraph( + Data( + edge_index=edge_index, + time=self.data.node_time.clone(), + num_nodes=self.num_fo_nodes, + ), + mapping=self.fo_mapping, + ) + return self._temporal_graph + + @property + def num_fo_nodes(self) -> int: + return self._num_fo_nodes + + @property + def num_events(self) -> int: + return self.n + + def event_time(self, i: int) -> Union[int, float]: + return self.data.node_time[i].item() + + def event_endpoints(self, i: int) -> Tuple: + u, v = self.data.node_sequence[i].tolist() + return self.fo_mapping.to_id(u), self.fo_mapping.to_id(v), self.data.node_time[i].item() + + def continuations(self, i: int) -> list: + out = [] + for nxt in self.get_successors(i): + nxt = int(nxt.item()) + out.append((nxt, self.data.edge_delta[self.edge_to_index[(i, nxt)]].item())) + return out + + def shortest_paths(self) -> Tuple[np.ndarray, np.ndarray]: + # TODO: This is wasteful, since we already have the lifted edge index + # Modify `temporal_shortest_paths` to take in an optional pre-computed + # edge_index? + return temporal_shortest_paths(self.to_temporal_graph(), self.delta) \ No newline at end of file diff --git a/tests/core/test_event_graph.py b/tests/core/test_event_graph.py new file mode 100644 index 00000000..600f143b --- /dev/null +++ b/tests/core/test_event_graph.py @@ -0,0 +1,280 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch +from scipy.sparse.csgraph import dijkstra +from torch_geometric.data import Data +from torch_geometric.utils import to_scipy_sparse_matrix + +from pathpyG.algorithms.temporal import lift_order_temporal, temporal_shortest_paths +from pathpyG.core.index_map import IndexMap +from pathpyG.core.temporal_graph import TemporalGraph +from pathpyG.core.multi_order_model import MultiOrderModel +from pathpyG.core.event_graph import EventGraph + + +DELTA = 2 + + +# Example temporal graph used in the tests: +# +# a +# | t=1 +# v +# b -------- t=5 --------> d +# | t=2 +# v +# c +# | t=3 +# v +# e +# +# Event graph corresponding to the above (with DELTA = 2): +# +# 0 --gap 1--> 1 --gap 1--> 2 +# 3 (isolated event) +# +# where +# +# event 0: (a->b)@1 +# event 1: (b->c)@2 +# event 2: (c->e)@3 +# event 3: (b->d)@5 + +@pytest.fixture +def temporal_graph() -> TemporalGraph: + return TemporalGraph.from_edge_list( + [ + ("a", "b", 1), + ("b", "c", 2), + ("b", "d", 5), + ("c", "e", 3), + ] + ) + + +@pytest.fixture +def existing(temporal_graph): + # Properties of `temporal_graph` computed using the existing API. + ho = lift_order_temporal(temporal_graph, DELTA) # (2, 2) + m = temporal_graph.data.time.numel() # 4 - number of events + n = temporal_graph.n # 5 - number of FO nodes + node_time = temporal_graph.data.time + node_sequence = temporal_graph.data.edge_index.as_tensor().t() # (m, 2) + + edge_delta = node_time[ho[1]] - node_time[ho[0]] + adj = to_scipy_sparse_matrix(ho, edge_attr=edge_delta, num_nodes=m) + fastest = dijkstra(adj, directed=True) # (m, m) + + dist_fo, pred_fo = temporal_shortest_paths(temporal_graph, DELTA) # (n, n) + + return { + "ho": ho, + "m": m, + "n": n, + "node_time": node_time, + "node_sequence": node_sequence, + "edge_delta": edge_delta, + "fastest": fastest, + "dist_fo": dist_fo, + "pred_fo": pred_fo, + } + + +@pytest.fixture +def event_graph(temporal_graph) -> EventGraph: + return EventGraph.from_temporal_graph(temporal_graph, delta=DELTA) + + +def test_basic(event_graph, existing): + assert event_graph.delta == DELTA + assert len(event_graph) == existing["m"] == 4 + assert event_graph.num_events == existing["m"] == 4 + assert event_graph.n == existing["m"] == 4 + assert event_graph.num_fo_nodes == existing["n"] == 5 + + +def test_str(event_graph): + assert ( + str(event_graph) + == "EventGraph (delta=2)\na->b@1\nb->c@2\nc->e@3\nb->d@5" + ) + + +def test_node_time(event_graph, existing): + assert torch.equal(event_graph.data.node_time, existing["node_time"]) + assert event_graph.data.node_time.tolist() == [1, 2, 3, 5] + + +def test_node_sequence(event_graph, existing): + assert torch.equal(event_graph.data.node_sequence, existing["node_sequence"]) + assert event_graph.data.node_sequence.tolist() == [[0, 1], [1, 2], [2, 4], [1, 3]] + + +def test_fo_mapping(event_graph, temporal_graph): + fo = event_graph.fo_mapping + assert fo.num_ids() == 5 + for node in "abcde": + assert fo.to_id(fo.to_idx(node)) == node + assert fo.to_idx(node) == temporal_graph.mapping.to_idx(node) + + +def test_continuation_edge_index_matches_existing(event_graph, existing): + got = event_graph.data.edge_index.as_tensor() + got_set = {tuple(c) for c in got.t().tolist()} + assert got_set == {tuple(c) for c in existing["ho"].t().tolist()} + assert got_set == {(0, 1), (1, 2)} + + +def test_event_time(event_graph, existing): + for i in range(len(event_graph)): + assert event_graph.event_time(i) == existing["node_time"][i].item() + + +def test_getitem(event_graph): + # (u, v, t) for each event + assert event_graph[0] == ("a", "b", 1) + assert event_graph[1] == ("b", "c", 2) + assert event_graph[2] == ("c", "e", 3) + assert event_graph[3] == ("b", "d", 5) + + +def test_isolated_events(event_graph): + isolated = [ + i + for i in range(event_graph.num_events) + if event_graph.get_successors(i).numel() == 0 and event_graph.get_predecessors(i).numel() == 0 + ] + assert isolated == [3] + + +def test_continuations_and_gaps(event_graph): + cont = {i: event_graph.continuations(i) for i in range(event_graph.num_events)} + assert cont[0] == [(1, 1)] # (a->b)@1 -> (b->c)@2, gap 1 + assert cont[1] == [(2, 1)] # (b->c)@2 -> (c->e)@3, gap 1 + assert cont[2] == [] + assert cont[3] == [] + + +def test_continuation_deltas(event_graph): + for i in range(event_graph.num_events): + for _nxt, gap in event_graph.continuations(i): + assert 0 < gap <= event_graph.delta + + +def test_edge_delta_matches_existing(event_graph, existing): + got = { + tuple(c): d + for c, d in zip( + event_graph.data.edge_index.as_tensor().t().tolist(), + event_graph.data.edge_delta.tolist(), + ) + } + expected = { + tuple(c): d + for c, d in zip(existing["ho"].t().tolist(), existing["edge_delta"].tolist()) + } + assert got == expected + assert got == {(0, 1): 1, (1, 2): 1} + + +def test_shortest_paths_distances(event_graph, existing): + dist, _pred = event_graph.shortest_paths() + np.testing.assert_array_equal(dist, existing["dist_fo"]) + + +def test_shortest_paths_predecessors(event_graph, existing): + _dist, pred = event_graph.shortest_paths() + np.testing.assert_array_equal(pred, existing["pred_fo"]) + + +def test_shortest_paths_a_to_d_is_unreachable(event_graph): + """a -> d needs a->b@1 then b->d@5, gap of 4 > delta""" + dist, _pred = event_graph.shortest_paths() + assert dist[0, 3] == np.inf + + +def test_fastest_path_distances(event_graph, existing): + fastest = dijkstra(event_graph.sparse_adj_matrix(edge_attr="edge_delta"), directed=True) + np.testing.assert_array_equal(fastest, existing["fastest"]) + + +def test_to_temporal_graph_round_trip(event_graph, temporal_graph): + # An EventGraph can be converted to a TemporalGraph and back again. + rebuilt = event_graph.to_temporal_graph() + assert isinstance(rebuilt, TemporalGraph) + assert torch.equal( + rebuilt.data.edge_index.as_tensor(), + temporal_graph.data.edge_index.as_tensor(), + ) + assert torch.equal(rebuilt.data.time, temporal_graph.data.time) + assert rebuilt.n == temporal_graph.n + for node in ("a", "b", "c", "d", "e"): + assert rebuilt.mapping.to_idx(node) == temporal_graph.mapping.to_idx(node) + + +def test_multi_order_model_construction(event_graph, temporal_graph): + # A MultiOrderModel can be constructed from an EventGraph or a TemporalGraph. + with pytest.raises(AttributeError): + # no such attribute yet + mom_eg = MultiOrderModel.from_event_graph(event_graph, max_order=2) + mom_tg = MultiOrderModel.from_temporal_graph(temporal_graph, delta=DELTA, max_order=2) + + for k in (1, 2): + assert torch.equal( + mom_eg.layers[k].data.edge_index.as_tensor(), + mom_tg.layers[k].data.edge_index.as_tensor(), + ) + assert torch.equal( + mom_eg.layers[k].data.edge_weight, + mom_tg.layers[k].data.edge_weight, + ) + + +def test_to_device(event_graph): + # Moving an EventGraph to a different device moves the underlying TemporalGraph too. + moved = event_graph.to(torch.device("cpu")) + assert isinstance(moved, EventGraph) + assert moved is event_graph + assert moved.to_temporal_graph().data.edge_index.device.type == "cpu" + + +""" +The following tests illustrate that an EventGraph can be constructed from a raw +`torch_geometric.data.Data` object. +""" +@pytest.fixture +def event_data() -> Data: + return Data( + edge_index=torch.tensor([[0, 1], [1, 2]]), + num_nodes=4, + node_sequence=torch.tensor([[0, 1], [1, 2], [2, 4], [1, 3]]), + node_time=torch.tensor([1, 2, 3, 5]), + ) + + +def test_construct_from_data(event_data): + eg = EventGraph(event_data, delta=DELTA) + got = { + tuple(c): d + for c, d in zip( + eg.data.edge_index.as_tensor().t().tolist(), eg.data.edge_delta.tolist() + ) + } + assert got == {(0, 1): 1, (1, 2): 1} + + +def test_construct_from_data_to_temporal_graph(event_data): + # An EventGraph constructed from raw `torch_geometric.data.Data` can still give us + # a `TemporalGraph` using `.to_temporal_graph()` + fo = IndexMap(["a", "b", "c", "d", "e"]) + eg = EventGraph(event_data, delta=DELTA, fo_mapping=fo, num_fo_nodes=5) + tg = eg.to_temporal_graph() + assert isinstance(tg, TemporalGraph) + assert tg.n == 5 + assert torch.equal( + tg.data.edge_index.as_tensor(), + torch.tensor([[0, 1, 2, 1], [1, 2, 4, 3]]), + ) + assert tg.data.time.tolist() == [1, 2, 3, 5] \ No newline at end of file diff --git a/tests/core/test_graph.py b/tests/core/test_graph.py index f27071bb..2ae1c060 100644 --- a/tests/core/test_graph.py +++ b/tests/core/test_graph.py @@ -77,6 +77,7 @@ def test_from_edge_list(): ("b", "c"), ] g = Graph.from_edge_list(edge_list) + # Since node labels are strings, they are sorted lexicographically. assert g.mapping.to_idx("a") == 0 assert g.mapping.to_idx("b") == 1 assert g.mapping.to_idx("c") == 2 @@ -86,12 +87,15 @@ def test_from_edge_list(): (2, 1), ] g = Graph.from_edge_list(edge_list) + # Since node labels are ints, they are sorted numerically. assert g.mapping.to_idx(1) == 0 assert g.mapping.to_idx(2) == 1 assert g.mapping.to_idx(12) == 2 edge_list = [("1", "12"), ("2", "1"), ("21", "3")] g = Graph.from_edge_list(edge_list) + # Node labels are strings, but since ALL of them are numeric, they are sorted + # numerically, not lexicographically. assert g.mapping.to_idx("1") == 0 assert g.mapping.to_idx("2") == 1 assert g.mapping.to_idx("3") == 2 @@ -280,8 +284,8 @@ def test_add_operator_wo_indices(): g1 = Graph.from_edge_index(torch.IntTensor([[0, 1, 1], [1, 2, 3]]), num_nodes=4) g2 = Graph.from_edge_index(torch.IntTensor([[0, 1, 1], [1, 2, 3]]), num_nodes=4) g = g1 + g2 - assert g.n == g1.n - assert g.m == g1.m + g2.m + assert g.n == g1.n # No relabeling - number of nodes stays the same + assert g.m == g1.m + g2.m # No deduplication - number of edges goes up assert torch.equal(g.data.edge_index, torch.tensor([[0, 0, 1, 1, 1, 1], [1, 1, 2, 3, 2, 3]])) g3 = Graph.from_edge_index(torch.IntTensor([[0, 2, 3], [2, 3, 4]]), num_nodes=5) @@ -296,8 +300,8 @@ def test_add_operator_complete_overlap(): g1 = Graph.from_edge_index(torch.IntTensor([[0, 1, 1], [1, 2, 3]]), mapping=IndexMap(["a", "b", "c", "d"])) g2 = Graph.from_edge_index(torch.IntTensor([[0, 1, 1], [1, 2, 3]]), mapping=IndexMap(["a", "b", "c", "d"])) g = g1 + g2 - assert g.n == g1.n - assert g.m == g1.m + g2.m + assert g.n == g1.n # No relabeling - number of nodes stays the same + assert g.m == g1.m + g2.m # No deduplication - number of edges goes up assert torch.equal(g.data.edge_index, torch.tensor([[0, 0, 1, 1, 1, 1], [1, 1, 2, 3, 2, 3]])) diff --git a/tests/core/test_index_map.py b/tests/core/test_index_map.py index 6cb91e1c..c7130730 100644 --- a/tests/core/test_index_map.py +++ b/tests/core/test_index_map.py @@ -122,7 +122,7 @@ def test_bulk_ids(): node_ids = np.array(["a", "c", "b", "d", "a", "e"]) with pytest.raises(ValueError): mapping = IndexMap(node_ids) - mapping = IndexMap(np.unique(node_ids)) + mapping = IndexMap(np.unique(node_ids)) # sorts assert mapping.to_idx("a") == 0 assert mapping.to_idx("c") == 2 From 10f273adf0499833e9cfd8c690d0ebbe5114ede8 Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Thu, 9 Jul 2026 11:42:20 -0400 Subject: [PATCH 2/2] ruff fixes --- src/pathpyG/core/event_graph.py | 20 +++++++++++++++++++- tests/core/test_event_graph.py | 33 +++++++++++++++++++++++---------- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/src/pathpyG/core/event_graph.py b/src/pathpyG/core/event_graph.py index 97402aa9..69909308 100644 --- a/src/pathpyG/core/event_graph.py +++ b/src/pathpyG/core/event_graph.py @@ -1,8 +1,12 @@ +"""Event graph representation of a temporal graph and related operations.""" from __future__ import annotations + from typing import Tuple, Union + import numpy as np import torch from torch_geometric.data import Data + from pathpyG.algorithms.temporal import lift_order_temporal, temporal_shortest_paths from pathpyG.core.graph import Graph from pathpyG.core.index_map import IndexMap @@ -10,6 +14,8 @@ class EventGraph(Graph): + """A directed acyclic graph whose nodes are time-stamped events.""" + def __init__( self, data: Data, @@ -18,7 +24,7 @@ def __init__( num_fo_nodes: int | None = None, mapping: IndexMap | None = None, ) -> None: - + """Create an EventGraph from a `Data` object carrying per-event `node_time`.""" if "node_time" not in data: raise ValueError("EventGraph requires a per-event `node_time` node attribute.") @@ -38,6 +44,7 @@ def __init__( @classmethod def from_temporal_graph(cls, g: TemporalGraph, delta: Union[int, float] = 1) -> "EventGraph": + """Build an EventGraph from a temporal graph by lifting its edges into events.""" ho_index = lift_order_temporal(g, delta) m = g.data.time.size(0) # number of events (== number of first-order edges) node_sequence = g.data.edge_index.as_tensor().t().contiguous() # [m, 2] @@ -57,6 +64,7 @@ def from_temporal_graph(cls, g: TemporalGraph, delta: Union[int, float] = 1) -> return eg def __str__(self) -> str: + """Return a human-readable summary listing the delta and all events.""" events_str = "" for i in range(self.n): u_id, v_id, t = self.event_endpoints(i) @@ -67,20 +75,24 @@ def __str__(self) -> str: ) def __len__(self): + """Return the number of events in the graph.""" return self.n def __getitem__(self, key): + """Return the (u, v, t) endpoints for an integer key, else delegate to `Graph`.""" if isinstance(key, (int, np.integer)) and not isinstance(key, bool): return self.event_endpoints(int(key)) return super().__getitem__(key) def to(self, device: torch.device) -> "EventGraph": + """Move the event graph and its underlying temporal graph to the given device.""" super().to(device) if self._temporal_graph is not None: self._temporal_graph.to(device) return self def to_temporal_graph(self) -> TemporalGraph: + """Return the underlying temporal graph, reconstructing it if necessary.""" if self._temporal_graph is None: edge_index = self.data.node_sequence.t().contiguous() # [2, num_events] self._temporal_graph = TemporalGraph( @@ -95,20 +107,25 @@ def to_temporal_graph(self) -> TemporalGraph: @property def num_fo_nodes(self) -> int: + """Number of distinct first-order nodes underlying the events.""" return self._num_fo_nodes @property def num_events(self) -> int: + """Number of events (nodes) in the event graph.""" return self.n def event_time(self, i: int) -> Union[int, float]: + """Return the timestamp of the i-th event.""" return self.data.node_time[i].item() def event_endpoints(self, i: int) -> Tuple: + """Return the (source id, target id, time) of the i-th event.""" u, v = self.data.node_sequence[i].tolist() return self.fo_mapping.to_id(u), self.fo_mapping.to_id(v), self.data.node_time[i].item() def continuations(self, i: int) -> list: + """Return each successor event of i paired with its time gap.""" out = [] for nxt in self.get_successors(i): nxt = int(nxt.item()) @@ -116,6 +133,7 @@ def continuations(self, i: int) -> list: return out def shortest_paths(self) -> Tuple[np.ndarray, np.ndarray]: + """Return first-order shortest-path distances and predecessors respecting delta.""" # TODO: This is wasteful, since we already have the lifted edge index # Modify `temporal_shortest_paths` to take in an optional pre-computed # edge_index? diff --git a/tests/core/test_event_graph.py b/tests/core/test_event_graph.py index 600f143b..767cd915 100644 --- a/tests/core/test_event_graph.py +++ b/tests/core/test_event_graph.py @@ -8,11 +8,10 @@ from torch_geometric.utils import to_scipy_sparse_matrix from pathpyG.algorithms.temporal import lift_order_temporal, temporal_shortest_paths +from pathpyG.core.event_graph import EventGraph from pathpyG.core.index_map import IndexMap -from pathpyG.core.temporal_graph import TemporalGraph from pathpyG.core.multi_order_model import MultiOrderModel -from pathpyG.core.event_graph import EventGraph - +from pathpyG.core.temporal_graph import TemporalGraph DELTA = 2 @@ -88,6 +87,7 @@ def event_graph(temporal_graph) -> EventGraph: def test_basic(event_graph, existing): + """Basic counts (delta, events, nodes) match the source temporal graph.""" assert event_graph.delta == DELTA assert len(event_graph) == existing["m"] == 4 assert event_graph.num_events == existing["m"] == 4 @@ -96,6 +96,7 @@ def test_basic(event_graph, existing): def test_str(event_graph): + """The string representation lists the delta and all events.""" assert ( str(event_graph) == "EventGraph (delta=2)\na->b@1\nb->c@2\nc->e@3\nb->d@5" @@ -103,16 +104,19 @@ def test_str(event_graph): def test_node_time(event_graph, existing): + """Each event node carries the timestamp of its underlying edge.""" assert torch.equal(event_graph.data.node_time, existing["node_time"]) assert event_graph.data.node_time.tolist() == [1, 2, 3, 5] def test_node_sequence(event_graph, existing): + """Each event node stores the (source, target) first-order node pair.""" assert torch.equal(event_graph.data.node_sequence, existing["node_sequence"]) assert event_graph.data.node_sequence.tolist() == [[0, 1], [1, 2], [2, 4], [1, 3]] def test_fo_mapping(event_graph, temporal_graph): + """The first-order node mapping round-trips and matches the temporal graph.""" fo = event_graph.fo_mapping assert fo.num_ids() == 5 for node in "abcde": @@ -121,6 +125,7 @@ def test_fo_mapping(event_graph, temporal_graph): def test_continuation_edge_index_matches_existing(event_graph, existing): + """The continuation edge index matches the existing lift-order result.""" got = event_graph.data.edge_index.as_tensor() got_set = {tuple(c) for c in got.t().tolist()} assert got_set == {tuple(c) for c in existing["ho"].t().tolist()} @@ -128,12 +133,13 @@ def test_continuation_edge_index_matches_existing(event_graph, existing): def test_event_time(event_graph, existing): + """event_time(i) returns the timestamp of the i-th event.""" for i in range(len(event_graph)): assert event_graph.event_time(i) == existing["node_time"][i].item() def test_getitem(event_graph): - # (u, v, t) for each event + """Indexing an EventGraph yields the (u, v, t) tuple for each event.""" assert event_graph[0] == ("a", "b", 1) assert event_graph[1] == ("b", "c", 2) assert event_graph[2] == ("c", "e", 3) @@ -141,6 +147,7 @@ def test_getitem(event_graph): def test_isolated_events(event_graph): + """Events with no predecessors or successors are correctly identified.""" isolated = [ i for i in range(event_graph.num_events) @@ -150,6 +157,7 @@ def test_isolated_events(event_graph): def test_continuations_and_gaps(event_graph): + """continuations(i) returns each successor event with its time gap.""" cont = {i: event_graph.continuations(i) for i in range(event_graph.num_events)} assert cont[0] == [(1, 1)] # (a->b)@1 -> (b->c)@2, gap 1 assert cont[1] == [(2, 1)] # (b->c)@2 -> (c->e)@3, gap 1 @@ -158,12 +166,14 @@ def test_continuations_and_gaps(event_graph): def test_continuation_deltas(event_graph): + """Every continuation gap lies within (0, delta].""" for i in range(event_graph.num_events): for _nxt, gap in event_graph.continuations(i): assert 0 < gap <= event_graph.delta def test_edge_delta_matches_existing(event_graph, existing): + """Per-edge time deltas match the existing lift-order result.""" got = { tuple(c): d for c, d in zip( @@ -180,28 +190,31 @@ def test_edge_delta_matches_existing(event_graph, existing): def test_shortest_paths_distances(event_graph, existing): + """shortest_paths() distances match the existing first-order result.""" dist, _pred = event_graph.shortest_paths() np.testing.assert_array_equal(dist, existing["dist_fo"]) def test_shortest_paths_predecessors(event_graph, existing): + """shortest_paths() predecessors match the existing first-order result.""" _dist, pred = event_graph.shortest_paths() np.testing.assert_array_equal(pred, existing["pred_fo"]) def test_shortest_paths_a_to_d_is_unreachable(event_graph): - """a -> d needs a->b@1 then b->d@5, gap of 4 > delta""" + """Transition a->d needs a->b@1 then b->d@5, gap of 4 > delta.""" dist, _pred = event_graph.shortest_paths() assert dist[0, 3] == np.inf def test_fastest_path_distances(event_graph, existing): + """Fastest-path distances over edge deltas match the existing result.""" fastest = dijkstra(event_graph.sparse_adj_matrix(edge_attr="edge_delta"), directed=True) np.testing.assert_array_equal(fastest, existing["fastest"]) def test_to_temporal_graph_round_trip(event_graph, temporal_graph): - # An EventGraph can be converted to a TemporalGraph and back again. + """An EventGraph converts back to an equivalent TemporalGraph.""" rebuilt = event_graph.to_temporal_graph() assert isinstance(rebuilt, TemporalGraph) assert torch.equal( @@ -215,7 +228,7 @@ def test_to_temporal_graph_round_trip(event_graph, temporal_graph): def test_multi_order_model_construction(event_graph, temporal_graph): - # A MultiOrderModel can be constructed from an EventGraph or a TemporalGraph. + """A MultiOrderModel built from an EventGraph matches one from a TemporalGraph.""" with pytest.raises(AttributeError): # no such attribute yet mom_eg = MultiOrderModel.from_event_graph(event_graph, max_order=2) @@ -233,7 +246,7 @@ def test_multi_order_model_construction(event_graph, temporal_graph): def test_to_device(event_graph): - # Moving an EventGraph to a different device moves the underlying TemporalGraph too. + """Moving an EventGraph moves its underlying TemporalGraph too.""" moved = event_graph.to(torch.device("cpu")) assert isinstance(moved, EventGraph) assert moved is event_graph @@ -255,6 +268,7 @@ def event_data() -> Data: def test_construct_from_data(event_data): + """An EventGraph can be built from a raw Data object with correct edge deltas.""" eg = EventGraph(event_data, delta=DELTA) got = { tuple(c): d @@ -266,8 +280,7 @@ def test_construct_from_data(event_data): def test_construct_from_data_to_temporal_graph(event_data): - # An EventGraph constructed from raw `torch_geometric.data.Data` can still give us - # a `TemporalGraph` using `.to_temporal_graph()` + """An EventGraph built from raw Data still converts to a TemporalGraph.""" fo = IndexMap(["a", "b", "c", "d", "e"]) eg = EventGraph(event_data, delta=DELTA, fo_mapping=fo, num_fo_nodes=5) tg = eg.to_temporal_graph()