k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/__init__.py

@@ -0,0 +1,149 @@
+from .message_passing import MessagePassing
+from .simple_conv import SimpleConv
+from .gcn_conv import GCNConv
+from .cheb_conv import ChebConv
+from .sage_conv import SAGEConv
+from .graph_conv import GraphConv
+from .gated_graph_conv import GatedGraphConv
+from .res_gated_graph_conv import ResGatedGraphConv
+from .gat_conv import GATConv, FusedGATConv
+from .gatv2_conv import GATv2Conv
+from .transformer_conv import TransformerConv
+from .agnn_conv import AGNNConv
+from .tag_conv import TAGConv
+from .gin_conv import GINConv, GINEConv
+from .arma_conv import ARMAConv
+from .sg_conv import SGConv
+from .ssg_conv import SSGConv
+from .appnp import APPNP
+from .appnp_conv import APPNPConv
+from .mf_conv import MFConv
+from .rgcn_conv import RGCNConv, FastRGCNConv, CuGraphRGCNConv
+from .rgat_conv import RGATConv
+from .signed_conv import SignedConv
+from .dir_gnn_conv import DirGNNConv
+from .antisymmetric_conv import AntiSymmetricConv
+from .mixhop_conv import MixHopConv
+from .pdn_conv import PDNConv
+from .fa_conv import FAConv
+from .film_conv import FiLMConv
+from .supergat_conv import SuperGATConv
+from .eg_conv import EGConv
+from .pan_conv import PANConv
+from .gen_conv import GENConv
+from .pna_conv import PNAConv
+from .le_conv import LEConv
+from .cluster_gcn_conv import ClusterGCNConv
+from .gcn2_conv import GCN2Conv
+from .lg_conv import LGConv
+from .nn_conv import NNConv
+from .cg_conv import CGConv
+from .edge_conv import EdgeConv, DynamicEdgeConv
+from .general_conv import GeneralConv
+from .point_conv import PointNetConv, PointConv
+from .point_transformer_conv import PointTransformerConv
+from .point_gnn_conv import PointGNNConv
+from .ppf_conv import PPFConv
+from .feast_conv import FeaStConv
+from .gmm_conv import GMMConv
+from .gravnet_conv import GravNetConv
+from .meshcnn_conv import MeshCNNConv
+from .x_conv import XConv
+from .spline_conv import SplineConv
+from .hetero_conv import HeteroConv
+from .hgt_conv import HGTConv
+from .han_conv import HANConv
+from .heat_conv import HEATConv
+from .hypergraph_conv import HypergraphConv
+from .dna_conv import DNAConv
+from .wl_conv import WLConv, WLConvContinuous
+from .gps_conv import GPSConv
+from .cugraph import CuGraphGATConv, CuGraphSAGEConv
+
+ECConv = NNConv
+
+# Legacy Spektral imports
+from .crystal_conv import CrystalConv
+from .diffusion_conv import DiffusionConv
+from .gcn import GraphConvolution
+from .graph_attention import GraphAttention
+from .ppnp import PPNPPropagation
+
+__all__ = [
+    "MessagePassing",
+    "SimpleConv",
+    "GCNConv",
+    "ChebConv",
+    "SAGEConv",
+    "GraphConv",
+    "GatedGraphConv",
+    "ResGatedGraphConv",
+    "GATConv",
+    "FusedGATConv",
+    "GATv2Conv",
+    "TransformerConv",
+    "AGNNConv",
+    "TAGConv",
+    "GINConv",
+    "GINEConv",
+    "ARMAConv",
+    "SGConv",
+    "SSGConv",
+    "APPNP",
+    "APPNPConv",
+    "MFConv",
+    "RGCNConv",
+    "FastRGCNConv",
+    "CuGraphRGCNConv",
+    "RGATConv",
+    "SignedConv",
+    "DirGNNConv",
+    "AntiSymmetricConv",
+    "MixHopConv",
+    "PDNConv",
+    "FAConv",
+    "FiLMConv",
+    "SuperGATConv",
+    "EGConv",
+    "PANConv",
+    "GENConv",
+    "PNAConv",
+    "LEConv",
+    "ClusterGCNConv",
+    "GCN2Conv",
+    "LGConv",
+    "NNConv",
+    "ECConv",
+    "CGConv",
+    "EdgeConv",
+    "DynamicEdgeConv",
+    "GeneralConv",
+    "PointNetConv",
+    "PointConv",
+    "PointTransformerConv",
+    "PointGNNConv",
+    "PPFConv",
+    "FeaStConv",
+    "GMMConv",
+    "GravNetConv",
+    "MeshCNNConv",
+    "XConv",
+    "SplineConv",
+    "HeteroConv",
+    "HGTConv",
+    "HANConv",
+    "HEATConv",
+    "HypergraphConv",
+    "DNAConv",
+    "WLConv",
+    "WLConvContinuous",
+    "GPSConv",
+    "CuGraphGATConv",
+    "CuGraphSAGEConv",
+    # Legacy Spektral
+    "CrystalConv",
+    "DiffusionConv",
+    "GraphConvolution",
+    "GraphAttention",
+    "PPNPPropagation",
+]

k3_node/layers/conv/agnn_conv.py

@@ -0,0 +1,120 @@
+from typing import Optional
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import (
+    add_self_loops,
+    extend_mask_for_self_loops,
+    mask_edge_logits,
+    remove_self_loops_masked,
+    softmax,
+)
+
+
+class AGNNConv(MessagePassing):
+    r"""The graph attentional propagation layer from the
+    `"Attention-based Graph Neural Network for Semi-Supervised Learning"
+    <https://arxiv.org/abs/1803.03735>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import AGNNConv
+
+        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 = AGNNConv(requires_grad=True)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        requires_grad: bool = True,
+        add_self_loops: bool = True,
+        trainable: Optional[bool] = None,
+        aggregate: str = "add",
+        activation=None,
+        **kwargs,
+    ):
+        if trainable is not None:
+            requires_grad = trainable
+        kwargs.setdefault("aggr", aggregate)
+        super().__init__(activation=activation, **kwargs)
+
+        self.requires_grad = requires_grad
+        self.add_self_loops = add_self_loops
+
+        if requires_grad:
+            self.beta = self.add_weight(
+                shape=(1,),
+                initializer="ones",
+                name="beta",
+            )
+        else:
+            self.beta = None
+
+    def build(self, input_shape=None):
+        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
+
+        if not self.built:
+            self.build()
+
+        # Check for legacy 2D matrix
+        is_legacy = False
+        if hasattr(edge_index, "shape") and len(edge_index.shape) == 2:
+            if (
+                edge_index.shape[0] is not None
+                and edge_index.shape[1] is not None
+                and 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
+
+        x_norm = x / (ops.norm(x, axis=-1, keepdims=True) + 1e-12)
+
+        if is_legacy:
+            out = self.propagate(x, edge_index, x_norm=x_norm)
+        else:
+            num_nodes = x.shape[0] if hasattr(x, "shape") and x.shape[0] is not None else ops.shape(x)[0]
+            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)
+            out = self.propagate(edge_index, x=x, x_norm=x_norm, keep_mask=keep_mask, size=(num_nodes, num_nodes))
+
+        if self.activation is not None:
+            out = self.activation(out)
+
+        return out
+
+    def message(self, x=None, x_j=None, x_norm=None, x_norm_i=None, x_norm_j=None, index=None, size_i=None, keep_mask=None):
+        beta = self.beta if self.beta is not None else 1.0
+
+        # Legacy Spektral path
+        if x_j is None and x is not None:
+            x_j = self.get_sources(x)
+            x_norm_i = self.get_targets(x_norm)
+            x_norm_j = self.get_sources(x_norm)
+            alpha = beta * ops.sum(x_norm_i * x_norm_j, axis=-1)
+            alpha = softmax(alpha, self.index_targets, num_nodes=self.n_nodes, dim=0)
+            return ops.expand_dims(alpha, -1) * x_j
+
+        # PyG path
+        alpha = beta * ops.sum(x_norm_i * x_norm_j, axis=-1)
+        alpha = mask_edge_logits(alpha, keep_mask)
+        alpha = softmax(alpha, index, num_nodes=size_i, dim=0)
+        return x_j * ops.expand_dims(alpha, -1)

k3_node/layers/conv/antisymmetric_conv.py

@@ -0,0 +1,94 @@
+from typing import Optional, Union, Callable
+from keras import ops, activations
+from keras.layers import Layer
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.gcn_conv import GCNConv
+
+
+class AntiSymmetricConv(Layer):
+    r"""The anti-symmetric graph convolutional operator from the
+    `"Anti-Symmetric DGN: a Continuous approach to Deep Graph Neural Networks"
+    <https://arxiv.org/abs/2202.13085>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import AntiSymmetricConv
+
+        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 = AntiSymmetricConv(in_channels=8)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        phi: Optional[MessagePassing] = None,
+        num_iters: int = 1,
+        epsilon: float = 0.1,
+        gamma: float = 0.1,
+        act: Union[str, Callable, None] = "tanh",
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.num_iters = num_iters
+        self.gamma = gamma
+        self.epsilon = epsilon
+        self.act = activations.get(act) if act is not None else None
+
+        if phi is None:
+            phi = GCNConv(in_channels, in_channels, bias=False)
+        self.phi = phi
+
+        self.W = self.add_weight(
+            shape=(in_channels, in_channels),
+            initializer="glorot_uniform",
+            name="W",
+        )
+
+        if bias:
+            self.bias = self.add_weight(
+                shape=(in_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape=None):
+        if hasattr(self.phi, "build"):
+            self.phi.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
+
+        eye = ops.eye(self.in_channels, dtype=self.W.dtype)
+        antisymmetric_W = self.W - ops.transpose(self.W) - self.gamma * eye
+
+        for _ in range(self.num_iters):
+            h = self.phi(x, edge_index)
+            h = ops.matmul(x, ops.transpose(antisymmetric_W)) + h
+
+            if self.bias is not None:
+                h = h + self.bias
+
+            if self.act is not None:
+                h = self.act(h)
+
+            x = x + self.epsilon * h
+
+        return x

k3_node/layers/conv/appnp.py

@@ -0,0 +1,105 @@
+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 APPNP(MessagePassing):
+    r"""The approximate personalized propagation of neural predictions (APPNP)
+    operator from the `"Predict then Propagate: Combining Neural Networks with
+    Personalized PageRank for Classification on Graphs"
+    <https://arxiv.org/abs/1810.05997>`_ paper.
+
+    Args:
+        K: Number of iterations :math:`K`.
+        alpha: Teleport probability :math:`\alpha`.
+        dropout: Dropout probability of edges or features during propagation.
+            (default: ``0.0``)
+        cached: If set to :obj:`True`, the layer 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 APPNP
+
+        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 = APPNP(K=2, alpha=0.1)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(
+        self,
+        K: int,
+        alpha: float,
+        dropout: float = 0.0,
+        cached: bool = False,
+        add_self_loops: bool = True,
+        normalize: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.K = K
+        self.alpha = alpha
+        self.dropout_rate = dropout
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.normalize = normalize
+        self._cached_edge_index = None
+        self._cached_norm = None
+        self.dropout = layers.Dropout(dropout) if dropout > 0.0 else None
+
+    def build(self, input_shape=None):
+        if self.dropout is not None and hasattr(self.dropout, "build"):
+            self.dropout.build(input_shape)
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, training=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        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 = x
+        for _ in range(self.K):
+            if self.dropout is not None:
+                h = self.dropout(h, training=training)
+            h = self.propagate(edge_index, x=h, edge_weight=edge_weight)
+            h = (1.0 - self.alpha) * h + self.alpha * x
+
+        return h
+
+    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
+
+
+# Backward-compatible alias
+APPNPConv = APPNP
+

k3_node/layers/conv/appnp_conv.py

@@ -0,0 +1,157 @@
+# ported from spektral
+
+from keras import activations
+from keras import activations, ops
+from keras.layers import Dense, Dropout
+from keras.models import Sequential
+
+from k3_node.layers.conv.conv import Conv
+from k3_node.ops import gcn_filter, modal_dot
+
+
+class APPNPConv(Conv):
+    """
+        `k3_node.layers.APPNPConv`
+        Implementation of Approximate Personalized Propagation of Neural Predictions
+
+        Args:
+            channels: The number of output channels.
+            alpha: The teleport probability.
+            propagations: The number of propagation steps.
+            mlp_hidden: A list of hidden channels for the MLP.
+            mlp_activation: The activation function to use in the MLP.
+            dropout_rate: The dropout rate for the MLP.
+            activation: The activation function to use in the layer.
+            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 keyword arguments.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import APPNPConv
+
+        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 = APPNPConv(channels=16, alpha=0.1, propagations=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        channels,
+        alpha=0.2,
+        propagations=1,
+        mlp_hidden=None,
+        mlp_activation="relu",
+        dropout_rate=0.0,
+        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__(
+            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.mlp_hidden = mlp_hidden if mlp_hidden else []
+        self.alpha = alpha
+        self.propagations = propagations
+        self.mlp_activation = activations.get(mlp_activation)
+        self.dropout_rate = dropout_rate
+
+    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,
+        )
+        mlp_layers = []
+        for channels in self.mlp_hidden:
+            mlp_layers.extend(
+                [
+                    Dropout(self.dropout_rate),
+                    Dense(channels, self.mlp_activation, **layer_kwargs),
+                ]
+            )
+        mlp_layers.append(Dense(self.channels, "linear", **layer_kwargs))
+        self.mlp = Sequential(mlp_layers)
+        if input_shape is not None:
+            feat_shape = input_shape[0] if isinstance(input_shape, (list, tuple)) else input_shape
+            self.mlp.build(feat_shape)
+        self.built = True
+
+    def call(self, inputs, mask=None):
+        x, a = inputs
+    def call(self, inputs, a=None, mask=None):
+        if a is not None:
+            x = inputs
+        elif isinstance(inputs, (list, tuple)) and len(inputs) == 2:
+            x, a = inputs
+        else:
+            x = inputs
+            a = None
+
+        if not self.built:
+            self.build(getattr(x, "shape", 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] if len(ops.shape(x)) >= 2 else ops.shape(x)[0]
+            row, col = ops.cast(a[0], "int32"), ops.cast(a[1], "int32")
+            idx = ops.stack([row, col], axis=-1)
+            zeros = ops.zeros((num_nodes, num_nodes), dtype=x.dtype)
+            a = ops.scatter_update(zeros, idx, ops.ones((ops.shape(a)[1],), dtype=x.dtype))
+
+        mlp_out = self.mlp(x)
+        output = mlp_out
+        if a is not None:
+            for _ in range(self.propagations):
+                output = (1 - self.alpha) * modal_dot(a, output) + self.alpha * mlp_out
+        if mask is not None and isinstance(mask, (list, tuple)) and len(mask) > 0 and mask[0] is not None:
+            output *= mask[0]
+        output = self.activation(output)
+
+        return output
+
+    @property
+    def config(self):
+        return {
+            "channels": self.channels,
+            "alpha": self.alpha,
+            "propagations": self.propagations,
+            "mlp_hidden": self.mlp_hidden,
+            "mlp_activation": activations.serialize(self.mlp_activation),
+            "dropout_rate": self.dropout_rate,
+        }
+
+    @staticmethod
+    def preprocess(a):
+        return gcn_filter(a)