Track: Track1; Team name: Oversmooth operators; Model: r-ℓMPNN

configs/model/graph/loopy.yaml

@@ -0,0 +1,41 @@
+_target_: topobench.model.TBModel
+
+model_name: loopy
+model_domain: graph
+
+feature_encoder:
+  _target_: topobench.nn.encoders.${model.feature_encoder.encoder_name}
+  encoder_name: AllCellFeatureEncoder
+  in_channels: ${infer_in_channels:${dataset},${oc.select:transforms,null}}
+  out_channels: 64
+  proj_dropout: 0.0
+
+backbone:
+  _target_: topobench.nn.backbones.graph.Loopy
+  in_channels: ${model.feature_encoder.out_channels}
+  hidden_channels: ${model.feature_encoder.out_channels}
+  num_layers: 2
+  dropout: 0.0
+  r: 2 # Maximal neighbourhood order; must match the r_neighbourhood transform
+  nonlinearity: relu
+  norm: BatchNorm1d
+  shared: false # If true, one convolution is shared across the orders
+
+backbone_wrapper:
+  _target_: topobench.nn.wrappers.LoopyWrapper
+  _partial_: true
+  wrapper_name: LoopyWrapper
+  out_channels: ${model.feature_encoder.out_channels}
+  num_cell_dimensions: ${infer_num_cell_dimensions:${oc.select:model.feature_encoder.selected_dimensions,null},${model.feature_encoder.in_channels}}
+
+readout:
+  _target_: topobench.nn.readouts.${model.readout.readout_name}
+  readout_name: NoReadOut #  Use <NoReadOut> in case readout is not needed Options: PropagateSignalDown
+  num_cell_dimensions: ${infer_num_cell_dimensions:${oc.select:model.feature_encoder.selected_dimensions,null},${model.feature_encoder.in_channels}} # The highest order of cell dimensions to consider
+  hidden_dim: ${model.feature_encoder.out_channels}
+  out_channels: ${dataset.parameters.num_classes}
+  task_level: ${define_task_level:${dataset.parameters.task_level},${dataset.split_params.learning_setting}} # Handles the edge case of node-inductive task
+  pooling_type: sum
+
+# compile model for faster training with pytorch 2.0
+compile: false

configs/transforms/data_manipulations/r_neighbourhood.yaml

@@ -0,0 +1,3 @@
+transform_name: "RNeighbourhood"
+transform_type: "data manipulation"
+r: 2 # Maximal neighbourhood order; must match model.backbone.r

configs/transforms/model_defaults/loopy.yaml

@@ -0,0 +1,2 @@
+defaults:
+  - data_manipulations@RNeighbourhood: r_neighbourhood

test/conftest.py

@@ -308,13 +308,22 @@ def random_graph_input():
     """
     num_nodes = 8
     d_feat = 12
-    x = torch.randn(num_nodes, 12)
-    edges_1 = torch.randint(0, num_nodes, (2, num_nodes*2))
-    edges_2 = torch.randint(0, num_nodes, (2, num_nodes*2))
+    # Use a dedicated generator so the fixture does not depend on the
+    # ambient RNG state, i.e. on which tests ran before. Consumers assume
+    # the edges cover every node (a dense adjacency is built from them),
+    # which only holds for some global seeds.
+    generator = torch.Generator().manual_seed(0)
+    x = torch.randn(num_nodes, 12, generator=generator)
+    edges_1 = torch.randint(
+        0, num_nodes, (2, num_nodes * 2), generator=generator
+    )
+    edges_2 = torch.randint(
+        0, num_nodes, (2, num_nodes * 2), generator=generator
+    )
 
     d_feat_1, d_feat_2 = 5, 17
 
-    x_1 = torch.randn(num_nodes*2, d_feat_1)
-    x_2 = torch.randn(num_nodes*2, d_feat_2)
+    x_1 = torch.randn(num_nodes*2, d_feat_1, generator=generator)
+    x_2 = torch.randn(num_nodes*2, d_feat_2, generator=generator)
 
     return x, x_1, x_2, edges_1, edges_2

test/nn/backbones/graph/test_loopy.py

@@ -0,0 +1,197 @@
+"""Unit tests for the Loopy backbone."""
+
+import pytest
+import torch
+from torch_geometric.data import Data
+
+from topobench.dataloader.dataload_dataset import DataloadDataset
+from topobench.dataloader.utils import collate_fn
+from topobench.nn.backbones.graph.loopy import (
+    ACTIVATIONS,
+    CustomGINConv,
+    Loopy,
+    LoopyLayer,
+    MLP,
+    _path_propagate,
+    get_activation,
+)
+from topobench.nn.wrappers.graph.loopy_wrapper import LoopyWrapper
+from topobench.transforms.data_manipulations.r_neighbourhood import (
+    RNeighbourhood,
+)
+
+TRIANGLE_TAIL = ([[0, 1], [1, 2], [2, 0], [2, 3]], 4)
+SQUARE = ([[0, 1], [1, 2], [2, 3], [3, 0]], 4)
+
+
+def _assembled(graphs, hidden, r=2):
+    """Return ``(x, loopy_n, loopy_a, num_nodes)`` for direct layer tests."""
+    datas = []
+    for edges, n in graphs:
+        edge_index = torch.tensor(edges, dtype=torch.long).t()
+        edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+        data = Data(
+            x_0=torch.randn(n, hidden),
+            edge_index=edge_index,
+            y=torch.zeros(n, dtype=torch.long),
+            num_nodes=n,
+        )
+        data.batch_0 = torch.zeros(n, dtype=torch.long)
+        datas.append(RNeighbourhood(r=r, transform_name="RN")(data))
+    dataset = DataloadDataset(datas)
+    batch = collate_fn([dataset.get(i) for i in range(len(datas))])
+    loopy_n, loopy_a = LoopyWrapper._assemble_paths(batch)
+    return batch.x_0, loopy_n, loopy_a, batch.x_0.shape[0]
+
+
+class TestGetActivation:
+    """Test the activation resolver."""
+
+    @pytest.mark.parametrize("name", sorted(ACTIVATIONS))
+    def test_known(self, name):
+        act = get_activation(name)
+        assert isinstance(act, torch.nn.Module)
+        assert act(torch.randn(3, 2)).shape == (3, 2)
+
+    def test_invalid(self):
+        with pytest.raises(ValueError, match="Unsupported activation"):
+            get_activation("nope")
+
+
+class TestPathPropagate:
+    """Test the path-neighbour convolution."""
+
+    def test_documented_example(self):
+        x = torch.tensor([[1.0], [5.0]]).unsqueeze(1)  # (2, 1, 1)
+        out = _path_propagate(x)
+        assert torch.equal(out.squeeze(), torch.tensor([5.0, 1.0]))
+
+    def test_shape_preserved(self):
+        x = torch.randn(4, 6, 8)
+        assert _path_propagate(x).shape == x.shape
+
+    def test_middle_node_sums_both_neighbours(self):
+        x = torch.tensor([[1.0], [2.0], [4.0]]).unsqueeze(1)  # (3, 1, 1)
+        out = _path_propagate(x).squeeze()
+        assert out[1] == 5.0  # 1 + 4
+
+
+class TestMLP:
+    """Test the internal MLP."""
+
+    def test_forward_shape(self):
+        assert MLP(8, 5)(torch.randn(6, 8)).shape == (6, 5)
+
+    @pytest.mark.parametrize("num_layers", [2, 3])
+    def test_num_layers(self, num_layers):
+        mlp = MLP(8, 8, num_layers=num_layers)
+        assert len(mlp.lins) == num_layers
+
+    def test_batchnorm(self):
+        mlp = MLP(8, 8, norm="BatchNorm1d")
+        assert isinstance(mlp.norm, torch.nn.BatchNorm1d)
+        assert mlp(torch.randn(4, 8)).shape == (4, 8)
+
+    def test_reset_parameters(self):
+        MLP(8, 8, norm="BatchNorm1d").reset_parameters()
+
+
+class TestCustomGINConv:
+    """Test the path GIN convolution."""
+
+    def test_forward_shape(self):
+        conv = CustomGINConv(MLP(8, 8), in_channels=8, num_embeddings=4)
+        x = torch.randn(3, 5, 8)  # (path_length, num_paths, channels)
+        atomic = torch.randint(0, 4, (3, 5))
+        assert conv(x, atomic).shape == (5, 8)
+
+    def test_reset_parameters(self):
+        CustomGINConv(MLP(8, 8), 8, 4).reset_parameters()
+
+
+class TestLoopyLayer:
+    """Test a single loopy layer."""
+
+    def test_forward_shape(self):
+        x, ln, la, n = _assembled([TRIANGLE_TAIL], hidden=8)
+        layer = LoopyLayer(8, 8, r=2)
+        assert layer(x, ln, la, n).shape == (n, 8)
+
+    def test_shared_uses_single_conv(self):
+        assert len(LoopyLayer(8, 8, r=3, shared=True).convs) == 1
+        assert len(LoopyLayer(8, 8, r=3, shared=False).convs) == 3
+
+    def test_chunk_size_invariant(self):
+        x, ln, la, n = _assembled([SQUARE, TRIANGLE_TAIL], hidden=8)
+        torch.manual_seed(0)
+        big = LoopyLayer(8, 8, r=2, path_chunk_size=10**9).eval()
+        torch.manual_seed(0)
+        small = LoopyLayer(8, 8, r=2, path_chunk_size=1).eval()
+        out_big = big(x, ln, la, n)
+        out_small = small(x, ln, la, n)
+        assert torch.allclose(out_big, out_small, atol=1e-5)
+
+    def test_checkpoint_backward(self):
+        x, ln, la, n = _assembled([SQUARE], hidden=8)
+        x = x.clone().requires_grad_(True)
+        layer = LoopyLayer(8, 8, r=2, path_chunk_size=1).train()
+        layer(x, ln, la, n).sum().backward()
+        assert x.grad is not None
+
+
+class TestLoopy:
+    """Test the full backbone."""
+
+    def _run(self, graphs, hidden=8, **kw):
+        x, ln, la, n = _assembled(graphs, hidden=hidden, r=kw.get("r", 2))
+        model = Loopy(hidden, hidden, **kw)
+        return model, model(x, None, loopy_n=ln, loopy_a=la), n
+
+    def test_init_attributes(self):
+        model = Loopy(8, 8, num_layers=3, r=2)
+        assert model.out_channels == 8
+        assert model.r == 2
+        assert len(model.layers) == 3
+
+    def test_forward_shape(self):
+        _, out, n = self._run([TRIANGLE_TAIL])
+        assert out.shape == (n, 8)
+        assert torch.isfinite(out).all()
+
+    def test_backward_all_params(self):
+        x, ln, la, n = _assembled([SQUARE, TRIANGLE_TAIL], hidden=8)
+        model = Loopy(8, 8, num_layers=2, r=2).train()
+        model(x, None, loopy_n=ln, loopy_a=la).sum().backward()
+        for name, p in model.named_parameters():
+            assert p.grad is not None, f"no gradient for {name}"
+
+    @pytest.mark.parametrize("r", [1, 2, 3])
+    def test_forward_different_r(self, r):
+        _, out, n = self._run([SQUARE], r=r)
+        assert out.shape == (n, 8)
+
+    def test_dropout_and_kwargs_ignored(self):
+        model, out, n = self._run(
+            [TRIANGLE_TAIL], dropout=0.5, unused="x"
+        )
+        assert out.shape == (n, 8)
+
+    def test_eval_is_deterministic(self):
+        x, ln, la, n = _assembled([SQUARE], hidden=8)
+        model = Loopy(8, 8, num_layers=2, r=2, dropout=0.5).eval()
+        a = model(x, None, loopy_n=ln, loopy_a=la)
+        b = model(x, None, loopy_n=ln, loopy_a=la)
+        assert torch.allclose(a, b)
+
+    def test_isolated_nodes(self):
+        # Graph with an isolated node (no paths touch it).
+        x, ln, la, n = _assembled([([[0, 1], [1, 2], [2, 0]], 5)], hidden=8)
+        out = Loopy(8, 8, num_layers=1, r=2)(x, None, loopy_n=ln, loopy_a=la)
+        assert out.shape == (n, 8)
+        assert torch.isfinite(out).all()
+
+    def test_end_to_end_via_wrapper(self):
+        x, ln, la, n = _assembled([TRIANGLE_TAIL, SQUARE], hidden=8)
+        model = Loopy(8, 8, num_layers=2, r=2)
+        out = model(x, None, loopy_n=ln, loopy_a=la)
+        assert out.shape == (n, 8)

test/nn/wrappers/graph/test_loopy_wrapper.py

@@ -0,0 +1,116 @@
+"""Unit tests for the LoopyWrapper."""
+
+import torch
+from torch_geometric.data import Data
+
+from topobench.dataloader.dataload_dataset import DataloadDataset
+from topobench.dataloader.utils import collate_fn
+from topobench.nn.backbones.graph.loopy import Loopy
+from topobench.nn.wrappers.graph.loopy_wrapper import LoopyWrapper
+from topobench.transforms.data_manipulations.r_neighbourhood import (
+    RNeighbourhood,
+)
+
+TRIANGLE_TAIL = ([[0, 1], [1, 2], [2, 0], [2, 3]], 4)
+SQUARE_ISOLATED = ([[0, 1], [1, 2], [2, 3], [3, 0]], 5)  # node 4 isolated
+
+
+def _graph(edges, num_nodes, feat, r):
+    """Build a single transformed graph with a per-graph batch vector."""
+    edge_index = torch.tensor(edges, dtype=torch.long).t()
+    edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+    data = Data(
+        x_0=torch.randn(num_nodes, feat),
+        edge_index=edge_index,
+        y=torch.zeros(num_nodes, dtype=torch.long),
+        num_nodes=num_nodes,
+    )
+    data.batch_0 = torch.zeros(num_nodes, dtype=torch.long)
+    return RNeighbourhood(r=r, transform_name="RNeighbourhood")(data)
+
+
+def _batch(graphs, feat=8, r=2):
+    """Collate several graphs the way the TopoBench dataloader does."""
+    datas = [_graph(edges, n, feat, r) for edges, n in graphs]
+    dataset = DataloadDataset(datas)
+    return collate_fn([dataset.get(i) for i in range(len(datas))])
+
+
+class TestAssemblePaths:
+    """Test the graph-local to batch-global index reconstruction."""
+
+    def test_shapes_and_transpose(self):
+        batch = _batch([TRIANGLE_TAIL])
+        loopy_n, loopy_a = LoopyWrapper._assemble_paths(batch)
+        for order in range(3):
+            if loopy_n[order].numel():
+                assert loopy_n[order].shape[0] == order + 2
+                assert loopy_n[order].shape == loopy_a[order].shape
+
+    def test_single_graph_indices_unchanged(self):
+        batch = _batch([TRIANGLE_TAIL])
+        loopy_n, _ = LoopyWrapper._assemble_paths(batch)
+        # With a single graph the offset is zero, so indices stay in range.
+        assert loopy_n[1].max() < batch.x_0.shape[0]
+
+    def test_two_graphs_offset(self):
+        batch = _batch([TRIANGLE_TAIL, SQUARE_ISOLATED])
+        loopy_n, _ = LoopyWrapper._assemble_paths(batch)
+        # The square lives in the second graph -> its order-2 paths must
+        # reference the second graph's node block (indices 4..8).
+        counts = batch["loopyNcount2"]
+        graph_of_path = torch.repeat_interleave(torch.arange(2), counts)
+        square_paths = loopy_n[2].t()[graph_of_path == 1]
+        assert square_paths.numel() > 0
+        assert square_paths.min() >= 4
+
+    def test_indices_within_own_graph(self):
+        batch = _batch([TRIANGLE_TAIL, SQUARE_ISOLATED])
+        loopy_n, _ = LoopyWrapper._assemble_paths(batch)
+        node_counts = torch.bincount(batch.batch_0)
+        node_ptr = torch.cat([node_counts.new_zeros(1), node_counts.cumsum(0)])
+        for order in range(3):
+            if not loopy_n[order].numel():
+                continue
+            counts = batch[f"loopyNcount{order}"]
+            gid = torch.repeat_interleave(torch.arange(2), counts)
+            glob = loopy_n[order].t()
+            lo = node_ptr[gid].unsqueeze(1)
+            hi = node_ptr[gid + 1].unsqueeze(1)
+            assert torch.all((glob >= lo) & (glob < hi))
+
+    def test_empty_order_handled(self):
+        # A pure triangle has no order-2 (length-4) paths.
+        batch = _batch([([[0, 1], [1, 2], [2, 0]], 3)])
+        loopy_n, loopy_a = LoopyWrapper._assemble_paths(batch)
+        assert loopy_n[2].shape[1] == 0
+        assert loopy_a[2].shape[1] == 0
+
+
+class TestLoopyWrapperForward:
+    """Test the wrapper forward pass with a real backbone."""
+
+    def _wrapper(self, hidden=8):
+        backbone = Loopy(
+            in_channels=hidden, hidden_channels=hidden, num_layers=2, r=2
+        )
+        return LoopyWrapper(
+            backbone, out_channels=hidden, num_cell_dimensions=1
+        )
+
+    def test_output_keys(self):
+        batch = _batch([TRIANGLE_TAIL, SQUARE_ISOLATED])
+        out = self._wrapper()(batch)
+        assert set(out.keys()) >= {"labels", "batch_0", "x_0"}
+
+    def test_output_shape(self):
+        batch = _batch([TRIANGLE_TAIL, SQUARE_ISOLATED])
+        out = self._wrapper(hidden=8)(batch)
+        assert out["x_0"].shape == (batch.x_0.shape[0], 8)
+        assert torch.isfinite(out["x_0"]).all()
+
+    def test_labels_and_batch_preserved(self):
+        batch = _batch([TRIANGLE_TAIL])
+        out = self._wrapper()(batch)
+        assert torch.equal(out["batch_0"], batch.batch_0)
+        assert torch.equal(out["labels"], batch.y)