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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
|
import os
from typing import Callable, List, Optional
import torch
from torch_geometric.data import (
Data,
InMemoryDataset,
download_url,
extract_zip,
)
from torch_geometric.io import fs
class GDELTLite(InMemoryDataset):
r"""The (reduced) version of the Global Database of Events, Language, and
Tone (GDELT) dataset used in the `"Do We Really Need Complicated Model
Architectures for Temporal Networks?" <http://arxiv.org/abs/2302.11636>`_
paper, consisting of events collected from 2016 to 2020.
Each node (actor) holds a 413-dimensional multi-hot feature vector that
represents CAMEO codes attached to the corresponding actor to server.
Each edge (event) holds a timestamp and a 186-dimensional multi-hot vector
representing CAMEO codes attached to the corresponding event to server.
Args:
root (str): Root directory where the dataset should be saved.
transform (callable, optional): A function/transform that takes in an
:obj:`torch_geometric.data.Data` object and returns a transformed
version. The data object will be transformed before every access.
(default: :obj:`None`)
pre_transform (callable, optional): A function/transform that takes in
an :obj:`torch_geometric.data.Data` object and returns a
transformed version. The data object will be transformed before
being saved to disk. (default: :obj:`None`)
force_reload (bool, optional): Whether to re-process the dataset.
(default: :obj:`False`)
**STATS:**
.. list-table::
:widths: 10 10 10 10
:header-rows: 1
* - #nodes
- #edges
- #features
- #classes
* - 8,831
- 1,912,909
- 413
-
"""
url = 'https://data.pyg.org/datasets/gdelt_lite.zip'
def __init__(
self,
root: str,
transform: Optional[Callable] = None,
pre_transform: Optional[Callable] = None,
force_reload: bool = False,
) -> None:
super().__init__(root, transform, pre_transform,
force_reload=force_reload)
self.load(self.processed_paths[0])
@property
def raw_file_names(self) -> List[str]:
return ['node_features.pt', 'edges.csv', 'edge_features.pt']
@property
def processed_file_names(self) -> str:
return 'data.pt'
def download(self) -> None:
path = download_url(self.url, self.raw_dir)
extract_zip(path, self.raw_dir)
os.unlink(path)
def process(self) -> None:
import pandas as pd
x = fs.torch_load(self.raw_paths[0])
df = pd.read_csv(self.raw_paths[1])
edge_attr = fs.torch_load(self.raw_paths[2])
row = torch.from_numpy(df['src'].values)
col = torch.from_numpy(df['dst'].values)
edge_index = torch.stack([row, col], dim=0)
time = torch.from_numpy(df['time'].values).to(torch.long)
data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, time=time)
data = data if self.pre_transform is None else self.pre_transform(data)
self.save([data], self.processed_paths[0])
|