k3-node 1.0.0__py3-none-any.whl

k3_node/io/npz.py

@@ -0,0 +1,45 @@
+from typing import Any, Dict
+import numpy as np
+import scipy.sparse as sp
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.layers.conv.utils import remove_self_loops
+from k3_node.transforms.utils import to_undirected as to_undirected_fn
+
+
+def read_npz(path: str, to_undirected: bool = True) -> Data:
+    """Reads a `.npz` graph archive (e.g. Amazon, Coauthor) and returns a `Data` object."""
+    with np.load(path) as f:
+        return parse_npz(f, to_undirected=to_undirected)
+
+
+def parse_npz(f: Dict[str, Any], to_undirected: bool = True) -> Data:
+    x_sp = sp.csr_matrix(
+        (f["attr_data"], f["attr_indices"], f["attr_indptr"]),
+        shape=tuple(f["attr_shape"]),
+    )
+    x = np.array(x_sp.todense(), dtype=np.float32)
+    x[x > 0] = 1.0
+
+    adj = sp.csr_matrix(
+        (f["adj_data"], f["adj_indices"], f["adj_indptr"]),
+        shape=tuple(f["adj_shape"]),
+    ).tocoo()
+
+    row = np.array(adj.row, dtype=np.int64)
+    col = np.array(adj.col, dtype=np.int64)
+    edge_index = np.stack([row, col], axis=0)
+
+    edge_index, _ = remove_self_loops(edge_index)
+    if to_undirected:
+        edge_index = to_undirected_fn(edge_index, num_nodes=x.shape[0])
+
+    y = np.array(f["labels"], dtype=np.int64)
+
+    return Data(
+        x=ops.convert_to_tensor(x, dtype="float32"),
+        edge_index=ops.convert_to_tensor(edge_index, dtype="int64"),
+        y=ops.convert_to_tensor(y, dtype="int64"),
+    )
+

k3_node/io/off.py

@@ -0,0 +1,29 @@
+from typing import List
+
+import numpy as np
+
+from k3_node.data import Data
+
+
+def parse_off(src: List[str]) -> Data:
+    r"""Parses the lines of an OFF (Object File Format) mesh into ``Data(pos, face)``;
+    quadrilaterals are split into two triangles."""
+    if src[0] == 'OFF':
+        src = src[1:]
+    else:  # some files lack the line break after "OFF"
+        src[0] = src[0][3:]
+    num_nodes, num_faces = (int(item) for item in src[0].split()[:2])
+    pos = np.array([[float(v) for v in line.split()[:3]] for line in src[1:1 + num_nodes]], dtype=np.float32)
+    faces = [[int(v) for v in line.strip().split()] for line in src[1 + num_nodes:1 + num_nodes + num_faces]]
+    tri = [f[1:4] for f in faces if f[0] == 3]
+    for f in faces:
+        if f[0] == 4:
+            tri += [[f[1], f[2], f[3]], [f[1], f[3], f[4]]]
+    face = np.array(tri, dtype=np.int64).reshape(-1, 3).T
+    return Data(pos=pos, face=face)
+
+
+def read_off(path: str) -> Data:
+    r"""Reads an OFF (Object File Format) mesh file into ``Data(pos, face)``."""
+    with open(path) as f:
+        return parse_off(f.read().split('\n')[:-1])

k3_node/io/planetoid.py

@@ -0,0 +1,98 @@
+import os.path as osp
+import pickle
+import warnings
+from typing import Dict, List, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.io.txt_array import read_txt_array
+from k3_node.layers.conv.utils import remove_self_loops
+from k3_node.utils.graph import coalesce
+
+
+def index_to_mask(index, size: int):
+    """Converts 1D index array into a boolean mask."""
+    mask = np.zeros(size, dtype=bool)
+    idx_np = ops.convert_to_numpy(index)
+    mask[idx_np] = True
+    return ops.convert_to_tensor(mask, dtype="bool")
+
+
+def edge_index_from_dict(graph_dict: Dict[int, List[int]], num_nodes: Optional[int] = None):
+    rows: List[int] = []
+    cols: List[int] = []
+    for key, val_list in graph_dict.items():
+        for val in val_list:
+            rows.append(key)
+            cols.append(val)
+    if len(rows) == 0:
+        return ops.zeros((2, 0), dtype="int64")
+
+    edge_index = np.array([rows, cols], dtype=np.int64)
+    edge_index, _ = remove_self_loops(edge_index)
+    edge_index, _ = coalesce(edge_index, num_nodes=num_nodes, sort_by_row=False)
+    return ops.convert_to_tensor(edge_index, dtype="int64")
+
+
+def read_file(folder: str, prefix: str, name: str):
+    path = osp.join(folder, f"ind.{prefix.lower()}.{name}")
+    if name == "test.index":
+        return read_txt_array(path, dtype="int64")
+
+    with open(path, "rb") as f:
+        warnings.filterwarnings("ignore", ".*`scipy.sparse.csr` name.*")
+        out = pickle.load(f, encoding="latin1")
+
+    if name == "graph":
+        return out
+
+    if hasattr(out, "todense"):
+        out = out.todense()
+    return np.array(out, dtype=np.float32)
+
+
+def read_planetoid_data(folder: str, prefix: str) -> Data:
+    """Reads planetoid citation graph files and returns a `Data` object."""
+    names = ["x", "tx", "allx", "y", "ty", "ally", "graph", "test.index"]
+    items = [read_file(folder, prefix, name) for name in names]
+    x, tx, allx, y, ty, ally, graph, test_index = items
+
+    test_index_np = ops.convert_to_numpy(test_index).astype(np.int64)
+    sorted_test_index = np.sort(test_index_np)
+
+    train_index = np.arange(y.shape[0], dtype=np.int64)
+    val_index = np.arange(y.shape[0], y.shape[0] + 500, dtype=np.int64)
+
+    if prefix.lower() == "citeseer":
+        len_test_indices = int(np.max(test_index_np) - np.min(test_index_np)) + 1
+        tx_ext = np.zeros((len_test_indices, tx.shape[1]), dtype=tx.dtype)
+        tx_ext[sorted_test_index - np.min(test_index_np), :] = tx
+        ty_ext = np.zeros((len_test_indices, ty.shape[1]), dtype=ty.dtype)
+        ty_ext[sorted_test_index - np.min(test_index_np), :] = ty
+        tx, ty = tx_ext, ty_ext
+
+    x = np.concatenate([allx, tx], axis=0)
+    x[test_index_np] = x[sorted_test_index]
+
+    y_cat = np.concatenate([ally, ty], axis=0)
+    y = np.argmax(y_cat, axis=1).astype(np.int64)
+    y[test_index_np] = y[sorted_test_index]
+
+    num_nodes = y.shape[0]
+    train_mask = index_to_mask(train_index, size=num_nodes)
+    val_mask = index_to_mask(val_index, size=num_nodes)
+    test_mask = index_to_mask(test_index_np, size=num_nodes)
+
+    edge_index = edge_index_from_dict(graph, num_nodes=num_nodes)
+
+    data = Data(
+        x=ops.convert_to_tensor(x, dtype="float32"),
+        edge_index=edge_index,
+        y=ops.convert_to_tensor(y, dtype="int64"),
+    )
+    data.train_mask = train_mask
+    data.val_mask = val_mask
+    data.test_mask = test_mask
+
+    return data

k3_node/io/tu.py

@@ -0,0 +1,137 @@
+import glob
+import os.path as osp
+from typing import Any, Dict, List, Optional, Tuple
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data
+from k3_node.io.txt_array import read_txt_array
+from k3_node.layers.conv.utils import remove_self_loops
+from k3_node.utils.graph import coalesce
+
+
+def np_one_hot(arr: np.ndarray) -> np.ndarray:
+    """One-hot encodes a 1D integer array."""
+    arr = arr.astype(np.int64)
+    if arr.size == 0:
+        return np.zeros((0, 0), dtype=np.float32)
+    min_val = np.min(arr)
+    arr = arr - min_val
+    max_val = np.max(arr)
+    out = np.zeros((arr.shape[0], max_val + 1), dtype=np.float32)
+    out[np.arange(arr.shape[0]), arr] = 1.0
+    return out
+
+
+def read_tu_data(folder: str, prefix: str) -> Tuple[Data, Dict[str, Any], Dict[str, int]]:
+    """Reads a TU benchmark dataset from disk and returns `(data, slices, sizes)`."""
+    files = sorted(glob.glob(osp.join(folder, f"{prefix}_*.txt")))
+    names = [osp.basename(f)[len(prefix) + 1 : -4] for f in files]
+
+    def read_file(name: str, dtype: str = "float32"):
+        path = osp.join(folder, f"{prefix}_{name}.txt")
+        return ops.convert_to_numpy(read_txt_array(path, sep=",", dtype=dtype))
+
+    # Adjacency: 1-indexed in raw TU files
+    edge_index_np = read_file("A", dtype="int64")
+    if edge_index_np.ndim == 1:
+        edge_index_np = edge_index_np.reshape(-1, 2)
+    edge_index_np = edge_index_np.T - 1  # 0-indexed (2, E)
+
+    # Graph indicator: which graph each node belongs to (1-indexed)
+    batch_np = read_file("graph_indicator", dtype="int64") - 1  # 0-indexed (N,)
+
+    node_attribute = np.empty((batch_np.shape[0], 0), dtype=np.float32)
+    if "node_attributes" in names:
+        node_attribute = read_file("node_attributes", dtype="float32")
+        if node_attribute.ndim == 1:
+            node_attribute = node_attribute[:, None]
+
+    node_label = np.empty((batch_np.shape[0], 0), dtype=np.float32)
+    if "node_labels" in names:
+        node_label_raw = read_file("node_labels", dtype="int64")
+        if node_label_raw.ndim == 1:
+            node_label_raw = node_label_raw[:, None]
+        encoded = [np_one_hot(node_label_raw[:, i]) for i in range(node_label_raw.shape[1])]
+        node_label = np.concatenate(encoded, axis=-1)
+
+    edge_attribute = np.empty((edge_index_np.shape[1], 0), dtype=np.float32)
+    if "edge_attributes" in names:
+        edge_attribute = read_file("edge_attributes", dtype="float32")
+        if edge_attribute.ndim == 1:
+            edge_attribute = edge_attribute[:, None]
+
+    edge_label = np.empty((edge_index_np.shape[1], 0), dtype=np.float32)
+    if "edge_labels" in names:
+        edge_label_raw = read_file("edge_labels", dtype="int64")
+        if edge_label_raw.ndim == 1:
+            edge_label_raw = edge_label_raw[:, None]
+        encoded = [np_one_hot(edge_label_raw[:, i]) for i in range(edge_label_raw.shape[1])]
+        edge_label = np.concatenate(encoded, axis=-1)
+
+    # Combine attributes and one-hot labels
+    x_list = [arr for arr in [node_attribute, node_label] if arr.shape[1] > 0]
+    x_np = np.concatenate(x_list, axis=-1) if len(x_list) > 0 else None
+
+    edge_attr_list = [arr for arr in [edge_attribute, edge_label] if arr.shape[1] > 0]
+    edge_attr_np = np.concatenate(edge_attr_list, axis=-1) if len(edge_attr_list) > 0 else None
+
+    y_np = None
+    if "graph_attributes" in names:
+        y_np = read_file("graph_attributes", dtype="float32")
+    elif "graph_labels" in names:
+        y_raw = read_file("graph_labels", dtype="int64")
+        _, y_inv = np.unique(y_raw, return_inverse=True)
+        y_np = y_inv.astype(np.int64)
+
+    num_nodes = x_np.shape[0] if x_np is not None else int(np.max(edge_index_np)) + 1
+    edge_index_np, edge_attr_np = remove_self_loops(edge_index_np, edge_attr_np)
+    edge_index_t, edge_attr_t = coalesce(edge_index_np, edge_attr_np, num_nodes=num_nodes)
+    edge_index_np = ops.convert_to_numpy(edge_index_t)
+    edge_attr_np = ops.convert_to_numpy(edge_attr_t) if edge_attr_t is not None else None
+
+    # Convert to tensors
+    edge_index = ops.convert_to_tensor(edge_index_np, dtype="int64")
+    x = ops.convert_to_tensor(x_np, dtype="float32") if x_np is not None else None
+    edge_attr = ops.convert_to_tensor(edge_attr_np, dtype="float32") if edge_attr_np is not None else None
+    y = ops.convert_to_tensor(y_np, dtype="float32" if "graph_attributes" in names else "int64") if y_np is not None else None
+
+    # Compute graph slices
+    num_graphs = int(np.max(batch_np)) + 1
+    node_counts = np.bincount(batch_np, minlength=num_graphs)
+    node_slice = np.pad(np.cumsum(node_counts), (1, 0))
+
+    row = edge_index_np[0]
+    edge_batch = batch_np[row]
+    edge_counts = np.bincount(edge_batch, minlength=num_graphs)
+    edge_slice = np.pad(np.cumsum(edge_counts), (1, 0))
+
+    # Shift edge indices so each graph starts at 0
+    shift = node_slice[edge_batch]
+    edge_index_shifted = edge_index_np - shift[None, :]
+    edge_index = ops.convert_to_tensor(edge_index_shifted, dtype="int64")
+
+    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
+    slices: Dict[str, Any] = {
+        "edge_index": ops.convert_to_tensor(edge_slice, dtype="int64"),
+    }
+    if x is not None:
+        slices["x"] = ops.convert_to_tensor(node_slice, dtype="int64")
+    else:
+        data.num_nodes = int(batch_np.shape[0])
+    if edge_attr is not None:
+        slices["edge_attr"] = ops.convert_to_tensor(edge_slice, dtype="int64")
+    if y is not None:
+        if y_np.shape[0] == batch_np.shape[0]:
+            slices["y"] = ops.convert_to_tensor(node_slice, dtype="int64")
+        else:
+            slices["y"] = ops.convert_to_tensor(np.arange(num_graphs + 1, dtype=np.int64), dtype="int64")
+
+    sizes = {
+        "num_node_attributes": node_attribute.shape[-1],
+        "num_node_labels": node_label.shape[-1],
+        "num_edge_attributes": edge_attribute.shape[-1],
+        "num_edge_labels": edge_label.shape[-1],
+    }
+
+    return data, slices, sizes

k3_node/io/txt_array.py

@@ -0,0 +1,58 @@
+from typing import List, Optional, Union
+import numpy as np
+from keras import ops
+
+
+def parse_txt_array(
+    src: List[str],
+    sep: Optional[str] = None,
+    start: int = 0,
+    end: Optional[int] = None,
+    dtype: Optional[str] = None,
+):
+    """Parses a list of string rows into a tensor."""
+    lines = [line.strip() for line in src if line.strip()]
+    if len(lines) == 0:
+        return ops.zeros((0,), dtype=dtype or "float32")
+
+    split_lines = []
+    is_float = False
+    for line in lines:
+        parts = line.split(sep)[start:end]
+        row = []
+        for x in parts:
+            x_str = x.strip()
+            if not x_str:
+                continue
+            if "." in x_str or "e" in x_str.lower():
+                is_float = True
+                row.append(float(x_str))
+            else:
+                try:
+                    row.append(int(x_str))
+                except ValueError:
+                    is_float = True
+                    row.append(float(x_str))
+        split_lines.append(row)
+
+    if dtype is None:
+        dtype = "float32" if is_float else "int64"
+
+    arr = np.array(split_lines, dtype=dtype)
+    if arr.ndim > 1 and arr.shape[1] == 1:
+        arr = np.squeeze(arr, axis=1)
+    return ops.convert_to_tensor(arr, dtype=dtype)
+
+
+def read_txt_array(
+    path: str,
+    sep: Optional[str] = None,
+    start: int = 0,
+    end: Optional[int] = None,
+    dtype: Optional[str] = None,
+):
+    """Reads a text array from a file and returns a tensor."""
+    with open(path, "r", encoding="utf-8") as f:
+        src = f.read().split("\n")
+    return parse_txt_array(src, sep, start, end, dtype)
+

k3_node/layers/__init__.py

@@ -0,0 +1,14 @@
+"""
+`k3_node.layers` module provides access to various layers for building 
+graph neural networks.
+"""
+
+from .conv import *
+from .norm import *
+from .aggr import *
+from .attention import *
+from .dense import *
+from .pool import *
+from .unpool import *
+from .kge import *
+from .functional import *

k3_node/layers/aggr/__init__.py

@@ -0,0 +1,70 @@
+r"""Aggregation operators for graph neural networks."""
+
+from .base import Aggregation, from_dense_batch, ptr2index, to_dense_adj, to_dense_batch
+from .basic import (
+    MaxAggregation,
+    MeanAggregation,
+    MinAggregation,
+    MulAggregation,
+    PowerMeanAggregation,
+    SoftmaxAggregation,
+    StdAggregation,
+    SumAggregation,
+    VarAggregation,
+)
+from .quantile import MedianAggregation, QuantileAggregation
+from .attention import AttentionalAggregation
+from .set2set import Set2Set
+from .scaler import DegreeScalerAggregation
+from .sort import SortAggregation
+from .multi import MultiAggregation
+from .deep_sets import DeepSetsAggregation
+from .mlp import MLPAggregation
+from .lstm import LSTMAggregation
+from .gru import GRUAggregation
+from .set_transformer import SetTransformerAggregation
+from .gmt import GraphMultisetTransformer
+from .variance_preserving import VariancePreservingAggregation
+from .patch_transformer import PatchTransformerAggregation
+from .lcm import LCMAggregation
+from .equilibrium import EquilibriumAggregation
+from .fused import FusedAggregation
+from .resolver import aggregation_resolver
+
+__all__ = [
+    "Aggregation",
+    "ptr2index",
+    "to_dense_batch",
+    "from_dense_batch",
+    "to_dense_adj",
+    "SumAggregation",
+    "MeanAggregation",
+    "MaxAggregation",
+    "MinAggregation",
+    "MulAggregation",
+    "VarAggregation",
+    "StdAggregation",
+    "SoftmaxAggregation",
+    "PowerMeanAggregation",
+    "QuantileAggregation",
+    "MedianAggregation",
+    "AttentionalAggregation",
+    "Set2Set",
+    "DegreeScalerAggregation",
+    "SortAggregation",
+    "MultiAggregation",
+    "DeepSetsAggregation",
+    "MLPAggregation",
+    "LSTMAggregation",
+    "GRUAggregation",
+    "SetTransformerAggregation",
+    "GraphMultisetTransformer",
+    "VariancePreservingAggregation",
+    "PatchTransformerAggregation",
+    "LCMAggregation",
+    "EquilibriumAggregation",
+    "FusedAggregation",
+    "aggregation_resolver",
+]
+
+classes = __all__

k3_node/layers/aggr/attention.py

@@ -0,0 +1,77 @@
+from typing import Optional
+from keras import ops
+
+from .base import Aggregation
+from k3_node.ops.segment import segment_max, segment_sum
+
+
+class AttentionalAggregation(Aggregation):
+    r"""The soft attention aggregation layer from the `"Graph Matching Networks
+    for Learning the Similarity of Graph Structured Objects"
+    <https://arxiv.org/abs/1904.12787>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import AttentionalAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = AttentionalAggregation(gate_nn=keras.layers.Dense(1), nn=keras.layers.Dense(16))
+        out = aggr(x, index=index, dim_size=2)  # attention-weighted sum per set
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        gate_nn,
+        nn: Optional[any] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.gate_nn = gate_nn
+        self.nn = nn
+
+    def reset_parameters(self):
+        if hasattr(self.gate_nn, "reset_parameters"):
+            self.gate_nn.reset_parameters()
+        if self.nn is not None and hasattr(self.nn, "reset_parameters"):
+            self.nn.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        gate = self.gate_nn(x)
+        if self.nn is not None:
+            x = self.nn(x)
+
+        if ptr is not None and index is None:
+            from .base import ptr2index
+            index = ptr2index(ptr)
+
+        index = ops.cast(index, dtype="int32")
+        if dim_size is None:  # a tensor while tracing; don't test its truth value
+            dim_size = int(ops.max(index)) + 1 if ops.shape(index)[0] > 0 else 0
+
+        # Graph-wise softmax over groups
+        max_val = segment_max(gate, index, num_segments=dim_size)
+        max_exp = ops.take(max_val, index, axis=0)
+        exp_gate = ops.exp(gate - max_exp)
+        sum_exp = segment_sum(exp_gate, index, num_segments=dim_size)
+        sum_exp_exp = ops.take(sum_exp, index, axis=0)
+        alpha = exp_gate / ops.maximum(sum_exp_exp, 1e-12)
+
+        return self.reduce(alpha * x, index, ptr, dim_size, dim, reduce="sum")
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(gate_nn={self.gate_nn}, nn={self.nn})"
+