k3-node 1.0.0__py3-none-any.whl

k3_node/layers/dense/linear.py

@@ -0,0 +1,327 @@
+import math
+from typing import Any, Dict, List, Optional, Union
+from keras import initializers, layers, ops
+
+
+def _get_weight_initializer(initializer: Optional[str], in_channels: int):
+    if initializer in ('glorot', 'xavier_uniform'):
+        return initializers.GlorotUniform()
+    elif initializer in ('kaiming_uniform', 'he_uniform'):
+        return initializers.HeUniform()
+    elif initializer == 'uniform' or initializer is None:
+        if in_channels > 0:
+            bound = 1.0 / math.sqrt(in_channels)
+            return initializers.RandomUniform(minval=-bound, maxval=bound)
+        return initializers.GlorotUniform()
+    return initializers.get(initializer)
+
+
+def _get_bias_initializer(initializer: Optional[str], in_channels: int):
+    if initializer == 'zeros':
+        return initializers.Zeros()
+    elif initializer is None:
+        if in_channels > 0:
+            bound = 1.0 / math.sqrt(in_channels)
+            return initializers.RandomUniform(minval=-bound, maxval=bound)
+        return initializers.Zeros()
+    return initializers.get(initializer)
+
+
+class Linear(layers.Layer):
+    r"""Applies a linear transformation to the incoming data:
+
+    .. math::
+        \mathbf{x}^{\prime} = \mathbf{x} \mathbf{W} + \mathbf{b}
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        weight_initializer (str, optional): The initializer for the weight
+            matrix (:obj:`"glorot"`, :obj:`"uniform"`, :obj:`"kaiming_uniform"`
+            or :obj:`None`). (default: :obj:`None`)
+        bias_initializer (str, optional): The initializer for the bias vector
+            (:obj:`"zeros"` or :obj:`None`). (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import Linear
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+
+        layer = Linear(in_channels=8, out_channels=16)
+        print(tuple(layer(x).shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        bias: bool = True,
+        weight_initializer: Optional[str] = None,
+        bias_initializer: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.use_bias = bias
+        self.weight_initializer = weight_initializer
+        self.bias_initializer = bias_initializer
+        self.weight = None
+        self.bias = None
+
+        if self.in_channels > 0:
+            self._build_weights(self.in_channels)
+
+    def _build_weights(self, in_features: int):
+        self.in_channels = in_features
+        w_init = _get_weight_initializer(self.weight_initializer, in_features)
+        b_init = _get_bias_initializer(self.bias_initializer, in_features)
+
+        self.weight = self.add_weight(
+            shape=(in_features, self.out_channels),
+            initializer=w_init,
+            trainable=True,
+            name="weight",
+        )
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.out_channels,),
+                initializer=b_init,
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape):
+        if not hasattr(self, 'weight') or self.weight is None:
+            self._build_weights(input_shape[-1])
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        if hasattr(self, 'weight') and self.weight is not None:
+            w_init = _get_weight_initializer(self.weight_initializer, self.in_channels)
+            self.weight.assign(w_init(self.weight.shape, dtype=self.weight.dtype))
+        if hasattr(self, 'bias') and self.bias is not None:
+            b_init = _get_bias_initializer(self.bias_initializer, self.in_channels)
+            self.bias.assign(b_init(self.bias.shape, dtype=self.bias.dtype))
+
+    def call(self, x):
+        if not hasattr(self, 'weight') or self.weight is None:
+            self._build_weights(ops.shape(x)[-1])
+
+        out = ops.matmul(x, self.weight)
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def compute_output_shape(self, input_shape):
+        return (*input_shape[:-1], self.out_channels)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, bias={self.use_bias})')
+
+
+class HeteroLinear(layers.Layer):
+    r"""Applies separate linear transformations to the incoming data according
+    to types.
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \mathbf{x}_i \mathbf{\Theta}_{\kappa_i} +
+        \mathbf{b}_{\kappa_i}
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        num_types (int): The number of types.
+        is_sorted (bool, optional): If set to :obj:`True`, assumes that
+            :obj:`type_vec` is sorted. (default: :obj:`False`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        weight_initializer (str, optional): The initializer for the weight
+            matrix (:obj:`"glorot"`, :obj:`"uniform"`, :obj:`"kaiming_uniform"`
+            or :obj:`None`). (default: :obj:`None`)
+        bias_initializer (str, optional): The initializer for the bias vector
+            (:obj:`"zeros"` or :obj:`None`). (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HeteroLinear
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        node_type = np.random.randint(0, 3, size=(10,))
+
+        layer = HeteroLinear(in_channels=8, out_channels=16, num_types=3)  # separate weights per type
+        out = layer(x, node_type)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_types: int,
+        is_sorted: bool = False,
+        bias: bool = True,
+        weight_initializer: Optional[str] = None,
+        bias_initializer: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_types = num_types
+        self.is_sorted = is_sorted
+        self.use_bias = bias
+        self.weight_initializer = weight_initializer
+        self.bias_initializer = bias_initializer
+        self.weight = None
+        self.bias = None
+
+        if self.in_channels > 0:
+            self._build_weights(self.in_channels)
+
+    def _build_weights(self, in_features: int):
+        self.in_channels = in_features
+        w_init = _get_weight_initializer(self.weight_initializer, in_features)
+        b_init = _get_bias_initializer(self.bias_initializer, in_features)
+
+        self.weight = self.add_weight(
+            shape=(self.num_types, in_features, self.out_channels),
+            initializer=w_init,
+            trainable=True,
+            name="weight",
+        )
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.num_types, self.out_channels),
+                initializer=b_init,
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape):
+        if self.in_channels <= 0:
+            self._build_weights(input_shape[-1])
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        if self.in_channels > 0 and self.weight is not None:
+            w_init = _get_weight_initializer(self.weight_initializer, self.in_channels)
+            self.weight.assign(w_init(self.weight.shape, dtype=self.weight.dtype))
+            if self.bias is not None:
+                b_init = _get_bias_initializer(self.bias_initializer, self.in_channels)
+                self.bias.assign(b_init(self.bias.shape, dtype=self.bias.dtype))
+
+    def call(self, x, type_vec):
+        if self.in_channels <= 0 and not self.built:
+            self.build(ops.shape(x))
+
+        num_nodes = ops.shape(x)[0]
+        type_vec = ops.cast(type_vec, dtype="int32")
+
+        w_selected = ops.take(self.weight, type_vec, axis=0)
+        x_exp = ops.expand_dims(x, axis=1)
+        out = ops.squeeze(ops.matmul(x_exp, w_selected), axis=1)
+
+        if self.bias is not None:
+            b_selected = ops.take(self.bias, type_vec, axis=0)
+            out = out + b_selected
+
+        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}, num_types={self.num_types}, '
+                f'bias={self.use_bias})')
+
+
+class HeteroDictLinear(layers.Layer):
+    r"""Applies separate linear transformations to the incoming data
+    dictionary.
+
+    For key :math:`\kappa`, it computes
+
+    .. math::
+        \mathbf{x}^{\prime}_{\kappa} = \mathbf{x}_{\kappa}
+        \mathbf{W}_{\kappa} + \mathbf{b}_{\kappa}.
+
+    Args:
+        in_channels (int or Dict[Any, int]): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        types (List[Any], optional): The keys of the input dictionary.
+            (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HeteroDictLinear
+
+        x_dict = {"author": np.random.rand(3, 8).astype("float32"), "paper": np.random.rand(4, 12).astype("float32")}
+        layer = HeteroDictLinear(in_channels={"author": 8, "paper": 12}, out_channels=16)
+        out_dict = layer(x_dict)
+        print(tuple(out_dict["author"].shape), tuple(out_dict["paper"].shape))  # (3, 16) (4, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: Union[int, Dict[Any, int]],
+        out_channels: int,
+        types: Optional[List[Any]] = None,
+        **kwargs
+    ):
+        layer_kwargs = {k: v for k, v in kwargs.items() if k in ["name", "trainable", "dtype", "autocast"]}
+        super().__init__(**layer_kwargs)
+
+        if isinstance(in_channels, dict):
+            self.types = list(in_channels.keys())
+            if types is not None and set(self.types) != set(types):
+                raise ValueError("The provided 'types' do not match with the "
+                                 "keys in the 'in_channels' dictionary")
+        else:
+            if types is None:
+                raise ValueError("Please provide a list of 'types' if passing "
+                                 "'in_channels' as an integer")
+            self.types = types
+            in_channels = {node_type: in_channels for node_type in types}
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kwargs = kwargs
+
+        lin_kwargs = {k: v for k, v in kwargs.items() if k not in ["name"]}
+        self.lins = {
+            str(key): Linear(channels, self.out_channels, name=f"lin_{key}", **lin_kwargs)
+            for key, channels in self.in_channels.items()
+        }
+
+    def reset_parameters(self):
+        for lin in self.lins.values():
+            lin.reset_parameters()
+
+    def call(self, x_dict: Dict[str, Any]) -> Dict[str, Any]:
+        out_dict = {}
+        for key, x in x_dict.items():
+            str_key = str(key)
+            if str_key in self.lins:
+                out_dict[key] = self.lins[str_key](x)
+        return out_dict
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, bias={self.kwargs.get("bias", True)})')

k3_node/layers/dense/mincut_pool.py

@@ -0,0 +1,92 @@
+from typing import Optional, Tuple
+from keras import ops
+
+
+def dense_mincut_pool(
+    x,
+    adj,
+    s,
+    mask: Optional[any] = None,
+    temp: float = 1.0,
+) -> Tuple[any, any, any, any]:
+    r"""The MinCut pooling operator from the `"Spectral Clustering in Graph
+    Neural Networks for Graph Pooling" <https://arxiv.org/abs/1907.00481>`_
+    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)
+        temp: Temperature parameter for softmax function. (default: 1.0)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import dense_mincut_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, mincut_loss, ortho_loss = dense_mincut_pool(x, adj, s)
+        print(tuple(x_pool.shape), tuple(adj_pool.shape))  # (2, 3, 8) (2, 3, 3)
+        ```
+    """
+    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]
+    k = ops.shape(s)[-1]
+
+    if temp != 1.0:
+        s = s / temp
+    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)
+
+    # MinCut regularization
+    mincut_num = ops.sum(ops.diagonal(out_adj, axis1=1, axis2=2), axis=-1)
+    d_flat = ops.sum(adj, axis=-1)
+    d = ops.expand_dims(d_flat, axis=-1) * ops.expand_dims(ops.eye(num_nodes, dtype=d_flat.dtype), axis=0)
+    s_d_s = ops.matmul(ops.matmul(s_t, d), s)
+    mincut_den = ops.sum(ops.diagonal(s_d_s, axis1=1, axis2=2), axis=-1)
+    mincut_loss = ops.mean(-(mincut_num / mincut_den))
+
+    # Orthogonality regularization
+    ss = ops.matmul(s_t, s)
+    norm_ss = ops.sqrt(ops.sum(ops.power(ss, 2), axis=(-1, -2), keepdims=True))
+    i_s = ops.eye(k, dtype=ss.dtype)
+    norm_is = ops.sqrt(ops.cast(k, 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))))
+
+    # Fix and normalize coarsened adjacency matrix
+    eye_k = ops.expand_dims(ops.eye(k, dtype=out_adj.dtype), axis=0)
+    out_adj = out_adj * (1.0 - eye_k)
+    d_out = ops.sum(out_adj, axis=-1, keepdims=True)
+    d_sqrt = ops.sqrt(d_out) + 1e-15
+    out_adj = (out_adj / d_sqrt) / ops.transpose(d_sqrt, (0, 2, 1))
+
+    return out, out_adj, mincut_loss, ortho_loss
+