k3-node 1.0.0__py3-none-any.whl

k3_node/layers/dense/dense_gin_conv.py

@@ -0,0 +1,88 @@
+from keras import initializers, layers, ops
+
+
+class DenseGINConv(layers.Layer):
+    r"""Applies the dense Graph Isomorphism Network (GIN) convolutional operator
+    from the `"How Powerful are Graph Neural Networks?"
+    <https://arxiv.org/abs/1810.00826>`_ paper.
+
+    .. math::
+        \mathbf{X}^{\prime} = \mathrm{MLP} \left( \left( \mathbf{A} + (1 + \epsilon)
+        \cdot \mathbf{I} \right) \cdot \mathbf{X} \right)
+
+    Args:
+        nn (callable): A neural network layer (e.g. MLP or Dense) mapping feature
+            representations to output representations.
+        eps (float, optional): (Initial) :math:`\epsilon`-value. (default: :obj:`0.0`)
+        train_eps (bool, optional): If set to :obj:`True`, :math:`\epsilon` will
+            be a trainable parameter. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import DenseGINConv
+
+        x = np.random.rand(2, 10, 8).astype("float32")  # batch of 2 graphs, 10 nodes, 8 features
+        adj = (np.random.rand(2, 10, 10) > 0.7).astype("float32")  # dense adjacency matrices
+
+        layer = DenseGINConv(keras.Sequential([keras.layers.Dense(16, activation="relu"), keras.layers.Dense(16)]))
+        out = layer(x, adj)
+        print(tuple(out.shape))  # (2, 10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        nn,
+        eps: float = 0.0,
+        train_eps: bool = False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.nn = nn
+        self.initial_eps = eps
+        self.train_eps = train_eps
+
+        self.eps = self.add_weight(
+            shape=(1,),
+            initializer=initializers.Constant(eps),
+            trainable=train_eps,
+            name="eps",
+        )
+
+    def reset_parameters(self):
+        if hasattr(self.nn, 'reset_parameters'):
+            self.nn.reset_parameters()
+        self.eps.assign(ops.cast(ops.convert_to_tensor([self.initial_eps]), dtype=self.eps.dtype))
+
+    def call(self, x, adj, mask=None, add_loop: bool = True):
+        is_2d_x = (len(ops.shape(x)) == 2)
+        if is_2d_x:
+            x = ops.expand_dims(x, axis=0)
+        if len(ops.shape(adj)) == 2:
+            adj = ops.expand_dims(adj, axis=0)
+
+        N = ops.shape(adj)[1]
+
+        out = ops.matmul(adj, x)
+        if add_loop:
+            out = (1.0 + self.eps) * x + out
+
+        out = self.nn(out)
+
+        if mask is not None:
+            out = out * ops.cast(ops.reshape(mask, (-1, N, 1)), x.dtype)
+
+        if is_2d_x and ops.shape(out)[0] == 1:
+            out = ops.squeeze(out, axis=0)
+
+        return out
+
+    def compute_output_shape(self, input_shape):
+        if hasattr(self.nn, 'compute_output_shape'):
+            return self.nn.compute_output_shape(input_shape)
+        return input_shape
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(nn={self.nn})'
+

k3_node/layers/dense/dense_graph_conv.py

@@ -0,0 +1,95 @@
+from typing import Optional
+from keras import layers, ops
+from .linear import Linear
+
+
+class DenseGraphConv(layers.Layer):
+    r"""See :class:`torch_geometric.nn.conv.GraphConv`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DenseGraphConv
+
+        x = np.random.rand(2, 10, 8).astype("float32")  # batch of 2 graphs, 10 nodes, 8 features
+        adj = (np.random.rand(2, 10, 10) > 0.7).astype("float32")  # dense adjacency matrices
+
+        layer = DenseGraphConv(in_channels=8, out_channels=16)
+        out = layer(x, adj)
+        print(tuple(out.shape))  # (2, 10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        aggr: str = 'add',
+        bias: bool = True,
+        **kwargs
+    ):
+        if aggr not in ('add', 'mean', 'max'):
+            raise ValueError(f"aggr must be one of 'add', 'mean', 'max' (got '{aggr}')")
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.aggr = aggr
+        self.use_bias = bias
+
+        self.lin_rel = Linear(in_channels, out_channels, bias=bias, name="lin_rel")
+        self.lin_root = Linear(in_channels, out_channels, bias=False, name="lin_root")
+
+    def reset_parameters(self):
+        self.lin_rel.reset_parameters()
+        self.lin_root.reset_parameters()
+
+    def build(self, input_shape):
+        if isinstance(input_shape, (tuple, list)):
+            x_shape = input_shape[0]
+        else:
+            x_shape = input_shape
+        self.lin_rel.build(x_shape)
+        self.lin_root.build(x_shape)
+        super().build(input_shape)
+
+    def call(self, x, adj, mask=None):
+        is_2d_x = (len(ops.shape(x)) == 2)
+        if is_2d_x:
+            x = ops.expand_dims(x, axis=0)
+        if len(ops.shape(adj)) == 2:
+            adj = ops.expand_dims(adj, axis=0)
+
+        N = ops.shape(x)[1]
+
+        if self.aggr == 'add':
+            out = ops.matmul(adj, x)
+        elif self.aggr == 'mean':
+            deg = ops.maximum(ops.sum(adj, axis=-1, keepdims=True), 1.0)
+            out = ops.matmul(adj, x) / deg
+        elif self.aggr == 'max':
+            x_src = ops.expand_dims(x, axis=1)  # [B, 1, N, C]
+            adj_expanded = ops.expand_dims(adj, axis=-1)  # [B, N, N, 1]
+            masked = ops.where(adj_expanded != 0, x_src, -1e9)
+            out = ops.max(masked, axis=2)  # [B, N, C]
+            out = ops.where(out <= -1e8, 0.0, out)
+
+        out = self.lin_rel(out) + self.lin_root(x)
+
+        if mask is not None:
+            out = out * ops.cast(ops.reshape(mask, (-1, N, 1)), x.dtype)
+
+        if is_2d_x and ops.shape(out)[0] == 1:
+            out = ops.squeeze(out, axis=0)
+
+        return out
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(input_shape, (tuple, list)):
+            x_shape = input_shape[0]
+        else:
+            x_shape = input_shape
+        return (*x_shape[:-1], self.out_channels)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels})')

k3_node/layers/dense/dense_sage_conv.py

@@ -0,0 +1,85 @@
+from keras import layers, ops
+from .linear import Linear
+
+
+class DenseSAGEConv(layers.Layer):
+    r"""See :class:`torch_geometric.nn.conv.SAGEConv`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DenseSAGEConv
+
+        x = np.random.rand(2, 10, 8).astype("float32")  # batch of 2 graphs, 10 nodes, 8 features
+        adj = (np.random.rand(2, 10, 10) > 0.7).astype("float32")  # dense adjacency matrices
+
+        layer = DenseSAGEConv(in_channels=8, out_channels=16)
+        out = layer(x, adj)
+        print(tuple(out.shape))  # (2, 10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        normalize: bool = False,
+        bias: bool = True,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.normalize = normalize
+        self.use_bias = bias
+
+        self.lin_rel = Linear(in_channels, out_channels, bias=False, name="lin_rel")
+        self.lin_root = Linear(in_channels, out_channels, bias=bias, name="lin_root")
+
+    def reset_parameters(self):
+        self.lin_rel.reset_parameters()
+        self.lin_root.reset_parameters()
+
+    def build(self, input_shape):
+        if isinstance(input_shape, (tuple, list)):
+            x_shape = input_shape[0]
+        else:
+            x_shape = input_shape
+        self.lin_rel.build(x_shape)
+        self.lin_root.build(x_shape)
+        super().build(input_shape)
+
+    def call(self, x, adj, mask=None):
+        is_2d_x = (len(ops.shape(x)) == 2)
+        if is_2d_x:
+            x = ops.expand_dims(x, axis=0)
+        if len(ops.shape(adj)) == 2:
+            adj = ops.expand_dims(adj, axis=0)
+
+        N = ops.shape(adj)[1]
+
+        deg = ops.maximum(ops.sum(adj, axis=-1, keepdims=True), 1.0)
+        out = ops.matmul(adj, x) / deg
+        out = self.lin_rel(out) + self.lin_root(x)
+
+        if self.normalize:
+            norm = ops.maximum(ops.sqrt(ops.sum(ops.power(out, 2), axis=-1, keepdims=True)), 1e-12)
+            out = out / norm
+
+        if mask is not None:
+            out = out * ops.cast(ops.reshape(mask, (-1, N, 1)), x.dtype)
+
+        if is_2d_x and ops.shape(out)[0] == 1:
+            out = ops.squeeze(out, axis=0)
+
+        return out
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(input_shape, (tuple, list)):
+            x_shape = input_shape[0]
+        else:
+            x_shape = input_shape
+        return (*x_shape[:-1], self.out_channels)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels})')

k3_node/layers/dense/diff_pool.py

@@ -0,0 +1,76 @@
+from typing import Optional, Tuple
+from keras import ops
+
+
+def dense_diff_pool(
+    x,
+    adj,
+    s,
+    mask: Optional[any] = None,
+    normalize: bool = True,
+) -> Tuple[any, any, any, any]:
+    r"""The differentiable pooling operator from the `"Hierarchical Graph
+    Representation Learning with Differentiable Pooling"
+    <https://arxiv.org/abs/1806.08804>`_ paper.
+
+    .. math::
+        \mathbf{X}^{\prime} &= {\mathrm{softmax}(\mathbf{S})}^{\top} \cdot
+        \mathbf{X}
+
+        \mathbf{A}^{\prime} &= {\mathrm{softmax}(\mathbf{S})}^{\top} \cdot
+        \mathbf{A} \cdot \mathrm{softmax}(\mathbf{S})
+
+    Args:
+        x: Node feature tensor [B, N, F] or [N, F].
+        adj: Adjacency tensor [B, N, N] or [N, N].
+        s: Assignment tensor [B, N, C] or [N, C].
+        mask: Mask tensor [B, N] indicating valid nodes. (default: None)
+        normalize: If set to False, link prediction loss is not divided by total elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import dense_diff_pool
+
+        x = np.random.rand(2, 10, 8).astype("float32")  # batch of 2 graphs, 10 nodes, 8 features
+        adj = (np.random.rand(2, 10, 10) > 0.7).astype("float32")  # dense adjacency matrices
+        s = np.random.rand(2, 10, 3).astype("float32")  # assignment scores for 3 clusters
+
+        x_pool, adj_pool, link_loss, entropy_loss = dense_diff_pool(x, adj, s)
+        print(tuple(x_pool.shape), tuple(adj_pool.shape))  # (2, 3, 8) (2, 3, 3)
+        ```
+    """
+    # Plain NumPy inputs cannot be mixed with backend tensors (e.g. `ndarray @ torch.Tensor`).
+    x, adj, s = ops.convert_to_tensor(x), ops.convert_to_tensor(adj), ops.convert_to_tensor(s)
+    if len(ops.shape(x)) == 2:
+        x = ops.expand_dims(x, axis=0)
+    if len(ops.shape(adj)) == 2:
+        adj = ops.expand_dims(adj, axis=0)
+    if len(ops.shape(s)) == 2:
+        s = ops.expand_dims(s, axis=0)
+
+    batch_size = ops.shape(x)[0]
+    num_nodes = ops.shape(x)[1]
+
+    s = ops.softmax(s, axis=-1)
+
+    if mask is not None:
+        mask_m = ops.cast(ops.reshape(mask, (batch_size, num_nodes, 1)), x.dtype)
+        x = x * mask_m
+        s = s * mask_m
+
+    s_t = ops.transpose(s, (0, 2, 1))
+
+    out = ops.matmul(s_t, x)
+    out_adj = ops.matmul(ops.matmul(s_t, adj), s)
+
+    link = adj - ops.matmul(s, s_t)
+    link_loss = ops.sqrt(ops.sum(ops.power(link, 2)))
+    if normalize:
+        numel = ops.cast(ops.prod(ops.shape(adj)), link_loss.dtype)
+        link_loss = link_loss / numel
+
+    ent_loss = ops.mean(ops.sum(-s * ops.log(s + 1e-15), axis=-1))
+
+    return out, out_adj, link_loss, ent_loss
+

k3_node/layers/dense/dmon_pool.py

@@ -0,0 +1,223 @@
+from typing import List, Optional, Tuple, Union
+from keras import layers, ops
+
+from .linear import Linear
+
+
+class MLP(layers.Layer):
+    r"""A simple Multi-Layer Perceptron (MLP) model."""
+    def __init__(
+        self,
+        channel_list: List[int],
+        act: Optional[Union[str, any]] = None,
+        norm: Optional[Union[str, any]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.channel_list = channel_list
+        self.lins = [
+            Linear(in_c, out_c)
+            for in_c, out_c in zip(channel_list[:-1], channel_list[1:])
+        ]
+        self.act = layers.Activation(act) if act is not None else None
+
+    @property
+    def in_channels(self) -> int:
+        r"""Size of each input sample."""
+        return self.channel_list[0]
+
+    @property
+    def out_channels(self) -> int:
+        r"""Size of each output sample."""
+        return self.channel_list[-1]
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for lin in self.lins:
+            if hasattr(lin, "reset_parameters"):
+                lin.reset_parameters()
+
+    def call(self, x):
+        for i, lin in enumerate(self.lins):
+            x = lin(x)
+            if self.act is not None and i < len(self.lins) - 1:
+                x = self.act(x)
+        return x
+
+
+class DMoNPooling(layers.Layer):
+    r"""The spectral modularity pooling operator from the `"Graph Clustering
+    with Graph Neural Networks" <https://arxiv.org/abs/2006.16904>`_ paper.
+
+    .. math::
+        \mathbf{X}^{\prime} &= {\mathrm{softmax}(\mathbf{S})}^{\top} \cdot
+        \mathbf{X}
+
+        \mathbf{A}^{\prime} &= {\mathrm{softmax}(\mathbf{S})}^{\top} \cdot
+        \mathbf{A} \cdot \mathrm{softmax}(\mathbf{S})
+
+    Args:
+        channels (int or List[int]): Size of each input sample. If given as a
+            list, will construct an MLP based on the given feature sizes.
+        k (int): The number of clusters.
+        dropout (float, optional): Dropout probability. (default: :obj:`0.0`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DMoNPooling
+
+        x = np.random.rand(2, 10, 8).astype("float32")  # batch of 2 graphs, 10 nodes, 8 features
+        adj = (np.random.rand(2, 10, 10) > 0.7).astype("float32")  # dense adjacency matrices
+
+        layer = DMoNPooling(channels=8, k=3)  # 3 clusters
+        s, x_pool, adj_pool, spectral_loss, ortho_loss, cluster_loss = layer(x, adj)
+        print(tuple(s.shape))  # (2, 10, 3): soft cluster assignments
+        print(tuple(x_pool.shape), tuple(adj_pool.shape))  # (2, 3, 8) (2, 3, 3)
+        ```
+    """
+    def __init__(
+        self,
+        channels: Union[int, List[int]],
+        k: int,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if isinstance(channels, int):
+            channels = [channels]
+
+        self.channels = channels
+        self.k = k
+        self.dropout = dropout
+        self.mlp = MLP(channels + [k], act=None, norm=None)
+        self.drop = layers.Dropout(dropout) if dropout > 0.0 else None
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.mlp.reset_parameters()
+
+    def call(
+        self,
+        x,
+        adj,
+        mask: Optional[any] = None,
+        training: bool = False,
+    ) -> Tuple[any, any, any, any, any, any]:
+        r"""Forward pass.
+
+        Args:
+            x: Node feature tensor [B, N, F] or [N, F].
+            adj: Adjacency tensor [B, N, N] or [N, N].
+            mask: Mask tensor [B, N] indicating valid nodes. (default: None)
+            training: Whether layer is in training mode.
+        """
+        if len(ops.shape(x)) == 2:
+            x = ops.expand_dims(x, axis=0)
+        if len(ops.shape(adj)) == 2:
+            adj = ops.expand_dims(adj, axis=0)
+
+        s = self.mlp(x)
+        if self.drop is not None:
+            s = self.drop(s, training=training)
+        s = ops.softmax(s, axis=-1)
+        if mask is not None:
+            batch_size = ops.shape(x)[0]
+            num_nodes = ops.shape(x)[1]
+            mask_m = ops.cast(ops.reshape(mask, (batch_size, num_nodes, 1)), x.dtype)
+            s = s * mask_m
+        out, out_adj, spectral_loss, ortho_loss, cluster_loss = dense_dmon_pool(x, adj, s, mask=mask)
+        return s, out, out_adj, spectral_loss, ortho_loss, cluster_loss
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.mlp.in_channels}, '
+                f'num_clusters={self.mlp.out_channels})')
+
+
+def dense_dmon_pool(
+    x,
+    adj,
+    s,
+    mask: Optional[any] = None,
+) -> Tuple[any, any, any, any, any]:
+    r"""Functional dense DMoN pooling operator.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import dense_dmon_pool
+
+        x = np.random.rand(2, 10, 8).astype("float32")  # batch of 2 graphs, 10 nodes, 8 features
+        adj = (np.random.rand(2, 10, 10) > 0.7).astype("float32")  # dense adjacency matrices
+        s = np.random.rand(2, 10, 3).astype("float32")  # assignment scores for 3 clusters
+
+        x_pool, adj_pool, spectral_loss, ortho_loss, cluster_loss = dense_dmon_pool(x, adj, s)
+        print(tuple(x_pool.shape), tuple(adj_pool.shape))  # (2, 3, 8) (2, 3, 3)
+        ```
+    """
+    # Plain NumPy inputs cannot be mixed with backend tensors (e.g. `ndarray @ torch.Tensor`).
+    x, adj, s = ops.convert_to_tensor(x), ops.convert_to_tensor(adj), ops.convert_to_tensor(s)
+    if len(ops.shape(x)) == 2:
+        x = ops.expand_dims(x, axis=0)
+    if len(ops.shape(adj)) == 2:
+        adj = ops.expand_dims(adj, axis=0)
+    if len(ops.shape(s)) == 2:
+        s = ops.expand_dims(s, axis=0)
+
+    batch_size = ops.shape(x)[0]
+    num_nodes = ops.shape(x)[1]
+    C = ops.shape(s)[-1]
+
+    if mask is None:
+        mask = ops.ones((batch_size, num_nodes, 1), dtype=x.dtype)
+    else:
+        mask = ops.cast(ops.reshape(mask, (batch_size, num_nodes, 1)), x.dtype)
+
+    x = x * mask
+    s = s * mask
+
+    out = ops.selu(ops.matmul(ops.swapaxes(s, 1, 2), x))
+    out_adj = ops.matmul(ops.matmul(ops.swapaxes(s, 1, 2), adj), s)
+
+    # Spectral loss:
+    degrees = ops.sum(adj, axis=-1, keepdims=True) * mask
+    degrees_t = ops.swapaxes(degrees, 1, 2)
+
+    m = ops.sum(degrees, axis=(1, 2)) / 2.0
+    m_expand = ops.broadcast_to(ops.reshape(m, (-1, 1, 1)), (batch_size, C, C))
+
+    ca = ops.matmul(ops.swapaxes(s, 1, 2), degrees)
+    cb = ops.matmul(degrees_t, s)
+
+    normalizer = ops.matmul(ca, cb) / 2.0 / (m_expand + 1e-15)
+    decompose = out_adj - normalizer
+    tr = ops.sum(ops.diagonal(decompose, axis1=1, axis2=2), axis=-1)
+    spectral_loss = ops.mean(-tr / 2.0 / (m + 1e-15))
+
+    # Orthogonality regularization:
+    ss = ops.matmul(ops.swapaxes(s, 1, 2), s)
+    i_s = ops.eye(C, dtype=ss.dtype)
+    norm_ss = ops.sqrt(ops.sum(ops.power(ss, 2), axis=(-1, -2), keepdims=True) + 1e-15)
+    norm_is = ops.sqrt(ops.cast(C, ss.dtype))
+    diff = (ss / norm_ss) - (ops.expand_dims(i_s, axis=0) / norm_is)
+    ortho_loss = ops.mean(ops.sqrt(ops.sum(ops.power(diff, 2), axis=(-1, -2)) + 1e-15))
+
+    # Cluster loss:
+    cluster_size = ops.sum(s, axis=1)
+    cluster_loss = ops.sqrt(ops.sum(ops.power(cluster_size, 2), axis=1) + 1e-15)
+    cluster_loss = ops.mean(cluster_loss / (ops.sum(mask, axis=1) + 1e-15) * norm_is - 1.0)
+
+    # Fix and normalize coarsened adjacency matrix:
+    eye_c = ops.expand_dims(ops.eye(C, dtype=out_adj.dtype), axis=0)
+    out_adj = out_adj * (1.0 - eye_c)
+    d = ops.sum(out_adj, axis=-1, keepdims=False)
+    d = ops.expand_dims(ops.sqrt(d), axis=1) + 1e-15
+    out_adj = (out_adj / d) / ops.swapaxes(d, 1, 2)
+
+    return out, out_adj, spectral_loss, ortho_loss, cluster_loss
+
+
+dmon_pool = dense_dmon_pool
+
+