Track: Track1; Team name: Prasanna28Devadiga; Model: Sheaf Neural Networks with Connection Laplacians

2026_tdl_challenge/run_ablation_nsd_bundle.py

@@ -0,0 +1,87 @@
+"""Run the full GraphUniverse challenge grid for upstream NSD-bundle.
+
+Apples-to-apples ablation against Conn-NSD: same harness, same grid, same
+trainer budget, same hidden_dim — but with the *learned* bundle maps of
+Bodnar et al. (NSDEncoder, sheaf_type="bundle") instead of the
+deterministic Algorithm-1 maps of Conn-NSD.
+
+Hyperparameter overrides force NSD-bundle to use stalk dim ``d=4`` and
+``dropout=0.0`` to match the Conn-NSD config; everything else is the
+upstream default (``num_layers=2``, ``hidden_dim=64``, etc.).
+
+Output mirrors the conn_nsd_full layout::
+
+    2026_tdl_challenge/outputs/nsd_bundle_full/
+        results.json
+        heatmap_community_detection_accuracy.png
+        heatmap_triangle_mse_over_triangles.png
+        OOD/
+            OOD_{low,mid,high}_homophily__{community_detection,triangle_counting}.png
+
+Run with::
+
+    cd 2026_tdl_challenge
+    WANDB_MODE=offline python run_ablation_nsd_bundle.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+_HERE = Path(__file__).resolve().parent
+_REPO = _HERE.parent
+if str(_REPO) not in sys.path:
+    sys.path.insert(0, str(_REPO))
+if str(_HERE) not in sys.path:
+    sys.path.insert(0, str(_HERE))
+
+from utils import (  # noqa: E402  — sys.path setup must precede this import
+    resolve_project_root,
+    run_challenge_grid,
+    save_challenge_artifacts,
+)
+
+MODEL_CONFIG = "graph/nsd"
+# Override NSD-bundle to match the Conn-NSD config (stalk_dim=4, dropout=0).
+# Anything else left at upstream defaults — same num_layers, same hidden_dim.
+EXTRA_OVERRIDES = [
+    "model.backbone.d=4",
+    "model.backbone.dropout=0.0",
+    "model.backbone.sheaf_type=bundle",
+]
+OUTPUT_DIR_NAME = "nsd_bundle_full"
+
+
+def main() -> None:
+    """Run the 72-cell grid for NSD-bundle and save artefacts."""
+    project_root = resolve_project_root(_REPO)
+    print(f"Project root: {project_root}", flush=True)
+    print(f"Model: {MODEL_CONFIG}", flush=True)
+    print(f"Overrides: {EXTRA_OVERRIDES}", flush=True)
+
+    results, study_id = run_challenge_grid(
+        project_root=project_root,
+        model_config=MODEL_CONFIG,
+        extra_overrides=EXTRA_OVERRIDES,
+        quiet=True,
+    )
+    print(
+        f"\nGrid finished: {len(results)} runs; study_id={study_id}",
+        flush=True,
+    )
+
+    out_dir = _HERE / "outputs" / OUTPUT_DIR_NAME
+    out = save_challenge_artifacts(
+        results,
+        out_dir=out_dir,
+        model_config=MODEL_CONFIG,
+        study_id=study_id,
+    )
+    print("Artefacts:")
+    for k, v in out.items():
+        print(f"  {k}: {v}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

configs/model/graph/conn_nsd.yaml

@@ -0,0 +1,58 @@
+# Sheaf Neural Network with Connection Laplacians (Conn-NSD).
+#
+# Barbero et al., "Sheaf Neural Networks with Connection Laplacians,"
+# ICML 2022 TAG-ML Workshop, arXiv:2206.08702.
+#
+# The sheaf is constructed deterministically by Algorithm 1 of the paper:
+# local PCA over each node's 1-hop neighbourhood gives an orthonormal
+# tangent basis O_v; the restriction map F_{vu} = U V^T is the orthogonal
+# Procrustes solution from the SVD of O_v^T O_u. There are no learnable
+# sheaf parameters — only the diffusion weights W_1, W_2 (left/right) and
+# the residual gates epsilon.
+
+_target_: topobench.model.TBModel
+
+model_name: conn_nsd
+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.ConnNSDEncoder
+  input_dim: ${model.feature_encoder.out_channels}
+  hidden_dim: ${model.feature_encoder.out_channels}   # must be divisible by stalk_dim
+  num_layers: 2                                       # SWEEP candidate: [2, 4]
+  stalk_dim: 4                                        # SWEEP candidate: [2, 3, 4]
+  dropout: 0.0                                        # SWEEP candidate: [0.0, 0.25, 0.5]
+  input_dropout: 0.0
+  connection_features: raw                            # features as seen by the backbone, after TopoBench feature encoding
+
+backbone_wrapper:
+  _target_: topobench.nn.wrappers.${model.backbone_wrapper.wrapper_name}
+  _partial_: true
+  wrapper_name: GNNWrapper
+  out_channels: ${model.feature_encoder.out_channels}
+  residual_connections: false
+  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: MLPReadout
+  num_cell_dimensions: ${infer_num_cell_dimensions:${oc.select:model.feature_encoder.selected_dimensions,null},${model.feature_encoder.in_channels}}
+  in_channels: ${model.feature_encoder.out_channels}
+  hidden_layers: [16]
+  out_channels: ${dataset.parameters.num_classes}
+  task_level: ${define_task_level:${dataset.parameters.task_level},${dataset.split_params.learning_setting}}
+  pooling_type: sum
+  dropout: 0.2
+  act: "relu"
+  norm: null
+  final_act: null
+
+# compile model for faster training with pytorch 2.0
+compile: false