k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/general_conv.py

@@ -0,0 +1,218 @@
+import math
+from typing import Optional, Union, Tuple
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import softmax
+
+
+class GeneralConv(MessagePassing):
+    r"""A general GNN layer adapted from the `"Design Space for Graph Neural
+    Networks" <https://arxiv.org/abs/2011.08843>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GeneralConv
+
+        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
+        edge_attr = np.random.rand(30, 3).astype("float32")  # 3 features per edge
+
+        layer = GeneralConv(in_channels=8, out_channels=16, in_edge_channels=3)
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: Optional[int] = None,
+        in_edge_channels: Optional[int] = None,
+        aggr: str = "add",
+        skip_linear: bool = False,
+        directed_msg: bool = True,
+        heads: int = 1,
+        attention: bool = False,
+        attention_type: str = "additive",
+        l2_normalize: bool = False,
+        bias: bool = True,
+        # Spektral compatibility arguments:
+        channels: Optional[int] = None,
+        batch_norm: Optional[bool] = None,
+        dropout: Optional[float] = None,
+        aggregate: Optional[str] = None,
+        activation: Optional[str] = None,
+        use_bias: Optional[bool] = None,
+        **kwargs,
+    ):
+        if out_channels is None:
+            if channels is not None:
+                out_channels = channels
+                in_channels = -1
+            else:
+                out_channels = in_channels
+                in_channels = -1
+
+        if aggregate is not None:
+            aggr = aggregate
+        if use_bias is not None:
+            bias = use_bias
+
+        kwargs.setdefault("aggr", aggr)
+        super().__init__(node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.in_edge_channels = in_edge_channels
+        self.skip_linear = skip_linear
+        self.directed_msg = directed_msg
+        self.heads = heads
+        self.attention = attention
+        self.attention_type = attention_type
+        self.l2_normalize = l2_normalize
+        self.use_bias = bias
+
+        if isinstance(in_channels, int):
+            self.in_channels_l = in_channels
+            self.in_channels_r = in_channels
+        else:
+            self.in_channels_l, self.in_channels_r = in_channels
+
+        self.lin_msg = Dense(out_channels * heads, use_bias=bias)
+        if not directed_msg:
+            self.lin_msg_i = Dense(out_channels * heads, use_bias=bias)
+        else:
+            self.lin_msg_i = None
+
+        if skip_linear or self.in_channels_r != out_channels:
+            self.lin_self = Dense(out_channels, use_bias=bias)
+        else:
+            self.lin_self = None
+
+        if in_edge_channels is not None:
+            self.lin_edge = Dense(out_channels * heads, use_bias=bias)
+        else:
+            self.lin_edge = None
+
+        if attention:
+            if attention_type == "additive":
+                self.att_msg = self.add_weight(
+                    shape=(1, heads, out_channels),
+                    initializer="glorot_uniform",
+                    name="att_msg",
+                )
+            elif attention_type == "dot_product":
+                self.scaler = math.sqrt(out_channels)
+            else:
+                raise ValueError(f"Attention type '{attention_type}' not supported")
+        else:
+            self.att_msg = None
+
+    def build(self, input_shape=None):
+        if input_shape is not None:
+            if isinstance(input_shape, (list, tuple)) and len(input_shape) > 0 and isinstance(input_shape[0], (list, tuple)):
+                dim = input_shape[0][-1]
+            elif isinstance(input_shape, (list, tuple)) and len(input_shape) > 0 and input_shape[0] is not None and not isinstance(input_shape[0], (int, type(None))):
+                dim = getattr(input_shape[0], "shape", [None, None])[-1]
+            else:
+                dim = input_shape[-1]
+            if (self.in_channels_l is None or self.in_channels_l == -1) and dim is not None:
+                self.in_channels_l = self.in_channels_r = dim
+        if self.in_channels_l is not None and self.in_channels_l != -1:
+            self.lin_msg.build((None, self.in_channels_l))
+            if self.lin_msg_i is not None:
+                self.lin_msg_i.build((None, self.in_channels_r))
+            if self.lin_self is not None:
+                self.lin_self.build((None, self.in_channels_r))
+            if self.lin_edge is not None:
+                self.lin_edge.build((None, self.in_edge_channels))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_attr=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, edge_index, edge_attr = inputs
+                elif len(inputs) == 2:
+                    x, edge_index = inputs
+                else:
+                    raise ValueError(f"Unexpected inputs length {len(inputs)}")
+            else:
+                raise ValueError("Expected (x, edge_index) or x and edge_index")
+        else:
+            x = inputs
+
+        if isinstance(x, (list, tuple)):
+            x_l, x_r = x
+        else:
+            x_l = x_r = x
+
+        if self.in_channels_l is None or self.in_channels_l == -1:
+            dim = ops.shape(x_l)[-1]
+            self.in_channels_l = self.in_channels_r = dim
+            self.build((None, dim))
+
+        # Check legacy adj matrix
+        is_legacy = False
+        if hasattr(edge_index, "shape") and len(edge_index.shape) == 2:
+            if edge_index.shape[0] != 2 and edge_index.shape[0] == edge_index.shape[1]:
+                is_legacy = True
+        elif not hasattr(edge_index, "shape"):
+            is_legacy = True
+
+        if is_legacy:
+            if hasattr(edge_index, "indices") and not callable(edge_index.indices):
+                edge_index = ops.transpose(edge_index.indices)
+            else:
+                adj = edge_index
+                row, col = ops.where(adj > 0)
+                edge_index = ops.stack([row, col], axis=0)
+
+        num_nodes = ops.shape(x_r)[0]
+        size = (ops.shape(x_l)[0], num_nodes)
+
+        out = self.propagate(
+            edge_index,
+            x=(x_l, x_r),
+            edge_attr=edge_attr,
+            size=size,
+        )
+        out = ops.mean(out, axis=1)  # aggregate heads
+
+        if self.lin_self is not None:
+            out = out + self.lin_self(x_r)
+        else:
+            out = out + x_r
+
+        if self.l2_normalize:
+            out = out / (ops.norm(out, axis=-1, keepdims=True) + 1e-12)
+
+        return out
+
+    def _message_basic(self, x_i, x_j, edge_attr=None):
+        if self.directed_msg:
+            x_j = self.lin_msg(x_j)
+        else:
+            x_j = self.lin_msg(x_j) + self.lin_msg_i(x_i)
+        if edge_attr is not None and self.lin_edge is not None:
+            x_j = x_j + self.lin_edge(edge_attr)
+        return x_j
+
+    def message(self, x_i, x_j, edge_attr=None, index=None, size_i=None):
+        x_j_out = self._message_basic(x_i, x_j, edge_attr)
+        x_j_out = ops.reshape(x_j_out, (-1, self.heads, self.out_channels))
+
+        if self.attention:
+            if self.attention_type == "dot_product":
+                x_i_out = self._message_basic(x_j, x_i, edge_attr)
+                x_i_out = ops.reshape(x_i_out, (-1, self.heads, self.out_channels))
+                alpha = ops.sum(x_i_out * x_j_out, axis=-1) / self.scaler
+            else:
+                alpha = ops.sum(x_j_out * self.att_msg, axis=-1)
+            alpha = ops.leaky_relu(alpha, negative_slope=0.2)
+            alpha = softmax(alpha, index, num_nodes=size_i, dim=0)
+            return x_j_out * ops.expand_dims(alpha, -1)
+        else:
+            return x_j_out

k3_node/layers/conv/gin_conv.py

@@ -0,0 +1,218 @@
+from typing import Callable, Optional, Union, List
+import keras
+from keras import layers, ops, activations
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+def _apply_nn(nn, x, training):
+    # Forward `training` explicitly: on the JAX backend Keras does not propagate it to nested
+    # layers, so dropout / batch norm inside `nn` would otherwise never be in training mode.
+    if isinstance(nn, keras.layers.Layer):
+        return nn(x, training=training)
+    return nn(x)
+
+class GINConv(MessagePassing):
+    r"""The graph isomorphism operator from the `"How Powerful are Graph
+    Neural Networks?" <https://arxiv.org/abs/1810.00826>`_ paper.
+
+    Args:
+        nn: A neural network :math:`h_{\mathbf{\Theta}}` that maps node
+            features to new embeddings (e.g. a :class:`keras.Sequential` or callable).
+            Also accepts an integer `channels` for backward compatibility.
+        eps: (Initial) :math:`\epsilon`-value. (default: ``0.0``)
+        train_eps: If set to :obj:`True`, :math:`\epsilon` will be a learnable
+            parameter. (default: ``False``)
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import GINConv
+
+        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
+
+        mlp = keras.Sequential([keras.layers.Dense(16, activation="relu"), keras.layers.Dense(16)])
+        layer = GINConv(mlp)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        nn: Union[Callable, int],
+        eps: float = 0.0,
+        train_eps: bool = False,
+        epsilon: Optional[float] = None,
+        mlp_hidden: Optional[List[int]] = None,
+        mlp_activation: str = "relu",
+        mlp_batchnorm: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr=kwargs.pop("aggr", kwargs.pop("aggregate", "add")), **kwargs)
+
+        if epsilon is not None:
+            eps = epsilon
+
+        self.initial_eps = eps
+        self.train_eps = train_eps
+
+        # Backward compatibility if nn is int (channels)
+        if isinstance(nn, int):
+            channels = nn
+            mlp_hidden = mlp_hidden or []
+            act = activations.get(mlp_activation)
+            seq_layers = []
+            for h in mlp_hidden:
+                seq_layers.append(layers.Dense(h, activation=act))
+                if mlp_batchnorm:
+                    seq_layers.append(layers.BatchNormalization(momentum=0.9, epsilon=1e-5))
+            seq_layers.append(layers.Dense(channels, activation=kwargs.get("activation", None)))
+            self.nn = keras.Sequential(seq_layers)
+        else:
+            self.nn = nn
+
+        if train_eps:
+            self.eps = self.add_weight(
+                shape=(1,),
+                initializer=keras.initializers.Constant(eps),
+                name="eps",
+            )
+        else:
+            self.eps = ops.cast(eps, "float32")
+
+    def build(self, input_shape):
+        if hasattr(self.nn, "build") and not getattr(self.nn, "built", False):
+            self.nn.build(input_shape)
+        self.built = True
+
+    def call(self, x, edge_index=None, size=None, training=None, **kwargs):
+        # Handle legacy calling: conv((x, adj))
+        if edge_index is None and isinstance(x, (tuple, list)) and len(x) == 2:
+            arg0, arg1 = x[0], x[1]
+            s1 = getattr(arg1, "shape", None)
+            if (
+                s1 is not None
+                and len(s1) == 2
+                and s1[0] is not None
+                and s1[1] is not None
+                and s1[0] > 2
+                and s1[0] == s1[1]
+            ):
+                where_adj = ops.where(arg1 != 0)
+                where_adj = where_adj if not isinstance(where_adj, list) else where_adj
+                edge_index = ops.stack([where_adj[0], where_adj[1]], axis=0)
+                x = arg0
+            elif s1 is not None and len(s1) >= 1 and s1[0] == 2:
+                edge_index = arg1
+                x = arg0
+
+        if not isinstance(x, (tuple, list)):
+            x_src, x_dst = x, x
+        else:
+            x_src, x_dst = x[0], x[1]
+
+        out = self.propagate(edge_index, x=(x_src, x_dst), size=size)
+
+        if x_dst is not None:
+            out = out + (1.0 + self.eps) * x_dst
+
+        return _apply_nn(self.nn, out, training)
+
+    def message(self, x_j):
+        return x_j
+
+
+class GINEConv(MessagePassing):
+    r"""The modified :class:`GINConv` operator from the `"Strategies for
+    Pre-training Graph Neural Networks" <https://arxiv.org/abs/1905.12265>`_
+    paper, which is able to incorporate edge features into aggregation.
+
+    Args:
+        nn: A neural network :math:`h_{\mathbf{\Theta}}`.
+        eps: (Initial) :math:`\epsilon`-value. (default: ``0.0``)
+        train_eps: If set to :obj:`True`, :math:`\epsilon` will be a learnable
+            parameter. (default: ``False``)
+        edge_dim: Edge feature dimensionality. (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import GINEConv
+
+        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
+        edge_attr = np.random.rand(30, 3).astype("float32")  # 3 features per edge
+
+        mlp = keras.Sequential([keras.layers.Dense(16, activation="relu"), keras.layers.Dense(16)])
+        layer = GINEConv(mlp, edge_dim=3)
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        nn: Callable,
+        eps: float = 0.0,
+        train_eps: bool = False,
+        edge_dim: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(aggr=kwargs.pop("aggr", kwargs.pop("aggregate", "add")), **kwargs)
+        self.nn = nn
+        self.initial_eps = eps
+        self.train_eps = train_eps
+        self.edge_dim = edge_dim
+
+        if train_eps:
+            self.eps = self.add_weight(
+                shape=(1,),
+                initializer=keras.initializers.Constant(eps),
+                name="eps",
+            )
+        else:
+            self.eps = ops.cast(eps, "float32")
+
+        if edge_dim is not None:
+            self.lin = layers.Dense(edge_dim, use_bias=True)
+        else:
+            self.lin = None
+
+    def build(self, input_shape):
+        if self.lin is not None:
+            if isinstance(input_shape, (tuple, list)) and len(input_shape) > 0 and isinstance(input_shape[0], (tuple, list)):
+                node_dim = input_shape[0][-1]
+            elif isinstance(input_shape, (tuple, list)):
+                node_dim = input_shape[-1]
+            else:
+                node_dim = 16
+            self.lin.units = node_dim
+            self.lin.build((None, self.edge_dim))
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_attr=None, size=None, training=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]
+
+        out = self.propagate(edge_index, x=(x_src, x_dst), edge_attr=edge_attr, size=size)
+
+        if x_dst is not None:
+            out = out + (1.0 + self.eps) * x_dst
+
+        return _apply_nn(self.nn, out, training)
+
+    def message(self, x_j, edge_attr=None):
+        if edge_attr is None:
+            return ops.relu(x_j)
+        if self.lin is not None:
+            edge_attr = self.lin(edge_attr)
+        return ops.relu(x_j + edge_attr)

k3_node/layers/conv/gmm_conv.py

@@ -0,0 +1,172 @@
+from typing import Union, Tuple
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class GMMConv(MessagePassing):
+    r"""The gaussian mixture model convolutional operator from the `"Geometric
+    Deep Learning on Graphs and Manifolds using Mixture Model CNNs"
+    <https://arxiv.org/abs/1611.08402>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GMMConv
+
+        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
+        edge_attr = np.random.rand(30, 3).astype("float32")  # 3 features per edge
+
+        layer = GMMConv(in_channels=8, out_channels=16, dim=3, kernel_size=2)
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: int,
+        dim: int,
+        kernel_size: int,
+        separate_gaussians: bool = False,
+        aggr: str = "mean",
+        root_weight: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.dim = dim
+        self.kernel_size = kernel_size
+        self.separate_gaussians = separate_gaussians
+        self.root_weight = root_weight
+        self.use_bias = bias
+
+        if isinstance(in_channels, int):
+            self.in_channels_l = in_channels
+            self.in_channels_r = in_channels
+        else:
+            self.in_channels_l, self.in_channels_r = in_channels
+
+        self.g = self.add_weight(
+            shape=(self.in_channels_l, out_channels * kernel_size),
+            initializer="glorot_uniform",
+            name="g",
+        )
+
+        if not separate_gaussians:
+            self.mu = self.add_weight(
+                shape=(kernel_size, dim),
+                initializer="glorot_uniform",
+                name="mu",
+            )
+            self.sigma = self.add_weight(
+                shape=(kernel_size, dim),
+                initializer="ones",
+                name="sigma",
+            )
+        else:
+            self.mu = self.add_weight(
+                shape=(self.in_channels_l, out_channels, kernel_size, dim),
+                initializer="glorot_uniform",
+                name="mu",
+            )
+            self.sigma = self.add_weight(
+                shape=(self.in_channels_l, out_channels, kernel_size, dim),
+                initializer="ones",
+                name="sigma",
+            )
+
+        if root_weight:
+            self.root = Dense(out_channels, use_bias=False)
+        else:
+            self.root = None
+
+        if bias:
+            self.bias = self.add_weight(
+                shape=(out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape=None):
+        if self.root is not None:
+            self.root.build((None, self.in_channels_r))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_attr=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, edge_index, edge_attr = inputs
+                elif len(inputs) == 2:
+                    x, edge_index = inputs
+                else:
+                    raise ValueError(f"Unexpected input length {len(inputs)}")
+            else:
+                raise ValueError("Expected (x, edge_index) or x and edge_index")
+        else:
+            x = inputs
+
+        if not self.built:
+            self.build()
+
+        if isinstance(x, (list, tuple)):
+            x_src, x_dst = x
+        else:
+            x_src = x_dst = x
+
+        num_nodes = ops.shape(x_dst)[0]
+
+        if not self.separate_gaussians:
+            x_l = ops.matmul(x_src, self.g)
+            out = self.propagate(
+                edge_index,
+                x=(x_l, x_dst),
+                edge_attr=edge_attr,
+                size=(ops.shape(x_src)[0], num_nodes),
+            )
+        else:
+            out = self.propagate(
+                edge_index,
+                x=(x_src, x_dst),
+                edge_attr=edge_attr,
+                size=(ops.shape(x_src)[0], num_nodes),
+            )
+
+        if self.root is not None:
+            out = out + self.root(x_dst)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_j, edge_attr):
+        EPS = 1e-15
+        M = self.out_channels
+        K = self.kernel_size
+
+        if not self.separate_gaussians:
+            # edge_attr: (E, D), mu: (K, D), sigma: (K, D)
+            diff = ops.expand_dims(edge_attr, 1) - ops.expand_dims(self.mu, 0)
+            gaussian = -0.5 * ops.square(diff) / (EPS + ops.square(ops.expand_dims(self.sigma, 0)))
+            gaussian = ops.exp(ops.sum(gaussian, axis=-1))  # (E, K)
+
+            x_j_reshaped = ops.reshape(x_j, (-1, K, M))
+            return ops.sum(x_j_reshaped * ops.expand_dims(gaussian, -1), axis=1)
+        else:
+            F = self.in_channels_l
+            diff = ops.expand_dims(ops.expand_dims(ops.expand_dims(edge_attr, 1), 1), 1) - ops.expand_dims(self.mu, 0)
+            gaussian = -0.5 * ops.square(diff) / (EPS + ops.square(ops.expand_dims(self.sigma, 0)))
+            gaussian = ops.exp(ops.sum(gaussian, axis=-1))  # (E, F, M, K)
+            g_reshaped = ops.reshape(self.g, (1, F, M, K))
+            gaussian = ops.sum(gaussian * g_reshaped, axis=-1)  # (E, F, M)
+            return ops.sum(ops.expand_dims(x_j, -1) * gaussian, axis=1)
+