k3-node 1.0.0__py3-none-any.whl

k3_node/layers/aggr/basic.py

@@ -0,0 +1,412 @@
+from typing import Optional
+from keras import initializers, ops
+
+from .base import Aggregation
+from k3_node.ops.segment import segment_max, segment_sum
+from k3_node.ops.creation import full
+
+
+class SumAggregation(Aggregation):
+    r"""An aggregation operator that sums up features across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SumAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = SumAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        return self.reduce(x, index, ptr, dim_size, dim, reduce="sum")
+
+
+class MeanAggregation(Aggregation):
+    r"""An aggregation operator that averages features across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MeanAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = MeanAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        return self.reduce(x, index, ptr, dim_size, dim, reduce="mean")
+
+
+class MaxAggregation(Aggregation):
+    r"""An aggregation operator that takes the feature-wise maximum across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MaxAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = MaxAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        return self.reduce(x, index, ptr, dim_size, dim, reduce="max")
+
+
+class MinAggregation(Aggregation):
+    r"""An aggregation operator that takes the feature-wise minimum across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MinAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = MinAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        return self.reduce(x, index, ptr, dim_size, dim, reduce="min")
+
+
+class MulAggregation(Aggregation):
+    r"""An aggregation operator that multiplies features across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MulAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = MulAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        self.assert_index_present(index)
+        return self.reduce(x, index, ptr, dim_size, dim, reduce="mul")
+
+
+class VarAggregation(Aggregation):
+    r"""An aggregation operator that takes the feature-wise variance across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import VarAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = VarAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(self, semi_grad: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        self.semi_grad = semi_grad
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        mean = self.reduce(x, index, ptr, dim_size, dim, reduce="mean")
+        x_sq = ops.power(x, 2)
+        if self.semi_grad:
+            x_sq = ops.stop_gradient(x_sq)
+        mean2 = self.reduce(x_sq, index, ptr, dim_size, dim, reduce="mean")
+        return mean2 - ops.power(mean, 2)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(semi_grad={self.semi_grad})"
+
+
+class StdAggregation(Aggregation):
+    r"""An aggregation operator that takes the feature-wise standard deviation across a set of elements.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import StdAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = StdAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(self, semi_grad: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        self.semi_grad = semi_grad
+        self.var_aggr = VarAggregation(semi_grad=semi_grad)
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        var = self.var_aggr(x, index, ptr, dim_size, dim)
+        out = ops.sqrt(ops.maximum(var, 1e-5))
+        out = ops.where(out <= (1e-5**0.5), 0.0, out)
+        return out
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(semi_grad={self.semi_grad})"
+
+
+class SoftmaxAggregation(Aggregation):
+    r"""The softmax aggregation operator based on a temperature term.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SoftmaxAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = SoftmaxAggregation(learn=True)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        t: float = 1.0,
+        learn: bool = False,
+        semi_grad: bool = False,
+        channels: int = 1,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if learn and semi_grad:
+            raise ValueError(
+                f"Cannot enable 'semi_grad' in '{self.__class__.__name__}' in "
+                f"case the temperature term 't' is learnable"
+            )
+
+        if not learn and channels != 1:
+            raise ValueError(
+                f"Cannot set 'channels' greater than '1' in case "
+                f"'{self.__class__.__name__}' is not trainable"
+            )
+
+        self._init_t = t
+        self.learn = learn
+        self.semi_grad = semi_grad
+        self.channels = channels
+
+        if learn:
+            self.t = self.add_weight(
+                shape=(channels,),
+                initializer=initializers.Constant(t),
+                trainable=True,
+                name="t",
+            )
+        else:
+            self.t = t
+
+    def reset_parameters(self):
+        if self.learn:
+            self.t.assign(full((self.channels,), self._init_t, dtype=self.t.dtype))
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        t = self.t
+        if self.channels != 1:
+            self.assert_two_dimensional_input(x, dim)
+            t = ops.reshape(t, (1, self.channels))
+
+        alpha = x
+        if self.learn or t != 1.0:
+            alpha = x * t
+
+        if not self.learn and self.semi_grad:
+            alpha = ops.stop_gradient(alpha)
+
+        # Graph-wise softmax over segments
+        index = ops.cast(index, dtype="int32")
+        if dim_size is None:  # a tensor while tracing; don't test its truth value
+            dim_size = int(ops.max(index)) + 1 if ops.shape(index)[0] > 0 else 0
+
+        max_val = segment_max(alpha, index, num_segments=dim_size)
+        max_exp = ops.take(max_val, index, axis=0)
+        exp_alpha = ops.exp(alpha - max_exp)
+        sum_exp = segment_sum(exp_alpha, index, num_segments=dim_size)
+        sum_exp_taken = ops.take(sum_exp, index, axis=0)
+        alpha_sm = exp_alpha / ops.maximum(sum_exp_taken, 1e-12)
+
+        return self.reduce(x * alpha_sm, index, ptr, dim_size, dim, reduce="sum")
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(learn={self.learn})"
+
+
+class PowerMeanAggregation(Aggregation):
+    r"""The powermean aggregation operator based on a power term.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PowerMeanAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = PowerMeanAggregation(learn=True)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        p: float = 1.0,
+        learn: bool = False,
+        channels: int = 1,
+        clamp_min: Optional[float] = 1e-4,
+        clamp_max: Optional[float] = 100.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if not learn and channels != 1:
+            raise ValueError(
+                f"Cannot set 'channels' greater than '1' in case "
+                f"'{self.__class__.__name__}' is not trainable"
+            )
+
+        self._init_p = p
+        self.learn = learn
+        self.channels = channels
+        self.min_value = clamp_min if clamp_min is not None else 1e-4
+        self.max_value = clamp_max if clamp_max is not None else 100.0
+
+        if learn:
+            self.p = self.add_weight(
+                shape=(channels,),
+                initializer=initializers.Constant(p),
+                trainable=True,
+                name="p",
+            )
+        else:
+            self.p = p
+
+    def reset_parameters(self):
+        if self.learn:
+            self.p.assign(full((self.channels,), self._init_p, dtype=self.p.dtype))
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        p = self.p
+        if self.channels != 1:
+            self.assert_two_dimensional_input(x, dim)
+            p = ops.reshape(p, (-1, self.channels))
+
+        if self.learn or p != 1.0:
+            x = ops.power(ops.clip(x, self.min_value, self.max_value), p)
+
+        out = self.reduce(x, index, ptr, dim_size, dim, reduce="mean")
+
+        if self.learn or p != 1.0:
+            out = ops.power(ops.clip(out, self.min_value, self.max_value), 1.0 / p)
+
+        return out
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(learn={self.learn})"

k3_node/layers/aggr/deep_sets.py

@@ -0,0 +1,65 @@
+from typing import Optional
+from .base import Aggregation
+
+
+class DeepSetsAggregation(Aggregation):
+    r"""Performs Deep Sets aggregation in which the elements to aggregate are
+    first transformed by a Multi-Layer Perceptron (MLP)
+    :math:`\phi_{\mathbf{\Theta}}`, summed, and then transformed by another MLP
+    :math:`\rho_{\mathbf{\Theta}}`.
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import DeepSetsAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = DeepSetsAggregation(local_nn=keras.layers.Dense(16), global_nn=keras.layers.Dense(16))
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        local_nn: Optional[any] = None,
+        global_nn: Optional[any] = None,
+        local_mlp: Optional[any] = None,
+        global_mlp: Optional[any] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.local_nn = local_nn if local_nn is not None else local_mlp
+        self.global_nn = global_nn if global_nn is not None else global_mlp
+
+    def reset_parameters(self):
+        if self.local_nn is not None and hasattr(self.local_nn, "reset_parameters"):
+            self.local_nn.reset_parameters()
+        if self.global_nn is not None and hasattr(self.global_nn, "reset_parameters"):
+            self.global_nn.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        if self.local_nn is not None:
+            x = self.local_nn(x)
+
+        x = self.reduce(x, index=index, ptr=ptr, dim_size=dim_size, dim=dim, reduce="sum")
+
+        if self.global_nn is not None:
+            x = self.global_nn(x)
+
+        return x
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(local_nn={self.local_nn}, global_nn={self.global_nn})"
+

k3_node/layers/aggr/deepsets.py

@@ -0,0 +1,29 @@
+from keras import ops
+
+from k3_node.layers.aggr import Aggregation
+
+
+
+class DeepSetsAggregation(Aggregation):
+    def __init__(
+        self,
+        local_mlp=None,
+        global_mlp=None,
+    ):
+        super().__init__()
+
+        self.local_mlp = local_mlp
+        self.global_mlp = global_mlp
+
+
+    def call(self, x, index=None,  axis=-2):
+
+        if self.local_mlp is not None:
+            x = self.local_mlp(x) 
+
+        x = self.reduce(x, index=index, axis=axis, reduce_fn=ops.segment_sum)
+
+        if self.global_mlp is not None:
+            x = self.global_mlp(x)  # Assuming batch handling within MLP
+
+        return x

k3_node/layers/aggr/equilibrium.py

@@ -0,0 +1,107 @@
+from typing import List, Optional
+from keras import layers, ops
+
+from .base import Aggregation
+
+
+class ResNetPotential(layers.Layer):
+    def __init__(self, in_channels: int, out_channels: int, num_layers: List[int], **kwargs):
+        super().__init__(**kwargs)
+        sizes = [in_channels] + num_layers + [out_channels]
+        self.layers_list = [
+            layers.Dense(out_size, activation="tanh")
+            for in_size, out_size in zip(sizes[:-2], sizes[1:-1])
+        ]
+        self.final_layer = layers.Dense(sizes[-1])
+        self.res_trans = [
+            layers.Dense(layer_size)
+            for layer_size in num_layers + [out_channels]
+        ]
+
+    def reset_parameters(self):
+        for l in self.layers_list:
+            l.reset_parameters()
+        self.final_layer.reset_parameters()
+        for r in self.res_trans:
+            r.reset_parameters()
+
+    def call(self, inp):
+        h = inp
+        for layer, res in zip(self.layers_list + [self.final_layer], self.res_trans):
+            h_next = layer(h)
+            h = res(inp) + h_next
+        return h
+
+
+class EquilibriumAggregation(Aggregation):
+    r"""The equilibrium aggregation layer from the `"Equilibrium Aggregation:
+    Encoding Sets via Optimization" <https://arxiv.org/abs/2202.12795>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import EquilibriumAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = EquilibriumAggregation(in_channels=8, out_channels=16, num_layers=[8])
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: List[int],
+        grad_iter: int = 5,
+        lamb: float = 0.1,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.grad_iter = grad_iter
+        self.initial_lamb = lamb
+
+        self.potential = ResNetPotential(in_channels + out_channels, out_channels, num_layers)
+        self.proj = layers.Dense(out_channels)
+
+    def reset_parameters(self):
+        self.potential.reset_parameters()
+        self.proj.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        self.assert_index_present(index)
+        index = ops.cast(index, dtype="int32")
+        if dim_size is None:
+            dim_size = ops.max(index) + 1
+
+        # Initial mean aggregation as starting state
+        x_mean = self.reduce(x, index, ptr, dim_size, dim, reduce="mean")
+        y = self.proj(x_mean)
+
+        # Unrolled iterative equilibrium updates
+        for _ in range(self.grad_iter):
+            y_expanded = ops.take(y, index, axis=0)
+            inp = ops.concatenate([x, y_expanded], axis=-1)
+            pot = self.potential(inp)
+            pot_mean = self.reduce(pot, index, ptr, dim_size, dim, reduce="mean")
+            y = y + 0.1 * pot_mean - self.initial_lamb * y
+
+        return y
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+

k3_node/layers/aggr/fused.py

@@ -0,0 +1,43 @@
+from typing import List, Union
+from .base import Aggregation
+
+
+class FusedAggregation(Aggregation):
+    r"""Helper class to fuse computation of multiple aggregations together.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import FusedAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = FusedAggregation(aggrs=["sum", "mean", "max"])
+        outs = aggr(x, index=index, dim_size=2)  # one result per aggregation, computed together
+        print(len(outs), tuple(outs[0].shape))  # 3 (2, 8)
+        ```
+    """
+
+    def __init__(self, aggrs: List[Union[Aggregation, str]], **kwargs):
+        super().__init__(**kwargs)
+        from .resolver import aggregation_resolver
+
+        self.aggrs = [aggregation_resolver(aggr) for aggr in aggrs]
+
+    def reset_parameters(self):
+        for aggr in self.aggrs:
+            if hasattr(aggr, "reset_parameters"):
+                aggr.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index=None,
+        ptr=None,
+        dim_size=None,
+        dim=-2,
+        **kwargs,
+    ):
+        return [aggr(x, index=index, ptr=ptr, dim_size=dim_size, dim=dim, **kwargs) for aggr in self.aggrs]
+