Skip to content

Commit 3b34104

Browse files
committed
Merge branch 'main' into ci/perf
2 parents f9e062c + 6a2bea6 commit 3b34104

59 files changed

Lines changed: 4613 additions & 1060 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/integration.yml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ jobs:
2020
if: github.actor == 'jacob-chmura' || github.actor == 'shenyanghuang'
2121
env:
2222
TGM_CI_LOG_BASE: $SLURM_TMPDIR/tgm_ci
23+
DATA_ROOT: $HOME/tgb_datasets
2324

2425
steps:
2526
- name: Checkout repository
@@ -45,6 +46,11 @@ jobs:
4546
- name: Install dependencies
4647
run: uv sync --group examples
4748

49+
- name: Download TGB datasets
50+
run: |
51+
echo "Preparing TGB datasets in $DATA_ROOT..."
52+
./scripts/download_tgb_datasets.sh "$DATA_ROOT"
53+
4854
- name: Run Integration Tests
4955
run: uv run pytest -m "integration" -n 2 -vvv
5056

.github/workflows/publish.yml

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
name: Publish Python Package
2+
3+
on: # yamllint disable-line rule:truthy
4+
push:
5+
tags:
6+
- "v*" # Trigger on tags like v0.1.0 (PyPI)
7+
- "test-v*" # Trigger ont ags like test-v0.1.0 (Test-PyPI)
8+
9+
jobs:
10+
publish:
11+
name: Build and Publish
12+
runs-on: ubuntu-latest
13+
environment: pypi
14+
permissions:
15+
id-token: write # Required for PyPI Trusted Publishing
16+
17+
steps:
18+
- name: Checkout repository
19+
uses: actions/checkout@v4
20+
21+
- name: Set up python
22+
id: setup-python
23+
uses: actions/setup-python@v5
24+
with:
25+
python-version-file: ".python-version"
26+
27+
- name: Set up uv
28+
run: curl -LsSf https://astral.sh/uv/${{ env.UV_VERSION }}/install.sh | sh
29+
30+
- name: Check package metadata
31+
run: uv version
32+
33+
- name: Build the package
34+
run: uv build
35+
36+
- name: Sanity check that package installs and imports
37+
run: |
38+
echo "Verifying package installs and imports correctly..."
39+
uv venv .venv
40+
uv pip install dist/*.whl
41+
uv run python -c "import tgm; print(tgm.__version__)"
42+
43+
- name: Publish package
44+
uses: pypa/gh-action-pypi-publish@release/v1
45+
with:
46+
repository-url: >-
47+
${{ startsWith(github.ref_name, 'test-v') && 'https://test.pypi.org/legacy/' || 'https://upload.pypi.org/legacy/' }}

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
![image](./docs/img/logo.svg)
44

55
<div align="center">
6-
<h3 style="font-size: 22px">Efficient and Modular ML on Dynamic Graphs</h3>
6+
<h3 style="font-size: 22px">Efficient and Modular ML on Temporal Graphs</h3>
77
<a href="https://tgm.readthedocs.io/en/latest"/><strong style="font-size: 18px;">Read Our Docs»</strong></a>
88
<a href="https://github.com/tgm-team/tgm"/><strong style="font-size: 18px;">Read Our Paper»</strong></a>
99
<br/>

docs/tutorials/dgraph_tutorial.md

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
# Constructing and Accessing Properties in DGraph
2+
3+
This tutorial shows how to construct a `DGraph` object in `tgm` and explore its properties.
4+
5+
The `DGraph` class is defined in [`tgm/graph.py`](https://github.com/tgm-team/tgm/blob/main/tgm/graph.py).
6+
7+
______________________________________________________________________
8+
9+
## Construct `DGraph` from TGB Datasets
10+
11+
The [Temporal Graph Benchmark (TGB)](https://tgb.complexdatalab.com/) provides a suite of temporal graph datasets with diverse scales, properties, and tasks.
12+
13+
TGM supports constructing `DGraph` objects from TGB [LinkPropPrediction](https://tgb.complexdatalab.com/docs/linkprop/) and [NodePropPrediction](https://tgb.complexdatalab.com/docs/nodeprop/) datasets. Temporal knowledge graphs (TKG) and temporal hypergraphs (THG) are not yet supported.
14+
15+
You can specify the time granularity with [`TimeDeltaDG`](https://github.com/tgm-team/tgm/blob/main/tgm/timedelta.py). For example, `'r'` (relative) means timestamps define only the ordering of edges and are not used for time conversion. Other granularities allow unit-based conversions. The default time granularity is `'r'`, and the default device is `'cpu'`.
16+
17+
```python
18+
from tgm import DGraph
19+
20+
train_dg = DGraph.from_tgb('tgbl-wiki', time_delta='r', split='train', device='cpu')
21+
```
22+
23+
> **Note:** Time granularity features are experimental and may change.
24+
25+
______________________________________________________________________
26+
27+
## Create a Custom `DGraph`
28+
29+
You can also define a `DGraph` from your own data.
30+
31+
### Construct `DGraph` From Raw Tensors
32+
33+
#### Define Temporal Edges
34+
35+
- `edge_index`: shape `[num_edge_events, 2]`
36+
- `edge_timestamps`: shape `[num_edge_events]`
37+
- `edge_feats`: shape `[num_edge_events, D_edge]` (optional)
38+
39+
```python
40+
import torch
41+
42+
edge_index = torch.LongTensor([[2, 2], [2, 4], [1, 8]])
43+
edge_timestamps = torch.LongTensor([1, 5, 20])
44+
edge_feats = torch.rand(3, 5) # optional edge features
45+
```
46+
47+
#### Define Node Events (Optional)
48+
49+
- `node_timestamps`: shape `[num_node_events]`
50+
- `node_ids`: shape `[num_node_events]`
51+
- `dynamic_node_feats`: shape `[num_node_events, D_node_dynamic]`
52+
- `static_node_feats`: shape `[num_nodes, D_node_static]` (optional)
53+
54+
```python
55+
node_timestamps = torch.LongTensor([1, 2, 3])
56+
node_ids = torch.LongTensor([2, 4, 6])
57+
dynamic_node_feats = torch.rand([3, 5])
58+
static_node_feats = torch.rand(9, 11)
59+
```
60+
61+
#### Construct the `DGraph`
62+
63+
```python
64+
from tgm import DGraph
65+
66+
dg = DGraph.from_raw(
67+
edge_timestamps=edge_timestamps,
68+
edge_index=edge_index,
69+
edge_feats=edge_feats,
70+
node_timestamps=node_timestamps,
71+
node_ids=node_ids,
72+
dynamic_node_feats=dynamic_node_feats,
73+
static_node_feats=static_node_feats,
74+
time_delta='s', # second-wise granularity
75+
device='cuda', # move graph to GPU
76+
)
77+
```
78+
79+
### Construct `DGraph` from Pandas DataFrames
80+
81+
```python
82+
import pandas as pd
83+
84+
edge_df = pd.DataFrame({
85+
'src': [2, 2, 1],
86+
'dst': [2, 4, 8],
87+
't': [1, 5, 10],
88+
'edge_feat': [torch.rand(5).tolist() for _ in range(3)],
89+
})
90+
91+
dynamic_node_df = pd.DataFrame({
92+
'node': [2, 4, 6],
93+
't': [1, 2, 3],
94+
'dynamic_node_feat': [torch.rand(5).tolist() for _ in range(3)],
95+
})
96+
97+
static_node_df = pd.DataFrame({
98+
'static_node_feat': [torch.rand(11).tolist() for _ in range(9)]
99+
})
100+
101+
dg = DGraph.from_pandas(
102+
edge_df=edge_df,
103+
edge_src_col='src',
104+
edge_dst_col='dst',
105+
edge_time_col='t',
106+
edge_feats_col='edge_feat',
107+
node_df=dynamic_node_df,
108+
node_id_col='node',
109+
node_time_col='t',
110+
dynamic_node_feats_col='dynamic_node_feat',
111+
static_node_feats_df=static_node_df,
112+
static_node_feats_col='static_node_feat',
113+
time_delta='s', # second-wise granularity
114+
device='cuda', # move graph to GPU
115+
)
116+
```
117+
118+
### Construct `DGraph` from CSV Files
119+
120+
To load graph data from CSV files, use `DGraph.from_csv()`.
121+
See [`tgm/graph.py`](https://github.com/tgm-team/tgm/blob/main/tgm/graph.py) for details.
122+
123+
______________________________________________________________________
124+
125+
## Accessing `DGraph` Properties
126+
127+
`DGraph` objects act as views over the underlying data. You can access properties and perform slicing operations.
128+
129+
The number of nodes is computed as `max(node_ids) + 1`. If the `DGraph` is empty, `start_time` and `end_time` are `None`.
130+
131+
```python
132+
print('=== Graph Properties ===')
133+
print(f'Start time : {dg.start_time}') # 1
134+
print(f'End time : {dg.end_time}') # 10
135+
print(f'Number of nodes : {dg.num_nodes}') # 9
136+
print(f'Number of edge events : {dg.num_edges}') # 3
137+
print(f'Number of timestamps : {dg.num_timestamps}') # or len(dg); 5
138+
print(f'Total events (edge+node) : {dg.num_events}') # 6
139+
print(f'Edge feature dimension : {dg.edge_feats_dim}') # 5
140+
print(f'Static node feature dim : {dg.static_node_feats_dim}') # 11
141+
print(f'Dynamic node feature dim : {dg.dynamic_node_feats_dim}') # 5
142+
print('==========================')
143+
```
144+
145+
______________________________________________________________________
146+
147+
## Slicing `DGraph`
148+
149+
You can slice temporal data using `slice_time()`. This returns a new `DGraph` containing only events within the specified time range (end time exclusive). Slicing is a lightweight operation since the underlying data storage is shared across `DGraph` instances.
150+
151+
```python
152+
start_time, end_time = 5, 10
153+
sliced_dg = dg.slice_time(start_time, end_time)
154+
```
155+
156+
______________________________________________________________________
157+
158+
## Using `DGDataLoader` and Hooks
159+
160+
TGM integrates operations like **negative sampling** and **neighbor sampling** into the data loader via hooks.
161+
162+
```python
163+
from tgm.loader import DGDataLoader
164+
from tgm.hooks import NegativeEdgeSamplerHook, RecencyNeighborHook
165+
166+
neg_hook = NegativeEdgeSamplerHook(low=0, high=train_dg.num_nodes)
167+
168+
# Sample 20 1-hop neighbors per node
169+
nbr_hook = RecencyNeighborHook(num_nbrs=[20], num_nodes=train_dg.num_nodes)
170+
171+
train_loader = DGDataLoader(
172+
train_dg,
173+
hook=[neg_hook, nbr_hook],
174+
batch_size=200,
175+
)
176+
```
177+
178+
You can also iterate over time windows (instead of event counts) by specifying a `TimeDeltaDG` in the data loader constructor. For this to work, the underlying graph must be non-ordered (e.g. time granularity `'s'`).
179+
180+
______________________________________________________________________
181+
182+
## Iterating Over Batches with `DGBatch`
183+
184+
Each batch from `DGDataLoader` is a `DGBatch`.
185+
186+
```python
187+
iter_loader = iter(train_loader)
188+
batch = next(iter_loader)
189+
190+
print('=== Batch of 200 edges ===')
191+
print(f'Source nodes shape : {batch.src.shape}')
192+
print(f'Destination nodes shape : {batch.dst.shape}')
193+
print(f'Timestamps shape : {batch.time.shape}')
194+
print(f'Negative destinations shape: {batch.neg.shape}')
195+
print(f'1-hop neighbors shape : {batch.nbr_nids[1].shape}')
196+
print('===========================')
197+
```
198+
199+
The materialized batch will contain, at a minimum the batch of edges, features and nodes on the appropriate device. Hooks inject additional attributes at runtime (e.g. `batch.neg`).
200+
201+
______________________________________________________________________
202+
203+
## Feedback
204+
205+
Please feel free to reach out to us if anything is unclear or unintuitive. We are happy to discuss and improve your experience with TGM.

examples/linkproppred/TGB/edgebank.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from tgb.linkproppred.evaluate import Evaluator
77
from tqdm import tqdm
88

9-
from tgm import DGraph
9+
from tgm import DGData, DGraph
1010
from tgm.hooks import TGBNegativeEdgeSamplerHook
1111
from tgm.loader import DGDataLoader
1212
from tgm.nn import EdgeBankPredictor
@@ -48,8 +48,8 @@ def eval(
4848
y_pred = model(query_src, query_dst)
4949
# compute MRR
5050
input_dict = {
51-
'y_pred_pos': np.array([y_pred[0]]),
52-
'y_pred_neg': np.array(y_pred[1:]),
51+
'y_pred_pos': y_pred[0].detach().cpu().numpy(),
52+
'y_pred_neg': y_pred[1:].detach().cpu().numpy(),
5353
'eval_metric': [eval_metric],
5454
}
5555
perf_list.append(evaluator.eval(input_dict)[eval_metric])
@@ -69,9 +69,10 @@ def eval(
6969
dataset.load_val_ns()
7070
dataset.load_test_ns()
7171

72-
train_dg = DGraph(args.dataset, split='train')
73-
val_dg = DGraph(args.dataset, split='val')
74-
test_dg = DGraph(args.dataset, split='test')
72+
train_data, val_data, test_data = DGData.from_tgb(args.dataset).split()
73+
train_dg = DGraph(train_data)
74+
val_dg = DGraph(val_data)
75+
test_dg = DGraph(test_data)
7576

7677
train_data = train_dg.materialize(materialize_features=False)
7778
val_loader = DGDataLoader(

0 commit comments

Comments
 (0)