k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/hypergraph_conv.py

@@ -0,0 +1,182 @@
+import keras
+from keras import ops
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import scatter, softmax
+from k3_node.layers.conv.utils import is_tracing
+
+
+class HypergraphConv(MessagePassing):
+    r"""The hypergraph convolutional operator from the `"Hypergraph Convolution
+    and Hypergraph Attention" <https://arxiv.org/abs/1901.08150>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        use_attention (bool, optional): Whether to use hypergraph attention. (default: :obj:`False`)
+        attention_mode (str, optional): Attention mode (:obj:`"node"` or :obj:`"edge"`). (default: :obj:`"node"`)
+        heads (int, optional): Number of multi-head-attentions. (default: :obj:`1`)
+        concat (bool, optional): Whether to concatenate heads. (default: :obj:`True`)
+        negative_slope (float, optional): LeakyReLU angle. (default: :obj:`0.2`)
+        bias (bool, optional): Whether to learn an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HypergraphConv
+
+        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
+
+        # Row 0: node index, row 1: hyperedge index (nodes 0-5 form two hyperedges)
+        hyperedge_index = np.array([[0, 1, 2, 3, 4, 5], [0, 0, 0, 1, 1, 1]])
+        layer = HypergraphConv(in_channels=8, out_channels=16)
+        out = layer(x, hyperedge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        use_attention: bool = False,
+        attention_mode: str = "node",
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        bias: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(node_dim=0, **kwargs)
+
+        assert attention_mode in ["node", "edge"]
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.use_attention = use_attention
+        self.attention_mode = attention_mode
+        self.negative_slope = negative_slope
+        self.use_bias = bias
+
+        if self.use_attention:
+            self.heads = heads
+            self.concat = concat
+            self.lin = keras.layers.Dense(heads * out_channels, use_bias=False)
+        else:
+            self.heads = 1
+            self.concat = True
+            self.lin = keras.layers.Dense(out_channels, use_bias=False)
+
+    def build(self, input_shape=None):
+        if not self.lin.built and self.in_channels > 0:
+            self.lin.build((None, self.in_channels))
+
+        if self.use_attention:
+            self.att = self.add_weight(
+                shape=(1, self.heads, 2 * self.out_channels),
+                initializer="glorot_uniform",
+                trainable=True,
+                name="att",
+            )
+
+        out_dim = self.heads * self.out_channels if self.concat else self.out_channels
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(out_dim,),
+                initializer="zeros",
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+        super().build(input_shape)
+
+    def call(
+        self,
+        x,
+        hyperedge_index,
+        hyperedge_weight=None,
+        hyperedge_attr=None,
+        num_edges=None,
+    ):
+        if not self.built:
+            self.build((None, self.in_channels))
+
+        num_nodes = ops.shape(x)[0]
+        if num_edges is None:
+            if not is_tracing(hyperedge_index):
+                try:
+                    num_edges = int(ops.max(hyperedge_index[1])) + 1
+                except Exception:
+                    num_edges = None
+
+        if hyperedge_weight is None:
+            edge_w = ops.ones((ops.shape(hyperedge_index)[1],), dtype=x.dtype)
+        else:
+            edge_w = ops.take(hyperedge_weight, hyperedge_index[1], axis=0)
+
+        x = self.lin(x)
+
+        alpha = None
+        if self.use_attention:
+            assert hyperedge_attr is not None
+            x = ops.reshape(x, (-1, self.heads, self.out_channels))
+            hyperedge_attr = self.lin(hyperedge_attr)
+            hyperedge_attr = ops.reshape(hyperedge_attr, (-1, self.heads, self.out_channels))
+
+            x_i = ops.take(x, hyperedge_index[0], axis=0)
+            x_j = ops.take(hyperedge_attr, hyperedge_index[1], axis=0)
+
+            alpha = ops.sum(ops.concatenate([x_i, x_j], axis=-1) * self.att, axis=-1)
+            alpha = ops.leaky_relu(alpha, negative_slope=self.negative_slope)
+
+            if self.attention_mode == "node":
+                alpha = softmax(alpha, index=hyperedge_index[1], num_nodes=num_edges)
+            else:
+                alpha = softmax(alpha, index=hyperedge_index[0], num_nodes=num_nodes)
+
+        D = scatter(edge_w, hyperedge_index[0], dim=0, dim_size=num_nodes, reduce="sum")
+        D = ops.where(ops.equal(D, 0), 0.0, 1.0 / D)
+
+        ones_e = ops.ones((ops.shape(hyperedge_index)[1],), dtype=x.dtype)
+        B = scatter(ones_e, hyperedge_index[1], dim=0, dim_size=num_edges, reduce="sum")
+        B = ops.where(ops.equal(B, 0), 0.0, 1.0 / B)
+
+        # 1. nodes -> hyperedges
+        out = self.propagate(
+            hyperedge_index,
+            x=x,
+            norm=B,
+            alpha=alpha,
+            size=(num_nodes, num_edges),
+        )
+        # 2. hyperedges -> nodes
+        flipped_edge_index = ops.stack([hyperedge_index[1], hyperedge_index[0]], axis=0)
+        out = self.propagate(
+            flipped_edge_index,
+            x=out,
+            norm=D,
+            alpha=alpha,
+            size=(num_edges, num_nodes),
+        )
+
+        if self.concat:
+            out = ops.reshape(out, (-1, self.heads * self.out_channels))
+        else:
+            out = ops.mean(out, axis=1)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_j, norm_i, alpha=None):
+        H, F = self.heads, self.out_channels
+        norm_i_exp = ops.expand_dims(ops.expand_dims(norm_i, axis=-1), axis=-1)
+        out = norm_i_exp * ops.reshape(x_j, (-1, H, F))
+        if alpha is not None:
+            out = ops.expand_dims(alpha, axis=-1) * out
+        return out
+

k3_node/layers/conv/le_conv.py

@@ -0,0 +1,81 @@
+from typing import Union, Tuple
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class LEConv(MessagePassing):
+    r"""The local extremum graph convolutional operator from the
+    `"ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph
+    Representations" <https://arxiv.org/abs/1911.07979>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample, or a tuple for bipartite graphs.
+        out_channels: Size of each output sample.
+        bias: If set to :obj:`False`, the layer will not learn an additive bias.
+            (default: ``True``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import LEConv
+
+        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
+
+        layer = LEConv(in_channels=8, out_channels=16)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: int,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.use_bias = bias
+
+        self.lin1 = layers.Dense(out_channels, use_bias=bias)
+        self.lin2 = layers.Dense(out_channels, use_bias=False)
+        self.lin3 = layers.Dense(out_channels, use_bias=bias)
+
+    def build(self, input_shape):
+        if isinstance(input_shape, (tuple, list)) and len(input_shape) > 0 and isinstance(input_shape[0], (tuple, list)):
+            in_channels_src = input_shape[0][-1]
+            in_channels_dst = input_shape[1][-1] if len(input_shape) > 1 and input_shape[1] is not None else in_channels_src
+        else:
+            in_channels_src = input_shape[-1]
+            in_channels_dst = input_shape[-1]
+
+        self.lin1.build((None, in_channels_src))
+        self.lin2.build((None, in_channels_dst))
+        self.lin3.build((None, in_channels_dst))
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        if not isinstance(x, (tuple, list)):
+            x_src, x_dst = x, x
+        else:
+            x_src, x_dst = x[0], x[1]
+
+        a = self.lin1(x_src)
+        b = self.lin2(x_dst)
+
+        out = self.propagate(edge_index, a=a, b=b, edge_weight=edge_weight)
+        return out + self.lin3(x_dst)
+
+    def message(self, a_j, b_i, edge_weight=None):
+        out = a_j - b_i
+        if edge_weight is None:
+            return out
+        return out * ops.expand_dims(edge_weight, -1)
+

k3_node/layers/conv/lg_conv.py

@@ -0,0 +1,58 @@
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm
+
+
+class LGConv(MessagePassing):
+    r"""The LightGCN operator from the `"LightGCN: Simplifying and Powering
+    Graph Convolution Network for Recommendation"
+    <https://arxiv.org/abs/2002.02126>`_ paper.
+
+    Args:
+        normalize: Whether to apply symmetric normalization. (default: ``True``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import LGConv
+
+        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
+
+        layer = LGConv()
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(self, normalize: bool = True, **kwargs):
+        super().__init__(aggr="add", **kwargs)
+        self.normalize = normalize
+
+    def build(self, input_shape):
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        if self.normalize:
+            num_nodes = x.shape[self.node_dim] if hasattr(x, "shape") and x.shape[self.node_dim] is not None else ops.shape(x)[self.node_dim]
+            edge_index, edge_weight = gcn_norm(
+                edge_index,
+                edge_weight,
+                num_nodes=num_nodes,
+                add_self_loops=False,
+                flow=self.flow,
+                dtype=x.dtype,
+            )
+
+        return self.propagate(edge_index, x=x, edge_weight=edge_weight)
+
+    def message(self, x_j, edge_weight=None):
+        if edge_weight is None:
+            return x_j
+        return ops.expand_dims(edge_weight, -1) * x_j

k3_node/layers/conv/meshcnn_conv.py

@@ -0,0 +1,84 @@
+from typing import Optional, List
+import keras
+from keras import ops
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class MeshCNNConv(MessagePassing):
+    r"""The MeshCNN convolutional operator from the `"MeshCNN: A Network With An Edge"
+    <https://arxiv.org/abs/1809.05910>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        kernels (List[keras.layers.Layer], optional): A list of 5 neural network layers
+            that transform edge representations. (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MeshCNNConv
+
+        # In MeshCNN the "nodes" are mesh edges; each mesh edge has exactly 4 neighboring edges,
+        # listed in a fixed order in edge_index (4 incoming entries per mesh edge).
+        x = np.random.rand(4, 8).astype("float32")  # features of 4 mesh edges
+        edge_index = np.array([
+            [1, 2, 3, 0, 0, 2, 3, 1, 0, 1, 3, 2, 0, 1, 2, 3],  # neighboring mesh edge
+            [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],  # mesh edge being updated
+        ])
+        layer = MeshCNNConv(in_channels=8, out_channels=16)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (4, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernels: Optional[List[keras.layers.Layer]] = None,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if kernels is None:
+            self.kernels = [
+                keras.layers.Dense(out_channels, use_bias=True)
+                for _ in range(5)
+            ]
+        else:
+            assert len(kernels) == 5, "kernels must be a list of 5 layers"
+            self.kernels = kernels
+
+    def build(self, input_shape=None):
+        for k in self.kernels:
+            if hasattr(k, "build") and not k.built:
+                k.build((None, self.in_channels))
+        super().build(input_shape)
+
+    def call(self, x, edge_index, **kwargs):
+        return self.propagate(edge_index, x=x)
+
+    def message(self, x_j):
+        n_a = x_j[0::4]
+        n_b = x_j[1::4]
+        n_c = x_j[2::4]
+        n_d = x_j[3::4]
+
+        m1 = self.kernels[1](ops.abs(n_a - n_c))
+        m2 = self.kernels[2](n_a + n_c)
+        m3 = self.kernels[3](ops.abs(n_b - n_d))
+        m4 = self.kernels[4](n_b + n_d)
+
+        return ops.reshape(
+            ops.stack([m1, m2, m3, m4], axis=1),
+            (-1, self.out_channels),
+        )
+
+    def update(self, inputs, x):
+        return self.kernels[0](x) + inputs
+