k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/nn_conv.py

@@ -0,0 +1,110 @@
+from typing import Callable, Union, Tuple
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class NNConv(MessagePassing):
+    r"""The continuous kernel-based convolutional operator from the
+    `"Neural Message Passing for Quantum Chemistry"
+    <https://arxiv.org/abs/1704.01212>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample, or a tuple for bipartite graphs.
+        out_channels: Size of each output sample.
+        nn: A neural network :math:`h_{\mathbf{\Theta}}` that maps edge features
+            to shape :obj:`[-1, in_channels * out_channels]`.
+        aggr: The aggregation scheme to use (``"add"``, ``"mean"``, ``"max"``).
+            (default: ``"add"``)
+        root_weight: If set to :obj:`False`, the layer will not add the
+            transformed root node features. (default: ``True``)
+        bias: If set to :obj:`False`, the layer will not learn an additive bias.
+            (default: ``True``)
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import NNConv
+
+        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
+
+        # `nn` maps each edge's features to an [in_channels * out_channels] weight matrix
+        layer = NNConv(in_channels=8, out_channels=16, nn=keras.layers.Dense(8 * 16))
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: int,
+        nn: Callable,
+        aggr: str = "add",
+        root_weight: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr=aggr, **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.nn = nn
+        self.root_weight = root_weight
+        self.use_bias = bias
+
+        if isinstance(in_channels, int):
+            self.in_channels_src = in_channels
+            self.in_channels_dst = in_channels
+        else:
+            self.in_channels_src = in_channels[0]
+            self.in_channels_dst = in_channels[1]
+
+        if root_weight:
+            self.lin_root = layers.Dense(out_channels, use_bias=False)
+        else:
+            self.lin_root = None
+
+    def build(self, input_shape):
+        if self.lin_root is not None:
+            self.lin_root.build((None, self.in_channels_dst))
+
+        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 call(self, x, edge_index=None, edge_attr=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]
+
+        weight = self.nn(edge_attr)
+        weight = ops.reshape(weight, (-1, self.in_channels_src, self.out_channels))
+
+        out = self.propagate(edge_index, x=x_src, weight=weight, size=size)
+
+        if self.root_weight and self.lin_root is not None and x_dst is not None:
+            out = out + self.lin_root(x_dst)
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def message(self, x_j, weight):
+        # x_j: [E, in_channels_src], weight: [E, in_channels_src, out_channels]
+        x_j = ops.expand_dims(x_j, axis=1)  # [E, 1, in_channels_src]
+        msg = ops.matmul(x_j, weight)       # [E, 1, out_channels]
+        return ops.squeeze(msg, axis=1)     # [E, out_channels]
+

k3_node/layers/conv/pan_conv.py

@@ -0,0 +1,100 @@
+from keras import ops
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.ops.creation import scatter
+
+
+class PANConv(MessagePassing):
+    r"""The path integral based convolution operator from the
+    `"Path Integral Based Convolution and Pooling for Graph Neural Networks"
+    <https://arxiv.org/abs/2004.14805>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PANConv
+
+        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 = PANConv(in_channels=8, out_channels=16, filter_size=2)
+        out, weights = layer(x, edge_index)  # also returns the learned path weights
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        filter_size: int,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_size = filter_size
+
+        self.lin = Dense(out_channels, use_bias=True)
+        self.weight = self.add_weight(
+            shape=(filter_size + 1,),
+            initializer="ones",
+            name="weight",
+        )
+
+    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()
+
+        num_nodes = ops.shape(x)[0]
+
+        # Construct adjacency matrix
+        if hasattr(edge_index, "shape") and len(edge_index.shape) == 2 and edge_index.shape[0] == 2:
+            row, col = edge_index[0], edge_index[1]
+            adj = scatter(
+                ops.stack([row, col], axis=-1),
+                ops.ones((ops.shape(row)[0],), dtype=x.dtype),
+                shape=(num_nodes, num_nodes),
+            )
+        else:
+            adj = ops.cast(edge_index, x.dtype)
+
+        # PAN entropy / path calculation
+        # M = sum_{k=0}^filter_size weight[k] * A^k
+        # M = sum_{k=0}^filter_size exp(-E(k)/T) * A^k
+        w = ops.softplus(self.weight)
+        eye = ops.eye(num_nodes, dtype=x.dtype)
+        M = self.weight[0] * eye
+        M = w[0] * eye
+        curr_adj = eye
+        for k in range(1, self.filter_size + 1):
+            curr_adj = ops.matmul(curr_adj, adj)
+            M = M + self.weight[k] * curr_adj
+            M = M + w[k] * curr_adj
+
+        deg = ops.sum(M, axis=1)
+        deg_inv_sqrt = ops.power(ops.maximum(deg, 1e-12), -0.5)
+        # Avoid inf / nan
+        deg_inv_sqrt = ops.where(ops.isfinite(deg_inv_sqrt), deg_inv_sqrt, 0.0)
+
+        M_norm = ops.expand_dims(deg_inv_sqrt, 0) * M * ops.expand_dims(deg_inv_sqrt, 1)
+        M_norm = ops.expand_dims(deg_inv_sqrt, 1) * M * ops.expand_dims(deg_inv_sqrt, 0)
+
+        out = ops.matmul(M_norm, x)
+        out = self.lin(out)
+
+        return out, M_norm

k3_node/layers/conv/pdn_conv.py

@@ -0,0 +1,109 @@
+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 PDNConv(MessagePassing):
+    r"""The pathfinder discovery network convolutional operator from the
+    `"Pathfinder Discovery Networks for Neural Message Passing"
+    <https://arxiv.org/abs/2010.12878>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PDNConv
+
+        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 = PDNConv(in_channels=8, out_channels=16, edge_dim=3, hidden_channels=16)
+        out = layer(x, edge_index, edge_attr)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        edge_dim: int,
+        hidden_channels: int = 16,
+        add_self_loops: bool = True,
+        normalize: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.edge_dim = edge_dim
+        self.hidden_channels = hidden_channels
+        self.add_self_loops = add_self_loops
+        self.normalize = normalize
+        self.use_bias = bias
+
+        self.mlp_1 = Dense(hidden_channels, activation="relu")
+        self.mlp_2 = Dense(1, activation="sigmoid")
+        self.lin = Dense(out_channels, use_bias=False)
+
+        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.mlp_1.build((None, self.edge_dim))
+        self.mlp_2.build((None, self.hidden_channels))
+        self.lin.build((None, self.in_channels))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_attr=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, edge_index, edge_attr = 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 not self.built:
+            self.build()
+
+        if edge_attr is not None:
+            edge_attr = self.mlp_1(edge_attr)
+            edge_attr = ops.squeeze(self.mlp_2(edge_attr), -1)
+
+        num_nodes = ops.shape(x)[0]
+        if self.normalize:
+            edge_index, edge_attr = gcn_norm(
+                edge_index,
+                edge_attr,
+                num_nodes=num_nodes,
+                add_self_loops=self.add_self_loops,
+                dtype=x.dtype,
+            )
+
+        x = self.lin(x)
+        out = self.propagate(edge_index, x=x, edge_weight=edge_attr)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_j, edge_weight=None):
+        return x_j if edge_weight is None else ops.expand_dims(edge_weight, -1) * x_j
+

k3_node/layers/conv/pna_conv.py

@@ -0,0 +1,177 @@
+from typing import Optional, Union, List, Callable
+from keras import ops, activations
+from keras.layers import Dense
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.aggr import DegreeScalerAggregation
+from k3_node.ops.creation import repeat
+
+
+class PNAConv(MessagePassing):
+    r"""The Principal Neighbourhood Aggregation graph convolutional operator
+    from the `"Principal Neighbourhood Aggregation for Graph Nets"
+    <https://arxiv.org/abs/2004.05718>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PNAConv
+
+        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
+
+        deg = np.array([0, 2, 4, 3, 1])  # in-degree histogram of the training graphs
+        layer = PNAConv(in_channels=8, out_channels=16, aggregators=["mean", "max", "min", "std"],
+                        scalers=["identity", "amplification"], deg=deg)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    @staticmethod
+    def get_degree_histogram(graphs):
+        r"""Returns the histogram of in-degrees over ``graphs`` (a dataset, a list of graphs or a
+        loader), which PNA uses to normalize its degree scalers:
+        ``deg[d]`` is the number of nodes with ``d`` incoming edges."""
+        import numpy as np
+
+        counts = []
+        for graph in graphs:
+            target = np.asarray(ops.convert_to_numpy(graph.edge_index))[1]
+            counts.append(np.bincount(np.bincount(target, minlength=graph.num_nodes)))
+        hist = np.zeros(max(len(c) for c in counts), dtype=np.int64)
+        for c in counts:
+            hist[: len(c)] += c
+        return hist
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        aggregators: List[str],
+        scalers: List[str],
+        deg,
+        edge_dim: Optional[int] = None,
+        towers: int = 1,
+        pre_layers: int = 1,
+        post_layers: int = 1,
+        divide_input: bool = False,
+        act: Union[str, Callable, None] = "relu",
+        train_norm: bool = False,
+        **kwargs,
+    ):
+        aggr = DegreeScalerAggregation(aggregators, scalers, deg, train_norm)
+        super().__init__(aggr=aggr, node_dim=0, **kwargs)
+
+        if divide_input:
+            assert in_channels % towers == 0
+        assert out_channels % towers == 0
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.aggregators = aggregators
+        self.scalers = scalers
+        self.edge_dim = edge_dim
+        self.towers = towers
+        self.divide_input = divide_input
+        self.act = activations.get(act) if act is not None else None
+
+        self.F_in = in_channels // towers if divide_input else in_channels
+        self.F_out = out_channels // towers
+
+        if edge_dim is not None:
+            self.edge_encoder = Dense(self.F_in, use_bias=False)
+        else:
+            self.edge_encoder = None
+
+        self.pre_nns = []
+        self.post_nns = []
+        for _ in range(towers):
+            pre_dim = (3 if edge_dim else 2) * self.F_in
+            pre_layers_list = [Dense(self.F_in, use_bias=True)]
+            for _ in range(pre_layers - 1):
+                pre_layers_list.append(Dense(self.F_in, activation=self.act, use_bias=True))
+            self.pre_nns.append(pre_layers_list)
+
+            post_in_dim = (len(aggregators) * len(scalers) + 1) * self.F_in
+            post_layers_list = [Dense(self.F_out, use_bias=True)]
+            for _ in range(post_layers - 1):
+                post_layers_list.append(Dense(self.F_out, activation=self.act, use_bias=True))
+            self.post_nns.append(post_layers_list)
+
+        self.lin = Dense(out_channels, use_bias=True)
+
+    def build(self, input_shape=None):
+        if self.edge_encoder is not None:
+            self.edge_encoder.build((None, self.edge_dim))
+        for pre_list in self.pre_nns:
+            dim = (3 if self.edge_dim else 2) * self.F_in
+            for l in pre_list:
+                l.build((None, dim))
+                dim = self.F_in
+        for post_list in self.post_nns:
+            dim = (len(self.aggregators) * len(self.scalers) + 1) * self.F_in
+            for l in post_list:
+                l.build((None, dim))
+                dim = self.F_out
+        self.lin.build((None, self.out_channels))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, edge_attr=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, edge_index, edge_attr = 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 not self.built:
+            self.build()
+
+        num_nodes = ops.shape(x)[0]
+        if self.divide_input:
+            x_towers = ops.reshape(x, (-1, self.towers, self.F_in))
+        else:
+            x_towers = repeat(ops.expand_dims(x, 1), self.towers, axis=1)
+
+        out = self.propagate(
+            edge_index,
+            x=x_towers,
+            edge_attr=edge_attr,
+            size=(num_nodes, num_nodes),
+        )
+
+        out = ops.concatenate([x_towers, out], axis=-1)
+
+        outs = []
+        for i, post_list in enumerate(self.post_nns):
+            h_i = out[:, i]
+            for l in post_list:
+                h_i = l(h_i)
+            outs.append(h_i)
+
+        out = ops.concatenate(outs, axis=-1)
+        return self.lin(out)
+
+    def message(self, x_i, x_j, edge_attr=None):
+        if edge_attr is not None and self.edge_encoder is not None:
+            edge_attr = self.edge_encoder(edge_attr)
+            edge_attr = repeat(ops.expand_dims(edge_attr, 1), self.towers, axis=1)
+            h = ops.concatenate([x_i, x_j, edge_attr], axis=-1)
+        else:
+            h = ops.concatenate([x_i, x_j], axis=-1)
+
+        hs = []
+        for i, pre_list in enumerate(self.pre_nns):
+            h_i = h[:, i]
+            for l in pre_list:
+                h_i = l(h_i)
+            hs.append(h_i)
+
+        return ops.stack(hs, axis=1)
+

k3_node/layers/conv/point_conv.py

@@ -0,0 +1,101 @@
+from typing import Optional, Callable, Tuple
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import remove_self_loops, add_self_loops
+
+
+class PointNetConv(MessagePassing):
+    r"""The PointNet set abstraction layer from the `"PointNet++: Deep
+    Hierarchical Feature Learning on Point Sets in a Metric Space"
+    <https://arxiv.org/abs/1706.02413>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import PointNetConv
+
+        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 = np.random.rand(10, 3).astype("float32")  # 3D node positions
+
+        local_nn = keras.Sequential([keras.layers.Dense(16, activation="relu"), keras.layers.Dense(16)])
+        layer = PointNetConv(local_nn=local_nn)
+        out = layer(x, pos, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        local_nn: Optional[Callable] = None,
+        global_nn: Optional[Callable] = None,
+        add_self_loops: bool = True,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "max")
+        super().__init__(**kwargs)
+
+        self.local_nn = local_nn
+        self.global_nn = global_nn
+        self.add_self_loops = add_self_loops
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(self, inputs, pos=None, edge_index=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, pos, edge_index = inputs
+                elif len(inputs) == 2:
+                    # x can be None or pos
+                    x, edge_index = inputs
+                else:
+                    raise ValueError(f"Unexpected input length {len(inputs)}")
+            else:
+                raise ValueError("Expected (x, pos, edge_index) or (x, edge_index)")
+        else:
+            x = inputs
+
+        if not self.built:
+            self.build()
+
+        if isinstance(pos, (list, tuple)):
+            pos_src, pos_dst = pos
+        else:
+            pos_src = pos_dst = pos
+
+        if isinstance(x, (list, tuple)):
+            x_src, x_dst = x
+        else:
+            x_src = x_dst = x
+
+        num_nodes = ops.shape(pos_dst)[0]
+        if self.add_self_loops:
+            edge_index, _ = remove_self_loops(edge_index)
+            edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+
+        out = self.propagate(
+            edge_index,
+            x=(x_src, x_dst),
+            pos=(pos_src, pos_dst),
+            size=(ops.shape(pos_src)[0], num_nodes),
+        )
+
+        if self.global_nn is not None:
+            out = self.global_nn(out)
+
+        return out
+
+    def message(self, x_j=None, pos_i=None, pos_j=None):
+        msg = pos_j - pos_i
+        if x_j is not None:
+            msg = ops.concatenate([x_j, msg], axis=-1)
+        if self.local_nn is not None:
+            msg = self.local_nn(msg)
+        return msg
+
+
+PointConv = PointNetConv
+

k3_node/layers/conv/point_gnn_conv.py

@@ -0,0 +1,90 @@
+from typing import Callable
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class PointGNNConv(MessagePassing):
+    r"""The PointGNN graph convolutional operator from the
+    `"Point-GNN: Graph Neural Network for 3D Object Detection in a Point Cloud"
+    <https://arxiv.org/abs/2003.01251>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import PointGNNConv
+
+        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 = np.random.rand(10, 3).astype("float32")  # 3D node positions
+
+        layer = PointGNNConv(
+            mlp_h=keras.layers.Dense(3),  # predicts a position offset per node
+            mlp_f=keras.layers.Dense(16),  # edge feature network
+            mlp_g=keras.layers.Dense(8),  # node update network (output size = input features)
+        )
+        out = layer(x, pos, edge_index)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        mlp_h: Callable,
+        mlp_f: Callable,
+        mlp_g: Callable,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "max")
+        super().__init__(node_dim=0, **kwargs)
+
+        self.mlp_h = mlp_h
+        self.mlp_f = mlp_f
+        self.mlp_g = mlp_g
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(self, inputs, pos=None, edge_index=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)):
+                if len(inputs) == 3:
+                    x, pos, edge_index = inputs
+                elif len(inputs) == 2:
+                    x, edge_index = inputs
+                else:
+                    raise ValueError(f"Unexpected input length {len(inputs)}")
+            else:
+                raise ValueError("Expected (x, pos, 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 isinstance(pos, (list, tuple)):
+            pos_src, pos_dst = pos
+        else:
+            pos_src = pos_dst = pos
+
+        num_nodes = ops.shape(pos_dst)[0]
+        out = self.propagate(
+            edge_index,
+            x=(x_src, x_dst),
+            pos=(pos_src, pos_dst),
+            size=(ops.shape(pos_src)[0], num_nodes),
+        )
+
+        out = self.mlp_g(out)
+        return x_dst + out
+
+    def message(self, pos_j, pos_i, x_i, x_j):
+        delta = self.mlp_h(x_i)
+        e = ops.concatenate([pos_j - pos_i + delta, x_j], axis=-1)
+        return self.mlp_f(e)
+