k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/gcn.py

@@ -0,0 +1,144 @@
+# ported from stellargraph
+from keras import ops
+from keras import activations, initializers, constraints, regularizers
+from keras.layers import Layer, dot
+from k3_node.ops.creation import repeat
+
+
+class GraphConvolution(Layer):
+    """
+    `k3_node.layers.GraphConvolution` 
+    Implementation of Graph Convolution (GCN) layer
+
+    Args:
+        units: Positive integer, dimensionality of the output space.
+        activation: Activation function to use.
+        use_bias: Whether to add a bias to the linear transformation.
+        final_layer: Deprecated, use tf.gather or GatherIndices instead.
+        input_dim: Deprecated, use `keras.layers.Input` with `input_shape` instead.
+        kernel_initializer: Initializer for the `kernel` weights matrix.
+        kernel_regularizer: Regularizer for the `kernel` weights matrix.
+        kernel_constraint: Constraint for the `kernel` weights matrix.
+        bias_initializer: Initializer for the bias vector.
+        bias_regularizer: Regularizer for the bias vector.
+        bias_constraint: Constraint for the bias vector.
+        **kwargs: Additional arguments to pass to the `Layer` superclass.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GraphConvolution
+
+        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 = GraphConvolution(units=16, activation="relu")
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        units,
+        activation=None,
+        use_bias=True,
+        final_layer=None,
+        input_dim=None,
+        kernel_initializer="glorot_uniform",
+        kernel_regularizer=None,
+        kernel_constraint=None,
+        bias_initializer="zeros",
+        bias_regularizer=None,
+        bias_constraint=None,
+        **kwargs,
+    ):
+        if isinstance(activation, int):
+            # Called as GraphConvolution(in_channels, out_channels)
+            self.in_channels = units
+            units = activation
+            activation = kwargs.pop("activation", None)
+
+        if "input_shape" not in kwargs and input_dim is not None:
+            kwargs["input_shape"] = (input_dim,)
+
+        self.units = units
+        self.activation = activations.get(activation)
+        self.use_bias = use_bias
+        if final_layer is not None:
+            raise ValueError(
+                "'final_layer' is not longer supported, use 'tf.gather' or 'GatherIndices' separately"
+            )
+
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.kernel_regularizer = regularizers.get(kernel_regularizer)
+        self.kernel_constraint = constraints.get(kernel_constraint)
+        self.bias_initializer = initializers.get(bias_initializer)
+        self.bias_regularizer = regularizers.get(bias_regularizer)
+        self.bias_constraint = constraints.get(bias_constraint)
+
+        super().__init__(**kwargs)
+
+    def build(self, input_shapes):
+        if isinstance(input_shapes, (list, tuple)) and len(input_shapes) > 0 and isinstance(input_shapes[0], (list, tuple)):
+            feat_shape = input_shapes[0]
+        else:
+            feat_shape = input_shapes
+        input_dim = int(feat_shape[-1]) if feat_shape is not None and feat_shape[-1] is not None else 8
+
+        self.kernel = self.add_weight(
+            shape=(1, input_dim, self.units),
+            initializer=self.kernel_initializer,
+            name="kernel",
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+        )
+
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.units,),
+                initializer=self.bias_initializer,
+                name="bias",
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+            )
+        else:
+            self.bias = None
+        self.built = True
+
+    def call(self, inputs, A=None, **kwargs):
+        if A is not None:
+            features = inputs
+        elif isinstance(inputs, (list, tuple)):
+            features, A = inputs
+        else:
+            features, A = inputs, None
+
+        if A is not None and hasattr(A, "shape") and len(A.shape) == 2 and A.shape[0] == 2 and A.shape[1] != 2:
+            num_nodes = ops.shape(features)[-2]
+            a_dense = ops.zeros((num_nodes, num_nodes), dtype=features.dtype)
+            indices = ops.transpose(A, axes=[1, 0])
+            updates = ops.ones(shape=(ops.shape(A)[1],), dtype=features.dtype)
+            A = ops.scatter_update(a_dense, indices, updates)
+
+        was_2d = len(ops.shape(features)) == 2
+        if was_2d:
+            features = ops.expand_dims(features, 0)
+        if len(ops.shape(A)) == 2:
+            A = ops.expand_dims(A, 0)
+
+        # Calculate the layer operation of GCN
+
+        h_graph = dot((A, features), axes=1)
+        b = ops.shape(h_graph)[0]
+        kernel = repeat(self.kernel, b, axis=0)
+        output = dot((h_graph, kernel), axes=(-1, 1))
+
+        # Add optional bias & apply activation
+        if self.bias is not None:
+            output += self.bias
+        output = self.activation(output)
+
+        if was_2d:
+            output = ops.squeeze(output, 0)
+
+        return output

k3_node/layers/conv/gcn2_conv.py

@@ -0,0 +1,126 @@
+import math
+from typing import Optional, Union, Tuple
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm, is_tracing
+
+
+class GCN2Conv(MessagePassing):
+    r"""The graph convolutional operator from the `"Simple and Deep Graph
+    Convolutional Networks" <https://arxiv.org/abs/2007.02133>`_ paper.
+
+    Args:
+        channels: Size of each input and output sample.
+        alpha: The strength of the initial residual connection :math:`\alpha`.
+        theta: The hyperparameter for the identity mapping :math:`\theta`.
+            (default: :obj:`None`)
+        layer: The layer index :math:`l`. (default: :obj:`None`)
+        shared_weights: If set to :obj:`True`, will use the same weights
+            for :math:`\mathbf{X}` and :math:`\mathbf{X}_0`. (default: :obj:`True`)
+        cached: If set to :obj:`True`, will cache the computation of normalization
+            coefficients. (default: ``False``)
+        add_self_loops: If set to :obj:`False`, will not add self-loops.
+            (default: ``True``)
+        normalize: Whether to apply symmetric normalization. (default: ``True``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GCN2Conv
+
+        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
+
+        x_0 = x  # initial node representations (from the first layer)
+        layer = GCN2Conv(channels=8, alpha=0.1, theta=0.5, layer=1)
+        out = layer(x, x_0, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(
+        self,
+        channels: int,
+        alpha: float,
+        theta: Optional[float] = None,
+        layer: Optional[int] = None,
+        shared_weights: bool = True,
+        cached: bool = False,
+        add_self_loops: bool = True,
+        normalize: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.channels = channels
+        self.alpha = alpha
+        self.beta = 1.0
+        if theta is not None and layer is not None:
+            self.beta = math.log(theta / layer + 1.0)
+        self.cached = cached
+        self.normalize = normalize
+        self.add_self_loops = add_self_loops
+        self.shared_weights = shared_weights
+
+        self._cached_edge_index = None
+        self._cached_norm = None
+
+    def build(self, input_shape):
+        self.weight1 = self.add_weight(
+            shape=(self.channels, self.channels),
+            initializer="glorot_uniform",
+            name="weight1",
+        )
+        if not self.shared_weights:
+            self.weight2 = self.add_weight(
+                shape=(self.channels, self.channels),
+                initializer="glorot_uniform",
+                name="weight2",
+            )
+        else:
+            self.weight2 = None
+        self.built = True
+
+    def call(self, x, x_0, edge_index=None, edge_weight=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, x_0, edge_index = x[0], x[1], x[2]
+
+        if self.normalize:
+            if self.cached and self._cached_edge_index is not None:
+                edge_index = self._cached_edge_index
+                edge_weight = self._cached_norm
+            else:
+                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=self.add_self_loops,
+                    flow=self.flow,
+                    dtype=x.dtype,
+                )
+                if self.cached and not is_tracing(edge_index):
+                    self._cached_edge_index = edge_index
+                    self._cached_norm = edge_weight
+
+        h = self.propagate(edge_index, x=x, edge_weight=edge_weight)
+        h = (1.0 - self.alpha) * h
+        h_0 = self.alpha * x_0
+
+        if self.weight2 is None:
+            combined = h + h_0
+            out = (1.0 - self.beta) * combined + self.beta * ops.matmul(combined, self.weight1)
+        else:
+            term1 = (1.0 - self.beta) * h + self.beta * ops.matmul(h, self.weight1)
+            term2 = (1.0 - self.beta) * h_0 + self.beta * ops.matmul(h_0, self.weight2)
+            out = term1 + term2
+
+        return out
+
+    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/gcn_conv.py

@@ -0,0 +1,135 @@
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm, is_tracing
+
+
+class GCNConv(MessagePassing):
+    r"""The graph convolutional operator from the `"Semi-supervised
+    Classification with Graph Convolutional Networks"
+    <https://arxiv.org/abs/1609.02907>`_ paper.
+
+    .. math::
+        \mathbf{X}^{\prime} = \mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}}
+        \mathbf{\hat{D}}^{-1/2} \mathbf{X} \mathbf{\Theta}
+
+    Args:
+        in_channels: Size of each input sample.
+        out_channels: Size of each output sample.
+        improved: If set to :obj:`True`, the layer computes
+            :math:`\mathbf{\hat{A}} = \mathbf{A} + 2 \mathbf{I}`.
+            (default: :obj:`False`)
+        cached: If set to :obj:`True`, the layer will cache the computation of
+            :math:`\mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}} \mathbf{\hat{D}}^{-1/2}`.
+            (default: :obj:`False`)
+        add_self_loops: If set to :obj:`False`, will not add
+            self-loops to the input graph. (default: :obj:`True`)
+        normalize: Whether to add self-loops and compute
+            symmetric normalization coefficients on the fly.
+            (default: :obj:`True`)
+        bias: If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GCNConv
+
+        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 = GCNConv(in_channels=8, out_channels=16)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        improved: bool = False,
+        cached: bool = False,
+        add_self_loops: bool = True,
+        normalize: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.improved = improved
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.normalize = normalize
+        self.use_bias = bias
+
+        self.lin = layers.Dense(out_channels, use_bias=False)
+        self.bias = None
+        self._cached_edge_index = None
+        self._cached_norm = None
+
+    def build(self, input_shape):
+        if isinstance(input_shape, (tuple, list)) and len(input_shape) > 0 and isinstance(input_shape[0], (tuple, list)):
+            feat_shape = input_shape[0]
+        else:
+            feat_shape = input_shape
+        self.lin.build(feat_shape)
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, **kwargs):
+        # Handle legacy call conv((x, adj))
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        if not self.built:
+            feat_shape = x.shape if hasattr(x, "shape") and x.shape is not None else (None, self.in_channels)
+            self.build(feat_shape)
+
+        # Handle dense adjacency [N, N]
+        e_shape = getattr(edge_index, "shape", None)
+        if e_shape is not None and len(e_shape) == 2 and e_shape[0] is not None and e_shape[1] is not None and e_shape[0] > 2 and e_shape[0] == e_shape[1]:
+            where_adj = ops.where(edge_index != 0)
+            where_adj = where_adj if not isinstance(where_adj, list) else where_adj
+            edge_weight = ops.take(edge_index, where_adj[0] * e_shape[1] + where_adj[1]) if edge_weight is None else edge_weight
+            edge_index = ops.stack([where_adj[0], where_adj[1]], axis=0)
+
+        if self.normalize:
+            if self.cached and self._cached_edge_index is not None:
+                edge_index = self._cached_edge_index
+                edge_weight = self._cached_norm
+            else:
+                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,
+                    improved=self.improved,
+                    add_self_loops=self.add_self_loops,
+                    flow=self.flow,
+                    dtype=x.dtype,
+                )
+                if self.cached and not is_tracing(edge_index):
+                    self._cached_edge_index = edge_index
+                    self._cached_norm = edge_weight
+
+        x = self.lin(x)
+        out = self.propagate(edge_index, x=x, edge_weight=edge_weight)
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    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/gen_conv.py

@@ -0,0 +1,163 @@
+from typing import Optional, Union, Tuple
+import keras
+from keras import ops
+from keras.layers import Dense, BatchNormalization, LayerNormalization
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.aggr import SoftmaxAggregation, PowerMeanAggregation
+
+
+class GENConv(MessagePassing):
+    r"""The generalized graph convolution operator from the `"DeeperGCN: All
+    You Need to Train Deeper GCNs" <https://arxiv.org/abs/2006.07739>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GENConv
+
+        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 = GENConv(in_channels=8, out_channels=16, edge_dim=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: int,
+        aggr: str = "softmax",
+        t: float = 1.0,
+        learn_t: bool = False,
+        p: float = 1.0,
+        learn_p: bool = False,
+        msg_norm: bool = False,
+        learn_msg_scale: bool = False,
+        norm: Optional[str] = "batch",
+        num_layers: int = 2,
+        expansion: int = 2,
+        eps: float = 1e-7,
+        bias: bool = False,
+        edge_dim: Optional[int] = None,
+        **kwargs,
+    ):
+        if aggr in ("softmax", "softmax_sg"):
+            aggr_module = SoftmaxAggregation(t=t, learn=learn_t)
+        elif aggr in ("power", "powermean"):
+            aggr_module = PowerMeanAggregation(p=p, learn=learn_p)
+        else:
+            aggr_module = aggr
+
+        super().__init__(aggr=aggr_module, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.eps = eps
+        self.edge_dim = edge_dim
+        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
+
+        if self.in_channels_l != out_channels:
+            self.lin_src = Dense(out_channels, use_bias=bias)
+        else:
+            self.lin_src = None
+
+        if edge_dim is not None and edge_dim != out_channels:
+            self.lin_edge = Dense(out_channels, use_bias=bias)
+        else:
+            self.lin_edge = None
+
+        if self.in_channels_r != out_channels:
+            self.lin_dst = Dense(out_channels, use_bias=bias)
+        else:
+            self.lin_dst = None
+
+        # MLP
+        self.mlp_layers = []
+        channels = [out_channels]
+        for _ in range(num_layers - 1):
+            channels.append(out_channels * expansion)
+        channels.append(out_channels)
+
+        for i in range(len(channels) - 1):
+            self.mlp_layers.append(Dense(channels[i + 1], use_bias=bias))
+            if i < len(channels) - 2:
+                if norm == "batch":
+                    self.mlp_layers.append(BatchNormalization(momentum=0.9, epsilon=1e-5))
+                elif norm == "layer":
+                    self.mlp_layers.append(LayerNormalization())
+                self.mlp_layers.append(keras.layers.ReLU())
+
+    def build(self, input_shape=None):
+        if self.lin_src is not None:
+            self.lin_src.build((None, self.in_channels_l))
+        if self.lin_dst is not None:
+            self.lin_dst.build((None, self.in_channels_r))
+        if self.lin_edge is not None:
+            self.lin_edge.build((None, self.edge_dim))
+        curr_dim = self.out_channels
+        for layer in self.mlp_layers:
+            if hasattr(layer, "build"):
+                layer.build((None, curr_dim))
+            if hasattr(layer, "units"):
+                curr_dim = layer.units
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_attr=None, training=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_l, x_r = x
+        else:
+            x_l = x_r = x
+
+        if self.lin_src is not None:
+            x_l = self.lin_src(x_l)
+
+        num_nodes = ops.shape(x_r)[0]
+        out = self.propagate(
+            edge_index,
+            x=(x_l, x_r),
+            edge_attr=edge_attr,
+            size=(ops.shape(x_l)[0], num_nodes),
+        )
+
+        x_dst = x_r
+        if self.lin_dst is not None:
+            x_dst = self.lin_dst(x_dst)
+        out = out + x_dst
+
+        for layer in self.mlp_layers:
+            # Batch norm needs `training` explicitly (Keras does not propagate it on JAX).
+            out = layer(out, training=training) if isinstance(layer, BatchNormalization) else layer(out)
+
+        return out
+
+    def message(self, x_j, edge_attr=None):
+        if edge_attr is not None and self.lin_edge is not None:
+            edge_attr = self.lin_edge(edge_attr)
+        msg = x_j if edge_attr is None else x_j + edge_attr
+        return ops.relu(msg) + self.eps
+