k3-node 1.0.0__py3-none-any.whl

k3_node/datasets/fake.py

@@ -0,0 +1,256 @@
+import random
+from collections import defaultdict
+from itertools import product
+from typing import Callable, Dict, List, Optional, Tuple, Union
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, HeteroData, InMemoryDataset
+from k3_node.layers.conv.utils import remove_self_loops
+from k3_node.transforms.utils import to_undirected
+from k3_node.utils.graph import coalesce
+
+
+def get_num_nodes(avg_num_nodes: int, avg_degree: float) -> int:
+    min_num_nodes = max(3 * avg_num_nodes // 4, int(avg_degree))
+    max_num_nodes = 5 * avg_num_nodes // 4
+    return random.randint(min_num_nodes, max_num_nodes)
+
+
+def get_num_channels(num_channels: int) -> int:
+    min_num_channels = 3 * num_channels // 4
+    max_num_channels = 5 * num_channels // 4
+    return random.randint(min_num_channels, max_num_channels)
+
+
+def get_edge_index(
+    num_src_nodes: int,
+    num_dst_nodes: int,
+    avg_degree: float,
+    is_undirected: bool = False,
+    remove_loops: bool = False,
+):
+    num_edges = int(num_src_nodes * avg_degree)
+    row = np.random.randint(0, num_src_nodes, size=(num_edges,), dtype=np.int64)
+    col = np.random.randint(0, num_dst_nodes, size=(num_edges,), dtype=np.int64)
+    edge_index = np.stack([row, col], axis=0)
+
+    if remove_loops:
+        edge_index, _ = remove_self_loops(edge_index)
+
+    num_nodes = max(num_src_nodes, num_dst_nodes)
+    if is_undirected:
+        edge_index = to_undirected(edge_index, num_nodes=num_nodes)
+    else:
+        edge_index, _ = coalesce(edge_index, num_nodes=num_nodes)
+
+    return ops.convert_to_tensor(edge_index, dtype="int64")
+
+
+class FakeDataset(InMemoryDataset):
+    r"""A fake dataset that returns randomly generated `k3_node.data.Data` objects."""
+
+    def __init__(
+        self,
+        num_graphs: int = 1,
+        avg_num_nodes: int = 1000,
+        avg_degree: float = 10.0,
+        num_channels: int = 64,
+        edge_dim: int = 0,
+        num_classes: int = 10,
+        task: str = "auto",
+        is_undirected: bool = True,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        **kwargs: Union[int, Tuple[int, ...]],
+    ):
+        super().__init__(None, transform)
+
+        if task == "auto":
+            task = "graph" if num_graphs > 1 else "node"
+        assert task in ["node", "graph"]
+
+        self.num_graphs_val = num_graphs
+        self.avg_num_nodes = max(avg_num_nodes, int(avg_degree))
+        self.avg_degree = max(avg_degree, 1)
+        self.num_channels = num_channels
+        self.edge_dim = edge_dim
+        self._num_classes = num_classes
+        self.task = task
+        self.is_undirected = is_undirected
+        self.kwargs = kwargs
+
+        data_list = [self.generate_data() for _ in range(max(num_graphs, 1))]
+        self.data, self.slices = self.collate(data_list)
+
+    def __repr__(self) -> str:
+        return f"FakeDataset({self.num_graphs_val})" if self.num_graphs_val > 1 else "FakeDataset()"
+
+    def generate_data(self) -> Data:
+        num_nodes = get_num_nodes(self.avg_num_nodes, self.avg_degree)
+        data = Data()
+
+        if self._num_classes > 0 and self.task == "node":
+            data.y = ops.convert_to_tensor(
+                np.random.randint(0, self._num_classes, size=(num_nodes,), dtype=np.int64),
+                dtype="int64",
+            )
+        elif self._num_classes > 0 and self.task == "graph":
+            data.y = ops.convert_to_tensor(
+                np.array([random.randint(0, self._num_classes - 1)], dtype=np.int64),
+                dtype="int64",
+            )
+
+        data.edge_index = get_edge_index(
+            num_nodes, num_nodes, self.avg_degree, self.is_undirected, remove_loops=True
+        )
+
+        if self.num_channels > 0:
+            x = np.random.randn(num_nodes, self.num_channels).astype(np.float32)
+            if self._num_classes > 0 and self.task == "node":
+                y_np = ops.convert_to_numpy(data.y)
+                x = x + y_np[:, None]
+            elif self._num_classes > 0 and self.task == "graph":
+                y_np = ops.convert_to_numpy(data.y)
+                x = x + y_np
+            data.x = ops.convert_to_tensor(x, dtype="float32")
+        else:
+            data.num_nodes = num_nodes
+
+        num_edges = int(ops.shape(data.edge_index)[1])
+        if self.edge_dim > 1:
+            data.edge_attr = ops.convert_to_tensor(
+                np.random.rand(num_edges, self.edge_dim).astype(np.float32),
+                dtype="float32",
+            )
+        elif self.edge_dim == 1:
+            data.edge_weight = ops.convert_to_tensor(
+                np.random.rand(num_edges).astype(np.float32),
+                dtype="float32",
+            )
+
+        for feature_name, feature_shape in self.kwargs.items():
+            shape = (feature_shape,) if isinstance(feature_shape, int) else feature_shape
+            setattr(
+                data,
+                feature_name,
+                ops.convert_to_tensor(np.random.randn(*shape).astype(np.float32), dtype="float32"),
+            )
+
+        return data
+
+
+class FakeHeteroDataset(InMemoryDataset):
+    r"""A fake dataset that returns randomly generated `k3_node.data.HeteroData` objects."""
+
+    def __init__(
+        self,
+        num_graphs: int = 1,
+        num_node_types: int = 3,
+        num_edge_types: int = 6,
+        avg_num_nodes: int = 1000,
+        avg_degree: float = 10.0,
+        avg_num_channels: int = 64,
+        edge_dim: int = 0,
+        num_classes: int = 10,
+        task: str = "auto",
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        **kwargs: Union[int, Tuple[int, ...]],
+    ):
+        super().__init__(None, transform)
+
+        if task == "auto":
+            task = "graph" if num_graphs > 1 else "node"
+        assert task in ["node", "graph"]
+
+        self.num_graphs_val = num_graphs
+        self.node_types = [f"v{i}" for i in range(max(num_node_types, 1))]
+
+        edge_types: List[Tuple[str, str]] = []
+        edge_type_product = list(product(self.node_types, self.node_types))
+        while len(edge_types) < max(num_edge_types, 1):
+            edge_types.extend(edge_type_product)
+        random.shuffle(edge_types)
+
+        self.edge_types: List[Tuple[str, str, str]] = []
+        count: Dict[Tuple[str, str], int] = defaultdict(int)
+        for edge_type in edge_types[: max(num_edge_types, 1)]:
+            rel = f"e{count[edge_type]}"
+            count[edge_type] += 1
+            self.edge_types.append((edge_type[0], rel, edge_type[1]))
+
+        self.avg_num_nodes = max(avg_num_nodes, int(avg_degree))
+        self.avg_degree = max(avg_degree, 1)
+        self.avg_num_channels = avg_num_channels
+        self.edge_dim = edge_dim
+        self._num_classes = num_classes
+        self.task = task
+        self.kwargs = kwargs
+
+        data_list = [self.generate_data() for _ in range(max(num_graphs, 1))]
+        self.data, self.slices = self.collate(data_list)
+
+    def __repr__(self) -> str:
+        return f"FakeHeteroDataset({self.num_graphs_val})" if self.num_graphs_val > 1 else "FakeHeteroDataset()"
+
+    def generate_data(self) -> HeteroData:
+        data = HeteroData()
+
+        for node_type in self.node_types:
+            num_nodes = get_num_nodes(self.avg_num_nodes, self.avg_degree)
+            num_channels = get_num_channels(self.avg_num_channels)
+            store = data[node_type]
+
+            if self.avg_num_channels > 0:
+                store.x = ops.convert_to_tensor(
+                    np.random.randn(num_nodes, num_channels).astype(np.float32),
+                    dtype="float32",
+                )
+            else:
+                store.num_nodes = num_nodes
+
+            if self._num_classes > 0 and self.task == "node":
+                store.y = ops.convert_to_tensor(
+                    np.random.randint(0, self._num_classes, size=(num_nodes,), dtype=np.int64),
+                    dtype="int64",
+                )
+
+        for edge_type in self.edge_types:
+            src, rel, dst = edge_type
+            store = data[edge_type]
+            store.edge_index = get_edge_index(
+                data[src].num_nodes,
+                data[dst].num_nodes,
+                self.avg_degree,
+                is_undirected=False,
+                remove_loops=False,
+            )
+
+            num_edges = int(ops.shape(store.edge_index)[1])
+            if self.edge_dim > 1:
+                store.edge_attr = ops.convert_to_tensor(
+                    np.random.rand(num_edges, self.edge_dim).astype(np.float32),
+                    dtype="float32",
+                )
+            elif self.edge_dim == 1:
+                store.edge_weight = ops.convert_to_tensor(
+                    np.random.rand(num_edges).astype(np.float32),
+                    dtype="float32",
+                )
+
+        if self._num_classes > 0 and self.task == "graph":
+            data.y = ops.convert_to_tensor(
+                np.array([random.randint(0, self._num_classes - 1)], dtype=np.int64),
+                dtype="int64",
+            )
+
+        for feature_name, feature_shape in self.kwargs.items():
+            shape = (feature_shape,) if isinstance(feature_shape, int) else feature_shape
+            setattr(
+                data,
+                feature_name,
+                ops.convert_to_tensor(np.random.randn(*shape).astype(np.float32), dtype="float32"),
+            )
+
+        return data

k3_node/datasets/freebase.py

@@ -0,0 +1,90 @@
+from typing import Callable, List, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.io import fs
+
+
+class FB15k_237(InMemoryDataset):
+    r"""The FB15K237 dataset containing 14,541 entities, 237 relations and 310,116 fact triples.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): "train", "val", or "test". (default: "train")
+        transform (callable, optional): Transform function.
+        pre_transform (callable, optional): Pre-transform function.
+        force_reload (bool, optional): Whether to re-process the dataset.
+    """
+
+    url = "https://raw.githubusercontent.com/villmow/datasets_knowledge_embedding/master/FB15k-237"
+
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        if split not in {"train", "val", "test"}:
+            raise ValueError(f"Invalid split argument (got {split})")
+        self.split = split
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        idx = ["train", "val", "test"].index(split)
+        self.load(self.processed_paths[idx])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ["train.txt", "valid.txt", "test.txt"]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ["train_data.pt", "val_data.pt", "test_data.pt"]
+
+    def download(self):
+        for filename in self.raw_file_names:
+            fs.cp(f"{self.url}/{filename}", self.raw_dir)
+
+    def process(self):
+        # Map entities and relations to integer IDs
+        entities, relations = {}, {}
+        for path in self.raw_paths:
+            with open(path) as f:
+                lines = f.read().split("\n")[:-1]
+                for line in lines:
+                    parts = line.split()
+                    if len(parts) >= 3:
+                        s, r, d = parts[0], parts[1], parts[2]
+                        if s not in entities:
+                            entities[s] = len(entities)
+                        if d not in entities:
+                            entities[d] = len(entities)
+                        if r not in relations:
+                            relations[r] = len(relations)
+
+        for in_path, out_path in zip(self.raw_paths, self.processed_paths):
+            srcs, dsts, rels = [], [], []
+            with open(in_path) as f:
+                lines = f.read().split("\n")[:-1]
+                for line in lines:
+                    parts = line.split()
+                    if len(parts) >= 3:
+                        srcs.append(entities[parts[0]])
+                        rels.append(relations[parts[1]])
+                        dsts.append(entities[parts[2]])
+
+            edge_index = np.array([srcs, dsts], dtype=np.int64)
+            edge_type = np.array(rels, dtype=np.int64)
+
+            data = Data(
+                edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+                edge_type=ops.convert_to_tensor(edge_type, dtype="int64"),
+                num_nodes=len(entities),
+            )
+
+            if self.pre_transform is not None:
+                data = self.pre_transform(data)
+
+            self.save([data], out_path)
+

k3_node/datasets/geometric_shapes.py

@@ -0,0 +1,69 @@
+import glob
+import os
+import os.path as osp
+from typing import Callable, List, Optional
+
+import numpy as np
+
+from k3_node.data import InMemoryDataset
+from k3_node.data.download import download_url
+from k3_node.data.extract import extract_zip
+from k3_node.io.off import read_off
+
+
+class GeometricShapes(InMemoryDataset):
+    r"""Synthetic meshes of 40 geometric shapes such as cubes, spheres or pyramids (one training and
+    one test mesh per shape), as in PyG.
+
+    The graphs hold mesh faces (``face``) and vertex positions (``pos``) but no edges. Use
+    :class:`~k3_node.transforms.FaceToEdge` to turn a mesh into a graph, or
+    :class:`~k3_node.transforms.SamplePoints` to sample a point cloud from its surface.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        train (bool, optional): Loads the training meshes if ``True``, else the test meshes.
+        transform (callable, optional): A function applied to each graph when it is accessed.
+        pre_transform (callable, optional): A function applied to each graph before saving.
+        pre_filter (callable, optional): A function deciding which graphs to keep.
+        force_reload (bool, optional): Whether to re-process the dataset. (default: ``False``)
+    """
+
+    url = 'https://github.com/Yannick-S/geometric_shapes/raw/master/raw.zip'
+
+    def __init__(self, root: str, train: bool = True, transform: Optional[Callable] = None,
+                 pre_transform: Optional[Callable] = None, pre_filter: Optional[Callable] = None,
+                 force_reload: bool = False):
+        super().__init__(root, transform, pre_transform, pre_filter, force_reload=force_reload)
+        self.load(self.processed_paths[0] if train else self.processed_paths[1])
+
+    @property
+    def raw_file_names(self) -> str:
+        return '2d_circle'
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['training.pt', 'test.pt']
+
+    def download(self):
+        path = download_url(self.url, self.root)
+        extract_zip(path, self.root)
+        os.unlink(path)
+
+    def process(self):
+        self.save(self._process_set('train'), self.processed_paths[0])
+        self.save(self._process_set('test'), self.processed_paths[1])
+
+    def _process_set(self, split: str):
+        categories = sorted(x.split(os.sep)[-2] for x in glob.glob(osp.join(self.raw_dir, '*', '')))
+        data_list = []
+        for target, category in enumerate(categories):
+            for path in sorted(glob.glob(osp.join(self.raw_dir, category, split, '*.off'))):
+                data = read_off(path)
+                data.pos = data.pos - data.pos.mean(axis=0, keepdims=True)
+                data.y = np.array([target], dtype=np.int64)
+                data_list.append(data)
+        if self.pre_filter is not None:
+            data_list = [d for d in data_list if self.pre_filter(d)]
+        if self.pre_transform is not None:
+            data_list = [self.pre_transform(d) for d in data_list]
+        return data_list

k3_node/datasets/github.py

@@ -0,0 +1,51 @@
+from typing import Callable, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.io import fs
+
+
+class GitHub(InMemoryDataset):
+    r"""The GitHub Web and ML Developers dataset."""
+
+    url = "https://graphmining.ai/datasets/ptg/github.npz"
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> str:
+        return "github.npz"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        fs.cp(self.url, self.raw_dir)
+
+    def process(self):
+        data = np.load(self.raw_paths[0], allow_pickle=True)
+        x = data["features"].astype(np.float32)
+        y = data["target"].astype(np.int64)
+        edge_index = data["edges"].astype(np.int64).T
+
+        data_obj = Data(
+            x=ops.convert_to_tensor(x, dtype="float32"),
+            y=ops.convert_to_tensor(y, dtype="int64"),
+            edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+        )
+
+        if self.pre_transform is not None:
+            data_obj = self.pre_transform(data_obj)
+
+        self.save([data_obj], self.processed_paths[0])
+

k3_node/datasets/graph_generator/__init__.py

@@ -0,0 +1,6 @@
+from .base import GraphGenerator
+from .ba_graph import BAGraph
+from .er_graph import ERGraph
+
+__all__ = ["GraphGenerator", "BAGraph", "ERGraph"]
+

k3_node/datasets/graph_generator/ba_graph.py

@@ -0,0 +1,20 @@
+from k3_node.data import Data
+from k3_node.datasets.graph_generator.base import GraphGenerator
+from k3_node.utils.random import barabasi_albert_graph
+
+
+class BAGraph(GraphGenerator):
+    r"""Generates random Barabasi-Albert (BA) graphs."""
+
+    def __init__(self, num_nodes: int, num_edges: int):
+        super().__init__()
+        self.num_nodes = num_nodes
+        self.num_edges = num_edges
+
+    def __call__(self) -> Data:
+        edge_index = barabasi_albert_graph(self.num_nodes, self.num_edges)
+        return Data(num_nodes=self.num_nodes, edge_index=edge_index)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(num_nodes={self.num_nodes}, num_edges={self.num_edges})"
+

k3_node/datasets/graph_generator/base.py

@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from typing import Any
+from k3_node.data import Data
+
+
+class GraphGenerator(ABC):
+    r"""An abstract base class for generating synthetic graphs."""
+
+    @abstractmethod
+    def __call__(self) -> Data:
+        raise NotImplementedError
+
+    @staticmethod
+    def resolve(query: Any, *args: Any, **kwargs: Any) -> "GraphGenerator":
+        if isinstance(query, GraphGenerator):
+            return query
+        if isinstance(query, str):
+            query = query.lower()
+            if query in ["ba", "bagraph", "barabasi_albert"]:
+                from k3_node.datasets.graph_generator.ba_graph import BAGraph
+                return BAGraph(*args, **kwargs)
+            elif query in ["er", "ergraph", "erdos_renyi"]:
+                from k3_node.datasets.graph_generator.er_graph import ERGraph
+                return ERGraph(*args, **kwargs)
+        raise ValueError(f"Could not resolve graph generator: {query}")
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+

k3_node/datasets/graph_generator/er_graph.py

@@ -0,0 +1,21 @@
+from k3_node.data import Data
+from k3_node.datasets.graph_generator.base import GraphGenerator
+from k3_node.utils.random import erdos_renyi_graph
+
+
+class ERGraph(GraphGenerator):
+    r"""Generates random Erdos-Renyi (ER) graphs."""
+
+    def __init__(self, num_nodes: int, edge_prob: float, directed: bool = False):
+        super().__init__()
+        self.num_nodes = num_nodes
+        self.edge_prob = edge_prob
+        self.directed = directed
+
+    def __call__(self) -> Data:
+        edge_index = erdos_renyi_graph(self.num_nodes, self.edge_prob, directed=self.directed)
+        return Data(num_nodes=self.num_nodes, edge_index=edge_index)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(num_nodes={self.num_nodes}, edge_prob={self.edge_prob})"
+

k3_node/datasets/icews.py

@@ -0,0 +1,58 @@
+from typing import Callable, List, Optional
+
+import numpy as np
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.data.download import download_url
+
+
+class ICEWS18(InMemoryDataset):
+    r"""The ICEWS18 temporal knowledge graph (Integrated Crisis Early Warning System, events from
+    1/1/2018 to 10/31/2018 at a daily resolution), used by RE-Net: every graph is one event with
+    subject ``sub``, relation ``rel``, object ``obj`` and day ``t``.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): ``"train"``, ``"val"`` or ``"test"``. (default: ``"train"``)
+        transform (callable, optional): A function applied to each event when it is accessed.
+        pre_transform (callable, optional): A function applied to the events in time order before
+            saving, e.g. :meth:`~k3_node.models.RENet.pre_transform`.
+        force_reload (bool, optional): Whether to re-process the dataset. (default: ``False``)
+    """
+
+    url = 'https://github.com/INK-USC/RE-Net/raw/master/data/ICEWS18'
+    splits = [0, 373018, 419013, 468558]
+    num_nodes, num_rels = 23033, 256
+
+    def __init__(self, root: str, split: str = 'train', transform: Optional[Callable] = None,
+                 pre_transform: Optional[Callable] = None, force_reload: bool = False):
+        assert split in ['train', 'val', 'test']
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        self.load(self.processed_paths[['train', 'val', 'test'].index(split)])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [f'{name}.txt' for name in ['train', 'valid', 'test']]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['train.pt', 'val.pt', 'test.pt']
+
+    def download(self):
+        for filename in self.raw_file_names:
+            download_url(f'{self.url}/{filename}', self.raw_dir)
+
+    def process(self):
+        events = np.concatenate([np.loadtxt(path, delimiter='\t', usecols=range(4), dtype=np.int64)
+                                 for path in self.raw_paths])
+        events[:, 3] //= 24  # hours -> days
+        events -= events.min(axis=0, keepdims=True)
+        data_list = []
+        for sub, rel, obj, t in events.tolist():
+            data = Data(sub=sub, rel=rel, obj=obj, t=t)
+            if self.pre_transform is not None:
+                data = self.pre_transform(data)
+            data_list.append(data)
+        s = self.splits
+        for i in range(3):
+            self.save(data_list[s[i]:s[i + 1]], self.processed_paths[i])