k3-node 1.0.0__py3-none-any.whl

k3_node/datasets/dblp.py

@@ -0,0 +1,106 @@
+import os
+import os.path as osp
+from itertools import product
+from typing import Callable, List, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import HeteroData, InMemoryDataset
+from k3_node.io import fs
+
+
+class DBLP(InMemoryDataset):
+    r"""A subset of the DBLP computer science bibliography website containing
+    four types of entities: authors, papers, terms, and conferences.
+    """
+
+    url = "https://www.dropbox.com/s/yh4grpeks87ugr2/DBLP_processed.zip?dl=1"
+
+    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) -> List[str]:
+        return [
+            "adjM.npz",
+            "features_0.npz",
+            "features_1.npz",
+            "features_2.npy",
+            "labels.npy",
+            "node_types.npy",
+            "train_val_test_idx.npz",
+        ]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        zip_path = osp.join(self.raw_dir, "DBLP_processed.zip")
+        fs.cp(self.url, zip_path, extract=True)
+        if osp.exists(zip_path):
+            fs.rm(zip_path)
+
+    def process(self):
+        import scipy.sparse as sp
+
+        data = HeteroData()
+        node_types = ["author", "paper", "term", "conference"]
+
+        for i, node_type in enumerate(node_types[:2]):
+            feat_path = osp.join(self.raw_dir, f"features_{i}.npz")
+            x = sp.load_npz(feat_path)
+            data[node_type].x = ops.convert_to_tensor(
+                np.array(x.todense(), dtype=np.float32), dtype="float32"
+            )
+
+        x2 = np.load(osp.join(self.raw_dir, "features_2.npy"))
+        data["term"].x = ops.convert_to_tensor(x2.astype(np.float32), dtype="float32")
+
+        node_type_idx = np.load(osp.join(self.raw_dir, "node_types.npy"))
+        data["conference"].num_nodes = int((node_type_idx == 3).sum())
+
+        y = np.load(osp.join(self.raw_dir, "labels.npy"))
+        data["author"].y = ops.convert_to_tensor(y.astype(np.int64), dtype="int64")
+
+        split = np.load(osp.join(self.raw_dir, "train_val_test_idx.npz"))
+        for name in ["train", "val", "test"]:
+            idx = split[f"{name}_idx"]
+            mask = np.zeros(data["author"].num_nodes, dtype=bool)
+            mask[idx] = True
+            data["author"][f"{name}_mask"] = ops.convert_to_tensor(mask, dtype="bool")
+
+        s = {}
+        N_a = data["author"].num_nodes
+        N_p = data["paper"].num_nodes
+        N_t = data["term"].num_nodes
+        N_c = data["conference"].num_nodes
+        s["author"] = (0, N_a)
+        s["paper"] = (N_a, N_a + N_p)
+        s["term"] = (N_a + N_p, N_a + N_p + N_t)
+        s["conference"] = (N_a + N_p + N_t, N_a + N_p + N_t + N_c)
+
+        A = sp.load_npz(osp.join(self.raw_dir, "adjM.npz"))
+        for src, dst in product(node_types, node_types):
+            A_sub = A[s[src][0] : s[src][1], s[dst][0] : s[dst][1]].tocoo()
+            if A_sub.nnz > 0:
+                row = np.array(A_sub.row, dtype=np.int64)
+                col = np.array(A_sub.col, dtype=np.int64)
+                edge_index = np.stack([row, col], axis=0)
+                data[src, dst].edge_index = ops.convert_to_tensor(edge_index, dtype="int64")
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+

k3_node/datasets/digits.py

@@ -0,0 +1,63 @@
+from typing import Callable, Optional
+import numpy as np
+
+from k3_node.data import Data, InMemoryDataset
+
+
+class Digits(InMemoryDataset):
+    r"""The handwritten digits of scikit-learn (1,797 images of 8x8 pixels) as graphs.
+
+    A small, offline stand-in for PyG's ``MNISTSuperpixels``. Every non-blank pixel becomes a node
+    with its intensity (in ``[0, 1]``) as the only feature ``x`` and its 2D position ``pos``, scaled
+    to ``[0, 28)`` like MNIST. Neighboring pixels (including diagonals) are connected in both
+    directions. ``y`` holds the digit (10 classes).
+
+    Args:
+        train (bool): If ``True``, loads the first 1,500 images, otherwise the remaining 297.
+            (default: ``True``)
+        transform (callable, optional): A function applied to each graph when it is accessed.
+
+    Example:
+        ```python
+        from k3_node.datasets import Digits
+
+        dataset = Digits(train=True)
+        print(len(dataset), dataset.num_classes)  # 1500 10
+        ```
+    """
+
+    num_train = 1500
+
+    def __init__(self, train: bool = True, transform: Optional[Callable] = None):
+        super().__init__(None, transform)
+        from sklearn.datasets import load_digits
+
+        images, labels = load_digits(return_X_y=True)
+        images = images.reshape(-1, 8, 8) / 16.0
+        split = slice(0, self.num_train) if train else slice(self.num_train, None)
+        self.data, self.slices = self.collate(
+            [self._to_graph(img, y) for img, y in zip(images[split], labels[split])]
+        )
+
+    @staticmethod
+    def _to_graph(image, label):
+        rows, cols = np.nonzero(image)
+        index = -np.ones((8, 8), dtype=np.int64)
+        index[rows, cols] = np.arange(len(rows))
+        src, dst = [], []
+        for dr in (-1, 0, 1):
+            for dc in (-1, 0, 1):
+                if dr == 0 and dc == 0:
+                    continue
+                r, c = rows + dr, cols + dc
+                ok = (r >= 0) & (r < 8) & (c >= 0) & (c < 8)
+                nbr = index[r[ok], c[ok]]
+                keep = nbr >= 0
+                src.append(np.nonzero(ok)[0][keep])
+                dst.append(nbr[keep])
+        return Data(
+            x=image[rows, cols][:, None].astype(np.float32),
+            pos=(np.stack([cols, rows], axis=1) * 3.5 + 1.75).astype(np.float32),
+            edge_index=np.stack([np.concatenate(src), np.concatenate(dst)]).astype(np.int64),
+            y=np.array([label], dtype=np.int64),
+        )

k3_node/datasets/email_eu_core.py

@@ -0,0 +1,60 @@
+import os
+import os.path as osp
+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 EmailEUCore(InMemoryDataset):
+    r"""An e-mail communication network of a large European research institution."""
+
+    urls = [
+        "https://snap.stanford.edu/data/email-Eu-core.txt.gz",
+        "https://snap.stanford.edu/data/email-Eu-core-department-labels.txt.gz",
+    ]
+
+    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) -> List[str]:
+        return ["email-Eu-core.txt", "email-Eu-core-department-labels.txt"]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        for url in self.urls:
+            filename = osp.basename(url)
+            gz_path = osp.join(self.raw_dir, filename)
+            fs.cp(url, gz_path, extract=True)
+            if osp.exists(gz_path):
+                fs.rm(gz_path)
+
+    def process(self):
+        edge_index = np.loadtxt(self.raw_paths[0], dtype=np.int64).T
+        labels = np.loadtxt(self.raw_paths[1], dtype=np.int64)
+        y = labels[:, 1]
+
+        data = Data(
+            edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+            y=ops.convert_to_tensor(y, dtype="int64"),
+            num_nodes=int(y.shape[0]),
+        )
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
+

k3_node/datasets/entities.py

@@ -0,0 +1,158 @@
+import logging
+import os
+import os.path as osp
+from collections import Counter
+from typing import Any, Callable, List, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, HeteroData, InMemoryDataset
+from k3_node.io import fs
+
+
+class Entities(InMemoryDataset):
+    r"""The relational entities networks "AIFB", "MUTAG", "BGS" and "AM"."""
+
+    url = "https://data.dgl.ai/dataset/{}.tgz"
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        hetero: bool = False,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        self.name = name.lower()
+        self.hetero = hetero
+        assert self.name in ["aifb", "am", "mutag", "bgs"]
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_dir(self) -> str:
+        return osp.join(self.root, self.name, "raw")
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root, self.name, "processed")
+
+    @property
+    def num_relations(self) -> int:
+        return int(ops.convert_to_numpy(self._data.edge_type).max()) + 1
+
+    @property
+    def num_classes(self) -> int:
+        return int(ops.convert_to_numpy(self._data.train_y).max()) + 1
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [
+            f"{self.name}_stripped.nt.gz",
+            "completeDataset.tsv",
+            "trainingSet.tsv",
+            "testSet.tsv",
+        ]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "hetero_data.pt" if self.hetero else "data.pt"
+
+    def download(self):
+        tgz_path = osp.join(self.raw_dir, f"{self.name}.tgz")  # extracted next to it, into raw_dir
+        fs.cp(self.url.format(self.name), tgz_path, extract=True)
+        if osp.exists(tgz_path):
+            fs.rm(tgz_path)
+
+    def process(self):
+        import gzip
+        import rdflib as rdf
+
+        graph_file, task_file, train_file, test_file = self.raw_paths
+
+        g = rdf.Graph()
+        with gzip.open(graph_file, "rb") as f:
+            g.parse(file=f, format="nt")
+
+        freq = Counter(g.predicates())
+        relations = sorted(set(g.predicates()), key=lambda p: -freq.get(p, 0))
+        subjects = set(g.subjects())
+        objects = set(g.objects())
+        nodes = list(subjects.union(objects))
+
+        N = len(nodes)
+        R = 2 * len(relations)
+
+        relations_dict = {rel: i for i, rel in enumerate(relations)}
+        nodes_dict = {str(node): i for i, node in enumerate(nodes)}
+
+        edges = []
+        for s, p, o in g.triples((None, None, None)):
+            src, dst = nodes_dict[str(s)], nodes_dict[str(o)]
+            rel = relations_dict[p]
+            edges.append([src, dst, 2 * rel])
+            edges.append([dst, src, 2 * rel + 1])
+
+        edge = np.array(edges, dtype=np.int64).T
+        sort_key = N * R * edge[0] + R * edge[1] + edge[2]
+        perm = np.argsort(sort_key)
+        edge = edge[:, perm]
+
+        edge_index, edge_type = edge[:2], edge[2]
+
+        if self.name == "am":
+            label_header = "label_cateogory"
+            nodes_header = "proxy"
+        elif self.name == "aifb":
+            label_header = "label_affiliation"
+            nodes_header = "person"
+        elif self.name == "mutag":
+            label_header = "label_mutagenic"
+            nodes_header = "bond"
+        elif self.name == "bgs":
+            label_header = "label_lithogenesis"
+            nodes_header = "rock"
+
+        import pandas as pd
+
+        labels_df = pd.read_csv(task_file, sep="\t")
+        labels_set = set(labels_df[label_header].values.tolist())
+        labels_dict = {lab: i for i, lab in enumerate(list(labels_set))}
+
+        train_labels_df = pd.read_csv(train_file, sep="\t")
+        train_indices, train_labels = [], []
+        for nod, lab in zip(train_labels_df[nodes_header].values, train_labels_df[label_header].values):
+            train_indices.append(nodes_dict[nod])
+            train_labels.append(labels_dict[lab])
+
+        train_idx = np.array(train_indices, dtype=np.int64)
+        train_y = np.array(train_labels, dtype=np.int64)
+
+        test_labels_df = pd.read_csv(test_file, sep="\t")
+        test_indices, test_labels = [], []
+        for nod, lab in zip(test_labels_df[nodes_header].values, test_labels_df[label_header].values):
+            test_indices.append(nodes_dict[nod])
+            test_labels.append(labels_dict[lab])
+
+        test_idx = np.array(test_indices, dtype=np.int64)
+        test_y = np.array(test_labels, 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"),
+            train_idx=ops.convert_to_tensor(train_idx, dtype="int64"),
+            train_y=ops.convert_to_tensor(train_y, dtype="int64"),
+            test_idx=ops.convert_to_tensor(test_idx, dtype="int64"),
+            test_y=ops.convert_to_tensor(test_y, dtype="int64"),
+            num_nodes=N,
+        )
+
+        if self.hetero:
+            data = data.to_heterogeneous(node_type_names=["v"])
+
+        self.save([data], self.processed_paths[0])
+
+    def __repr__(self) -> str:
+        return f"{self.name.upper()}{self.__class__.__name__}()"
+

k3_node/datasets/explainer_dataset.py

@@ -0,0 +1,101 @@
+from typing import Any, Callable, Dict, Optional, Union
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.datasets.graph_generator import GraphGenerator
+from k3_node.datasets.motif_generator import MotifGenerator
+
+
+class ExplainerDataset(InMemoryDataset):
+    r"""Generates a synthetic dataset for evaluating explainability algorithms,
+    as described in the "GNNExplainer: Generating Explanations for Graph Neural Networks" paper.
+
+    Args:
+        graph_generator (GraphGenerator or str): The graph generator to use.
+        motif_generator (MotifGenerator or str): The motif generator to use.
+        num_motifs (int): The number of motifs to attach to the graph.
+        num_graphs (int, optional): The number of graphs to generate. (default: 1)
+        graph_generator_kwargs (dict, optional): Keyword arguments for graph generator.
+        motif_generator_kwargs (dict, optional): Keyword arguments for motif generator.
+        transform (callable, optional): Transform function.
+    """
+
+    def __init__(
+        self,
+        graph_generator: Union[GraphGenerator, str],
+        motif_generator: Union[MotifGenerator, str],
+        num_motifs: int,
+        num_graphs: int = 1,
+        graph_generator_kwargs: Optional[Dict[str, Any]] = None,
+        motif_generator_kwargs: Optional[Dict[str, Any]] = None,
+        transform: Optional[Callable] = None,
+    ):
+        super().__init__(root=None, transform=transform)
+
+        if num_motifs <= 0:
+            raise ValueError(f"At least one motif needs to be attached (got {num_motifs})")
+
+        self.graph_generator = GraphGenerator.resolve(
+            graph_generator, **(graph_generator_kwargs or {})
+        )
+        self.motif_generator = MotifGenerator.resolve(
+            motif_generator, **(motif_generator_kwargs or {})
+        )
+        self.num_motifs = num_motifs
+
+        data_list = [self.get_graph() for _ in range(num_graphs)]
+        self.data, self.slices = self.collate(data_list)
+
+    def get_graph(self) -> Data:
+        data = self.graph_generator()
+        edge_index_np = ops.convert_to_numpy(data.edge_index)
+        num_nodes = data.num_nodes
+        num_edges = edge_index_np.shape[1]
+
+        edge_indices = [edge_index_np]
+        node_masks = [np.zeros(num_nodes, dtype=np.float32)]
+        edge_masks = [np.zeros(num_edges, dtype=np.float32)]
+        ys = [np.zeros(num_nodes, dtype=np.int64)]
+
+        connecting_nodes = np.random.permutation(num_nodes)[: self.num_motifs]
+        for i in connecting_nodes.tolist():
+            motif = self.motif_generator()
+            motif_ei = ops.convert_to_numpy(motif.edge_index)
+            motif_num_nodes = motif.num_nodes
+            motif_num_edges = motif_ei.shape[1]
+
+            edge_indices.append(motif_ei + num_nodes)
+            node_masks.append(np.ones(motif_num_nodes, dtype=np.float32))
+            edge_masks.append(np.ones(motif_num_edges, dtype=np.float32))
+
+            j = int(np.random.randint(0, motif_num_nodes)) + num_nodes
+            edge_indices.append(np.array([[i, j], [j, i]], dtype=np.int64))
+            edge_masks.append(np.zeros(2, dtype=np.float32))
+
+            if hasattr(motif, "y") and motif.y is not None:
+                motif_y = ops.convert_to_numpy(motif.y)
+                if np.min(motif_y) == 0:
+                    ys.append(motif_y + 1)
+                else:
+                    ys.append(motif_y)
+            else:
+                ys.append(np.ones(motif_num_nodes, dtype=np.int64))
+
+            num_nodes += motif_num_nodes
+
+        return Data(
+            edge_index=ops.convert_to_tensor(np.concatenate(edge_indices, axis=1), dtype="int64"),
+            y=ops.convert_to_tensor(np.concatenate(ys, axis=0), dtype="int64"),
+            edge_mask=ops.convert_to_tensor(np.concatenate(edge_masks, axis=0), dtype="float32"),
+            node_mask=ops.convert_to_tensor(np.concatenate(node_masks, axis=0), dtype="float32"),
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({len(self)}, "
+            f"graph_generator={self.graph_generator}, "
+            f"motif_generator={self.motif_generator}, "
+            f"num_motifs={self.num_motifs})"
+        )
+

k3_node/datasets/facebook.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 FacebookPagePage(InMemoryDataset):
+    r"""The Facebook Page-Page network dataset."""
+
+    url = "https://graphmining.ai/datasets/ptg/facebook.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 "facebook.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])
+