k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/crystal_conv.py

@@ -0,0 +1,140 @@
+# ported from spektral
+
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class CrystalConv(MessagePassing):
+    """
+    `k3_node.layers.CrystalConv`
+    Implementation of Crystal Graph Convolutional Neural Networks (CGCNN) layer
+
+    Args:
+        channels: The number of channels/units (optional).
+        edge_dim: The dimensionality of edge features (optional).
+        aggregate: Aggregation function to use (one of 'sum', 'mean', 'max').
+        activation: Activation function to use.
+        use_bias: Whether to add a bias to the linear transformation.
+        kernel_initializer: Initializer for the `kernel` weights matrix.
+        bias_initializer: Initializer for the bias vector.
+        kernel_regularizer: Regularizer for the `kernel` weights matrix.
+        bias_regularizer: Regularizer for the bias vector.
+        activity_regularizer: Regularizer for the output.
+        kernel_constraint: Constraint for the `kernel` weights matrix.
+        bias_constraint: Constraint for the bias vector.
+        **kwargs: Additional arguments to pass to the `MessagePassing` superclass. 
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import CrystalConv
+
+        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 = CrystalConv(channels=8, edge_dim=3)
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        channels=None,
+        edge_dim=None,
+        aggregate="sum",
+        activation=None,
+        use_bias=True,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        kernel_regularizer=None,
+        bias_regularizer=None,
+        activity_regularizer=None,
+        kernel_constraint=None,
+        bias_constraint=None,
+        **kwargs,
+    ):
+        
+        super().__init__(
+            aggregate=aggregate,
+            activation=activation,
+            use_bias=use_bias,
+            kernel_initializer=kernel_initializer,
+            bias_initializer=bias_initializer,
+            kernel_regularizer=kernel_regularizer,
+            bias_regularizer=bias_regularizer,
+            activity_regularizer=activity_regularizer,
+            kernel_constraint=kernel_constraint,
+            bias_constraint=bias_constraint,
+            **kwargs,
+        )
+        self.channels = channels
+        self.edge_dim = edge_dim
+
+    def build(self, input_shape=None):
+        layer_kwargs = dict(
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+            kernel_regularizer=self.kernel_regularizer,
+            bias_regularizer=self.bias_regularizer,
+            kernel_constraint=self.kernel_constraint,
+            bias_constraint=self.bias_constraint,
+            dtype=self.dtype,
+        )
+        if self.channels is not None:
+            channels = self.channels
+        elif input_shape is not None:
+            if isinstance(input_shape, (list, tuple)) and len(input_shape) > 0 and isinstance(input_shape[0], (list, tuple)):
+                channels = input_shape[0][-1]
+            elif isinstance(input_shape, (list, tuple)) and len(input_shape) > 0 and isinstance(input_shape[-1], int):
+                channels = input_shape[-1]
+            else:
+                channels = 16
+        else:
+            channels = 16
+
+        self.channels = channels
+        self.dense_f = Dense(channels, activation="sigmoid", **layer_kwargs)
+        self.dense_s = Dense(channels, activation=self.activation, **layer_kwargs)
+
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_attr=None, **kwargs):
+        if not self.built:
+            x_shape = getattr(x, "shape", None)
+            self.build(x_shape)
+
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x_in, a, e = self.get_inputs(x)
+            return self.propagate(x_in, a, e, **kwargs)
+
+        if edge_attr is None:
+            edge_attr = kwargs.get("e", None)
+
+        return self.propagate(edge_index, x=x, edge_attr=edge_attr)
+
+    def message(self, x_i=None, x_j=None, edge_attr=None, x=None, e=None, **kwargs):
+        if x_i is None and x is not None:
+            x_i = self.get_targets(x)
+            x_j = self.get_sources(x)
+        if edge_attr is None:
+            edge_attr = e
+
+        to_concat = [x_i, x_j]
+        if e is not None:
+            to_concat += [e]
+        if edge_attr is not None:
+            to_concat.append(edge_attr)
+        z = ops.concatenate(to_concat, axis=-1)
+        output = self.dense_s(z) * self.dense_f(z)
+
+        return output
+
+    def update(self, embeddings, x=None, **kwargs):
+        if x is None:
+            return embeddings
+        if isinstance(x, (tuple, list)):
+            x = x[1] if x[1] is not None else x[0]
+        return x + embeddings

k3_node/layers/conv/cugraph.py

@@ -0,0 +1,84 @@
+from k3_node.layers.conv.gat_conv import GATConv
+from k3_node.layers.conv.sage_conv import SAGEConv
+from k3_node.layers.conv.rgcn_conv import CuGraphRGCNConv
+
+
+class CuGraphGATConv(GATConv):
+    r"""An optimized / multi-backend compatible version of :class:`GATConv`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import CuGraphGATConv
+
+        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 = CuGraphGATConv(8, 16, heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 32)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            heads=heads,
+            concat=concat,
+            negative_slope=negative_slope,
+            bias=bias,
+            **kwargs,
+        )
+
+
+class CuGraphSAGEConv(SAGEConv):
+    r"""An optimized / multi-backend compatible version of :class:`SAGEConv`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import CuGraphSAGEConv
+
+        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 = CuGraphSAGEConv(8, 16)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        aggr: str = "mean",
+        normalize: bool = False,
+        root_weight: bool = True,
+        project: bool = False,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            aggr=aggr,
+            normalize=normalize,
+            root_weight=root_weight,
+            project=project,
+            bias=bias,
+            **kwargs,
+        )
+
+
+__all__ = ["CuGraphGATConv", "CuGraphSAGEConv", "CuGraphRGCNConv"]

k3_node/layers/conv/diffusion_conv.py

@@ -0,0 +1,144 @@
+from keras import layers, ops
+
+from k3_node.layers.conv.conv import Conv
+from k3_node.ops import normalized_adjacency, polyval
+
+
+class DiffuseFeatures(layers.Layer):
+    def __init__(
+        self,
+        num_diffusion_steps,
+        kernel_initializer,
+        kernel_regularizer,
+        kernel_constraint,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.K = num_diffusion_steps
+        self.kernel_initializer = kernel_initializer
+        self.kernel_regularizer = kernel_regularizer
+        self.kernel_constraint = kernel_constraint
+
+    def build(self, input_shape):
+        # Initializing the kernel vector (R^K) (theta in paper)
+        self.kernel = self.add_weight(
+            shape=(self.K,),
+            name="kernel",
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+        )
+
+    def call(self, inputs):
+        x, a = inputs
+
+        diffusion_matrix = polyval(self.kernel, a)
+        diffused_features = ops.matmul(diffusion_matrix, x)
+        H = ops.sum(diffused_features, axis=-1)
+        return ops.expand_dims(H, -1)
+
+
+class DiffusionConv(Conv):
+    """
+    `k3_node.layers.DiffusionConv`
+    Implementation of Diffusion Convolutional Neural Networks (DCNN) layer
+
+    Args:
+        channels: The number of output channels.
+        K: The number of diffusion steps.
+        activation: Activation function to use.
+        kernel_initializer: Initializer for the `kernel` weights matrix.
+        kernel_regularizer: Regularizer for the `kernel` weights matrix.
+        kernel_constraint: Constraint for the `kernel` weights matrix.
+        **kwargs: Additional arguments to pass to the `Conv` superclass.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DiffusionConv
+
+        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 = DiffusionConv(channels=8, out_channels=16, K=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        channels,
+        out_channels=None,
+        K=6,
+        activation="tanh",
+        kernel_initializer="glorot_uniform",
+        kernel_regularizer=None,
+        kernel_constraint=None,
+        **kwargs,
+    ):
+        if out_channels is not None:
+            self.in_channels = channels
+            channels = out_channels
+        super().__init__(
+            activation=activation,
+            kernel_initializer=kernel_initializer,
+            kernel_regularizer=kernel_regularizer,
+            kernel_constraint=kernel_constraint,
+            **kwargs,
+        )
+
+        self.channels = channels
+        self.K = K + 1
+
+    def build(self, input_shape):
+        self.filters = [
+            DiffuseFeatures(
+                num_diffusion_steps=self.K,
+                kernel_initializer=self.kernel_initializer,
+                kernel_regularizer=self.kernel_regularizer,
+                kernel_constraint=self.kernel_constraint,
+            )
+            for _ in range(self.channels)
+        ]
+        for f in self.filters:
+            f.build(None)
+        super().build(input_shape)
+
+    def apply_filters(self, x, a):
+        diffused_features = []
+
+        for diffusion in self.filters:
+            diffused_feature = diffusion((x, a))
+            diffused_features.append(diffused_feature)
+
+        return ops.concatenate(diffused_features, -1)
+
+    def call(self, inputs, a=None, **kwargs):
+        if a is not None:
+            x = inputs
+        elif isinstance(inputs, (list, tuple)):
+            x, a = inputs
+        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)[0]
+            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)
+
+        output = self.apply_filters(x, a)
+
+        output = self.activation(output)
+
+        return output
+
+    @property
+    def config(self):
+        return {"channels": self.channels, "K": self.K - 1}
+
+    @staticmethod
+    def preprocess(a):
+        return normalized_adjacency(a)

k3_node/layers/conv/dir_gnn_conv.py

@@ -0,0 +1,93 @@
+import copy
+from keras import ops
+from keras.layers import Layer, Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class DirGNNConv(Layer):
+    r"""A directed graph neural network operator from the
+    `"Directed Graph Neural Networks" <https://arxiv.org/abs/2301.07663>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DirGNNConv, 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
+
+        # Runs the wrapped convolution on incoming and outgoing edges separately
+        layer = DirGNNConv(GCNConv(in_channels=8, out_channels=16))
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        conv: MessagePassing,
+        alpha: float = 0.5,
+        root_weight: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.alpha = alpha
+        self.root_weight = root_weight
+        self.conv = conv
+
+        try:
+            self.conv_in = conv.__class__.from_config(conv.get_config())
+            self.conv_out = conv.__class__.from_config(conv.get_config())
+        except Exception:
+            self.conv_in = copy.deepcopy(conv)
+            self.conv_out = copy.deepcopy(conv)
+
+        if hasattr(self.conv_in, "add_self_loops"):
+            self.conv_in.add_self_loops = False
+            self.conv_out.add_self_loops = False
+        if hasattr(self.conv_in, "root_weight"):
+            self.conv_in.root_weight = False
+            self.conv_out.root_weight = False
+
+        in_channels = getattr(conv, "in_channels", None)
+        out_channels = getattr(conv, "out_channels", None)
+        if isinstance(in_channels, (list, tuple)):
+            in_channels = in_channels[0]
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if root_weight and out_channels is not None:
+            self.lin = Dense(out_channels, use_bias=True)
+        else:
+            self.lin = None
+
+    def build(self, input_shape=None):
+        if self.in_channels is not None:
+            if hasattr(self.conv_in, "build"):
+                self.conv_in.build((None, self.in_channels))
+            if hasattr(self.conv_out, "build"):
+                self.conv_out.build((None, self.in_channels))
+            if self.lin is not None:
+                self.lin.build((None, self.in_channels))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)) and len(inputs) == 2:
+                x, edge_index = inputs
+            else:
+                raise ValueError("Expected (x, edge_index) or x and edge_index")
+        else:
+            x = inputs
+
+        x_in = self.conv_in(x, edge_index)
+        edge_index_rev = ops.stack([edge_index[1], edge_index[0]], axis=0)
+        x_out = self.conv_out(x, edge_index_rev)
+
+        out = self.alpha * x_out + (1.0 - self.alpha) * x_in
+
+        if self.lin is not None:
+            out = out + self.lin(x)
+
+        return out

k3_node/layers/conv/dna_conv.py

@@ -0,0 +1,192 @@
+import math
+import keras
+from keras import ops
+from k3_node.layers.conv.gcn_conv import gcn_norm
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+def restricted_softmax(src, axis: int = -1, margin: float = 0.0):
+    src_max = ops.maximum(ops.max(src, axis=axis, keepdims=True), 0.0)
+    out = ops.exp(src - src_max)
+    denom = ops.sum(out, axis=axis, keepdims=True) + ops.exp(margin - src_max)
+    return out / denom
+
+
+class GroupedLinear(keras.layers.Layer):
+    def __init__(self, in_channels: int, out_channels: int, groups: int = 1, bias: bool = True, **kwargs):
+        super().__init__(**kwargs)
+        assert in_channels % groups == 0 and out_channels % groups == 0
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.groups = groups
+        self.use_bias = bias
+
+    def build(self, input_shape=None):
+        self.weight = self.add_weight(
+            shape=(self.groups, self.in_channels // self.groups, self.out_channels // self.groups),
+            initializer="glorot_uniform",
+            trainable=True,
+            name="weight",
+        )
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.out_channels,),
+                initializer="zeros",
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.bias = None
+        super().build(input_shape)
+
+    def call(self, src):
+        if self.groups > 1:
+            orig_shape = ops.shape(src)
+            # Flatten batch dims
+            src_flat = ops.reshape(src, (-1, self.groups, self.in_channels // self.groups))
+            src_trans = ops.transpose(src_flat, (1, 0, 2))
+            out = ops.matmul(src_trans, self.weight)
+            out = ops.transpose(out, (1, 0, 2))
+            out_shape = tuple(orig_shape[:-1]) + (self.out_channels,)
+            out = ops.reshape(out, out_shape)
+        else:
+            out = ops.matmul(src, self.weight[0])
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+class DNAMultiHead(keras.layers.Layer):
+    def __init__(self, in_channels: int, out_channels: int, heads: int = 1, groups: int = 1, bias: bool = True, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.groups = groups
+        self.use_bias = bias
+
+        self.lin_q = GroupedLinear(in_channels, out_channels, groups, bias)
+        self.lin_k = GroupedLinear(in_channels, out_channels, groups, bias)
+        self.lin_v = GroupedLinear(in_channels, out_channels, groups, bias)
+
+    def build(self, input_shape=None):
+        self.lin_q.build()
+        self.lin_k.build()
+        self.lin_v.build()
+        super().build(input_shape)
+
+    def call(self, query, key, value):
+        q = self.lin_q(query)
+        k = self.lin_k(key)
+        v = self.lin_v(value)
+
+        # q: (E, 1, C) -> (E, heads, 1, C // heads)
+        E = ops.shape(q)[0]
+        H = self.heads
+        D = self.out_channels // H
+
+        q_entries = ops.shape(q)[1]
+        k_entries = ops.shape(k)[1]
+
+        q = ops.transpose(ops.reshape(q, (E, q_entries, H, D)), (0, 2, 1, 3))
+        k = ops.transpose(ops.reshape(k, (E, k_entries, H, D)), (0, 2, 1, 3))
+        v = ops.transpose(ops.reshape(v, (E, k_entries, H, D)), (0, 2, 1, 3))
+
+        # score: (E, H, q_entries, k_entries)
+        score = ops.matmul(q, ops.transpose(k, (0, 1, 3, 2))) / math.sqrt(D)
+        score = restricted_softmax(score, axis=-1)
+
+        out = ops.matmul(score, v)  # (E, H, q_entries, D)
+        out = ops.transpose(out, (0, 2, 1, 3))  # (E, q_entries, H, D)
+        return ops.reshape(out, (E, q_entries, self.out_channels))
+
+
+class DNAConv(MessagePassing):
+    r"""The dynamic neighborhood aggregation operator from the `"Just Jump:
+    Towards Dynamic Neighborhood Aggregation in Graph Neural Networks"
+    <https://arxiv.org/abs/1904.04849>`_ paper.
+
+    Args:
+        channels (int): Size of each input/output sample.
+        heads (int, optional): Number of multi-head-attentions. (default: :obj:`1`)
+        groups (int, optional): Number of groups for linear projections. (default: :obj:`1`)
+        dropout (float, optional): Dropout probability. (default: :obj:`0.0`)
+        cached (bool, optional): Whether to cache GCN normalization. (default: :obj:`False`)
+        normalize (bool, optional): Whether to apply symmetric normalization. (default: :obj:`True`)
+        add_self_loops (bool, optional): Whether to add self-loops. (default: :obj:`True`)
+        bias (bool, optional): Whether to learn an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DNAConv
+
+        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
+
+        # DNAConv attends over the representations of all previous layers: [num_nodes, num_layers, channels]
+        x_all = np.random.rand(10, 3, 8).astype("float32")
+        layer = DNAConv(channels=8, heads=2, groups=2)
+        out = layer(x_all, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        heads: int = 1,
+        groups: int = 1,
+        dropout: float = 0.0,
+        cached: bool = False,
+        normalize: bool = True,
+        add_self_loops: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(node_dim=0, **kwargs)
+
+        self.channels = channels
+        self.heads = heads
+        self.groups = groups
+        self.dropout = dropout
+        self.cached = cached
+        self.normalize = normalize
+        self.add_self_loops = add_self_loops
+        self.use_bias = bias
+
+        self.multi_head = DNAMultiHead(channels, channels, heads, groups, bias)
+
+    def build(self, input_shape=None):
+        self.multi_head.build()
+        super().build(input_shape)
+
+    def call(self, x, edge_index, edge_weight=None):
+        if not self.built:
+            self.build()
+
+        if len(ops.shape(x)) == 2:
+            x = ops.expand_dims(x, axis=1)
+
+        num_nodes = ops.shape(x)[0]
+        if self.normalize:
+            edge_index, edge_weight = gcn_norm(
+                edge_index,
+                edge_weight,
+                num_nodes=num_nodes,
+                add_self_loops=self.add_self_loops,
+                dtype=x.dtype,
+            )
+
+        return self.propagate(edge_index, x=x, edge_weight=edge_weight)
+
+    def message(self, x_i, x_j, edge_weight=None):
+        q = x_i[:, -1:]  # (E, 1, C)
+        out = self.multi_head(q, x_j, x_j)  # (E, 1, C)
+        out = ops.squeeze(out, axis=1)  # (E, C)
+        if edge_weight is not None:
+            out = ops.expand_dims(edge_weight, axis=-1) * out
+        return out
+