k3-node 1.0.0__py3-none-any.whl

k3_node/loader/test_samplers.py

@@ -0,0 +1,96 @@
+import numpy as np
+import pytest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from k3_node.data import Data
+from k3_node.loader import (
+    CachedLoader,
+    DataLoader,
+    DynamicBatchSampler,
+    ImbalancedSampler,
+    PrefetchLoader,
+    ZipLoader,
+)
+
+
+def create_dummy_data(num_nodes=4, num_features=8, label=0):
+    x = np.random.randn(num_nodes, num_features).astype(np.float32)
+    edge_index = np.array([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=np.int64)
+    y = np.array([label], dtype=np.int64)
+    if torch is not None:
+        x = torch.from_numpy(x)
+        edge_index = torch.from_numpy(edge_index)
+        y = torch.from_numpy(y)
+    return Data(x=x, edge_index=edge_index, y=y)
+
+
+def test_dynamic_batch_sampler():
+    dataset = [create_dummy_data(num_nodes=i + 2) for i in range(5)]
+    sampler = DynamicBatchSampler(dataset, max_num=15, mode='node')
+
+    loader = DataLoader(dataset, batch_sampler=sampler)
+    count = 0
+    for batch in loader:
+        assert batch.num_nodes <= 15
+        count += 1
+    assert count > 0
+
+    with pytest.raises(ValueError, match="length of 'DynamicBatchSampler'"):
+        len(sampler)
+    assert len(DynamicBatchSampler(dataset, max_num=15, num_steps=2)) == 2
+
+
+def test_imbalanced_sampler():
+    # 8 samples with label 0, 2 samples with label 1
+    dataset = [create_dummy_data(label=0) for _ in range(8)] + [create_dummy_data(label=1) for _ in range(2)]
+    sampler = ImbalancedSampler(dataset, num_samples=10)
+
+    sampled_indices = list(sampler)
+    assert len(sampled_indices) == 10
+
+    loader = DataLoader(dataset, batch_size=5, sampler=sampler)
+    batches = list(loader)
+    assert len(batches) == 2
+
+
+def test_zip_loader():
+    from k3_node.loader import NeighborLoader
+    data = create_dummy_data(num_nodes=10)
+
+    loader1 = NeighborLoader(data, num_neighbors=[2], input_nodes=[0, 1, 2, 3])
+    loader2 = NeighborLoader(data, num_neighbors=[2], input_nodes=[4, 5, 6, 7])
+
+    zip_loader = ZipLoader([loader1, loader2], batch_size=2)
+    for batch1, batch2 in zip_loader:
+        assert batch1.batch_size == 2
+        assert batch2.batch_size == 2
+
+
+def test_cached_loader():
+    dataset = [create_dummy_data(num_nodes=3) for _ in range(4)]
+    loader = DataLoader(dataset, batch_size=2)
+
+    cached_loader = CachedLoader(loader)
+    epoch1 = list(cached_loader)
+    epoch2 = list(cached_loader)
+
+    assert len(epoch1) == 2
+    assert len(epoch2) == 2
+    assert len(cached_loader) == 2
+
+    cached_loader.clear()
+    assert len(cached_loader._cache) == 0
+
+
+def test_prefetch_loader():
+    dataset = [create_dummy_data(num_nodes=3) for _ in range(4)]
+    loader = DataLoader(dataset, batch_size=2)
+
+    prefetch = PrefetchLoader(loader)
+    batches = list(prefetch)
+    assert len(batches) == 2
+    assert len(prefetch) == 2

k3_node/loader/test_subgraph_loaders.py

@@ -0,0 +1,89 @@
+import numpy as np
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from k3_node.data import Data
+from k3_node.loader import (
+    ClusterData,
+    ClusterLoader,
+    GraphSAINTEdgeSampler,
+    GraphSAINTNodeSampler,
+    GraphSAINTRandomWalkSampler,
+    NeighborSampler,
+    ShaDowKHopSampler,
+)
+
+
+def get_graph(num_nodes=12):
+    # Circular ring graph
+    row = np.arange(num_nodes)
+    col = (row + 1) % num_nodes
+    edge_index = np.stack([row, col], axis=0).astype(np.int64)
+    x = np.random.randn(num_nodes, 8).astype(np.float32)
+    y = np.random.randint(0, 2, size=(num_nodes,)).astype(np.int64)
+
+    if torch is not None:
+        edge_index = torch.from_numpy(edge_index)
+        x = torch.from_numpy(x)
+        y = torch.from_numpy(y)
+
+    return Data(x=x, edge_index=edge_index, y=y)
+
+
+def test_cluster_loader():
+    data = get_graph(num_nodes=12)
+    cluster_data = ClusterData(data, num_parts=3)
+    assert len(cluster_data) == 3
+
+    sub0 = cluster_data[0]
+    assert isinstance(sub0, Data)
+    assert sub0.num_nodes > 0
+
+    loader = ClusterLoader(cluster_data, batch_size=2, shuffle=False)
+    batches = list(loader)
+    assert len(batches) == 2  # 3 parts with batch_size 2 => 2 batches
+
+
+def test_graph_saint():
+    data = get_graph(num_nodes=12)
+
+    # Node sampler
+    node_sampler = GraphSAINTNodeSampler(data, batch_size=4, num_steps=2, sample_coverage=1)
+    batch = next(iter(node_sampler))
+    assert isinstance(batch, Data)
+    assert hasattr(batch, 'node_norm')
+
+    # Edge sampler
+    edge_sampler = GraphSAINTEdgeSampler(data, batch_size=4, num_steps=2)
+    batch = next(iter(edge_sampler))
+    assert isinstance(batch, Data)
+
+    # Random walk sampler
+    rw_sampler = GraphSAINTRandomWalkSampler(data, batch_size=4, walk_length=2, num_steps=2)
+    batch = next(iter(rw_sampler))
+    assert isinstance(batch, Data)
+
+
+def test_shadow_k_hop_sampler():
+    data = get_graph(num_nodes=12)
+    loader = ShaDowKHopSampler(data, depth=2, num_neighbors=2, batch_size=3)
+
+    batch = next(iter(loader))
+    assert batch.num_graphs == 3
+    assert hasattr(batch, 'root_n_id')
+
+
+def test_legacy_neighbor_sampler():
+    data = get_graph(num_nodes=12)
+    loader = NeighborSampler(data.edge_index, sizes=[2, 2], batch_size=3)
+
+    batch_size, n_id, adjs = next(iter(loader))
+    assert batch_size == 3
+    assert len(n_id) >= 3
+    assert len(adjs) == 2
+    for adj in adjs:
+        assert hasattr(adj, 'edge_index')
+        assert hasattr(adj, 'size')
+

k3_node/loader/utils.py

@@ -0,0 +1,232 @@
+import copy
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+try:
+    import torch
+except ImportError:
+    torch = None
+    Tensor = type(None)
+
+from k3_node.data import Data, HeteroData
+from k3_node.data.storage import NodeStorage, EdgeStorage
+
+
+def to_numpy_or_tensor(x):
+    if torch is not None and isinstance(x, torch.Tensor):
+        return x
+    if isinstance(x, np.ndarray):
+        return x
+    if hasattr(x, '__array__'):
+        return np.asarray(x)
+    return x
+
+
+def to_numpy(x, dtype=None):
+    r"""Converts any tensor (PyTorch, TensorFlow, JAX) or array to a CPU NumPy ndarray."""
+    if hasattr(x, "cpu"):
+        x = x.cpu()
+    if hasattr(x, "detach"):
+        x = x.detach()
+    if hasattr(x, "numpy") and callable(x.numpy):
+        x = x.numpy()
+    return np.asarray(x, dtype=dtype)
+
+
+def index_select(value: Any, index: Any, dim: int = 0) -> Any:
+    r"""Indexes the :obj:`value` tensor along dimension :obj:`dim` using the
+    entries in :obj:`index`. Supports PyTorch, TensorFlow, JAX, and NumPy.
+    """
+    if torch is not None and isinstance(value, torch.Tensor):
+        if not isinstance(index, torch.Tensor):
+            index = torch.as_tensor(index, dtype=torch.long, device=value.device)
+        else:
+            index = index.to(dtype=torch.long, device=value.device)
+        return torch.index_select(value, dim, index)
+
+    # NumPy / Keras / JAX / TensorFlow array:
+    if hasattr(value, 'numpy'):
+        is_tf_or_jax = True
+        np_val = value.numpy()
+    elif isinstance(value, np.ndarray):
+        is_tf_or_jax = False
+        np_val = value
+    elif hasattr(value, '__array__'):
+        is_tf_or_jax = False
+        np_val = np.asarray(value)
+    else:
+        return value
+
+    if torch is not None and isinstance(index, torch.Tensor):
+        np_idx = index.cpu().numpy()
+    else:
+        np_idx = np.asarray(index, dtype=np.int64)
+
+    res = np.take(np_val, np_idx, axis=dim)
+    if is_tf_or_jax:
+        import keras
+        return keras.ops.convert_to_tensor(res)
+    return res
+
+
+def filter_node_store_(store: NodeStorage, out_store: NodeStorage, index: Any):
+    for key, value in store.items():
+        if key == 'num_nodes':
+            numel = index.numel() if hasattr(index, 'numel') else len(index)
+            out_store.num_nodes = numel
+        elif store.is_node_attr(key):
+            dim = 0
+            if hasattr(store, '_parent') and store._parent() is not None:
+                dim = store._parent().__cat_dim__(key, value, store)
+            out_store[key] = index_select(value, index, dim=dim)
+
+
+def filter_edge_store_(
+    store: EdgeStorage,
+    out_store: EdgeStorage,
+    row: Any,
+    col: Any,
+    index: Optional[Any],
+    perm: Optional[Any] = None,
+):
+    for key, value in store.items():
+        if key == 'edge_index':
+            if torch is not None and isinstance(row, torch.Tensor):
+                edge_index = torch.stack([row, col], dim=0)
+            else:
+                edge_index = np.stack([np.asarray(row), np.asarray(col)], axis=0)
+            out_store.edge_index = edge_index
+        elif store.is_edge_attr(key):
+            if index is None:
+                out_store[key] = None
+                continue
+            dim = 0
+            if hasattr(store, '_parent') and store._parent() is not None:
+                dim = store._parent().__cat_dim__(key, value, store)
+            if perm is None:
+                out_store[key] = index_select(value, index, dim=dim)
+            else:
+                sel_idx = perm[index] if hasattr(perm, '__getitem__') else index
+                out_store[key] = index_select(value, sel_idx, dim=dim)
+
+
+def filter_data(data: Data, node: Any, row: Any, col: Any, edge: Optional[Any] = None, perm: Optional[Any] = None) -> Data:
+    out = copy.copy(data)
+    out._store = copy.copy(data._store)
+    filter_node_store_(data._store, out._store, node)
+    filter_edge_store_(data._store, out._store, row, col, edge, perm)
+    return out
+
+
+def filter_hetero_data(
+    data: HeteroData,
+    node_dict: Dict[str, Any],
+    row_dict: Dict[Tuple[str, str, str], Any],
+    col_dict: Dict[Tuple[str, str, str], Any],
+    edge_dict: Dict[Tuple[str, str, str], Optional[Any]],
+    perm_dict: Optional[Dict[Tuple[str, str, str], Optional[Any]]] = None,
+) -> HeteroData:
+    out = copy.copy(data)
+    out._node_store_dict = {k: copy.copy(v) for k, v in data._node_store_dict.items()}
+    out._edge_store_dict = {k: copy.copy(v) for k, v in data._edge_store_dict.items()}
+
+    for node_type in out.node_types:
+        if node_type not in node_dict:
+            node_dict[node_type] = torch.empty(0, dtype=torch.long) if torch is not None else np.empty(0, dtype=np.int64)
+        filter_node_store_(data[node_type], out[node_type], node_dict[node_type])
+
+    for edge_type in out.edge_types:
+        canonical = data._to_canonical(*edge_type) if hasattr(data, '_to_canonical') else edge_type
+        if canonical not in row_dict:
+            empty_arr = torch.empty(0, dtype=torch.long) if torch is not None else np.empty(0, dtype=np.int64)
+            row_dict[canonical] = empty_arr
+            col_dict[canonical] = empty_arr
+            edge_dict[canonical] = empty_arr
+
+        filter_edge_store_(
+            data[edge_type],
+            out[edge_type],
+            row_dict[canonical],
+            col_dict[canonical],
+            edge_dict[canonical],
+            perm_dict.get(canonical, None) if perm_dict else None,
+        )
+
+    return out
+
+
+def get_input_nodes(
+    data: Union[Data, HeteroData],
+    input_nodes: Any,
+    input_id: Optional[Any] = None,
+) -> Tuple[Optional[str], Any, Optional[Any]]:
+    def to_index(nodes, in_id):
+        if torch is not None and isinstance(nodes, torch.Tensor):
+            if nodes.dtype == torch.bool:
+                nodes = nodes.nonzero(as_tuple=False).view(-1)
+                in_id = nodes if in_id is None else in_id
+            return nodes, in_id
+        if isinstance(nodes, np.ndarray) and nodes.dtype == bool:
+            nodes = np.nonzero(nodes)[0]
+            in_id = nodes if in_id is None else in_id
+            return nodes, in_id
+        if torch is not None and not isinstance(nodes, torch.Tensor):
+            nodes = torch.tensor(nodes, dtype=torch.long)
+        elif torch is None:
+            nodes = np.asarray(nodes, dtype=np.int64)
+        return nodes, in_id
+
+    if isinstance(data, Data):
+        if input_nodes is None:
+            nodes = torch.arange(data.num_nodes) if torch is not None else np.arange(data.num_nodes)
+            return None, nodes, None
+        return None, *to_index(input_nodes, input_id)
+
+    elif isinstance(data, HeteroData):
+        assert input_nodes is not None
+        if isinstance(input_nodes, str):
+            num_nodes = data[input_nodes].num_nodes
+            nodes = torch.arange(num_nodes) if torch is not None else np.arange(num_nodes)
+            return input_nodes, nodes, None
+
+        assert isinstance(input_nodes, (list, tuple)) and len(input_nodes) == 2
+        node_type, input_nodes = input_nodes
+        if input_nodes is None:
+            num_nodes = data[node_type].num_nodes
+            nodes = torch.arange(num_nodes) if torch is not None else np.arange(num_nodes)
+            return node_type, nodes, None
+        return node_type, *to_index(input_nodes, input_id)
+
+    raise TypeError(f"Invalid data type: {type(data)}")
+
+
+def get_edge_label_index(
+    data: Union[Data, HeteroData],
+    edge_label_index: Any,
+) -> Tuple[Optional[Tuple[str, str, str]], Any]:
+    if isinstance(data, Data):
+        if edge_label_index is None:
+            return None, data.edge_index
+        return None, edge_label_index
+
+    if isinstance(data, HeteroData):
+        assert edge_label_index is not None
+        if isinstance(edge_label_index, (list, tuple)) and len(edge_label_index) == 3 and isinstance(edge_label_index[0], str):
+            edge_type = data._to_canonical(*edge_label_index)
+            return edge_type, data[edge_type].edge_index
+
+        assert isinstance(edge_label_index, (list, tuple)) and len(edge_label_index) == 2
+        edge_type, edge_index = edge_label_index
+        edge_type = data._to_canonical(*edge_type)
+        if edge_index is None:
+            return edge_type, data[edge_type].edge_index
+        return edge_type, edge_index
+
+    raise TypeError(f"Invalid data type: {type(data)}")
+
+
+def infer_filter_per_worker(data: Any) -> bool:
+    out = True
+    if hasattr(data, 'is_cuda') and data.is_cuda:
+        out = False
+    return out

k3_node/loader/zip_loader.py

@@ -0,0 +1,88 @@
+from typing import Any, Iterator, List, Optional, Tuple, Union
+
+try:
+    import torch
+    from torch import Tensor
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    Tensor = type(None)
+    BaseDataLoader = object
+
+from k3_node.data import Data, HeteroData
+from k3_node.loader.base import DataLoaderIterator
+from k3_node.loader.utils import infer_filter_per_worker
+
+
+class ZipLoader(BaseDataLoader):
+    r"""A loader that returns a tuple of data objects by sampling from multiple
+    loader instances.
+
+    Args:
+        loaders (List[Any]): The loader instances.
+        filter_per_worker (bool, optional): If set to :obj:`True`, will filter
+            the returned data in each worker's subprocess. (default: :obj:`None`)
+        **kwargs (optional): Additional arguments of :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(
+        self,
+        loaders: List[Any],
+        filter_per_worker: Optional[bool] = None,
+        **kwargs,
+    ):
+        if filter_per_worker is None:
+            first_data = getattr(loaders[0], 'data', None)
+            filter_per_worker = infer_filter_per_worker(first_data) if first_data is not None else True
+
+        kwargs.pop('dataset', None)
+        kwargs.pop('collate_fn', None)
+
+        for loader in loaders:
+            if not callable(getattr(loader, 'collate_fn', None)):
+                raise ValueError(f"'{loader.__class__.__name__}' does not have a 'collate_fn' method")
+            if not callable(getattr(loader, 'filter_fn', None)):
+                raise ValueError(f"'{loader.__class__.__name__}' does not have a 'filter_fn' method")
+            loader.filter_per_worker = filter_per_worker
+
+        lens = []
+        for loader in loaders:
+            if hasattr(loader, 'dataset'):
+                lens.append(len(loader.dataset))
+            elif hasattr(loader, '__len__'):
+                lens.append(len(loader))
+            else:
+                lens.append(0)
+
+        iterator = range(min(lens) if lens else 0)
+
+        self.loaders = loaders
+        self.filter_per_worker = filter_per_worker
+
+        if torch is not None:
+            super().__init__(iterator, collate_fn=self.collate_fn, **kwargs)
+        else:
+            self.dataset = iterator
+            self.collate_fn = self.collate_fn
+
+    def __call__(self, index: Union[Tensor, List[int]]) -> Union[Tuple[Data, ...], Tuple[HeteroData, ...]]:
+        out = self.collate_fn(index)
+        if not self.filter_per_worker:
+            out = self.filter_fn(out)
+        return out
+
+    def collate_fn(self, index: List[int]) -> Tuple[Any, ...]:
+        if torch is not None and not isinstance(index, Tensor):
+            index = torch.tensor(index, dtype=torch.long)
+        return tuple(loader.collate_fn(index) for loader in self.loaders)
+
+    def filter_fn(self, outs: Tuple[Any, ...]) -> Tuple[Union[Data, HeteroData], ...]:
+        return tuple(loader.filter_fn(v) for loader, v in zip(self.loaders, outs))
+
+    def _get_iterator(self) -> Iterator:
+        if self.filter_per_worker:
+            return super()._get_iterator()
+        return DataLoaderIterator(super()._get_iterator(), self.filter_fn)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(loaders={self.loaders})'
+

k3_node/metrics.py

@@ -0,0 +1,94 @@
+"""Metrics for graph learning tasks."""
+import numpy as np
+from typing import Optional
+
+import keras
+from keras import ops
+
+
+@keras.saving.register_keras_serializable(package="k3_node")
+class F1Score(keras.metrics.F1Score):
+    r"""F1 score that also accepts logits, e.g. for multi-label node classification.
+
+    Same as :class:`keras.metrics.F1Score`, plus ``from_logits``: when :obj:`True`, a sigmoid is
+    applied to the predictions first, so models can output logits (as used with
+    ``BinaryCrossentropy(from_logits=True)``).
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.metrics import F1Score
+
+        f1 = F1Score(average="micro", from_logits=True)
+        y_true = np.array([[1, 0, 1], [0, 1, 0]], dtype="float32")
+        logits = np.array([[2.0, -1.0, 0.5], [-3.0, 1.5, 0.2]], dtype="float32")
+        f1.update_state(y_true, logits)
+        print(round(float(f1.result()), 2))  # 0.86
+        ```
+    """
+
+    def __init__(self, average=None, threshold=0.5, from_logits=False, name="f1_score", dtype=None):
+        super().__init__(average=average, threshold=threshold, name=name, dtype=dtype)
+        self.from_logits = from_logits
+
+    def update_state(self, y_true, y_pred, sample_weight=None):
+        if self.from_logits:
+            y_pred = ops.sigmoid(y_pred)
+        return super().update_state(y_true, y_pred, sample_weight=sample_weight)
+
+    def get_config(self):
+        return {**super().get_config(), "from_logits": self.from_logits}
+
+
+def precision_recall_at_k(user_emb, item_emb, train_edge_index, test_edge_index, k: int = 20,
+                          num_users: Optional[int] = None, batch_size: int = 8192):
+    r"""Top-``k`` recommendation quality, as in PyG's LightGCN example.
+
+    Every user's items are ranked by the dot product of the embeddings, excluding the items the
+    user interacted with during training. Returns the precision@k and recall@k averaged over the
+    users with at least one test interaction.
+
+    Args:
+        user_emb: User embeddings ``[num_users, dim]``.
+        item_emb: Item embeddings ``[num_items, dim]``.
+        train_edge_index: Training interactions ``(user, item)``; item ids may be offset by
+            ``num_users`` (as in a homogeneous user-item graph).
+        test_edge_index: Test interactions ``(user, item)``, with the same numbering.
+        k (int): The number of recommendations per user. (default: ``20``)
+        num_users (int, optional): The item-id offset. (default: ``len(user_emb)``)
+        batch_size (int): Users scored at once. (default: ``8192``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.metrics import precision_recall_at_k
+
+        users, items = np.eye(2, 4, dtype="float32"), np.eye(3, 4, dtype="float32")  # toy embeddings
+        train = np.array([[0], [2]])  # user 0 already has item 2 (id 2 + num_users = 4 in the graph)
+        test = np.array([[0, 1], [0 + 2, 1 + 2]])  # user 0 likes item 0, user 1 likes item 1
+        print(precision_recall_at_k(users, items, train + [[0], [2]], test, k=1))  # (1.0, 1.0)
+        ```
+    """
+    from keras import ops as _ops
+
+    user_emb, item_emb = (np.asarray(_ops.convert_to_numpy(e)) for e in (user_emb, item_emb))
+    num_users = len(user_emb) if num_users is None else num_users
+    train = np.asarray(_ops.convert_to_numpy(train_edge_index)).astype(np.int64)
+    test = np.asarray(_ops.convert_to_numpy(test_edge_index)).astype(np.int64)
+    train = train[:, train[0] < num_users]
+    precision = recall = examples = 0.0
+    for start in range(0, len(user_emb), batch_size):
+        end = min(start + batch_size, len(user_emb))
+        logits = user_emb[start:end] @ item_emb.T
+        m = (train[0] >= start) & (train[0] < end)
+        logits[train[0, m] - start, train[1, m] - num_users] = -np.inf  # skip already known items
+        truth = np.zeros_like(logits, dtype=bool)
+        m = (test[0] >= start) & (test[0] < end)
+        truth[test[0, m] - start, test[1, m] - num_users] = True
+        count = truth.sum(axis=1)
+        top = np.argpartition(-logits, k - 1, axis=1)[:, :k]
+        hits = np.take_along_axis(truth, top, axis=1).sum(axis=1)
+        precision += float((hits / k)[count > 0].sum())
+        recall += float((hits / np.maximum(count, 1e-6))[count > 0].sum())
+        examples += int((count > 0).sum())
+    return precision / examples, recall / examples