k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/edge_conv.py

@@ -0,0 +1,107 @@
+from typing import Callable, Optional, Union, Tuple
+import keras
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.pool.knn import knn_graph
+
+
+class EdgeConv(MessagePassing):
+    r"""The edge convolutional operator from the `"Dynamic Graph CNN for
+    Learning on Point Clouds" <https://arxiv.org/abs/1801.07829>`_ paper.
+
+    Args:
+        nn: A neural network :math:`h_{\mathbf{\Theta}}` that maps
+            pair-wise node features to new edge representations.
+        aggr: The aggregation scheme to use (``"max"``, ``"mean"``, ``"sum"``).
+            (default: ``"max"``)
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import EdgeConv
+
+        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
+
+        nn = keras.Sequential([keras.layers.Dense(16, activation="relu"), keras.layers.Dense(16)])
+        layer = EdgeConv(nn)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(self, nn: Callable, aggr: str = "max", **kwargs):
+        super().__init__(aggr=aggr, **kwargs)
+        self.nn = nn
+
+    def build(self, input_shape):
+        if hasattr(self.nn, "build") and not getattr(self.nn, "built", False):
+            if isinstance(input_shape, (tuple, list)) and len(input_shape) > 0 and isinstance(input_shape[0], (tuple, list)):
+                c = input_shape[0][-1]
+            elif isinstance(input_shape, (tuple, list)):
+                c = input_shape[-1]
+            else:
+                c = None
+            nn_shape = (None, 2 * c) if c is not None else None
+            self.nn.build(nn_shape)
+        self.built = True
+
+    def call(self, x, edge_index=None, training=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]
+
+        return self.propagate(edge_index, x=(x_src, x_dst), training=training, **kwargs)
+
+    def message(self, x_i, x_j, training=None):
+        h = ops.concatenate([x_i, x_j - x_i], axis=-1)
+        # Forward `training` explicitly: Keras does not propagate it to nested layers on JAX.
+        return self.nn(h, training=training) if isinstance(self.nn, keras.layers.Layer) else self.nn(h)
+
+
+class DynamicEdgeConv(EdgeConv):
+    r"""The dynamic edge convolutional operator from the `"Dynamic Graph CNN
+    for Learning on Point Clouds" <https://arxiv.org/abs/1801.07829>`_ paper,
+    which dynamically constructs a graph using :math:`k`-NN at each layer.
+
+    Args:
+        nn: A neural network :math:`h_{\mathbf{\Theta}}`.
+        k: Number of nearest neighbors. (default: ``6``)
+        aggr: The aggregation scheme to use (``"max"``, ``"mean"``, ``"sum"``).
+            (default: ``"max"``)
+        num_workers: Number of workers (ignored in Keras backend).
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import DynamicEdgeConv
+
+        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
+
+        nn = keras.Sequential([keras.layers.Dense(16, activation="relu"), keras.layers.Dense(16)])
+        layer = DynamicEdgeConv(nn, k=3)
+        out = layer(x)  # k-NN graph is built from x
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(self, nn: Callable, k: int = 6, aggr: str = "max", num_workers: int = 1, **kwargs):
+        super().__init__(nn=nn, aggr=aggr, **kwargs)
+        self.k = k
+
+    def call(self, x, batch=None, training=None, **kwargs):
+        if isinstance(x, (tuple, list)):
+            x_src = x[0]
+        else:
+            x_src = x
+
+        edge_index = knn_graph(x_src, k=self.k, batch=batch, loop=False, flow=self.flow)
+        return super().call(x, edge_index=edge_index, training=training, **kwargs)

k3_node/layers/conv/eg_conv.py

@@ -0,0 +1,155 @@
+from typing import Optional, List
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm, scatter
+
+
+class EGConv(MessagePassing):
+    r"""The Efficient Graph Convolution from the `"Adaptive Filters and
+    Aggregator Fusion for Efficient Graph Convolutions"
+    <https://arxiv.org/abs/2104.01481>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import EGConv
+
+        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 = EGConv(in_channels=8, out_channels=16, num_heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        aggregators: Optional[List[str]] = None,
+        num_heads: int = 8,
+        num_bases: int = 4,
+        cached: bool = False,
+        add_self_loops: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(node_dim=0, **kwargs)
+
+        if out_channels % num_heads != 0:
+            raise ValueError(
+                f"'out_channels' ({out_channels}) must be divisible by num_heads ({num_heads})"
+            )
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_heads = num_heads
+        self.num_bases = num_bases
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.aggregators = aggregators or ["symnorm"]
+        self.use_bias = bias
+
+        self.bases_lin = Dense(
+            (out_channels // num_heads) * num_bases, use_bias=False
+        )
+        self.comb_lin = Dense(
+            num_heads * num_bases * len(self.aggregators), use_bias=True
+        )
+
+        if bias:
+            self.bias = self.add_weight(
+                shape=(out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape=None):
+        self.bases_lin.build((None, self.in_channels))
+        self.comb_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
+
+        if not self.built:
+            self.build()
+
+        num_nodes = ops.shape(x)[0]
+        symnorm_weight = None
+        if "symnorm" in self.aggregators:
+            edge_index, symnorm_weight = gcn_norm(
+                edge_index,
+                edge_weight=None,
+                num_nodes=num_nodes,
+                add_self_loops=self.add_self_loops,
+                dtype=x.dtype,
+            )
+
+        bases = self.bases_lin(x)
+        weightings = self.comb_lin(x)
+
+        aggregated = self.propagate(
+            edge_index,
+            x=bases,
+            symnorm_weight=symnorm_weight,
+            size=(num_nodes, num_nodes),
+        )
+
+        weightings = ops.reshape(
+            weightings,
+            (-1, self.num_heads, self.num_bases * len(self.aggregators)),
+        )
+        aggregated = ops.reshape(
+            aggregated,
+            (
+                -1,
+                len(self.aggregators) * self.num_bases,
+                self.out_channels // self.num_heads,
+            ),
+        )
+
+        out = ops.matmul(weightings, aggregated)
+        out = ops.reshape(out, (-1, self.out_channels))
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_j):
+        return x_j
+
+    def aggregate(self, inputs, edge_index=None, index=None, dim_size=None, symnorm_weight=None, **kwargs):
+        if index is None and edge_index is not None:
+            index = edge_index[1]
+
+        outs = []
+        for aggr in self.aggregators:
+            if aggr == "symnorm":
+                inp = inputs if symnorm_weight is None else inputs * ops.expand_dims(symnorm_weight, -1)
+                out = scatter(inp, index, dim=0, dim_size=dim_size, reduce="sum")
+            elif aggr in ("var", "std"):
+                mean = scatter(inputs, index, dim=0, dim_size=dim_size, reduce="mean")
+                mean_sq = scatter(inputs * inputs, index, dim=0, dim_size=dim_size, reduce="mean")
+                out = mean_sq - mean * mean
+                if aggr == "std":
+                    out = ops.sqrt(ops.maximum(out, 1e-5))
+            else:
+                out = scatter(inputs, index, dim=0, dim_size=dim_size, reduce=aggr)
+            outs.append(out)
+
+        if len(outs) > 1:
+            return ops.stack(outs, axis=1)
+        return outs[0]
+

k3_node/layers/conv/fa_conv.py

@@ -0,0 +1,107 @@
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm
+
+
+class FAConv(MessagePassing):
+    r"""The Frequency Adaptive Graph Convolution operator from the
+    `"Beyond Low-Frequency Information in Graph Convolutional Networks"
+    <https://arxiv.org/abs/2101.00797>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import FAConv
+
+        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
+
+        x_0 = x  # initial node representations
+        layer = FAConv(channels=8, eps=0.1)
+        out = layer(x, x_0, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        channels: int,
+        eps: float = 0.1,
+        dropout: float = 0.0,
+        cached: bool = False,
+        add_self_loops: bool = True,
+        normalize: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+
+        self.channels = channels
+        self.eps = eps
+        self.dropout_rate = dropout
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.normalize = normalize
+
+        self.att_l = Dense(1, use_bias=False)
+        self.att_r = Dense(1, use_bias=False)
+
+    def build(self, input_shape=None):
+        self.att_l.build((None, self.channels))
+        self.att_r.build((None, self.channels))
+        self.built = True
+
+    def call(self, inputs, x_0=None, edge_index=None, edge_weight=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 4:
+                    x, x_0, edge_index, edge_weight = inputs
+                elif len(inputs) == 3:
+                    x, x_0, edge_index = inputs
+                elif len(inputs) == 2:
+                    x, edge_index = inputs
+                    x_0 = x
+                else:
+                    raise ValueError(f"Unexpected input length {len(inputs)}")
+            else:
+                raise ValueError("Expected inputs with edge_index")
+        else:
+            x = inputs
+            if x_0 is None:
+                x_0 = x
+
+        if not self.built:
+            self.build()
+
+        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,
+            )
+
+        alpha_l = self.att_l(x)
+        alpha_r = self.att_r(x)
+
+        out = self.propagate(
+            edge_index,
+            x=x,
+            alpha=(alpha_l, alpha_r),
+            edge_weight=edge_weight,
+        )
+
+        if self.eps != 0.0:
+            out = out + self.eps * x_0
+
+        return out
+
+    def message(self, x_j, alpha_j, alpha_i, edge_weight=None):
+        alpha = ops.tanh(alpha_j + alpha_i)
+        if edge_weight is not None:
+            alpha = alpha * ops.expand_dims(edge_weight, -1)
+        return x_j * alpha
+

k3_node/layers/conv/feast_conv.py

@@ -0,0 +1,126 @@
+from keras import ops
+from keras.layers import Dense
+
+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 FeaStConv(MessagePassing):
+    r"""The (fault-tolerant) feature-steered graph convolution operator from
+    the `"FeaStNet: Feature-Steered Graph Convolutions for 3D Shape Analysis"
+    <https://arxiv.org/abs/1706.05206>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import FeaStConv
+
+        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 = FeaStConv(in_channels=8, out_channels=16, heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        heads: int = 1,
+        add_self_loops: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "mean")
+        super().__init__(node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.add_self_loops = add_self_loops
+        self.use_bias = bias
+
+        self.lin = Dense(heads * out_channels, use_bias=False)
+        self.u = self.add_weight(
+            shape=(in_channels, heads),
+            initializer="glorot_uniform",
+            name="u",
+        )
+        self.c = self.add_weight(
+            shape=(heads,),
+            initializer="zeros",
+            name="c",
+        )
+
+        if bias:
+            self.bias = self.add_weight(
+                shape=(out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape=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
+
+        if not self.built:
+            self.build()
+
+        if isinstance(x, (list, tuple)):
+            x_src, x_dst = x
+        else:
+            x_src = x_dst = x
+
+        num_nodes = ops.shape(x_dst)[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_src, x_dst),
+            keep_mask=keep_mask,
+            size=(ops.shape(x_src)[0], num_nodes),
+        )
+        if keep_mask is not None and self.aggr == "mean":
+            # Masked messages are zero but still counted by the mean; rescale to the kept count.
+            col = ops.cast(edge_index[1], "int32")
+            count_all = degree(col, num_nodes=num_nodes, dtype=out.dtype)
+            count_kept = segment_sum(ops.cast(keep_mask, out.dtype), col, num_segments=num_nodes)
+            out = out * ops.expand_dims(count_all / ops.maximum(count_kept, 1.0), -1)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_i, x_j, keep_mask=None):
+        q = ops.matmul(x_j - x_i, self.u) + self.c
+        q = ops.softmax(q, axis=-1)
+        x_j_mapped = ops.reshape(self.lin(x_j), (-1, self.heads, self.out_channels))
+        msg = ops.sum(x_j_mapped * ops.expand_dims(q, -1), axis=1)
+        # For max/min a duplicated self-loop message is harmless, so only sum/mean need masking.
+        if keep_mask is not None and self.aggr in ("add", "sum", "mean"):
+            msg = msg * ops.expand_dims(ops.cast(keep_mask, msg.dtype), -1)
+        return msg
+

k3_node/layers/conv/film_conv.py

@@ -0,0 +1,143 @@
+import copy
+from typing import Optional, Union, Tuple, Callable
+from keras import ops, activations
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class FiLMConv(MessagePassing):
+    r"""The FiLM graph convolutional operator from the
+    `"GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation"
+    <https://arxiv.org/abs/1906.12192>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import FiLMConv
+
+        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 = FiLMConv(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]],
+        out_channels: int,
+        num_relations: int = 1,
+        nn: Optional[Callable] = None,
+        act: Optional[Union[str, Callable]] = "relu",
+        aggr: str = "mean",
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_relations = max(num_relations, 1)
+        self.act = activations.get(act) if act is not None else None
+
+        if isinstance(in_channels, int):
+            self.in_channels_l = in_channels
+            self.in_channels_r = in_channels
+        else:
+            self.in_channels_l, self.in_channels_r = in_channels
+
+        self.lins = [
+            Dense(out_channels, use_bias=False) for _ in range(self.num_relations)
+        ]
+        self.films = []
+        for _ in range(self.num_relations):
+            if nn is None:
+                self.films.append(Dense(2 * out_channels, use_bias=True))
+            else:
+                self.films.append(copy.deepcopy(nn))
+
+        self.lin_skip = Dense(out_channels, use_bias=False)
+        if nn is None:
+            self.film_skip = Dense(2 * out_channels, use_bias=False)
+        else:
+            self.film_skip = copy.deepcopy(nn)
+
+    def build(self, input_shape=None):
+        for lin in self.lins:
+            lin.build((None, self.in_channels_l))
+        for film in self.films:
+            film.build((None, self.in_channels_r))
+        self.lin_skip.build((None, self.in_channels_r))
+        self.film_skip.build((None, self.in_channels_r))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_type=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, edge_index, edge_type = inputs
+                elif len(inputs) == 2:
+                    x, edge_index = inputs
+                else:
+                    raise ValueError(f"Unexpected input length {len(inputs)}")
+            else:
+                raise ValueError("Expected inputs with edge_index")
+        else:
+            x = inputs
+
+        if not self.built:
+            self.build()
+
+        if isinstance(x, (list, tuple)):
+            x_l, x_r = x
+        else:
+            x_l = x_r = x
+
+        edge_index = ops.cast(edge_index, "int32")
+        if edge_type is not None:
+            edge_type = ops.cast(edge_type, "int32")
+
+        # Skip connection
+        film_s = self.film_skip(x_r)
+        beta_s, gamma_s = ops.split(film_s, 2, axis=-1)
+        out = gamma_s * self.lin_skip(x_r) + beta_s
+        if self.act is not None:
+            out = self.act(out)
+
+        # Message passing per relation
+        num_nodes = ops.shape(x_r)[0]
+        size = (ops.shape(x_l)[0], num_nodes)
+
+        for i in range(self.num_relations):
+            if edge_type is not None:
+                mask = ops.equal(edge_type, i)
+                where_mask = ops.where(mask)
+                idx = where_mask[0] if isinstance(where_mask, (list, tuple)) else where_mask
+                idx = ops.reshape(idx, (-1,))
+                idx = ops.cast(idx, "int32")
+                if ops.shape(idx)[0] == 0:
+                    continue
+                edge_index_i = ops.take(edge_index, idx, axis=1)
+            else:
+                edge_index_i = edge_index
+
+            film_val = self.films[i](x_r)
+            beta, gamma = ops.split(film_val, 2, axis=-1)
+            h = self.propagate(
+                edge_index_i,
+                x=self.lins[i](x_l),
+                beta=beta,
+                gamma=gamma,
+                size=size,
+            )
+            out = out + h
+
+        return out
+
+    def message(self, x_j, beta_i, gamma_i):
+        out = gamma_i * x_j + beta_i
+        if self.act is not None:
+            out = self.act(out)
+        return out
+