k3-node 1.0.0__py3-none-any.whl

k3_node/data/dataset.py

@@ -0,0 +1,182 @@
+import copy
+import os
+import os.path as osp
+from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from keras import ops
+
+from k3_node.data.data import BaseData
+from k3_node.data.storage import is_tensor_like, to_numpy
+
+
+class Dataset:
+    """Dataset base class for creating graph datasets."""
+
+    def __init__(
+        self,
+        root: Optional[str] = None,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        log: bool = True,
+        force_reload: bool = False,
+    ):
+        self.root = root
+        self.transform = transform
+        self.pre_transform = pre_transform
+        self.pre_filter = pre_filter
+        self.log = log
+        self.force_reload = force_reload
+        self._indices: Optional[Sequence] = None
+
+        if self.has_download:
+            self._download()
+
+        if self.has_process:
+            self._process()
+
+    @property
+    def raw_file_names(self) -> Union[str, List[str], Tuple[str, ...]]:
+        return []
+
+    @property
+    def processed_file_names(self) -> Union[str, List[str], Tuple[str, ...]]:
+        return []
+
+    @property
+    def raw_dir(self) -> str:
+        return osp.join(self.root or "", "raw")
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root or "", "processed")
+
+    @property
+    def raw_paths(self) -> List[str]:
+        files = self.raw_file_names
+        if isinstance(files, str):
+            files = [files]
+        return [osp.join(self.raw_dir, f) for f in files]
+
+    @property
+    def processed_paths(self) -> List[str]:
+        files = self.processed_file_names
+        if isinstance(files, str):
+            files = [files]
+        return [osp.join(self.processed_dir, f) for f in files]
+
+    @property
+    def has_download(self) -> bool:
+        return any("download" in cls.__dict__ for cls in self.__class__.__mro__ if cls is not Dataset)
+
+    @property
+    def has_process(self) -> bool:
+        return any("process" in cls.__dict__ for cls in self.__class__.__mro__ if cls is not Dataset)
+
+    def download(self):
+        pass
+
+    def process(self):
+        pass
+
+    def _download(self):
+        if all(osp.exists(p) for p in self.raw_paths if p):
+            return
+        os.makedirs(self.raw_dir, exist_ok=True)
+        self.download()
+
+    def _process(self):
+        if not self.force_reload and len(self.processed_paths) > 0 and all(osp.exists(p) for p in self.processed_paths):
+            return
+        os.makedirs(self.processed_dir, exist_ok=True)
+        self.process()
+
+    def len(self) -> int:
+        raise NotImplementedError
+
+    def get(self, idx: int) -> BaseData:
+        raise NotImplementedError
+
+    def indices(self) -> Sequence:
+        return range(self.len()) if self._indices is None else self._indices
+
+    def __len__(self) -> int:
+        return len(self.indices())
+
+    def __getitem__(self, idx: Any) -> Any:
+        if isinstance(idx, (int, np.integer)):
+            data = self.get(self.indices()[idx])
+            if hasattr(data, "to_backend"):
+                data = data.to_backend()
+            data = data if self.transform is None else self.transform(data)
+            return data
+        else:
+            return self.index_select(idx)
+
+    def index_select(self, idx: Any) -> "Dataset":
+        indices = list(self.indices())
+        if isinstance(idx, slice):
+            # Allow fractional slicing as in PyG, e.g. dataset[:0.9] for the first 90%
+            start, stop = idx.start, idx.stop
+            start = round(start * len(indices)) if isinstance(start, float) else start
+            stop = round(stop * len(indices)) if isinstance(stop, float) else stop
+            indices = indices[slice(start, stop, idx.step)]
+        elif is_tensor_like(idx):
+            idx_np = to_numpy(idx)
+            if idx_np.dtype == bool:
+                indices = [indices[i] for i in np.where(idx_np)[0]]
+            else:
+                indices = [indices[int(i)] for i in idx_np]
+        elif isinstance(idx, Sequence):
+            indices = [indices[i] for i in idx]
+        else:
+            raise IndexError(f"Invalid index type {type(idx)}")
+
+        dataset = copy.copy(self)
+        dataset._indices = indices
+        return dataset
+
+    def shuffle(self, return_perm: bool = False):
+        r"""Randomly shuffles the examples in the dataset (as in PyG).
+
+        Args:
+            return_perm (bool): If :obj:`True`, also returns the permutation used.
+        """
+        perm = np.random.permutation(len(self))
+        dataset = self.index_select(perm.tolist())
+        return (dataset, perm) if return_perm else dataset
+
+    def __iter__(self) -> Iterator[BaseData]:
+        for i in range(len(self)):
+            yield self[i]
+
+    @property
+    def num_node_features(self) -> int:
+        data = self[0]
+        return getattr(data, "num_node_features", 0)
+
+    @property
+    def num_features(self) -> int:
+        return self.num_node_features
+
+    @property
+    def num_edge_features(self) -> int:
+        data = self[0]
+        return getattr(data, "num_edge_features", 0)
+
+    @property
+    def num_classes(self) -> int:
+        y_list = [d.y for d in self if hasattr(d, "y") and d.y is not None]
+        if len(y_list) == 0:
+            return 0
+        y_np_list = [to_numpy(y) for y in y_list]
+        y = np.concatenate(y_np_list, axis=0) if len(y_np_list) > 1 else y_np_list[0]
+        if np.issubdtype(y.dtype, np.integer):
+            return int(np.max(y)) + 1
+        return len(np.unique(y))
+
+    def __repr__(self) -> str:
+        arg_repr = str(len(self)) if len(self) > 1 else ""
+        return f"{self.__class__.__name__}({arg_repr})"
+

k3_node/data/download.py

@@ -0,0 +1,49 @@
+import os
+import os.path as osp
+import ssl
+import sys
+import urllib.request
+from typing import Optional
+
+
+def download_url(
+    url: str,
+    folder: str,
+    log: bool = True,
+    filename: Optional[str] = None,
+) -> str:
+    """Downloads the content of a URL to a specific folder."""
+    if filename is None:
+        filename = url.rpartition("/")[2]
+        filename = filename if filename[0] == "?" else filename.split("?")[0]
+
+    path = osp.join(folder, filename)
+    if osp.exists(path):
+        return path
+
+    if log and "PYTEST_CURRENT_TEST" not in os.environ:
+        print(f"Downloading {url}", file=sys.stderr)
+
+    os.makedirs(folder, exist_ok=True)
+    context = ssl._create_unverified_context()
+    with urllib.request.urlopen(url, context=context) as response:
+        with open(path, "wb") as f:
+            while True:
+                chunk = response.read(10 * 1024 * 1024)
+                if not chunk:
+                    break
+                f.write(chunk)
+
+    return path
+
+
+def download_google_url(
+    id: str,
+    folder: str,
+    filename: str,
+    log: bool = True,
+) -> str:
+    """Downloads the content of a Google Drive ID to a specific folder."""
+    url = f"https://drive.usercontent.google.com/download?id={id}&confirm=t"
+    return download_url(url, folder, log, filename)
+

k3_node/data/extract.py

@@ -0,0 +1,45 @@
+import bz2
+import gzip
+import os
+import os.path as osp
+import sys
+import tarfile
+import zipfile
+
+
+def maybe_log(path: str, log: bool = True):
+    if log and "PYTEST_CURRENT_TEST" not in os.environ:
+        print(f"Extracting {path}", file=sys.stderr)
+
+
+def extract_tar(path: str, folder: str, mode: str = "r:gz", log: bool = True):
+    """Extracts a tar archive to a specific folder."""
+    maybe_log(path, log)
+    with tarfile.open(path, mode) as f:
+        f.extractall(folder)
+
+
+def extract_zip(path: str, folder: str, log: bool = True):
+    """Extracts a zip archive to a specific folder."""
+    maybe_log(path, log)
+    with zipfile.ZipFile(path, "r") as f:
+        f.extractall(folder)
+
+
+def extract_bz2(path: str, folder: str, log: bool = True):
+    """Extracts a bz2 archive to a specific folder."""
+    maybe_log(path, log)
+    path = osp.abspath(path)
+    with bz2.open(path, "r") as r:
+        with open(osp.join(folder, ".".join(path.split(".")[:-1])), "wb") as w:
+            w.write(r.read())
+
+
+def extract_gz(path: str, folder: str, log: bool = True):
+    """Extracts a gz archive to a specific folder."""
+    maybe_log(path, log)
+    path = osp.abspath(path)
+    with gzip.open(path, "r") as r:
+        with open(osp.join(folder, ".".join(path.split(".")[:-1])), "wb") as w:
+            w.write(r.read())
+

k3_node/data/feature_store.py

@@ -0,0 +1,70 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+class _FieldStatus(Enum):
+    UNSET = None
+
+
+@dataclass
+class TensorAttr:
+    """Defines the attributes of a FeatureStore tensor."""
+
+    group_name: Optional[Any] = _FieldStatus.UNSET
+    attr_name: Optional[str] = _FieldStatus.UNSET
+    index: Optional[Any] = _FieldStatus.UNSET
+
+    def is_set(self, key: str) -> bool:
+        return getattr(self, key) != _FieldStatus.UNSET
+
+    def is_fully_specified(self) -> bool:
+        return all(self.is_set(k) for k in ("group_name", "attr_name", "index"))
+
+
+class FeatureStore(ABC):
+    """Abstract base class for feature stores."""
+
+    def __init__(self):
+        self._feat_dict: Dict[Tuple[Any, str], Any] = {}
+
+    @abstractmethod
+    def _put_tensor(self, tensor: Any, attr: TensorAttr) -> bool:
+        pass
+
+    @abstractmethod
+    def _get_tensor(self, attr: TensorAttr) -> Optional[Any]:
+        pass
+
+    @abstractmethod
+    def _remove_tensor(self, attr: TensorAttr) -> bool:
+        pass
+
+    def put_tensor(self, tensor: Any, group_name: Any = None, attr_name: Optional[str] = None, index: Any = None) -> bool:
+        attr = TensorAttr(group_name=group_name, attr_name=attr_name, index=index)
+        return self._put_tensor(tensor, attr)
+
+    def get_tensor(self, group_name: Any = None, attr_name: Optional[str] = None, index: Any = None) -> Optional[Any]:
+        attr = TensorAttr(group_name=group_name, attr_name=attr_name, index=index)
+        return self._get_tensor(attr)
+
+    def remove_tensor(self, group_name: Any = None, attr_name: Optional[str] = None, index: Any = None) -> bool:
+        attr = TensorAttr(group_name=group_name, attr_name=attr_name, index=index)
+        return self._remove_tensor(attr)
+
+    def __getitem__(self, key: Any) -> Any:
+        if isinstance(key, tuple):
+            group_name, attr_name = key[:2]
+            index = key[2] if len(key) > 2 else None
+            return self.get_tensor(group_name=group_name, attr_name=attr_name, index=index)
+        return self.get_tensor(group_name=key)
+
+    def __setitem__(self, key: Any, value: Any):
+        if isinstance(key, tuple):
+            group_name, attr_name = key[:2]
+            index = key[2] if len(key) > 2 else None
+            self.put_tensor(value, group_name=group_name, attr_name=attr_name, index=index)
+        else:
+            self.put_tensor(value, group_name=key)
+

k3_node/data/graph_store.py

@@ -0,0 +1,92 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+class EdgeLayout(Enum):
+    COO = "coo"
+    CSC = "csc"
+    CSR = "csr"
+
+
+@dataclass
+class EdgeAttr:
+    """Defines the attributes of a GraphStore edge."""
+
+    edge_type: Any
+    layout: EdgeLayout
+    is_sorted: bool = False
+    size: Optional[Tuple[int, int]] = None
+
+    def __init__(
+        self,
+        edge_type: Any,
+        layout: Union[EdgeLayout, str] = EdgeLayout.COO,
+        is_sorted: bool = False,
+        size: Optional[Tuple[int, int]] = None,
+    ):
+        if isinstance(layout, str):
+            layout = EdgeLayout(layout.lower())
+        self.edge_type = edge_type
+        self.layout = layout
+        self.is_sorted = is_sorted
+        self.size = size
+
+
+class GraphStore(ABC):
+    """Abstract base class for graph edge stores."""
+
+    @abstractmethod
+    def _put_edge_index(self, edge_index: Any, edge_attr: EdgeAttr) -> bool:
+        pass
+
+    @abstractmethod
+    def _get_edge_index(self, edge_attr: EdgeAttr) -> Optional[Any]:
+        pass
+
+    @abstractmethod
+    def _remove_edge_index(self, edge_attr: EdgeAttr) -> bool:
+        pass
+
+    def put_edge_index(
+        self,
+        edge_index: Any,
+        edge_type: Any,
+        layout: Union[EdgeLayout, str] = EdgeLayout.COO,
+        is_sorted: bool = False,
+        size: Optional[Tuple[int, int]] = None,
+    ) -> bool:
+        attr = EdgeAttr(edge_type=edge_type, layout=layout, is_sorted=is_sorted, size=size)
+        return self._put_edge_index(edge_index, attr)
+
+    def get_edge_index(
+        self,
+        edge_type: Any,
+        layout: Union[EdgeLayout, str] = EdgeLayout.COO,
+        is_sorted: bool = False,
+    ) -> Optional[Any]:
+        attr = EdgeAttr(edge_type=edge_type, layout=layout, is_sorted=is_sorted)
+        return self._get_edge_index(attr)
+
+    def remove_edge_index(
+        self,
+        edge_type: Any,
+        layout: Union[EdgeLayout, str] = EdgeLayout.COO,
+    ) -> bool:
+        attr = EdgeAttr(edge_type=edge_type, layout=layout)
+        return self._remove_edge_index(attr)
+
+    def __getitem__(self, key: Any) -> Any:
+        if isinstance(key, tuple) and len(key) == 2:
+            edge_type, layout = key
+            return self.get_edge_index(edge_type=edge_type, layout=layout)
+        return self.get_edge_index(edge_type=key)
+
+    def __setitem__(self, key: Any, value: Any):
+        if isinstance(key, tuple) and len(key) == 2:
+            edge_type, layout = key
+            self.put_edge_index(value, edge_type=edge_type, layout=layout)
+        else:
+            self.put_edge_index(value, edge_type=key)
+