k3-node 1.0.0__py3-none-any.whl

k3_node/datasets/planetoid.py

@@ -0,0 +1,137 @@
+import os.path as osp
+from typing import Callable, List, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import InMemoryDataset
+from k3_node.io import fs, read_planetoid_data
+
+
+class Planetoid(InMemoryDataset):
+    r"""The citation network datasets "Cora", "CiteSeer" and "PubMed" from the
+    "Revisiting Semi-Supervised Learning with Graph Embeddings" paper.
+    Nodes represent documents and edges represent citation links.
+    Training, validation and test splits are given by binary masks.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        name (str): The name of the dataset ("Cora", "CiteSeer", "PubMed").
+        split (str, optional): The type of dataset split ("public", "full", "geom-gcn", "random").
+            (default: "public")
+        num_train_per_class (int, optional): The number of training samples per class for "random" split.
+            (default: 20)
+        num_val (int, optional): The number of validation samples for "random" split. (default: 500)
+        num_test (int, optional): The number of test samples for "random" split. (default: 1000)
+        transform (callable, optional): A function/transform that takes in a Data object and returns a transformed version.
+        pre_transform (callable, optional): A function/transform that takes in a Data object and returns a transformed version.
+        force_reload (bool, optional): Whether to re-process the dataset. (default: False)
+    """
+
+    url = "https://github.com/kimiyoung/planetoid/raw/master/data"
+    geom_gcn_url = "https://raw.githubusercontent.com/graphdml-uiuc-jlu/geom-gcn/master"
+
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        split: str = "public",
+        num_train_per_class: int = 20,
+        num_val: int = 500,
+        num_test: int = 1000,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        self.name = name
+        self.split = split.lower()
+        assert self.split in ["public", "full", "geom-gcn", "random"]
+
+        super().__init__(root, transform, pre_transform, force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+        if self.split == "full":
+            data = self.get(0)
+            val_m = ops.convert_to_numpy(data.val_mask)
+            test_m = ops.convert_to_numpy(data.test_mask)
+            train_mask = np.ones(data.num_nodes, dtype=bool)
+            train_mask[val_m | test_m] = False
+            data.train_mask = ops.convert_to_tensor(train_mask, dtype="bool")
+            self.data, self.slices = self.collate([data])
+
+        elif self.split == "random":
+            data = self.get(0)
+            num_nodes = data.num_nodes
+            y_np = ops.convert_to_numpy(data.y)
+            train_mask = np.zeros(num_nodes, dtype=bool)
+            for c in range(self.num_classes):
+                idx = np.where(y_np == c)[0]
+                perm = np.random.permutation(len(idx))
+                train_mask[idx[perm[:num_train_per_class]]] = True
+
+            remaining = np.where(~train_mask)[0]
+            remaining = remaining[np.random.permutation(len(remaining))]
+
+            val_mask = np.zeros(num_nodes, dtype=bool)
+            val_mask[remaining[:num_val]] = True
+
+            test_mask = np.zeros(num_nodes, dtype=bool)
+            test_mask[remaining[num_val : num_val + num_test]] = True
+
+            data.train_mask = ops.convert_to_tensor(train_mask, dtype="bool")
+            data.val_mask = ops.convert_to_tensor(val_mask, dtype="bool")
+            data.test_mask = ops.convert_to_tensor(test_mask, dtype="bool")
+
+            self.data, self.slices = self.collate([data])
+
+    @property
+    def raw_dir(self) -> str:
+        if self.split == "geom-gcn":
+            return osp.join(self.root, self.name, "geom-gcn", "raw")
+        return osp.join(self.root, self.name, "raw")
+
+    @property
+    def processed_dir(self) -> str:
+        if self.split == "geom-gcn":
+            return osp.join(self.root, self.name, "geom-gcn", "processed")
+        return osp.join(self.root, self.name, "processed")
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        names = ["x", "tx", "allx", "y", "ty", "ally", "graph", "test.index"]
+        return [f"ind.{self.name.lower()}.{name}" for name in names]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        for name in self.raw_file_names:
+            fs.cp(f"{self.url}/{name}", self.raw_dir)
+        if self.split == "geom-gcn":
+            for i in range(10):
+                url = f"{self.geom_gcn_url}/splits/{self.name.lower()}"
+                fs.cp(f"{url}_split_0.6_0.2_{i}.npz", self.raw_dir)
+
+    def process(self):
+        data = read_planetoid_data(self.raw_dir, self.name)
+
+        if self.split == "geom-gcn":
+            train_masks, val_masks, test_masks = [], [], []
+            for i in range(10):
+                name = f"{self.name.lower()}_split_0.6_0.2_{i}.npz"
+                splits = np.load(osp.join(self.raw_dir, name))
+                train_masks.append(splits["train_mask"])
+                val_masks.append(splits["val_mask"])
+                test_masks.append(splits["test_mask"])
+            data.train_mask = ops.convert_to_tensor(np.stack(train_masks, axis=1), dtype="bool")
+            data.val_mask = ops.convert_to_tensor(np.stack(val_masks, axis=1), dtype="bool")
+            data.test_mask = ops.convert_to_tensor(np.stack(test_masks, axis=1), dtype="bool")
+
+        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}()"
+

k3_node/datasets/polblogs.py

@@ -0,0 +1,63 @@
+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 PolBlogs(InMemoryDataset):
+    r"""The Political Blogs dataset containing 1,490 vertices and 19,025 edges.
+
+    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://netset.telecom-paris.fr/datasets/polblogs.tar.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 ["adjacency.tsv", "labels.tsv"]
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        tar_path = osp.join(self.raw_dir, "polblogs.tar.gz")
+        fs.cp(self.url, tar_path, extract=True)
+        if osp.exists(tar_path):
+            fs.rm(tar_path)
+
+    def process(self):
+        adj = np.loadtxt(self.raw_paths[0], delimiter="\t", usecols=(0, 1), dtype=np.int64)
+        edge_index = adj.T
+
+        y = np.loadtxt(self.raw_paths[1], delimiter="\t", usecols=(1,), dtype=np.int64)
+
+        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/ppi.py

@@ -0,0 +1,189 @@
+import json
+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 (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+from k3_node.layers.conv.utils import remove_self_loops
+
+
+class PPI(InMemoryDataset):
+    r"""The protein-protein interaction networks from the `"Predicting
+    Multicellular Function through Multi-layer Tissue Networks"
+    <https://arxiv.org/abs/1707.04638>`_ paper, containing positional gene
+    sets, motif gene sets and immunological signatures as features (50 in
+    total) and gene ontology sets as labels (121 in total).
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
+        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 10
+        :header-rows: 1
+
+        * - #graphs
+          - #nodes
+          - #edges
+          - #features
+          - #tasks
+        * - 20
+          - ~2,245.3
+          - ~61,318.4
+          - 50
+          - 121
+    """
+
+    url = "https://data.dgl.ai/dataset/ppi.zip"
+
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+        assert split.lower() in ["train", "val", "valid", "test"], f"Invalid split '{split}'"
+        self.split = "val" if split.lower() == "valid" else split.lower()
+
+        super().__init__(
+            root,
+            transform,
+            pre_transform,
+            pre_filter,
+            force_reload=force_reload,
+        )
+
+        if self.split == "train":
+            self.load(self.processed_paths[0])
+        elif self.split == "val":
+            self.load(self.processed_paths[1])
+        elif self.split == "test":
+            self.load(self.processed_paths[2])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        splits = ["train", "valid", "test"]
+        files = ["feats.npy", "graph_id.npy", "graph.json", "labels.npy"]
+        return [f"{split}_{name}" for split, name in product(splits, files)]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ["train.pt", "val.pt", "test.pt"]
+
+    def download(self) -> None:
+        path = download_url(self.url, self.root)
+        extract_zip(path, self.raw_dir)
+        if osp.exists(path):
+            os.unlink(path)
+
+    def process(self) -> None:
+        try:
+            import networkx as nx
+            from networkx.readwrite import json_graph
+            has_networkx = True
+        except ImportError:
+            has_networkx = False
+
+        for s, split in enumerate(["train", "valid", "test"]):
+            path = osp.join(self.raw_dir, f"{split}_graph.json")
+            with open(path, "r", encoding="utf-8") as f:
+                graph_json = json.load(f)
+
+            if has_networkx:
+                try:
+                    G = nx.DiGraph(json_graph.node_link_graph(graph_json, edges="links"))
+                except TypeError:
+                    G = nx.DiGraph(json_graph.node_link_graph(graph_json))
+            else:
+                links = graph_json.get("links", [])
+                src_all = np.array([link["source"] for link in links], dtype=np.int64)
+                dst_all = np.array([link["target"] for link in links], dtype=np.int64)
+
+            x_np = np.load(osp.join(self.raw_dir, f"{split}_feats.npy"))
+            y_np = np.load(osp.join(self.raw_dir, f"{split}_labels.npy"))
+
+            data_list = []
+            path = osp.join(self.raw_dir, f"{split}_graph_id.npy")
+            idx = np.load(path)
+            idx = idx - idx.min()
+
+            for i in range(int(idx.max()) + 1):
+                mask = (idx == i)
+                node_indices = np.where(mask)[0]
+
+                if has_networkx:
+                    G_s = G.subgraph(node_indices.tolist())
+                    edges = list(G_s.edges)
+                    if len(edges) > 0:
+                        edge_index = np.array(edges, dtype=np.int64).T
+                        edge_index = edge_index - edge_index.min()
+                    else:
+                        edge_index = np.zeros((2, 0), dtype=np.int64)
+                else:
+                    min_node, max_node = node_indices.min(), node_indices.max()
+                    edge_mask = (
+                        (src_all >= min_node)
+                        & (src_all <= max_node)
+                        & (dst_all >= min_node)
+                        & (dst_all <= max_node)
+                    )
+                    sub_src = src_all[edge_mask] - min_node
+                    sub_dst = dst_all[edge_mask] - min_node
+                    edge_index = np.stack([sub_src, sub_dst], axis=0)
+
+                edge_index = ops.convert_to_tensor(edge_index, dtype="int64")
+                edge_index, _ = remove_self_loops(edge_index)
+
+                x = ops.convert_to_tensor(x_np[mask], dtype="float32")
+                y = ops.convert_to_tensor(y_np[mask], dtype="float32")
+
+                data = Data(edge_index=edge_index, x=x, y=y)
+
+                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[s])
+
+    @property
+    def num_classes(self) -> int:
+        data = self.get(0)
+        y = getattr(data, "y", None)
+        if y is not None and len(y.shape) > 1:
+            return y.shape[-1]
+        return super().num_classes
+

k3_node/datasets/qm7.py

@@ -0,0 +1,65 @@
+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 QM7b(InMemoryDataset):
+    r"""The QM7b dataset consisting of 7,211 molecules with 14 regression targets."""
+
+    url = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/qm7b.mat"
+
+    def __init__(
+        self,
+        root: str,
+        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])
+
+    @property
+    def raw_file_names(self) -> str:
+        return "qm7b.mat"
+
+    @property
+    def processed_file_names(self) -> str:
+        return "data.pt"
+
+    def download(self):
+        fs.cp(self.url, self.raw_dir)
+
+    def process(self):
+        from scipy.io import loadmat
+
+        data = loadmat(self.raw_paths[0])
+        coulomb_matrix = data["X"]
+        target = data["T"].astype(np.float32)
+
+        data_list = []
+        for i in range(target.shape[0]):
+            nz = np.nonzero(coulomb_matrix[i])
+            edge_index = np.stack([nz[0], nz[1]], axis=0).astype(np.int64)
+            edge_attr = coulomb_matrix[i, edge_index[0], edge_index[1]].astype(np.float32)
+            y = target[i].reshape(1, -1)
+            num_nodes = int(np.max(edge_index)) + 1 if edge_index.size > 0 else 0
+
+            d = Data(
+                edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+                edge_attr=ops.convert_to_tensor(edge_attr, dtype="float32"),
+                y=ops.convert_to_tensor(y, dtype="float32"),
+                num_nodes=num_nodes,
+            )
+            data_list.append(d)
+
+        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/qm9.py

@@ -0,0 +1,132 @@
+import os
+import os.path as osp
+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
+from k3_node.data.extract import extract_zip
+
+HAR2EV = 27.211386246
+KCALMOL2EV = 0.04336414
+CONVERSION = np.array([1., 1., HAR2EV, HAR2EV, HAR2EV, 1., HAR2EV, HAR2EV, HAR2EV, HAR2EV, HAR2EV,
+                       1., KCALMOL2EV, KCALMOL2EV, KCALMOL2EV, KCALMOL2EV, 1., 1., 1.], dtype=np.float32)
+ATOMREFS = {
+    6: [0., 0., 0., 0., 0.],
+    7: [-13.61312172, -1029.86312267, -1485.30251237, -2042.61123593, -2713.48485589],
+    8: [-13.5745904, -1029.82456413, -1485.26398105, -2042.5727046, -2713.44632457],
+    9: [-13.54887564, -1029.79887659, -1485.2382935, -2042.54701705, -2713.42063702],
+    10: [-13.90303183, -1030.25891228, -1485.71166277, -2043.01812778, -2713.88796536],
+    11: [0., 0., 0., 0., 0.],
+}
+
+
+class QM9(InMemoryDataset):
+    r"""The QM9 dataset: about 130,000 small organic molecules with their 3D structure and 19
+    regression targets (dipole moment, HOMO/LUMO energies, internal energy, ...), as in PyG.
+
+    Every molecule has 11 atom features ``x`` (one-hot H/C/N/O/F, atomic number, aromatic,
+    sp/sp2/sp3 hybridization, number of hydrogens), atomic numbers ``z``, positions ``pos``,
+    one-hot bond types ``edge_attr`` and the targets ``y`` of shape ``[1, 19]`` (energies in eV).
+    Requires RDKit to process the raw files.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        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``)
+    """
+
+    raw_url = 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/molnet_publish/qm9.zip'
+    raw_url2 = 'https://ndownloader.figshare.com/files/3195404'
+
+    def __init__(self, root: str, 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])
+
+    def mean(self, target: int) -> float:
+        return float(np.asarray(self._data.y)[:, target].mean())
+
+    def std(self, target: int) -> float:
+        return float(np.asarray(self._data.y)[:, target].std(ddof=1))
+
+    def atomref(self, target: int) -> Optional[np.ndarray]:
+        r"""Per-element reference energies (a ``[100, 1]`` array indexed by atomic number), or
+        ``None`` for targets without them."""
+        if target not in ATOMREFS:
+            return None
+        out = np.zeros((100, 1), dtype=np.float32)
+        out[[1, 6, 7, 8, 9], 0] = ATOMREFS[target]
+        return out
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['gdb9.sdf', 'gdb9.sdf.csv', 'uncharacterized.txt']
+
+    @property
+    def processed_file_names(self) -> str:
+        return 'data_v3.pt'
+
+    def download(self):
+        path = download_url(self.raw_url, self.raw_dir)
+        extract_zip(path, self.raw_dir)
+        os.unlink(path)
+        download_url(self.raw_url2, self.raw_dir)
+        os.rename(osp.join(self.raw_dir, '3195404'), osp.join(self.raw_dir, 'uncharacterized.txt'))
+
+    def process(self):
+        from rdkit import Chem, RDLogger
+        from rdkit.Chem.rdchem import BondType as BT
+        from rdkit.Chem.rdchem import HybridizationType as HT
+
+        RDLogger.DisableLog('rdApp.*')
+        types = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4}
+        bonds = {BT.SINGLE: 0, BT.DOUBLE: 1, BT.TRIPLE: 2, BT.AROMATIC: 3}
+
+        with open(self.raw_paths[1]) as f:
+            target = np.array([[float(x) for x in line.split(',')[1:20]] for line in f.read().split('\n')[1:-1]],
+                              dtype=np.float32)
+        target = np.concatenate([target[:, 3:], target[:, :3]], axis=-1) * CONVERSION
+        with open(self.raw_paths[2]) as f:
+            skip = {int(x.split()[0]) - 1 for x in f.read().split('\n')[9:-2]}
+
+        data_list = []
+        for i, mol in enumerate(Chem.SDMolSupplier(self.raw_paths[0], removeHs=False, sanitize=False)):
+            if i in skip:
+                continue
+            N = mol.GetNumAtoms()
+            pos = mol.GetConformer().GetPositions().astype(np.float32)
+            atoms = list(mol.GetAtoms())
+            z = np.array([a.GetAtomicNum() for a in atoms], dtype=np.int64)
+            rows, cols, edge_types = [], [], []
+            for bond in mol.GetBonds():
+                s, e = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+                rows += [s, e]
+                cols += [e, s]
+                edge_types += 2 * [bonds[bond.GetBondType()]]
+            edge_index = np.array([rows, cols], dtype=np.int64).reshape(2, -1)
+            edge_type = np.array(edge_types, dtype=np.int64)
+            perm = np.argsort(edge_index[0] * N + edge_index[1], kind="stable")
+            edge_index, edge_type = edge_index[:, perm], edge_type[perm]
+
+            num_hs = np.zeros(N, dtype=np.float32)
+            np.add.at(num_hs, edge_index[1], (z == 1).astype(np.float32)[edge_index[0]])
+            hybrid = [a.GetHybridization() for a in atoms]
+            x = np.concatenate([
+                np.eye(len(types), dtype=np.float32)[[types[a.GetSymbol()] for a in atoms]],
+                np.stack([z, [a.GetIsAromatic() for a in atoms], [h == HT.SP for h in hybrid],
+                          [h == HT.SP2 for h in hybrid], [h == HT.SP3 for h in hybrid], num_hs], axis=1).astype(np.float32),
+            ], axis=1)
+            data = Data(x=x, z=z, pos=pos, edge_index=edge_index,
+                        edge_attr=np.eye(len(bonds), dtype=np.float32)[edge_type].reshape(-1, len(bonds)),
+                        y=target[i][None], smiles=Chem.MolToSmiles(mol, isomericSmiles=True),
+                        name=mol.GetProp('_Name'), idx=i)
+            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])