k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/arma_conv.py

@@ -0,0 +1,231 @@
+from typing import Optional, Union, Callable
+from keras import ops, activations
+from keras.layers import Dropout
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm
+
+
+class ARMAConv(MessagePassing):
+    r"""The ARMA graph convolutional operator from the `"Graph Neural Networks
+    with Convolutional ARMA Filters" <https://arxiv.org/abs/1901.01343>`_
+    paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import ARMAConv
+
+        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 = ARMAConv(in_channels=8, out_channels=16, num_stacks=1, num_layers=1)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        num_stacks: int = 1,
+        num_layers: int = 1,
+        shared_weights: bool = False,
+        act: Union[str, Callable, None] = "relu",
+        dropout: float = 0.0,
+        bias: bool = True,
+        # Spektral compatibility arguments:
+        channels: Optional[int] = None,
+        order: Optional[int] = None,
+        iterations: Optional[int] = None,
+        share_weights: Optional[bool] = None,
+        gcn_activation: Optional[str] = None,
+        dropout_rate: Optional[float] = None,
+        activation: Optional[str] = None,
+        use_bias: Optional[bool] = None,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(node_dim=0, **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 order is not None:
+            num_stacks = order
+        if iterations is not None:
+            num_layers = iterations
+        if share_weights is not None:
+            shared_weights = share_weights
+        if gcn_activation is not None:
+            act = gcn_activation
+        if dropout_rate is not None:
+            dropout = dropout_rate
+        if use_bias is not None:
+            bias = use_bias
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_stacks = num_stacks
+        self.num_layers = num_layers
+        self.shared_weights = shared_weights
+        self.act = activations.get(act) if act is not None else None
+        self.dropout_rate = dropout
+        self.use_bias = bias
+
+        K, T, F_in, F_out = num_stacks, num_layers, in_channels, out_channels
+        T_w = 1 if self.shared_weights else T
+
+        self.weight = self.add_weight(
+            shape=(max(1, T_w - 1), K, F_out, F_out),
+            initializer="glorot_uniform",
+            name="weight",
+        )
+        if in_channels is not None and in_channels != -1:
+            self.init_weight = self.add_weight(
+                shape=(K, F_in, F_out),
+                initializer="glorot_uniform",
+                name="init_weight",
+            )
+            self.root_weight = self.add_weight(
+                shape=(T_w, K, F_in, F_out),
+                initializer="glorot_uniform",
+                name="root_weight",
+            )
+        else:
+            self.init_weight = None
+            self.root_weight = None
+
+        if bias:
+            self.bias = self.add_weight(
+                shape=(T_w, K, 1, F_out),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+        self.dropout = Dropout(dropout)
+
+    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 is None or self.in_channels == -1) and dim is not None:
+                self.in_channels = dim
+        if self.in_channels is not None and self.in_channels != -1:
+            K, T, F_in, F_out = self.num_stacks, self.num_layers, self.in_channels, self.out_channels
+            T_w = 1 if self.shared_weights else T
+            if self.init_weight is None:
+                self.init_weight = self.add_weight(
+                    shape=(K, F_in, F_out),
+                    initializer="glorot_uniform",
+                    name="init_weight",
+                )
+            if self.root_weight is None:
+                self.root_weight = self.add_weight(
+                    shape=(T_w, K, F_in, F_out),
+                    initializer="glorot_uniform",
+                    name="root_weight",
+                )
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_weight=None, training=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, edge_index, edge_weight = 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 self.in_channels is None or self.in_channels == -1:
+            self.build((None, ops.shape(x)[-1]))
+
+        # Legacy adj check
+        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_weight = edge_index.values
+                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)
+                edge_weight = ops.take(adj, row * ops.shape(adj)[1] + col)
+
+        num_nodes = ops.shape(x)[0]
+        edge_index, edge_weight = gcn_norm(
+            edge_index,
+            edge_weight,
+            num_nodes=num_nodes,
+            add_self_loops=False,
+            dtype=x.dtype,
+        )
+
+        # PyG: x = x.unsqueeze(-3) -> (1, N, F_in)
+        # out = x
+        # for t in range(num_layers):
+        #   if t == 0: out = out @ init_weight   (K, F_in, F_out) -> (K, N, F_out)
+        #   else: out = out @ weight[t-1]       (K, F_out, F_out) -> (K, N, F_out)
+        #   out = propagate(edge_index, x=out, edge_weight=edge_weight)
+        #   root = dropout(x) @ root_weight[t]  (K, F_in, F_out) -> (K, N, F_out)
+        #   out = out + root
+        #   if bias: out = out + bias[t]
+        #   if act: out = act(out)
+        # return out.mean(dim=-3)
+        K, T, F_in, F_out = self.num_stacks, self.num_layers, self.in_channels, self.out_channels
+
+        # out shape: (N, K, F_out)
+        out = None
+        for t in range(self.num_layers):
+            w_idx = 0 if self.shared_weights else t
+            if t == 0:
+                # x: (N, F_in), init_weight: (K, F_in, F_out)
+                # out: (K, N, F_out)
+                out = ops.einsum("nf,kfo->kno", x, self.init_weight)
+            else:
+                w = self.weight[0 if self.shared_weights else t - 1]
+                out = ops.einsum("kno,kof->knf", out, w)
+
+            # Transpose to (N, K, F_out) so node_dim=0
+            out_n = ops.transpose(out, (1, 0, 2))
+            out_prop = self.propagate(edge_index, x=out_n, edge_weight=edge_weight, size=(num_nodes, num_nodes))
+            out = ops.transpose(out_prop, (1, 0, 2))  # (K, N, F_out)
+
+            root_x = self.dropout(x, training=training)
+            root = ops.einsum("nf,kfo->kno", root_x, self.root_weight[w_idx])
+            out = out + root
+
+            if self.bias is not None:
+                # bias[w_idx]: (K, 1, F_out)
+                out = out + self.bias[w_idx]
+
+            if self.act is not None:
+                out = self.act(out)
+
+        # out: (K, N, F_out) -> mean over K -> (N, F_out)
+        return ops.mean(out, axis=0)
+
+    def message(self, x_j, edge_weight=None):
+        return x_j if edge_weight is None else ops.expand_dims(ops.expand_dims(edge_weight, -1), -1) * x_j

k3_node/layers/conv/cg_conv.py

@@ -0,0 +1,92 @@
+from typing import Union, Tuple
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class CGConv(MessagePassing):
+    r"""The Crystal Graph Convolutional operator from the
+    `"Crystal Graph Convolutional Neural Networks for an Accurate and
+    Interpretable Prediction of Material Properties"
+    <https://arxiv.org/abs/1710.10324>`_ paper.
+
+    Args:
+        channels: Size of each input sample, or a tuple for bipartite graphs.
+        dim: Edge feature dimensionality. (default: ``0``)
+        aggr: The aggregation scheme to use (``"add"``, ``"mean"``, ``"max"``).
+            (default: ``"add"``)
+        batch_norm: If set to :obj:`True`, will apply batch normalization.
+            (default: ``False``)
+        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 CGConv
+
+        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 = CGConv(channels=8, dim=3)
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        channels: Union[int, Tuple[int, int]],
+        dim: int = 0,
+        aggr: str = "add",
+        batch_norm: bool = False,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+        self.channels = channels
+        self.dim = dim
+        self.batch_norm = batch_norm
+        self.use_bias = bias
+
+        if isinstance(channels, int):
+            self.in_channels_src = channels
+            self.in_channels_dst = channels
+        else:
+            self.in_channels_src = channels[0]
+            self.in_channels_dst = channels[1]
+
+        in_dim = self.in_channels_src + self.in_channels_dst + dim
+        self.lin_f = layers.Dense(self.in_channels_dst, use_bias=bias)
+        self.lin_s = layers.Dense(self.in_channels_dst, use_bias=bias)
+        self.bn = layers.BatchNormalization(momentum=0.9, epsilon=1e-5) if batch_norm else None
+
+    def build(self, input_shape):
+        in_dim = self.in_channels_src + self.in_channels_dst + self.dim
+        self.lin_f.build((None, in_dim))
+        self.lin_s.build((None, in_dim))
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_attr=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)
+        if self.bn is not None:
+            out = self.bn(out)
+        out = x_dst + out
+        return out
+
+    def message(self, x_i, x_j, edge_attr=None):
+        if edge_attr is None:
+            z = ops.concatenate([x_i, x_j], axis=-1)
+        else:
+            z = ops.concatenate([x_i, x_j, edge_attr], axis=-1)
+        return ops.sigmoid(self.lin_f(z)) * ops.softplus(self.lin_s(z))
+

k3_node/layers/conv/cheb_conv.py

@@ -0,0 +1,137 @@
+from typing import Optional, List
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import get_laplacian
+
+
+class ChebConv(MessagePassing):
+    r"""The Chebyshev spectral graph convolutional operator from the
+    `"Convolutional Neural Networks on Graphs with Fast Localized Spectral
+    Filtering" <https://arxiv.org/abs/1606.09375>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample.
+        out_channels: Size of each output sample.
+        K: Chebyshev filter size :math:`K`.
+        normalization: The normalization scheme for the graph
+            Laplacian (``"sym"``, ``"rw"`` or :obj:`None`). (default: ``"sym"``)
+        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 ChebConv
+
+        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 = ChebConv(in_channels=8, out_channels=16, K=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        K: int,
+        normalization: Optional[str] = "sym",
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        if K <= 0:
+            raise ValueError(f"K must be a positive integer, got {K}")
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.K = K
+        self.normalization = normalization
+        self.use_bias = bias
+
+        self.lins = [layers.Dense(out_channels, use_bias=False) for _ in range(K)]
+
+    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
+        for lin in self.lins:
+            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 __norm__(
+        self,
+        edge_index,
+        num_nodes: Optional[int],
+        edge_weight=None,
+        normalization: Optional[str] = "sym",
+        lambda_max=None,
+        dtype=None,
+    ):
+        edge_index, edge_weight = get_laplacian(
+            edge_index, edge_weight, normalization, dtype, num_nodes
+        )
+        if lambda_max is None:
+            lambda_max = 2.0 * ops.max(edge_weight)
+        else:
+            lambda_max = ops.convert_to_tensor(lambda_max, dtype=edge_weight.dtype)
+
+        edge_weight = (2.0 * edge_weight) / lambda_max
+        edge_weight = ops.where(
+            ops.isinf(edge_weight) | ops.isnan(edge_weight), 0.0, edge_weight
+        )
+
+        loop_mask = edge_index[0] == edge_index[1]
+        edge_weight = ops.where(loop_mask, edge_weight - 1.0, edge_weight)
+        return edge_index, edge_weight
+
+    def call(self, x, edge_index=None, edge_weight=None, lambda_max=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        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, norm = self.__norm__(
+            edge_index,
+            num_nodes,
+            edge_weight,
+            self.normalization,
+            lambda_max,
+            dtype=x.dtype,
+        )
+
+        Tx_0 = x
+        Tx_1 = x
+        out = self.lins[0](Tx_0)
+
+        if len(self.lins) > 1:
+            Tx_1 = self.propagate(edge_index, x=x, norm=norm)
+            out = out + self.lins[1](Tx_1)
+
+        for lin in self.lins[2:]:
+            Tx_2 = self.propagate(edge_index, x=Tx_1, norm=norm)
+            Tx_2 = 2.0 * Tx_2 - Tx_0
+            out = out + lin(Tx_2)
+            Tx_0, Tx_1 = Tx_1, Tx_2
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def message(self, x_j, norm=None):
+        if norm is None:
+            return x_j
+        return ops.expand_dims(norm, -1) * x_j
+

k3_node/layers/conv/cluster_gcn_conv.py

@@ -0,0 +1,102 @@
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import (
+    add_self_loops,
+    degree,
+    extend_mask_for_self_loops,
+    remove_self_loops_masked,
+)
+from k3_node.ops.segment import segment_sum
+
+
+class ClusterGCNConv(MessagePassing):
+    r"""The ClusterGCN graph convolutional operator from the
+    `"Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph
+    Convolutional Networks" <https://arxiv.org/abs/1905.07953>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample.
+        out_channels: Size of each output sample.
+        diag_lambda: Diagonal enhancement coefficient :math:`\lambda`.
+            (default: ``0.0``)
+        add_self_loops: If set to :obj:`False`, will not add self-loops.
+            (default: ``True``)
+        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 ClusterGCNConv
+
+        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 = ClusterGCNConv(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,
+        diag_lambda: float = 0.0,
+        add_self_loops: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.diag_lambda = diag_lambda
+        self.add_self_loops = add_self_loops
+        self.use_bias = bias
+
+        self.lin_out = layers.Dense(out_channels, use_bias=bias)
+        self.lin_root = layers.Dense(out_channels, use_bias=False)
+
+    def build(self, input_shape):
+        feat_shape = input_shape[0] if isinstance(input_shape, (tuple, list)) and isinstance(input_shape[0], (tuple, list)) else input_shape
+        self.lin_out.build(feat_shape)
+        self.lin_root.build(feat_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]
+
+        num_nodes = ops.shape(x)[self.node_dim]
+
+        keep_mask = None
+        if self.add_self_loops:
+            edge_index, _, keep_mask = remove_self_loops_masked(edge_index)
+            edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+            keep_mask = extend_mask_for_self_loops(keep_mask, num_nodes)
+
+        row, col = edge_index[0], edge_index[1]
+        col_cast = ops.cast(col, "int32")
+        if keep_mask is None:
+            deg = degree(col_cast, num_nodes=num_nodes)
+        else:
+            deg = segment_sum(ops.cast(keep_mask, x.dtype), col_cast, num_segments=num_nodes)
+        deg_inv = 1.0 / ops.maximum(ops.cast(deg, x.dtype), 1.0)
+
+        edge_weight = ops.take(deg_inv, col_cast, axis=0)
+        loop_mask = ops.equal(row, col)
+        edge_weight = ops.where(loop_mask, edge_weight + self.diag_lambda * ops.take(deg_inv, col_cast, axis=0), edge_weight)
+        if keep_mask is not None:
+            edge_weight = edge_weight * ops.cast(keep_mask, edge_weight.dtype)
+
+        out = self.propagate(edge_index, x=x, edge_weight=edge_weight)
+        return self.lin_out(out) + self.lin_root(x)
+
+    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/conv.py

@@ -0,0 +1,100 @@
+# ported from spektral
+
+import warnings
+from functools import wraps
+
+from keras import ops, backend
+from keras.layers import Layer
+
+from k3_node.utils import (
+    is_keras_kwarg,
+    is_layer_kwarg,
+    deserialize_kwarg,
+    serialize_kwarg,
+)
+
+
+class Conv(Layer):
+    def __init__(self, **kwargs):
+        unknown = sorted(k for k in kwargs if not (is_keras_kwarg(k) or is_layer_kwarg(k)))
+        if unknown:
+            raise TypeError(f"{type(self).__name__}() got unexpected keyword argument(s): {', '.join(unknown)}")
+        super().__init__(**{k: v for k, v in kwargs.items() if is_keras_kwarg(k)})
+        self.supports_masking = True
+        self.kwargs_keys = []
+        for key in kwargs:
+            if is_layer_kwarg(key):
+                attr = kwargs[key]
+                attr = deserialize_kwarg(key, attr)
+                self.kwargs_keys.append(key)
+                setattr(self, key, attr)
+        self.call = check_dtypes_decorator(self.call)
+
+    def build(self, input_shape):
+        self.built = True
+
+    def call(self, inputs):
+        raise NotImplementedError
+
+    def get_config(self):
+        base_config = super().get_config()
+        keras_config = {}
+        for key in self.kwargs_keys:
+            keras_config[key] = serialize_kwarg(key, getattr(self, key))
+        return {**base_config, **keras_config, **self.config}
+
+    @property
+    def config(self):
+        return {}
+
+    @staticmethod
+    def preprocess(a):
+        return a
+
+
+def check_dtypes_decorator(call):
+    @wraps(call)
+    def _inner_check_dtypes(*args, **kwargs):
+        if len(args) == 0:
+            return call(**kwargs)
+        elif len(args) == 1:
+            inputs = check_dtypes(args[0])
+            return call(inputs, **kwargs)
+        else:
+            checked = check_dtypes(list(args))
+            if isinstance(checked, (list, tuple)) and len(checked) == len(args):
+                return call(*checked, **kwargs)
+            return call(*args, **kwargs)
+
+    return _inner_check_dtypes
+
+
+def check_dtypes(inputs):
+    if not isinstance(inputs, (list, tuple)):
+        return inputs
+    for value in inputs:
+        if not hasattr(value, "dtype"):
+            # It's not a valid tensor.
+            return inputs
+
+    if len(inputs) == 2:
+        x, a = inputs
+        e = None
+    elif len(inputs) == 3:
+        x, a, e = inputs
+    else:
+        return inputs
+
+    # If 'a' is an edge_index of shape (2, E), it must remain integer
+    if hasattr(a, "shape") and len(a.shape) == 2 and a.shape[0] == 2 and a.shape[1] != 2:
+        pass
+    elif backend.is_int_dtype(a.dtype) and backend.is_float_dtype(x.dtype):
+        warnings.warn(
+            f"The adjacency matrix of dtype {a.dtype} is incompatible with the dtype "
+            f"of the node features {x.dtype} and has been automatically cast to "
+            f"{x.dtype}."
+        )
+        a = ops.cast(a, x.dtype)
+
+    output = [_ for _ in [x, a, e] if _ is not None]
+    return output