k3-node 1.0.0__py3-none-any.whl

k3_node/datasets/actor.py

@@ -0,0 +1,101 @@
+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
+from k3_node.utils.graph import coalesce
+
+
+class Actor(InMemoryDataset):
+    r"""The actor-only induced subgraph of the film-director-actor-writer network
+    used in the "Geom-GCN: Geometric Graph Convolutional Networks" paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        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/graphdml-uiuc-jlu/geom-gcn/master"
+
+    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 [
+            "out1_node_feature_label.txt",
+            "out1_graph_edges.txt",
+            "film_split_0.6_0.2_0.npz",
+            "film_split_0.6_0.2_1.npz",
+            "film_split_0.6_0.2_2.npz",
+            "film_split_0.6_0.2_3.npz",
+            "film_split_0.6_0.2_4.npz",
+            "film_split_0.6_0.2_5.npz",
+            "film_split_0.6_0.2_6.npz",
+            "film_split_0.6_0.2_7.npz",
+            "film_split_0.6_0.2_8.npz",
+            "film_split_0.6_0.2_9.npz",
+        ]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        for f in self.raw_file_names[:2]:
+            fs.cp(f"{self.url}/new_data/film/{f}", self.raw_dir)
+        for f in self.raw_file_names[2:]:
+            fs.cp(f"{self.url}/splits/{f}", self.raw_dir)
+
+    def process(self):
+        with open(self.raw_paths[0], "r") as f:
+            lines = f.read().split("\n")[1:-1]
+        xs = [[float(v) for v in r.split("\t")[1].split(",")] for r in lines]
+        ys = [int(r.split("\t")[2]) for r in lines]
+
+        x = np.array(xs, dtype=np.float32)
+        y = np.array(ys, dtype=np.int64)
+
+        with open(self.raw_paths[1], "r") as f:
+            lines = f.read().split("\n")[1:-1]
+        edges = [[int(v) for v in r.split("\t")] for r in lines]
+        edge_index = np.array(edges, dtype=np.int64).T
+        edge_index_t, _ = coalesce(edge_index, num_nodes=x.shape[0])
+        edge_index = ops.convert_to_numpy(edge_index_t)
+
+        train_masks, val_masks, test_masks = [], [], []
+        for filepath in self.raw_paths[2:]:
+            masks = np.load(filepath)
+            train_masks.append(masks["train_mask"])
+            val_masks.append(masks["val_mask"])
+            test_masks.append(masks["test_mask"])
+
+        train_mask = np.stack(train_masks, axis=1)
+        val_mask = np.stack(val_masks, axis=1)
+        test_mask = np.stack(test_masks, axis=1)
+
+        data = Data(
+            x=ops.convert_to_tensor(x, dtype="float32"),
+            edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+            y=ops.convert_to_tensor(y, dtype="int64"),
+            train_mask=ops.convert_to_tensor(train_mask, dtype="bool"),
+            val_mask=ops.convert_to_tensor(val_mask, dtype="bool"),
+            test_mask=ops.convert_to_tensor(test_mask, dtype="bool"),
+        )
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
+

k3_node/datasets/airports.py

@@ -0,0 +1,84 @@
+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
+from k3_node.utils.graph import coalesce
+
+
+class Airports(InMemoryDataset):
+    r"""The Airports dataset from the "struc2vec: Learning Node Representations from Structural Identity" paper."""
+
+    edge_url = "https://github.com/leoribeiro/struc2vec/raw/master/graph/{}-airports.edgelist"
+    label_url = "https://github.com/leoribeiro/struc2vec/raw/master/graph/labels-{}-airports.txt"
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        self.name = name.lower()
+        assert self.name in ["usa", "brazil", "europe"]
+        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 raw_file_names(self) -> List[str]:
+        return [f"{self.name}-airports.edgelist", f"labels-{self.name}-airports.txt"]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        fs.cp(self.edge_url.format(self.name), self.raw_dir)
+        fs.cp(self.label_url.format(self.name), self.raw_dir)
+
+    def process(self):
+        index_map, ys = {}, []
+        with open(self.raw_paths[1]) as f:
+            rows = f.read().split("\n")[1:-1]
+            for i, row in enumerate(rows):
+                idx, label = row.split()
+                index_map[int(idx)] = i
+                ys.append(int(label))
+        y = np.array(ys, dtype=np.int64)
+        x = np.eye(len(ys), dtype=np.float32)
+
+        edge_indices = []
+        with open(self.raw_paths[0]) as f:
+            rows = f.read().split("\n")[:-1]
+            for row in rows:
+                src, dst = row.split()
+                edge_indices.append([index_map[int(src)], index_map[int(dst)]])
+        edge_index = np.array(edge_indices, dtype=np.int64).T
+        edge_index_t, _ = coalesce(edge_index, num_nodes=len(ys))
+        edge_index = ops.convert_to_numpy(edge_index_t)
+
+        data = Data(
+            x=ops.convert_to_tensor(x, dtype="float32"),
+            edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+            y=ops.convert_to_tensor(y, 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.name.capitalize()}Airports()"
+

k3_node/datasets/amazon.py

@@ -0,0 +1,66 @@
+import os.path as osp
+from typing import Callable, Optional
+
+from k3_node.data import InMemoryDataset
+from k3_node.io import fs, read_npz
+
+
+class Amazon(InMemoryDataset):
+    r"""The Amazon Computers and Amazon Photo networks from the
+    "Pitfalls of Graph Neural Network Evaluation" paper.
+    Nodes represent goods and edges represent that two goods are frequently
+    bought together.
+    Given product reviews as bag-of-words node features, the task is to
+    map goods to their respective product category.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset ("Computers", "Photo").
+        transform (callable, optional): Transform function.
+        pre_transform (callable, optional): Pre-transform function.
+        force_reload (bool, optional): Whether to re-process the dataset.
+    """
+
+    url = "https://github.com/shchur/gnn-benchmark/raw/master/data/npz/"
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        self.name = name.lower()
+        assert self.name in ["computers", "photo"]
+        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.capitalize(), "raw")
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root, self.name.capitalize(), "processed")
+
+    @property
+    def raw_file_names(self) -> str:
+        return f"amazon_electronics_{self.name.lower()}.npz"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        fs.cp(f"{self.url}{self.raw_file_names}", self.raw_dir)
+
+    def process(self):
+        data = read_npz(self.raw_paths[0], to_undirected=True)
+        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__}{self.name.capitalize()}()"
+

k3_node/datasets/ba2motif_dataset.py

@@ -0,0 +1,73 @@
+import pickle
+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 BA2MotifDataset(InMemoryDataset):
+    r"""The synthetic BA-2motifs graph classification dataset for evaluating
+    explainability algorithms, containing 1000 random Barabasi-Albert graphs.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): Transform function.
+        pre_transform (callable, optional): Pre-transform function.
+        force_reload (bool, optional): Whether to re-process the dataset.
+    """
+
+    url = "https://github.com/flyingdoog/PGExplainer/raw/master/dataset"
+    filename = "BA-2motif.pkl"
+
+    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 self.filename
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        fs.cp(f"{self.url}/{self.filename}", self.raw_dir)
+
+    def process(self):
+        with open(self.raw_paths[0], "rb") as f:
+            adj, x, y = pickle.load(f)
+
+        adj_np = np.array(adj)
+        x_np = np.array(x, dtype=np.float32)
+        y_np = np.array(y)
+
+        data_list: List[Data] = []
+        for i in range(x_np.shape[0]):
+            nz = np.nonzero(adj_np[i])
+            edge_index = np.stack([nz[0], nz[1]], axis=0).astype(np.int64)
+            xi = x_np[i]
+            y_nz = np.nonzero(y_np[i])[0]
+            yi = int(y_nz[0]) if len(y_nz) > 0 else 0
+
+            data = Data(
+                x=ops.convert_to_tensor(xi, dtype="float32"),
+                edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+                y=ops.convert_to_tensor(np.array([yi], dtype=np.int64), dtype="int64"),
+            )
+
+            if self.pre_transform is not None:
+                data = self.pre_transform(data)
+
+            data_list.append(data)
+
+        self.save(data_list, self.processed_paths[0])
+

k3_node/datasets/ba_shapes.py

@@ -0,0 +1,81 @@
+from typing import Callable, Optional, Tuple
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.utils.random import barabasi_albert_graph
+
+
+def house() -> Tuple[np.ndarray, np.ndarray]:
+    edge_index = np.array(
+        [[0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 4, 4], [1, 3, 4, 4, 2, 0, 1, 3, 2, 0, 0, 1]],
+        dtype=np.int64,
+    )
+    label = np.array([1, 1, 2, 2, 3], dtype=np.int64)
+    return edge_index, label
+
+
+class BAShapes(InMemoryDataset):
+    r"""The BA-Shapes dataset from the "GNNExplainer: Generating Explanations
+    for Graph Neural Networks" paper, containing a Barabasi-Albert (BA) graph
+    with 300 nodes and a set of 80 "house"-structured graphs connected to it.
+
+    Args:
+        connection_distribution (str, optional): Specifies how the houses and
+            the BA graph get connected ("random", "uniform"). (default: "random")
+        transform (callable, optional): Transform function.
+    """
+
+    def __init__(
+        self,
+        connection_distribution: str = "random",
+        transform: Optional[Callable] = None,
+    ):
+        super().__init__(None, transform)
+        assert connection_distribution in ["random", "uniform"]
+
+        num_nodes = 300
+        edge_index = ops.convert_to_numpy(barabasi_albert_graph(num_nodes, num_edges=5))
+        edge_label = np.zeros(edge_index.shape[1], dtype=np.int64)
+        node_label = np.zeros(num_nodes, dtype=np.int64)
+
+        num_houses = 80
+        if connection_distribution == "random":
+            connecting_nodes = np.random.permutation(num_nodes)[:num_houses]
+        else:
+            step = num_nodes // num_houses
+            connecting_nodes = np.arange(0, num_nodes, step)
+
+        edge_indices = [edge_index]
+        edge_labels = [edge_label]
+        node_labels = [node_label]
+
+        for i in range(num_houses):
+            house_edge_index, house_label = house()
+            edge_indices.append(house_edge_index + num_nodes)
+            edge_indices.append(
+                np.array([[int(connecting_nodes[i]), num_nodes], [num_nodes, int(connecting_nodes[i])]], dtype=np.int64)
+            )
+            edge_labels.append(np.ones(house_edge_index.shape[1], dtype=np.int64))
+            edge_labels.append(np.zeros(2, dtype=np.int64))
+            node_labels.append(house_label)
+            num_nodes += 5
+
+        edge_index = np.concatenate(edge_indices, axis=1)
+        edge_label = np.concatenate(edge_labels, axis=0)
+        node_label = np.concatenate(node_labels, axis=0)
+
+        x = np.ones((num_nodes, 10), dtype=np.float32)
+        expl_mask = np.zeros(num_nodes, dtype=bool)
+        expl_mask[np.arange(400, num_nodes, 5)] = True
+
+        data = Data(
+            x=ops.convert_to_tensor(x, dtype="float32"),
+            edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+            y=ops.convert_to_tensor(node_label, dtype="int64"),
+            expl_mask=ops.convert_to_tensor(expl_mask, dtype="bool"),
+            edge_label=ops.convert_to_tensor(edge_label, dtype="int64"),
+        )
+
+        self.data, self.slices = self.collate([data])
+

k3_node/datasets/bitcoin_otc.py

@@ -0,0 +1,77 @@
+import datetime
+import os
+from typing import Callable, Optional
+
+import numpy as np
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.data.download import download_url
+from k3_node.data.extract import extract_gz
+
+
+class BitcoinOTC(InMemoryDataset):
+    r"""The Bitcoin-OTC dataset from the `"EvolveGCN: Evolving Graph Convolutional Networks for
+    Dynamic Graphs" <https://arxiv.org/abs/1902.10191>`_ paper: 138 who-trusts-whom networks of
+    consecutive time steps among 6,005 users. ``edge_attr`` holds the trust rating of every edge,
+    from -10 (total distrust) to +10 (total trust).
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        edge_window_size (int, optional): The number of time steps an edge stays in the graph
+            sequence after it was created. (default: ``10``)
+        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.
+        force_reload (bool, optional): Whether to re-process the dataset. (default: ``False``)
+    """
+
+    url = "https://snap.stanford.edu/data/soc-sign-bitcoinotc.csv.gz"
+
+    def __init__(self, root: str, edge_window_size: int = 10, transform: Optional[Callable] = None,
+                 pre_transform: Optional[Callable] = None, force_reload: bool = False):
+        self.edge_window_size = edge_window_size
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> str:
+        return "soc-sign-bitcoinotc.csv"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    @property
+    def num_nodes(self) -> int:
+        return int(np.max(np.asarray(self._data.edge_index))) + 1
+
+    def download(self):
+        path = download_url(self.url, self.raw_dir)
+        extract_gz(path, self.raw_dir)
+        os.unlink(path)
+
+    def process(self):
+        with open(self.raw_paths[0]) as f:
+            lines = [line.split(",") for line in f.read().split("\n")[:-1]]
+        edge_index = np.array([[int(line[0]), int(line[1])] for line in lines], dtype=np.int64)
+        edge_index = (edge_index - edge_index.min()).T
+        num_nodes = int(edge_index.max()) + 1
+        edge_attr = np.array([int(line[2]) for line in lines], dtype=np.int64)
+        stamps = [datetime.datetime.fromtimestamp(int(float(line[3]))) for line in lines]
+
+        offset = datetime.timedelta(days=13.8)  # results in 138 time steps
+        graph_indices, factor = [], 1
+        for t in stamps:
+            factor = factor if t < stamps[0] + factor * offset else factor + 1
+            graph_indices.append(factor - 1)
+        graph_idx = np.array(graph_indices, dtype=np.int64)
+
+        data_list = []
+        for i in range(int(graph_idx.max()) + 1):
+            mask = (graph_idx > (i - self.edge_window_size)) & (graph_idx <= i)
+            data_list.append(Data(edge_index=edge_index[:, mask], edge_attr=edge_attr[mask], num_nodes=num_nodes))
+
+        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]
+        self.save(data_list, self.processed_paths[0])

k3_node/datasets/citation_full.py

@@ -0,0 +1,81 @@
+import os.path as osp
+from typing import Callable, Optional
+
+from k3_node.data import InMemoryDataset
+from k3_node.io import fs, read_npz
+
+
+class CitationFull(InMemoryDataset):
+    r"""The full citation network datasets from the
+    "Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via
+    Ranking" paper.
+    Nodes represent documents and edges represent citation links.
+    Datasets include "Cora", "Cora_ML", "CiteSeer", "DBLP", "PubMed".
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset ("Cora", "Cora_ML", "CiteSeer", "DBLP", "PubMed").
+        transform (callable, optional): Transform function.
+        pre_transform (callable, optional): Pre-transform function.
+        to_undirected (bool, optional): Whether the original graph is converted to undirected.
+        force_reload (bool, optional): Whether to re-process the dataset.
+    """
+
+    url = "https://github.com/abojchevski/graph2gauss/raw/master/data/{}.npz"
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        to_undirected: bool = True,
+        force_reload: bool = False,
+    ):
+        self.name = name.lower()
+        self.to_undirected = to_undirected
+        assert self.name in ["cora", "cora_ml", "citeseer", "dblp", "pubmed"]
+        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 raw_file_names(self) -> str:
+        return f"{self.name}.npz"
+
+    @property
+    def processed_file_names(self) -> str:
+        suffix = "undirected" if self.to_undirected else "directed"
+        return f"data_{suffix}.pt"
+
+    def download(self):
+        fs.cp(self.url.format(self.name), self.raw_dir)
+
+    def process(self):
+        data = read_npz(self.raw_paths[0], to_undirected=self.to_undirected)
+        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.name.capitalize()}Full()"
+
+
+class CoraFull(CitationFull):
+    r"""Alias for CitationFull with name="Cora"."""
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+    ):
+        super().__init__(root, "cora", transform, pre_transform)
+

k3_node/datasets/coauthor.py

@@ -0,0 +1,66 @@
+import os.path as osp
+from typing import Callable, Optional
+
+from k3_node.data import InMemoryDataset
+from k3_node.io import fs, read_npz
+
+
+class Coauthor(InMemoryDataset):
+    r"""The Coauthor CS and Coauthor Physics networks from the
+    "Pitfalls of Graph Neural Network Evaluation" paper.
+    Nodes represent authors that are connected by an edge if they co-authored a
+    paper.
+    Given paper keywords for each author's papers, the task is to map authors
+    to their respective field of study.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset ("CS", "Physics").
+        transform (callable, optional): Transform function.
+        pre_transform (callable, optional): Pre-transform function.
+        force_reload (bool, optional): Whether to re-process the dataset.
+    """
+
+    url = "https://github.com/shchur/gnn-benchmark/raw/master/data/npz/"
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        assert name.lower() in ["cs", "physics"]
+        self.name = "CS" if name.lower() == "cs" else "Physics"
+        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 raw_file_names(self) -> str:
+        return f"ms_academic_{self.name[:3].lower()}.npz"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        fs.cp(f"{self.url}{self.raw_file_names}", self.raw_dir)
+
+    def process(self):
+        data = read_npz(self.raw_paths[0], to_undirected=True)
+        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__}{self.name}()"
+