k3-node 1.0.0__py3-none-any.whl

k3_node/layers/norm/test_norm.py

@@ -0,0 +1,275 @@
+import numpy as np
+import pytest
+from keras import ops, random
+
+from k3_node.layers.norm import (
+    BatchNorm,
+    DiffGroupNorm,
+    GraphNorm,
+    GraphSizeNorm,
+    HeteroBatchNorm,
+    HeteroLayerNorm,
+    InstanceNorm,
+    LayerNorm,
+    MeanSubtractionNorm,
+    MessageNorm,
+    PairNorm,
+)
+
+
+def test_graph_size_norm():
+    x = random.normal((100, 16))
+    batch = ops.repeat(ops.arange(10, dtype="int32"), 10)
+
+    norm = GraphSizeNorm()
+
+    out1 = norm(x)
+    assert ops.shape(out1) == (100, 16)
+
+    out2 = norm(x, batch)
+    assert ops.shape(out2) == (100, 16)
+
+    # Tuple input
+    out3 = norm((x, batch))
+    assert np.allclose(ops.convert_to_numpy(out2), ops.convert_to_numpy(out3), atol=1e-6)
+
+
+@pytest.mark.parametrize("scale_individually", [False, True])
+def test_pair_norm(scale_individually):
+    x = random.normal((100, 16))
+    batch = ops.zeros((100,), dtype="int32")
+
+    norm = PairNorm(scale_individually=scale_individually)
+
+    out1 = norm(x)
+    assert ops.shape(out1) == (100, 16)
+
+    x2 = ops.concatenate([x, x], axis=0)
+    batch2 = ops.concatenate([batch, batch + 1], axis=0)
+    out2 = norm(x2, batch2)
+
+    assert np.allclose(
+        ops.convert_to_numpy(out1),
+        ops.convert_to_numpy(out2[:100]),
+        atol=1e-5,
+    )
+    assert np.allclose(
+        ops.convert_to_numpy(out1),
+        ops.convert_to_numpy(out2[100:]),
+        atol=1e-5,
+    )
+
+
+def test_mean_subtraction_norm():
+    x = random.normal((6, 16))
+    batch = ops.convert_to_tensor(np.array([0, 0, 1, 1, 1, 2], dtype=np.int32))
+
+    norm = MeanSubtractionNorm()
+
+    out1 = norm(x)
+    assert ops.shape(out1) == (6, 16)
+    assert np.allclose(ops.convert_to_numpy(ops.mean(out1)), 0.0, atol=1e-6)
+
+    out2 = norm(x, batch)
+    assert ops.shape(out2) == (6, 16)
+    assert np.allclose(ops.convert_to_numpy(ops.mean(out2[0:2], axis=0)), 0.0, atol=1e-6)
+    assert np.allclose(ops.convert_to_numpy(ops.mean(out2[2:5], axis=0)), 0.0, atol=1e-6)
+    assert np.allclose(ops.convert_to_numpy(ops.mean(out2[5:6], axis=0)), 0.0, atol=1e-6)
+
+
+@pytest.mark.parametrize("learn_scale", [False, True])
+def test_message_norm(learn_scale):
+    norm = MessageNorm(learn_scale=learn_scale)
+    x = random.normal((100, 16))
+    msg = random.normal((100, 16))
+
+    out = norm(x, msg)
+    assert ops.shape(out) == (100, 16)
+
+    # Tuple calling
+    out_tuple = norm((x, msg))
+    assert np.allclose(ops.convert_to_numpy(out), ops.convert_to_numpy(out_tuple), atol=1e-6)
+
+    norm.reset_parameters()
+    assert np.allclose(ops.convert_to_numpy(norm.scale), 1.0)
+
+
+def test_graph_norm():
+    x = random.normal((200, 16))
+    batch = ops.repeat(ops.arange(4, dtype="int32"), 50)
+
+    norm = GraphNorm(16)
+
+    out = norm(x)
+    assert ops.shape(out) == (200, 16)
+    assert np.allclose(ops.convert_to_numpy(ops.mean(out, axis=0)), 0.0, atol=1e-5)
+    assert np.allclose(ops.convert_to_numpy(ops.std(out, axis=0)), 1.0, atol=1e-3)
+
+    out_b = norm(x, batch)
+    assert ops.shape(out_b) == (200, 16)
+    assert np.allclose(ops.convert_to_numpy(ops.mean(out_b[:50], axis=0)), 0.0, atol=1e-5)
+    assert np.allclose(ops.convert_to_numpy(ops.std(out_b[:50], axis=0)), 1.0, atol=1e-3)
+
+    norm.reset_parameters()
+    assert np.allclose(ops.convert_to_numpy(norm.weight), 1.0)
+    assert np.allclose(ops.convert_to_numpy(norm.bias), 0.0)
+    assert np.allclose(ops.convert_to_numpy(norm.mean_scale), 1.0)
+
+
+@pytest.mark.parametrize("conf", [True, False])
+def test_instance_norm(conf):
+    batch = ops.zeros((100,), dtype="int32")
+    x1 = random.normal((100, 16))
+    x2 = random.normal((100, 16))
+
+    norm1 = InstanceNorm(16, affine=conf, track_running_stats=conf)
+    norm2 = InstanceNorm(16, affine=conf, track_running_stats=conf)
+
+    out1 = norm1(x1)
+    out2 = norm2(x1, batch)
+    assert ops.shape(out1) == (100, 16)
+    assert np.allclose(ops.convert_to_numpy(out1), ops.convert_to_numpy(out2), atol=1e-5)
+
+    if conf:
+        assert np.allclose(
+            ops.convert_to_numpy(norm1.running_mean),
+            ops.convert_to_numpy(norm2.running_mean),
+            atol=1e-5,
+        )
+        assert np.allclose(
+            ops.convert_to_numpy(norm1.running_var),
+            ops.convert_to_numpy(norm2.running_var),
+            atol=1e-5,
+        )
+
+    out1_eval = norm1(x1, training=False)
+    out2_eval = norm2(x1, batch, training=False)
+    assert np.allclose(ops.convert_to_numpy(out1_eval), ops.convert_to_numpy(out2_eval), atol=1e-5)
+
+    norm1.reset_parameters()
+    norm1.reset_running_stats()
+
+
+@pytest.mark.parametrize("affine", [True, False])
+@pytest.mark.parametrize("mode", ["graph", "node"])
+def test_layer_norm(affine, mode):
+    x = random.normal((100, 16))
+    batch = ops.zeros((100,), dtype="int32")
+
+    norm = LayerNorm(16, affine=affine, mode=mode)
+
+    out1 = norm(x)
+    assert ops.shape(out1) == (100, 16)
+
+    if mode == "graph":
+        out1_b = norm(x, batch)
+        assert np.allclose(ops.convert_to_numpy(out1), ops.convert_to_numpy(out1_b), atol=1e-5)
+
+        x2 = ops.concatenate([x, x], axis=0)
+        batch2 = ops.concatenate([batch, batch + 1], axis=0)
+        out2 = norm(x2, batch2)
+        assert np.allclose(ops.convert_to_numpy(out1), ops.convert_to_numpy(out2[:100]), atol=1e-5)
+        assert np.allclose(ops.convert_to_numpy(out1), ops.convert_to_numpy(out2[100:]), atol=1e-5)
+    else:
+        mean = ops.mean(out1, axis=-1)
+        std = ops.std(out1, axis=-1)
+        assert np.allclose(ops.convert_to_numpy(mean), 0.0, atol=1e-4)
+        if not affine:
+            assert np.allclose(ops.convert_to_numpy(std), 1.0, atol=1e-3)
+
+
+@pytest.mark.parametrize("affine", [False, True])
+def test_hetero_layer_norm(affine):
+    x = random.normal((100, 16))
+    expected = LayerNorm(16, affine=affine, mode="node")(x)
+
+    type_vec = ops.zeros((100,), dtype="int32")
+    type_ptr = [0, 100]
+
+    norm = HeteroLayerNorm(16, num_types=1, affine=affine)
+
+    out = norm(x, type_vec)
+    assert ops.shape(out) == (100, 16)
+    assert np.allclose(ops.convert_to_numpy(out), ops.convert_to_numpy(expected), atol=1e-3)
+
+    out_ptr = norm(x, type_ptr=type_ptr)
+    assert np.allclose(ops.convert_to_numpy(out_ptr), ops.convert_to_numpy(expected), atol=1e-3)
+
+
+@pytest.mark.parametrize("conf", [True, False])
+def test_batch_norm(conf):
+    x = random.normal((100, 16))
+    norm = BatchNorm(16, affine=conf, track_running_stats=conf)
+    norm.reset_running_stats()
+    norm.reset_parameters()
+
+    assert norm.in_channels == 16
+    assert norm.eps == 1e-5
+    assert norm.momentum == 0.1
+    assert norm.affine == conf
+    assert norm.track_running_stats == conf
+    assert (norm.weight is not None) == conf
+    assert (norm.bias is not None) == conf
+
+    out = norm(x)
+    assert ops.shape(out) == (100, 16)
+
+
+def test_batch_norm_single_element():
+    x = random.normal((1, 16))
+
+    norm = BatchNorm(16)
+    with pytest.raises(ValueError, match="Expected more than 1 value"):
+        norm(x, training=True)
+
+    with pytest.raises(ValueError, match="requires 'track_running_stats'"):
+        BatchNorm(16, track_running_stats=False, allow_single_element=True)
+
+    norm = BatchNorm(16, track_running_stats=True, allow_single_element=True)
+    out = norm(x, training=True)
+    assert np.allclose(ops.convert_to_numpy(out), ops.convert_to_numpy(x), atol=1e-5)
+
+
+@pytest.mark.parametrize("conf", [True, False])
+def test_hetero_batch_norm(conf):
+    x = random.normal((100, 16))
+
+    norm = BatchNorm(16, affine=conf, track_running_stats=conf)
+    expected = norm(x)
+
+    type_vec = ops.zeros((100,), dtype="int32")
+    h_norm = HeteroBatchNorm(16, num_types=1, affine=conf, track_running_stats=conf)
+    h_norm.reset_running_stats()
+    h_norm.reset_parameters()
+
+    out = h_norm(x, type_vec)
+    assert ops.shape(out) == (100, 16)
+    assert np.allclose(ops.convert_to_numpy(out), ops.convert_to_numpy(expected), atol=1e-4)
+
+
+def test_diff_group_norm():
+    x = random.normal((100, 16))
+    norm = DiffGroupNorm(16, groups=4, lamda=0.01)
+
+    out = norm(x)
+    assert ops.shape(out) == (100, 16)
+
+    norm.reset_parameters()
+
+
+def test_group_distance_ratio():
+    x = random.normal((6, 16))
+    y = ops.convert_to_tensor(np.array([0, 1, 0, 1, 1, 1], dtype=np.int64))
+
+    assert DiffGroupNorm.group_distance_ratio(x, y) > 0
+
+
+def test_batch_norm_defaults_to_inference_when_training_is_none():
+    # Keras semantics: without an explicit or propagated `training`, use running statistics
+    # (otherwise `predict` output would depend on the batch).
+    x = random.normal((8, 4)) * 3.0 + 5.0
+    norm = BatchNorm(4)
+    np.testing.assert_allclose(
+        ops.convert_to_numpy(norm(x)), ops.convert_to_numpy(norm(x, training=False)), rtol=1e-6, atol=1e-6
+    )
+    assert not np.allclose(ops.convert_to_numpy(norm(x)), ops.convert_to_numpy(norm(x, training=True)))

k3_node/layers/pool/__init__.py

@@ -0,0 +1,83 @@
+r"""Pooling package for graph neural networks."""
+
+from .select import Select, SelectOutput, SelectTopK, topk
+from .connect import Connect, ConnectOutput, FilterEdges, filter_adj
+
+from .glob import global_add_pool, global_max_pool, global_mean_pool
+from .consecutive import consecutive_cluster
+from .pool import pool_batch, pool_edge, pool_pos
+from .max_pool import max_pool, max_pool_neighbor_x, max_pool_x
+from .avg_pool import avg_pool, avg_pool_neighbor_x, avg_pool_x
+
+from .topk_pool import TopKPooling
+from .sag_pool import SAGPooling
+from .edge_pool import EdgePooling
+from .cluster_pool import ClusterPooling
+from .asap import ASAPooling, LEConv
+from .pan_pool import PANPooling
+from .mem_pool import MemPooling
+
+from .voxel_grid import voxel_grid
+from .graclus import graclus
+from .decimation import decimation_indices
+from .knn import (
+    KNNIndex,
+    L2KNNIndex,
+    MIPSKNNIndex,
+    ApproxL2KNNIndex,
+    ApproxMIPSKNNIndex,
+    knn,
+    knn_graph,
+)
+from .approx_knn import approx_knn, approx_knn_graph
+from .point_cloud import fps, nearest, radius, radius_graph
+
+__all__ = [
+    "global_add_pool",
+    "global_mean_pool",
+    "global_max_pool",
+    "KNNIndex",
+    "L2KNNIndex",
+    "MIPSKNNIndex",
+    "ApproxL2KNNIndex",
+    "ApproxMIPSKNNIndex",
+    "TopKPooling",
+    "SAGPooling",
+    "EdgePooling",
+    "ClusterPooling",
+    "ASAPooling",
+    "PANPooling",
+    "MemPooling",
+    "max_pool",
+    "avg_pool",
+    "max_pool_x",
+    "max_pool_neighbor_x",
+    "avg_pool_x",
+    "avg_pool_neighbor_x",
+    "graclus",
+    "voxel_grid",
+    "fps",
+    "knn",
+    "knn_graph",
+    "approx_knn",
+    "approx_knn_graph",
+    "radius",
+    "radius_graph",
+    "nearest",
+    # Additional primitives
+    "Select",
+    "SelectOutput",
+    "SelectTopK",
+    "topk",
+    "Connect",
+    "ConnectOutput",
+    "FilterEdges",
+    "filter_adj",
+    "consecutive_cluster",
+    "pool_batch",
+    "pool_edge",
+    "pool_pos",
+    "decimation_indices",
+]
+
+classes = __all__

k3_node/layers/pool/approx_knn.py

@@ -0,0 +1,101 @@
+from typing import Optional
+from keras import ops
+import numpy as np
+
+from .knn import knn
+from k3_node.ops.host import to_numpy
+
+
+def approx_knn(
+    x,
+    y,
+    k: int,
+    batch_x: Optional[any] = None,
+    batch_y: Optional[any] = None,
+):
+    r"""Finds for each element in `y` the `k` approximated nearest points in `x`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import approx_knn
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+        query = np.random.rand(4, 3).astype("float32")  # 4 query points
+
+        assign = approx_knn(pos, query, k=3)  # approximate (NN-descent) nearest neighbors
+        print(tuple(assign.shape))  # (2, 12)
+        ```
+    """
+    try:
+        from pynndescent import NNDescent
+
+        if batch_x is None:
+            batch_x = ops.zeros((ops.shape(x)[0],), dtype="int32")
+        if batch_y is None:
+            batch_y = ops.zeros((ops.shape(y)[0],), dtype="int32")
+
+        x_np = to_numpy(x)
+        y_np = to_numpy(y)
+        batch_x_np = to_numpy(batch_x)
+        batch_y_np = to_numpy(batch_y)
+
+        min_xy = min(np.min(x_np), np.min(y_np))
+        x_np = x_np - min_xy
+        y_np = y_np - min_xy
+
+        max_xy = max(np.max(x_np), np.max(y_np))
+        if max_xy > 0:
+            x_np = x_np / max_xy
+            y_np = y_np / max_xy
+
+        x_aug = np.concatenate([x_np, 2.0 * x_np.shape[1] * batch_x_np[:, None]], axis=-1)
+        y_aug = np.concatenate([y_np, 2.0 * y_np.shape[1] * batch_y_np[:, None]], axis=-1)
+
+        index = NNDescent(x_aug)
+        col, dist = index.query(y_aug, k=k)
+        row = np.repeat(np.arange(y_np.shape[0]), k)
+        col = col.reshape(-1)
+        dist = dist.reshape(-1)
+
+        valid = ~np.isinf(dist)
+        return ops.convert_to_tensor(np.stack([row[valid], col[valid]], axis=0), dtype="int64")
+    except ImportError:
+        return knn(x, y, k, batch_x=batch_x, batch_y=batch_y)
+
+
+def approx_knn_graph(
+    x,
+    k: int,
+    batch: Optional[any] = None,
+    loop: bool = False,
+    flow: str = "source_to_target",
+):
+    r"""Computes graph edges to the nearest approximated `k` points.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import approx_knn_graph
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+
+        edge_index = approx_knn_graph(pos, k=3)
+        print(edge_index.shape[0])  # 2
+        ```
+    """
+    assert flow in ["source_to_target", "target_to_source"]
+    edge_index = approx_knn(x, x, k if loop else k + 1, batch, batch)
+    edge_index_np = to_numpy(edge_index)
+
+    if flow == "source_to_target":
+        edge_index_np = np.flip(edge_index_np, axis=0)
+
+    if not loop:
+        mask = edge_index_np[0] != edge_index_np[1]
+        edge_index_np = edge_index_np[:, mask]
+
+    return ops.convert_to_tensor(edge_index_np, dtype=edge_index.dtype)
+

k3_node/layers/pool/asap.py

@@ -0,0 +1,173 @@
+from typing import Callable, Optional, Tuple, Union
+from keras import layers, ops
+import numpy as np
+
+from .connect.filter_edges import FilterEdges
+from .select.topk import SelectTopK
+from k3_node.ops.segment import segment_max, segment_sum
+
+
+class LEConv(layers.Layer):
+    r"""The local extremum graph neural network operator."""
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.lin1 = layers.Dense(out_channels, use_bias=bias, name="lin1")
+        self.lin2 = layers.Dense(out_channels, use_bias=False, name="lin2")
+        self.lin3 = layers.Dense(out_channels, use_bias=bias, name="lin3")
+
+    def reset_parameters(self):
+        pass
+
+    def call(self, x, edge_index, edge_weight: Optional[any] = None):
+        a = self.lin1(x)
+        b = self.lin2(x)
+        row = ops.cast(edge_index[0], dtype="int32")
+        col = ops.cast(edge_index[1], dtype="int32")
+        num_nodes = ops.shape(x)[0]
+
+        a_j = ops.take(a, row, axis=0)
+        b_i = ops.take(b, col, axis=0)
+        msg = a_j - b_i
+        if edge_weight is not None:
+            msg = msg * ops.reshape(edge_weight, (-1, 1))
+
+        out = segment_sum(msg, col, num_segments=num_nodes)
+        return out + self.lin3(x)
+
+
+class ASAPooling(layers.Layer):
+    r"""The Adaptive Structure Aware Pooling operator from the
+    `"ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical
+    Graph Representations" <https://arxiv.org/abs/1911.07979>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import ASAPooling
+
+        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
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        layer = ASAPooling(in_channels=8, ratio=0.5)  # keep half of the nodes of each graph
+        out = layer(x, edge_index, batch=batch)
+        x_pool, edge_index_pool, edge_attr_pool, batch_pool = out[0], out[1], out[2], out[3]
+        print(tuple(x_pool.shape))  # (6, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        ratio: Union[float, int] = 0.5,
+        GNN: Optional[Callable] = None,
+        dropout: float = 0.0,
+        negative_slope: float = 0.2,
+        add_self_loops: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.ratio = ratio
+        self.negative_slope = negative_slope
+        self.dropout_rate = dropout
+        self.drop = layers.Dropout(dropout) if dropout > 0.0 else None
+        self.add_self_loops = add_self_loops
+
+        self.lin = layers.Dense(in_channels, use_bias=True, name="lin")
+        self.att = layers.Dense(1, use_bias=True, name="att")
+        self.gnn_score = LEConv(in_channels, 1, name="gnn_score")
+
+        if GNN is not None:
+            self.gnn_intra_cluster = GNN(in_channels, in_channels, **kwargs)
+        else:
+            self.gnn_intra_cluster = None
+
+        self.select = SelectTopK(1, ratio)
+
+    def reset_parameters(self):
+        self.select.reset_parameters()
+
+    def call(
+        self,
+        x,
+        edge_index,
+        edge_weight: Optional[any] = None,
+        batch: Optional[any] = None,
+        training: bool = False,
+    ) -> Tuple[any, any, Optional[any], any, any]:
+        r"""Forward pass."""
+        N = ops.shape(x)[0]
+        if batch is None:
+            batch = ops.zeros((N,), dtype="int32")
+        else:
+            batch = ops.cast(batch, dtype="int32")
+
+        if len(ops.shape(x)) == 1:
+            x = ops.expand_dims(x, axis=-1)
+
+        # Add self-loops to edge_index and edge_weight
+        loop_idx = ops.arange(N, dtype=edge_index.dtype)
+        loop_edge = ops.stack([loop_idx, loop_idx], axis=0)
+        edge_index = ops.concatenate([edge_index, loop_edge], axis=1)
+
+        num_edges = ops.shape(edge_index)[1]
+        if edge_weight is None:
+            edge_weight = ops.ones((num_edges,), dtype=x.dtype)
+        else:
+            edge_weight = ops.concatenate([edge_weight, ops.ones((N,), dtype=x.dtype)], axis=0)
+
+        x_pool = x
+        if self.gnn_intra_cluster is not None:
+            x_pool = self.gnn_intra_cluster(x=x, edge_index=edge_index, edge_weight=edge_weight)
+
+        row = ops.cast(edge_index[0], dtype="int32")
+        col = ops.cast(edge_index[1], dtype="int32")
+
+        x_pool_j = ops.take(x_pool, row, axis=0)
+        x_q = segment_max(x_pool_j, col, num_segments=N)
+        x_q = ops.take(self.lin(x_q), col, axis=0)
+
+        score = ops.reshape(self.att(ops.concatenate([x_q, x_pool_j], axis=-1)), (-1,))
+        score = ops.leaky_relu(score, negative_slope=self.negative_slope)
+
+        # Softmax over col
+        score_max = segment_max(score, col, num_segments=N)
+        score_max_exp = ops.take(score_max, col, axis=0)
+        exp_score = ops.exp(score - score_max_exp)
+        sum_exp = segment_sum(exp_score, col, num_segments=N)
+        sum_exp_exp = ops.take(sum_exp, col, axis=0)
+        score = exp_score / (sum_exp_exp + 1e-12)
+
+        if self.drop is not None:
+            score = self.drop(score, training=training)
+
+        v_j = ops.take(x, row, axis=0) * ops.reshape(score, (-1, 1))
+        x_new = segment_sum(v_j, col, num_segments=N)
+
+        fitness = ops.reshape(ops.sigmoid(self.gnn_score(x_new, edge_index)), (-1,))
+        select_out = self.select(fitness, batch)
+        perm = select_out.node_index
+
+        x_out = ops.take(x_new, perm, axis=0) * ops.reshape(ops.take(fitness, perm, axis=0), (-1, 1))
+
+        connect = FilterEdges()
+        connect_out = connect(select_out, edge_index, edge_weight, batch)
+
+        edge_index_out = connect_out.edge_index
+        edge_weight_out = connect_out.edge_attr
+        batch_out = connect_out.batch
+
+        return x_out, edge_index_out, edge_weight_out, batch_out, perm
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels}, ratio={self.ratio})"
+