[Model] Add CoED-GNN model with multi-dataset node classification support

examples/coed/geom_planetoid.py

@@ -0,0 +1,53 @@
+"""Helpers for applying Geom-GCN 10-split evaluation to GammaGL Planetoid."""
+
+import os
+import os.path as osp
+
+import numpy as np
+import tensorlayerx as tlx
+
+from gammagl.data import download_url
+from gammagl.datasets import Planetoid
+
+
+GEOM_GCN_URL = "https://raw.githubusercontent.com/graphdml-uiuc-jlu/geom-gcn/master/splits"
+
+
+def _geom_raw_dir(root, name):
+    return osp.join(root, name.lower(), "geom-gcn", "raw")
+
+
+def _split_file(name, split_id):
+    return "{}_split_0.6_0.2_{}.npz".format(name.lower(), split_id)
+
+
+def ensure_geom_gcn_splits(root, name, num_splits=10):
+    """Ensure Geom-GCN split files exist under the GammaGL dataset directory."""
+    raw_dir = _geom_raw_dir(root, name)
+    os.makedirs(raw_dir, exist_ok=True)
+    for split_id in range(num_splits):
+        filename = _split_file(name, split_id)
+        path = osp.join(raw_dir, filename)
+        if not osp.exists(path):
+            download_url("{}/{}".format(GEOM_GCN_URL, filename), raw_dir)
+    return raw_dir
+
+
+def load_planetoid_with_geom_splits(root, name, num_splits=10):
+    """Load Planetoid data and replace masks with Geom-GCN fixed splits."""
+    dataset = Planetoid(root=root, name=name)
+    graph = dataset[0]
+    raw_dir = ensure_geom_gcn_splits(root, name, num_splits=num_splits)
+
+    train_masks, val_masks, test_masks = [], [], []
+    for split_id in range(num_splits):
+        split_path = osp.join(raw_dir, _split_file(name, split_id))
+        split_data = np.load(split_path)
+        train_masks.append(split_data["train_mask"])
+        val_masks.append(split_data["val_mask"])
+        test_masks.append(split_data["test_mask"])
+
+    graph.train_mask = tlx.convert_to_tensor(np.stack(train_masks, axis=1), dtype=tlx.bool)
+    graph.val_mask = tlx.convert_to_tensor(np.stack(val_masks, axis=1), dtype=tlx.bool)
+    graph.test_mask = tlx.convert_to_tensor(np.stack(test_masks, axis=1), dtype=tlx.bool)
+    return dataset, graph

examples/coed/readme.md

@@ -0,0 +1,80 @@
+# CoED-GNN Node Classification
+
+- Paper link: [https://arxiv.org/abs/2410.14109](https://arxiv.org/abs/2410.14109)
+- Author's code repo: [https://github.com/hormoz-lab/coed-gnn](https://github.com/hormoz-lab/coed-gnn)
+
+## Dataset Statics
+
+| Dataset    | # Nodes | # Edges | # Classes |
+|------------|---------|---------|-----------|
+| Cora       | 2,708   | 10,556  | 7         |
+| Texas      | 183     | 309     | 5         |
+| Wisconsin  | 251     | 515     | 5         |
+| Chameleon  | 2,277   | 36,101  | 5         |
+| Squirrel   | 5,201   | 217,073 | 5         |
+
+All datasets use the `Geom-GCN` 10 fixed splits for evaluation.
+
+## Files
+
+- `examples/coed/coed_trainer.py`: Multi-dataset training and evaluation entry
+- `gammagl/models/coed.py`: CoED-GNN backbone model
+- `gammagl/layers/conv/coed_conv.py`: CoED directional convolution layer
+
+## Results
+
+### Cora
+
+```bash
+TL_BACKEND="torch" python examples/coed/coed_trainer.py --dataset cora
+```
+
+| Metric     | Paper      | Our(torch)           |
+|------------|------------|----------------------|
+| Test Acc   | 86.42      | 87.00 +/- 1.44      |
+
+### Texas
+
+```bash
+TL_BACKEND="torch" python examples/coed/coed_trainer.py --dataset texas
+```
+
+| Metric     | Paper      | Our(torch)           |
+|------------|------------|----------------------|
+| Test Acc   |            |                      |
+
+### Wisconsin
+
+```bash
+TL_BACKEND="torch" python examples/coed/coed_trainer.py --dataset wisconsin
+```
+
+| Metric     | Paper      | Our(torch)           |
+|------------|------------|----------------------|
+| Test Acc   |            |                      |
+
+### Chameleon
+
+```bash
+TL_BACKEND="torch" python examples/coed/coed_trainer.py --dataset chameleon
+```
+
+| Metric     | Paper      | Our(torch)           |
+|------------|------------|----------------------|
+| Test Acc   |            |                      |
+
+### Squirrel
+
+```bash
+TL_BACKEND="torch" python examples/coed/coed_trainer.py --dataset squirrel
+```
+
+| Metric     | Paper      | Our(torch)           |
+|------------|------------|----------------------|
+| Test Acc   |            |                      |
+
+## Notes
+
+- The default setup uses `hidden_dim=64`, `num_layers=2`, `lr=1e-3`, `l2_coef=0.0`, `alpha=0.5`, `self_loop=True`, `normalize=False`, `self_feature_transform=False`, `patience=100`, `n_epoch=5000`.
+- The implementation evaluates all 10 Geom-GCN fixed splits and reports mean +/- std test accuracy.
+- The model and convolution layers are registered in `gammagl/models/__init__.py` and `gammagl/layers/conv/__init__.py` and can be imported via standard GammaGL paths.

examples/coed/reproduce_cora.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+set -e
+
+ROOT_DIR="$(cd "$(dirname "$0")/../.." && pwd)"
+source /home/mr/venv/gammagl-py311-cpu/bin/activate
+export TL_BACKEND=torch
+export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH}"
+
+python "${ROOT_DIR}/examples/coed/coed_trainer.py" "$@"

examples/coed/run_coed_cora.py

@@ -0,0 +1,19 @@
+"""Launcher for the CoED-GNN Cora reproduction."""
+
+import os
+import subprocess
+import sys
+
+
+def main():
+    root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+    env = os.environ.copy()
+    env.setdefault("TL_BACKEND", "torch")
+    env["PYTHONPATH"] = root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
+
+    cmd = [sys.executable, os.path.join(os.path.dirname(__file__), "coed_trainer.py")] + sys.argv[1:]
+    raise SystemExit(subprocess.call(cmd, env=env, cwd=root))
+
+
+if __name__ == "__main__":
+    main()

gammagl/layers/conv/__init__.py

@@ -36,6 +36,7 @@
 from .dhn_conv import DHNConv
 from .dna_conv import DNAConv
 from .rohehan_conv import RoheHANConv
+from .coed_conv import CoEDConv
 
 __all__ = [
     'MessagePassing',
@@ -75,7 +76,8 @@
     'HEATlayer',
     'DHNConv',
     'DNAConv',
-    'RoheHANConv'
+    'RoheHANConv',
+    'CoEDConv'
 ]
 
 classes = __all__

gammagl/layers/conv/coed_conv.py

@@ -0,0 +1,120 @@
+"""CoED directional convolution layer.
+
+This module implements the directional message passing operator used in
+`"Co-Embedding of Edges and Directions for Graph Neural Networks"
+<https://arxiv.org/abs/2410.14109>`_.
+"""
+
+import tensorlayerx as tlx
+from tensorlayerx.nn import Linear
+
+from gammagl.layers.conv import MessagePassing
+
+
+class CoEDConv(MessagePassing):
+    r"""The directional convolution operator used by CoED-GNN.
+
+    The layer separately aggregates messages for two directional channels and
+    optionally applies an additional self-feature transformation.
+
+    Parameters
+    ----------
+    in_channels: int
+        Size of each input sample.
+    out_channels: int
+        Size of each output sample.
+    self_feature_transform: bool, optional
+        If set to :obj:`True`, adds an extra linear transform on the input node
+        features and combines it with directional messages.
+    bias: bool, optional
+        If set to :obj:`False`, the layer will not learn additive bias terms.
+
+    """
+
+    def __init__(self, in_channels, out_channels, self_feature_transform=True, bias=True):
+        super().__init__()
+        self.self_feature_transform = self_feature_transform
+
+        self.lin_src_to_dst = Linear(
+            in_features=in_channels,
+            out_features=out_channels,
+            W_init="xavier_uniform",
+            b_init=None,
+        )
+        self.lin_dst_to_src = Linear(
+            in_features=in_channels,
+            out_features=out_channels,
+            W_init="xavier_uniform",
+            b_init=None,
+        )
+
+        if self_feature_transform:
+            self.lin_self = Linear(
+                in_features=in_channels,
+                out_features=out_channels,
+                W_init="xavier_uniform",
+                b_init=None,
+            )
+        else:
+            self.lin_self = None
+
+        if bias:
+            zeros = tlx.initializers.Zeros()
+            self.bias_src_to_dst = self._get_weights("bias_src_to_dst", shape=(out_channels,), init=zeros)
+            self.bias_dst_to_src = self._get_weights("bias_dst_to_src", shape=(out_channels,), init=zeros)
+            self.bias_self = (
+                self._get_weights("bias_self", shape=(out_channels,), init=zeros)
+                if self_feature_transform
+                else None
+            )
+        else:
+            self.bias_src_to_dst = None
+            self.bias_dst_to_src = None
+            self.bias_self = None
+
+    def forward(self, x, edge_index, edge_weight=None, num_nodes=None):
+        """Compute directional node representations."""
+        if num_nodes is None:
+            num_nodes = tlx.get_tensor_shape(x)[0]
+
+        if isinstance(edge_weight, (tuple, list)):
+            edge_weight_src_to_dst, edge_weight_dst_to_src = edge_weight
+        else:
+            edge_weight_src_to_dst = edge_weight
+            edge_weight_dst_to_src = edge_weight
+
+        x_src_to_dst = self.propagate(
+            x=x,
+            edge_index=edge_index,
+            edge_weight=edge_weight_src_to_dst,
+            num_nodes=num_nodes,
+        )
+        x_dst_to_src = self.propagate(
+            x=x,
+            edge_index=edge_index,
+            edge_weight=edge_weight_dst_to_src,
+            num_nodes=num_nodes,
+        )
+
+        x_src_to_dst = self.lin_src_to_dst.forward(x_src_to_dst)
+        x_dst_to_src = self.lin_dst_to_src.forward(x_dst_to_src)
+
+        if self.bias_src_to_dst is not None:
+            x_src_to_dst = x_src_to_dst + self.bias_src_to_dst
+        if self.bias_dst_to_src is not None:
+            x_dst_to_src = x_dst_to_src + self.bias_dst_to_src
+
+        if self.self_feature_transform:
+            x_self = self.lin_self.forward(x)
+            if self.bias_self is not None:
+                x_self = x_self + self.bias_self
+            return x_src_to_dst, x_dst_to_src, x_self
+
+        return x_src_to_dst, x_dst_to_src
+
+    def message(self, x, edge_index, edge_weight=None):
+        """Construct messages on each edge."""
+        msg = tlx.gather(x, edge_index[0, :])
+        if edge_weight is None:
+            return msg
+        return msg * tlx.reshape(edge_weight, (-1, 1))

gammagl/models/__init__.py

@@ -67,6 +67,7 @@
 from .sgformer import SGFormerModel
 from .adagad import PreModel, ReModel
 from .nodeid import NodeIDGNN
+from .coed import CoEDModel
 
 __all__ = [
     'HeCo',
@@ -142,7 +143,8 @@
     'sgformer',
     'PreModel',
     'ReModel'
-    , 'NodeIDGNN'
+    , 'NodeIDGNN',
+    'CoEDModel'
 ]
 
 classes = __all__