k3-node 1.0.0__py3-none-any.whl

k3_node/layers/unpool/test_unpool.py

@@ -0,0 +1,31 @@
+from keras import ops
+
+from k3_node.layers.unpool import knn_interpolate
+
+
+def test_knn_interpolate():
+    x = ops.convert_to_tensor(
+        [[1.0], [10.0], [100.0], [-1.0], [-10.0], [-100.0]], dtype="float32"
+    )
+    pos_x = ops.convert_to_tensor([
+        [-1.0, 0.0], [0.0, 0.0], [1.0, 0.0],
+        [-2.0, 0.0], [0.0, 0.0], [2.0, 0.0],
+    ], dtype="float32")
+    pos_y = ops.convert_to_tensor([
+        [-1.0, -1.0], [1.0, 1.0], [-2.0, -2.0], [2.0, 2.0],
+    ], dtype="float32")
+    batch_x = ops.convert_to_tensor([0, 0, 0, 1, 1, 1], dtype="int64")
+    batch_y = ops.convert_to_tensor([0, 0, 1, 1], dtype="int64")
+
+    y = knn_interpolate(x, pos_x, pos_y, batch_x, batch_y, k=2)
+    assert ops.shape(y) == (4, 1)
+    assert ops.convert_to_numpy(y).tolist() == [[4.0], [70.0], [-4.0], [-70.0]]
+
+
+def test_knn_interpolate_no_batch():
+    x = ops.convert_to_tensor([[1.0], [10.0], [100.0]], dtype="float32")
+    pos_x = ops.convert_to_tensor([[-1.0, 0.0], [0.0, 0.0], [1.0, 0.0]], dtype="float32")
+    pos_y = ops.convert_to_tensor([[-1.0, -1.0], [1.0, 1.0]], dtype="float32")
+
+    y = knn_interpolate(x, pos_x, pos_y, k=2)
+    assert ops.shape(y) == (2, 1)

k3_node/loader/__init__.py

@@ -0,0 +1,62 @@
+from .base import DataLoaderIterator
+from .cache import CachedLoader
+from .cluster import ClusterData, ClusterLoader
+from .data_list_loader import DataListLoader
+from .dataloader import Collater, DataLoader
+from .dense_data_loader import DenseDataLoader
+from .dynamic_batch_sampler import DynamicBatchSampler
+from .keras_dataset import FullGraphDataset
+from .graph_saint import (
+    GraphSAINTEdgeSampler,
+    GraphSAINTNodeSampler,
+    GraphSAINTRandomWalkSampler,
+    GraphSAINTSampler,
+)
+from .hgt_loader import HGTLoader
+from .imbalanced_sampler import ImbalancedSampler
+from .link_loader import EdgeSamplerInput, LinkLoader
+from .link_neighbor_loader import LinkNeighborLoader
+from .mixin import AffinityMixin, LogMemoryMixin, MultithreadingMixin
+from .neighbor_loader import NeighborLoader
+from .neighbor_sampler import Adj, EdgeIndex, NeighborSampler
+from .node_loader import HeteroSamplerOutput, NodeLoader, NodeSamplerInput, SamplerOutput
+from .prefetch import DeviceHelper, PrefetchLoader
+from .random_node_loader import RandomNodeLoader
+from .shadow import ShaDowKHopSampler
+from .temporal_dataloader import TemporalDataLoader
+from .zip_loader import ZipLoader
+from .utils import to_numpy
+
+__all__ = [
+    'to_numpy',
+    'DataLoader',
+    'NodeLoader',
+    'LinkLoader',
+    'NeighborLoader',
+    'LinkNeighborLoader',
+    'HGTLoader',
+    'ClusterData',
+    'ClusterLoader',
+    'GraphSAINTSampler',
+    'GraphSAINTNodeSampler',
+    'GraphSAINTEdgeSampler',
+    'GraphSAINTRandomWalkSampler',
+    'ShaDowKHopSampler',
+    'RandomNodeLoader',
+    'ZipLoader',
+    'DataListLoader',
+    'DenseDataLoader',
+    'TemporalDataLoader',
+    'NeighborSampler',
+    'ImbalancedSampler',
+    'DynamicBatchSampler',
+    'PrefetchLoader',
+    'CachedLoader',
+    'AffinityMixin',
+    'MultithreadingMixin',
+    'LogMemoryMixin',
+    'Collater',
+    'DataLoaderIterator',
+    'FullGraphDataset',
+]
+

k3_node/loader/base.py

@@ -0,0 +1,69 @@
+from typing import Any, Callable
+
+try:
+    import torch
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    class BaseDataLoader:
+        r"""Fallback BaseDataLoader when PyTorch is not installed."""
+        def __init__(self, dataset=None, batch_size=1, shuffle=False, **kwargs):
+            self.dataset = dataset
+            self.batch_size = batch_size
+            self.shuffle = shuffle
+
+        def __iter__(self):
+            collate_fn = getattr(self, 'collate_fn', None) or (lambda x: x)
+            dataset = getattr(self, 'dataset', [])
+            batch_size = getattr(self, 'batch_size', 1)
+            shuffle = getattr(self, 'shuffle', False)
+            indices = list(range(len(dataset)))
+            if shuffle:
+                import random
+                random.shuffle(indices)
+            for i in range(0, len(indices), batch_size):
+                batch_indices = indices[i:i + batch_size]
+                batch = [dataset[idx] for idx in batch_indices]
+                yield collate_fn(batch)
+
+        def __len__(self):
+            dataset = getattr(self, 'dataset', [])
+            batch_size = getattr(self, 'batch_size', 1)
+            return (len(dataset) + batch_size - 1) // batch_size if len(dataset) > 0 else 0
+
+try:
+    from torch.utils.data.dataloader import (
+        _BaseDataLoaderIter,
+        _MultiProcessingDataLoaderIter,
+    )
+except ImportError:
+    class _BaseDataLoaderIter:
+        pass
+
+    class _MultiProcessingDataLoaderIter:
+        pass
+
+
+class DataLoaderIterator:
+    r"""A data loader iterator extended by a post transformation function
+    :meth:`transform_fn`.
+    """
+    def __init__(self, iterator: Any, transform_fn: Callable):
+        self.iterator = iterator
+        self.transform_fn = transform_fn
+
+    def __iter__(self) -> 'DataLoaderIterator':
+        return self
+
+    def _reset(self, loader: Any, first_iter: bool = False):
+        if hasattr(self.iterator, '_reset'):
+            self.iterator._reset(loader, first_iter)
+
+    def __len__(self) -> int:
+        return len(self.iterator)
+
+    def __next__(self) -> Any:
+        return self.transform_fn(next(self.iterator))
+
+    def __del__(self) -> Any:
+        if isinstance(self.iterator, _MultiProcessingDataLoaderIter):
+            self.iterator.__del__()

k3_node/loader/cache.py

@@ -0,0 +1,68 @@
+from collections.abc import Mapping, Sequence
+from typing import Any, Callable, List, Optional
+
+try:
+    import torch
+    from torch.utils.data import DataLoader
+except ImportError:
+    torch = None
+    DataLoader = object
+
+
+def to_device(inputs: Any, device: Optional[Any] = None) -> Any:
+    if device is None:
+        return inputs
+    if hasattr(inputs, 'to'):
+        return inputs.to(device)
+    elif isinstance(inputs, Mapping):
+        return {key: to_device(value, device) for key, value in inputs.items()}
+    elif isinstance(inputs, tuple) and hasattr(inputs, '_fields'):
+        return type(inputs)(*(to_device(s, device) for s in zip(*inputs)))
+    elif isinstance(inputs, Sequence) and not isinstance(inputs, str):
+        return [to_device(s, device) for s in zip(*inputs)]
+    return inputs
+
+
+class CachedLoader:
+    r"""A loader to cache mini-batch outputs in memory across epochs.
+
+    Args:
+        loader (DataLoader): The data loader.
+        device (torch.device, optional): The device to load the data to. (default: :obj:`None`)
+        transform (callable, optional): A function that takes in a sampled mini-batch and returns a transformed version. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        loader: Any,
+        device: Optional[Any] = None,
+        transform: Optional[Callable] = None,
+    ):
+        self.loader = loader
+        self.device = device
+        self.transform = transform
+        self._cache: List[Any] = []
+
+    def clear(self):
+        r"""Clears the cache."""
+        self._cache = []
+
+    def __iter__(self) -> Any:
+        if len(self._cache) > 0:
+            for batch in self._cache:
+                yield batch
+            return
+
+        for batch in self.loader:
+            if self.transform is not None:
+                batch = self.transform(batch)
+
+            batch = to_device(batch, self.device)
+            self._cache.append(batch)
+            yield batch
+
+    def __len__(self) -> int:
+        return len(self.loader)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.loader})'
+

k3_node/loader/cluster.py

@@ -0,0 +1,127 @@
+import copy
+from typing import Any, List, Optional, Union
+
+import numpy as np
+
+try:
+    import torch
+    import torch.utils.data
+    from torch import Tensor
+    BaseDataset = torch.utils.data.Dataset
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    Tensor = type(None)
+    BaseDataset = object
+    BaseDataLoader = object
+
+from k3_node.data import Data
+from k3_node.loader.sampler_utils import partition_graph
+from k3_node.loader.keras_dataset import loader_bases
+
+
+class ClusterData(BaseDataset):
+    r"""Clusters/partitions a graph data object into multiple subgraphs, as
+    motivated by the "Cluster-GCN" paper.
+
+    Args:
+        data (Data): The graph data object.
+        num_parts (int): The number of partitions.
+        recursive (bool, optional): Multilevel recursive bisection if True. (default: :obj:`False`)
+        save_dir (str, optional): Directory to save partitioned data. (default: :obj:`None`)
+        log (bool, optional): If set to :obj:`False`, will not log. (default: :obj:`True`)
+        keep_inter_cluster_edges (bool, optional): Keep inter-cluster connections. (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        data: Data,
+        num_parts: int,
+        recursive: bool = False,
+        save_dir: Optional[str] = None,
+        filename: Optional[str] = None,
+        log: bool = True,
+        keep_inter_cluster_edges: bool = False,
+        sparse_format: str = 'csr',
+    ):
+        assert data.edge_index is not None
+
+        self.num_parts = num_parts
+        self.recursive = recursive
+        self.keep_inter_cluster_edges = keep_inter_cluster_edges
+        self.sparse_format = sparse_format
+        self.data = data
+
+        loaded = False
+        if save_dir is not None:
+            import os.path as osp
+            recursive_str = '_recursive' if recursive else ''
+            root_dir = osp.join(save_dir, f'part_{num_parts}{recursive_str}')
+            path = osp.join(root_dir, filename or 'metis.pt')
+            if osp.exists(path):
+                try:
+                    import torch
+                    part = torch.load(path, map_location="cpu", weights_only=False)
+                    if hasattr(part, "partptr") and hasattr(part, "node_perm"):
+                        partptr = np.asarray(part.partptr)
+                        node_perm = np.asarray(part.node_perm)
+                        self.part_nodes = [node_perm[partptr[i]:partptr[i+1]] for i in range(num_parts)]
+                        self.cluster = np.zeros(data.num_nodes, dtype=np.int64)
+                        for i in range(num_parts):
+                            self.cluster[self.part_nodes[i]] = i
+                        loaded = True
+                except Exception:
+                    pass
+
+        if not loaded:
+            self.cluster = partition_graph(data.edge_index, data.num_nodes, num_parts)
+            from k3_node.loader.utils import to_numpy
+            cluster_np = to_numpy(self.cluster)
+            sort_idx = np.argsort(cluster_np)
+            sorted_cluster = cluster_np[sort_idx]
+            split_idx = np.searchsorted(sorted_cluster, np.arange(num_parts + 1))
+            self.part_nodes = [sort_idx[split_idx[i]:split_idx[i+1]] for i in range(num_parts)]
+
+    def __len__(self) -> int:
+        return self.num_parts
+
+    def __getitem__(self, idx: int) -> Data:
+        nodes = self.part_nodes[idx]
+        return self.data.subgraph(nodes)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.num_parts})'
+
+
+class ClusterLoader(*loader_bases(BaseDataLoader)):
+    r"""The data loader scheme from Cluster-GCN which merges partitioned
+    subgraphs to form a mini-batch.
+
+    Args:
+        cluster_data (ClusterData): The already partitioned data object.
+        **kwargs (optional): Additional arguments of :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(self, cluster_data: ClusterData, **kwargs):
+        self.cluster_data = cluster_data
+        kwargs.pop('collate_fn', None)
+        iterator = range(len(cluster_data))
+
+        if torch is not None:
+            super().__init__(iterator, collate_fn=self._collate, **kwargs)
+        else:
+            self.dataset = iterator
+            self.collate_fn = self._collate
+
+    def _collate(self, batch: List[int]) -> Data:
+        all_nodes = []
+        is_torch = torch is not None and isinstance(self.cluster_data.cluster, Tensor)
+
+        for part_id in batch:
+            all_nodes.append(self.cluster_data.part_nodes[part_id])
+
+        if is_torch and all_nodes and isinstance(all_nodes[0], Tensor):
+            nodes = torch.cat(all_nodes, dim=0)
+        else:
+            nodes = np.concatenate([np.asarray(x) for x in all_nodes], axis=0)
+
+        return self.cluster_data.data.subgraph(nodes)
+

k3_node/loader/data_list_loader.py

@@ -0,0 +1,45 @@
+from typing import List, Union
+
+try:
+    import torch
+    import torch.utils.data
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    BaseDataLoader = object
+
+from k3_node.data import Dataset
+from k3_node.data.data import BaseData
+
+
+def collate_fn(data_list):
+    return data_list
+
+
+class DataListLoader(BaseDataLoader):
+    r"""A data loader which batches data objects from a
+    :class:`k3_node.data.Dataset` to a Python list without batch collation.
+    """
+    def __init__(
+        self,
+        dataset: Union[Dataset, List[BaseData]],
+        batch_size: int = 1,
+        shuffle: bool = False,
+        **kwargs,
+    ):
+        kwargs.pop('collate_fn', None)
+
+        if torch is not None:
+            super().__init__(
+                dataset,
+                batch_size=batch_size,
+                shuffle=shuffle,
+                collate_fn=collate_fn,
+                **kwargs,
+            )
+        else:
+            self.dataset = dataset
+            self.batch_size = batch_size
+            self.shuffle = shuffle
+            self.collate_fn = collate_fn
+

k3_node/loader/dataloader.py

@@ -0,0 +1,117 @@
+from collections.abc import Mapping, Sequence
+from typing import Any, List, Optional, Union
+
+import numpy as np
+
+try:
+    import torch
+    import torch.utils.data
+    from torch.utils.data.dataloader import default_collate
+except ImportError:
+    torch = None
+    default_collate = None
+
+from k3_node.data import Batch
+from k3_node.data.data import BaseData
+from k3_node.data.dataset import Dataset
+from k3_node.loader.keras_dataset import loader_bases
+
+
+class Collater:
+    r"""Collates a list of graph data objects or primitives into a mini-batch."""
+    def __init__(
+        self,
+        dataset: Optional[Union[Dataset, Sequence[BaseData]]] = None,
+        follow_batch: Optional[List[str]] = None,
+        exclude_keys: Optional[List[str]] = None,
+    ):
+        self.dataset = dataset
+        self.follow_batch = follow_batch
+        self.exclude_keys = exclude_keys
+
+    def __call__(self, batch: List[Any]) -> Any:
+        elem = batch[0]
+        if isinstance(elem, BaseData):
+            return Batch.from_data_list(
+                batch,
+                follow_batch=self.follow_batch,
+                exclude_keys=self.exclude_keys,
+            )
+        elif torch is not None and isinstance(elem, torch.Tensor):
+            return default_collate(batch)
+        elif isinstance(elem, np.ndarray):
+            return np.stack(batch, axis=0)
+        elif hasattr(elem, '__array__') and not isinstance(elem, (str, bytes)):
+            import keras
+            np_batch = np.stack([np.asarray(x) for x in batch], axis=0)
+            return keras.ops.convert_to_tensor(np_batch)
+        elif isinstance(elem, float):
+            if torch is not None:
+                return torch.tensor(batch, dtype=torch.float)
+            return np.array(batch, dtype=np.float32)
+        elif isinstance(elem, int):
+            if torch is not None:
+                return torch.tensor(batch, dtype=torch.long)
+            return np.array(batch, dtype=np.int64)
+        elif isinstance(elem, str):
+            return batch
+        elif isinstance(elem, Mapping):
+            return {key: self([data[key] for data in batch]) for key in elem}
+        elif isinstance(elem, tuple) and hasattr(elem, '_fields'):
+            return type(elem)(*(self(s) for s in zip(*batch)))
+        elif isinstance(elem, Sequence) and not isinstance(elem, str):
+            return [self(s) for s in zip(*batch)]
+
+        raise TypeError(f"DataLoader found invalid type: '{type(elem)}'")
+
+
+BaseDataLoader = torch.utils.data.DataLoader if torch is not None else object
+
+
+class DataLoader(*loader_bases(BaseDataLoader)):
+    r"""A data loader which merges data objects from a
+    :class:`k3_node.data.Dataset` to a mini-batch.
+    Data objects can be either of type :class:`~k3_node.data.Data` or
+    :class:`~k3_node.data.HeteroData`.
+
+    Args:
+        dataset (Dataset): The dataset from which to load the data.
+        batch_size (int, optional): How many samples per batch to load.
+            (default: :obj:`1`)
+        shuffle (bool, optional): If set to :obj:`True`, the data will be
+            reshuffled at every epoch. (default: :obj:`False`)
+        follow_batch (List[str], optional): Creates assignment batch
+            vectors for each key in the list. (default: :obj:`None`)
+        exclude_keys (List[str], optional): Will exclude each key in the
+            list. (default: :obj:`None`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(
+        self,
+        dataset: Union[Dataset, Sequence[BaseData]],
+        batch_size: int = 1,
+        shuffle: bool = False,
+        follow_batch: Optional[List[str]] = None,
+        exclude_keys: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        kwargs.pop('collate_fn', None)
+
+        self.follow_batch = follow_batch
+        self.exclude_keys = exclude_keys
+
+        if torch is not None:
+            super().__init__(
+                dataset,
+                batch_size=batch_size,
+                shuffle=shuffle,
+                collate_fn=Collater(dataset, follow_batch, exclude_keys),
+                **kwargs,
+            )
+        else:
+            self.dataset = dataset
+            self.batch_size = batch_size
+            self.shuffle = shuffle
+            self.collate_fn = Collater(dataset, follow_batch, exclude_keys)
+

k3_node/loader/dense_data_loader.py

@@ -0,0 +1,62 @@
+from typing import List, Union
+
+import numpy as np
+
+try:
+    import torch
+    import torch.utils.data
+    from torch.utils.data.dataloader import default_collate
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    default_collate = None
+    BaseDataLoader = object
+
+from k3_node.data import Batch, Data, Dataset
+from k3_node.loader.keras_dataset import loader_bases
+
+
+def collate_fn(data_list: List[Data]) -> Batch:
+    batch = Batch()
+    for key in data_list[0].keys():
+        vals = [data[key] for data in data_list]
+        if default_collate is not None and isinstance(vals[0], torch.Tensor):
+            batch[key] = default_collate(vals)
+        elif isinstance(vals[0], np.ndarray):
+            batch[key] = np.stack(vals, axis=0)
+        elif hasattr(vals[0], '__array__'):
+            import keras
+            batch[key] = keras.ops.stack(vals, axis=0)
+        else:
+            batch[key] = vals
+    return batch
+
+
+class DenseDataLoader(*loader_bases(BaseDataLoader)):
+    r"""A data loader which batches data objects from a
+    :class:`k3_node.data.Dataset` to a :class:`k3_node.data.Batch`
+    object by stacking all attributes in a new dimension.
+    """
+    def __init__(
+        self,
+        dataset: Union[Dataset, List[Data]],
+        batch_size: int = 1,
+        shuffle: bool = False,
+        **kwargs,
+    ):
+        kwargs.pop('collate_fn', None)
+
+        if torch is not None:
+            super().__init__(
+                dataset,
+                batch_size=batch_size,
+                shuffle=shuffle,
+                collate_fn=collate_fn,
+                **kwargs,
+            )
+        else:
+            self.dataset = dataset
+            self.batch_size = batch_size
+            self.shuffle = shuffle
+            self.collate_fn = collate_fn
+

k3_node/loader/dynamic_batch_sampler.py

@@ -0,0 +1,93 @@
+from typing import Iterator, List, Optional
+
+import numpy as np
+
+try:
+    import torch
+    import torch.utils.data.sampler
+    BaseSampler = torch.utils.data.sampler.Sampler
+except ImportError:
+    torch = None
+    BaseSampler = object
+
+from k3_node.data import Dataset
+
+
+class DynamicBatchSampler(BaseSampler):
+    r"""Dynamically adds samples to a mini-batch up to a maximum size (either
+    based on number of nodes or number of edges).
+
+    Args:
+        dataset (Dataset): Dataset to sample from.
+        max_num (int): Size of mini-batch to aim for in number of nodes or edges.
+        mode (str, optional): :obj:`"node"` or :obj:`"edge"` to measure batch size. (default: :obj:`"node"`)
+        shuffle (bool, optional): If set to :obj:`True`, will have the data reshuffled at every epoch. (default: :obj:`False`)
+        skip_too_big (bool, optional): If set to :obj:`True`, skip samples which cannot fit in a batch by itself. (default: :obj:`False`)
+        num_steps (int, optional): The number of mini-batches to draw for a single epoch. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        dataset: Dataset,
+        max_num: int,
+        mode: str = 'node',
+        shuffle: bool = False,
+        skip_too_big: bool = False,
+        num_steps: Optional[int] = None,
+    ):
+        if max_num <= 0:
+            raise ValueError(f"`max_num` should be a positive integer value (got {max_num})")
+        if mode not in ['node', 'edge']:
+            raise ValueError(f"`mode` choice should be either 'node' or 'edge' (got '{mode}')")
+
+        self.dataset = dataset
+        self.max_num = max_num
+        self.mode = mode
+        self.shuffle = shuffle
+        self.skip_too_big = skip_too_big
+        self.num_steps = num_steps
+        self.max_steps = num_steps or len(dataset)
+
+    def __iter__(self) -> Iterator[List[int]]:
+        if self.shuffle:
+            if torch is not None:
+                indices = torch.randperm(len(self.dataset)).tolist()
+            else:
+                indices = np.random.permutation(len(self.dataset)).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+        samples: List[int] = []
+        current_num: int = 0
+        num_steps: int = 0
+        num_processed: int = 0
+
+        while num_processed < len(self.dataset) and num_steps < self.max_steps:
+            for i in indices[num_processed:]:
+                data = self.dataset[i]
+                num = data.num_nodes if self.mode == 'node' else data.num_edges
+
+                if current_num + num > self.max_num:
+                    if current_num == 0:
+                        if self.skip_too_big:
+                            num_processed += 1
+                            continue
+                    else:
+                        break
+
+                samples.append(i)
+                num_processed += 1
+                current_num += num
+
+            yield samples
+            samples = []
+            current_num = 0
+            num_steps += 1
+
+    def __len__(self) -> int:
+        if self.num_steps is None:
+            raise ValueError(
+                f"The length of '{self.__class__.__name__}' is undefined since the number of steps per epoch "
+                f"is ambiguous. Either specify `num_steps` or use a static batch sampler."
+            )
+        return self.num_steps
+