k3-node 1.0.0__py3-none-any.whl

k3_node/data/temporal.py

@@ -0,0 +1,154 @@
+import copy
+from typing import Any, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from keras import ops
+
+from k3_node.data.data import BaseData, size_repr
+from k3_node.data.storage import BaseStorage, GlobalStorage, get_shape, is_tensor_like
+
+
+class TemporalData(BaseData):
+    """A data object composed of a stream of events describing a temporal graph."""
+
+    def __init__(
+        self,
+        src: Optional[Any] = None,
+        dst: Optional[Any] = None,
+        t: Optional[Any] = None,
+        msg: Optional[Any] = None,
+        y: Optional[Any] = None,
+        **kwargs,
+    ):
+        self.__dict__["_store"] = GlobalStorage(_parent=self)
+        if src is not None:
+            self.src = src
+        if dst is not None:
+            self.dst = dst
+        if t is not None:
+            self.t = t
+        if msg is not None:
+            self.msg = msg
+        if y is not None:
+            self.y = y
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    @classmethod
+    def from_dict(cls, mapping: Dict[str, Any]) -> "TemporalData":
+        return cls(**mapping)
+
+    @property
+    def num_events(self) -> int:
+        for key in ("src", "dst", "t", "msg"):
+            if key in self._store and self._store[key] is not None:
+                return get_shape(self._store[key])[0]
+        return 0
+
+    def train_val_test_split(self, val_ratio: float = 0.15, test_ratio: float = 0.15):
+        r"""Splits the events chronologically into training, validation and test events (the
+        last ``val_ratio + test_ratio`` of the time span go to validation and test)."""
+        t = np.asarray(ops.convert_to_numpy(self.t))
+        val_time, test_time = np.quantile(t, [1.0 - val_ratio - test_ratio, 1.0 - test_ratio])
+        val_idx, test_idx = int((t <= val_time).sum()), int((t <= test_time).sum())
+        return self[:val_idx], self[val_idx:test_idx], self[test_idx:]
+
+    @property
+    def num_nodes(self) -> int:
+        nodes = []
+        if "src" in self._store and self._store.src is not None:
+            src_np = ops.convert_to_numpy(self._store.src)
+            if src_np.size > 0:
+                nodes.append(int(np.max(src_np)))
+        if "dst" in self._store and self._store.dst is not None:
+            dst_np = ops.convert_to_numpy(self._store.dst)
+            if dst_np.size > 0:
+                nodes.append(int(np.max(dst_np)))
+        return max(nodes) + 1 if len(nodes) > 0 else 0
+
+    def __len__(self) -> int:
+        return self.num_events
+
+    def __getitem__(self, idx: Any) -> Any:
+        if isinstance(idx, str):
+            return self._store[idx]
+        data = copy.copy(self)
+        num_events = self.num_events
+        for key, value in data._store.items():
+            if is_tensor_like(value) and get_shape(value)[0] == num_events:
+                data[key] = value[idx]
+        return data
+
+    def __setitem__(self, key: str, value: Any):
+        self._store[key] = value
+
+    def __delitem__(self, key: str):
+        if key in self._store:
+            del self._store[key]
+
+    def __getattr__(self, key: str) -> Any:
+        if "_store" not in self.__dict__:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{key}'")
+        try:
+            return getattr(self._store, key)
+        except AttributeError:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{key}'") from None
+
+    def __setattr__(self, key: str, value: Any):
+        if key == "_store":
+            self.__dict__["_store"] = value
+        elif "_store" in self.__dict__:
+            setattr(self._store, key, value)
+        else:
+            self.__dict__[key] = value
+
+    def __delattr__(self, key: str):
+        if key == "_store":
+            del self.__dict__["_store"]
+        elif "_store" in self.__dict__:
+            delattr(self._store, key)
+        else:
+            del self.__dict__[key]
+
+    def __copy__(self):
+        out = self.__class__.__new__(self.__class__)
+        for key, value in self.__dict__.items():
+            out.__dict__[key] = value
+        out.__dict__["_store"] = copy.copy(self._store)
+        out._store._parent = out
+        return out
+
+    def __deepcopy__(self, memo=None):
+        out = self.__class__.__new__(self.__class__)
+        for key, value in self.__dict__.items():
+            out.__dict__[key] = copy.deepcopy(value, memo)
+        out._store._parent = out
+        return out
+
+    @property
+    def stores(self) -> List[BaseStorage]:
+        return [self._store]
+
+    @property
+    def node_stores(self) -> List[Any]:
+        return [self._store]
+
+    @property
+    def edge_stores(self) -> List[Any]:
+        return [self._store]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return self._store.to_dict()
+
+    def to_namedtuple(self) -> NamedTuple:
+        fields = sorted(list(self.keys()))
+        import collections
+
+        TemporalTuple = collections.namedtuple("TemporalTuple", fields)
+        return TemporalTuple(**{f: self[f] for f in fields})
+
+    def __repr__(self) -> str:
+        cls = self.__class__.__name__
+        attrs = [size_repr(k, v) for k, v in self._store.items()]
+        return f"{cls}({', '.join(attrs)})"
+

k3_node/data/test_batch.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+from keras import ops
+
+from k3_node.data import Batch, Data, HeteroData
+
+
+def test_batch_homogeneous():
+    d1 = Data(
+        x=ops.convert_to_tensor([[1.0, 2.0], [3.0, 4.0]]),
+        edge_index=ops.convert_to_tensor([[0, 1], [1, 0]], dtype="int64"),
+        y=ops.convert_to_tensor([0]),
+    )
+    d2 = Data(
+        x=ops.convert_to_tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]),
+        edge_index=ops.convert_to_tensor([[0, 1, 2], [1, 2, 0]], dtype="int64"),
+        y=ops.convert_to_tensor([1]),
+    )
+
+    batch = Batch.from_data_list([d1, d2])
+    assert batch.num_graphs == 2
+    assert batch.num_nodes == 5
+    assert batch.num_edges == 5
+
+    # Check batch vector
+    batch_vec = ops.convert_to_numpy(batch.batch)
+    assert np.array_equal(batch_vec, [0, 0, 1, 1, 1])
+
+    # Check ptr vector
+    ptr_vec = ops.convert_to_numpy(batch.ptr)
+    assert np.array_equal(ptr_vec, [0, 2, 5])
+
+    # Check offset edge_index
+    ei = ops.convert_to_numpy(batch.edge_index)
+    assert np.array_equal(ei[:, :2], [[0, 1], [1, 0]])
+    assert np.array_equal(ei[:, 2:], [[2, 3, 4], [3, 4, 2]])
+
+    # Separate back
+    rec1 = batch[0]
+    rec2 = batch[1]
+    assert rec1.num_nodes == 2
+    assert rec2.num_nodes == 3
+    assert np.allclose(ops.convert_to_numpy(rec1.x), ops.convert_to_numpy(d1.x))
+    assert np.allclose(ops.convert_to_numpy(rec2.x), ops.convert_to_numpy(d2.x))
+
+    data_list = batch.to_data_list()
+    assert len(data_list) == 2
+
+
+def test_batch_heterogeneous():
+    h1 = HeteroData()
+    h1["v"].x = ops.convert_to_tensor([[1.0], [2.0]])
+    h1["v", "e", "v"].edge_index = ops.convert_to_tensor([[0], [1]], dtype="int64")
+
+    h2 = HeteroData()
+    h2["v"].x = ops.convert_to_tensor([[3.0], [4.0], [5.0]])
+    h2["v", "e", "v"].edge_index = ops.convert_to_tensor([[0, 1], [1, 2]], dtype="int64")
+
+    batch = Batch.from_data_list([h1, h2])
+    assert batch.num_graphs == 2
+    assert batch["v"].num_nodes == 5
+    assert batch["v", "e", "v"].num_edges == 3
+
+    rec1 = batch[0]
+    assert rec1["v"].num_nodes == 2
+    assert rec1["v", "e", "v"].num_edges == 1
+

k3_node/data/test_data.py

@@ -0,0 +1,68 @@
+import numpy as np
+import pytest
+from keras import ops
+
+from k3_node.data import Data
+
+
+def test_data_basic():
+    x = ops.convert_to_tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype="int64")
+    y = ops.convert_to_tensor([0, 1, 0], dtype="int64")
+
+    data = Data(x=x, edge_index=edge_index, y=y)
+    assert data.num_nodes == 3
+    assert data.num_edges == 4
+    assert data.num_features == 2
+    assert data.num_classes == 2
+    assert data.is_undirected()
+    assert not data.is_directed()
+    assert not data.has_self_loops()
+    assert not data.has_isolated_nodes()
+
+    # Attribute and dict access
+    assert "x" in data
+    assert "edge_index" in data
+    assert "edge_attr" not in data
+    assert len(data.keys()) == 3
+    assert ops.convert_to_numpy(data["x"]).shape == (3, 2)
+    assert ops.convert_to_numpy(data.x).shape == (3, 2)
+
+    # Clone
+    clone = data.clone()
+    assert clone.num_nodes == 3
+    assert clone.num_edges == 4
+
+    # Dict and namedtuple
+    d = data.to_dict()
+    assert "x" in d and "edge_index" in d and "y" in d
+    nt = data.to_namedtuple()
+    assert hasattr(nt, "x") and hasattr(nt, "edge_index")
+
+    # String repr
+    repr_str = str(data)
+    assert "Data(" in repr_str and "x=" in repr_str and "edge_index=" in repr_str
+
+
+def test_data_subgraph():
+    x = ops.convert_to_tensor([[10.0], [20.0], [30.0], [40.0]])
+    edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+    data = Data(x=x, edge_index=edge_index)
+
+    subset = ops.convert_to_tensor([0, 2], dtype="int64")
+    sub = data.subgraph(subset)
+    assert sub.num_nodes == 2
+    assert ops.convert_to_numpy(sub.x).tolist() == [[10.0], [30.0]]
+
+
+def test_data_to_heterogeneous():
+    x = ops.convert_to_tensor([[1.0], [2.0]])
+    edge_index = ops.convert_to_tensor([[0, 1], [1, 0]], dtype="int64")
+    data = Data(x=x, edge_index=edge_index)
+
+    hetero = data.to_heterogeneous(node_type="v", edge_type=("v", "e", "v"))
+    assert "v" in hetero.node_types
+    assert ("v", "e", "v") in hetero.edge_types
+    assert ops.convert_to_numpy(hetero["v"].x).shape == (2, 1)
+    assert ops.convert_to_numpy(hetero["v", "e", "v"].edge_index).shape == (2, 2)
+

k3_node/data/test_dataset_and_stores.py

@@ -0,0 +1,111 @@
+import os
+import tempfile
+import numpy as np
+import pytest
+from keras import ops
+
+from k3_node.data import (
+    Data,
+    EdgeAttr,
+    EdgeLayout,
+    FeatureStore,
+    GraphStore,
+    InMemoryDataset,
+    OnDiskDataset,
+    SQLiteDatabase,
+    TensorAttr,
+)
+
+
+class MyInMemoryDataset(InMemoryDataset):
+    def __init__(self, data_list, root=None):
+        super().__init__(root)
+        self._data_list = data_list
+
+
+def test_in_memory_dataset():
+    d1 = Data(x=ops.convert_to_tensor([[1.0]]), edge_index=ops.convert_to_tensor([[0], [0]], dtype="int64"))
+    d2 = Data(x=ops.convert_to_tensor([[2.0]]), edge_index=ops.convert_to_tensor([[0], [0]], dtype="int64"))
+
+    ds = MyInMemoryDataset([d1, d2])
+    assert len(ds) == 2
+    assert ops.convert_to_numpy(ds[0].x).item() == 1.0
+    assert ops.convert_to_numpy(ds[1].x).item() == 2.0
+
+    # Slicing
+    sub_ds = ds[1:]
+    assert len(sub_ds) == 1
+    assert ops.convert_to_numpy(sub_ds[0].x).item() == 2.0
+
+
+def test_sqlite_database_and_on_disk_dataset():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        db_path = os.path.join(tmp_dir, "test.db")
+        db = SQLiteDatabase(path=db_path)
+        d1 = Data(x=ops.convert_to_tensor([[1.0]]))
+        d2 = Data(x=ops.convert_to_tensor([[2.0]]))
+        db[0] = d1
+        db[1] = d2
+        assert len(db) == 2
+        rec1 = db[0]
+        assert ops.convert_to_numpy(rec1.x).item() == 1.0
+        db.close()
+
+        # Test OnDiskDataset
+        ds = OnDiskDataset(root=tmp_dir, backend="sqlite")
+        ds.append(d1)
+        ds.append(d2)
+        assert len(ds) == 2
+        rec = ds[0]
+        assert ops.convert_to_numpy(rec.x).item() == 1.0
+        ds.close()
+
+
+class SimpleFeatureStore(FeatureStore):
+    def _put_tensor(self, tensor, attr):
+        self._feat_dict[(attr.group_name, attr.attr_name)] = tensor
+        return True
+
+    def _get_tensor(self, attr):
+        return self._feat_dict.get((attr.group_name, attr.attr_name))
+
+    def _remove_tensor(self, attr):
+        return self._feat_dict.pop((attr.group_name, attr.attr_name), None) is not None
+
+
+class SimpleGraphStore(GraphStore):
+    def __init__(self):
+        self._store = {}
+
+    def _put_edge_index(self, edge_index, edge_attr):
+        self._store[(edge_attr.edge_type, edge_attr.layout)] = edge_index
+        return True
+
+    def _get_edge_index(self, edge_attr):
+        return self._store.get((edge_attr.edge_type, edge_attr.layout))
+
+    def _remove_edge_index(self, edge_attr):
+        return self._store.pop((edge_attr.edge_type, edge_attr.layout), None) is not None
+
+
+def test_feature_and_graph_store():
+    fs = SimpleFeatureStore()
+    x = ops.convert_to_tensor([[1.0, 2.0]])
+    fs.put_tensor(x, group_name="user", attr_name="feat")
+    ret = fs.get_tensor(group_name="user", attr_name="feat")
+    assert ops.convert_to_numpy(ret).shape == (1, 2)
+
+    gs = SimpleGraphStore()
+    ei = ops.convert_to_tensor([[0, 1], [1, 0]], dtype="int64")
+    gs.put_edge_index(ei, edge_type=("u", "follows", "u"), layout=EdgeLayout.COO)
+    ret_ei = gs.get_edge_index(edge_type=("u", "follows", "u"), layout=EdgeLayout.COO)
+    assert ops.convert_to_numpy(ret_ei).shape == (2, 2)
+
+
+
+def test_fractional_slicing():
+    from k3_node.datasets import FakeDataset
+
+    dataset = FakeDataset(num_graphs=10)
+    assert len(dataset[:0.9]) == 9 and len(dataset[0.9:]) == 1
+    assert len(dataset[0.2:0.5]) == 3

k3_node/data/test_hetero_data.py

@@ -0,0 +1,33 @@
+import numpy as np
+import pytest
+from keras import ops
+
+from k3_node.data import HeteroData
+
+
+def test_hetero_data_basic():
+    data = HeteroData()
+
+    data["paper"].x = ops.convert_to_tensor([[1.0, 2.0], [3.0, 4.0]])
+    data["author"].x = ops.convert_to_tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]])
+    data["author", "writes", "paper"].edge_index = ops.convert_to_tensor(
+        [[0, 1, 2], [0, 1, 1]], dtype="int64"
+    )
+
+    assert set(data.node_types) == {"paper", "author"}
+    assert set(data.edge_types) == {("author", "writes", "paper")}
+    assert data["paper"].num_nodes == 2
+    assert data["author"].num_nodes == 3
+    assert data["author", "writes", "paper"].num_edges == 3
+
+    meta = data.metadata()
+    assert len(meta[0]) == 2
+    assert len(meta[1]) == 1
+
+    # Homogeneous conversion
+    homo = data.to_homogeneous()
+    assert homo.num_nodes == 5
+    assert homo.num_edges == 3
+    assert ops.convert_to_numpy(homo.x).shape == (5, 2)
+    assert ops.convert_to_numpy(homo.edge_index).shape == (2, 3)
+

k3_node/data/test_temporal_and_hyper.py

@@ -0,0 +1,32 @@
+import numpy as np
+import pytest
+from keras import ops
+
+from k3_node.data import HypergraphData, TemporalData
+
+
+def test_temporal_data():
+    src = ops.convert_to_tensor([0, 1, 2, 3], dtype="int64")
+    dst = ops.convert_to_tensor([1, 2, 3, 4], dtype="int64")
+    t = ops.convert_to_tensor([10, 20, 30, 40], dtype="int64")
+    msg = ops.convert_to_tensor([[1.0], [2.0], [3.0], [4.0]])
+
+    data = TemporalData(src=src, dst=dst, t=t, msg=msg)
+    assert data.num_events == 4
+    assert data.num_nodes == 5
+    assert len(data) == 4
+
+    sub = data[:2]
+    assert sub.num_events == 2
+    assert ops.convert_to_numpy(sub.src).tolist() == [0, 1]
+
+
+def test_hypergraph_data():
+    x = ops.convert_to_tensor([[1.0], [2.0], [3.0], [4.0], [5.0]])
+    edge_index = ops.convert_to_tensor(
+        [[0, 1, 2, 1, 2, 3, 4], [0, 0, 0, 1, 1, 1, 1]], dtype="int64"
+    )
+    data = HypergraphData(x=x, edge_index=edge_index)
+    assert data.num_nodes == 5
+    assert data.num_edges == 2
+

k3_node/data/view.py

@@ -0,0 +1,43 @@
+from typing import Any, Iterator, List, Mapping, Tuple
+
+
+class MappingView:
+    def __init__(self, mapping: Mapping[str, Any], *args: str):
+        self._mapping = mapping
+        self._args = args
+
+    def _keys(self) -> List[str]:
+        if len(self._args) == 0:
+            return list(self._mapping.keys())
+        else:
+            return [arg for arg in self._args if arg in self._mapping]
+
+    def __len__(self) -> int:
+        return len(self._keys())
+
+    def __contains__(self, item: Any) -> bool:
+        return item in self._keys()
+
+    def __repr__(self) -> str:
+        mapping = {key: self._mapping[key] for key in self._keys()}
+        return f"{self.__class__.__name__}({mapping})"
+
+    __class_getitem__ = classmethod(type([]))
+
+
+class KeysView(MappingView):
+    def __iter__(self) -> Iterator[str]:
+        yield from self._keys()
+
+
+class ValuesView(MappingView):
+    def __iter__(self) -> Iterator[Any]:
+        for key in self._keys():
+            yield self._mapping[key]
+
+
+class ItemsView(MappingView):
+    def __iter__(self) -> Iterator[Tuple[str, Any]]:
+        for key in self._keys():
+            yield (key, self._mapping[key])
+

k3_node/datasets/__init__.py

@@ -0,0 +1,88 @@
+from .karate import KarateClub
+from .fake import FakeDataset, FakeHeteroDataset
+from .planetoid import Planetoid
+from .tu_dataset import TUDataset
+from .citation_full import CitationFull, CoraFull
+from .amazon import Amazon
+from .coauthor import Coauthor
+from .wikics import WikiCS
+from .webkb import WebKB
+from .actor import Actor
+from .polblogs import PolBlogs
+from .airports import Airports
+from .email_eu_core import EmailEUCore
+from .github import GitHub
+from .facebook import FacebookPagePage
+from .lastfm_asia import LastFMAsia
+from .twitch import Twitch
+from .ba_shapes import BAShapes
+from .ba2motif_dataset import BA2MotifDataset
+from .sbm_dataset import StochasticBlockModelDataset, RandomPartitionGraphDataset
+from .explainer_dataset import ExplainerDataset
+from .entities import Entities
+from .word_net import WordNet18, WordNet18RR
+from .freebase import FB15k_237
+from .dblp import DBLP
+from .imdb import IMDB
+from .qm7 import QM7b
+from .molecule_net import MoleculeNet
+from .ppi import PPI
+from .reddit import Reddit
+from .digits import Digits
+from .seal import SEALDataset
+from .bitcoin_otc import BitcoinOTC
+from .geometric_shapes import GeometricShapes
+from .shape_scenes import ShapeScenes
+from .mesh_correspondence import MeshCorrespondence
+from .movielens import MovieLens100K
+from .qm9 import QM9
+from .jodie import JODIEDataset
+from .icews import ICEWS18
+
+__all__ = [
+    "Digits",
+    "SEALDataset",
+    "BitcoinOTC",
+    "GeometricShapes",
+    "ShapeScenes",
+    "MeshCorrespondence",
+    "MovieLens100K",
+    "QM9",
+    "JODIEDataset",
+    "ICEWS18",
+    "KarateClub",
+    "FakeDataset",
+    "FakeHeteroDataset",
+    "Planetoid",
+    "TUDataset",
+    "CitationFull",
+    "CoraFull",
+    "Amazon",
+    "Coauthor",
+    "WikiCS",
+    "WebKB",
+    "Actor",
+    "PolBlogs",
+    "Airports",
+    "EmailEUCore",
+    "GitHub",
+    "FacebookPagePage",
+    "LastFMAsia",
+    "Twitch",
+    "BAShapes",
+    "BA2MotifDataset",
+    "StochasticBlockModelDataset",
+    "RandomPartitionGraphDataset",
+    "ExplainerDataset",
+    "Entities",
+    "WordNet18",
+    "WordNet18RR",
+    "FB15k_237",
+    "DBLP",
+    "IMDB",
+    "QM7b",
+    "MoleculeNet",
+    "PPI",
+    "Reddit",
+]
+