import marimo as mo
Causality-Aware Graph Neural Networks¶
Prerequisites¶
First, we need to set up our Python environment that has PyTorch, PyTorch Geometric and PathpyG installed. Depending on where you are executing this notebook, this might already be (partially) done. E.g. Google Colab has PyTorch installed by default so we only need to install the remaining dependencies. The DevContainer that is part of our GitHub Repository on the other hand already has all of the necessary dependencies installed.
In the following, we install the packages for usage in Google Colab using Jupyter magic commands. For other environments comment in or out the commands as necessary. For more details on how to install pathpyG especially if you want to install it with GPU-support, we refer to our documentation. Note that %%capture discards the full output of the cell to not clutter this tutorial with unnecessary installation details. If you want to print the output, you can comment %%capture out.
# magic command not supported in marimo; please file an issue to add support
# %%capture
# # !pip install torch
# # !pip install torch_geometric
# # !pip install git+https://github.com/pathpy/pathpyG.git
Motivation and Learning Objectives¶
In previous tutorials, we have introduced causal paths in temporal graphs, and how we can use them to generate higher-order De Bruijn graph models that capture temporal-topological patterns in time series data. In this tutorial, we will show how we can use De Bruijn Graph Neural Networks, a causality-aware deep learning architecture for temporal graph data. The details of this approach are introduced in this paper. The architecture is implemented in pathpyG and can be readily applied to temporal graph data.
Below we illustrate this method in a supervised node classification task, i.e. given a temporal graph we will use the temporal-topological patterns in the graph to classify nodes.
We start by importing a few modules:
from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
import scipy as sp
import torch
from sklearn.manifold import TSNE
from sklearn.metrics import balanced_accuracy_score
from torch_geometric.transforms import RandomNodeSplit
import pathpyG as pp
from pathpyG.nn.dbgnn import DBGNN
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Temporal-Topological Clusters in Temporal Graphs¶
We use the synthetic temporal graph with 60.000 time-stamped interactions between 30 nodes and three planted temporal communities, which we have already explored in Time-Respecting Paths, Centralities and Communities. We load it from a file containing edges with discrete time-stamps.
Dataset Availability
Depending on how you are executing this notebook, the dataset may not be available locally. We use <code>pp.io.example_data</code> to resolve the file name: it returns the path of a local copy in the <code>docs/data</code> folder of the pathpyG repository, or - if you are running outside of a clone, e.g. in Google Colab - a URL pointing to the GitHub repository. Both can be passed directly to the <code>pp.io.read_csv_*</code> functions.t = pp.io.read_csv_temporal_graph(pp.io.example_data('temporal_clusters.tedges'))
t = t.to(device)
This example has created in such a way that the nodes naturally form three clusters, which are highlighted in the interactive visualization below:
style = {}
style["node_color"] = ["green"] * 10 + ["red"] * 10 + ["blue"] * 10
pp.plot(t, **style, show_labels=False);
Modelling Causal Structures with Higher-Order De Bruijn Graphs¶
But what is the origin for the cluster pattern? In the visualization above, you will notice that the time-stamped edges randomly interconnect nodes within and across clusters, actually there is no correlation whatsoever between the topology of links and the cluster membership of the nodes. Hence, the notion of clusters does not correspond to the common idea of cluster patterns in static graphs, which we can highlight further by plotting the static time-aggregated network:
pp.plot(t.to_static_graph(), **style, show_labels=False);
In fact, the topology of this graph corresponds to that of a random graph, i.e. there are no patterns whatsoever in the topology of links. Nevertheless, the temporal graph contains a cluster pattern in the topology of causal or time-respecting paths. In particular, the temporal ordering of time-stamped edges is such that nodes with the same cluster label are more frequently connected by time-respecting paths than nodes with different cluster labels. Hence, nodes within the same clusters can more strongly influence each other in a causal way, i.e. via multiple interactions that follow the arrow of time.
Traditional (temporal) graph neural networks will not be able to learn from this pattern, as it is due to the specific microscopic temporal ordering of edges. Using higher-order De Bruijn graph models implemented in pathpyG, we can learn from temporal graph data that contains such patterns. Let us explain this step by step.
Building on the previous notebooks on time-respecting paths in temporal graphs, we directly construct a multi-order model whose second-order layer captures the causal structure of the temporal graph. We only consider two time-stamped edges $(u,v;t)$ and $(v,w;t')$ to contribute to a causal path iff $0 < t'-t \leq 1$, i.e. we use the default maximum time difference $\delta=1$.
# Cell tags: skip-execution
# Hide the progress output (originally done via `%%capture` in Jupyter)
with mo.capture_stdout(), mo.capture_stderr():
m = pp.MultiOrderModel.from_temporal_graph(t, max_order=2)
# Cell tags: skip-execution
print(m)
MultiOrderModel with max. order 2
We can get the first and second order networks from the MultiOrderModel object. The first order network is the network of nodes and edges, while the second-order network uses first-order edges as nodes and edges represent time-respecting paths of length two. The second order network is a De Bruijn graph that captures temporal-topological patterns in the network.
# Cell tags: skip-execution
g = m.layers[1]
g2 = m.layers[2]
# Cell tags: skip-execution
pp.plot(g, edge_size=1, show_labels=False);
Since it does not consider patterns in the causal topology of the temporal graph, this is not a meaningful model. We can instead use a second-order De Bruijn graph model, which we can easily fit to the paths:
# Cell tags: skip-execution
layout_style = {}
layout_style['layout'] = 'Fruchterman-Reingold'
layout_style['seed'] = 1
layout_style['k'] = 0.5
layout_style['iterations'] = 300
_layout = pp.layout(g2, **layout_style)
pp.plot(g2, backend='matplotlib', layout=_layout, edge_size=0.5, node_size=3, show_labels=False);
In this graph, every node is a link and links correspond to causal paths of length two, i.e. temporally ordered sequences consisting of two edges that overlap in the center node. In this graph, we clearly see a cluster pattern that is due to the way in which temporal edges are ordered in time. In particular, we see three clusters, where the edges in three of the clusters correspond to causal paths of length two that connect nodes within each of the three clusters. The edges in the fourth cluster (in the center of the visualization) represent causal paths that connect nodes in different clusters.
Comparison to Temporal Graph with Shuffled Time Stamps¶
You may wonder whether this pattern is really due to the temporal ordering of time-stamped edges. It is easy to check this. We can simply randomly shuffle the time stamps of all edges, which will break any correlations in the temporal ordering that lead to patterns in the causal topology.
We repeat the path calculation for this shuffled temporal graph and construct the second-order De Bruijn Graph model again:
# Cell tags: skip-execution
t_shuffled = deepcopy(t)
t_shuffled.shuffle_time()
# Cell tags: skip-execution
# Hide the progress output (originally done via `%%capture` in Jupyter)
with mo.capture_stdout(), mo.capture_stderr():
g2_shuffled = pp.MultiOrderModel.from_temporal_graph(t_shuffled, max_order=2).layers[2]
# Cell tags: skip-execution
print(g2_shuffled)
Higher-order graph of order 2 with 557 nodes and 2050 edges
(over 30 first-order nodes)
{ 'Edge Attributes': {'edge_weight': "<class 'torch.Tensor'> -> torch.Size([2050])"},
'Graph Attributes': {'inverse_idx': "<class 'torch.Tensor'> -> torch.Size([60000])", 'num_nodes': "<class 'int'>"},
'Node Attributes': {}}
# Cell tags: skip-execution
_layout = pp.layout(g2_shuffled, **layout_style)
pp.plot(g2_shuffled, backend='matplotlib', layout=_layout, edge_size=0.5, node_size=3, show_labels=False);
We now find that the cluster pattern in the second-order graph has vanished. In fact, there is no pattern whatsoever since the underlying (static) graph topology is random and the random shuffling of time stamps leads to random causal paths.
Spectral clustering with second-order graph Laplacian¶
To take a different perspective on cluster patterns, we can actually use pathpyG to apply a spectral analysis to the higher-order graph. We can simply calculate a generalization of the Laplacian matrix to the second-order graph both for the actual temporal graph and its shuffled counterpart:
# Cell tags: skip-execution
L = g2.laplacian(normalization='rw', edge_attr='edge_weight')
L_shuffled= g2_shuffled.laplacian(normalization='rw',edge_attr='edge_weight')
We then calculate the eigenvalues and eigenvectors of the Laplacians, and compute the Fiedler vector, i.e. the eigenvector that corresponds to the second-smallest eigenvalue of the Laplacian.
# Cell tags: skip-execution
w,v = sp.linalg.eig(L.todense(),left= False, right = True)
w_shuffled, v_shuffled = sp.linalg.eig(L_shuffled.todense())
# Cell tags: skip-execution
fiedler = v[:,np.argsort(w)[1]]
fiedler_shuffled = v_shuffled[:,np.argsort(w_shuffled)[1]]
Below, we show that the clusters in the causal topology of the temporal graph correspond to clusters in the distribution of entries in the Fiedler vector, while there is no such pattern for the Fiedler vector of the second-order graph constructed from the shuffled temporal graph:
def higher_order_class_assignment(ho_node_id):
"""Assign class labels based on the higher-order node ids.
There are three classes that were assigned based on the original node ids:
- Class 0: node ids 0-9
- Class 1: node ids 10-19
- Class 2: node ids 20-29
The higher-order patterns were constructed such that the clusters are formed by second-order nodes whose first-order node ids belong to the same class.
We therefore assign higher-order class labels based on the first-order node ids in the higher-order node id tuple.
"""
if ho_node_id[0] < 10 and ho_node_id[1] < 10:
return 0
elif ho_node_id[0] < 20 and ho_node_id[0] >= 10 and ho_node_id[1] < 20 and ho_node_id[1] >= 10:
return 1
elif ho_node_id[0] < 30 and ho_node_id[0] >= 20 and ho_node_id[1] < 30 and ho_node_id[1] >= 20:
return 2
else:
return 3
colors = {0: 'green', 1: 'red', 2: 'blue', 3: 'gray'}
opacities = {0: 0.6, 1: 0.6, 2: 0.6, 3: 0.1}
In the plots below, we have colored those entries of the Fiedler vectors that correspond to edges connecting nodes within one of the three clusters shown above. The Fiedler vector shows a clear pattern, which translates to the cluster pattern in the causal topology that we have planted into our synthetic temporal graph.
# Cell tags: skip-execution
ho_class_ids = list(map(higher_order_class_assignment, g2.nodes))
plt.scatter(
range(g2.n), np.real(fiedler), c=[colors[i] for i in ho_class_ids], alpha=[opacities[i] for i in ho_class_ids]
)
plt.ylim(-0.25, 0.25)
No such pattern exists in the Fiedler vector of the second-order graph corresponding to the shuffled TemporalGraph.
# Cell tags: skip-execution
shuffled_ho_class_ids = list(map(higher_order_class_assignment, g2_shuffled.nodes))
plt.scatter(
range(g2_shuffled.n),
np.real(fiedler_shuffled),
c=[colors[i] for i in shuffled_ho_class_ids],
alpha=[opacities[i] for i in shuffled_ho_class_ids],
)
plt.ylim(-0.25, 0.25)
Node Classification with Causality-Aware Graph Neural Networks¶
Let us now explore how we can develop a causality-aware deep graph learning architecture that utilizes this pattern in the causal topology. We will follow the architecture introduced in this work. The architecture actually performs message passing in higher-order models with multiple orders at once. In a final message passing step, a bipartite graph is used to obtain vector-space representations of actual nodes in the temporal graph.
We now set up a pytorch_geometric.Data object that contains all of the information needed to train the DBGNN model. For this, we can use a convenience function of the MultiOrderModel class in pathpyG. Combining a first- and a second-order model, this uses the edge indices and the weight tensors for a message passing scheme. it further constructs an edge_index of a bipartite graph that uses the last node in a second-order node to map messages back to first-order nodes.
# Cell tags: skip-execution
data = m.to_dbgnn_data(max_order=2, mapping="last")
data.y = torch.tensor([int(i) // 10 for i in t.nodes], device=device)
Training the model¶
We are now ready to train and evaluate our causality-aware graph neural network. We will frist create a random split of the nodes, set the optimizer and the hyperparameters of our model.
# Cell tags: skip-execution
data_1 = RandomNodeSplit(num_val=0, num_test=0.3)(data)
model = DBGNN(num_features=[g.n, g2.n], num_classes=len(data_1.y.unique()), hidden_dims=[16, 32, 8], p_dropout=0.4).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
loss_function = torch.nn.CrossEntropyLoss()
The following function evaluates the prediction of our model based on the balanced accuracy score for categorical predictions.
# Cell tags: skip-execution
def test(model, data):
"""Evaluate the model on training and test data."""
model.eval()
_, pred = model(data).max(dim=1)
metrics_train = balanced_accuracy_score(data.y[data.train_mask].cpu(), pred[data.train_mask].cpu().numpy())
metrics_test = balanced_accuracy_score(data.y[data.test_mask].cpu(), pred[data.test_mask].cpu().numpy())
return metrics_train, metrics_test
/media/vineetb/delta/projects/pathpyG/.venv/lib/python3.13/site-packages/coverage/control.py:958: CoverageWarning: No data was collected. (no-data-collected); see https://coverage.readthedocs.io/en/7.13.4/messages.html#warning-no-data-collected
self._warn("No data was collected.", slug="no-data-collected")
E ERROR: Coverage failure: total of 0 is less than fail-under=85 [100%] ==================================== ERRORS ==================================== ____________________________ ERROR at setup of test ____________________________ file /tmp/marimo_35541/__marimo__cell_lgWD_.py, line 2 def test(model, data): E fixture 'model' not found > available fixtures: anyio_backend, anyio_backend_name, anyio_backend_options, benchmark, benchmark_weave, capfd, capfdbinary, caplog, capsys, capsysbinary, capteesys, class_mocker, cov, doctest_namespace, free_tcp_port, free_tcp_port_factory, free_udp_port, free_udp_port_factory, mocker, module_mocker, monkeypatch, no_cover, package_mocker, pytestconfig, record_property, record_testsuite_property, record_xml_attribute, recwarn, session_mocker, subtests, tmp_path, tmp_path_factory, tmpdir, tmpdir_factory > use 'pytest --fixtures [testpath]' for help on them. /tmp/marimo_35541/__marimo__cell_lgWD_.py:2 ================================ tests coverage ================================ _______________ coverage: platform linux, python 3.13.12-final-0 _______________ Name Stmts Miss Cover Missing ----------------------------------------------------------------------------------------- src/pathpyG/__init__.py 15 15 0% 5-23 src/pathpyG/algorithms/__init__.py 6 6 0% 33-39 src/pathpyG/algorithms/centrality.py 141 141 0% 32-357 src/pathpyG/algorithms/components.py 19 19 0% 3-53 src/pathpyG/algorithms/generative_models.py 195 195 0% 14-578 src/pathpyG/algorithms/lift_order.py 54 54 0% 3-208 src/pathpyG/algorithms/rolling_time_window.py 19 19 0% 4-61 src/pathpyG/algorithms/shortest_paths.py 15 15 0% 7-52 src/pathpyG/algorithms/temporal.py 34 34 0% 3-74 src/pathpyG/algorithms/weisfeiler_leman.py 37 37 0% 3-75 src/pathpyG/core/event_graph.py 130 130 0% 2-291 src/pathpyG/core/graph.py 268 268 0% 3-818 src/pathpyG/core/higher_order_graph.py 224 224 0% 3-676 src/pathpyG/core/index_map.py 77 77 0% 3-398 src/pathpyG/core/multi_order_model.py 203 203 0% 3-580 src/pathpyG/core/path_data.py 46 46 0% 3-198 src/pathpyG/core/temporal_graph.py 110 110 0% 3-365 src/pathpyG/io/__init__.py 4 4 0% 3-18 src/pathpyG/io/datasets.py 13 13 0% 3-32 src/pathpyG/io/netzschleuder.py 68 68 0% 3-182 src/pathpyG/io/pandas.py 195 195 0% 3-602 src/pathpyG/nn/__init__.py 2 2 0% 3-5 src/pathpyG/nn/dbgnn.py 46 46 0% 25-151 src/pathpyG/statistics/__init__.py 4 4 0% 20-33 src/pathpyG/statistics/clustering.py 25 25 0% 3-88 src/pathpyG/statistics/degrees.py 52 52 0% 3-326 src/pathpyG/statistics/node_similarities.py 55 55 0% 3-255 src/pathpyG/utils/__init__.py 1 1 0% 7 src/pathpyG/utils/config.py 16 16 0% 3-37 src/pathpyG/utils/convert.py 17 17 0% 3-34 src/pathpyG/utils/logger.py 12 12 0% 12-38 src/pathpyG/utils/progress.py 20 20 0% 12-51 src/pathpyG/visualisations/__init__.py 2 2 0% 13-14 src/pathpyG/visualisations/_d3js/backend.py 113 113 0% 14-377 src/pathpyG/visualisations/_manim/backend.py 59 59 0% 25-222 src/pathpyG/visualisations/_manim/temporal_graph_scene.py 65 65 0% 7-181 src/pathpyG/visualisations/_matplotlib/backend.py 113 113 0% 7-407 src/pathpyG/visualisations/_tikz/backend.py 142 142 0% 36-508 src/pathpyG/visualisations/layout.py 71 71 0% 58-269 src/pathpyG/visualisations/network_plot.py 158 158 0% 26-385 src/pathpyG/visualisations/pathpy_plot.py 13 13 0% 7-39 src/pathpyG/visualisations/plot_backend.py 8 8 0% 20-51 src/pathpyG/visualisations/plot_function.py 63 63 0% 54-252 src/pathpyG/visualisations/temporal_network_plot.py 109 109 0% 8-245 src/pathpyG/visualisations/unfolded_network_plot.py 41 41 0% 7-108 src/pathpyG/visualisations/utils.py 60 60 0% 48-395 ----------------------------------------------------------------------------------------- TOTAL 3140 3140 0% 5 files skipped due to complete coverage. FAIL Required test coverage of 85.0% not reached. Total coverage: 0.00% =================================== Overview =================================== Summary: Total: 1, Passed: 0, Failed: 0, Errors: 1, Skipped: 0 =========================== short test summary info ============================ ERROR docs/tutorial/advanced/dbgnn.py::test
# Cell tags: skip-execution
losses = []
for epoch in range(50):
output = model(data_1)
loss = loss_function(output[data_1.train_mask], data_1.y[data_1.train_mask])
loss.backward()
optimizer.step()
optimizer.zero_grad()
losses.append(loss)
if epoch % 10 == 0:
train_ba, test_ba = test(model, data_1)
print(f'Epoch: {epoch}, Loss: {loss}, Train balanced accuracy: {train_ba}, Test balanced accuracy: {test_ba}')
/media/vineetb/delta/projects/pathpyG/.venv/lib/python3.13/site-packages/sklearn/metrics/_classification.py:2924: UserWarning: y_pred contains classes not in y_true
warnings.warn("y_pred contains classes not in y_true")
Epoch: 0, Loss: 1.3474797010421753, Train balanced accuracy: 0.3333333333333333, Test balanced accuracy: 0.0 Epoch: 10, Loss: 0.47045913338661194, Train balanced accuracy: 1.0, Test balanced accuracy: 1.0 Epoch: 20, Loss: 0.056760646402835846, Train balanced accuracy: 1.0, Test balanced accuracy: 1.0 Epoch: 30, Loss: 0.00212463247589767, Train balanced accuracy: 1.0, Test balanced accuracy: 1.0 Epoch: 40, Loss: 0.00036114014801569283, Train balanced accuracy: 1.0, Test balanced accuracy: 1.0
Causality-aware latent space representation of nodes¶
We can inspect the model by plotting a latent space representation of the edges generated by the second-order layer of our architecture.
# Cell tags: skip-execution
model.eval()
_latent = model.higher_order_layers[0].forward(data_1.x_h, data_1.edge_index_higher_order).detach()
_latent = model.higher_order_layers[1].forward(_latent, data_1.edge_index_higher_order).detach()
_node_embedding = TSNE(n_components=2, learning_rate='auto', init='random').fit_transform(_latent.cpu())
_embedding_layout = {v: _node_embedding[g2.mapping.to_idx(v)] for v in g2.nodes}
pp.plot(g2, backend='matplotlib', layout=_embedding_layout, show_labels=False, edge_size=0.3, node_size=3, edge_opacity=0.1, node_color=[colors[i] for i in ho_class_ids]);
We can further generate latent space representations of the nodes generated by the last bipartite layer of our architecture:
# Cell tags: skip-execution
model.eval()
_latent = model.forward(data_1).detach()
_node_embedding = TSNE(n_components=2, learning_rate='auto', init='random', perplexity=10).fit_transform(_latent.cpu())
_embedding_layout = {v: _node_embedding[g.mapping.to_idx(v)] for v in g.nodes}
pp.plot(g, backend='matplotlib', layout=_embedding_layout, show_labels=False, edge_size=0.3, node_size=3, edge_opacity=0.1, **style);
