k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/gps_conv.py

@@ -0,0 +1,153 @@
+from typing import Optional
+import numpy as np
+import keras
+from keras import ops
+from k3_node.ops.segment import segment_sum
+
+
+def to_dense_batch(x, batch=None, batch_size=None):
+    """Differentiable dense batching; see :func:`k3_node.layers.aggr.to_dense_batch`."""
+    from k3_node.layers.aggr.base import to_dense_batch as _to_dense_batch
+
+    if batch is None:
+        return ops.expand_dims(x, axis=0), ops.ones((1, ops.shape(x)[0]), dtype="bool")
+    return _to_dense_batch(x, batch, dim_size=batch_size)
+
+
+class GPSConv(keras.layers.Layer):
+    r"""The general, powerful, scalable (GPS) graph transformer layer from the
+    `"Recipe for a General, Powerful, Scalable Graph Transformer"
+    <https://arxiv.org/abs/2205.12454>`_ paper.
+
+    Args:
+        channels (int): Size of each input sample.
+        conv (keras.layers.Layer, optional): The local message passing layer.
+        heads (int, optional): Number of multi-head-attentions. (default: :obj:`1`)
+        dropout (float, optional): Dropout probability. (default: :obj:`0.0`)
+        act (str, optional): Activation function. (default: :obj:`"relu"`)
+        norm (str, optional): Normalization function. (default: :obj:`"batch_norm"`)
+
+    Call arguments: ``x``, ``edge_index``, ``batch`` (the graph of every node) and ``batch_size``
+    (the number of graphs; pass it when training compiled so that shapes are static), plus any
+    arguments of the local ``conv`` such as ``edge_attr``.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GPSConv, 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
+
+        batch = np.repeat([0, 1], 5)  # two graphs with 5 nodes each
+        layer = GPSConv(channels=8, conv=GCNConv(8, 8), heads=2)
+        out = layer(x, edge_index, batch=batch)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        conv: Optional[keras.layers.Layer] = None,
+        heads: int = 1,
+        dropout: float = 0.0,
+        act: str = "relu",
+        norm: Optional[str] = "batch_norm",
+        **kwargs,
+    ):
+        conv = kwargs.pop("local_gnn", conv)
+        super().__init__(**kwargs)
+
+        self.channels = channels
+        self.conv = conv
+        self.heads = heads
+        self.dropout_rate = dropout
+        self.act = act
+        self.norm_name = norm
+
+        self.attn = keras.layers.MultiHeadAttention(
+            num_heads=heads,
+            key_dim=channels // heads,
+            value_dim=channels // heads,
+            output_shape=channels,
+        )
+
+        self.mlp_l1 = keras.layers.Dense(channels * 2)
+        self.mlp_l2 = keras.layers.Dense(channels)
+        self.dropout = keras.layers.Dropout(dropout)
+
+        if norm == "batch_norm":
+            self.norm1 = keras.layers.BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5) if conv is not None else None
+            self.norm2 = keras.layers.BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
+            self.norm3 = keras.layers.BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
+        elif norm == "layer_norm":
+            self.norm1 = keras.layers.LayerNormalization(axis=-1) if conv is not None else None
+            self.norm2 = keras.layers.LayerNormalization(axis=-1)
+            self.norm3 = keras.layers.LayerNormalization(axis=-1)
+        else:
+            self.norm1 = None
+            self.norm2 = None
+            self.norm3 = None
+
+    def build(self, input_shape=None):
+        if self.conv is not None and hasattr(self.conv, "build") and not self.conv.built:
+            self.conv.build(input_shape)
+        if not self.attn.built:
+            self.attn.build((None, None, self.channels), (None, None, self.channels))
+        if not self.mlp_l1.built:
+            self.mlp_l1.build((None, self.channels))
+        if not self.mlp_l2.built:
+            self.mlp_l2.build((None, self.channels * 2))
+        super().build(input_shape)
+
+    def call(self, x, edge_index, batch=None, batch_size=None, training=None, **kwargs):
+        if not self.built:
+            self.build((None, self.channels))
+
+        # `training` is forwarded explicitly: Keras does not propagate it to nested layers on JAX.
+        hs = []
+        if self.conv is not None:
+            h = self.conv(x, edge_index, training=training, **kwargs)
+            h = self.dropout(h, training=training)
+            h = h + x
+            if self.norm1 is not None:
+                h = self.norm1(h, training=training)
+            hs.append(h)
+
+        # Global attention
+        # `batch_size` (the number of graphs) keeps shapes static when compiled
+        h_dense, mask = to_dense_batch(x, batch, batch_size)
+        # Attention mask for Keras: shape (B, 1, max_nodes)
+        attn_mask = ops.expand_dims(mask, axis=1)
+        attn_out = self.attn(h_dense, h_dense, attention_mask=attn_mask, training=training)
+
+        # Unpack dense batch to original flat shape (static shapes: jit-friendly)
+        if batch is None:
+            h_global = attn_out[0]
+        else:
+            from k3_node.layers.aggr.base import from_dense_batch
+
+            h_global = from_dense_batch(attn_out, batch)
+
+        h_global = self.dropout(h_global, training=training)
+        h_global = h_global + x
+        if self.norm2 is not None:
+            h_global = self.norm2(h_global, training=training)
+        hs.append(h_global)
+
+        # Combine local and global
+        if len(hs) > 1:
+            out = hs[0] + hs[1]
+        else:
+            out = hs[0]
+
+        # MLP
+        mlp_h = self.dropout(ops.relu(self.mlp_l1(out)), training=training)
+        mlp_out = self.dropout(self.mlp_l2(mlp_h), training=training)
+        out = out + mlp_out
+
+        if self.norm3 is not None:
+            out = self.norm3(out, training=training)
+
+        return out

k3_node/layers/conv/graph_attention.py

@@ -0,0 +1,262 @@
+# ported from stellargraph
+from keras import ops
+from keras import activations, constraints, initializers, regularizers
+from keras.layers import Layer, LeakyReLU, Dropout
+
+
+class GraphAttention(Layer):
+    """
+    `k3_node.layers.GraphAttention`
+    Implementation of Graph Attention (GAT) layer
+
+    Args:
+        units: Positive integer, dimensionality of the output space.
+        attn_heads: Positive integer, number of attention heads.
+        attn_heads_reduction: {'concat', 'average'} Method for reducing attention heads.
+        in_dropout_rate: Dropout rate applied to the input (node features).
+        attn_dropout_rate: Dropout rate applied to attention coefficients.
+        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.
+        saliency_map_support: Whether to support saliency map calculations.
+        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.
+        attn_kernel_initializer: Initializer for the attention kernel weights matrix.
+        attn_kernel_regularizer: Regularizer for the attention kernel weights matrix.
+        attn_kernel_constraint: Constraint for the attention kernel weights matrix.
+        **kwargs: Additional arguments to pass to the `Layer` superclass.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GraphAttention
+
+        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 = GraphAttention(units=16, attn_heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 2)
+        ```
+    """
+    def __init__(
+        self,
+        units,
+        attn_heads=1,
+        attn_heads_reduction="concat",  # {'concat', 'average'}
+        in_dropout_rate=0.0,
+        attn_dropout_rate=0.0,
+        activation="relu",
+        use_bias=True,
+        final_layer=None,
+        saliency_map_support=False,
+        kernel_initializer="glorot_uniform",
+        kernel_regularizer=None,
+        kernel_constraint=None,
+        bias_initializer="zeros",
+        bias_regularizer=None,
+        bias_constraint=None,
+        attn_kernel_initializer="glorot_uniform",
+        attn_kernel_regularizer=None,
+        attn_kernel_constraint=None,
+        **kwargs,
+    ):
+        if attn_heads_reduction not in {"concat", "average"}:
+            raise ValueError(
+                "{}: Possible heads reduction methods: concat, average; received {}".format(
+                    type(self).__name__, attn_heads_reduction
+                )
+            )
+
+        if isinstance(attn_heads, int) and attn_heads > 1 and "out_channels" not in kwargs:
+            self.in_channels = units
+            units = attn_heads
+            attn_heads = 1
+
+        self.units = units  # Number of output features (F' in the paper)
+        self.attn_heads = attn_heads  # Number of attention heads (K in the paper)
+        self.attn_heads_reduction = attn_heads_reduction  # Eq. 5 and 6 in the paper
+        self.in_dropout_rate = in_dropout_rate  # dropout rate for node features
+        self.attn_dropout_rate = attn_dropout_rate  # dropout rate for attention coefs
+        self.activation = activations.get(activation)  # Eq. 4 in the paper
+        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.saliency_map_support = saliency_map_support
+        # Populated by build()
+        self.kernels = []  # Layer kernels for attention heads
+        self.biases = []  # Layer biases for attention heads
+        self.attn_kernels = []  # Attention kernels for attention heads
+
+        if attn_heads_reduction == "concat":
+            # Output will have shape (..., K * F')
+            self.output_dim = self.units * self.attn_heads
+        else:
+            # Output will have shape (..., F')
+            self.output_dim = self.units
+
+        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)
+        self.attn_kernel_initializer = initializers.get(attn_kernel_initializer)
+        self.attn_kernel_regularizer = regularizers.get(attn_kernel_regularizer)
+        self.attn_kernel_constraint = constraints.get(attn_kernel_constraint)
+
+        super().__init__(**kwargs)
+        self.in_dropout = Dropout(in_dropout_rate)  # created once, applied with `training`
+        self.attn_dropout = Dropout(attn_dropout_rate)
+
+    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
+
+        # Variables to support integrated gradients
+        self.delta = self.add_weight(
+            name="ig_delta", shape=(), trainable=False, initializer=initializers.ones()
+        )
+        self.non_exist_edge = self.add_weight(
+            name="ig_non_exist_edge",
+            shape=(),
+            trainable=False,
+            initializer=initializers.zeros(),
+        )
+
+        # Initialize weights for each attention head
+        for head in range(self.attn_heads):
+            # Layer kernel
+            kernel = self.add_weight(
+                shape=(input_dim, self.units),
+                initializer=self.kernel_initializer,
+                regularizer=self.kernel_regularizer,
+                constraint=self.kernel_constraint,
+                name="kernel_{}".format(head),
+            )
+            self.kernels.append(kernel)
+
+            # # Layer bias
+            if self.use_bias:
+                bias = self.add_weight(
+                    shape=(self.units,),
+                    initializer=self.bias_initializer,
+                    regularizer=self.bias_regularizer,
+                    constraint=self.bias_constraint,
+                    name="bias_{}".format(head),
+                )
+                self.biases.append(bias)
+
+            # Attention kernels
+            attn_kernel_self = self.add_weight(
+                shape=(self.units, 1),
+                initializer=self.attn_kernel_initializer,
+                regularizer=self.attn_kernel_regularizer,
+                constraint=self.attn_kernel_constraint,
+                name="attn_kernel_self_{}".format(head),
+            )
+            attn_kernel_neighs = self.add_weight(
+                shape=(self.units, 1),
+                initializer=self.attn_kernel_initializer,
+                regularizer=self.attn_kernel_regularizer,
+                constraint=self.attn_kernel_constraint,
+                name="attn_kernel_neigh_{}".format(head),
+            )
+            self.attn_kernels.append([attn_kernel_self, attn_kernel_neighs])
+        self.built = True
+
+    def call(self, inputs, A=None, training=None, **kwargs):
+        if A is not None:
+            X = inputs
+        elif isinstance(inputs, (list, tuple)):
+            X = inputs[0]
+            A = inputs[1]
+        else:
+            X, 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(X)[-2]
+            a_dense = ops.zeros((num_nodes, num_nodes), dtype=X.dtype)
+            indices = ops.transpose(A, axes=[1, 0])
+            updates = ops.ones(shape=(ops.shape(A)[1],), dtype=X.dtype)
+            A = ops.scatter_update(a_dense, indices, updates)
+
+        assert len(ops.shape(A)) == 2, f"Adjacency matrix A should be 2-D"
+        N = ops.shape(A)[-1]
+
+        outputs = []
+        for head in range(self.attn_heads):
+            kernel = self.kernels[head]  # W in the paper (F x F')
+            attention_kernel = self.attn_kernels[
+                head
+            ]  # Attention kernel a in the paper (2F' x 1)
+
+            # Compute inputs to attention network
+
+            features = ops.dot(X, kernel)  # (N x F')
+
+            # Compute feature combinations
+            # Note: [[a_1], [a_2]]^T [[Wh_i], [Wh_2]] = [a_1]^T [Wh_i] + [a_2]^T [Wh_j]
+            attn_for_self = ops.dot(
+                features, attention_kernel[0]
+            )  # (N x 1), [a_1]^T [Wh_i]
+            attn_for_neighs = ops.dot(
+                features, attention_kernel[1]
+            )  # (N x 1), [a_2]^T [Wh_j]
+
+            # Attention head a(Wh_i, Wh_j) = a^T [[Wh_i], [Wh_j]]
+            dense = attn_for_self + ops.transpose(
+                attn_for_neighs
+            )  # (N x N) via broadcasting
+
+            dense = LeakyReLU(0.2)(dense)
+
+            if not self.saliency_map_support:
+                mask = -10e9 * (1.0 - A)
+                dense += mask
+                dense = ops.softmax(dense)  # (N x N), Eq. 3 of the paper
+
+            else:
+                # dense = dense - tf.reduce_max(dense)
+                # GAT with support for saliency calculations
+                W = (self.delta * A) * ops.exp(
+                    dense - ops.max(dense, axis=1, keepdims=True)
+                ) * (1 - self.non_exist_edge) + self.non_exist_edge * (
+                    A + self.delta * (ops.ones((N, N)) - A) + ops.eye(N)
+                ) * ops.exp(
+                    dense - ops.max(dense, axis=1, keepdims=True)
+                )
+                dense = W / ops.sum(W, axis=1, keepdims=True)
+
+            # Apply dropout to features and attention coefficients
+            dropout_feat = self.in_dropout(features, training=training)  # (N x F')
+            dropout_attn = self.attn_dropout(dense, training=training)  # (N x N)
+
+            # Linear combination with neighbors' features [YT: see Eq. 4]
+            node_features = ops.dot(dropout_attn, dropout_feat)  # (N x F')
+
+            if self.use_bias:
+                node_features = ops.add(node_features, self.biases[head])
+
+            # Add output of attention head to final output
+            outputs.append(node_features)
+
+        # Aggregate the heads' output according to the reduction method
+        if self.attn_heads_reduction == "concat":
+            output = ops.concatenate(outputs, axis=1)  # (N x KF')
+        else:
+            output = ops.mean(ops.stack(outputs), axis=0)  # N x F')
+
+        output = self.activation(output)
+
+        return output

k3_node/layers/conv/graph_conv.py

@@ -0,0 +1,84 @@
+from typing import Union, Tuple
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class GraphConv(MessagePassing):
+    r"""The graph neural network operator from the `"Weisfeiler and Leman Go
+    Neural: Higher-order Graph Neural Networks"
+    <https://arxiv.org/abs/1810.02244>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample, or a tuple for bipartite graphs.
+        out_channels: Size of each output sample.
+        aggr: The aggregation scheme to use (``"add"``, ``"mean"``, ``"max"``).
+            (default: ``"add"``)
+        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 GraphConv
+
+        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 = GraphConv(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: Union[int, Tuple[int, int]],
+        out_channels: int,
+        aggr: str = "add",
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.use_bias = bias
+
+        self.lin_rel = layers.Dense(out_channels, use_bias=bias)
+        self.lin_root = layers.Dense(out_channels, use_bias=False)
+
+    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_l = input_shape[0][-1]
+            in_channels_r = input_shape[1][-1] if len(input_shape) > 1 and input_shape[1] is not None else in_channels_l
+        else:
+            in_channels_l = input_shape[-1]
+            in_channels_r = input_shape[-1]
+
+        self.lin_rel.build((None, in_channels_l))
+        self.lin_root.build((None, in_channels_r))
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, size=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
+            x_dst = x
+        else:
+            x_src, x_dst = x[0], x[1]
+
+        out = self.propagate(edge_index, x=(x_src, x_dst), edge_weight=edge_weight, size=size)
+        out = self.lin_rel(out)
+        if x_dst is not None:
+            out = out + self.lin_root(x_dst)
+        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/gravnet_conv.py

@@ -0,0 +1,93 @@
+from typing import Optional, Union, Tuple
+import keras
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.pool.knn import knn_graph
+
+
+class GravNetConv(MessagePassing):
+    r"""The GravNet operator from the `"Learning Representations of Irregular
+    Particle-Detector Geometry with Distance-Weighted Graph Networks"
+    <https://arxiv.org/abs/1902.07987>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GravNetConv
+
+        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
+
+        # Neighbors are found by k-NN in a learned space, so no edge_index is needed
+        layer = GravNetConv(in_channels=8, out_channels=16, space_dimensions=3, propagate_dimensions=4, k=3)
+        out = layer(x)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        space_dimensions: int,
+        propagate_dimensions: int,
+        k: int,
+        num_workers: Optional[int] = None,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "mean")
+        super().__init__(node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.space_dimensions = space_dimensions
+        self.propagate_dimensions = propagate_dimensions
+        self.k = k
+
+        self.lin_s = Dense(space_dimensions, use_bias=True)
+        self.lin_h = Dense(propagate_dimensions, use_bias=True)
+        self.lin_out1 = Dense(out_channels, use_bias=True)
+        self.lin_out2 = Dense(out_channels, use_bias=True)
+
+    def build(self, input_shape=None):
+        self.lin_s.build((None, self.in_channels))
+        self.lin_h.build((None, self.in_channels))
+        self.lin_out1.build((None, self.in_channels))
+        self.lin_out2.build((None, self.propagate_dimensions))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, **kwargs):
+        if isinstance(inputs, (list, tuple)):
+            x = inputs[0]
+        else:
+            x = inputs
+
+        if not self.built:
+            self.build()
+
+        s = self.lin_s(x)
+        h = self.lin_h(x)
+
+        if edge_index is None:
+            edge_index = knn_graph(s, k=self.k)
+
+        edge_index = ops.cast(edge_index, "int32")
+        s_src = ops.take(s, edge_index[0], axis=0)
+        s_dst = ops.take(s, edge_index[1], axis=0)
+        dist_sq = ops.sum(ops.square(s_src - s_dst), axis=-1)
+        edge_weight = ops.exp(-10.0 * dist_sq)
+
+        num_nodes = ops.shape(x)[0]
+        out = self.propagate(
+            edge_index,
+            x=h,
+            edge_weight=edge_weight,
+            size=(num_nodes, num_nodes),
+        )
+
+        return self.lin_out1(x) + self.lin_out2(out)
+
+    def message(self, x_j, edge_weight):
+        return ops.expand_dims(edge_weight, -1) * x_j
+