k3-node 1.0.0__py3-none-any.whl

k3_node/datasets/reddit.py

@@ -0,0 +1,121 @@
+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,
+    download_url,
+    extract_zip,
+)
+from k3_node.utils import coalesce
+
+
+class Reddit(InMemoryDataset):
+    r"""The Reddit dataset from the `"Inductive Representation Learning on
+    Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper, containing
+    Reddit posts belonging to different communities.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`k3_node.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:`k3_node.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in an
+            :obj:`k3_node.data.Data` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (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
+        * - 232,965
+          - 114,615,892
+          - 602
+          - 41
+    """
+
+    url = "https://data.dgl.ai/dataset/reddit.zip"
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+        super().__init__(
+            root,
+            transform,
+            pre_transform,
+            pre_filter,
+            force_reload=force_reload,
+        )
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ["reddit_data.npz", "reddit_graph.npz"]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    @property
+    def num_classes(self) -> int:
+        return 41
+
+    def download(self) -> None:
+        path = download_url(self.url, self.raw_dir)
+        extract_zip(path, self.raw_dir)
+        if osp.exists(path):
+            os.unlink(path)
+
+    def process(self) -> None:
+        import scipy.sparse as sp
+
+        data = np.load(osp.join(self.raw_dir, "reddit_data.npz"))
+        x = ops.convert_to_tensor(data["feature"], dtype="float32")
+        y = ops.convert_to_tensor(data["label"], dtype="int64")
+        split = data["node_types"]
+
+        adj = sp.load_npz(osp.join(self.raw_dir, "reddit_graph.npz"))
+        if not hasattr(adj, "row"):
+            adj = adj.tocoo()
+        row = adj.row.astype(np.int64)
+        col = adj.col.astype(np.int64)
+        edge_index = np.stack([row, col], axis=0)
+        edge_index = ops.convert_to_tensor(edge_index, dtype="int64")
+        edge_index, _ = coalesce(edge_index, num_nodes=int(data["feature"].shape[0]))
+
+        data = Data(x=x, edge_index=edge_index, y=y)
+        data.train_mask = ops.convert_to_tensor(split == 1, dtype="bool")
+        data.val_mask = ops.convert_to_tensor(split == 2, dtype="bool")
+        data.test_mask = ops.convert_to_tensor(split == 3, dtype="bool")
+
+        if self.pre_filter is not None and not self.pre_filter(data):
+            return
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
+

k3_node/datasets/sbm_dataset.py

@@ -0,0 +1,165 @@
+import os.path as osp
+from typing import Any, Callable, List, Optional, Union
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+from k3_node.utils.random import stochastic_blockmodel_graph
+
+
+class StochasticBlockModelDataset(InMemoryDataset):
+    r"""A synthetic graph dataset generated by the stochastic block model.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        block_sizes ([int] or array): The sizes of blocks.
+        edge_probs ([[float]] or array): The density of edges between blocks.
+        num_graphs (int, optional): The number of graphs. (default: 1)
+        num_channels (int, optional): The number of node features. (default: None)
+        is_undirected (bool, optional): Whether the graph is undirected. (default: True)
+        transform (callable, optional): Transform function.
+        pre_transform (callable, optional): Pre-transform function.
+        force_reload (bool, optional): Whether to re-process the dataset.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        block_sizes: Union[List[int], np.ndarray],
+        edge_probs: Union[List[List[float]], np.ndarray],
+        num_graphs: int = 1,
+        num_channels: Optional[int] = None,
+        is_undirected: bool = True,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+        **kwargs: Any,
+    ):
+        self.block_sizes = np.array(block_sizes, dtype=np.int64)
+        self.edge_probs = np.array(edge_probs, dtype=np.float32)
+        assert num_graphs > 0
+
+        self.num_graphs = num_graphs
+        self.num_channels = num_channels
+        self.is_undirected = is_undirected
+
+        self.kwargs = {
+            "n_informative": num_channels,
+            "n_redundant": 0,
+            "flip_y": 0.0,
+            "shuffle": False,
+        }
+        self.kwargs.update(kwargs)
+
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root, self.__class__.__name__, "processed")
+
+    @property
+    def processed_file_names(self) -> str:
+        bs = self.block_sizes.flatten().tolist()
+        hash1 = "-".join([f"{x:.1f}" for x in bs])
+        ep = self.edge_probs.flatten().tolist()
+        hash2 = "-".join([f"{x:.1f}" for x in ep])
+        return f"data_{self.num_channels}_{hash1}_{hash2}_{self.num_graphs}.pt"
+
+    def process(self):
+        try:
+            from sklearn.datasets import make_classification
+        except ImportError:
+            make_classification = None
+
+        edge_index = stochastic_blockmodel_graph(
+            self.block_sizes, self.edge_probs, directed=not self.is_undirected
+        )
+
+        num_samples = int(self.block_sizes.sum())
+        num_classes = len(self.block_sizes)
+
+        data_list = []
+        for _ in range(self.num_graphs):
+            x = None
+            if self.num_channels is not None:
+                if make_classification is not None:
+                    x_raw, y_not_sorted = make_classification(
+                        n_samples=num_samples,
+                        n_features=self.num_channels,
+                        n_classes=num_classes,
+                        weights=self.block_sizes / num_samples,
+                        **self.kwargs,
+                    )
+                    x_raw = x_raw[np.argsort(y_not_sorted)]
+                    x = ops.convert_to_tensor(x_raw.astype(np.float32), dtype="float32")
+                else:
+                    x = ops.convert_to_tensor(
+                        np.random.randn(num_samples, self.num_channels).astype(np.float32),
+                        dtype="float32",
+                    )
+
+            y_np = np.repeat(np.arange(num_classes, dtype=np.int64), self.block_sizes)
+            y = ops.convert_to_tensor(y_np, dtype="int64")
+
+            data = Data(x=x, edge_index=edge_index, y=y)
+            if self.pre_transform is not None:
+                data = self.pre_transform(data)
+            data_list.append(data)
+
+        self.save(data_list, self.processed_paths[0])
+
+
+class RandomPartitionGraphDataset(StochasticBlockModelDataset):
+    r"""The random partition graph dataset from the "How to Find Your Friendly
+    Neighborhood: Graph Attention Design with Self-Supervision" paper.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        num_classes: int,
+        num_nodes_per_class: int,
+        node_homophily_ratio: float,
+        average_degree: float,
+        num_graphs: int = 1,
+        num_channels: Optional[int] = None,
+        is_undirected: bool = True,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        **kwargs: Any,
+    ):
+        self._num_classes = num_classes
+        self.num_nodes_per_class = num_nodes_per_class
+        self.node_homophily_ratio = node_homophily_ratio
+        self.average_degree = average_degree
+
+        ec_over_v2 = average_degree / num_nodes_per_class
+        p_in = node_homophily_ratio * ec_over_v2
+        p_out = (ec_over_v2 - p_in) / (num_classes - 1)
+
+        block_sizes = [num_nodes_per_class for _ in range(num_classes)]
+        edge_probs = [[p_out for _ in range(num_classes)] for _ in range(num_classes)]
+        for r in range(num_classes):
+            edge_probs[r][r] = p_in
+
+        super().__init__(
+            root,
+            block_sizes,
+            edge_probs,
+            num_graphs=num_graphs,
+            num_channels=num_channels,
+            is_undirected=is_undirected,
+            transform=transform,
+            pre_transform=pre_transform,
+            **kwargs,
+        )
+
+    @property
+    def processed_file_names(self) -> str:
+        return (
+            f"data_{self.num_channels}_{self._num_classes}_"
+            f"{self.num_nodes_per_class}_{self.node_homophily_ratio:.1f}_"
+            f"{self.average_degree:.1f}_{self.num_graphs}.pt"
+        )
+

k3_node/datasets/seal.py

@@ -0,0 +1,74 @@
+from typing import Optional
+
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+
+
+class SEALDataset(InMemoryDataset):
+    r"""Enclosing subgraphs for link prediction, as in SEAL
+    (`"Link Prediction Based on Graph Neural Networks" <https://arxiv.org/abs/1802.09691>`_).
+
+    Every link to classify becomes a small graph: the ``num_hops``-hop neighborhood of its two
+    nodes, without the link itself. Nodes are described only by their double-radius labels
+    (their distances to the two nodes, one-hot encoded as ``x``); ``y`` is 1 for a true link and
+    0 for a non-edge. This turns link prediction into graph classification.
+
+    Args:
+        data (Data): One split from :class:`~k3_node.transforms.RandomLinkSplit`: message passing
+            edges in ``edge_index`` and the links in ``edge_label_index`` / ``edge_label`` (or
+            ``pos_edge_label_index`` / ``neg_edge_label_index``).
+        num_hops (int): Size of the neighborhood around each link. (default: ``2``)
+        num_labels (int, optional): Size of the one-hot node labels. Use the training set's
+            ``num_labels`` for the validation and test sets. (default: the largest label + 1)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.data import Data
+        from k3_node.datasets import SEALDataset
+        from k3_node.transforms import RandomLinkSplit
+
+        data = Data(edge_index=np.random.randint(0, 30, size=(2, 120)), num_nodes=30)
+        train_data, val_data, test_data = RandomLinkSplit(num_val=0.1, num_test=0.1)(data)
+        train_dataset = SEALDataset(train_data, num_hops=2)
+        print(len(train_dataset) == train_data.edge_label_index.shape[1])  # True: one graph per link
+        ```
+    """
+
+    def __init__(self, data, num_hops: int = 2, num_labels: Optional[int] = None):
+        super().__init__(None)
+        from k3_node.utils.graph import drnl_node_labeling, k_hop_subgraph
+
+        links, labels = self._links(data)
+        edge_index = np.asarray(ops.convert_to_numpy(data.edge_index)).astype(np.int64)
+        graphs = []
+        for (src, dst), y in zip(links.T, labels):
+            nodes, sub_edge_index, mapping, _ = k_hop_subgraph(
+                [src, dst], num_hops, edge_index, relabel_nodes=True, num_nodes=data.num_nodes)
+            s, d = (int(m) for m in mapping)
+            keep = ~(((sub_edge_index[0] == s) & (sub_edge_index[1] == d))
+                     | ((sub_edge_index[0] == d) & (sub_edge_index[1] == s)))
+            sub_edge_index = sub_edge_index[:, keep]  # hide the link to predict
+            z = drnl_node_labeling(sub_edge_index, s, d, num_nodes=len(nodes))
+            graphs.append((z, sub_edge_index, y))
+
+        self.num_labels = num_labels or max(int(z.max()) for z, _, _ in graphs) + 1
+        self.data, self.slices = self.collate([
+            Data(x=np.eye(self.num_labels, dtype=np.float32)[np.minimum(z, self.num_labels - 1)],
+                 edge_index=e, y=np.array([y], dtype=np.float32))
+            for z, e, y in graphs
+        ])
+
+    @staticmethod
+    def _links(data):
+        def arr(x):
+            return np.asarray(ops.convert_to_numpy(x))
+
+        if getattr(data, "edge_label_index", None) is not None:
+            return arr(data.edge_label_index).astype(np.int64), arr(data.edge_label)
+        pos = arr(data.pos_edge_label_index).astype(np.int64)
+        neg = getattr(data, "neg_edge_label_index", None)
+        neg = np.zeros((2, 0), np.int64) if neg is None else arr(neg).astype(np.int64)
+        return np.concatenate([pos, neg], axis=1), np.concatenate([np.ones(pos.shape[1]), np.zeros(neg.shape[1])])

k3_node/datasets/shape_scenes.py

@@ -0,0 +1,92 @@
+from typing import Callable, Optional
+
+import numpy as np
+
+from k3_node.data import Data, InMemoryDataset
+
+
+class ShapeScenes(InMemoryDataset):
+    r"""Small point cloud scenes for semantic segmentation, built from :class:`GeometricShapes`
+    meshes (a compact stand-in for ShapeNet part segmentation).
+
+    Every scene contains ``shapes_per_scene`` objects, each a cube, sphere, cone or torus that is
+    randomly scaled, rotated and placed apart from the others. ``num_points`` points are sampled on
+    their surfaces: ``pos`` holds the positions (scaled to the unit ball), ``x`` the surface
+    normals, and ``y`` the type of the object each point lies on (4 classes). Training scenes use
+    the training meshes, test scenes the test meshes; the scenes are generated with a fixed seed.
+
+    Args:
+        root (str): Directory where GeometricShapes is (or will be) downloaded.
+        train (bool, optional): Build training scenes if ``True``, else test scenes.
+        num_scenes (int, optional): Number of scenes. (default: ``200`` for training, ``50`` for test)
+        shapes_per_scene (int, optional): Objects per scene. (default: ``3``)
+        num_points (int, optional): Points per scene. (default: ``1024``)
+        transform (callable, optional): A function applied to each scene when it is accessed.
+        seed (int, optional): Random seed. (default: ``0``)
+    """
+
+    categories = ["3d_cube", "3d_sphere", "3d_cone", "3d_torus"]
+
+    def __init__(self, root: str, train: bool = True, num_scenes: Optional[int] = None, shapes_per_scene: int = 3,
+                 num_points: int = 1024, transform: Optional[Callable] = None, seed: int = 0):
+        super().__init__(None, transform)
+        from k3_node.datasets.geometric_shapes import GeometricShapes
+        from k3_node.transforms.spatial import _np
+
+        shapes = GeometricShapes(root, train=train)
+        names = sorted(__import__("os").listdir(shapes.raw_dir))
+        meshes = {}
+        for i in range(len(shapes)):
+            data = shapes[i]
+            name = names[int(_np(data.y)[0])]
+            if name in self.categories:
+                meshes[self.categories.index(name)] = (_np(data.pos).astype(np.float64), _np(data.face))
+
+        rng = np.random.default_rng(seed + (0 if train else 1))
+        num_scenes = num_scenes or (200 if train else 50)
+        self.data, self.slices = self.collate(
+            [self._scene(rng, meshes, shapes_per_scene, num_points) for _ in range(num_scenes)])
+
+    @staticmethod
+    def _sample(rng, pos, face, num):
+        a, b, c = pos[face[0]], pos[face[1]], pos[face[2]]
+        cross = np.cross(b - a, c - a)
+        area = np.linalg.norm(cross, axis=1)
+        tri = rng.choice(len(area), size=num, p=area / area.sum())
+        u, v = rng.random((2, num))
+        flip = u + v > 1
+        u[flip], v[flip] = 1 - u[flip], 1 - v[flip]
+        points = a[tri] + u[:, None] * (b - a)[tri] + v[:, None] * (c - a)[tri]
+        normals = cross[tri] / np.maximum(area[tri, None], 1e-12)
+        return points, normals
+
+    @staticmethod
+    def _rotation(rng):
+        q = rng.normal(size=4)
+        q /= np.linalg.norm(q)
+        w, x, y, z = q
+        return np.array([[1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
+                         [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
+                         [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]])
+
+    def _scene(self, rng, meshes, shapes_per_scene, num_points):
+        counts = np.full(shapes_per_scene, num_points // shapes_per_scene)
+        counts[: num_points - counts.sum()] += 1
+        angle = rng.random() * 2 * np.pi
+        positions, normals, labels = [], [], []
+        for i in range(shapes_per_scene):
+            label = int(rng.integers(len(self.categories)))
+            pos, face = meshes[label]
+            points, normal = self._sample(rng, pos, face, counts[i])
+            points = points / np.abs(points).max()
+            rot = self._rotation(rng)
+            theta = angle + 2 * np.pi * i / shapes_per_scene  # objects around a circle, apart
+            offset = 2.5 * np.array([np.cos(theta), np.sin(theta), 0.0])
+            positions.append(points @ rot.T * rng.uniform(0.6, 1.0) + offset)
+            normals.append(normal @ rot.T)
+            labels.append(np.full(counts[i], label))
+        pos = np.concatenate(positions)
+        pos = pos - pos.mean(axis=0)
+        pos = pos / np.linalg.norm(pos, axis=1).max()
+        return Data(pos=pos.astype(np.float32), x=np.concatenate(normals).astype(np.float32),
+                    y=np.concatenate(labels).astype(np.int64))