k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/sage_conv.py

@@ -0,0 +1,154 @@
+from typing import Optional, Union, Tuple, List
+
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.aggr.base import Aggregation
+
+
+class SAGEConv(MessagePassing):
+    r"""The GraphSAGE operator from the `"Inductive Representation Learning on
+    Large Graphs" <https://arxiv.org/abs/1706.02216>`_ 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 (``"mean"``, ``"max"``,
+            ``"lstm"``, etc.). (default: ``"mean"``)
+        normalize: If set to :obj:`True`, output features will be
+            :math:`\ell_2`-normalized. (default: :obj:`False`)
+        root_weight: If set to :obj:`False`, the layer will not add
+            the transformed root node features. (default: :obj:`True`)
+        project: If set to :obj:`True`, the layer will apply a linear
+            transformation followed by an activation to source node features.
+            (default: :obj:`False`)
+        bias: If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SAGEConv
+
+        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 = SAGEConv(in_channels=8, out_channels=16)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int], None],
+        out_channels: Optional[int] = None,
+        aggr: Optional[Union[str, List[str], Aggregation]] = "mean",
+        normalize: bool = False,
+        root_weight: bool = True,
+        project: bool = False,
+        bias: bool = True,
+        **kwargs,
+    ):
+        if out_channels is None:
+            # Backward compatibility: SAGEConv(out_channels, ...)
+            out_channels = in_channels
+            in_channels = None
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.normalize = normalize
+        self.root_weight = root_weight
+        self.project = project
+        self.use_bias = bias
+
+        super().__init__(aggr=aggr, **kwargs)
+
+        if self.project:
+            self.lin_proj = layers.Dense(
+                in_channels[0] if isinstance(in_channels, (tuple, list)) else in_channels,
+                activation="relu",
+                use_bias=True,
+            )
+        else:
+            self.lin_proj = None
+
+        self.lin_l = layers.Dense(out_channels, use_bias=bias)
+        if self.root_weight:
+            self.lin_r = layers.Dense(out_channels, use_bias=False)
+        else:
+            self.lin_r = None
+
+    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_l.build((None, self.out_channels if self.project else in_channels_l))
+        if self.lin_r is not None:
+            self.lin_r.build((None, in_channels_r))
+        if self.lin_proj is not None:
+            self.lin_proj.build((None, in_channels_l))
+        self.built = True
+
+    def call(self, x, edge_index=None, size=None, **kwargs):
+        # Handle legacy calling: conv(x, adj) where adj is [N, N]
+        shape = getattr(edge_index, "shape", None)
+        if (
+            shape is not None
+            and len(shape) == 2
+            and shape[0] is not None
+            and shape[1] is not None
+            and shape[0] > 2
+            and shape[0] == shape[1]
+        ):
+            where_adj = ops.where(edge_index != 0)
+            where_adj = where_adj if not isinstance(where_adj, list) else where_adj
+            edge_index = ops.stack([where_adj[0], where_adj[1]], axis=0)
+
+        # Handle legacy calling: conv((x, adj))
+        if edge_index is None and isinstance(x, (tuple, list)) and len(x) == 2:
+            arg0, arg1 = x[0], x[1]
+            s1 = getattr(arg1, "shape", None)
+            if (
+                s1 is not None
+                and len(s1) == 2
+                and s1[0] is not None
+                and s1[1] is not None
+                and s1[0] > 2
+                and s1[0] == s1[1]
+            ):
+                where_adj = ops.where(arg1 != 0)
+                where_adj = where_adj if not isinstance(where_adj, list) else where_adj
+                edge_index = ops.stack([where_adj[0], where_adj[1]], axis=0)
+                x = arg0
+            elif s1 is not None and len(s1) >= 1 and s1[0] == 2:
+                edge_index = arg1
+                x = arg0
+
+        if not isinstance(x, (tuple, list)):
+            x_src = x
+            x_dst = x
+        else:
+            x_src, x_dst = x[0], x[1]
+
+        if self.project and self.lin_proj is not None:
+            x_src = self.lin_proj(x_src)
+
+        out = self.propagate(edge_index, x=(x_src, x_dst), size=size)
+        out = self.lin_l(out)
+
+        if self.root_weight and self.lin_r is not None and x_dst is not None:
+            out = out + self.lin_r(x_dst)
+
+        if self.normalize:
+            norm = ops.sqrt(ops.sum(ops.square(out), axis=-1, keepdims=True) + 1e-12)
+            out = out / norm
+
+        return out
+
+    def message(self, x_j):
+        return x_j

k3_node/layers/conv/sg_conv.py

@@ -0,0 +1,96 @@
+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 SGConv(MessagePassing):
+    r"""The simple graph convolutional operator from the `"Simplifying Graph
+    Convolutional Networks" <https://arxiv.org/abs/1902.07153>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample.
+        out_channels: Size of each output sample.
+        K: Number of hops :math:`K`. (default: ``1``)
+        cached: If set to :obj:`True`, the layer will cache the computation of
+            :math:`\mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}} \mathbf{\hat{D}}^{-1/2}`.
+            (default: ``False``)
+        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 SGConv
+
+        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 = SGConv(in_channels=8, out_channels=16, K=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        K: int = 1,
+        cached: bool = False,
+        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.K = K
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.use_bias = bias
+
+        self.lin = layers.Dense(out_channels, use_bias=bias)
+        self._cached_edge_index = None
+        self._cached_norm = None
+
+    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.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]
+
+        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
+
+        for _ in range(self.K):
+            x = self.propagate(edge_index, x=x, edge_weight=edge_weight)
+
+        return self.lin(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/signed_conv.py

@@ -0,0 +1,100 @@
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class SignedConv(MessagePassing):
+    r"""The signed graph convolutional operator from the `"Signed Graph
+    Convolutional Network" <https://arxiv.org/abs/1808.06354>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SignedConv
+
+        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
+
+        pos_edge_index = edge_index[:, :15]  # positive (e.g. "trust") edges
+        neg_edge_index = edge_index[:, 15:]  # negative (e.g. "distrust") edges
+        layer = SignedConv(in_channels=8, out_channels=16, first_aggr=True)
+        out = layer(x, pos_edge_index, neg_edge_index)
+        print(tuple(out.shape))  # (10, 32): positive and negative embeddings, concatenated
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        first_aggr: bool,
+        bias: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "mean")
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.first_aggr = first_aggr
+        self.use_bias = bias
+
+        in_pos = in_channels if first_aggr else 2 * in_channels
+        self.lin_pos_l = Dense(out_channels, use_bias=False)
+        self.lin_pos_r = Dense(out_channels, use_bias=bias)
+        self.lin_neg_l = Dense(out_channels, use_bias=False)
+        self.lin_neg_r = Dense(out_channels, use_bias=bias)
+
+    def build(self, input_shape=None):
+        in_dim = self.in_channels if self.first_aggr else 2 * self.in_channels
+        self.lin_pos_l.build((None, in_dim))
+        self.lin_pos_r.build((None, self.in_channels))
+        self.lin_neg_l.build((None, in_dim))
+        self.lin_neg_r.build((None, self.in_channels))
+        self.built = True
+
+    def call(self, inputs, pos_edge_index=None, neg_edge_index=None, **kwargs):
+        if pos_edge_index is None:
+            if isinstance(inputs, (list, tuple)) and len(inputs) == 3:
+                x, pos_edge_index, neg_edge_index = inputs
+            else:
+                raise ValueError("Expected (x, pos_edge_index, neg_edge_index)")
+        else:
+            x = inputs
+
+        if not self.built:
+            self.build()
+
+        if isinstance(x, (list, tuple)):
+            x_src, x_dst = x
+        else:
+            x_src = x_dst = x
+
+        if self.first_aggr:
+            out_pos = self.propagate(pos_edge_index, x=(x_src, x_dst))
+            out_pos = self.lin_pos_l(out_pos) + self.lin_pos_r(x_dst)
+
+            out_neg = self.propagate(neg_edge_index, x=(x_src, x_dst))
+            out_neg = self.lin_neg_l(out_neg) + self.lin_neg_r(x_dst)
+
+            return ops.concatenate([out_pos, out_neg], axis=-1)
+        else:
+            F_in = self.in_channels
+            x_src_1, x_src_2 = x_src[..., :F_in], x_src[..., F_in:]
+            x_dst_1, x_dst_2 = x_dst[..., :F_in], x_dst[..., F_in:]
+
+            out_pos1 = self.propagate(pos_edge_index, x=(x_src_1, x_dst_1))
+            out_pos2 = self.propagate(neg_edge_index, x=(x_src_2, x_dst_2))
+            out_pos = ops.concatenate([out_pos1, out_pos2], axis=-1)
+            out_pos = self.lin_pos_l(out_pos) + self.lin_pos_r(x_dst_1)
+
+            out_neg1 = self.propagate(pos_edge_index, x=(x_src_2, x_dst_2))
+            out_neg2 = self.propagate(neg_edge_index, x=(x_src_1, x_dst_1))
+            out_neg = ops.concatenate([out_neg1, out_neg2], axis=-1)
+            out_neg = self.lin_neg_l(out_neg) + self.lin_neg_r(x_dst_2)
+
+            return ops.concatenate([out_pos, out_neg], axis=-1)
+
+    def message(self, x_j):
+        return x_j
+

k3_node/layers/conv/simple_conv.py

@@ -0,0 +1,75 @@
+from typing import Optional, Union, List
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.aggr.base import Aggregation
+
+
+class SimpleConv(MessagePassing):
+    r"""A simple, parameter-free message passing operator.
+
+    Args:
+        aggr: The aggregation scheme to use (``"sum"``, ``"mean"``,
+            ``"min"``, ``"max"``, ``"mul"``). (default: ``"sum"``)
+        combine_root: The way to combine root node features with the
+            aggregated output (``"sum"``, ``"cat"``, ``"self_loop"``,
+            or :obj:`None`). (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SimpleConv
+
+        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 = SimpleConv(aggr="mean")
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        aggr: Union[str, List[str], Aggregation, None] = "sum",
+        combine_root: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+        self.combine_root = combine_root
+        if combine_root is not None and combine_root not in ["sum", "cat", "self_loop"]:
+            raise ValueError(
+                f"combine_root must be 'sum', 'cat', 'self_loop', or None, got {combine_root}"
+            )
+
+    def build(self, input_shape):
+        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_dst = x, x
+        else:
+            x_src, x_dst = x[0], x[1]
+
+        if self.combine_root == "self_loop":
+            from k3_node.layers.conv.utils import add_self_loops
+            num_nodes = ops.shape(x_src)[0]
+            edge_index, edge_weight = add_self_loops(edge_index, edge_weight, num_nodes=num_nodes)
+
+        out = self.propagate(edge_index, x=(x_src, x_dst), edge_weight=edge_weight, size=size)
+
+        if self.combine_root is not None and x_dst is not None and self.combine_root != "self_loop":
+            if self.combine_root == "sum":
+                out = out + x_dst
+            elif self.combine_root == "cat":
+                out = ops.concatenate([out, x_dst], axis=-1)
+
+        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/spline_conv.py

@@ -0,0 +1,182 @@
+from typing import List, Tuple, Union
+import keras
+from keras import ops
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class SplineConv(MessagePassing):
+    r"""The spline-based convolutional operator from the `"SplineCNN: Fast
+    Geometric Deep Learning with Continuous B-Spline Kernels"
+    <https://arxiv.org/abs/1711.08920>`_ paper.
+
+    Args:
+        in_channels (int or tuple): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        dim (int): Pseudo-coordinate dimensionality.
+        kernel_size (int or List[int]): Size of the convolving kernel.
+        is_open_spline (bool or List[bool], optional): If set to :obj:`False`,
+            uses closed B-spline basis. (default: :obj:`True`)
+        degree (int, optional): B-spline basis degree. (default: :obj:`1`)
+        aggr (str, optional): The aggregation scheme to use (:obj:`"mean"`,
+            :obj:`"add"`, :obj:`"max"`). (default: :obj:`"mean"`)
+        root_weight (bool, optional): Whether to add transformed root node
+            features. (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 SplineConv
+
+        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
+        pseudo = np.random.rand(30, 2).astype("float32")  # pseudo-coordinates in [0, 1]
+
+        layer = SplineConv(in_channels=8, out_channels=16, dim=2, kernel_size=3)
+        out = layer(x, edge_index, pseudo)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: int,
+        dim: int,
+        kernel_size: Union[int, List[int]],
+        is_open_spline: Union[bool, List[bool]] = True,
+        degree: int = 1,
+        aggr: str = "mean",
+        root_weight: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+
+        if isinstance(in_channels, int):
+            in_channels = (in_channels, in_channels)
+
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size] * dim
+        assert len(kernel_size) == dim
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.dim = dim
+        self.kernel_size = kernel_size
+        self.degree = degree
+        self.root_weight = root_weight
+        self.use_bias = bias
+
+        # Total number of spline basis elements
+        k_prod = 1
+        for k in kernel_size:
+            k_prod *= k
+        self.K = k_prod
+
+        # Compute strides for multidimensional indexing
+        strides = []
+        stride = 1
+        for k in reversed(kernel_size):
+            strides.insert(0, stride)
+            stride *= k
+        self.strides = strides
+
+    def build(self, input_shape=None):
+        in_dim = self.in_channels[0]
+        if in_dim == -1 and input_shape is not None:
+            if isinstance(input_shape, (list, tuple)) and isinstance(input_shape[0], (list, tuple)):
+                in_dim = input_shape[0][-1]
+            elif isinstance(input_shape, (list, tuple)):
+                in_dim = input_shape[-1]
+            self.in_channels = (in_dim, self.in_channels[1] if self.in_channels[1] != -1 else in_dim)
+
+        self.weight = self.add_weight(
+            shape=(self.K, self.in_channels[0], self.out_channels),
+            initializer="glorot_uniform",
+            trainable=True,
+            name="weight",
+        )
+
+        if self.root_weight:
+            self.root_lin = keras.layers.Dense(self.out_channels, use_bias=False)
+            self.root_lin.build((None, self.in_channels[1]))
+
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.out_channels,),
+                initializer="zeros",
+                trainable=True,
+                name="bias",
+            )
+
+        super().build(input_shape)
+
+    def _spline_basis_1d(self, e_d, k_d):
+        # e_d: (E,) in [0, 1]
+        u = ops.clip(e_d * (k_d - 1), 0.0, float(k_d - 1))
+        i0 = ops.cast(ops.floor(u), "int32")
+        i0 = ops.clip(i0, 0, k_d - 1)
+        i1 = ops.clip(i0 + 1, 0, k_d - 1)
+        w1 = u - ops.cast(i0, u.dtype)
+        w0 = 1.0 - w1
+        return [(i0, w0), (i1, w1)]
+
+    def _compute_kernel(self, edge_attr):
+        # edge_attr: (E, D)
+        E = ops.shape(edge_attr)[0]
+        # Basis product over dimensions
+        dim_bases = [
+            self._spline_basis_1d(edge_attr[:, d], self.kernel_size[d])
+            for d in range(self.dim)
+        ]
+
+        # Cartesian product of basis across dimensions
+        basis_combinations = [([], 1.0)]
+        for d in range(self.dim):
+            new_combinations = []
+            stride = self.strides[d]
+            for curr_indices, curr_weight in basis_combinations:
+                for idx, w in dim_bases[d]:
+                    new_idx = curr_indices + [idx * stride]
+                    new_w = curr_weight * w
+                    new_combinations.append((new_idx, new_w))
+            basis_combinations = new_combinations
+
+        # Construct edge-wise weight matrix W_eff: (E, C_in, C_out)
+        W_eff = ops.zeros((E, self.in_channels[0], self.out_channels), dtype=self.weight.dtype)
+        for idx_parts, basis_w in basis_combinations:
+            total_idx = idx_parts[0]
+            for part in idx_parts[1:]:
+                total_idx = total_idx + part
+            # total_idx: (E,)
+            basis_weights = ops.take(self.weight, total_idx, axis=0)  # (E, C_in, C_out)
+            basis_w_expanded = ops.expand_dims(ops.expand_dims(basis_w, axis=-1), axis=-1)
+            W_eff = W_eff + basis_w_expanded * basis_weights
+
+        return W_eff
+
+    def call(self, x, edge_index, edge_attr=None, **kwargs):
+        if isinstance(x, (list, tuple)):
+            x_src, x_dst = x[0], x[1]
+        else:
+            x_src = x_dst = x
+
+        out = self.propagate(edge_index, x=(x_src, x_dst), edge_attr=edge_attr)
+
+        if self.root_weight and x_dst is not None:
+            out = out + self.root_lin(x_dst)
+
+        if self.use_bias:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_j, edge_attr):
+        if edge_attr is None:
+            return x_j
+        W_eff = self._compute_kernel(edge_attr)  # (E, C_in, C_out)
+        # x_j: (E, C_in)
+        # m = sum_c x_j * W_eff
+        return ops.einsum("ec,ecd->ed", x_j, W_eff)
+