k3-node 1.0.0__py3-none-any.whl

k3_node/datasets/imdb.py

@@ -0,0 +1,96 @@
+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 IMDB(InMemoryDataset):
+    r"""A subset of the Internet Movie Database (IMDB) containing three types of entities:
+    movies, actors, and directors.
+    """
+
+    url = "https://www.dropbox.com/s/g0btk9ctr1es39x/IMDB_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.npz",
+            "labels.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, "IMDB_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 = ["movie", "director", "actor"]
+
+        for i, node_type in enumerate(node_types):
+            x = sp.load_npz(osp.join(self.raw_dir, f"features_{i}.npz"))
+            data[node_type].x = ops.convert_to_tensor(
+                np.array(x.todense(), dtype=np.float32), dtype="float32"
+            )
+
+        y = np.load(osp.join(self.raw_dir, "labels.npy"))
+        data["movie"].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["movie"].num_nodes, dtype=bool)
+            mask[idx] = True
+            data["movie"][f"{name}_mask"] = ops.convert_to_tensor(mask, dtype="bool")
+
+        s = {}
+        N_m = data["movie"].num_nodes
+        N_d = data["director"].num_nodes
+        N_a = data["actor"].num_nodes
+        s["movie"] = (0, N_m)
+        s["director"] = (N_m, N_m + N_d)
+        s["actor"] = (N_m + N_d, N_m + N_d + N_a)
+
+        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/jodie.py

@@ -0,0 +1,56 @@
+import os
+import os.path as osp
+
+import numpy as np
+
+
+class JODIEDataset:
+    r"""The temporal interaction datasets of `"JODIE: Predicting Dynamic Embedding Trajectory in
+    Temporal Interaction Networks" <https://cs.stanford.edu/~srijan/pubs/jodie-kdd2019.pdf>`_:
+    ``"wikipedia"``, ``"reddit"``, ``"mooc"`` and ``"lastfm"``.
+
+    ``dataset[0]`` is a :class:`~k3_node.data.TemporalData` event stream: user ``src`` interacts
+    with item ``dst`` (numbered after the users) at time ``t``, with features ``msg`` and label
+    ``y``. MOOC (7,144 nodes, 411,749 events, 4 features) is the smallest download (40 MB).
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset.
+    """
+
+    url = "https://snap.stanford.edu/jodie/{}.csv"
+    names = ["wikipedia", "reddit", "mooc", "lastfm"]
+
+    def __init__(self, root: str, name: str):
+        self.name = name.lower()
+        assert self.name in self.names
+        folder = osp.join(root, self.name)
+        cache = osp.join(folder, "processed.npz")
+        if not osp.exists(cache):
+            os.makedirs(folder, exist_ok=True)
+            csv = osp.join(folder, f"{self.name}.csv")
+            if not osp.exists(csv):
+                from k3_node.data.download import download_url
+
+                download_url(self.url.format(self.name), folder)
+            import pandas as pd
+
+            df = pd.read_csv(csv, skiprows=1, header=None)
+            src = df.iloc[:, 0].values.astype(np.int64)
+            dst = df.iloc[:, 1].values.astype(np.int64) + int(src.max()) + 1
+            np.savez(cache, src=src, dst=dst, t=df.iloc[:, 2].values.astype(np.int64),
+                     y=df.iloc[:, 3].values.astype(np.int64), msg=df.iloc[:, 4:].values.astype(np.float32))
+        self._arrays = dict(np.load(cache))
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, idx):
+        from k3_node.data import TemporalData
+
+        if idx != 0:
+            raise IndexError(idx)
+        return TemporalData(**self._arrays)
+
+    def __repr__(self):
+        return f"JODIEDataset({self.name})"

k3_node/datasets/karate.py

@@ -0,0 +1,56 @@
+from typing import Callable, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data, InMemoryDataset
+
+
+class KarateClub(InMemoryDataset):
+    r"""Zachary's karate club network from the `"An Information Flow Model for
+    Conflict and Fission in Small Groups" paper, containing 34 nodes and 156
+    undirected edges labeled into 4 community classes.
+    """
+
+    def __init__(self, transform: Optional[Callable] = None):
+        super().__init__(None, transform)
+
+        row = [
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1,
+            1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 4, 4, 4,
+            5, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8, 8, 8, 9, 9, 10, 10,
+            10, 11, 12, 12, 13, 13, 13, 13, 13, 14, 14, 15, 15, 16, 16, 17, 17,
+            18, 18, 19, 19, 19, 20, 20, 21, 21, 22, 22, 23, 23, 23, 23, 23, 24,
+            24, 24, 25, 25, 25, 26, 26, 27, 27, 27, 27, 28, 28, 28, 29, 29, 29,
+            29, 30, 30, 30, 30, 31, 31, 31, 31, 31, 31, 32, 32, 32, 32, 32, 32,
+            32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33,
+            33, 33, 33, 33, 33, 33
+        ]
+        col = [
+            1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 13, 17, 19, 21, 31, 0, 2, 3, 7,
+            13, 17, 19, 21, 30, 0, 1, 3, 7, 8, 9, 13, 27, 28, 32, 0, 1, 2, 7,
+            12, 13, 0, 6, 10, 0, 6, 10, 16, 0, 4, 5, 16, 0, 1, 2, 3, 0, 2, 30,
+            32, 33, 2, 33, 0, 4, 5, 0, 0, 3, 0, 1, 2, 3, 33, 32, 33, 32, 33, 5,
+            6, 0, 1, 32, 33, 0, 1, 33, 32, 33, 0, 1, 32, 33, 25, 27, 29, 32,
+            33, 25, 27, 31, 23, 24, 31, 29, 33, 2, 23, 24, 33, 2, 31, 33, 23,
+            26, 32, 33, 1, 8, 32, 33, 0, 24, 25, 28, 32, 33, 2, 8, 14, 15, 18,
+            20, 22, 23, 29, 30, 31, 33, 8, 9, 13, 14, 15, 18, 19, 20, 22, 23,
+            26, 27, 28, 29, 30, 31, 32
+        ]
+        edge_index = ops.convert_to_tensor(np.array([row, col], dtype=np.int64), dtype="int64")
+
+        y_np = np.array([
+            1, 1, 1, 1, 3, 3, 3, 1, 0, 1, 3, 1, 1, 1, 0, 0, 3, 1, 0, 1, 0, 1,
+            0, 0, 2, 2, 0, 0, 2, 0, 0, 2, 0, 0
+        ], dtype=np.int64)
+        y = ops.convert_to_tensor(y_np, dtype="int64")
+
+        x = ops.convert_to_tensor(np.eye(34, dtype=np.float32), dtype="float32")
+
+        train_mask_np = np.zeros(34, dtype=bool)
+        for i in range(int(np.max(y_np)) + 1):
+            train_mask_np[np.where(y_np == i)[0][0]] = True
+        train_mask = ops.convert_to_tensor(train_mask_np, dtype="bool")
+
+        data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask)
+        self.data, self.slices = self.collate([data])
+

k3_node/datasets/lastfm_asia.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 LastFMAsia(InMemoryDataset):
+    r"""The LastFM Asia Network dataset."""
+
+    url = "https://graphmining.ai/datasets/ptg/lastfm_asia.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 "lastfm_asia.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/mesh_correspondence.py

@@ -0,0 +1,50 @@
+import os
+from typing import Callable, Optional
+
+import numpy as np
+
+from k3_node.data import Data, InMemoryDataset
+
+
+class MeshCorrespondence(InMemoryDataset):
+    r"""A small shape correspondence dataset in the spirit of FAUST: randomly deformed copies of one
+    :class:`GeometricShapes` mesh (the monkey head by default). Every copy is smoothly bent and
+    slightly rotated; the task is to recognize every vertex, so ``y`` holds the vertex indices.
+
+    The graphs hold ``pos`` and mesh faces ``face``; use :class:`~k3_node.transforms.FaceToEdge`
+    to connect the vertices. As in FAUST, there are 80 training and 20 test meshes by default.
+
+    Args:
+        root (str): Directory where GeometricShapes is (or will be) downloaded.
+        train (bool, optional): Build the training meshes if ``True``, else the test meshes.
+        num_meshes (int, optional): Number of meshes. (default: ``80`` / ``20``)
+        shape (str, optional): The GeometricShapes mesh to deform. (default: ``"3d_monkey"``)
+        transform (callable, optional): A function applied to each mesh when it is accessed.
+        pre_transform (callable, optional): A function applied to each mesh once, when built.
+        seed (int, optional): Random seed. (default: ``0``)
+    """
+
+    def __init__(self, root: str, train: bool = True, num_meshes: Optional[int] = None, shape: str = "3d_monkey",
+                 transform: Optional[Callable] = None, pre_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=True)
+        names = sorted(os.listdir(shapes.raw_dir))
+        base = next(shapes[i] for i in range(len(shapes)) if names[int(_np(shapes[i].y)[0])] == shape)
+        pos, face = _np(base.pos).astype(np.float64), _np(base.face)
+        pos = pos / np.abs(pos).max()
+
+        rng = np.random.default_rng(seed + (0 if train else 1))
+        graphs = []
+        for _ in range(num_meshes or (80 if train else 20)):
+            w = rng.normal(size=(3, 3)) * 1.5
+            bent = pos + 0.12 * np.sin(pos @ w + rng.uniform(0, 2 * np.pi, size=3))  # smooth deformation
+            angle = np.deg2rad(rng.uniform(-15, 15))
+            c, s = np.cos(angle), np.sin(angle)
+            rot = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])
+            data = Data(pos=(bent @ rot.T).astype(np.float32), face=face.copy(),
+                        y=np.arange(len(pos), dtype=np.int64))
+            graphs.append(pre_transform(data) if pre_transform is not None else data)
+        self.data, self.slices = self.collate(graphs)

k3_node/datasets/molecule_net.py

@@ -0,0 +1,148 @@
+import os
+import os.path as osp
+import re
+import warnings
+from typing import Callable, Dict, Optional, Tuple, Union
+
+import numpy as np
+from keras import ops
+
+from k3_node.data import InMemoryDataset, download_url, extract_gz
+from k3_node.utils.smiles import from_smiles as default_from_smiles
+
+
+class MoleculeNet(InMemoryDataset):
+    r"""The `MoleculeNet <http://moleculenet.org/datasets-1>`_ benchmark collection
+    from the `"MoleculeNet: A Benchmark for Molecular Machine Learning"
+    <https://arxiv.org/abs/1703.00564>`_ paper, containing datasets from physical
+    chemistry, biophysics and physiology.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset (:obj:`"ESOL"`, :obj:`"FreeSolv"`,
+            :obj:`"Lipo"`, :obj:`"PCBA"`, :obj:`"MUV"`, :obj:`"HIV"`,
+            :obj:`"BACE"`, :obj:`"BBBP"`, :obj:`"Tox21"`, :obj:`"ToxCast"`,
+            :obj:`"SIDER"`, :obj:`"ClinTox"`).
+        transform (callable, optional): A function/transform that takes in a
+            :obj:`k3_node.data.Data` object and returns a transformed version.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in a
+            :obj:`k3_node.data.Data` object and returns a transformed version.
+            (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in a
+            :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`)
+        from_smiles (callable, optional): A custom function that takes a SMILES
+            string and outputs a :obj:`k3_node.data.Data` object.
+            (default: :obj:`None`)
+    """
+
+    url = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/{}"
+
+    # Format: name: (display_name, url_name, csv_name, smiles_idx, y_idx)
+    names: Dict[str, Tuple[str, str, str, int, Union[int, slice]]] = {
+        "esol": ("ESOL", "delaney-processed.csv", "delaney-processed", -1, -2),
+        "freesolv": ("FreeSolv", "SAMPL.csv", "SAMPL", 1, 2),
+        "lipo": ("Lipophilicity", "Lipophilicity.csv", "Lipophilicity", 2, 1),
+        "pcba": ("PCBA", "pcba.csv.gz", "pcba", -1, slice(0, 128)),
+        "muv": ("MUV", "muv.csv.gz", "muv", -1, slice(0, 17)),
+        "hiv": ("HIV", "HIV.csv", "HIV", 0, -1),
+        "bace": ("BACE", "bace.csv", "bace", 0, 2),
+        "bbbp": ("BBBP", "BBBP.csv", "BBBP", -1, -2),
+        "tox21": ("Tox21", "tox21.csv.gz", "tox21", -1, slice(0, 12)),
+        "toxcast": ("ToxCast", "toxcast_data.csv.gz", "toxcast_data", 0, slice(1, 618)),
+        "sider": ("SIDER", "sider.csv.gz", "sider", 0, slice(1, 28)),
+        "clintox": ("ClinTox", "clintox.csv.gz", "clintox", 0, slice(1, 3)),
+    }
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+        from_smiles: Optional[Callable] = None,
+    ) -> None:
+        self.name = name.lower()
+        if self.name not in self.names:
+            raise ValueError(
+                f"Unknown dataset name '{name}'. Available names: {list(self.names.keys())}"
+            )
+        self.from_smiles = from_smiles or default_from_smiles
+        super().__init__(
+            root,
+            transform,
+            pre_transform,
+            pre_filter,
+            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.names[self.name][2]}.csv"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self) -> None:
+        url = self.url.format(self.names[self.name][1])
+        path = download_url(url, self.raw_dir)
+        if self.names[self.name][1].endswith("gz"):
+            extract_gz(path, self.raw_dir)
+            os.unlink(path)
+
+    def process(self) -> None:
+        with open(self.raw_paths[0], "r", encoding="utf-8") as f:
+            dataset = f.read().split("\n")[1:-1]
+            dataset = [x for x in dataset if len(x) > 0]
+
+        data_list = []
+        for line in dataset:
+            line = re.sub(r'".*"', "", line)  # Replace quoted substrings
+            values = line.split(",")
+
+            smiles = values[self.names[self.name][3]]
+            labels = values[self.names[self.name][4]]
+            labels = labels if isinstance(labels, list) else [labels]
+
+            ys = [float(y) if len(y) > 0 else float("nan") for y in labels]
+            y = ops.convert_to_tensor(np.array(ys, dtype=np.float32).reshape(1, -1), dtype="float32")
+
+            data = self.from_smiles(smiles)
+            data.y = y
+
+            if data.num_nodes == 0:
+                warnings.warn(
+                    f"Skipping molecule '{smiles}' since it resulted in zero atoms",
+                    stacklevel=2,
+                )
+                continue
+
+            if self.pre_filter is not None and not self.pre_filter(data):
+                continue
+
+            if self.pre_transform is not None:
+                data = self.pre_transform(data)
+
+            data_list.append(data)
+
+        self.save(data_list, self.processed_paths[0])
+
+    def __repr__(self) -> str:
+        return f"{self.names[self.name][0]}({len(self)})"
+

k3_node/datasets/motif_generator/__init__.py

@@ -0,0 +1,7 @@
+from .base import MotifGenerator
+from .custom import CustomMotif
+from .house import HouseMotif
+from .cycle import CycleMotif
+
+__all__ = ["MotifGenerator", "CustomMotif", "HouseMotif", "CycleMotif"]
+

k3_node/datasets/motif_generator/base.py

@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from typing import Any
+from k3_node.data import Data
+
+
+class MotifGenerator(ABC):
+    r"""An abstract base class for generating motifs."""
+
+    @abstractmethod
+    def __call__(self) -> Data:
+        raise NotImplementedError
+
+    @staticmethod
+    def resolve(query: Any, *args: Any, **kwargs: Any) -> "MotifGenerator":
+        if isinstance(query, MotifGenerator):
+            return query
+        if isinstance(query, str):
+            query = query.lower()
+            if query in ["house", "housemotif"]:
+                from k3_node.datasets.motif_generator.house import HouseMotif
+                return HouseMotif(*args, **kwargs)
+            elif query in ["cycle", "cyclemotif"]:
+                from k3_node.datasets.motif_generator.cycle import CycleMotif
+                return CycleMotif(*args, **kwargs)
+        raise ValueError(f"Could not resolve motif generator: {query}")
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+

k3_node/datasets/motif_generator/custom.py

@@ -0,0 +1,17 @@
+from typing import Any, Optional
+from k3_node.data import Data
+from k3_node.datasets.motif_generator.base import MotifGenerator
+
+
+class CustomMotif(MotifGenerator):
+    r"""Generates a motif based on a custom structure coming from a Data object."""
+
+    def __init__(self, structure: Any):
+        super().__init__()
+        if not isinstance(structure, Data):
+            raise ValueError(f"Expected structure of type Data, got {type(structure)}")
+        self.structure = structure
+
+    def __call__(self) -> Data:
+        return self.structure
+

k3_node/datasets/motif_generator/cycle.py

@@ -0,0 +1,25 @@
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.datasets.motif_generator.custom import CustomMotif
+
+
+class CycleMotif(CustomMotif):
+    r"""Generates the cycle motif from the "GNNExplainer" paper."""
+
+    def __init__(self, num_nodes: int):
+        self.num_nodes = num_nodes
+
+        row = np.repeat(np.arange(num_nodes), 2)
+        col1 = np.mod(np.arange(-1, num_nodes - 1), num_nodes)
+        col2 = np.mod(np.arange(1, num_nodes + 1), num_nodes)
+        col = np.sort(np.stack([col1, col2], axis=1), axis=-1).flatten()
+
+        edge_index = ops.convert_to_tensor(np.stack([row, col], axis=0), dtype="int64")
+        structure = Data(num_nodes=num_nodes, edge_index=edge_index)
+        super().__init__(structure)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.num_nodes})"
+

k3_node/datasets/motif_generator/house.py

@@ -0,0 +1,27 @@
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.datasets.motif_generator.custom import CustomMotif
+
+
+class HouseMotif(CustomMotif):
+    r"""Generates the house-structured motif from the "GNNExplainer" paper,
+    containing 5 nodes and 6 undirected edges.
+    """
+
+    def __init__(self):
+        edge_index = ops.convert_to_tensor(
+            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,
+            ),
+            dtype="int64",
+        )
+        y = ops.convert_to_tensor(np.array([0, 0, 1, 1, 2], dtype=np.int64), dtype="int64")
+        structure = Data(num_nodes=5, edge_index=edge_index, y=y)
+        super().__init__(structure)
+

k3_node/datasets/movielens.py

@@ -0,0 +1,55 @@
+import os
+import os.path as osp
+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_zip
+
+
+class MovieLens100K(InMemoryDataset):
+    r"""The MovieLens 100K ratings (943 users, 1,682 movies, 100,000 ratings) as a user-movie graph
+    for recommendation, a small stand-in for PyG's ``AmazonBook``.
+
+    Every rating counts as an interaction. The graph is homogeneous: users are nodes
+    ``0 .. num_users - 1`` and movies ``num_users .. num_nodes - 1``. ``edge_index`` holds the
+    training interactions in both directions, ``edge_label_index`` the test interactions
+    (user, movie), from the official 80/20 split ``u1.base`` / ``u1.test``.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function applied to the graph when it is accessed.
+        force_reload (bool, optional): Whether to re-process the dataset. (default: ``False``)
+    """
+
+    url = "https://files.grouplens.org/datasets/movielens/ml-100k.zip"
+    num_users, num_items = 943, 1682
+
+    def __init__(self, root: str, transform: Optional[Callable] = None, force_reload: bool = False):
+        super().__init__(root, transform, force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self):
+        return [osp.join("ml-100k", "u1.base"), osp.join("ml-100k", "u1.test")]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        path = download_url(self.url, self.raw_dir)
+        extract_zip(path, self.raw_dir)
+        os.unlink(path)
+
+    def _read(self, path):
+        ratings = np.loadtxt(path, dtype=np.int64)[:, :2] - 1  # user id, movie id (1-based)
+        return np.stack([ratings[:, 0], ratings[:, 1] + self.num_users])
+
+    def process(self):
+        train, test = self._read(self.raw_paths[0]), self._read(self.raw_paths[1])
+        data = Data(edge_index=np.concatenate([train, train[::-1]], axis=1), edge_label_index=test,
+                    num_nodes=self.num_users + self.num_items)
+        self.save([data], self.processed_paths[0])