k3-node 1.0.0__py3-none-any.whl

k3_node/layers/pool/pool.py

@@ -0,0 +1,119 @@
+from typing import Optional, Tuple
+from keras import ops
+import numpy as np
+from k3_node.ops.segment import segment_max, segment_sum
+from k3_node.ops.host import to_numpy
+
+
+def pool_edge(
+    cluster,
+    edge_index,
+    edge_attr: Optional[any] = None,
+    reduce: str = "sum",
+) -> Tuple[any, Optional[any]]:
+    r"""Pools edge indices and attributes based on cluster assignments.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import pool_edge
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+        cluster = np.repeat(np.arange(5), 2)  # merge nodes pairwise into 5 clusters
+
+        edge_attr = np.random.rand(30, 3).astype("float32")
+        edge_index_pool, edge_attr_pool = pool_edge(cluster, edge_index, edge_attr)  # coarsened, deduplicated edges
+        print(edge_index_pool.shape[0], edge_attr_pool.shape[1])  # 2 3
+        ```
+    """
+    cluster_np = to_numpy(cluster)
+    edge_index_np = to_numpy(edge_index)
+
+    row = cluster_np[edge_index_np[0]]
+    col = cluster_np[edge_index_np[1]]
+
+    # Remove self-loops
+    non_loop = row != col
+    row = row[non_loop]
+    col = col[non_loop]
+
+    if len(row) == 0:
+        empty_ei = ops.zeros((2, 0), dtype=edge_index.dtype)
+        empty_ea = None if edge_attr is None else ops.zeros((0, *ops.shape(edge_attr)[1:]), dtype=edge_attr.dtype)
+        return empty_ei, empty_ea
+
+    edges = np.stack([row, col], axis=0)
+    # Coalesce duplicate edges
+    unique_edges, inv = np.unique(edges, axis=1, return_inverse=True)
+
+    out_edge_index = ops.convert_to_tensor(unique_edges, dtype=edge_index.dtype)
+
+    out_edge_attr = None
+    if edge_attr is not None:
+        ea_np = to_numpy(edge_attr)[non_loop]
+        num_unique = unique_edges.shape[1]
+        ea_tensor = ops.convert_to_tensor(ea_np, dtype=edge_attr.dtype)
+        inv_tensor = ops.convert_to_tensor(inv, dtype="int32")
+        if reduce == "sum":
+            out_edge_attr = segment_sum(ea_tensor, inv_tensor, num_segments=num_unique)
+        elif reduce == "mean":
+            sum_ea = segment_sum(ea_tensor, inv_tensor, num_segments=num_unique)
+            count = segment_sum(ops.ones_like(ea_tensor), inv_tensor, num_segments=num_unique)
+            out_edge_attr = sum_ea / ops.maximum(count, 1.0)
+        elif reduce == "max":
+            out_edge_attr = segment_max(ea_tensor, inv_tensor, num_segments=num_unique)
+
+    return out_edge_index, out_edge_attr
+
+
+def pool_batch(perm, batch):
+    r"""Pools batch vector given representative indices `perm`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import pool_batch
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        perm = np.array([0, 3, 6, 9])  # nodes kept after pooling
+        print(tuple(pool_batch(perm, batch).shape))  # (4,): graph id of each kept node
+        ```
+    """
+    return ops.take(batch, perm, axis=0)
+
+
+def pool_pos(cluster, pos):
+    r"""Pools node positions by computing average coordinates within each cluster.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import pool_pos
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+        cluster = np.repeat(np.arange(5), 2)  # merge nodes pairwise into 5 clusters
+
+        print(tuple(pool_pos(cluster, pos).shape))  # (5, 3): mean position of every cluster
+        ```
+    """
+    cluster = ops.cast(cluster, dtype="int32")
+    num_clusters = int(to_numpy(cluster).max()) + 1 if ops.shape(cluster)[0] > 0 else 0
+    sum_pos = segment_sum(pos, cluster, num_segments=num_clusters)
+    count = segment_sum(ops.ones_like(pos), cluster, num_segments=num_clusters)
+    return sum_pos / ops.maximum(count, 1.0)
+
+
+def as_mutable_graph(data):
+    """Turns a (read-only) loader batch into a ``Data`` object that pooling can update in place.
+
+    ``ptr`` is dropped because it no longer matches the nodes once they are pooled.
+    """
+    if hasattr(data, "_asdict"):
+        from k3_node.data import Data
+
+        return Data(**{k: v for k, v in data._asdict().items() if k != "ptr" and v is not None})
+    return data

k3_node/layers/pool/sag_pool.py

@@ -0,0 +1,174 @@
+from typing import Callable, Optional, Tuple, Union
+from keras import layers, ops
+
+from .connect.filter_edges import FilterEdges
+from .select.topk import SelectTopK
+from k3_node.ops.segment import segment_max, segment_sum
+
+
+class GraphConv(layers.Layer):
+    r"""Basic GraphConv layer for projection scoring in pooling layers."""
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        aggr: str = "add",
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.aggr = aggr
+
+        self.lin_rel = layers.Dense(out_channels, use_bias=bias, name="lin_rel")
+        self.lin_root = layers.Dense(out_channels, use_bias=False, name="lin_root")
+
+    def reset_parameters(self):
+        pass
+
+    def build(self, input_shape=None):
+        if hasattr(self.lin_rel, "built") and not self.lin_rel.built:
+            self.lin_rel.build((None, self.in_channels))
+        if hasattr(self.lin_root, "built") and not self.lin_root.built:
+            self.lin_root.build((None, self.in_channels))
+        self.built = True
+
+    def call(self, x, edge_index, edge_weight: Optional[any] = None):
+        row = ops.cast(edge_index[0], dtype="int32")
+        col = ops.cast(edge_index[1], dtype="int32")
+        num_nodes = ops.shape(x)[0]
+
+        msg = ops.take(x, row, axis=0)
+        if edge_weight is not None:
+            msg = msg * ops.reshape(edge_weight, (-1, 1))
+
+        if self.aggr == "add":
+            aggr_out = segment_sum(msg, col, num_segments=num_nodes)
+        elif self.aggr == "mean":
+            sum_val = segment_sum(msg, col, num_segments=num_nodes)
+            count = segment_sum(ops.ones_like(msg), col, num_segments=num_nodes)
+            aggr_out = sum_val / ops.maximum(count, 1.0)
+        elif self.aggr == "max":
+            aggr_out = segment_max(msg, col, num_segments=num_nodes)
+        else:
+            aggr_out = segment_sum(msg, col, num_segments=num_nodes)
+
+        return self.lin_rel(aggr_out) + self.lin_root(x)
+
+
+class SAGPooling(layers.Layer):
+    r"""The self-attention pooling operator from the `"Self-Attention Graph
+    Pooling" <https://arxiv.org/abs/1904.08082>`_ and `"Understanding
+    Attention and Generalization in Graph Neural Networks"
+    <https://arxiv.org/abs/1905.02850>`_ papers.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SAGPooling
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        layer = SAGPooling(in_channels=8, ratio=0.5)  # keep half of the nodes of each graph
+        out = layer(x, edge_index, batch=batch)
+        x_pool, edge_index_pool, edge_attr_pool, batch_pool = out[0], out[1], out[2], out[3]
+        print(tuple(x_pool.shape))  # (6, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        ratio: Union[float, int] = 0.5,
+        GNN: Optional[any] = None,
+        min_score: Optional[float] = None,
+        multiplier: float = 1.0,
+        nonlinearity: Union[str, Callable] = "tanh",
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.ratio = ratio
+        self.min_score = min_score
+        self.multiplier = multiplier
+        self.nonlinearity = nonlinearity
+
+        if GNN is None:
+            self.gnn = GraphConv(in_channels, 1, **kwargs)
+        elif callable(GNN):
+            self.gnn = GNN(in_channels, 1, **kwargs)
+        else:
+            self.gnn = GNN
+
+        self.select = SelectTopK(1, ratio, min_score, nonlinearity)
+        self.connect = FilterEdges()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        if hasattr(self.gnn, "reset_parameters"):
+            self.gnn.reset_parameters()
+        self.select.reset_parameters()
+
+    def build(self, input_shape=None):
+        if hasattr(self.gnn, "built") and not self.gnn.built:
+            self.gnn.build(input_shape)
+        if hasattr(self.select, "built") and not self.select.built:
+            self.select.build(None)
+        if hasattr(self.connect, "built") and not self.connect.built:
+            self.connect.build(None)
+        self.built = True
+
+    def call(
+        self,
+        x,
+        edge_index,
+        edge_attr: Optional[any] = None,
+        batch: Optional[any] = None,
+        attn: Optional[any] = None,
+    ) -> Tuple[any, any, Optional[any], Optional[any], any, any]:
+        r"""Forward pass."""
+        num_nodes = ops.shape(x)[0]
+        if batch is None:
+            batch = ops.zeros((num_nodes,), dtype="int32")
+
+        if attn is None:
+            attn = x
+        if len(ops.shape(attn)) == 1:
+            attn = ops.expand_dims(attn, axis=-1)
+
+        attn = self.gnn(attn, edge_index)
+
+        select_out = self.select(attn, batch)
+
+        perm = select_out.node_index
+        score = select_out.weight
+
+        x_pooled = ops.take(x, perm, axis=0) * ops.expand_dims(score, axis=-1)
+        if self.multiplier != 1.0:
+            x_pooled = x_pooled * self.multiplier
+
+        connect_out = self.connect(select_out, edge_index, edge_attr, batch)
+
+        return (
+            x_pooled,
+            connect_out.edge_index,
+            connect_out.edge_attr,
+            connect_out.batch,
+            perm,
+            score,
+        )
+
+    def __repr__(self) -> str:
+        if self.min_score is None:
+            ratio = f"ratio={self.ratio}"
+        else:
+            ratio = f"min_score={self.min_score}"
+        gnn_name = self.gnn.__class__.__name__
+        return (
+            f"{self.__class__.__name__}({gnn_name}, {self.in_channels}, "
+            f"{ratio}, multiplier={self.multiplier})"
+        )
+

k3_node/layers/pool/select/__init__.py

@@ -0,0 +1,10 @@
+from .base import Select, SelectOutput
+from .topk import SelectTopK, topk
+
+__all__ = [
+    "SelectOutput",
+    "Select",
+    "topk",
+    "SelectTopK",
+]
+

k3_node/layers/pool/select/base.py

@@ -0,0 +1,112 @@
+from dataclasses import dataclass
+from typing import Optional
+from keras import layers, ops
+
+
+@dataclass
+class SelectOutput:
+    r"""The output of the :class:`Select` method, which holds an assignment
+    from selected nodes to their respective cluster(s).
+
+    Args:
+        node_index: The indices of the selected nodes.
+        num_nodes: The number of nodes.
+        cluster_index: The indices of the clusters each node in
+            :obj:`node_index` is assigned to.
+        num_clusters: The number of clusters.
+        weight (optional): A weight vector, denoting the strength
+            of the assignment of a node to its cluster. (default: :obj:`None`)
+    """
+    node_index: any
+    num_nodes: int
+    cluster_index: any
+    num_clusters: int
+    weight: Optional[any] = None
+
+    def __post_init__(self):
+        shape_node = getattr(self.node_index, "shape", None)
+        shape_cluster = getattr(self.cluster_index, "shape", None)
+        if shape_node is not None and len(shape_node) != 1:
+            raise ValueError(
+                f"Expected 'node_index' to be one-dimensional "
+                f"(got {len(shape_node)} dimensions)"
+            )
+        if shape_cluster is not None and len(shape_cluster) != 1:
+            raise ValueError(
+                f"Expected 'cluster_index' to be one-dimensional "
+                f"(got {len(shape_cluster)} dimensions)"
+            )
+        if (
+            shape_node is not None
+            and shape_cluster is not None
+            and len(shape_node) > 0
+            and len(shape_cluster) > 0
+            and shape_node[0] is not None
+            and shape_cluster[0] is not None
+            and shape_node[0] != shape_cluster[0]
+        ):
+            raise ValueError(
+                f"Expected 'node_index' and 'cluster_index' to hold the same "
+                f"number of values (got {shape_node[0]} and "
+                f"{shape_cluster[0]} values)"
+            )
+        if self.weight is not None:
+            shape_weight = getattr(self.weight, "shape", None)
+            if shape_weight is not None and len(shape_weight) != 1:
+                raise ValueError(
+                    f"Expected 'weight' vector to be one-dimensional "
+                    f"(got {len(shape_weight)} dimensions)"
+                )
+            if (
+                shape_weight is not None
+                and shape_node is not None
+                and len(shape_weight) > 0
+                and len(shape_node) > 0
+                and shape_weight[0] is not None
+                and shape_node[0] is not None
+                and shape_weight[0] != shape_node[0]
+            ):
+                raise ValueError(
+                    f"Expected 'weight' to hold {shape_node[0]} "
+                    f"values (got {shape_weight[0]} values)"
+                )
+
+
+import keras
+
+if keras.config.backend() == "jax":
+    try:
+        import jax
+        from jax.tree_util import register_pytree_node
+
+        register_pytree_node(
+            SelectOutput,
+            lambda s: (
+                (s.node_index, s.cluster_index, s.weight),
+                (s.num_nodes, s.num_clusters),
+            ),
+            lambda aux, children: SelectOutput(
+                children[0], aux[0], children[1], aux[1], children[2]
+            ),
+        )
+    except Exception:
+        pass
+
+
+
+class Select(layers.Layer):
+    r"""An abstract base class for implementing custom node selections as
+    described in the `"Understanding Pooling in Graph Neural Networks"
+    <https://arxiv.org/abs/1905.05178>`_ paper, which maps the nodes of an
+    input graph to supernodes in the coarsened graph.
+    """
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        pass
+
+    def call(self, *args, **kwargs) -> SelectOutput:
+        raise NotImplementedError
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}()'
+

k3_node/layers/pool/select/topk.py

@@ -0,0 +1,206 @@
+import numpy as np
+from typing import Callable, Optional, Union
+from keras import initializers, layers, ops
+
+from .base import Select, SelectOutput
+
+
+from k3_node.layers.conv.utils import is_tracing
+from k3_node.ops.segment import segment_max, segment_sum
+from k3_node.ops.creation import full
+
+
+def topk(
+    x,
+    ratio: Optional[Union[float, int]],
+    batch,
+    min_score: Optional[float] = None,
+    tol: float = 1e-7,
+):
+    r"""Selects top-k items according to score and batch assignment.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import topk
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        score = np.random.rand(10).astype("float32")
+        perm = topk(score, ratio=0.5, batch=batch)  # indices of the top 50% nodes per graph
+        print(tuple(perm.shape))  # (6,)
+        ```
+    """
+    # Plain NumPy inputs cannot be mixed with backend tensors; convert them first.
+    x = ops.convert_to_tensor(x) if isinstance(x, np.ndarray) else x
+    batch = ops.convert_to_tensor(batch) if isinstance(batch, np.ndarray) else batch
+    if is_tracing(x) or is_tracing(batch):
+        return ops.arange(ops.shape(x)[0], dtype="int32")
+
+    batch = ops.cast(batch, dtype="int32")
+    num_nodes = ops.shape(x)[0]
+
+    num_graphs = ops.max(batch) + 1 if ops.shape(batch)[0] > 0 else 0
+    try:
+        num_graphs = int(num_graphs)
+    except (TypeError, ValueError):
+        pass
+
+    if min_score is not None:
+        scores_max = segment_max(x, batch, num_segments=num_graphs)
+        scores_max_expanded = ops.take(scores_max, batch, axis=0) - tol
+        scores_min = ops.minimum(scores_max_expanded, min_score)
+        mask = x > scores_min
+        perm = ops.where(mask)
+        if isinstance(perm, (tuple, list)):
+            perm = perm[0]
+        perm = ops.reshape(perm, (-1,))
+        return ops.cast(perm, "int32")
+
+    if ratio is not None:
+        ones = ops.ones((num_nodes,), dtype="int32")
+        num_nodes_per_graph = segment_sum(ones, batch, num_segments=num_graphs)
+
+        if ratio >= 1:
+            k = full(ops.shape(num_nodes_per_graph), int(ratio), dtype="int32")
+        else:
+            k = ops.cast(
+                ops.ceil(ratio * ops.cast(num_nodes_per_graph, x.dtype)),
+                dtype="int32",
+            )
+
+        # Composite key: sorts by batch ascending, then by score descending
+        score_span = ops.max(x) - ops.min(x) + 1.0
+        key = ops.cast(batch, x.dtype) * (score_span * 2.0) - x
+        perm = ops.argsort(key)
+
+        batch_sorted = ops.take(batch, perm, axis=0)
+        # ptr for cumsum
+        ptr = ops.concatenate(
+            [ops.zeros((1,), dtype="int32"), ops.cumsum(num_nodes_per_graph)[:-1]],
+            axis=0,
+        )
+        rank_in_graph = ops.arange(num_nodes, dtype="int32") - ops.take(ptr, batch_sorted, axis=0)
+        mask = rank_in_graph < ops.take(k, batch_sorted, axis=0)
+        valid_idx = ops.where(mask)
+        if isinstance(valid_idx, (tuple, list)):
+            valid_idx = valid_idx[0]
+        valid_idx = ops.reshape(valid_idx, (-1,))
+        return ops.cast(ops.take(perm, valid_idx, axis=0), "int32")
+
+    raise ValueError("At least one of 'ratio' and 'min_score' must be specified.")
+
+
+class SelectTopK(Select):
+    r"""Selects the top-:math:`k` nodes with highest projection scores.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SelectTopK
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        select = SelectTopK(in_channels=8, ratio=0.5)
+        out = select(x, batch)
+        print(tuple(out.node_index.shape))  # (6,): indices of the kept nodes
+        print(tuple(out.weight.shape))  # (6,): their scores
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        ratio: Union[int, float] = 0.5,
+        min_score: Optional[float] = None,
+        act: Union[str, Callable] = "tanh",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if ratio is None and min_score is None:
+            raise ValueError(
+                f"At least one of 'ratio' and 'min_score' must be specified in '{self.__class__.__name__}'"
+            )
+
+        self.in_channels = in_channels
+        self.ratio = ratio
+        self.min_score = min_score
+        self.act_fn = act if callable(act) else layers.Activation(act)
+
+        self.weight = self.add_weight(
+            shape=(1, in_channels),
+            initializer=initializers.RandomUniform(
+                minval=-1.0 / (in_channels**0.5), maxval=1.0 / (in_channels**0.5)
+            ),
+            trainable=True,
+            name="weight",
+        )
+
+    def reset_parameters(self):
+        limit = 1.0 / (self.in_channels**0.5)
+        init = initializers.RandomUniform(minval=-limit, maxval=limit)
+        self.weight.assign(init(self.weight.shape, dtype=self.weight.dtype))
+
+    def build(self, input_shape=None):
+        if hasattr(self.act_fn, "built") and not self.act_fn.built:
+            self.act_fn.build(input_shape)
+        self.built = True
+
+    def call(self, x, batch=None) -> SelectOutput:
+        num_nodes = ops.shape(x)[0]
+        if batch is None:
+            batch = ops.zeros((num_nodes,), dtype="int32")
+        else:
+            batch = ops.cast(batch, dtype="int32")
+
+        if len(ops.shape(x)) == 1:
+            x = ops.expand_dims(x, axis=-1)
+
+        score = ops.sum(x * self.weight, axis=-1)
+
+        if is_tracing(x) or is_tracing(batch):
+            node_index = ops.arange(num_nodes, dtype="int32")
+            return SelectOutput(
+                node_index=node_index,
+                num_nodes=num_nodes,
+                cluster_index=node_index,
+                num_clusters=num_nodes,
+                weight=score,
+            )
+
+        if self.min_score is None:
+            norm_w = ops.sqrt(ops.sum(ops.power(self.weight, 2), axis=-1))
+            score = self.act_fn(score / norm_w)
+        else:
+            # Graph-wise softmax
+            num_graphs = ops.max(batch) + 1 if ops.shape(batch)[0] > 0 else 0
+            try:
+                num_graphs = int(num_graphs)
+            except (TypeError, ValueError):
+                pass
+            score_max = segment_max(score, batch, num_segments=num_graphs)
+            score_max_exp = ops.take(score_max, batch, axis=0)
+            exp_score = ops.exp(score - score_max_exp)
+            exp_sum = segment_sum(exp_score, batch, num_segments=num_graphs)
+            exp_sum_exp = ops.take(exp_sum, batch, axis=0)
+            score = exp_score / (exp_sum_exp + 1e-12)
+
+        node_index = topk(score, self.ratio, batch, self.min_score)
+        num_selected = ops.shape(node_index)[0]
+
+        return SelectOutput(
+            node_index=node_index,
+            num_nodes=num_nodes,
+            cluster_index=ops.arange(num_selected, dtype="int32"),
+            num_clusters=num_selected,
+            weight=ops.take(score, node_index, axis=0),
+        )
+
+    def __repr__(self) -> str:
+        if self.min_score is None:
+            arg = f"ratio={self.ratio}"
+        else:
+            arg = f"min_score={self.min_score}"
+        return f"{self.__class__.__name__}({self.in_channels}, {arg})"