Upload folder using huggingface_hub

.gitignore

@@ -47,3 +47,14 @@ GraphUNets/data/*
 temp*
 
 *.tex
+
+src/data/*
+
+src/cora_seeds/*
+
+.venv/*
+
+
+.ipynb_checkpoints/
+
+*.dot

GCond/.gitignore

@@ -0,0 +1,140 @@
+# Auxillary file on MacOS
+.DS_Store
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+models/__pycache__/
+*.py[co
+dataset
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+data/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/d]
+*$py.class:
+
+def __init__(self):
+pass
+
+

GCond/KDD22_DosCond/README.md

@@ -0,0 +1,38 @@
+# DosCond
+
+[KDD 2022] A PyTorch Implementation for ["Condensing Graphs via One-Step Gradient Matching"](https://arxiv.org/abs/2206.07746) under node classification setting. For graph classification setting, please refer to [https://github.com/amazon-research/DosCond](https://github.com/amazon-research/DosCond). 
+
+
+Abstract
+----
+As training deep learning models on large dataset takes a lot of time and resources, it is desired to construct a small synthetic dataset with which we can train deep learning models sufficiently. There are recent works that have explored solutions on condensing image datasets through complex bi-level optimization. For instance, dataset condensation (DC) matches network gradients w.r.t. large-real data and small-synthetic data, where the network weights are optimized for multiple steps at each outer iteration. However, existing approaches have their inherent limitations: (1) they are not directly applicable to graphs where the data is discrete; and (2) the condensation process is computationally expensive due to the involved nested optimization. To bridge the gap, we investigate efficient dataset condensation tailored for graph datasets where we model the discrete graph structure as a probabilistic model. We further propose a one-step gradient matching scheme, which performs gradient matching for only one single step without training the network weights. 
+
+
+Here we do not implement the discrete structure learning, but only borrow the idea from ["Condensing Graphs via One-Step Gradient Matching"](https://arxiv.org/abs/2206.07746) to perform one-step gradient matching, which significantly speeds up the condensation process.
+
+
+Essentially, we can run the following commands:
+```
+python train_gcond_transduct.py --dataset citeseer --nlayers=2 --lr_feat=1e-3 --lr_adj=1e-3 --r=0.5  --sgc=0 --dis=mse --one_step=1  --epochs=3000
+python train_gcond_transduct.py --dataset cora --nlayers=2 --lr_feat=1e-3 --lr_adj=1e-3 --r=0.5  --sgc=0 --dis=mse --gpu_id=2 --one_step=1  --epochs=5000
+python train_gcond_transduct.py --dataset pubmed --nlayers=2 --lr_feat=1e-3 --lr_adj=1e-3 --r=0.5  --sgc=0 --dis=mse --gpu_id=2 --one_step=1  --epochs=2000
+python train_gcond_transduct.py --dataset ogbn-arxiv --nlayers=2 --lr_feat=1e-2 --lr_adj=2e-2 --r=0.001  --sgc=1 --dis=ours --gpu_id=2 --one_step=1  --epochs=1000
+python train_gcond_induct.py --dataset flickr --nlayers=2 --lr_feat=5e-3 --lr_adj=5e-3 --r=0.001  --sgc=0 --dis=mse --gpu_id=3 --one_step=1  --epochs=1000
+```
+Note that using smaller learning rate and larger epochs can get even higher performance.
+
+
+## Cite
+For more information, you can take a look at the [paper](https://arxiv.org/abs/2206.07746).
+
+If you find this repo to be useful, please cite our paper. Thank you.
+```
+@inproceedings{jin2022condensing,
+  title={Condensing Graphs via One-Step Gradient Matching},
+  author={Jin, Wei and Tang, Xianfeng and Jiang, Haoming and Li, Zheng and Zhang, Danqing and Tang, Jiliang and Yin, Bing},
+  booktitle={Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
+  pages={720--730},
+  year={2022}
+}
+```
+

GCond/README.md

@@ -0,0 +1,139 @@
+# GCond
+[ICLR 2022] The PyTorch implementation for ["Graph Condensation for Graph Neural Networks"](https://cse.msu.edu/~jinwei2/files/GCond.pdf) is provided under the main directory. 
+
+[KDD 2022] The implementation for ["Condensing Graphs via One-Step Gradient Matching"](https://arxiv.org/abs/2206.07746) is shown in the `KDD22_DosCond` directory. See [link](https://github.com/ChandlerBang/GCond/tree/main/KDD22_DosCond).
+
+[IJCAI 2024] Please read our recent survey ["A Comprehensive Survey on Graph Reduction: Sparsification, Coarsening, and Condensation"](https://arxiv.org/abs/2402.03358) for a detailed review of graph reduction techniques!
+
+[ArXiv 2024] We released a benchmarking framework for graph condensation ["GC4NC: A Benchmark Framework for Graph Condensation with New Insights"](https://arxiv.org/abs/2406.16715), including **robustness**, **privacy preservation**, NAS performance, property analysis, etc!
+
+Abstract
+----
+We propose and study the problem of graph condensation for graph neural networks (GNNs). Specifically, we aim to condense the large, original graph into a small, synthetic, and highly-informative graph, such that GNNs trained on the small graph and large graph have comparable performance. Extensive experiments have demonstrated the effectiveness of the proposed framework in condensing different graph datasets into informative smaller graphs. In particular, we are able to approximate the original test accuracy by 95.3% on Reddit, 99.8% on Flickr and 99.0% on Citeseer, while reducing their graph size by more than 99.9%, and the condensed graphs can be used to train various GNN architectures.
+
+
+![]()
+
+<div align=center><img src="https://github.com/ChandlerBang/GCond/blob/main/GCond.png" width="800"/></div>
+
+
+## A Nice Survey Paper
+Please check out our survey paper blew, which summarizes the recent advances in graph condensation.
+
+![image](https://github.com/CurryTang/Towards-Graph-Foundation-Models-New-perspective-/assets/15672123/89a23a37-71d4-47f7-8949-7d859a41e369)![image](https://github.com/CurryTang/Towards-Graph-Foundation-Models-New-perspective-/assets/15672123/89a23a37-71d4-47f7-8949-7d859a41e369)![image](https://github.com/CurryTang/Towards-Graph-Foundation-Models-New-perspective-/assets/15672123/89a23a37-71d4-47f7-8949-7d859a41e369)[[A Comprehensive Survey on Graph Reduction: Sparsification, Coarsening, and Condensation]](https://arxiv.org/abs/2402.03358)
+
+
+
+## Requirements
+Please see [requirements.txt](https://github.com/ChandlerBang/GCond/blob/main/requirements.txt).
+```
+torch==1.7.0
+torch_geometric==1.6.3
+scipy==1.6.2
+numpy==1.19.2
+ogb==1.3.0
+tqdm==4.59.0
+torch_sparse==0.6.9
+deeprobust==0.2.4
+scikit_learn==1.0.2
+```
+
+## Download Datasets
+For cora, citeseer and pubmed, the code will directly download them; so no extra script is needed.
+For reddit, flickr and arxiv, we use the datasets provided by [GraphSAINT](https://github.com/GraphSAINT/GraphSAINT). 
+They are available on [Google Drive link](https://drive.google.com/open?id=1zycmmDES39zVlbVCYs88JTJ1Wm5FbfLz) (alternatively, [BaiduYun link (code: f1ao)](https://pan.baidu.com/s/1SOb0SiSAXavwAcNqkttwcg)). Rename the folder to `data` at the root directory. Note that the links are provided by GraphSAINT team. 
+
+
+
+
+## Run the code
+For transductive setting, please run the following command:
+```
+python train_gcond_transduct.py --dataset cora --nlayers=2 --lr_feat=1e-4 --gpu_id=0  --lr_adj=1e-4 --r=0.5  
+```
+where `r` indicates the ratio of condensed samples to the labeled samples. For instance, there are only 140 labeled nodes in Cora dataset, so `r=0.5` indicates the number of condensed samples are 70, **which corresponds to  r=2.6%=70/2710 in our paper**. Thus, the parameter `r` is different from the real reduction rate in the paper for the transductive setting, please see the following table for the correspondence.
+
+|              | `r` in the code     | `r` in the paper  (real reduction rate)    |
+|--------------|-------------------|---------------------|
+| Transductive | Cora, r=0.5       | Cora, r=2.6%        |
+| Transductive | Citeseer, r=0.5   | Citeseer, r=1.8%    |
+| Transductive | Ogbn-arxiv, r= 0.005 | Ogbn-arxiv, r=0.25% |
+| Transductive | Pubmed, r=0.5     | Pubmed, r=0.3%      |
+| Inductive    | Flickr, r=0.01    | Flickr, r=1%        |
+| Inductive    | Reddit, r=0.001   | Reddit, r=0.1%      |
+
+For inductive setting, please run the following command:
+```
+python train_gcond_induct.py --dataset flickr --nlayers=2 --lr_feat=0.01 --gpu_id=0  --lr_adj=0.01 --r=0.005 --epochs=1000  --outer=10 --inner=1
+```
+
+## Reproduce the performance
+The generated graphs are saved in the folder `saved_ours`; you can directly load them to test the performance.
+
+For Table 2, run `bash scripts/run_main.sh`.
+
+For Table 3, run `bash scripts/run_cross.sh`.
+
+## [Faster Condensation!] One-Step Gradient Matching
+From the KDD'22 paper ["Condensing Graphs via One-Step Gradient Matching"](https://arxiv.org/abs/2206.07746), we know that performing gradient matching for only one step can also achieve a good performance while significantly accelerating the condensation process. Hence, we can run the following command to perform one-step gradient matching, which is essentially much faster than the original version:
+```
+python train_gcond_transduct.py --dataset citeseer --nlayers=2 --lr_feat=1e-2 --lr_adj=1e-2 --r=0.5 \
+    --sgc=0 --dis=mse --gpu_id=2 --one_step=1  --epochs=3000
+```
+For more commands, please go to [`KDD22_DosCond`](https://github.com/ChandlerBang/GCond/tree/main/KDD22_DosCond).
+
+**[Note]: I found that sometimes using MSE loss for gradient matching can be more stable than using `ours` loss**, and it gives more flexibility on the model used in condensation (using GCN as the backbone can also generate good condensed graphs). 
+
+
+## Whole Dataset Performance 
+When we do coreset selection, we need to first the model on the whole dataset. Thus we can obtain the performanceo of whole dataset by running `train_coreset.py` and `train_coreset_induct.py`:
+```
+python train_coreset.py --dataset cora --r=0.01  --method=random
+python train_coreset_induct.py --dataset flickr --r=0.01  --method=random
+```
+
+## Coreset Performance
+Run the following code to get the coreset performance for transductive setting.
+```
+python train_coreset.py --dataset cora --r=0.01  --method=herding
+python train_coreset.py --dataset cora --r=0.01  --method=random
+python train_coreset.py --dataset cora --r=0.01  --method=kcenter
+```
+Similarly, run the following code for the inductive setting.
+```
+python train_coreset_induct.py --dataset flickr --r=0.01  --method=kcenter
+```
+
+
+## Cite
+If you find this repo to be useful, please cite our three papers. Thank you!
+```
+@inproceedings{
+    jin2022graph,
+    title={Graph Condensation for Graph Neural Networks},
+    author={Wei Jin and Lingxiao Zhao and Shichang Zhang and Yozen Liu and Jiliang Tang and Neil Shah},
+    booktitle={International Conference on Learning Representations},
+    year={2022},
+    url={https://openreview.net/forum?id=WLEx3Jo4QaB}
+}
+```
+
+```
+@inproceedings{jin2022condensing,
+  title={Condensing Graphs via One-Step Gradient Matching},
+  author={Jin, Wei and Tang, Xianfeng and Jiang, Haoming and Li, Zheng and Zhang, Danqing and Tang, Jiliang and Yin, Bing},
+  booktitle={Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
+  pages={720--730},
+  year={2022}
+}
+```
+
+```
+@article{hashemi2024comprehensive,
+  title={A Comprehensive Survey on Graph Reduction: Sparsification, Coarsening, and Condensation},
+  author={Hashemi, Mohammad and Gong, Shengbo and Ni, Juntong and Fan, Wenqi and Prakash, B Aditya and Jin, Wei},
+  journal={International Joint Conference on Artificial Intelligence (IJCAI)},
+  year={2024}
+}
+```
+

GCond/configs.py

@@ -0,0 +1,24 @@
+'''Configuration'''
+
+def load_config(args):
+    dataset = args.dataset
+    if dataset in ['flickr']:
+        args.nlayers = 2
+        args.hidden = 256
+        args.weight_decay = 5e-3
+        args.dropout = 0.0
+
+    if dataset in ['reddit']:
+        args.nlayers = 2
+        args.hidden = 256
+        args.weight_decay = 0e-4
+        args.dropout = 0
+
+    if dataset in ['ogbn-arxiv']:
+        args.hidden = 256
+        args.weight_decay = 0
+        args.dropout = 0
+
+    return args
+
+

GCond/coreset/__init__.py

@@ -0,0 +1,3 @@
+from .all_methods import KCenter, Herding, Random, LRMC
+
+__all__ = ['KCenter', 'Herding', 'Random', 'LRMC']

GCond/coreset/all_methods.py

@@ -0,0 +1,212 @@
+import torch
+import numpy as np
+import json
+
+
+class Base:
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        self.data = data
+        self.args = args
+        self.device = device
+        n = int(data.feat_train.shape[0] * args.reduction_rate)
+        d = data.feat_train.shape[1]
+        self.nnodes_syn = n
+        self.labels_syn = torch.LongTensor(self.generate_labels_syn(data)).to(device)
+
+    def generate_labels_syn(self, data):
+        from collections import Counter
+        counter = Counter(data.labels_train)
+        num_class_dict = {}
+        n = len(data.labels_train)
+
+        sorted_counter = sorted(counter.items(), key=lambda x:x[1])
+        sum_ = 0
+        labels_syn = []
+        self.syn_class_indices = {}
+        for ix, (c, num) in enumerate(sorted_counter):
+            if ix == len(sorted_counter) - 1:
+                num_class_dict[c] = int(n * self.args.reduction_rate) - sum_
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+            else:
+                num_class_dict[c] = max(int(num * self.args.reduction_rate), 1)
+                sum_ += num_class_dict[c]
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+
+        self.num_class_dict = num_class_dict
+        return labels_syn
+
+    def select(self):
+        return
+
+class KCenter(Base):
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        super(KCenter, self).__init__(data, args, device='cuda', **kwargs)
+
+    def select(self, embeds, inductive=False):
+        # feature: embeds
+        # kcenter # class by class
+        num_class_dict = self.num_class_dict
+        if inductive:
+            idx_train = np.arange(len(self.data.idx_train))
+        else:
+            idx_train = self.data.idx_train
+        labels_train = self.data.labels_train
+        idx_selected = []
+
+        for class_id, cnt in num_class_dict.items():
+            idx = idx_train[labels_train==class_id]
+            feature = embeds[idx]
+            mean = torch.mean(feature, dim=0, keepdim=True)
+            # dis = distance(feature, mean)[:,0]
+            dis = torch.cdist(feature, mean)[:,0]
+            rank = torch.argsort(dis)
+            idx_centers = rank[:1].tolist()
+            for i in range(cnt-1):
+                feature_centers = feature[idx_centers]
+                dis_center = torch.cdist(feature, feature_centers)
+                dis_min, _ = torch.min(dis_center, dim=-1)
+                id_max = torch.argmax(dis_min).item()
+                idx_centers.append(id_max)
+
+            idx_selected.append(idx[idx_centers])
+        # return np.array(idx_selected).reshape(-1)
+        return np.hstack(idx_selected)
+
+
+class Herding(Base):
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        super(Herding, self).__init__(data, args, device='cuda', **kwargs)
+
+    def select(self, embeds, inductive=False):
+        num_class_dict = self.num_class_dict
+        if inductive:
+            idx_train = np.arange(len(self.data.idx_train))
+        else:
+            idx_train = self.data.idx_train
+        labels_train = self.data.labels_train
+        idx_selected = []
+
+        # herding # class by class
+        for class_id, cnt in num_class_dict.items():
+            idx = idx_train[labels_train==class_id]
+            features = embeds[idx]
+            mean = torch.mean(features, dim=0, keepdim=True)
+            selected = []
+            idx_left = np.arange(features.shape[0]).tolist()
+
+            for i in range(cnt):
+                det = mean*(i+1) - torch.sum(features[selected], dim=0)
+                dis = torch.cdist(det, features[idx_left])
+                id_min = torch.argmin(dis)
+                selected.append(idx_left[id_min])
+                del idx_left[id_min]
+            idx_selected.append(idx[selected])
+        # return np.array(idx_selected).reshape(-1)
+        return np.hstack(idx_selected)
+
+
+class Random(Base):
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        super(Random, self).__init__(data, args, device='cuda', **kwargs)
+
+    def select(self, embeds, inductive=False):
+        num_class_dict = self.num_class_dict
+        if inductive:
+            idx_train = np.arange(len(self.data.idx_train))
+        else:
+            idx_train = self.data.idx_train
+
+        labels_train = self.data.labels_train
+        idx_selected = []
+
+        for class_id, cnt in num_class_dict.items():
+            idx = idx_train[labels_train==class_id]
+            selected = np.random.permutation(idx)
+            idx_selected.append(selected[:cnt])
+
+        # return np.array(idx_selected).reshape(-1)
+        return np.hstack(idx_selected)
+
+
+class LRMC(Base):
+    """
+    Coreset selection using precomputed seed nodes from the Laplacian‑Integrated
+    Relaxed Maximal Clique (L‑RMC) algorithm.  Seed nodes are read from a JSON
+    file specified by ``args.lrmc_seeds_path`` and used to preferentially select
+    training examples.  Per‑class reduction counts are respected: if a class has
+    fewer seeds than required, random training nodes from that class are added
+    until the quota is met.
+    """
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        super(LRMC, self).__init__(data, args, device=device, **kwargs)
+        seeds_path = getattr(args, 'lrmc_seeds_path', None)
+        if seeds_path is None:
+            raise ValueError(
+                "LRMC method selected but no path to seed file provided. "
+                "Please specify --lrmc_seeds_path when running the training script."
+            )
+        self.seed_nodes = self._load_seed_nodes(seeds_path)
+
+    def _load_seed_nodes(self, path: str):
+        # Parse seed nodes from JSON file (supports 'seed_nodes' or 'members').
+        with open(path, 'r') as f:
+            js = json.load(f)
+        clusters = js.get('clusters', [])
+        if not clusters:
+            raise ValueError(f"No clusters found in L‑RMC seeds file {path}")
+        def _cluster_length(c):
+            nodes = c.get('seed_nodes') or c.get('members') or []
+            return len(nodes)
+        best_cluster = max(clusters, key=_cluster_length)
+        nodes = best_cluster.get('seed_nodes') or best_cluster.get('members') or []
+        seed_nodes = []
+        for u in nodes:
+            try:
+                uid = int(u)
+            except Exception:
+                continue
+            zero_idx = uid - 1
+            if zero_idx >= 0:
+                seed_nodes.append(zero_idx)
+            else:
+                if uid >= 0:
+                    seed_nodes.append(uid)
+        seed_nodes = sorted(set(seed_nodes))
+        return seed_nodes
+
+    def select(self, embeds, inductive=False):
+        # Determine training indices depending on the inductive setting.
+        if inductive:
+            idx_train = np.arange(len(self.data.idx_train))
+            labels_train = self.data.labels_train
+        else:
+            idx_train = self.data.idx_train
+            labels_train = self.data.labels_train
+        num_class_dict = self.num_class_dict
+        idx_selected = []
+        seed_set = set(self.seed_nodes)
+        # Pick seed nodes per class; fill remainder with random nodes if needed.
+        for class_id, cnt in num_class_dict.items():
+            class_mask = (labels_train == class_id)
+            class_indices = idx_train[class_mask]
+            seed_in_class = [u for u in class_indices if u in seed_set]
+            selected = seed_in_class[:min(len(seed_in_class), cnt)]
+            remaining_required = cnt - len(selected)
+            if remaining_required > 0:
+                remaining_candidates = [u for u in class_indices if u not in selected]
+                if len(remaining_candidates) <= remaining_required:
+                    additional = remaining_candidates
+                else:
+                    additional = np.random.choice(remaining_candidates, remaining_required, replace=False).tolist()
+                selected += additional
+            idx_selected.append(np.array(selected))
+        return np.hstack(idx_selected)
+
+

GCond/gcond_agent_induct.py

@@ -0,0 +1,327 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.nn import Parameter
+import torch.nn.functional as F
+from utils import match_loss, regularization, row_normalize_tensor
+import deeprobust.graph.utils as utils
+from copy import deepcopy
+import numpy as np
+from tqdm import tqdm
+from models.gcn import GCN
+from models.sgc import SGC
+from models.sgc_multi import SGC as SGC1
+from models.parametrized_adj import PGE
+import scipy.sparse as sp
+from torch_sparse import SparseTensor
+
+
+class GCond:
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        self.data = data
+        self.args = args
+        self.device = device
+
+        n = int(len(data.idx_train) * args.reduction_rate)
+        d = data.feat_train.shape[1]
+        self.nnodes_syn = n
+        self.feat_syn = nn.Parameter(torch.FloatTensor(n, d).to(device))
+        self.pge = PGE(nfeat=d, nnodes=n, device=device, args=args).to(device)
+
+        self.labels_syn = torch.LongTensor(self.generate_labels_syn(data)).to(device)
+        self.reset_parameters()
+        self.optimizer_feat = torch.optim.Adam([self.feat_syn], lr=args.lr_feat)
+        self.optimizer_pge = torch.optim.Adam(self.pge.parameters(), lr=args.lr_adj)
+        print('adj_syn:', (n,n), 'feat_syn:', self.feat_syn.shape)
+
+    def reset_parameters(self):
+        self.feat_syn.data.copy_(torch.randn(self.feat_syn.size()))
+
+    def generate_labels_syn(self, data):
+        from collections import Counter
+        counter = Counter(data.labels_train)
+        num_class_dict = {}
+        n = len(data.labels_train)
+
+        sorted_counter = sorted(counter.items(), key=lambda x:x[1])
+        sum_ = 0
+        labels_syn = []
+        self.syn_class_indices = {}
+
+        for ix, (c, num) in enumerate(sorted_counter):
+            if ix == len(sorted_counter) - 1:
+                num_class_dict[c] = int(n * self.args.reduction_rate) - sum_
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+            else:
+                num_class_dict[c] = max(int(num * self.args.reduction_rate), 1)
+                sum_ += num_class_dict[c]
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+
+        self.num_class_dict = num_class_dict
+        return labels_syn
+
+    def test_with_val(self, verbose=True):
+        res = []
+
+        data, device = self.data, self.device
+        feat_syn, pge, labels_syn = self.feat_syn.detach(), \
+                                self.pge, self.labels_syn
+        # with_bn = True if args.dataset in ['ogbn-arxiv'] else False
+        dropout = 0.5 if self.args.dataset in ['reddit'] else 0
+        model = GCN(nfeat=feat_syn.shape[1], nhid=self.args.hidden, dropout=dropout,
+                    weight_decay=5e-4, nlayers=2,
+                    nclass=data.nclass, device=device).to(device)
+
+        adj_syn = pge.inference(feat_syn)
+        args = self.args
+
+        if args.save:
+            torch.save(adj_syn, f'saved_ours/adj_{args.dataset}_{args.reduction_rate}_{args.seed}.pt')
+            torch.save(feat_syn, f'saved_ours/feat_{args.dataset}_{args.reduction_rate}_{args.seed}.pt')
+
+        noval = True
+        model.fit_with_val(feat_syn, adj_syn, labels_syn, data,
+                     train_iters=600, normalize=True, verbose=False, noval=noval)
+
+        model.eval()
+        labels_test = torch.LongTensor(data.labels_test).cuda()
+
+        output = model.predict(data.feat_test, data.adj_test)
+
+        loss_test = F.nll_loss(output, labels_test)
+        acc_test = utils.accuracy(output, labels_test)
+        res.append(acc_test.item())
+        if verbose:
+            print("Test set results:",
+                  "loss= {:.4f}".format(loss_test.item()),
+                  "accuracy= {:.4f}".format(acc_test.item()))
+        print(adj_syn.sum(), adj_syn.sum()/(adj_syn.shape[0]**2))
+
+        if False:
+            if self.args.dataset == 'ogbn-arxiv':
+                thresh = 0.6
+            elif self.args.dataset == 'reddit':
+                thresh = 0.91
+            else:
+                thresh = 0.7
+
+            labels_train = torch.LongTensor(data.labels_train).cuda()
+            output = model.predict(data.feat_train, data.adj_train)
+            # loss_train = F.nll_loss(output, labels_train)
+            # acc_train = utils.accuracy(output, labels_train)
+            loss_train = torch.tensor(0)
+            acc_train = torch.tensor(0)
+            if verbose:
+                print("Train set results:",
+                      "loss= {:.4f}".format(loss_train.item()),
+                      "accuracy= {:.4f}".format(acc_train.item()))
+            res.append(acc_train.item())
+        return res
+
+    def train(self, verbose=True):
+        args = self.args
+        data = self.data
+        feat_syn, pge, labels_syn = self.feat_syn, self.pge, self.labels_syn
+        features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        syn_class_indices = self.syn_class_indices
+        features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        feat_sub, adj_sub = self.get_sub_adj_feat(features)
+        self.feat_syn.data.copy_(feat_sub)
+
+        if utils.is_sparse_tensor(adj):
+            adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+        else:
+            adj_norm = utils.normalize_adj_tensor(adj)
+
+        adj = adj_norm
+        adj = SparseTensor(row=adj._indices()[0], col=adj._indices()[1],
+                value=adj._values(), sparse_sizes=adj.size()).t()
+
+
+        outer_loop, inner_loop = get_loops(args)
+
+        for it in range(args.epochs+1):
+            loss_avg = 0
+            if args.sgc==1:
+                model = SGC(nfeat=data.feat_train.shape[1], nhid=args.hidden,
+                            nclass=data.nclass, dropout=args.dropout,
+                            nlayers=args.nlayers, with_bn=False,
+                            device=self.device).to(self.device)
+            elif args.sgc==2:
+                model = SGC1(nfeat=data.feat_train.shape[1], nhid=args.hidden,
+                            nclass=data.nclass, dropout=args.dropout,
+                            nlayers=args.nlayers, with_bn=False,
+                            device=self.device).to(self.device)
+
+            else:
+                model = GCN(nfeat=data.feat_train.shape[1], nhid=args.hidden,
+                            nclass=data.nclass, dropout=args.dropout, nlayers=args.nlayers,
+                            device=self.device).to(self.device)
+
+            model.initialize()
+
+            model_parameters = list(model.parameters())
+
+            optimizer_model = torch.optim.Adam(model_parameters, lr=args.lr_model)
+            model.train()
+
+            for ol in range(outer_loop):
+                adj_syn = pge(self.feat_syn)
+                adj_syn_norm = utils.normalize_adj_tensor(adj_syn, sparse=False)
+                feat_syn_norm = feat_syn
+
+                BN_flag = False
+                for module in model.modules():
+                    if 'BatchNorm' in module._get_name(): #BatchNorm
+                        BN_flag = True
+                if BN_flag:
+                    model.train() # for updating the mu, sigma of BatchNorm
+                    output_real = model.forward(features, adj_norm)
+                    for module in model.modules():
+                        if 'BatchNorm' in module._get_name():  #BatchNorm
+                            module.eval() # fix mu and sigma of every BatchNorm layer
+
+                loss = torch.tensor(0.0).to(self.device)
+                for c in range(data.nclass):
+                    if c not in self.num_class_dict:
+                        continue
+
+                    batch_size, n_id, adjs = data.retrieve_class_sampler(
+                            c, adj, transductive=False, args=args)
+
+                    if args.nlayers == 1:
+                        adjs = [adjs]
+                    adjs = [adj.to(self.device) for adj in adjs]
+                    output = model.forward_sampler(features[n_id], adjs)
+                    loss_real = F.nll_loss(output, labels[n_id[:batch_size]])
+                    gw_real = torch.autograd.grad(loss_real, model_parameters)
+                    gw_real = list((_.detach().clone() for _ in gw_real))
+
+                    ind = syn_class_indices[c]
+                    if args.nlayers == 1:
+                        adj_syn_norm_list = [adj_syn_norm[ind[0]: ind[1]]]
+                    else:
+                        adj_syn_norm_list = [adj_syn_norm]*(args.nlayers-1) + \
+                                [adj_syn_norm[ind[0]: ind[1]]]
+
+                    output_syn = model.forward_sampler_syn(feat_syn, adj_syn_norm_list)
+                    loss_syn = F.nll_loss(output_syn, labels_syn[ind[0]: ind[1]])
+
+                    gw_syn = torch.autograd.grad(loss_syn, model_parameters, create_graph=True)
+                    coeff = self.num_class_dict[c] / max(self.num_class_dict.values())
+                    loss += coeff  * match_loss(gw_syn, gw_real, args, device=self.device)
+
+                loss_avg += loss.item()
+                # TODO: regularize
+                if args.alpha > 0:
+                    loss_reg = args.alpha * regularization(adj_syn, utils.tensor2onehot(labels_syn))
+                # else:
+                else:
+                    loss_reg = torch.tensor(0)
+
+                loss = loss + loss_reg
+
+                # update sythetic graph
+                self.optimizer_feat.zero_grad()
+                self.optimizer_pge.zero_grad()
+                loss.backward()
+                if it % 50 < 10:
+                    self.optimizer_pge.step()
+                else:
+                    self.optimizer_feat.step()
+
+                if args.debug and ol % 5 ==0:
+                    print('Gradient matching loss:', loss.item())
+
+                if ol == outer_loop - 1:
+                    # print('loss_reg:', loss_reg.item())
+                    # print('Gradient matching loss:', loss.item())
+                    break
+
+
+                feat_syn_inner = feat_syn.detach()
+                adj_syn_inner = pge.inference(feat_syn)
+                adj_syn_inner_norm = utils.normalize_adj_tensor(adj_syn_inner, sparse=False)
+                feat_syn_inner_norm = feat_syn_inner
+                for j in range(inner_loop):
+                    optimizer_model.zero_grad()
+                    output_syn_inner = model.forward(feat_syn_inner_norm, adj_syn_inner_norm)
+                    loss_syn_inner = F.nll_loss(output_syn_inner, labels_syn)
+                    loss_syn_inner.backward()
+                    optimizer_model.step() # update gnn param
+
+            loss_avg /= (data.nclass*outer_loop)
+            if it % 50 == 0:
+                print('Epoch {}, loss_avg: {}'.format(it, loss_avg))
+
+            eval_epochs = [100, 200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800, 2000, 3000, 4000, 5000]
+
+            if verbose and it in eval_epochs:
+            # if verbose and (it+1) % 500 == 0:
+                res = []
+                runs = 1 if args.dataset in ['ogbn-arxiv', 'reddit', 'flickr'] else 3
+                for i in range(runs):
+                    # self.test()
+                    res.append(self.test_with_val())
+                res = np.array(res)
+                print('Test:',
+                        repr([res.mean(0), res.std(0)]))
+
+
+
+    def get_sub_adj_feat(self, features):
+        data = self.data
+        args = self.args
+        idx_selected = []
+
+        from collections import Counter;
+        counter = Counter(self.labels_syn.cpu().numpy())
+
+        for c in range(data.nclass):
+            tmp = data.retrieve_class(c, num=counter[c])
+            tmp = list(tmp)
+            idx_selected = idx_selected + tmp
+        idx_selected = np.array(idx_selected).reshape(-1)
+        features = features[idx_selected]
+
+        # adj_knn = torch.zeros((data.nclass*args.nsamples, data.nclass*args.nsamples)).to(self.device)
+        # for i in range(data.nclass):
+        #     idx = np.arange(i*args.nsamples, i*args.nsamples+args.nsamples)
+        #     adj_knn[np.ix_(idx, idx)] = 1
+
+        from sklearn.metrics.pairwise import cosine_similarity
+        # features[features!=0] = 1
+        k = 2
+        sims = cosine_similarity(features.cpu().numpy())
+        sims[(np.arange(len(sims)), np.arange(len(sims)))] = 0
+        for i in range(len(sims)):
+            indices_argsort = np.argsort(sims[i])
+            sims[i, indices_argsort[: -k]] = 0
+        adj_knn = torch.FloatTensor(sims).to(self.device)
+        return features, adj_knn
+
+
+def get_loops(args):
+    # Get the two hyper-parameters of outer-loop and inner-loop.
+    # The following values are empirically good.
+    if args.one_step:
+        return 10, 0
+
+    if args.dataset in ['ogbn-arxiv']:
+        return 20, 0
+    if args.dataset in ['reddit']:
+        return args.outer, args.inner
+    if args.dataset in ['flickr']:
+        return args.outer, args.inner
+        # return 10, 1
+    if args.dataset in ['cora']:
+        return 20, 10
+    if args.dataset in ['citeseer']:
+        return 20, 5 # at least 200 epochs
+    else:
+        return 20, 5
+

GCond/gcond_agent_transduct.py

@@ -0,0 +1,326 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.nn import Parameter
+import torch.nn.functional as F
+from utils import match_loss, regularization, row_normalize_tensor
+import deeprobust.graph.utils as utils
+from copy import deepcopy
+import numpy as np
+from tqdm import tqdm
+from models.gcn import GCN
+from models.sgc import SGC
+from models.sgc_multi import SGC as SGC1
+from models.parametrized_adj import PGE
+import scipy.sparse as sp
+from torch_sparse import SparseTensor
+
+
+class GCond:
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        self.data = data
+        self.args = args
+        self.device = device
+
+        # n = data.nclass * args.nsamples
+        n = int(data.feat_train.shape[0] * args.reduction_rate)
+        # from collections import Counter; print(Counter(data.labels_train))
+
+        d = data.feat_train.shape[1]
+        self.nnodes_syn = n
+        self.feat_syn = nn.Parameter(torch.FloatTensor(n, d).to(device))
+        self.pge = PGE(nfeat=d, nnodes=n, device=device,args=args).to(device)
+
+        self.labels_syn = torch.LongTensor(self.generate_labels_syn(data)).to(device)
+
+        self.reset_parameters()
+        self.optimizer_feat = torch.optim.Adam([self.feat_syn], lr=args.lr_feat)
+        self.optimizer_pge = torch.optim.Adam(self.pge.parameters(), lr=args.lr_adj)
+        print('adj_syn:', (n,n), 'feat_syn:', self.feat_syn.shape)
+
+    def reset_parameters(self):
+        self.feat_syn.data.copy_(torch.randn(self.feat_syn.size()))
+
+    def generate_labels_syn(self, data):
+        from collections import Counter
+        counter = Counter(data.labels_train)
+        num_class_dict = {}
+        n = len(data.labels_train)
+
+        sorted_counter = sorted(counter.items(), key=lambda x:x[1])
+        sum_ = 0
+        labels_syn = []
+        self.syn_class_indices = {}
+        for ix, (c, num) in enumerate(sorted_counter):
+            if ix == len(sorted_counter) - 1:
+                num_class_dict[c] = int(n * self.args.reduction_rate) - sum_
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+            else:
+                num_class_dict[c] = max(int(num * self.args.reduction_rate), 1)
+                sum_ += num_class_dict[c]
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+
+        self.num_class_dict = num_class_dict
+        return labels_syn
+
+
+    def test_with_val(self, verbose=True):
+        res = []
+
+        data, device = self.data, self.device
+        feat_syn, pge, labels_syn = self.feat_syn.detach(), \
+                                self.pge, self.labels_syn
+
+        # with_bn = True if args.dataset in ['ogbn-arxiv'] else False
+        model = GCN(nfeat=feat_syn.shape[1], nhid=self.args.hidden, dropout=0.5,
+                    weight_decay=5e-4, nlayers=2,
+                    nclass=data.nclass, device=device).to(device)
+
+        if self.args.dataset in ['ogbn-arxiv']:
+            model = GCN(nfeat=feat_syn.shape[1], nhid=self.args.hidden, dropout=0.5,
+                        weight_decay=0e-4, nlayers=2, with_bn=False,
+                        nclass=data.nclass, device=device).to(device)
+
+        adj_syn = pge.inference(feat_syn)
+        args = self.args
+
+        if self.args.save:
+            torch.save(adj_syn, f'saved_ours/adj_{args.dataset}_{args.reduction_rate}_{args.seed}.pt')
+            torch.save(feat_syn, f'saved_ours/feat_{args.dataset}_{args.reduction_rate}_{args.seed}.pt')
+
+        if self.args.lr_adj == 0:
+            n = len(labels_syn)
+            adj_syn = torch.zeros((n, n))
+
+        model.fit_with_val(feat_syn, adj_syn, labels_syn, data,
+                     train_iters=600, normalize=True, verbose=False)
+
+        model.eval()
+        labels_test = torch.LongTensor(data.labels_test).cuda()
+
+        labels_train = torch.LongTensor(data.labels_train).cuda()
+        output = model.predict(data.feat_train, data.adj_train)
+        loss_train = F.nll_loss(output, labels_train)
+        acc_train = utils.accuracy(output, labels_train)
+        if verbose:
+            print("Train set results:",
+                  "loss= {:.4f}".format(loss_train.item()),
+                  "accuracy= {:.4f}".format(acc_train.item()))
+        res.append(acc_train.item())
+
+        # Full graph
+        output = model.predict(data.feat_full, data.adj_full)
+        loss_test = F.nll_loss(output[data.idx_test], labels_test)
+        acc_test = utils.accuracy(output[data.idx_test], labels_test)
+        res.append(acc_test.item())
+        if verbose:
+            print("Test set results:",
+                  "loss= {:.4f}".format(loss_test.item()),
+                  "accuracy= {:.4f}".format(acc_test.item()))
+        return res
+
+    def train(self, verbose=True):
+        args = self.args
+        data = self.data
+        feat_syn, pge, labels_syn = self.feat_syn, self.pge, self.labels_syn
+        features, adj, labels = data.feat_full, data.adj_full, data.labels_full
+        idx_train = data.idx_train
+
+        syn_class_indices = self.syn_class_indices
+
+        features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+
+        feat_sub, adj_sub = self.get_sub_adj_feat(features)
+        self.feat_syn.data.copy_(feat_sub)
+
+        if utils.is_sparse_tensor(adj):
+            adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+        else:
+            adj_norm = utils.normalize_adj_tensor(adj)
+
+        adj = adj_norm
+        adj = SparseTensor(row=adj._indices()[0], col=adj._indices()[1],
+                value=adj._values(), sparse_sizes=adj.size()).t()
+
+
+        outer_loop, inner_loop = get_loops(args)
+        loss_avg = 0
+
+        for it in range(args.epochs+1):
+            if args.dataset in ['ogbn-arxiv']:
+                model = SGC1(nfeat=feat_syn.shape[1], nhid=self.args.hidden,
+                            dropout=0.0, with_bn=False,
+                            weight_decay=0e-4, nlayers=2,
+                            nclass=data.nclass,
+                            device=self.device).to(self.device)
+            else:
+                if args.sgc == 1:
+                    model = SGC(nfeat=data.feat_train.shape[1], nhid=args.hidden,
+                                nclass=data.nclass, dropout=args.dropout,
+                                nlayers=args.nlayers, with_bn=False,
+                                device=self.device).to(self.device)
+                else:
+                    model = GCN(nfeat=data.feat_train.shape[1], nhid=args.hidden,
+                                nclass=data.nclass, dropout=args.dropout, nlayers=args.nlayers,
+                                device=self.device).to(self.device)
+
+
+            model.initialize()
+
+            model_parameters = list(model.parameters())
+
+            optimizer_model = torch.optim.Adam(model_parameters, lr=args.lr_model)
+            model.train()
+
+            for ol in range(outer_loop):
+                adj_syn = pge(self.feat_syn)
+                adj_syn_norm = utils.normalize_adj_tensor(adj_syn, sparse=False)
+                feat_syn_norm = feat_syn
+
+                BN_flag = False
+                for module in model.modules():
+                    if 'BatchNorm' in module._get_name(): #BatchNorm
+                        BN_flag = True
+                if BN_flag:
+                    model.train() # for updating the mu, sigma of BatchNorm
+                    output_real = model.forward(features, adj_norm)
+                    for module in model.modules():
+                        if 'BatchNorm' in module._get_name():  #BatchNorm
+                            module.eval() # fix mu and sigma of every BatchNorm layer
+
+                loss = torch.tensor(0.0).to(self.device)
+                for c in range(data.nclass):
+                    batch_size, n_id, adjs = data.retrieve_class_sampler(
+                            c, adj, transductive=True, args=args)
+                    if args.nlayers == 1:
+                        adjs = [adjs]
+
+                    adjs = [adj.to(self.device) for adj in adjs]
+                    output = model.forward_sampler(features[n_id], adjs)
+                    loss_real = F.nll_loss(output, labels[n_id[:batch_size]])
+
+                    gw_real = torch.autograd.grad(loss_real, model_parameters)
+                    gw_real = list((_.detach().clone() for _ in gw_real))
+                    output_syn = model.forward(feat_syn, adj_syn_norm)
+
+                    ind = syn_class_indices[c]
+                    loss_syn = F.nll_loss(
+                            output_syn[ind[0]: ind[1]],
+                            labels_syn[ind[0]: ind[1]])
+                    gw_syn = torch.autograd.grad(loss_syn, model_parameters, create_graph=True)
+                    coeff = self.num_class_dict[c] / max(self.num_class_dict.values())
+                    loss += coeff  * match_loss(gw_syn, gw_real, args, device=self.device)
+
+                loss_avg += loss.item()
+                # TODO: regularize
+                if args.alpha > 0:
+                    loss_reg = args.alpha * regularization(adj_syn, utils.tensor2onehot(labels_syn))
+                else:
+                    loss_reg = torch.tensor(0)
+
+                loss = loss + loss_reg
+
+                # update sythetic graph
+                self.optimizer_feat.zero_grad()
+                self.optimizer_pge.zero_grad()
+                loss.backward()
+                if it % 50 < 10:
+                    self.optimizer_pge.step()
+                else:
+                    self.optimizer_feat.step()
+
+                if args.debug and ol % 5 ==0:
+                    print('Gradient matching loss:', loss.item())
+
+                if ol == outer_loop - 1:
+                    # print('loss_reg:', loss_reg.item())
+                    # print('Gradient matching loss:', loss.item())
+                    break
+
+                feat_syn_inner = feat_syn.detach()
+                adj_syn_inner = pge.inference(feat_syn_inner)
+                adj_syn_inner_norm = utils.normalize_adj_tensor(adj_syn_inner, sparse=False)
+                feat_syn_inner_norm = feat_syn_inner
+                for j in range(inner_loop):
+                    optimizer_model.zero_grad()
+                    output_syn_inner = model.forward(feat_syn_inner_norm, adj_syn_inner_norm)
+                    loss_syn_inner = F.nll_loss(output_syn_inner, labels_syn)
+                    loss_syn_inner.backward()
+                    # print(loss_syn_inner.item())
+                    optimizer_model.step() # update gnn param
+
+
+            loss_avg /= (data.nclass*outer_loop)
+            if it % 50 == 0:
+                print('Epoch {}, loss_avg: {}'.format(it, loss_avg))
+
+            eval_epochs = [400, 600, 800, 1000, 1200, 1600, 2000, 3000, 4000, 5000]
+
+            if verbose and it in eval_epochs:
+            # if verbose and (it+1) % 50 == 0:
+                res = []
+                runs = 1 if args.dataset in ['ogbn-arxiv'] else 3
+                for i in range(runs):
+                    if args.dataset in ['ogbn-arxiv']:
+                        res.append(self.test_with_val())
+                    else:
+                        res.append(self.test_with_val())
+
+                res = np.array(res)
+                print('Train/Test Mean Accuracy:',
+                        repr([res.mean(0), res.std(0)]))
+
+    def get_sub_adj_feat(self, features):
+        data = self.data
+        args = self.args
+        idx_selected = []
+
+        from collections import Counter;
+        counter = Counter(self.labels_syn.cpu().numpy())
+
+        for c in range(data.nclass):
+            tmp = data.retrieve_class(c, num=counter[c])
+            tmp = list(tmp)
+            idx_selected = idx_selected + tmp
+        idx_selected = np.array(idx_selected).reshape(-1)
+        features = features[self.data.idx_train][idx_selected]
+
+        # adj_knn = torch.zeros((data.nclass*args.nsamples, data.nclass*args.nsamples)).to(self.device)
+        # for i in range(data.nclass):
+        #     idx = np.arange(i*args.nsamples, i*args.nsamples+args.nsamples)
+        #     adj_knn[np.ix_(idx, idx)] = 1
+
+        from sklearn.metrics.pairwise import cosine_similarity
+        # features[features!=0] = 1
+        k = 2
+        sims = cosine_similarity(features.cpu().numpy())
+        sims[(np.arange(len(sims)), np.arange(len(sims)))] = 0
+        for i in range(len(sims)):
+            indices_argsort = np.argsort(sims[i])
+            sims[i, indices_argsort[: -k]] = 0
+        adj_knn = torch.FloatTensor(sims).to(self.device)
+        return features, adj_knn
+
+
+def get_loops(args):
+    # Get the two hyper-parameters of outer-loop and inner-loop.
+    # The following values are empirically good.
+    if args.one_step:
+        if args.dataset =='ogbn-arxiv':
+            return 5, 0
+        return 1, 0
+    if args.dataset in ['ogbn-arxiv']:
+        return args.outer, args.inner
+    if args.dataset in ['cora']:
+        return 20, 15 # sgc
+    if args.dataset in ['citeseer']:
+        return 20, 15
+    if args.dataset in ['physics']:
+        return 20, 10
+    else:
+        return 20, 10
+

GCond/models/gat.py

@@ -0,0 +1,312 @@
+"""
+Extended from https://github.com/rusty1s/pytorch_geometric/tree/master/benchmark/citation
+"""
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from torch_geometric.nn import SGConv
+from torch_geometric.nn import APPNP as ModuleAPPNP
+# from torch_geometric.nn import GATConv
+from .mygatconv import GATConv
+import numpy as np
+import scipy.sparse as sp
+
+from torch.nn import Linear
+from itertools import repeat
+
+
+class GAT(torch.nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, heads=8, output_heads=1, dropout=0.5, lr=0.01,
+            weight_decay=5e-4, with_bias=True, device=None, **kwargs):
+
+        super(GAT, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.dropout = dropout
+        self.lr = lr
+        self.weight_decay = weight_decay
+
+        if 'dataset' in kwargs:
+            if kwargs['dataset'] in ['ogbn-arxiv']:
+                dropout = 0.7 # arxiv
+            elif kwargs['dataset'] in ['reddit']:
+                dropout = 0.05; self.dropout = 0.1; self.weight_decay = 5e-4
+                # self.weight_decay = 5e-2; dropout=0.05; self.dropout=0.1
+            elif kwargs['dataset'] in ['citeseer']:
+                dropout = 0.7
+                self.weight_decay = 5e-4
+            elif kwargs['dataset'] in ['flickr']:
+                dropout = 0.8
+                # nhid=8; heads=8
+                # self.dropout=0.1
+            else:
+                dropout = 0.7 # cora, citeseer, reddit
+        else:
+            dropout = 0.7
+        self.conv1 = GATConv(
+            nfeat,
+            nhid,
+            heads=heads,
+            dropout=dropout,
+            bias=with_bias)
+
+        self.conv2 = GATConv(
+            nhid * heads,
+            nclass,
+            heads=output_heads,
+            concat=False,
+            dropout=dropout,
+            bias=with_bias)
+
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+
+    # def forward(self, data):
+    #     x, edge_index = data.x, data.edge_index
+    #     x = F.dropout(x, p=self.dropout, training=self.training)
+    #     x = F.elu(self.conv1(x, edge_index))
+    #     x = F.dropout(x, p=self.dropout, training=self.training)
+    #     x = self.conv2(x, edge_index)
+    #     return F.log_softmax(x, dim=1)
+
+    def forward(self, data):
+        # x, edge_index = data.x, data.edge_index
+        x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = F.elu(self.conv1(x, edge_index, edge_weight=edge_weight))
+        # print(self.conv1.att_l.sum())
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.conv2(x, edge_index, edge_weight=edge_weight)
+        return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GAT.
+        """
+        self.conv1.reset_parameters()
+        self.conv2.reset_parameters()
+
+
+    def fit(self, feat, adj, labels, idx, data=None, train_iters=600, initialize=True, verbose=False, patience=None, noval=False, **kwargs):
+
+        data_train = GraphData(feat, adj, labels)
+        data_train = Dpr2Pyg(data_train)[0]
+
+        data_test = Dpr2Pyg(GraphData(data.feat_test, data.adj_test, None))[0]
+
+        if noval:
+            data_val = GraphData(data.feat_val, data.adj_val, None)
+            data_val = Dpr2Pyg(data_val)[0]
+        else:
+            data_val = GraphData(data.feat_full, data.adj_full, None)
+            data_val = Dpr2Pyg(data_val)[0]
+
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if initialize:
+            self.initialize()
+
+        if len(data_train.y.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+
+        data_train.y = data_train.y.float() if self.multi_label else data_train.y
+        # data_val.y = data_val.y.float() if self.multi_label else data_val.y
+
+        if verbose:
+            print('=== training gat model ===')
+
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+        best_acc_val = 0
+        best_loss_val = 100
+        for i in range(train_iters):
+            # if i == train_iters // 2:
+            if i in [1500]:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(data_train)
+            loss_train = self.loss(output, data_train.y)
+            loss_train.backward()
+            optimizer.step()
+
+            with torch.no_grad():
+                self.eval()
+
+                output = self.forward(data_val)
+                if noval:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+
+                if loss_val < best_loss_val:
+                    best_loss_val = loss_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+                    # print(best_acc_val)
+                    # output = self.forward(data_test)
+                    # labels_test = torch.LongTensor(data.labels_test).to(self.device)
+                    # loss_test = F.nll_loss(output, labels_test)
+                    # acc_test = utils.accuracy(output, labels_test)
+                    # print('acc_test:', acc_test.item())
+
+
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+    def test(self, data_test):
+        """Evaluate GCN performance
+        """
+        self.eval()
+        with torch.no_grad():
+            output = self.forward(data_test)
+        evaluate(output, data_test.y, self.args)
+
+    # @torch.no_grad()
+    # def predict(self, data):
+    #     self.eval()
+    #     return self.forward(data)
+    @torch.no_grad()
+    def predict(self, feat, adj):
+        self.eval()
+        data = GraphData(feat, adj, None)
+        data = Dpr2Pyg(data)[0]
+        return self.forward(data)
+
+    @torch.no_grad()
+    def predict_unnorm(self, feat, adj):
+        self.eval()
+        data = GraphData(feat, adj, None)
+        data = Dpr2Pyg(data)[0]
+
+        return self.forward(data)
+
+
+class GraphData:
+
+    def __init__(self, features, adj, labels, idx_train=None, idx_val=None, idx_test=None):
+        self.adj = adj
+        self.features = features
+        self.labels = labels
+        self.idx_train = idx_train
+        self.idx_val = idx_val
+        self.idx_test = idx_test
+
+
+from torch_geometric.data import InMemoryDataset, Data
+import scipy.sparse as sp
+
+class Dpr2Pyg(InMemoryDataset):
+
+    def __init__(self, dpr_data, transform=None, **kwargs):
+        root = 'data/' # dummy root; does not mean anything
+        self.dpr_data = dpr_data
+        super(Dpr2Pyg, self).__init__(root, transform)
+        pyg_data = self.process()
+        self.data, self.slices = self.collate([pyg_data])
+        self.transform = transform
+
+    def process____(self):
+        dpr_data = self.dpr_data
+        try:
+            edge_index = torch.LongTensor(dpr_data.adj.nonzero().cpu()).cuda().T
+        except:
+            edge_index = torch.LongTensor(dpr_data.adj.nonzero()).cuda()
+        # by default, the features in pyg data is dense
+        try:
+            x = torch.FloatTensor(dpr_data.features.cpu()).float().cuda()
+        except:
+            x = torch.FloatTensor(dpr_data.features).float().cuda()
+        try:
+            y = torch.LongTensor(dpr_data.labels.cpu()).cuda()
+        except:
+            y = dpr_data.labels
+
+        data = Data(x=x, edge_index=edge_index, y=y)
+        data.train_mask = None
+        data.val_mask = None
+        data.test_mask = None
+        return data
+
+    def process(self):
+        dpr_data = self.dpr_data
+        if type(dpr_data.adj) == torch.Tensor:
+            adj_selfloop = dpr_data.adj + torch.eye(dpr_data.adj.shape[0]).cuda()
+            edge_index_selfloop = adj_selfloop.nonzero().T
+            edge_index = edge_index_selfloop
+            edge_weight = adj_selfloop[edge_index_selfloop[0], edge_index_selfloop[1]]
+        else:
+            adj_selfloop = dpr_data.adj + sp.eye(dpr_data.adj.shape[0])
+            edge_index = torch.LongTensor(adj_selfloop.nonzero()).cuda()
+            edge_weight = torch.FloatTensor(adj_selfloop[adj_selfloop.nonzero()]).cuda()
+
+        # by default, the features in pyg data is dense
+        try:
+            x = torch.FloatTensor(dpr_data.features.cpu()).float().cuda()
+        except:
+            x = torch.FloatTensor(dpr_data.features).float().cuda()
+        try:
+            y = torch.LongTensor(dpr_data.labels.cpu()).cuda()
+        except:
+            y = dpr_data.labels
+
+
+        data = Data(x=x, edge_index=edge_index, y=y, edge_weight=edge_weight)
+        data.train_mask = None
+        data.val_mask = None
+        data.test_mask = None
+        return data
+
+    def get(self, idx):
+        data = self.data.__class__()
+
+        if hasattr(self.data, '__num_nodes__'):
+            data.num_nodes = self.data.__num_nodes__[idx]
+
+        for key in self.data.keys:
+            item, slices = self.data[key], self.slices[key]
+            s = list(repeat(slice(None), item.dim()))
+            s[self.data.__cat_dim__(key, item)] = slice(slices[idx],
+                                                   slices[idx + 1])
+            data[key] = item[s]
+        return data
+
+    @property
+    def raw_file_names(self):
+        return ['some_file_1', 'some_file_2', ...]
+
+    @property
+    def processed_file_names(self):
+        return ['data.pt']
+
+    def _download(self):
+        pass
+

GCond/models/gcn.py

@@ -0,0 +1,404 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+
+
+class GraphConvolution(Module):
+    """Simple GCN layer, similar to https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True):
+        super(GraphConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        self.bias = Parameter(torch.FloatTensor(out_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1. / math.sqrt(self.weight.T.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, adj):
+        """ Graph Convolutional Layer forward function
+        """
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight)
+        else:
+            support = torch.mm(input, self.weight)
+        if isinstance(adj, torch_sparse.SparseTensor):
+            output = torch_sparse.matmul(adj, support)
+        else:
+            output = torch.spmm(adj, support)
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+
+class GCN(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            with_relu=True, with_bias=True, with_bn=False, device=None):
+
+        super(GCN, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+
+        self.layers = nn.ModuleList([])
+
+        if nlayers == 1:
+            self.layers.append(GraphConvolution(nfeat, nclass, with_bias=with_bias))
+        else:
+            if with_bn:
+                self.bns = torch.nn.ModuleList()
+                self.bns.append(nn.BatchNorm1d(nhid))
+            self.layers.append(GraphConvolution(nfeat, nhid, with_bias=with_bias))
+            for i in range(nlayers-2):
+                self.layers.append(GraphConvolution(nhid, nhid, with_bias=with_bias))
+                if with_bn:
+                    self.bns.append(nn.BatchNorm1d(nhid))
+            self.layers.append(GraphConvolution(nhid, nclass, with_bias=with_bias))
+
+        self.dropout = dropout
+        self.lr = lr
+        if not with_relu:
+            self.weight_decay = 0
+        else:
+            self.weight_decay = weight_decay
+        self.with_relu = with_relu
+        self.with_bn = with_bn
+        self.with_bias = with_bias
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+
+    def forward(self, x, adj):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        # for ix, layer in enumerate(self.layers):
+        for ix, (adj, _, size) in enumerate(adjs):
+            x = self.layers[ix](x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        for ix, (adj) in enumerate(adjs):
+            x = self.layers[ix](x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        for layer in self.layers:
+            layer.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit(self, features, adj, labels, idx_train, idx_val=None, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, **kwargs):
+
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+
+        if idx_val is not None:
+            self._train_with_val2(labels, idx_train, idx_val, train_iters, verbose)
+        else:
+            self._train_without_val2(labels, idx_train, train_iters, verbose)
+
+    def _train_without_val2(self, labels, idx_train, train_iters, verbose):
+        self.train()
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output[idx_train], labels[idx_train])
+            loss_train.backward()
+            optimizer.step()
+            if verbose and i % 10 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+        self.eval()
+        output = self.forward(self.features, self.adj_norm)
+        self.output = output
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+        if noval:
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_acc_val = 0
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output, labels)
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            if utils.is_sparse_tensor(adj):
+                self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                self.adj_norm = utils.normalize_adj_tensor(adj)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+
+
+    def _train_with_val2(self, labels, idx_train, idx_val, train_iters, verbose):
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_loss_val = 100
+        best_acc_val = 0
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = F.nll_loss(output[idx_train], labels[idx_train])
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 10 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            self.eval()
+            output = self.forward(self.features, self.adj_norm)
+            loss_val = F.nll_loss(output[idx_val], labels[idx_val])
+            acc_val = utils.accuracy(output[idx_val], labels[idx_val])
+
+            if acc_val > best_acc_val:
+                best_acc_val = acc_val
+                self.output = output
+                weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)

GCond/models/myappnp.py

@@ -0,0 +1,344 @@
+"""multiple transformaiton and multiple propagation"""
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+
+
+class APPNP(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            ntrans=1, with_bias=True, with_bn=False, device=None):
+
+        super(APPNP, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+        self.alpha = 0.1
+
+        with_bn = False
+
+        self.layers = nn.ModuleList([])
+        if ntrans == 1:
+            self.layers.append(MyLinear(nfeat, nclass))
+        else:
+            self.layers.append(MyLinear(nfeat, nhid))
+            if with_bn:
+                self.bns = torch.nn.ModuleList()
+                self.bns.append(nn.BatchNorm1d(nhid))
+            for i in range(ntrans-2):
+                if with_bn:
+                    self.bns.append(nn.BatchNorm1d(nhid))
+                self.layers.append(MyLinear(nhid, nhid))
+            self.layers.append(MyLinear(nhid, nclass))
+
+        self.nlayers = nlayers
+        self.weight_decay = weight_decay
+        self.dropout = dropout
+        self.lr = lr
+        self.with_bn = with_bn
+        self.with_bias = with_bias
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+        self.sparse_dropout = SparseDropout(dprob=0)
+
+    def forward(self, x, adj):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        h = x
+        # here nlayers means K
+        for i in range(self.nlayers):
+            # adj_drop = self.sparse_dropout(adj, training=self.training)
+            adj_drop = adj
+            x = torch.spmm(adj_drop, x)
+            x = x * (1 - self.alpha)
+            x = x + self.alpha * h
+
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        h = x
+        for ix, (adj, _, size) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            # adj_drop = F.dropout(adj, p=self.dropout)
+            adj_drop = adj
+            h = h[: size[1]]
+            x = torch_sparse.matmul(adj_drop, x)
+            x = x * (1 - self.alpha)
+            x = x + self.alpha * h
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        for ix, (adj) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            x = torch_sparse.matmul(adj, x)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        for layer in self.layers:
+            layer.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+        if noval:
+            # self._train_without_val(labels, data, train_iters, verbose)
+            # self._train_without_val(labels, data, train_iters, verbose)
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_acc_val = 0
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output, labels)
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            if utils.is_sparse_tensor(adj):
+                self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                self.adj_norm = utils.normalize_adj_tensor(adj)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+
+
+
+class MyLinear(Module):
+    """Simple Linear layer, modified from https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True):
+        super(MyLinear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        if with_bias:
+            self.bias = Parameter(torch.FloatTensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 1. / math.sqrt(self.weight.T.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input):
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight)
+        else:
+            support = torch.mm(input, self.weight)
+        output = support
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+class SparseDropout(torch.nn.Module):
+    def __init__(self, dprob=0.5):
+        super(SparseDropout, self).__init__()
+        self.kprob=1-dprob
+
+    def forward(self, x, training):
+        if training:
+            mask=((torch.rand(x._values().size())+(self.kprob)).floor()).type(torch.bool)
+            rc=x._indices()[:,mask]
+            val=x._values()[mask]*(1.0/self.kprob)
+            return torch.sparse.FloatTensor(rc, val, x.size())
+        else:
+            return x

GCond/models/myappnp1.py

@@ -0,0 +1,348 @@
+"""multiple transformaiton and multiple propagation"""
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+
+
+class APPNP1(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            with_relu=True, with_bias=True, with_bn=False, device=None):
+
+        super(APPNP1, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+        self.alpha = 0.1
+
+        if with_bn:
+            self.bns = torch.nn.ModuleList()
+            self.bns.append(nn.BatchNorm1d(nhid))
+
+        self.layers = nn.ModuleList([])
+        # self.layers.append(MyLinear(nfeat, nclass))
+        self.layers.append(MyLinear(nfeat, nhid))
+        # self.layers.append(MyLinear(nhid, nhid))
+        self.layers.append(MyLinear(nhid, nclass))
+
+        # if nlayers == 1:
+        #     self.layers.append(nn.Linear(nfeat, nclass))
+        # else:
+        #     self.layers.append(nn.Linear(nfeat, nhid))
+        #     for i in range(nlayers-2):
+        #         self.layers.append(nn.Linear(nhid, nhid))
+        #     self.layers.append(nn.Linear(nhid, nclass))
+
+        self.nlayers = nlayers
+        self.dropout = dropout
+        self.lr = lr
+        if not with_relu:
+            self.weight_decay = 0
+        else:
+            self.weight_decay = weight_decay
+        self.with_relu = with_relu
+        self.with_bn = with_bn
+        self.with_bias = with_bias
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+        self.sparse_dropout = SparseDropout(dprob=0)
+
+    def forward(self, x, adj):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        h = x
+        # here nlayers means K
+        for i in range(self.nlayers):
+            # adj_drop = self.sparse_dropout(adj, training=self.training)
+            adj_drop = adj
+            x = torch.spmm(adj_drop, x)
+            x = x * (1 - self.alpha)
+            x = x + self.alpha * h
+
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        h = x
+        for ix, (adj, _, size) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            # adj_drop = F.dropout(adj, p=self.dropout)
+            adj_drop = adj
+            h = h[: size[1]]
+            x = torch_sparse.matmul(adj_drop, x)
+            x = x * (1 - self.alpha)
+            x = x + self.alpha * h
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        for ix, (adj) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            x = torch_sparse.matmul(adj, x)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        for layer in self.layers:
+            layer.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+
+        if noval:
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_acc_val = 0
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output, labels)
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            if utils.is_sparse_tensor(adj):
+                self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                self.adj_norm = utils.normalize_adj_tensor(adj)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+
+
+
+class MyLinear(Module):
+    """Simple Linear layer, modified from https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True):
+        super(MyLinear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        if with_bias:
+            self.bias = Parameter(torch.FloatTensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 1. / math.sqrt(self.weight.T.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input):
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight)
+        else:
+            support = torch.mm(input, self.weight)
+        output = support
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+class SparseDropout(torch.nn.Module):
+    def __init__(self, dprob=0.5):
+        super(SparseDropout, self).__init__()
+        self.kprob=1-dprob
+
+    def forward(self, x, training):
+        if training:
+            mask=((torch.rand(x._values().size())+(self.kprob)).floor()).type(torch.bool)
+            rc=x._indices()[:,mask]
+            val=x._values()[mask]*(1.0/self.kprob)
+            return torch.sparse.FloatTensor(rc, val, x.size())
+        else:
+            return x

GCond/models/mycheby.py

@@ -0,0 +1,417 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+from torch_geometric.nn.inits import zeros
+import scipy.sparse as sp
+import numpy as np
+
+
+class ChebConvolution(Module):
+    """Simple GCN layer, similar to https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True, single_param=True, K=2):
+        """set single_param to True to alleivate the overfitting issue"""
+        super(ChebConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.lins = torch.nn.ModuleList([
+           MyLinear(in_features, out_features, with_bias=False) for _ in range(K)])
+        # self.lins = torch.nn.ModuleList([
+        #    MyLinear(in_features, out_features, with_bias=True) for _ in range(K)])
+        if with_bias:
+            self.bias = Parameter(torch.Tensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.single_param = single_param
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for lin in self.lins:
+            lin.reset_parameters()
+        zeros(self.bias)
+
+    def forward(self, input, adj, size=None):
+        """ Graph Convolutional Layer forward function
+        """
+        # support = torch.mm(input, self.weight_l)
+        x = input
+        Tx_0 = x[:size[1]] if size is not None else x
+        Tx_1 = x # dummy
+        output = self.lins[0](Tx_0)
+
+        if len(self.lins) > 1:
+            if isinstance(adj, torch_sparse.SparseTensor):
+                Tx_1 = torch_sparse.matmul(adj, x)
+            else:
+                Tx_1 = torch.spmm(adj, x)
+
+            if self.single_param:
+                output = output + self.lins[0](Tx_1)
+            else:
+                output = output + self.lins[1](Tx_1)
+
+        for lin in self.lins[2:]:
+            if self.single_param:
+                lin = self.lins[0]
+            if isinstance(adj, torch_sparse.SparseTensor):
+                Tx_2 = torch_sparse.matmul(adj, Tx_1)
+            else:
+                Tx_2 = torch.spmm(adj, Tx_1)
+            Tx_2 = 2. * Tx_2 - Tx_0
+            output = output + lin.forward(Tx_2)
+            Tx_0, Tx_1 = Tx_1, Tx_2
+
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+
+class Cheby(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            with_relu=True, with_bias=True, with_bn=False, device=None):
+
+        super(Cheby, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+
+        self.layers = nn.ModuleList([])
+
+        if nlayers == 1:
+            self.layers.append(ChebConvolution(nfeat, nclass, with_bias=with_bias))
+        else:
+            if with_bn:
+                self.bns = torch.nn.ModuleList()
+                self.bns.append(nn.BatchNorm1d(nhid))
+            self.layers.append(ChebConvolution(nfeat, nhid, with_bias=with_bias))
+            for i in range(nlayers-2):
+                self.layers.append(ChebConvolution(nhid, nhid, with_bias=with_bias))
+                if with_bn:
+                    self.bns.append(nn.BatchNorm1d(nhid))
+            self.layers.append(ChebConvolution(nhid, nclass, with_bias=with_bias))
+
+        # self.lin = MyLinear(nhid, nclass, with_bias=True)
+
+        # dropout = 0.5
+        self.dropout = dropout
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.with_relu = with_relu
+        self.with_bn = with_bn
+        self.with_bias = with_bias
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+
+    def forward(self, x, adj):
+        for ix, layer in enumerate(self.layers):
+            # x = F.dropout(x, 0.2, training=self.training)
+            x  = layer(x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+                # x = F.dropout(x, 0.5, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        # TODO: do we need normalization?
+        # for ix, layer in enumerate(self.layers):
+        for ix, (adj, _, size) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            x = self.layers[ix](x, adj, size=size)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        for ix, (adj) in enumerate(adjs):
+            x = self.layers[ix](x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        for layer in self.layers:
+            layer.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        adj = adj - torch.eye(adj.shape[0]).to(self.device) # cheby
+        if normalize:
+            adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+        if noval:
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+        # adj_full = adj_full - sp.eye(adj_full.shape[0])
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_acc_val = 0
+        best_loss_val = 100
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output, labels)
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                # if loss_val < best_loss_val:
+                #     best_loss_val = loss_val
+                #     self.output = output
+                #     weights = deepcopy(self.state_dict())
+                #     print(best_loss_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+                    # print(best_acc_val)
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            # adj = adj-sp.eye(adj.shape[0])
+            # adj[0,0]=0
+
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            adj = utils.to_scipy(adj)
+
+            adj = adj-sp.eye(adj.shape[0])
+            mx = normalize_adj(adj)
+            adj = utils.sparse_mx_to_torch_sparse_tensor(mx).to(self.device)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+
+class MyLinear(Module):
+    """Simple Linear layer, modified from https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True):
+        super(MyLinear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        if with_bias:
+            self.bias = Parameter(torch.FloatTensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 1. / math.sqrt(self.weight.T.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input):
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight)
+        else:
+            support = torch.mm(input, self.weight)
+        output = support
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+
+
+def normalize_adj(mx):
+    """Normalize sparse adjacency matrix,
+    A' = (D + I)^-1/2 * ( A + I ) * (D + I)^-1/2
+    Row-normalize sparse matrix
+    Parameters
+    ----------
+    mx : scipy.sparse.csr_matrix
+        matrix to be normalized
+    Returns
+    -------
+    scipy.sprase.lil_matrix
+        normalized matrix
+    """
+
+    # TODO: maybe using coo format would be better?
+    if type(mx) is not sp.lil.lil_matrix:
+        mx = mx.tolil()
+    mx = mx + sp.eye(mx.shape[0])
+    rowsum = np.array(mx.sum(1))
+    r_inv = np.power(rowsum, -1/2).flatten()
+    r_inv[np.isinf(r_inv)] = 0.
+    r_mat_inv = sp.diags(r_inv)
+    mx = r_mat_inv.dot(mx)
+    mx = mx.dot(r_mat_inv)
+    return mx

GCond/models/mygatconv.py

@@ -0,0 +1,203 @@
+from typing import Union, Tuple, Optional
+from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
+                                    OptTensor)
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+from torch.nn import Parameter, Linear
+from torch_sparse import SparseTensor, set_diag
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import remove_self_loops, add_self_loops, softmax
+
+from torch_geometric.nn.inits import glorot, zeros
+
+
+class GATConv(MessagePassing):
+    r"""The graph attentional operator from the `"Graph Attention Networks"
+    <https://arxiv.org/abs/1710.10903>`_ paper
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}\mathbf{x}_{i} +
+        \sum_{j \in \mathcal{N}(i)} \alpha_{i,j}\mathbf{\Theta}\mathbf{x}_{j},
+
+    where the attention coefficients :math:`\alpha_{i,j}` are computed as
+
+    .. math::
+        \alpha_{i,j} =
+        \frac{
+        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
+        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_j]
+        \right)\right)}
+        {\sum_{k \in \mathcal{N}(i) \cup \{ i \}}
+        \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top}
+        [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_k]
+        \right)\right)}.
+
+    Args:
+        in_channels (int or tuple): Size of each input sample. A tuple
+            corresponds to the sizes of source and target dimensionalities.
+        out_channels (int): Size of each output sample.
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        concat (bool, optional): If set to :obj:`False`, the multi-head
+            attentions are averaged instead of concatenated.
+            (default: :obj:`True`)
+        negative_slope (float, optional): LeakyReLU angle of the negative
+            slope. (default: :obj:`0.2`)
+        dropout (float, optional): Dropout probability of the normalized
+            attention coefficients which exposes each node to a stochastically
+            sampled neighborhood during training. (default: :obj:`0`)
+        add_self_loops (bool, optional): If set to :obj:`False`, will not add
+            self-loops to the input graph. (default: :obj:`True`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    _alpha: OptTensor
+
+    def __init__(self, in_channels: Union[int, Tuple[int, int]],
+                 out_channels: int, heads: int = 1, concat: bool = True,
+                 negative_slope: float = 0.2, dropout: float = 0.0,
+                 add_self_loops: bool = True, bias: bool = True, **kwargs):
+        kwargs.setdefault('aggr', 'add')
+        super(GATConv, self).__init__(node_dim=0, **kwargs)
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.dropout = dropout
+        self.add_self_loops = add_self_loops
+
+        if isinstance(in_channels, int):
+            self.lin_l = Linear(in_channels, heads * out_channels, bias=False)
+            self.lin_r = self.lin_l
+        else:
+            self.lin_l = Linear(in_channels[0], heads * out_channels, False)
+            self.lin_r = Linear(in_channels[1], heads * out_channels, False)
+
+        self.att_l = Parameter(torch.Tensor(1, heads, out_channels))
+        self.att_r = Parameter(torch.Tensor(1, heads, out_channels))
+
+        if bias and concat:
+            self.bias = Parameter(torch.Tensor(heads * out_channels))
+        elif bias and not concat:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self._alpha = None
+
+        self.reset_parameters()
+        self.edge_weight = None
+
+    def reset_parameters(self):
+        glorot(self.lin_l.weight)
+        glorot(self.lin_r.weight)
+        glorot(self.att_l)
+        glorot(self.att_r)
+        zeros(self.bias)
+
+    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
+                size: Size = None, return_attention_weights=None, edge_weight=None):
+        # type: (Union[Tensor, OptPairTensor], Tensor, Size, NoneType) -> Tensor  # noqa
+        # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, NoneType) -> Tensor  # noqa
+        # type: (Union[Tensor, OptPairTensor], Tensor, Size, bool) -> Tuple[Tensor, Tuple[Tensor, Tensor]]  # noqa
+        # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, bool) -> Tuple[Tensor, SparseTensor]  # noqa
+        r"""
+        Args:
+            return_attention_weights (bool, optional): If set to :obj:`True`,
+                will additionally return the tuple
+                :obj:`(edge_index, attention_weights)`, holding the computed
+                attention weights for each edge. (default: :obj:`None`)
+        """
+        H, C = self.heads, self.out_channels
+
+        x_l: OptTensor = None
+        x_r: OptTensor = None
+        alpha_l: OptTensor = None
+        alpha_r: OptTensor = None
+        if isinstance(x, Tensor):
+            assert x.dim() == 2, 'Static graphs not supported in `GATConv`.'
+            x_l = x_r = self.lin_l(x).view(-1, H, C)
+            alpha_l = (x_l * self.att_l).sum(dim=-1)
+            alpha_r = (x_r * self.att_r).sum(dim=-1)
+        else:
+            x_l, x_r = x[0], x[1]
+            assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.'
+            x_l = self.lin_l(x_l).view(-1, H, C)
+            alpha_l = (x_l * self.att_l).sum(dim=-1)
+            if x_r is not None:
+                x_r = self.lin_r(x_r).view(-1, H, C)
+                alpha_r = (x_r * self.att_r).sum(dim=-1)
+
+
+        assert x_l is not None
+        assert alpha_l is not None
+
+        if self.add_self_loops:
+            if isinstance(edge_index, Tensor):
+                num_nodes = x_l.size(0)
+                if x_r is not None:
+                    num_nodes = min(num_nodes, x_r.size(0))
+                if size is not None:
+                    num_nodes = min(size[0], size[1])
+                edge_index, _ = remove_self_loops(edge_index)
+                edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+
+                if edge_weight is not None:
+                    if self.edge_weight is None:
+                        self.edge_weight = edge_weight
+
+                    if edge_index.size(1) != self.edge_weight.shape[0]:
+                        self.edge_weight = edge_weight
+
+            elif isinstance(edge_index, SparseTensor):
+                edge_index = set_diag(edge_index)
+
+        # propagate_type: (x: OptPairTensor, alpha: OptPairTensor)
+
+        out = self.propagate(edge_index, x=(x_l, x_r),
+                             alpha=(alpha_l, alpha_r), size=size)
+
+        alpha = self._alpha
+        self._alpha = None
+
+        if self.concat:
+            out = out.view(-1, self.heads * self.out_channels)
+        else:
+            out = out.mean(dim=1)
+
+        if self.bias is not None:
+            out += self.bias
+
+        if isinstance(return_attention_weights, bool):
+            assert alpha is not None
+            if isinstance(edge_index, Tensor):
+                return out, (edge_index, alpha)
+            elif isinstance(edge_index, SparseTensor):
+                return out, edge_index.set_value(alpha, layout='coo')
+        else:
+            return out
+
+    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
+                index: Tensor, ptr: OptTensor,
+                size_i: Optional[int]) -> Tensor:
+        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
+        alpha = F.leaky_relu(alpha, self.negative_slope)
+        alpha = softmax(alpha, index, ptr, size_i)
+        self._alpha = alpha
+        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+
+        if self.edge_weight is not None:
+            x_j = self.edge_weight.view(-1, 1, 1) * x_j
+        return x_j * alpha.unsqueeze(-1)
+
+    def __repr__(self):
+        return '{}({}, {}, heads={})'.format(self.__class__.__name__,
+                                             self.in_channels,
+                                             self.out_channels, self.heads)
+

GCond/models/mygraphsage.py

@@ -0,0 +1,353 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+from torch_geometric.data import NeighborSampler
+from torch_sparse import SparseTensor
+
+
+class SageConvolution(Module):
+    """Simple GCN layer, similar to https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True, root_weight=False):
+        super(SageConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight_l = Parameter(torch.FloatTensor(in_features, out_features))
+        self.bias_l = Parameter(torch.FloatTensor(out_features))
+        self.weight_r = Parameter(torch.FloatTensor(in_features, out_features))
+        self.bias_r = Parameter(torch.FloatTensor(out_features))
+        self.reset_parameters()
+        self.root_weight = root_weight
+        # self.weight = Parameter(torch.FloatTensor(out_features, in_features))
+        # self.linear = torch.nn.Linear(self.in_features, self.out_features)
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 1. / math.sqrt(self.weight_l.T.size(1))
+        self.weight_l.data.uniform_(-stdv, stdv)
+        self.bias_l.data.uniform_(-stdv, stdv)
+
+        stdv = 1. / math.sqrt(self.weight_r.T.size(1))
+        self.weight_r.data.uniform_(-stdv, stdv)
+        self.bias_r.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, adj, size=None):
+        """ Graph Convolutional Layer forward function
+        """
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight_l)
+        else:
+            support = torch.mm(input, self.weight_l)
+        if isinstance(adj, torch_sparse.SparseTensor):
+            output = torch_sparse.matmul(adj, support)
+        else:
+            output = torch.spmm(adj, support)
+        output = output + self.bias_l
+
+        if self.root_weight:
+            if size is not None:
+                output = output + input[:size[1]] @ self.weight_r + self.bias_r
+            else:
+                output = output + input @ self.weight_r + self.bias_r
+        else:
+            output = output
+
+        return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+
+class GraphSage(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            with_relu=True, with_bias=True, with_bn=False, device=None):
+
+        super(GraphSage, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+
+        self.layers = nn.ModuleList([])
+
+        if nlayers == 1:
+            self.layers.append(SageConvolution(nfeat, nclass, with_bias=with_bias))
+        else:
+            if with_bn:
+                self.bns = torch.nn.ModuleList()
+                self.bns.append(nn.BatchNorm1d(nhid))
+            self.layers.append(SageConvolution(nfeat, nhid, with_bias=with_bias))
+            for i in range(nlayers-2):
+                self.layers.append(SageConvolution(nhid, nhid, with_bias=with_bias))
+                if with_bn:
+                    self.bns.append(nn.BatchNorm1d(nhid))
+            self.layers.append(SageConvolution(nhid, nclass, with_bias=with_bias))
+
+        self.dropout = dropout
+        self.lr = lr
+        if not with_relu:
+            self.weight_decay = 0
+        else:
+            self.weight_decay = weight_decay
+        self.with_relu = with_relu
+        self.with_bn = with_bn
+        self.with_bias = with_bias
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+
+    def forward(self, x, adj):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        # TODO: do we need normalization?
+        # for ix, layer in enumerate(self.layers):
+        for ix, (adj, _, size) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            x = self.layers[ix](x, adj, size=size)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        for ix, (adj) in enumerate(adjs):
+            x = self.layers[ix](x, adj)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                if self.with_relu:
+                    x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        for layer in self.layers:
+            layer.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+        if noval:
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        adj_norm = self.adj_norm
+        node_idx = torch.arange(adj_norm.shape[0]).long()
+
+        edge_index = adj_norm.nonzero().T
+        adj_norm = SparseTensor(row=edge_index[0], col=edge_index[1],
+                value=adj_norm[edge_index[0], edge_index[1]], sparse_sizes=adj_norm.size()).t()
+        # edge_index = adj_norm._indices()
+        # adj_norm = SparseTensor(row=edge_index[0], col=edge_index[1],
+        #         value=adj_norm._values(), sparse_sizes=adj_norm.size()).t()
+
+        if adj_norm.density() > 0.5: # if the weighted graph is too dense, we need a larger neighborhood size
+            sizes = [30, 20]
+        else:
+            sizes = [5, 5]
+        train_loader = NeighborSampler(adj_norm,
+                            node_idx=node_idx,
+                            sizes=sizes, batch_size=len(node_idx),
+                            num_workers=0, return_e_id=False,
+                            num_nodes=adj_norm.size(0),
+                            shuffle=True)
+
+        best_acc_val = 0
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            # optimizer.zero_grad()
+            # output = self.forward(self.features, self.adj_norm)
+            # loss_train = self.loss(output, labels)
+            # loss_train.backward()
+            # optimizer.step()
+
+            for batch_size, n_id, adjs in train_loader:
+                adjs = [adj.to(self.device) for adj in adjs]
+                optimizer.zero_grad()
+                out = self.forward_sampler(self.features[n_id], adjs)
+                loss_train = F.nll_loss(out, labels[n_id[:batch_size]])
+                loss_train.backward()
+                optimizer.step()
+
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            if utils.is_sparse_tensor(adj):
+                self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                self.adj_norm = utils.normalize_adj_tensor(adj)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+

GCond/models/parametrized_adj.py

@@ -0,0 +1,88 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from itertools import product
+import numpy as np
+
+class PGE(nn.Module):
+
+    def __init__(self, nfeat, nnodes, nhid=128, nlayers=3, device=None, args=None):
+        super(PGE, self).__init__()
+        if args.dataset in ['ogbn-arxiv', 'arxiv', 'flickr']:
+           nhid = 256
+        if args.dataset in ['reddit']:
+           nhid = 256
+           if args.reduction_rate==0.01:
+               nhid = 128
+           nlayers = 3
+           # nhid = 128
+
+        self.layers = nn.ModuleList([])
+        self.layers.append(nn.Linear(nfeat*2, nhid))
+        self.bns = torch.nn.ModuleList()
+        self.bns.append(nn.BatchNorm1d(nhid))
+        for i in range(nlayers-2):
+            self.layers.append(nn.Linear(nhid, nhid))
+            self.bns.append(nn.BatchNorm1d(nhid))
+        self.layers.append(nn.Linear(nhid, 1))
+
+        edge_index = np.array(list(product(range(nnodes), range(nnodes))))
+        self.edge_index = edge_index.T
+        self.nnodes = nnodes
+        self.device = device
+        self.reset_parameters()
+        self.cnt = 0
+        self.args = args
+        self.nnodes = nnodes
+
+    def forward(self, x, inference=False):
+        if self.args.dataset == 'reddit' and self.args.reduction_rate >= 0.01:
+            edge_index = self.edge_index
+            n_part = 5
+            splits = np.array_split(np.arange(edge_index.shape[1]), n_part)
+            edge_embed = []
+            for idx in splits:
+                tmp_edge_embed = torch.cat([x[edge_index[0][idx]],
+                        x[edge_index[1][idx]]], axis=1)
+                for ix, layer in enumerate(self.layers):
+                    tmp_edge_embed = layer(tmp_edge_embed)
+                    if ix != len(self.layers) - 1:
+                        tmp_edge_embed = self.bns[ix](tmp_edge_embed)
+                        tmp_edge_embed = F.relu(tmp_edge_embed)
+                edge_embed.append(tmp_edge_embed)
+            edge_embed = torch.cat(edge_embed)
+        else:
+            edge_index = self.edge_index
+            edge_embed = torch.cat([x[edge_index[0]],
+                    x[edge_index[1]]], axis=1)
+            for ix, layer in enumerate(self.layers):
+                edge_embed = layer(edge_embed)
+                if ix != len(self.layers) - 1:
+                    edge_embed = self.bns[ix](edge_embed)
+                    edge_embed = F.relu(edge_embed)
+
+        adj = edge_embed.reshape(self.nnodes, self.nnodes)
+
+        adj = (adj + adj.T)/2
+        adj = torch.sigmoid(adj)
+        adj = adj - torch.diag(torch.diag(adj, 0))
+        return adj
+
+    @torch.no_grad()
+    def inference(self, x):
+        # self.eval()
+        adj_syn = self.forward(x, inference=True)
+        return adj_syn
+
+    def reset_parameters(self):
+        def weight_reset(m):
+            if isinstance(m, nn.Linear):
+                m.reset_parameters()
+            if isinstance(m, nn.BatchNorm1d):
+                m.reset_parameters()
+        self.apply(weight_reset)
+

GCond/models/sgc.py

@@ -0,0 +1,290 @@
+'''one transformation with multiple propagation'''
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+
+class GraphConvolution(Module):
+    """Simple GCN layer, similar to https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True):
+        super(GraphConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        self.bias = Parameter(torch.FloatTensor(out_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 1. / math.sqrt(self.weight.T.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, adj):
+        """ Graph Convolutional Layer forward function
+        """
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight)
+        else:
+            support = torch.mm(input, self.weight)
+        if isinstance(adj, torch_sparse.SparseTensor):
+            output = torch_sparse.matmul(adj, support)
+        else:
+            output = torch.spmm(adj, support)
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+
+class SGC(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            with_relu=True, with_bias=True, with_bn=False, device=None):
+
+        super(SGC, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+
+        self.conv = GraphConvolution(nfeat, nclass, with_bias=with_bias)
+
+        self.nlayers = nlayers
+        self.dropout = dropout
+        self.lr = lr
+        if not with_relu:
+            self.weight_decay = 0
+        else:
+            self.weight_decay = weight_decay
+        self.with_relu = with_relu
+        if with_bn:
+            print('Warning: SGC does not have bn!!!')
+        self.with_bn = False
+        self.with_bias = with_bias
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+
+    def forward(self, x, adj):
+        weight = self.conv.weight
+        bias = self.conv.bias
+        x = torch.mm(x, weight)
+        for i in range(self.nlayers):
+            x = torch.spmm(adj, x)
+        x = x + bias
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        weight = self.conv.weight
+        bias = self.conv.bias
+        x = torch.mm(x, weight)
+        for ix, (adj, _, size) in enumerate(adjs):
+            x = torch_sparse.matmul(adj, x)
+        x = x + bias
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        weight = self.conv.weight
+        bias = self.conv.bias
+        x = torch.mm(x, weight)
+        for ix, (adj) in enumerate(adjs):
+            if type(adj) == torch.Tensor:
+                x = adj @ x
+            else:
+                x = torch_sparse.matmul(adj, x)
+        x = x + bias
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        self.conv.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+        if noval:
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_acc_val = 0
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output, labels)
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            if utils.is_sparse_tensor(adj):
+                self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                self.adj_norm = utils.normalize_adj_tensor(adj)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+
+

GCond/models/sgc_multi.py

@@ -0,0 +1,315 @@
+"""multiple transformaiton and multiple propagation"""
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import torch
+import torch.optim as optim
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from deeprobust.graph import utils
+from copy import deepcopy
+from sklearn.metrics import f1_score
+from torch.nn import init
+import torch_sparse
+
+
+class SGC(nn.Module):
+
+    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0.5, lr=0.01, weight_decay=5e-4,
+            ntrans=2, with_bias=True, with_bn=False, device=None):
+
+        """nlayers indicates the number of propagations"""
+        super(SGC, self).__init__()
+
+        assert device is not None, "Please specify 'device'!"
+        self.device = device
+        self.nfeat = nfeat
+        self.nclass = nclass
+
+
+        self.layers = nn.ModuleList([])
+        if ntrans == 1:
+            self.layers.append(MyLinear(nfeat, nclass))
+        else:
+            self.layers.append(MyLinear(nfeat, nhid))
+            if with_bn:
+                self.bns = torch.nn.ModuleList()
+                self.bns.append(nn.BatchNorm1d(nhid))
+            for i in range(ntrans-2):
+                if with_bn:
+                    self.bns.append(nn.BatchNorm1d(nhid))
+                self.layers.append(MyLinear(nhid, nhid))
+            self.layers.append(MyLinear(nhid, nclass))
+
+        self.nlayers = nlayers
+        self.dropout = dropout
+        self.lr = lr
+        self.with_bn = with_bn
+        self.with_bias = with_bias
+        self.weight_decay = weight_decay
+        self.output = None
+        self.best_model = None
+        self.best_output = None
+        self.adj_norm = None
+        self.features = None
+        self.multi_label = None
+
+    def forward(self, x, adj):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        for i in range(self.nlayers):
+            x = torch.spmm(adj, x)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler(self, x, adjs):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        for ix, (adj, _, size) in enumerate(adjs):
+            # x_target = x[: size[1]]
+            # x = self.layers[ix]((x, x_target), edge_index)
+            # adj = adj.to(self.device)
+            x = torch_sparse.matmul(adj, x)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+    def forward_sampler_syn(self, x, adjs):
+        for ix, layer in enumerate(self.layers):
+            x = layer(x)
+            if ix != len(self.layers) - 1:
+                x = self.bns[ix](x) if self.with_bn else x
+                x = F.relu(x)
+                x = F.dropout(x, self.dropout, training=self.training)
+
+        for ix, (adj) in enumerate(adjs):
+            if type(adj) == torch.Tensor:
+                x = adj @ x
+            else:
+                x = torch_sparse.matmul(adj, x)
+
+        if self.multi_label:
+            return torch.sigmoid(x)
+        else:
+            return F.log_softmax(x, dim=1)
+
+
+    def initialize(self):
+        """Initialize parameters of GCN.
+        """
+        for layer in self.layers:
+            layer.reset_parameters()
+        if self.with_bn:
+            for bn in self.bns:
+                bn.reset_parameters()
+
+    def fit_with_val(self, features, adj, labels, data, train_iters=200, initialize=True, verbose=False, normalize=True, patience=None, noval=False, **kwargs):
+        '''data: full data class'''
+        if initialize:
+            self.initialize()
+
+        # features, adj, labels = data.feat_train, data.adj_train, data.labels_train
+        if type(adj) is not torch.Tensor:
+            features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device)
+        else:
+            features = features.to(self.device)
+            adj = adj.to(self.device)
+            labels = labels.to(self.device)
+
+        if normalize:
+            if utils.is_sparse_tensor(adj):
+                adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                adj_norm = utils.normalize_adj_tensor(adj)
+        else:
+            adj_norm = adj
+
+        if 'feat_norm' in kwargs and kwargs['feat_norm']:
+            from utils import row_normalize_tensor
+            features = row_normalize_tensor(features-features.min())
+
+        self.adj_norm = adj_norm
+        self.features = features
+
+        if len(labels.shape) > 1:
+            self.multi_label = True
+            self.loss = torch.nn.BCELoss()
+        else:
+            self.multi_label = False
+            self.loss = F.nll_loss
+
+        labels = labels.float() if self.multi_label else labels
+        self.labels = labels
+
+        if noval:
+            self._train_with_val(labels, data, train_iters, verbose, adj_val=True)
+        else:
+            self._train_with_val(labels, data, train_iters, verbose)
+
+    def _train_with_val(self, labels, data, train_iters, verbose, adj_val=False):
+        if adj_val:
+            feat_full, adj_full = data.feat_val, data.adj_val
+        else:
+            feat_full, adj_full = data.feat_full, data.adj_full
+
+        feat_full, adj_full = utils.to_tensor(feat_full, adj_full, device=self.device)
+        adj_full_norm = utils.normalize_adj_tensor(adj_full, sparse=True)
+        labels_val = torch.LongTensor(data.labels_val).to(self.device)
+
+        if verbose:
+            print('=== training gcn model ===')
+        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        best_acc_val = 0
+
+        for i in range(train_iters):
+            if i == train_iters // 2:
+                lr = self.lr*0.1
+                optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=self.weight_decay)
+
+            self.train()
+            optimizer.zero_grad()
+            output = self.forward(self.features, self.adj_norm)
+            loss_train = self.loss(output, labels)
+            loss_train.backward()
+            optimizer.step()
+
+            if verbose and i % 100 == 0:
+                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))
+
+            with torch.no_grad():
+                self.eval()
+                output = self.forward(feat_full, adj_full_norm)
+                if adj_val:
+                    loss_val = F.nll_loss(output, labels_val)
+                    acc_val = utils.accuracy(output, labels_val)
+                else:
+                    loss_val = F.nll_loss(output[data.idx_val], labels_val)
+                    acc_val = utils.accuracy(output[data.idx_val], labels_val)
+
+                if acc_val > best_acc_val:
+                    best_acc_val = acc_val
+                    self.output = output
+                    weights = deepcopy(self.state_dict())
+
+        if verbose:
+            print('=== picking the best model according to the performance on validation ===')
+        self.load_state_dict(weights)
+
+
+    def test(self, idx_test):
+        """Evaluate GCN performance on test set.
+        Parameters
+        ----------
+        idx_test :
+            node testing indices
+        """
+        self.eval()
+        output = self.predict()
+        # output = self.output
+        loss_test = F.nll_loss(output[idx_test], self.labels[idx_test])
+        acc_test = utils.accuracy(output[idx_test], self.labels[idx_test])
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+        return acc_test.item()
+
+
+    @torch.no_grad()
+    def predict(self, features=None, adj=None):
+        """By default, the inputs should be unnormalized adjacency
+        Parameters
+        ----------
+        features :
+            node features. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        adj :
+            adjcency matrix. If `features` and `adj` are not given, this function will use previous stored `features` and `adj` from training to make predictions.
+        Returns
+        -------
+        torch.FloatTensor
+            output (log probabilities) of GCN
+        """
+
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            if utils.is_sparse_tensor(adj):
+                self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True)
+            else:
+                self.adj_norm = utils.normalize_adj_tensor(adj)
+            return self.forward(self.features, self.adj_norm)
+
+    @torch.no_grad()
+    def predict_unnorm(self, features=None, adj=None):
+        self.eval()
+        if features is None and adj is None:
+            return self.forward(self.features, self.adj_norm)
+        else:
+            if type(adj) is not torch.Tensor:
+                features, adj = utils.to_tensor(features, adj, device=self.device)
+
+            self.features = features
+            self.adj_norm = adj
+            return self.forward(self.features, self.adj_norm)
+
+
+
+class MyLinear(Module):
+    """Simple Linear layer, modified from https://github.com/tkipf/pygcn
+    """
+
+    def __init__(self, in_features, out_features, with_bias=True):
+        super(MyLinear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
+        if with_bias:
+            self.bias = Parameter(torch.FloatTensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        stdv = 1. / math.sqrt(self.weight.T.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, input):
+        if input.data.is_sparse:
+            support = torch.spmm(input, self.weight)
+        else:
+            support = torch.mm(input, self.weight)
+        output = support
+        if self.bias is not None:
+            return output + self.bias
+        else:
+            return output
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + str(self.in_features) + ' -> ' \
+               + str(self.out_features) + ')'
+
+

GCond/requirements.txt

@@ -0,0 +1,11 @@
+torch
+torch_geometric
+scipy
+numpy
+ogb
+tqdm
+torch_sparse
+torch_vision
+configs
+deeprobust
+scikit_learn

GCond/res/cross/empty

@@ -0,0 +1 @@
+

GCond/script.sh

@@ -0,0 +1,4 @@
+python train_cond_tranduct_sampler.py --dataset cora --mlp=0 --nlayers=2 --sgc=1 --lr_feat=1e-4 --gpu_id=0  --lr_adj=1e-4 --r=0.5  --seed=1000
+
+python train_cond_tranduct_sampler.py --dataset ogbn-arxiv --mlp=0 --nlayers=2 --sgc=1 --lr_feat=0.01 --gpu_id=3  --lr_adj=0.01 --r=0.02   --seed=0 --inner=3  --epochs=1000  --save=0
+

GCond/scripts/run_cross.sh

@@ -0,0 +1,16 @@
+for r in 0.001 0.005 0.01
+do
+    bash scripts/script_cross.sh flickr ${r} 0  
+    bash scripts/script_cross.sh ogbn-arxiv ${r} 0 
+done
+
+for r in 0.25 0.5 1
+do 
+    bash scripts/script_cross.sh citeseer ${r} 0  
+    bash scripts/script_cross.sh cora ${r} 0  
+done
+
+for r in 0.001 0.0005 0.002
+do     
+    bash scripts/script_cross.sh reddit ${r} 0 
+done

GCond/scripts/run_main.sh

@@ -0,0 +1,26 @@
+for r in 0.25 0.5 1
+do
+python train_gcond_transduct.py --dataset cora --nlayers=2 --sgc=1 --lr_feat=1e-4 --gpu_id=0  --lr_adj=1e-4 --r=${r}  --seed=1 --epoch=600 --save=0
+done
+
+for r in 0.25 0.5 1
+do
+python train_gcond_transduct.py --dataset citeseer --nlayers=2 --sgc=1 --lr_feat=1e-4 --gpu_id=0  --lr_adj=1e-4 --r=${r}  --seed=1 --epoch=600 --save=0
+done
+
+
+for r in 0.001 0.005 0.01
+do
+python train_gcond_transduct.py --dataset ogbn-arxiv --nlayers=2 --sgc=1 --lr_feat=0.01 --gpu_id=3  --lr_adj=0.01 --r=${r}  --seed=1 --inner=3  --epochs=1000  --save=0
+done
+
+
+for r in 0.001 0.005 0.01
+do
+    python train_gcond_induct.py --dataset flickr --sgc=2 --nlayers=2 --lr_feat=0.005 --lr_adj=0.005  --r=${r} --seed=1 --gpu_id=0 --epochs=1000  --inner=1 --outer=10 --save=0 
+done
+
+for r in 0.001 0.005 0.0005 0.002
+do
+    python train_gcond_induct.py --dataset reddit --sgc=1 --nlayers=2 --lr_feat=0.1 --lr_adj=0.1  --r=${r} --seed=1 --gpu_id=0 --epochs=1000  --inner=1 --outer=10 --save=0 
+done

GCond/scripts/script_cross.sh

@@ -0,0 +1,7 @@
+dataset=${1}
+r=${2}
+gpu_id=${3}
+for s in 0 1 2 3 4
+do
+python test_other_arcs.py --dataset ${dataset} --gpu_id=${gpu_id} --r=${r} --seed=${s} --nruns=10  >> res/flickr/${1}_${2}.out 
+done

GCond/test_other_arcs.py

@@ -0,0 +1,55 @@
+import sys
+from deeprobust.graph.data import Dataset
+import numpy as np
+import random
+import time
+import argparse
+import torch
+from utils import *
+import torch.nn.functional as F
+from tester_other_arcs import Evaluator
+from utils_graphsaint import DataGraphSAINT
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu_id', type=int, default=0, help='gpu id')
+parser.add_argument('--dataset', type=str, default='cora')
+parser.add_argument('--nlayers', type=int, default=2)
+parser.add_argument('--hidden', type=int, default=256)
+parser.add_argument('--keep_ratio', type=float, default=1)
+parser.add_argument('--reduction_rate', type=float, default=1)
+parser.add_argument('--weight_decay', type=float, default=0.0)
+parser.add_argument('--dropout', type=float, default=0.0)
+parser.add_argument('--normalize_features', type=bool, default=True)
+parser.add_argument('--seed', type=int, default=0, help='Random seed.')
+parser.add_argument('--mlp', type=int, default=0)
+parser.add_argument('--inner', type=int, default=0)
+parser.add_argument('--epsilon', type=float, default=-1)
+parser.add_argument('--nruns', type=int, default=20)
+args = parser.parse_args()
+
+torch.cuda.set_device(args.gpu_id)
+
+# random seed setting
+random.seed(args.seed)
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed(args.seed)
+
+if args.dataset in ['cora', 'citeseer']:
+    args.epsilon = 0.05
+else:
+    args.epsilon = 0.01
+
+print(args)
+
+data_graphsaint = ['flickr', 'reddit', 'ogbn-arxiv']
+if args.dataset in data_graphsaint:
+    data = DataGraphSAINT(args.dataset)
+    data_full = data.data_full
+else:
+    data_full = get_dataset(args.dataset, args.normalize_features)
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+
+agent = Evaluator(data, args, device='cuda')
+agent.train()

GCond/tester_other_arcs.py

@@ -0,0 +1,258 @@
+import sys
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.nn import Parameter
+import torch.nn.functional as F
+from utils import match_loss, regularization, row_normalize_tensor
+import deeprobust.graph.utils as utils
+from copy import deepcopy
+import numpy as np
+from tqdm import tqdm
+from models.gcn import GCN
+from models.sgc import SGC
+from models.sgc_multi import SGC as SGC1
+from models.myappnp import APPNP
+from models.myappnp1 import APPNP1
+from models.mycheby import Cheby
+from models.mygraphsage import GraphSage
+from models.gat import GAT
+import scipy.sparse as sp
+
+
+class Evaluator:
+
+    def __init__(self, data, args, device='cuda', **kwargs):
+        self.data = data
+        self.args = args
+        self.device = device
+        n = int(data.feat_train.shape[0] * args.reduction_rate)
+        d = data.feat_train.shape[1]
+        self.nnodes_syn = n
+        self.adj_param= nn.Parameter(torch.FloatTensor(n, n).to(device))
+        self.feat_syn = nn.Parameter(torch.FloatTensor(n, d).to(device))
+        self.labels_syn = torch.LongTensor(self.generate_labels_syn(data)).to(device)
+        self.reset_parameters()
+        print('adj_param:', self.adj_param.shape, 'feat_syn:', self.feat_syn.shape)
+
+    def reset_parameters(self):
+        self.adj_param.data.copy_(torch.randn(self.adj_param.size()))
+        self.feat_syn.data.copy_(torch.randn(self.feat_syn.size()))
+
+    def generate_labels_syn(self, data):
+        from collections import Counter
+        counter = Counter(data.labels_train)
+        num_class_dict = {}
+        n = len(data.labels_train)
+
+        sorted_counter = sorted(counter.items(), key=lambda x:x[1])
+        sum_ = 0
+        labels_syn = []
+        self.syn_class_indices = {}
+        for ix, (c, num) in enumerate(sorted_counter):
+            if ix == len(sorted_counter) - 1:
+                num_class_dict[c] = int(n * self.args.reduction_rate) - sum_
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+            else:
+                num_class_dict[c] = max(int(num * self.args.reduction_rate), 1)
+                sum_ += num_class_dict[c]
+                self.syn_class_indices[c] = [len(labels_syn), len(labels_syn) + num_class_dict[c]]
+                labels_syn += [c] * num_class_dict[c]
+
+        self.num_class_dict = num_class_dict
+        return labels_syn
+
+
+    def test_gat(self, nlayers, model_type, verbose=False):
+        res = []
+        args = self.args
+
+        if args.dataset in ['cora', 'citeseer']:
+            args.epsilon = 0.5 # Make the graph sparser as GAT does not work well on dense graph
+        else:
+            args.epsilon = 0.01
+
+        print('======= testing %s' % model_type)
+        data, device = self.data, self.device
+
+
+        feat_syn, adj_syn, labels_syn = self.get_syn_data(model_type)
+        # with_bn = True if self.args.dataset in ['ogbn-arxiv'] else False
+        with_bn = False
+        if model_type == 'GAT':
+            model = GAT(nfeat=feat_syn.shape[1], nhid=16, heads=16, dropout=0.0,
+                        weight_decay=0e-4, nlayers=self.args.nlayers, lr=0.001,
+                        nclass=data.nclass, device=device, dataset=self.args.dataset).to(device)
+
+
+        noval = True if args.dataset in ['reddit', 'flickr'] else False
+        model.fit(feat_syn, adj_syn, labels_syn, np.arange(len(feat_syn)), noval=noval, data=data,
+                     train_iters=10000 if noval else 3000, normalize=True, verbose=verbose)
+
+        model.eval()
+        labels_test = torch.LongTensor(data.labels_test).cuda()
+
+        if args.dataset in ['reddit', 'flickr']:
+            output = model.predict(data.feat_test, data.adj_test)
+            loss_test = F.nll_loss(output, labels_test)
+            acc_test = utils.accuracy(output, labels_test)
+            res.append(acc_test.item())
+            if verbose:
+                print("Test set results:",
+                      "loss= {:.4f}".format(loss_test.item()),
+                      "accuracy= {:.4f}".format(acc_test.item()))
+
+        else:
+            # Full graph
+            output = model.predict(data.feat_full, data.adj_full)
+            loss_test = F.nll_loss(output[data.idx_test], labels_test)
+            acc_test = utils.accuracy(output[data.idx_test], labels_test)
+            res.append(acc_test.item())
+            if verbose:
+                print("Test set results:",
+                      "loss= {:.4f}".format(loss_test.item()),
+                      "accuracy= {:.4f}".format(acc_test.item()))
+
+        labels_train = torch.LongTensor(data.labels_train).cuda()
+        output = model.predict(data.feat_train, data.adj_train)
+        loss_train = F.nll_loss(output, labels_train)
+        acc_train = utils.accuracy(output, labels_train)
+        if verbose:
+            print("Train set results:",
+                  "loss= {:.4f}".format(loss_train.item()),
+                  "accuracy= {:.4f}".format(acc_train.item()))
+        res.append(acc_train.item())
+        return res
+
+    def get_syn_data(self, model_type=None):
+        data, device = self.data, self.device
+        feat_syn, adj_param, labels_syn = self.feat_syn.detach(), \
+                                self.adj_param.detach(), self.labels_syn
+
+        args = self.args
+        adj_syn = torch.load(f'saved_ours/adj_{args.dataset}_{args.reduction_rate}_{args.seed}.pt', map_location='cuda')
+        feat_syn = torch.load(f'saved_ours/feat_{args.dataset}_{args.reduction_rate}_{args.seed}.pt', map_location='cuda')
+
+        if model_type == 'MLP':
+            adj_syn = adj_syn.to(self.device)
+            adj_syn = adj_syn - adj_syn
+        else:
+            adj_syn = adj_syn.to(self.device)
+
+        print('Sum:', adj_syn.sum(), adj_syn.sum()/(adj_syn.shape[0]**2))
+        print('Sparsity:', adj_syn.nonzero().shape[0]/(adj_syn.shape[0]**2))
+
+        if self.args.epsilon > 0:
+            adj_syn[adj_syn < self.args.epsilon] = 0
+            print('Sparsity after truncating:', adj_syn.nonzero().shape[0]/(adj_syn.shape[0]**2))
+        feat_syn = feat_syn.to(self.device)
+
+        # edge_index = adj_syn.nonzero().T
+        # adj_syn = torch.sparse.FloatTensor(edge_index,  adj_syn[edge_index[0], edge_index[1]], adj_syn.size())
+
+        return feat_syn, adj_syn, labels_syn
+
+
+    def test(self, nlayers, model_type, verbose=True):
+        res = []
+
+        args = self.args
+        data, device = self.data, self.device
+
+        feat_syn, adj_syn, labels_syn = self.get_syn_data(model_type)
+
+        print('======= testing %s' % model_type)
+        if model_type == 'MLP':
+            model_class = GCN
+        else:
+            model_class = eval(model_type)
+        weight_decay = 5e-4
+        dropout = 0.5 if args.dataset in ['reddit'] else 0
+
+        model = model_class(nfeat=feat_syn.shape[1], nhid=args.hidden, dropout=dropout,
+                    weight_decay=weight_decay, nlayers=nlayers,
+                    nclass=data.nclass, device=device).to(device)
+
+        # with_bn = True if self.args.dataset in ['ogbn-arxiv'] else False
+        if args.dataset in ['ogbn-arxiv', 'arxiv']:
+            model = model_class(nfeat=feat_syn.shape[1], nhid=args.hidden, dropout=0.,
+                        weight_decay=weight_decay, nlayers=nlayers, with_bn=False,
+                        nclass=data.nclass, device=device).to(device)
+
+        noval = True if args.dataset in ['reddit', 'flickr'] else False
+        model.fit_with_val(feat_syn, adj_syn, labels_syn, data,
+                     train_iters=600, normalize=True, verbose=True, noval=noval)
+
+        model.eval()
+        labels_test = torch.LongTensor(data.labels_test).cuda()
+
+        if model_type == 'MLP':
+            output = model.predict_unnorm(data.feat_test, sp.eye(len(data.feat_test)))
+        else:
+            output = model.predict(data.feat_test, data.adj_test)
+
+        if args.dataset in ['reddit', 'flickr']:
+            loss_test = F.nll_loss(output, labels_test)
+            acc_test = utils.accuracy(output, labels_test)
+            res.append(acc_test.item())
+            if verbose:
+                print("Test set results:",
+                      "loss= {:.4f}".format(loss_test.item()),
+                      "accuracy= {:.4f}".format(acc_test.item()))
+
+        # if not args.dataset in ['reddit', 'flickr']:
+        else:
+            # Full graph
+            output = model.predict(data.feat_full, data.adj_full)
+            loss_test = F.nll_loss(output[data.idx_test], labels_test)
+            acc_test = utils.accuracy(output[data.idx_test], labels_test)
+            res.append(acc_test.item())
+            if verbose:
+                print("Test full set results:",
+                      "loss= {:.4f}".format(loss_test.item()),
+                      "accuracy= {:.4f}".format(acc_test.item()))
+
+            labels_train = torch.LongTensor(data.labels_train).cuda()
+            output = model.predict(data.feat_train, data.adj_train)
+            loss_train = F.nll_loss(output, labels_train)
+            acc_train = utils.accuracy(output, labels_train)
+            if verbose:
+                print("Train set results:",
+                      "loss= {:.4f}".format(loss_train.item()),
+                      "accuracy= {:.4f}".format(acc_train.item()))
+            res.append(acc_train.item())
+        return res
+
+    def train(self, verbose=True):
+        args = self.args
+        data = self.data
+
+        final_res = {}
+        runs = self.args.nruns
+
+        for model_type in ['GCN',  'GraphSage', 'SGC1', 'MLP', 'APPNP1', 'Cheby']:
+            res = []
+            nlayer = 2
+            for i in range(runs):
+                res.append(self.test(nlayer, verbose=False, model_type=model_type))
+            res = np.array(res)
+            print('Test/Train Mean Accuracy:',
+                    repr([res.mean(0), res.std(0)]))
+            final_res[model_type] = [res.mean(0), res.std(0)]
+
+
+        print('=== testing GAT')
+        res = []
+        nlayer = 2
+        for i in range(runs):
+            res.append(self.test_gat(verbose=True, nlayers=nlayer, model_type='GAT'))
+        res = np.array(res)
+        print('Layer:', nlayer)
+        print('Test/Full Test/Train Mean Accuracy:',
+                repr([res.mean(0), res.std(0)]))
+        final_res['GAT'] = [res.mean(0), res.std(0)]
+
+        print('Final result:', final_res)
+

GCond/train_coreset.py

@@ -0,0 +1,117 @@
+from deeprobust.graph.data import Dataset
+import numpy as np
+import random
+import time
+import argparse
+import torch
+import sys
+from deeprobust.graph.utils import *
+import torch.nn.functional as F
+from configs import load_config
+from utils import *
+from utils_graphsaint import DataGraphSAINT
+from models.gcn import GCN
+from coreset import KCenter, Herding, Random, LRMC
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu_id', type=int, default=0, help='gpu id')
+parser.add_argument('--dataset', type=str, default='cora')
+parser.add_argument('--hidden', type=int, default=256)
+parser.add_argument('--normalize_features', type=bool, default=True)
+parser.add_argument('--keep_ratio', type=float, default=1.0)
+parser.add_argument('--lr', type=float, default=0.01)
+parser.add_argument('--weight_decay', type=float, default=5e-4)
+parser.add_argument('--seed', type=int, default=15, help='Random seed.')
+parser.add_argument('--nlayers', type=int, default=2, help='Random seed.')
+parser.add_argument('--epochs', type=int, default=400)
+parser.add_argument('--inductive', type=int, default=1)
+parser.add_argument('--save', type=int, default=0)
+parser.add_argument('--method', type=str, choices=['kcenter', 'herding', 'random', 'lrmc'])
+parser.add_argument('--lrmc_seeds_path', type=str, default=None,
+                    help='Path to a JSON file containing L‑RMC seed nodes. Required when method=lrmc.')
+parser.add_argument('--reduction_rate', type=float, required=True)
+args = parser.parse_args()
+
+torch.cuda.set_device(args.gpu_id)
+args = load_config(args)
+print(args)
+
+# random seed setting
+random.seed(args.seed)
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed(args.seed)
+
+
+data_graphsaint = ['flickr', 'reddit', 'ogbn-arxiv']
+if args.dataset in data_graphsaint:
+    data = DataGraphSAINT(args.dataset)
+    data_full = data.data_full
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+else:
+    data_full = get_dataset(args.dataset, args.normalize_features)
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+
+features = data_full.features
+adj = data_full.adj
+labels = data_full.labels
+idx_train = data_full.idx_train
+idx_val = data_full.idx_val
+idx_test = data_full.idx_test
+
+# Setup GCN Model
+device = 'cuda'
+model = GCN(nfeat=features.shape[1], nhid=256, nclass=labels.max()+1, device=device, weight_decay=args.weight_decay)
+
+model = model.to(device)
+model.fit(features, adj, labels, idx_train, idx_val, train_iters=600, verbose=False)
+
+model.eval()
+# You can use the inner function of model to test
+model.test(idx_test)
+
+embeds = model.predict().detach()
+
+if args.method == 'kcenter':
+    agent = KCenter(data, args, device='cuda')
+elif args.method == 'herding':
+    agent = Herding(data, args, device='cuda')
+elif args.method == 'random':
+    agent = Random(data, args, device='cuda')
+elif args.method == 'lrmc':
+    if args.lrmc_seeds_path is None:
+        raise ValueError("--lrmc_seeds_path must be specified when method='lrmc'")
+    agent = LRMC(data, args, device='cuda')
+
+idx_selected = agent.select(embeds)
+
+
+feat_train = features[idx_selected]
+adj_train = adj[np.ix_(idx_selected, idx_selected)]
+
+labels_train = labels[idx_selected]
+
+if args.save:
+    np.save(f'saved/idx_{args.dataset}_{args.reduction_rate}_{args.method}_{args.seed}.npy', idx_selected)
+
+
+res = []
+runs = 10
+for _ in tqdm(range(runs)):
+    model.initialize()
+    model.fit_with_val(feat_train, adj_train, labels_train, data,
+                 train_iters=600, normalize=True, verbose=False)
+
+    model.eval()
+    labels_test = torch.LongTensor(data.labels_test).cuda()
+
+    # Full graph
+    output = model.predict(data.feat_full, data.adj_full)
+    loss_test = F.nll_loss(output[data.idx_test], labels_test)
+    acc_test = accuracy(output[data.idx_test], labels_test)
+    res.append(acc_test.item())
+
+res = np.array(res)
+print('Mean accuracy:', repr([res.mean(), res.std()]))
+

GCond/train_coreset_induct.py

@@ -0,0 +1,119 @@
+from deeprobust.graph.data import Dataset
+import numpy as np
+import random
+import time
+import argparse
+import torch
+import sys
+import deeprobust.graph.utils as utils
+import torch.nn.functional as F
+from configs import load_config
+from utils import *
+from utils_graphsaint import DataGraphSAINT
+from models.gcn import GCN
+from coreset import KCenter, Herding, Random, LRMC
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu_id', type=int, default=0, help='gpu id')
+parser.add_argument('--dataset', type=str, default='cora')
+parser.add_argument('--hidden', type=int, default=256)
+parser.add_argument('--normalize_features', type=bool, default=True)
+parser.add_argument('--keep_ratio', type=float, default=1.0)
+parser.add_argument('--lr', type=float, default=0.01)
+parser.add_argument('--weight_decay', type=float, default=5e-4)
+parser.add_argument('--dropout', type=float, default=0.5)
+parser.add_argument('--seed', type=int, default=15, help='Random seed.')
+parser.add_argument('--nlayers', type=int, default=2, help='Random seed.')
+parser.add_argument('--epochs', type=int, default=400)
+parser.add_argument('--inductive', type=int, default=1)
+parser.add_argument('--mlp', type=int, default=0)
+parser.add_argument('--method', type=str, choices=['kcenter', 'herding', 'random', 'lrmc'])
+parser.add_argument('--lrmc_seeds_path', type=str, default=None,
+                    help='Path to a JSON file containing L‑RMC seed nodes. Required when method=lrmc.')
+parser.add_argument('--reduction_rate', type=float, required=True)
+args = parser.parse_args()
+
+torch.cuda.set_device(args.gpu_id)
+args = load_config(args)
+print(args)
+
+# random seed setting
+random.seed(args.seed)
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed(args.seed)
+
+data_graphsaint = ['flickr', 'reddit', 'ogbn-arxiv']
+if args.dataset in data_graphsaint:
+    data = DataGraphSAINT(args.dataset)
+    data_full = data.data_full
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+else:
+    data_full = get_dataset(args.dataset, args.normalize_features)
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+
+feat_train = data.feat_train
+adj_train = data.adj_train
+labels_train = data.labels_train
+
+
+# Setup GCN Model
+device = 'cuda'
+model = GCN(nfeat=feat_train.shape[1], nhid=256, nclass=labels_train.max()+1, device=device, weight_decay=args.weight_decay)
+
+model = model.to(device)
+
+model.fit_with_val(feat_train, adj_train, labels_train, data,
+             train_iters=600, normalize=True, verbose=False)
+
+model.eval()
+labels_test = torch.LongTensor(data.labels_test).cuda()
+feat_test, adj_test = data.feat_test, data.adj_test
+
+embeds = model.predict().detach()
+
+output = model.predict(feat_test, adj_test)
+loss_test = F.nll_loss(output, labels_test)
+acc_test = utils.accuracy(output, labels_test)
+print("Test set results:",
+      "loss= {:.4f}".format(loss_test.item()),
+      "accuracy= {:.4f}".format(acc_test.item()))
+
+
+if args.method == 'kcenter':
+    agent = KCenter(data, args, device='cuda')
+elif args.method == 'herding':
+    agent = Herding(data, args, device='cuda')
+elif args.method == 'random':
+    agent = Random(data, args, device='cuda')
+elif args.method == 'lrmc':
+    if args.lrmc_seeds_path is None:
+        raise ValueError("--lrmc_seeds_path must be specified when method='lrmc'")
+    agent = LRMC(data, args, device='cuda')
+
+idx_selected = agent.select(embeds, inductive=True)
+
+feat_train = feat_train[idx_selected]
+adj_train = adj_train[np.ix_(idx_selected, idx_selected)]
+
+labels_train = labels_train[idx_selected]
+
+res = []
+print('shape of feat_train:', feat_train.shape)
+runs = 10
+for _ in tqdm(range(runs)):
+    model.initialize()
+    model.fit_with_val(feat_train, adj_train, labels_train, data,
+                 train_iters=600, normalize=True, verbose=False, noval=True)
+
+    model.eval()
+    labels_test = torch.LongTensor(data.labels_test).cuda()
+
+    output = model.predict(feat_test, adj_test)
+    loss_test = F.nll_loss(output, labels_test)
+    acc_test = utils.accuracy(output, labels_test)
+    res.append(acc_test.item())
+res = np.array(res)
+print('Mean accuracy:', repr([res.mean(), res.std()]))
+

GCond/train_gcond_induct.py

@@ -0,0 +1,61 @@
+from deeprobust.graph.data import Dataset
+import numpy as np
+import random
+import time
+import argparse
+import torch
+from utils import *
+import torch.nn.functional as F
+from gcond_agent_induct import GCond
+from utils_graphsaint import DataGraphSAINT
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu_id', type=int, default=0, help='gpu id')
+parser.add_argument('--dataset', type=str, default='cora')
+parser.add_argument('--dis_metric', type=str, default='ours')
+parser.add_argument('--epochs', type=int, default=600)
+parser.add_argument('--nlayers', type=int, default=3)
+parser.add_argument('--hidden', type=int, default=256)
+parser.add_argument('--lr_adj', type=float, default=0.01)
+parser.add_argument('--lr_feat', type=float, default=0.01)
+parser.add_argument('--lr_model', type=float, default=0.01)
+parser.add_argument('--weight_decay', type=float, default=0.0)
+parser.add_argument('--dropout', type=float, default=0.0)
+parser.add_argument('--normalize_features', type=bool, default=True)
+parser.add_argument('--keep_ratio', type=float, default=1.0)
+parser.add_argument('--reduction_rate', type=float, default=0.01)
+parser.add_argument('--seed', type=int, default=15, help='Random seed.')
+parser.add_argument('--alpha', type=float, default=0, help='regularization term.')
+parser.add_argument('--debug', type=int, default=0)
+parser.add_argument('--sgc', type=int, default=1)
+parser.add_argument('--inner', type=int, default=0)
+parser.add_argument('--outer', type=int, default=20)
+parser.add_argument('--option', type=int, default=0)
+parser.add_argument('--save', type=int, default=0)
+parser.add_argument('--label_rate', type=float, default=1)
+parser.add_argument('--one_step', type=int, default=0)
+args = parser.parse_args()
+
+torch.cuda.set_device(args.gpu_id)
+
+# random seed setting
+random.seed(args.seed)
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed(args.seed)
+
+print(args)
+
+data_graphsaint = ['flickr', 'reddit', 'ogbn-arxiv']
+if args.dataset in data_graphsaint:
+    # data = DataGraphSAINT(args.dataset)
+    data = DataGraphSAINT(args.dataset, label_rate=args.label_rate)
+    data_full = data.data_full
+else:
+    data_full = get_dataset(args.dataset, args.normalize_features)
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+
+agent = GCond(data, args, device='cuda')
+
+agent.train()

GCond/train_gcond_transduct.py

@@ -0,0 +1,57 @@
+from deeprobust.graph.data import Dataset
+import numpy as np
+import random
+import time
+import argparse
+import torch
+from utils import *
+import torch.nn.functional as F
+from gcond_agent_transduct import GCond
+from utils_graphsaint import DataGraphSAINT
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu_id', type=int, default=0, help='gpu id')
+parser.add_argument('--dataset', type=str, default='cora')
+parser.add_argument('--dis_metric', type=str, default='ours')
+parser.add_argument('--epochs', type=int, default=2000)
+parser.add_argument('--nlayers', type=int, default=3)
+parser.add_argument('--hidden', type=int, default=256)
+parser.add_argument('--lr_adj', type=float, default=0.01)
+parser.add_argument('--lr_feat', type=float, default=0.01)
+parser.add_argument('--lr_model', type=float, default=0.01)
+parser.add_argument('--weight_decay', type=float, default=0.0)
+parser.add_argument('--dropout', type=float, default=0.0)
+parser.add_argument('--normalize_features', type=bool, default=True)
+parser.add_argument('--keep_ratio', type=float, default=1.0)
+parser.add_argument('--reduction_rate', type=float, default=1)
+parser.add_argument('--seed', type=int, default=15, help='Random seed.')
+parser.add_argument('--alpha', type=float, default=0, help='regularization term.')
+parser.add_argument('--debug', type=int, default=0)
+parser.add_argument('--sgc', type=int, default=1)
+parser.add_argument('--inner', type=int, default=0)
+parser.add_argument('--outer', type=int, default=20)
+parser.add_argument('--save', type=int, default=0)
+parser.add_argument('--one_step', type=int, default=0)
+args = parser.parse_args()
+
+torch.cuda.set_device(args.gpu_id)
+
+# random seed setting
+random.seed(args.seed)
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed(args.seed)
+
+print(args)
+
+data_graphsaint = ['flickr', 'reddit', 'ogbn-arxiv']
+if args.dataset in data_graphsaint:
+    data = DataGraphSAINT(args.dataset)
+    data_full = data.data_full
+else:
+    data_full = get_dataset(args.dataset, args.normalize_features)
+    data = Transd2Ind(data_full, keep_ratio=args.keep_ratio)
+
+agent = GCond(data, args, device='cuda')
+
+agent.train()

GCond/utils.py

@@ -0,0 +1,383 @@
+import os.path as osp
+import numpy as np
+import scipy.sparse as sp
+import torch
+import torch_geometric.transforms as T
+from ogb.nodeproppred import PygNodePropPredDataset, Evaluator
+from deeprobust.graph.data import Dataset
+from deeprobust.graph.utils import get_train_val_test
+from torch_geometric.utils import train_test_split_edges
+from sklearn.model_selection import train_test_split
+from sklearn import metrics
+import numpy as np
+import torch.nn.functional as F
+from sklearn.preprocessing import StandardScaler
+from deeprobust.graph.utils import *
+from torch_geometric.data import NeighborSampler
+from torch_geometric.utils import add_remaining_self_loops, to_undirected
+from torch_geometric.datasets import Planetoid
+
+
+def get_dataset(name, normalize_features=False, transform=None, if_dpr=True):
+    path = osp.join(osp.dirname(osp.realpath(__file__)), 'data', name)
+    if name in ['cora', 'citeseer', 'pubmed']:
+        dataset = Planetoid(path, name)
+    elif name in ['ogbn-arxiv']:
+        dataset = PygNodePropPredDataset(name='ogbn-arxiv')
+    else:
+        raise NotImplementedError
+
+    if transform is not None and normalize_features:
+        dataset.transform = T.Compose([T.NormalizeFeatures(), transform])
+    elif normalize_features:
+        dataset.transform = T.NormalizeFeatures()
+    elif transform is not None:
+        dataset.transform = transform
+
+    dpr_data = Pyg2Dpr(dataset)
+    if name in ['ogbn-arxiv']:
+        # the features are different from the features provided by GraphSAINT
+        # normalize features, following graphsaint
+        feat, idx_train = dpr_data.features, dpr_data.idx_train
+        feat_train = feat[idx_train]
+        scaler = StandardScaler()
+        scaler.fit(feat_train)
+        feat = scaler.transform(feat)
+        dpr_data.features = feat
+
+    return dpr_data
+
+
+class Pyg2Dpr(Dataset):
+    def __init__(self, pyg_data, **kwargs):
+        try:
+            splits = pyg_data.get_idx_split()
+        except:
+            pass
+
+        dataset_name = pyg_data.name
+        pyg_data = pyg_data[0]
+        n = pyg_data.num_nodes
+
+        if dataset_name == 'ogbn-arxiv': # symmetrization
+            pyg_data.edge_index = to_undirected(pyg_data.edge_index, pyg_data.num_nodes)
+
+        self.adj = sp.csr_matrix((np.ones(pyg_data.edge_index.shape[1]),
+            (pyg_data.edge_index[0], pyg_data.edge_index[1])), shape=(n, n))
+
+        self.features = pyg_data.x.numpy()
+        self.labels = pyg_data.y.numpy()
+
+        if len(self.labels.shape) == 2 and self.labels.shape[1] == 1:
+            self.labels = self.labels.reshape(-1) # ogb-arxiv needs to reshape
+
+        if hasattr(pyg_data, 'train_mask'):
+            # for fixed split
+            self.idx_train = mask_to_index(pyg_data.train_mask, n)
+            self.idx_val = mask_to_index(pyg_data.val_mask, n)
+            self.idx_test = mask_to_index(pyg_data.test_mask, n)
+            self.name = 'Pyg2Dpr'
+        else:
+            try:
+                # for ogb
+                self.idx_train = splits['train']
+                self.idx_val = splits['valid']
+                self.idx_test = splits['test']
+                self.name = 'Pyg2Dpr'
+            except:
+                # for other datasets
+                self.idx_train, self.idx_val, self.idx_test = get_train_val_test(
+                        nnodes=n, val_size=0.1, test_size=0.8, stratify=self.labels)
+
+
+def mask_to_index(index, size):
+    all_idx = np.arange(size)
+    return all_idx[index]
+
+def index_to_mask(index, size):
+    mask = torch.zeros((size, ), dtype=torch.bool)
+    mask[index] = 1
+    return mask
+
+
+
+class Transd2Ind:
+    # transductive setting to inductive setting
+
+    def __init__(self, dpr_data, keep_ratio):
+        idx_train, idx_val, idx_test = dpr_data.idx_train, dpr_data.idx_val, dpr_data.idx_test
+        adj, features, labels = dpr_data.adj, dpr_data.features, dpr_data.labels
+        self.nclass = labels.max()+1
+        self.adj_full, self.feat_full, self.labels_full = adj, features, labels
+        self.idx_train = np.array(idx_train)
+        self.idx_val = np.array(idx_val)
+        self.idx_test = np.array(idx_test)
+
+        if keep_ratio < 1:
+            idx_train, _ = train_test_split(idx_train,
+                                            random_state=None,
+                                            train_size=keep_ratio,
+                                            test_size=1-keep_ratio,
+                                            stratify=labels[idx_train])
+
+        self.adj_train = adj[np.ix_(idx_train, idx_train)]
+        self.adj_val = adj[np.ix_(idx_val, idx_val)]
+        self.adj_test = adj[np.ix_(idx_test, idx_test)]
+        print('size of adj_train:', self.adj_train.shape)
+        print('#edges in adj_train:', self.adj_train.sum())
+
+        self.labels_train = labels[idx_train]
+        self.labels_val = labels[idx_val]
+        self.labels_test = labels[idx_test]
+
+        self.feat_train = features[idx_train]
+        self.feat_val = features[idx_val]
+        self.feat_test = features[idx_test]
+
+        self.class_dict = None
+        self.samplers = None
+        self.class_dict2 = None
+
+    def retrieve_class(self, c, num=256):
+        if self.class_dict is None:
+            self.class_dict = {}
+            for i in range(self.nclass):
+                self.class_dict['class_%s'%i] = (self.labels_train == i)
+        idx = np.arange(len(self.labels_train))
+        idx = idx[self.class_dict['class_%s'%c]]
+        return np.random.permutation(idx)[:num]
+
+    def retrieve_class_sampler(self, c, adj, transductive, num=256, args=None):
+        if self.class_dict2 is None:
+            self.class_dict2 = {}
+            for i in range(self.nclass):
+                if transductive:
+                    idx = self.idx_train[self.labels_train == i]
+                else:
+                    idx = np.arange(len(self.labels_train))[self.labels_train==i]
+                self.class_dict2[i] = idx
+
+        if args.nlayers == 1:
+            sizes = [15]
+        if args.nlayers == 2:
+            sizes = [10, 5]
+            # sizes = [-1, -1]
+        if args.nlayers == 3:
+            sizes = [15, 10, 5]
+        if args.nlayers == 4:
+            sizes = [15, 10, 5, 5]
+        if args.nlayers == 5:
+            sizes = [15, 10, 5, 5, 5]
+
+
+        if self.samplers is None:
+            self.samplers = []
+            for i in range(self.nclass):
+                node_idx = torch.LongTensor(self.class_dict2[i])
+                self.samplers.append(NeighborSampler(adj,
+                                    node_idx=node_idx,
+                                    sizes=sizes, batch_size=num,
+                                    num_workers=12, return_e_id=False,
+                                    num_nodes=adj.size(0),
+                                    shuffle=True))
+        batch = np.random.permutation(self.class_dict2[c])[:num]
+        out = self.samplers[c].sample(batch)
+        return out
+
+    def retrieve_class_multi_sampler(self, c, adj, transductive, num=256, args=None):
+        if self.class_dict2 is None:
+            self.class_dict2 = {}
+            for i in range(self.nclass):
+                if transductive:
+                    idx = self.idx_train[self.labels_train == i]
+                else:
+                    idx = np.arange(len(self.labels_train))[self.labels_train==i]
+                self.class_dict2[i] = idx
+
+
+        if self.samplers is None:
+            self.samplers = []
+            for l in range(2):
+                layer_samplers = []
+                sizes = [15] if l == 0 else [10, 5]
+                for i in range(self.nclass):
+                    node_idx = torch.LongTensor(self.class_dict2[i])
+                    layer_samplers.append(NeighborSampler(adj,
+                                        node_idx=node_idx,
+                                        sizes=sizes, batch_size=num,
+                                        num_workers=12, return_e_id=False,
+                                        num_nodes=adj.size(0),
+                                        shuffle=True))
+                self.samplers.append(layer_samplers)
+        batch = np.random.permutation(self.class_dict2[c])[:num]
+        out = self.samplers[args.nlayers-1][c].sample(batch)
+        return out
+
+
+
+def match_loss(gw_syn, gw_real, args, device):
+    dis = torch.tensor(0.0).to(device)
+
+    if args.dis_metric == 'ours':
+
+        for ig in range(len(gw_real)):
+            gwr = gw_real[ig]
+            gws = gw_syn[ig]
+            dis += distance_wb(gwr, gws)
+
+    elif args.dis_metric == 'mse':
+        gw_real_vec = []
+        gw_syn_vec = []
+        for ig in range(len(gw_real)):
+            gw_real_vec.append(gw_real[ig].reshape((-1)))
+            gw_syn_vec.append(gw_syn[ig].reshape((-1)))
+        gw_real_vec = torch.cat(gw_real_vec, dim=0)
+        gw_syn_vec = torch.cat(gw_syn_vec, dim=0)
+        dis = torch.sum((gw_syn_vec - gw_real_vec)**2)
+
+    elif args.dis_metric == 'cos':
+        gw_real_vec = []
+        gw_syn_vec = []
+        for ig in range(len(gw_real)):
+            gw_real_vec.append(gw_real[ig].reshape((-1)))
+            gw_syn_vec.append(gw_syn[ig].reshape((-1)))
+        gw_real_vec = torch.cat(gw_real_vec, dim=0)
+        gw_syn_vec = torch.cat(gw_syn_vec, dim=0)
+        dis = 1 - torch.sum(gw_real_vec * gw_syn_vec, dim=-1) / (torch.norm(gw_real_vec, dim=-1) * torch.norm(gw_syn_vec, dim=-1) + 0.000001)
+
+    else:
+        exit('DC error: unknown distance function')
+
+    return dis
+
+def distance_wb(gwr, gws):
+    shape = gwr.shape
+
+    # TODO: output node!!!!
+    if len(gwr.shape) == 2:
+        gwr = gwr.T
+        gws = gws.T
+
+    if len(shape) == 4: # conv, out*in*h*w
+        gwr = gwr.reshape(shape[0], shape[1] * shape[2] * shape[3])
+        gws = gws.reshape(shape[0], shape[1] * shape[2] * shape[3])
+    elif len(shape) == 3:  # layernorm, C*h*w
+        gwr = gwr.reshape(shape[0], shape[1] * shape[2])
+        gws = gws.reshape(shape[0], shape[1] * shape[2])
+    elif len(shape) == 2: # linear, out*in
+        tmp = 'do nothing'
+    elif len(shape) == 1: # batchnorm/instancenorm, C; groupnorm x, bias
+        gwr = gwr.reshape(1, shape[0])
+        gws = gws.reshape(1, shape[0])
+        return 0
+
+    dis_weight = torch.sum(1 - torch.sum(gwr * gws, dim=-1) / (torch.norm(gwr, dim=-1) * torch.norm(gws, dim=-1) + 0.000001))
+    dis = dis_weight
+    return dis
+
+
+
+def calc_f1(y_true, y_pred,is_sigmoid):
+    if not is_sigmoid:
+        y_pred = np.argmax(y_pred, axis=1)
+    else:
+        y_pred[y_pred > 0.5] = 1
+        y_pred[y_pred <= 0.5] = 0
+    return metrics.f1_score(y_true, y_pred, average="micro"), metrics.f1_score(y_true, y_pred, average="macro")
+
+def evaluate(output, labels, args):
+    data_graphsaint = ['yelp', 'ppi', 'ppi-large', 'flickr', 'reddit', 'amazon']
+    if args.dataset in data_graphsaint:
+        labels = labels.cpu().numpy()
+        output = output.cpu().numpy()
+        if len(labels.shape) > 1:
+            micro, macro = calc_f1(labels, output, is_sigmoid=True)
+        else:
+            micro, macro = calc_f1(labels, output, is_sigmoid=False)
+        print("Test set results:", "F1-micro= {:.4f}".format(micro),
+                "F1-macro= {:.4f}".format(macro))
+    else:
+        loss_test = F.nll_loss(output, labels)
+        acc_test = accuracy(output, labels)
+        print("Test set results:",
+              "loss= {:.4f}".format(loss_test.item()),
+              "accuracy= {:.4f}".format(acc_test.item()))
+    return
+
+
+from torchvision import datasets, transforms
+def get_mnist(data_path):
+    channel = 1
+    im_size = (28, 28)
+    num_classes = 10
+    mean = [0.1307]
+    std = [0.3081]
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
+    dst_train = datasets.MNIST(data_path, train=True, download=True, transform=transform) # no        augmentation
+    dst_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)
+    class_names = [str(c) for c in range(num_classes)]
+
+    labels = []
+    feat = []
+    for x, y in dst_train:
+        feat.append(x.view(1, -1))
+        labels.append(y)
+    feat = torch.cat(feat, axis=0).numpy()
+    from utils_graphsaint import GraphData
+    adj = sp.eye(len(feat))
+    idx = np.arange(len(feat))
+    dpr_data = GraphData(adj-adj, feat, labels, idx, idx, idx)
+    from deeprobust.graph.data import Dpr2Pyg
+    return Dpr2Pyg(dpr_data)
+
+def regularization(adj, x, eig_real=None):
+    # fLf
+    loss = 0
+    # loss += torch.norm(adj, p=1)
+    loss += feature_smoothing(adj, x)
+    return loss
+
+def maxdegree(adj):
+    n = adj.shape[0]
+    return F.relu(max(adj.sum(1))/n - 0.5)
+
+def sparsity2(adj):
+    n = adj.shape[0]
+    loss_degree = - torch.log(adj.sum(1)).sum() / n
+    loss_fro = torch.norm(adj) / n
+    return 0 * loss_degree + loss_fro
+
+def sparsity(adj):
+    n = adj.shape[0]
+    thresh = n * n * 0.01
+    return F.relu(adj.sum()-thresh)
+    # return F.relu(adj.sum()-thresh) / n**2
+
+def feature_smoothing(adj, X):
+    adj = (adj.t() + adj)/2
+    rowsum = adj.sum(1)
+    r_inv = rowsum.flatten()
+    D = torch.diag(r_inv)
+    L = D - adj
+
+    r_inv = r_inv  + 1e-8
+    r_inv = r_inv.pow(-1/2).flatten()
+    r_inv[torch.isinf(r_inv)] = 0.
+    r_mat_inv = torch.diag(r_inv)
+    # L = r_mat_inv @ L
+    L = r_mat_inv @ L @ r_mat_inv
+
+    XLXT = torch.matmul(torch.matmul(X.t(), L), X)
+    loss_smooth_feat = torch.trace(XLXT)
+    # loss_smooth_feat = loss_smooth_feat / (adj.shape[0]**2)
+    return loss_smooth_feat
+
+def row_normalize_tensor(mx):
+    rowsum = mx.sum(1)
+    r_inv = rowsum.pow(-1).flatten()
+    # r_inv[torch.isinf(r_inv)] = 0.
+    r_mat_inv = torch.diag(r_inv)
+    mx = r_mat_inv @ mx
+    return mx
+
+

GCond/utils_graphsaint.py

@@ -0,0 +1,220 @@
+import scipy.sparse as sp
+import numpy as np
+import sys
+import json
+import os
+from sklearn.preprocessing import StandardScaler
+from torch_geometric.data import InMemoryDataset, Data
+import torch
+from itertools import repeat
+from torch_geometric.data import NeighborSampler
+
+class DataGraphSAINT:
+    '''datasets used in GraphSAINT paper'''
+
+    def __init__(self, dataset, **kwargs):
+        dataset_str='data/'+dataset+'/'
+        adj_full = sp.load_npz(dataset_str+'adj_full.npz')
+        self.nnodes = adj_full.shape[0]
+        if dataset == 'ogbn-arxiv':
+            adj_full = adj_full + adj_full.T
+            adj_full[adj_full > 1] = 1
+
+        role = json.load(open(dataset_str+'role.json','r'))
+        idx_train = role['tr']
+        idx_test = role['te']
+        idx_val = role['va']
+
+        if 'label_rate' in kwargs:
+            label_rate = kwargs['label_rate']
+            if label_rate < 1:
+                idx_train = idx_train[:int(label_rate*len(idx_train))]
+
+        self.adj_train = adj_full[np.ix_(idx_train, idx_train)]
+        self.adj_val = adj_full[np.ix_(idx_val, idx_val)]
+        self.adj_test = adj_full[np.ix_(idx_test, idx_test)]
+
+        feat = np.load(dataset_str+'feats.npy')
+        # ---- normalize feat ----
+        feat_train = feat[idx_train]
+        scaler = StandardScaler()
+        scaler.fit(feat_train)
+        feat = scaler.transform(feat)
+
+        self.feat_train = feat[idx_train]
+        self.feat_val = feat[idx_val]
+        self.feat_test = feat[idx_test]
+
+        class_map = json.load(open(dataset_str + 'class_map.json','r'))
+        labels = self.process_labels(class_map)
+
+        self.labels_train = labels[idx_train]
+        self.labels_val = labels[idx_val]
+        self.labels_test = labels[idx_test]
+
+        self.data_full = GraphData(adj_full, feat, labels, idx_train, idx_val, idx_test)
+        self.class_dict = None
+        self.class_dict2 = None
+
+        self.adj_full = adj_full
+        self.feat_full = feat
+        self.labels_full = labels
+        self.idx_train = np.array(idx_train)
+        self.idx_val = np.array(idx_val)
+        self.idx_test = np.array(idx_test)
+        self.samplers = None
+
+    def process_labels(self, class_map):
+        """
+        setup vertex property map for output classests
+        """
+        num_vertices = self.nnodes
+        if isinstance(list(class_map.values())[0], list):
+            num_classes = len(list(class_map.values())[0])
+            self.nclass = num_classes
+            class_arr = np.zeros((num_vertices, num_classes))
+            for k,v in class_map.items():
+                class_arr[int(k)] = v
+        else:
+            class_arr = np.zeros(num_vertices, dtype=np.int)
+            for k, v in class_map.items():
+                class_arr[int(k)] = v
+            class_arr = class_arr - class_arr.min()
+            self.nclass = max(class_arr) + 1
+        return class_arr
+
+    def retrieve_class(self, c, num=256):
+        if self.class_dict is None:
+            self.class_dict = {}
+            for i in range(self.nclass):
+                self.class_dict['class_%s'%i] = (self.labels_train == i)
+        idx = np.arange(len(self.labels_train))
+        idx = idx[self.class_dict['class_%s'%c]]
+        return np.random.permutation(idx)[:num]
+
+    def retrieve_class_sampler(self, c, adj, transductive, num=256, args=None):
+        if args.nlayers == 1:
+            sizes = [30]
+        if args.nlayers == 2:
+            if args.dataset in ['reddit', 'flickr']:
+                if args.option == 0:
+                    sizes = [15, 8]
+                if args.option == 1:
+                    sizes = [20, 10]
+                if args.option == 2:
+                    sizes = [25, 10]
+            else:
+                sizes = [10, 5]
+
+        if self.class_dict2 is None:
+            print(sizes)
+            self.class_dict2 = {}
+            for i in range(self.nclass):
+                if transductive:
+                    idx_train = np.array(self.idx_train)
+                    idx = idx_train[self.labels_train == i]
+                else:
+                    idx = np.arange(len(self.labels_train))[self.labels_train==i]
+                self.class_dict2[i] = idx
+
+        if self.samplers is None:
+            self.samplers = []
+            for i in range(self.nclass):
+                node_idx = torch.LongTensor(self.class_dict2[i])
+                if len(node_idx) == 0:
+                    continue
+
+                self.samplers.append(NeighborSampler(adj,
+                                    node_idx=node_idx,
+                                    sizes=sizes, batch_size=num,
+                                    num_workers=8, return_e_id=False,
+                                    num_nodes=adj.size(0),
+                                    shuffle=True))
+        batch = np.random.permutation(self.class_dict2[c])[:num]
+        out = self.samplers[c].sample(batch)
+        return out
+
+
+class GraphData:
+
+    def __init__(self, adj, features, labels, idx_train, idx_val, idx_test):
+        self.adj = adj
+        self.features = features
+        self.labels = labels
+        self.idx_train = idx_train
+        self.idx_val = idx_val
+        self.idx_test = idx_test
+
+
+class Data2Pyg:
+
+    def __init__(self, data, device='cuda', transform=None, **kwargs):
+        self.data_train = Dpr2Pyg(data.data_train, transform=transform)[0].to(device)
+        self.data_val = Dpr2Pyg(data.data_val, transform=transform)[0].to(device)
+        self.data_test = Dpr2Pyg(data.data_test, transform=transform)[0].to(device)
+        self.nclass = data.nclass
+        self.nfeat = data.nfeat
+        self.class_dict = None
+
+    def retrieve_class(self, c, num=256):
+        if self.class_dict is None:
+            self.class_dict = {}
+            for i in range(self.nclass):
+                self.class_dict['class_%s'%i] = (self.data_train.y == i).cpu().numpy()
+        idx = np.arange(len(self.data_train.y))
+        idx = idx[self.class_dict['class_%s'%c]]
+        return np.random.permutation(idx)[:num]
+
+
+class Dpr2Pyg(InMemoryDataset):
+
+    def __init__(self, dpr_data, transform=None, **kwargs):
+        root = 'data/' # dummy root; does not mean anything
+        self.dpr_data = dpr_data
+        super(Dpr2Pyg, self).__init__(root, transform)
+        pyg_data = self.process()
+        self.data, self.slices = self.collate([pyg_data])
+        self.transform = transform
+
+    def process(self):
+        dpr_data = self.dpr_data
+        edge_index = torch.LongTensor(dpr_data.adj.nonzero())
+        # by default, the features in pyg data is dense
+        if sp.issparse(dpr_data.features):
+            x = torch.FloatTensor(dpr_data.features.todense()).float()
+        else:
+            x = torch.FloatTensor(dpr_data.features).float()
+        y = torch.LongTensor(dpr_data.labels)
+        data = Data(x=x, edge_index=edge_index, y=y)
+        data.train_mask = None
+        data.val_mask = None
+        data.test_mask = None
+        return data
+
+
+    def get(self, idx):
+        data = self.data.__class__()
+
+        if hasattr(self.data, '__num_nodes__'):
+            data.num_nodes = self.data.__num_nodes__[idx]
+
+        for key in self.data.keys:
+            item, slices = self.data[key], self.slices[key]
+            s = list(repeat(slice(None), item.dim()))
+            s[self.data.__cat_dim__(key, item)] = slice(slices[idx],
+                                                   slices[idx + 1])
+            data[key] = item[s]
+        return data
+
+    @property
+    def raw_file_names(self):
+        return ['some_file_1', 'some_file_2', ...]
+
+    @property
+    def processed_file_names(self):
+        return ['data.pt']
+
+    def _download(self):
+        pass
+
+

requirements.txt

@@ -1,4 +1,4 @@
-torch
 torch-scatter
 torch-sparse
-torch-geometric
+torch-geometric
+rich

src/2.1_lrmc_bilevel.py

@@ -1,6 +1,4 @@
-# lrmc_bilevel.py
 # Top-1 LRMC ablation: one-cluster pooling vs. plain GCN on Planetoid (e.g., Cora)
-# Requires: torch, torch_geometric, torch_scatter, torch_sparse
 #
 # Usage examples:
 #   python lrmc_bilevel.py --dataset Cora --seeds /path/to/lrmc_seeds.json --variant baseline
@@ -28,6 +26,8 @@ from torch_sparse import coalesce, spspmm
 from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import GCNConv
 
+from rich import print
+
 
 # ---------------------------
 # Utilities: edges and seeds

src/2.2_lrmc_bilevel.py

@@ -1,4 +1,3 @@
-# 2.1_lrmc_bilevel.py
 # Top-1 LRMC ablation with debug guards so seeds differences are visible.
 # Requires: torch, torch_geometric, torch_scatter, torch_sparse
 
@@ -16,6 +15,8 @@ from torch_sparse import coalesce, spspmm
 from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import GCNConv
 
+from rich import print
+
 
 # ---------------------------
 # Utilities: edges and seeds

src/2.3_lrmc_bilevel.py

@@ -1,4 +1,3 @@
-# 2.3_lrmc_bilevel.py
 # Top-1 LRMC ablation with: cluster refinement (k-core), gated residual fusion,
 # sparsified cluster graph (drop self-loops + per-row top-k), and A + γA² mix.
 # Requires: torch, torch_geometric, torch_scatter, torch_sparse
@@ -17,6 +16,8 @@ from torch_sparse import coalesce, spspmm
 from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import GCNConv
 
+from rich import print
+
 
 # ---------------------------
 # Utilities: edges and seeds

src/2.4_lrmc_bilevel.py

@@ -1,4 +1,3 @@
-# lrmc_bilevel.py
 # Top-1 LRMC ablation: one-cluster pooling vs. plain GCN on Planetoid (e.g., Cora)
 # Requires: torch, torch_geometric, torch_scatter, torch_sparse
 #
@@ -30,6 +29,8 @@ from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import GCNConv
 from torch_geometric.utils import subgraph, degree # Added for stability score
 
+from rich import print
+
 # ---------------------------
 # Utilities: edges and seeds
 # ---------------------------

src/2.5_lrmc_bilevel.py

@@ -0,0 +1,540 @@
+"""
+L-RMC Anchored GCN vs. Plain GCN (dynamic robustness evaluation)
+==============================================================
+
+This script trains a baseline two‑layer GCN and a new **anchor‑gated** GCN on
+Planetoid citation networks (Cora/Citeseer/Pubmed).  The anchor‑gated GCN uses
+the top‑1 L‑RMC cluster (loaded from a provided JSON file) as a *decentralized
+core*.  During message passing it blends standard neighborhood aggregation
+(`h_base`) with aggregation restricted to the core (`h_core`) via a per‑node
+gating network.  Cross‑boundary edges are optionally down‑weighted by a
+damping factor `γ`.
+
+After training on the static graph, the script evaluates *robustness over
+time*.  Starting from the original adjacency, it repeatedly performs random
+edge rewires (removes a fraction of existing edges and adds the same number
+of random new edges) and measures test accuracy at each step **without
+retraining**.  The area under the accuracy–time curve (AUC‑AT) is reported
+for both the baseline and the anchored model.  A higher AUC‑AT indicates
+longer resilience to graph churn.
+
+Usage examples::
+
+    # Train only baseline and report dynamic AUC
+    python 2.5_lrmc_bilevel.py --dataset Cora --seeds path/to/lrmc_seeds.json --variant baseline
+
+    # Train baseline and anchor models, evaluate AUC‑over‑time on 30 steps with 5% rewiring
+    python 2.5_lrmc_bilevel.py --dataset Cora --seeds path/to/lrmc_seeds.json --variant anchor \
+        --dynamic_steps 30 --flip_fraction 0.05 --gamma 0.8
+
+Notes
+-----
+* The seeds JSON must contain an entry ``"clusters"`` with a list of clusters; the
+  cluster with maximum (score, size) is chosen as the core.
+* For fairness, both models are trained on the identical training mask and
+  evaluated on the same dynamic perturbations.
+* Random rewiring is undirected: an edge (u,v) is treated as the same as (v,u).
+* Cross‑boundary damping and the gating network use only structural
+  information; features are left unchanged during perturbations.
+"""
+
+import argparse
+import json
+import random
+from pathlib import Path
+from typing import Tuple, List, Optional, Set
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch_geometric.datasets import Planetoid
+from torch_geometric.nn import GCNConv
+
+from rich import print
+
+# -----------------------------------------------------------------------------
+#  Utilities for loading LRMC core assignment
+# -----------------------------------------------------------------------------
+
+def _pick_top1_cluster(obj: dict) -> List[int]:
+    """
+    From LRMC JSON with structure {"clusters":[{"seed_nodes":[...],"score":float,...},...]}
+    choose the cluster with the highest (score, size) and return its members as
+    0‑indexed integers.  If no clusters exist, returns an empty list.
+    """
+    clusters = obj.get("clusters", [])
+    if not clusters:
+        return []
+    # Choose by highest score, tie‑break by size
+    best = max(clusters, key=lambda c: (float(c.get("score", 0.0)), len(c.get("seed_nodes", []))))
+    return [nid - 1 for nid in best.get("seed_nodes", [])]
+
+
+def load_top1_assignment(seeds_json: str, n_nodes: int) -> Tuple[Tensor, Tensor]:
+    """
+    Given a path to the LRMC seeds JSON and total number of nodes, returns:
+
+    * core_mask: bool Tensor of shape [N] where True indicates membership in the
+      top‑1 LRMC cluster.
+    * core_nodes: Long Tensor containing the indices of the core nodes.
+
+    Nodes not in the core form the periphery.  If the JSON has no clusters,
+    the core is empty.
+    """
+    obj = json.loads(Path(seeds_json).read_text())
+    core_list = _pick_top1_cluster(obj)
+    core_nodes = torch.tensor(sorted(set(core_list)), dtype=torch.long)
+    core_mask = torch.zeros(n_nodes, dtype=torch.bool)
+    if core_nodes.numel() > 0:
+        core_mask[core_nodes] = True
+    return core_mask, core_nodes
+
+
+# -----------------------------------------------------------------------------
+#  Baseline GCN: standard two‑layer GCN
+# -----------------------------------------------------------------------------
+
+class GCN2(nn.Module):
+    """Plain 2‑layer GCN (baseline)."""
+
+    def __init__(self, in_dim: int, hid_dim: int, out_dim: int, dropout: float = 0.5):
+        super().__init__()
+        self.conv1 = GCNConv(in_dim, hid_dim)
+        self.conv2 = GCNConv(hid_dim, out_dim)
+        self.dropout = dropout
+
+    def forward(self, x: Tensor, edge_index: Tensor, edge_weight: Optional[Tensor] = None) -> Tensor:
+        # Use self loops implicitly (GCNConv defaults add_self_loops=True)
+        x = F.relu(self.conv1(x, edge_index, edge_weight))
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.conv2(x, edge_index, edge_weight)
+        return x
+
+
+# -----------------------------------------------------------------------------
+#  Anchor‑gated GCN
+# -----------------------------------------------------------------------------
+
+class AnchorGCN(nn.Module):
+    """
+    A two‑layer GCN that injects a core‑restricted aggregation channel and
+    down‑weights edges crossing the core boundary.  After the first GCN layer
+    computes base features, a gating network mixes them with features
+    aggregated only among core neighbors.
+
+    Parameters
+    ----------
+    in_dim : int
+        Dimensionality of input node features.
+    hid_dim : int
+        Dimensionality of hidden layer.
+    out_dim : int
+        Number of output classes.
+    core_mask : Tensor[bool]
+        Boolean mask indicating which nodes belong to the L‑RMC core.
+    gamma : float, optional
+        Damping factor for edges that connect core and non‑core nodes.
+        Values <1.0 reduce the influence of boundary edges.  Default is 1.0
+        (no damping).
+    dropout : float, optional
+        Dropout probability applied after the first layer.
+    """
+
+    def __init__(self,
+                 in_dim: int,
+                 hid_dim: int,
+                 out_dim: int,
+                 core_mask: Tensor,
+                 gamma: float = 1.0,
+                 dropout: float = 0.5):
+        super().__init__()
+        self.core_mask = core_mask.clone().detach()
+        self.gamma = float(gamma)
+        self.dropout = dropout
+
+        # Base and core convolutions for the first layer
+        # Base conv uses self loops; core conv disables self loops to avoid
+        # spurious core contributions on non‑core nodes
+        self.base1 = GCNConv(in_dim, hid_dim, add_self_loops=True)
+        self.core1 = GCNConv(in_dim, hid_dim, add_self_loops=False)
+
+        # Second layer: standard GCN on mixed features
+        self.conv2 = GCNConv(hid_dim, out_dim)
+
+        # Gating network: maps structural features to α ∈ [0,1]
+        self.gate = nn.Sequential(
+            nn.Linear(3, 16),
+            nn.ReLU(),
+            nn.Linear(16, 1),
+            nn.Sigmoid(),
+        )
+
+    def _compute_edge_weights(self, edge_index: Tensor) -> Tensor:
+        """
+        Given an edge index (two‑row tensor), return a weight tensor of ones
+        multiplied by ``gamma`` for edges with exactly one endpoint in the core.
+        Self loops (if present) are untouched.  Edge weights are 1 for base
+        edges and <1 for cross‑boundary edges.
+        """
+        if self.gamma >= 1.0:
+            return torch.ones(edge_index.size(1), dtype=torch.float32, device=edge_index.device)
+        src, dst = edge_index[0], edge_index[1]
+        in_core_src = self.core_mask[src]
+        in_core_dst = self.core_mask[dst]
+        cross = in_core_src ^ in_core_dst
+        w = torch.ones(edge_index.size(1), dtype=torch.float32, device=edge_index.device)
+        w[cross] *= self.gamma
+        return w
+
+    def _compute_structural_features(self, edge_index: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Compute structural features used by the gating network:
+
+        * `in_core`       – 1 if node in core, else 0
+        * `frac_core_nbrs` – fraction of neighbors that are in the core
+        * `is_boundary`    – 1 if node has both core and non‑core neighbors
+
+        The features are returned as a tuple of three tensors of shape [N,1].
+        Nodes with zero degree get frac_core_nbrs=0 and is_boundary=0.
+        """
+        N = self.core_mask.size(0)
+        device = edge_index.device
+        # Degree and core neighbor counts
+        src = edge_index[0]
+        dst = edge_index[1]
+        deg = torch.zeros(N, dtype=torch.float32, device=device)
+        core_deg = torch.zeros(N, dtype=torch.float32, device=device)
+        # Count contributions of directed edges; duplicates will double‑count but
+        # the ratio remains stable if the graph is symmetric.
+        deg.index_add_(0, src, torch.ones_like(src, dtype=torch.float32))
+        # Count core neighbors: only increment source if destination is core
+        core_flags = self.core_mask[dst].float()
+        core_deg.index_add_(0, src, core_flags)
+        # Avoid division by zero
+        frac_core = torch.zeros(N, dtype=torch.float32, device=device)
+        nonzero = deg > 0
+        frac_core[nonzero] = core_deg[nonzero] / deg[nonzero]
+        # Determine boundary: at least one core neighbor AND at least one non‑core neighbor
+        has_core = core_deg > 0
+        has_non_core = (deg - core_deg) > 0
+        is_boundary = (has_core & has_non_core).float()
+        in_core = self.core_mask.float()
+        return in_core.view(-1, 1), frac_core.view(-1, 1), is_boundary.view(-1, 1)
+
+    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
+        # Compute dynamic edge weights (for base channels) using damping
+        w = self._compute_edge_weights(edge_index)
+        # First layer: base aggregation (standard neighbors with self loops)
+        h_base = self.base1(x, edge_index, w)
+        h_base = F.relu(h_base)
+
+        # First layer: core aggregation (only core neighbors, no self loops)
+        # Extract edges where both endpoints are core
+        src, dst = edge_index
+        mask_core_edges = self.core_mask[src] & self.core_mask[dst]
+        ei_core = edge_index[:, mask_core_edges]
+        # If no core edges exist, h_core will be zeros
+        if ei_core.numel() == 0:
+            h_core = torch.zeros_like(h_base)
+        else:
+            h_core = self.core1(x, ei_core)
+            h_core = F.relu(h_core)
+
+        # Structural features for gating
+        in_core, frac_core, is_boundary = self._compute_structural_features(edge_index)
+        feats = torch.cat([in_core, frac_core, is_boundary], dim=1)
+        alpha = self.gate(feats).view(-1)  # shape [N]
+        # Force α=0 for nodes with no core neighbors to avoid modifying true periphery.
+        # Nodes with frac_core == 0 have zero core neighbors by construction.
+        no_core_neighbors = (frac_core.view(-1) == 0)
+        alpha = torch.where(no_core_neighbors, torch.zeros_like(alpha), alpha)
+
+        # Mix base and core features; h_final = h_base + α (h_core - h_base)
+        # Equivalent to (1-α)*h_base + α*h_core
+        h1 = h_base + alpha.unsqueeze(1) * (h_core - h_base)
+
+        h1 = F.dropout(h1, p=self.dropout, training=self.training)
+        # Second layer: standard GCN with the same damping weights
+        out = self.conv2(h1, edge_index, w)
+        return out
+
+
+# The deg_for_division helper is no longer used but left here for completeness.
+def deg_for_division(edge_index: Tensor, num_nodes: int) -> Tensor:
+    src = edge_index[0]
+    deg = torch.zeros(num_nodes, dtype=torch.float32, device=edge_index.device)
+    deg.index_add_(0, src, torch.ones_like(src, dtype=torch.float32))
+    return deg
+
+
+# -----------------------------------------------------------------------------
+#  Training and evaluation routines
+# -----------------------------------------------------------------------------
+
+@torch.no_grad()
+def accuracy(logits: Tensor, y: Tensor, mask: Tensor) -> float:
+    """Compute accuracy of the predictions over the mask."""
+    pred = logits[mask].argmax(dim=1)
+    return (pred == y[mask]).float().mean().item()
+
+
+def train_model(model: nn.Module,
+                data,
+                epochs: int = 200,
+                lr: float = 0.01,
+                weight_decay: float = 5e-4) -> None:
+    """Standard training loop for either baseline or anchor models."""
+    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+    best_val = 0.0
+    best_state = None
+    for ep in range(1, epochs + 1):
+        model.train()
+        opt.zero_grad(set_to_none=True)
+        logits = model(data.x, data.edge_index)
+        loss = F.cross_entropy(logits[data.train_mask], data.y[data.train_mask])
+        loss.backward()
+        opt.step()
+        # Evaluate on validation
+        model.eval()
+        logits_val = model(data.x, data.edge_index)
+        val_acc = accuracy(logits_val, data.y, data.val_mask)
+        if val_acc > best_val:
+            best_val = val_acc
+            best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
+    if best_state is not None:
+        model.load_state_dict(best_state)
+    model.eval()
+
+
+def evaluate_model(model: nn.Module, data) -> dict:
+    """Evaluate a trained model on train, val, and test masks."""
+    model.eval()
+    logits = model(data.x, data.edge_index)
+    return {
+        "train": accuracy(logits, data.y, data.train_mask),
+        "val": accuracy(logits, data.y, data.val_mask),
+        "test": accuracy(logits, data.y, data.test_mask),
+    }
+
+
+# -----------------------------------------------------------------------------
+#  Dynamic graph perturbation utilities
+# -----------------------------------------------------------------------------
+
+def undirected_edge_set(edge_index: Tensor) -> Set[Tuple[int, int]]:
+    """
+    Convert a directed edge index into a set of undirected edges represented
+    as (u,v) tuples with u < v.  Self loops are ignored.
+    """
+    edges = set()
+    src = edge_index[0].tolist()
+    dst = edge_index[1].tolist()
+    for u, v in zip(src, dst):
+        if u == v:
+            continue
+        a, b = (u, v) if u < v else (v, u)
+        edges.add((a, b))
+    return edges
+
+
+def edge_set_to_index(edges: Set[Tuple[int, int]], num_nodes: int) -> Tensor:
+    """
+    Convert an undirected edge set into a directed edge_index tensor of shape
+    [2, 2*|edges|] by adding both (u,v) and (v,u) for each undirected edge.
+    Self loops are omitted; GCNConv adds them automatically.
+    """
+    if not edges:
+        return torch.empty(2, 0, dtype=torch.long)
+    src_list = []
+    dst_list = []
+    for u, v in edges:
+        src_list.extend([u, v])
+        dst_list.extend([v, u])
+    edge_index = torch.tensor([src_list, dst_list], dtype=torch.long)
+    return edge_index
+
+
+def random_rewire(edges: Set[Tuple[int, int]], num_nodes: int, n_changes: int, rng: random.Random) -> Set[Tuple[int, int]]:
+    """
+    Perform n_changes edge removals and n_changes edge additions on the given
+    undirected edge set.  For each change we remove a random existing edge and
+    add a random new edge (u,v) not already present.  Self loops are never
+    added.  Duplicate additions are skipped.
+    """
+    edges = set(edges)  # copy
+    # If there are fewer edges than n_changes, adjust
+    n_changes = min(n_changes, len(edges))
+    # Remove random edges
+    to_remove = rng.sample(list(edges), n_changes)
+    for e in to_remove:
+        edges.remove(e)
+    # Add random new edges
+    added = 0
+    attempts = 0
+    while added < n_changes and attempts < n_changes * 10:
+        u = rng.randrange(num_nodes)
+        v = rng.randrange(num_nodes)
+        if u == v:
+            attempts += 1
+            continue
+        a, b = (u, v) if u < v else (v, u)
+        if (a, b) not in edges:
+            edges.add((a, b))
+            added += 1
+        attempts += 1
+    return edges
+
+
+def auc_over_time(acc_list: List[float]) -> float:
+    """
+    Compute the area under an accuracy–time curve using the trapezoidal rule.
+    ``acc_list`` should contain the accuracies at t=0,1,...,T.  The AUC is
+    normalized by T so that a perfect score of 1.0 yields AUC=1.0.
+    """
+    if not acc_list:
+        return 0.0
+    area = 0.0
+    for i in range(1, len(acc_list)):
+        area += (acc_list[i] + acc_list[i-1]) / 2.0
+    return area / (len(acc_list) - 1)
+
+
+def evaluate_dynamic_auc(model: nn.Module,
+                         data,
+                         core_mask: Tensor,
+                         steps: int = 30,
+                         flip_fraction: float = 0.05,
+                         rng_seed: int = 1234) -> List[float]:
+    """
+    Evaluate a model's test accuracy over a sequence of random edge rewiring steps.
+
+    Parameters
+    ----------
+    model : nn.Module
+        A trained model that accepts (x, edge_index) and returns logits.
+    data : Data
+        PyG data object with attributes x, y, test_mask.  ``data.edge_index``
+        provides the initial adjacency.
+    core_mask : Tensor[bool]
+        Boolean mask indicating core nodes (used for gating during evaluation).
+        The baseline model ignores it.
+    steps : int, optional
+        Number of rewiring steps to perform.  The accuracy at t=0 is computed
+        before any rewiring.  Default: 30.
+    flip_fraction : float, optional
+        Fraction of edges to remove/add at each step.  For example, 0.05
+        rewires 5% of existing edges per step.  Default: 0.05.
+    rng_seed : int, optional
+        Random seed for reproducibility.  Default: 1234.
+
+    Returns
+    -------
+    List[float]
+        A list of length ``steps+1`` containing the test accuracy at each
+        iteration (including t=0).
+    """
+    # Convert initial edge_index to undirected edge set
+    base_edges = undirected_edge_set(data.edge_index)
+    num_edges = len(base_edges)
+    # Determine number of changes per step
+    n_changes = max(1, int(flip_fraction * num_edges))
+    # Clone model state so we don't accidentally update it during evaluation
+    model.eval()
+    # Random generator
+    rng = random.Random(rng_seed)
+    # Copy of edges for dynamic modification
+    cur_edges = set(base_edges)
+    accuracies = []
+    # Evaluate at t=0
+    ei = edge_set_to_index(cur_edges, data.num_nodes)
+    # Because PyG expects a tensor on the same device as data.x
+    ei = ei.to(data.x.device)
+    logits = model(data.x, ei)
+    accuracies.append(accuracy(logits, data.y, data.test_mask))
+    # Perform rewiring steps
+    for t in range(1, steps + 1):
+        cur_edges = random_rewire(cur_edges, data.num_nodes, n_changes, rng)
+        ei = edge_set_to_index(cur_edges, data.num_nodes).to(data.x.device)
+        logits = model(data.x, ei)
+        acc = accuracy(logits, data.y, data.test_mask)
+        accuracies.append(acc)
+    return accuracies
+
+
+# -----------------------------------------------------------------------------
+#  Main entrypoint
+# -----------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="L‑RMC anchored GCN vs. baseline with dynamic evaluation.")
+    parser.add_argument("--dataset", required=True, choices=["Cora", "Citeseer", "Pubmed"],
+                        help="Planetoid dataset to load.")
+    parser.add_argument("--seeds", required=True, help="Path to LRMC seeds JSON (for core extraction).")
+    parser.add_argument("--variant", choices=["baseline", "anchor"], default="anchor", help="Which variant to run.")
+    parser.add_argument("--hidden", type=int, default=64, help="Hidden dimension.")
+    parser.add_argument("--epochs", type=int, default=200, help="Number of training epochs.")
+    parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
+    parser.add_argument("--wd", type=float, default=5e-4, help="Weight decay (L2).")
+    parser.add_argument("--dropout", type=float, default=0.5, help="Dropout probability.")
+    parser.add_argument("--gamma", type=float, default=1.0, help="Damping factor γ for cross‑boundary edges (anchor only).")
+    parser.add_argument("--dynamic_steps", type=int, default=30, help="Number of dynamic rewiring steps for AUC evaluation.")
+    parser.add_argument("--flip_fraction", type=float, default=0.05, help="Fraction of edges rewired at each step.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for PyTorch.")
+    args = parser.parse_args()
+
+    # Set seeds
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+
+    # Load dataset
+    dataset = Planetoid(root=f"./data/{args.dataset}", name=args.dataset)
+    data = dataset[0]
+    in_dim = dataset.num_node_features
+    out_dim = dataset.num_classes
+    num_nodes = data.num_nodes
+
+    # Load core assignment
+    core_mask, core_nodes = load_top1_assignment(args.seeds, num_nodes)
+    print(f"Loaded core of size {core_nodes.numel()} from {args.seeds}.")
+
+    if args.variant == "baseline":
+        # Train baseline only
+        baseline = GCN2(in_dim, args.hidden, out_dim, dropout=args.dropout)
+        train_model(baseline, data, epochs=args.epochs, lr=args.lr, weight_decay=args.wd)
+        res = evaluate_model(baseline, data)
+        print(f"Baseline GCN: train={res['train']:.4f}  val={res['val']:.4f}  test={res['test']:.4f}")
+        # Evaluate dynamic AUC
+        accs = evaluate_dynamic_auc(baseline, data, core_mask, steps=args.dynamic_steps,
+                                    flip_fraction=args.flip_fraction, rng_seed=args.seed)
+        auc = auc_over_time(accs)
+        print(f"Baseline dynamic AUC‑AT (steps={args.dynamic_steps}, flip={args.flip_fraction}): {auc:.4f}")
+        return
+
+    # ----- Train both baseline and anchor variants -----
+    # Baseline
+    baseline = GCN2(in_dim, args.hidden, out_dim, dropout=args.dropout)
+    train_model(baseline, data, epochs=args.epochs, lr=args.lr, weight_decay=args.wd)
+    res_base = evaluate_model(baseline, data)
+    print(f"Baseline GCN: train={res_base['train']:.4f}  val={res_base['val']:.4f}  test={res_base['test']:.4f}")
+    # Anchor model
+    anchor = AnchorGCN(in_dim, args.hidden, out_dim,
+                       core_mask=core_mask,
+                       gamma=args.gamma,
+                       dropout=args.dropout)
+    train_model(anchor, data, epochs=args.epochs, lr=args.lr, weight_decay=args.wd)
+    res_anchor = evaluate_model(anchor, data)
+    print(f"Anchor‑GCN:  train={res_anchor['train']:.4f}  val={res_anchor['val']:.4f}  test={res_anchor['test']:.4f}")
+    # Dynamic evaluation
+    accs_base = evaluate_dynamic_auc(baseline, data, core_mask, steps=args.dynamic_steps,
+                                     flip_fraction=args.flip_fraction, rng_seed=args.seed)
+    accs_anchor = evaluate_dynamic_auc(anchor, data, core_mask, steps=args.dynamic_steps,
+                                       flip_fraction=args.flip_fraction, rng_seed=args.seed)
+    auc_base = auc_over_time(accs_base)
+    auc_anchor = auc_over_time(accs_anchor)
+    print(f"Dynamic AUC‑AT (steps={args.dynamic_steps}, flip={args.flip_fraction}):")
+    print(f"  Baseline : {auc_base:.4f}\n  Anchor   : {auc_anchor:.4f}")
+
+
+if __name__ == "__main__":
+    main()

src/2.6_lrmc_summary.py

@@ -0,0 +1,691 @@
+"""
+Train a GCN on an L‑RMC subgraph and compare to a full‑graph baseline.
+
+Modes:
+  - core_mode=forward : Train on core subgraph, then forward on full graph (your current approach).
+  - core_mode=appnp   : Train on core subgraph, then seed logits on core and APPNP‑propagate on full graph.
+
+Extras:
+  - --expand_core_with_train : Make sure all training labels lie inside the core
+                               (C' = C ∪ train_idx) for fair train‑time comparison.
+  - --warm_ft_epochs N       : Optional short finetune on the full graph starting
+                               from the core model's weights (measure time‑to‑target).
+
+It prints:
+  - Dataset stats
+  - Core size and coverage of train/val/test inside the core
+  - Train/Val/Test accuracy for baseline and core model
+  - Wall‑clock times
+"""
+
+import argparse
+import json
+import time
+import random
+from statistics import mean, stdev
+from pathlib import Path
+from typing import Dict
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch_geometric.datasets import Planetoid
+from torch_geometric.nn import GCNConv, APPNP
+from torch_geometric.utils import subgraph
+
+# ------------------------------------------------------------
+# Rich imports
+# ------------------------------------------------------------
+from rich.console import Console
+from rich.table import Table
+from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
+
+# Rich console instance
+console = Console()
+
+# ------------------------------------------------------------
+# Utilities
+# ------------------------------------------------------------
+def load_top1_assignment(seeds_json: str, n_nodes: int) -> torch.Tensor:
+    """
+    seeds_json format (expected):
+      {"clusters": [{"seed_nodes":[...], "score": float, ...}, ...]}
+    We pick the cluster with max (score, size) and return a boolean core mask.
+
+    Always assume that the seeds json nodes are 1-indexed.
+    """
+    obj = json.loads(Path(seeds_json).read_text())
+    clusters = obj.get("clusters", [])
+    if not clusters:
+        return torch.zeros(n_nodes, dtype=torch.bool)
+    best = max(clusters, key=lambda c: (float(c.get("score", 0.0)), len(c.get("seed_nodes", []))))
+    ids = best.get("seed_nodes", [])
+    ids = [int(x) - 1 for x in ids]  # Convert 1-indexed to 0-indexed
+    ids = sorted(set([i for i in ids if 0 <= i < n_nodes]))
+    mask = torch.zeros(n_nodes, dtype=torch.bool)
+    if ids:
+        mask[torch.tensor(ids, dtype=torch.long)] = True
+    return mask
+
+
+def coverage_counts(core_mask: torch.Tensor, train_mask: torch.Tensor,
+                    val_mask: torch.Tensor, test_mask: torch.Tensor) -> Dict[str, int]:
+    return {
+        "core_size": int(core_mask.sum().item()),
+        "train_in_core": int((core_mask & train_mask).sum().item()),
+        "val_in_core": int((core_mask & val_mask).sum().item()),
+        "test_in_core": int((core_mask & test_mask).sum().item()),
+    }
+
+
+def accuracy(logits: Tensor, y: Tensor, mask: Tensor) -> float:
+    pred = logits[mask].argmax(dim=1)
+    return (pred == y[mask]).float().mean().item()
+
+
+def set_seed(seed: int):
+    """Set random seeds for reproducibility across runs."""
+    random.seed(seed)
+    try:
+        import numpy as np  # optional
+        np.random.seed(seed)
+    except Exception:
+        pass
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    # Make CUDA/CuDNN deterministic where applicable
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+# ------------------------------------------------------------
+# Models
+# ------------------------------------------------------------
+class GCN2(nn.Module):
+    def __init__(self, in_dim: int, hid: int, out_dim: int, dropout: float = 0.5):
+        super().__init__()
+        self.c1 = GCNConv(in_dim, hid)
+        self.c2 = GCNConv(hid, out_dim)
+        self.dropout = dropout
+
+    def forward(self, x, ei):
+        x = self.c1(x, ei)
+        x = torch.relu(x)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.c2(x, ei)
+        return x
+
+
+# ------------------------------------------------------------
+# Training / evaluation
+# ------------------------------------------------------------
+@torch.no_grad()
+def eval_all(model: nn.Module, data) -> Dict[str, float]:
+    model.eval()
+    logits = model(data.x, data.edge_index)
+    return {
+        "train": accuracy(logits, data.y, data.train_mask),
+        "val": accuracy(logits, data.y, data.val_mask),
+        "test": accuracy(logits, data.y, data.test_mask),
+    }
+
+
+def train(model: nn.Module, data, epochs=200, lr=0.01, wd=5e-4, patience=100):
+    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
+    best, best_state, bad = -1.0, None, 0
+
+    # Optional progress bar
+    with Progress(
+        SpinnerColumn(),
+        "[progress.description]{task.description}",
+        TimeElapsedColumn(),
+        transient=True,
+    ) as progress:
+        task = progress.add_task("Training", total=epochs)
+
+        for ep in range(1, epochs + 1):
+            model.train()
+            opt.zero_grad(set_to_none=True)
+            out = model(data.x, data.edge_index)
+            loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
+            loss.backward()
+            opt.step()
+
+            # early stop on val
+            with torch.no_grad():
+                val = accuracy(model(data.x, data.edge_index), data.y, data.val_mask)
+
+            if val > best:
+                best, bad = val, 0
+                best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
+            else:
+                bad += 1
+                if bad >= patience:
+                    break
+
+            progress.update(task, advance=1, description=f"Epoch {ep} | val={val:.4f}")
+
+    if best_state is not None:
+        model.load_state_dict(best_state)
+    model.eval()
+
+
+def subset_data(data, nodes_idx: torch.Tensor):
+    """
+    Build an induced subgraph on 'nodes_idx'. Keeps x,y,masks restricted to that set.
+    Returns a shallow copy with edge_index/feature/labels/masks sliced.
+    """
+    nodes_idx = nodes_idx.to(torch.long)
+    sub_ei, _ = subgraph(nodes_idx, data.edge_index, relabel_nodes=True, num_nodes=data.num_nodes)
+    sub = type(data)()
+    sub.x = data.x[nodes_idx]
+    sub.y = data.y[nodes_idx]
+    sub.train_mask = data.train_mask[nodes_idx]
+    sub.val_mask = data.val_mask[nodes_idx]
+    sub.test_mask = data.test_mask[nodes_idx]
+    sub.edge_index = sub_ei
+    sub.num_nodes = sub.x.size(0)
+    return sub
+
+
+# ------------------------------------------------------------
+# APPNP seeding (Mode B)
+# ------------------------------------------------------------
+def appnp_seed_propagate(logits_seed: Tensor, edge_index: Tensor, K=10, alpha=0.1) -> Tensor:
+    """
+    logits_seed is [N, C] where rows outside the core are zeros.
+    We propagate these logits with APPNP to fill the graph.
+    """
+    appnp = APPNP(K=K, alpha=alpha)  # no trainable params
+    return appnp(logits_seed, edge_index)
+
+
+# ------------------------------------------------------------
+# Main
+# ------------------------------------------------------------
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dataset", required=True, choices=["Cora", "Citeseer", "Pubmed"])
+    p.add_argument("--seeds", required=True, help="Path to LRMC seeds JSON")
+    p.add_argument("--hidden", type=int, default=64)
+    p.add_argument("--dropout", type=float, default=0.5)
+    p.add_argument("--epochs", type=int, default=200)
+    p.add_argument("--lr", type=float, default=0.01)
+    p.add_argument("--wd", type=float, default=5e-4)
+    p.add_argument("--patience", type=int, default=100)
+    p.add_argument("--core_mode", choices=["forward", "appnp"], default="forward",
+                   help="How to evaluate the core model on the full graph.")
+    p.add_argument("--alpha", type=float, default=0.1, help="APPNP teleport prob (Mode B).")
+    p.add_argument("--K", type=int, default=10, help="APPNP steps (Mode B).")
+    p.add_argument("--expand_core_with_train", action="store_true",
+                   help="Expand LRMC core with all training nodes (C' = C ∪ train_idx).")
+    p.add_argument("--warm_ft_epochs", type=int, default=0,
+                   help="If >0, run a short finetune on the FULL graph starting from the core model.")
+    p.add_argument("--warm_ft_lr", type=float, default=0.005)
+    p.add_argument("--runs", type=int, default=1,
+                   help="Number of runs with different seeds to average results.")
+    p.add_argument("-o", "--output_json", type=str, default=None,
+                   help="If set, save all computed metrics and settings to this JSON file.")
+    args = p.parse_args()
+
+    # ------------------------------------------------------------
+    # Load data
+    # ------------------------------------------------------------
+    ds = Planetoid(root=f"./data/{args.dataset}", name=args.dataset)
+    data = ds[0]
+    n, e = data.num_nodes, data.edge_index.size(1) // 2
+
+    console.print(f"[bold cyan]Dataset: {args.dataset} | Nodes: {n} | Edges: {e}[/bold cyan]")
+
+    # Results accumulator for optional JSON output
+    results = {
+        "args": {
+            k: (float(v) if isinstance(v, float) else v)
+            for k, v in vars(args).items()
+            if k != "output_json"
+        },
+        "dataset": {
+            "name": args.dataset,
+            "num_nodes": int(n),
+            "num_edges": int(e),
+        },
+    }
+
+    def maybe_save_results():
+        """Write results to JSON if the user requested it."""
+        if not args.output_json:
+            return
+        out_path = Path(args.output_json)
+        try:
+            out_path.parent.mkdir(parents=True, exist_ok=True)
+        except Exception:
+            pass
+        with out_path.open("w") as f:
+            json.dump(results, f, indent=2)
+
+    # ------------------------------------------------------------
+    # Load LRMC core
+    # ------------------------------------------------------------
+    core_mask = load_top1_assignment(args.seeds, n)
+    if args.expand_core_with_train:
+        core_mask = core_mask | data.train_mask
+
+    C_idx = torch.nonzero(core_mask, as_tuple=False).view(-1)
+    frac = 100.0 * C_idx.numel() / n
+    cov = coverage_counts(core_mask, data.train_mask, data.val_mask, data.test_mask)
+
+    console.print(f"[bold green]Loaded LRMC core of size {cov['core_size']} (≈{frac:.2f}% of the graph) from {args.seeds}[/bold green]")
+
+    # Record core coverage info
+    results["core"] = {
+        "source": str(args.seeds),
+        "expanded_with_train": bool(args.expand_core_with_train),
+        "size": int(cov["core_size"]),
+        "fraction": float(frac / 100.0),
+        "coverage": {
+            "train_in_core": int(cov["train_in_core"]),
+            "val_in_core": int(cov["val_in_core"]),
+            "test_in_core": int(cov["test_in_core"]),
+        },
+    }
+
+    # Coverage table
+    cov_table = Table(title="LRMC Core Coverage")
+    cov_table.add_column("Metric", style="cyan")
+    cov_table.add_column("Count", style="magenta")
+    cov_table.add_row("Core Size", str(cov["core_size"]))
+    cov_table.add_row("Train in Core", str(cov["train_in_core"]))
+    cov_table.add_row("Val in Core", str(cov["val_in_core"]))
+    cov_table.add_row("Test in Core", str(cov["test_in_core"]))
+    console.print(cov_table)
+
+    # ------------------------------------------------------------
+    # Single-run or multi-run execution
+    # ------------------------------------------------------------
+    if args.runs == 1:
+        # ---------------------
+        # Baseline (full graph)
+        # ---------------------
+        set_seed(0)
+        t0 = time.perf_counter()
+        base = GCN2(in_dim=ds.num_node_features,
+                    hid=args.hidden,
+                    out_dim=ds.num_classes,
+                    dropout=args.dropout)
+        train(base, data, epochs=args.epochs, lr=args.lr, wd=args.wd, patience=args.patience)
+        base_metrics = eval_all(base, data)
+        t1 = time.perf_counter()
+
+        console.print("\n[bold]Baseline (trained on full graph):[/bold]")
+        base_table = Table(show_header=True, header_style="bold magenta")
+        base_table.add_column("Metric", style="cyan")
+        base_table.add_column("Value", style="magenta")
+        base_table.add_row("Train Accuracy", f"{base_metrics['train']:.4f}")
+        base_table.add_row("Validation Accuracy", f"{base_metrics['val']:.4f}")
+        base_table.add_row("Test Accuracy", f"{base_metrics['test']:.4f}")
+        base_table.add_row("Time (s)", f"{t1 - t0:.2f}")
+        console.print(base_table)
+
+        # Save baseline single-run metrics
+        results["single_run"] = {
+            "baseline": {
+                "train": float(base_metrics["train"]),
+                "val": float(base_metrics["val"]),
+                "test": float(base_metrics["test"]),
+                "time_s": float(t1 - t0),
+            }
+        }
+
+        # ---------------------
+        # Core model (train on subgraph)
+        # ---------------------
+        if C_idx.numel() == 0:
+            console.print("[bold yellow]LRMC core is empty; skipping core model.[/bold yellow]")
+            results["core_empty"] = True
+            maybe_save_results()
+            return
+
+        data_C = subset_data(data, C_idx)
+        mC = GCN2(in_dim=ds.num_node_features,
+                  hid=args.hidden,
+                  out_dim=ds.num_classes,
+                  dropout=args.dropout)
+
+        t2 = time.perf_counter()
+        train(mC, data_C, epochs=args.epochs, lr=args.lr, wd=args.wd, patience=args.patience)
+        t3 = time.perf_counter()
+
+        # Evaluate core model on FULL graph
+        if args.core_mode == "forward":
+            # Mode A: run a standard forward pass on the full graph
+            mC.eval()
+            logits_full = mC(data.x, data.edge_index)
+        else:
+            # Mode B: seed logits on core and propagate with APPNP
+            mC.eval()
+            with torch.no_grad():
+                logits_C = mC(data_C.x, data_C.edge_index)  # [|C|, num_classes]
+                logits_seed = torch.zeros(n, ds.num_classes, device=logits_C.device)
+                logits_seed[C_idx] = logits_C
+                logits_full = appnp_seed_propagate(logits_seed,
+                                                   data.edge_index,
+                                                   K=args.K,
+                                                   alpha=args.alpha)
+
+        core_metrics = {
+            "train": accuracy(logits_full, data.y, data.train_mask),
+            "val": accuracy(logits_full, data.y, data.val_mask),
+            "test": accuracy(logits_full, data.y, data.test_mask),
+        }
+
+        console.print("\n[bold]LRMC‑core model (trained on core, evaluated on full graph):[/bold]")
+        core_table = Table(show_header=True, header_style="bold magenta")
+        core_table.add_column("Metric", style="cyan")
+        core_table.add_column("Value", style="magenta")
+        core_table.add_row("Train Accuracy", f"{core_metrics['train']:.4f}")
+        core_table.add_row("Validation Accuracy", f"{core_metrics['val']:.4f}")
+        core_table.add_row("Test Accuracy", f"{core_metrics['test']:.4f}")
+        core_table.add_row("Core Training Time (s)", f"{t3 - t2:.2f}")
+        speedup = (t1 - t0) / (t3 - t2 + 1e-9)
+        core_table.add_row("Speedup vs. Baseline", f"{speedup:.2f}×")
+        console.print(core_table)
+
+        # Save core single-run metrics
+        results["single_run"]["core_model"] = {
+            "mode": str(args.core_mode),
+            "train": float(core_metrics["train"]),
+            "val": float(core_metrics["val"]),
+            "test": float(core_metrics["test"]),
+            "core_train_time_s": float(t3 - t2),
+            "speedup_vs_baseline": float(speedup),
+        }
+
+        # ------------------------------------------------------------------------------------
+
+        console.print("\n[bold]Model Comparison: Baseline vs. L-RMC-core[/bold]")
+
+        # Create comparison table
+        comparison_table = Table(title="Performance Comparison", show_header=True, header_style="bold magenta")
+        comparison_table.add_column("Metric", style="cyan")
+        comparison_table.add_column("Baseline", style="magenta")
+        comparison_table.add_column("L-RMC-core", style="green")
+        comparison_table.add_column("Speedup", style="yellow")
+
+        # Add performance metrics
+        for metric in ["train", "val", "test"]:
+            comparison_table.add_row(
+                f"{metric.capitalize()} Accuracy",
+                f"{base_metrics[metric]:.4f}",
+                f"{core_metrics[metric]:.4f}",
+                ""  # Speedup is not applicable for accuracy
+            )
+
+        # Add timing and speedup
+        baseline_time = t1 - t0
+        core_time = t3 - t2
+        speedup = baseline_time / core_time if core_time > 0 else float('inf')
+
+        comparison_table.add_row(
+            "Training Time (s)",
+            f"{baseline_time:.2f}",
+            f"{core_time:.2f}",
+            f"{speedup:.2f}x"
+        )
+
+        comparison_table.add_row(
+            "Speedup",
+            "1x",
+            f"{speedup:.2f}x",
+            ""
+        )
+
+        console.print(comparison_table)
+
+        # Optional warm‑start finetune (single run)
+        if args.warm_ft_epochs > 0:
+            warm = GCN2(in_dim=ds.num_node_features,
+                        hid=args.hidden,
+                        out_dim=ds.num_classes,
+                        dropout=args.dropout)
+            warm.load_state_dict(mC.state_dict())
+
+            t4 = time.perf_counter()
+            train(warm, data,
+                  epochs=args.warm_ft_epochs,
+                  lr=args.warm_ft_lr,
+                  wd=args.wd,
+                  patience=args.warm_ft_epochs + 1)
+            t5 = time.perf_counter()
+            warm_metrics = eval_all(warm, data)
+
+            console.print("\n[bold]Warm‑start finetune (start from core model, train on FULL graph):[/bold]")
+            warm_table = Table(show_header=True, header_style="bold magenta")
+            warm_table.add_column("Metric", style="cyan")
+            warm_table.add_column("Value", style="magenta")
+            warm_table.add_row("Train Accuracy", f"{warm_metrics['train']:.4f}")
+            warm_table.add_row("Validation Accuracy", f"{warm_metrics['val']:.4f}")
+            warm_table.add_row("Test Accuracy", f"{warm_metrics['test']:.4f}")
+            warm_table.add_row("Finetune Time (s)", f"{t5 - t4:.2f}")
+            warm_table.add_row("Total (core train + warm)", f"{(t3 - t2 + t5 - t4):.2f}s")
+            console.print(warm_table)
+
+            # Save warm single-run metrics
+            results["single_run"]["warm_finetune"] = {
+                "train": float(warm_metrics["train"]),
+                "val": float(warm_metrics["val"]),
+                "test": float(warm_metrics["test"]),
+                "finetune_time_s": float(t5 - t4),
+                "total_time_s": float((t3 - t2) + (t5 - t4)),
+            }
+
+        # Emit results for single-run
+        maybe_save_results()
+    else:
+        # --------------------------------------------------------
+        # Multi-run: average metrics across different seeds
+        # --------------------------------------------------------
+        runs = args.runs
+        console.print(f"\n[bold]Running {runs} seeds and averaging results[/bold]")
+
+        # Storage for metrics across runs
+        base_train, base_val, base_test, base_time = [], [], [], []
+        core_train, core_val, core_test, core_time = [], [], [], []
+        speedups = []
+
+        warm_train, warm_val, warm_test, warm_time, warm_total_time = [], [], [], [], []
+
+        data_C = subset_data(data, C_idx) if C_idx.numel() > 0 else None
+        results["core_empty"] = data_C is None
+
+        for r in range(runs):
+            set_seed(r)
+
+            # Baseline
+            t0 = time.perf_counter()
+            base = GCN2(in_dim=ds.num_node_features,
+                        hid=args.hidden,
+                        out_dim=ds.num_classes,
+                        dropout=args.dropout)
+            train(base, data, epochs=args.epochs, lr=args.lr, wd=args.wd, patience=args.patience)
+            bm = eval_all(base, data)
+            t1 = time.perf_counter()
+
+            base_train.append(bm["train"]) ; base_val.append(bm["val"]) ; base_test.append(bm["test"]) ; base_time.append(t1 - t0)
+
+            # Core model
+            if data_C is None:
+                continue  # no core available
+
+            t2 = time.perf_counter()
+            mC = GCN2(in_dim=ds.num_node_features,
+                      hid=args.hidden,
+                      out_dim=ds.num_classes,
+                      dropout=args.dropout)
+            train(mC, data_C, epochs=args.epochs, lr=args.lr, wd=args.wd, patience=args.patience)
+            t3 = time.perf_counter()
+
+            if args.core_mode == "forward":
+                mC.eval()
+                logits_full = mC(data.x, data.edge_index)
+            else:
+                mC.eval()
+                with torch.no_grad():
+                    logits_C = mC(data_C.x, data_C.edge_index)
+                    logits_seed = torch.zeros(n, ds.num_classes, device=logits_C.device)
+                    logits_seed[C_idx] = logits_C
+                    logits_full = appnp_seed_propagate(logits_seed,
+                                                       data.edge_index,
+                                                       K=args.K,
+                                                       alpha=args.alpha)
+
+            cm = {
+                "train": accuracy(logits_full, data.y, data.train_mask),
+                "val": accuracy(logits_full, data.y, data.val_mask),
+                "test": accuracy(logits_full, data.y, data.test_mask),
+            }
+
+            core_train.append(cm["train"]) ; core_val.append(cm["val"]) ; core_test.append(cm["test"]) ; core_time.append(t3 - t2)
+            speedups.append((t1 - t0) / (t3 - t2 + 1e-9))
+
+            # Optional warm finetune per run
+            if args.warm_ft_epochs > 0:
+                warm = GCN2(in_dim=ds.num_node_features,
+                            hid=args.hidden,
+                            out_dim=ds.num_classes,
+                            dropout=args.dropout)
+                warm.load_state_dict(mC.state_dict())
+
+                t4 = time.perf_counter()
+                train(warm, data,
+                      epochs=args.warm_ft_epochs,
+                      lr=args.warm_ft_lr,
+                      wd=args.wd,
+                      patience=args.warm_ft_epochs + 1)
+                t5 = time.perf_counter()
+                wm = eval_all(warm, data)
+                warm_train.append(wm["train"]) ; warm_val.append(wm["val"]) ; warm_test.append(wm["test"]) ; warm_time.append(t5 - t4)
+                warm_total_time.append((t3 - t2) + (t5 - t4))
+
+        # Helper to format mean ± std
+        def fmt(values, prec=4):
+            if not values:
+                return "n/a"
+            if len(values) == 1:
+                return f"{values[0]:.{prec}f}"
+            try:
+                return f"{mean(values):.{prec}f} ± {stdev(values):.{prec}f}"
+            except Exception:
+                m = sum(values) / len(values)
+                var = sum((v - m) ** 2 for v in values) / max(1, len(values) - 1)
+                return f"{m:.{prec}f} ± {var ** 0.5:.{prec}f}"
+
+        def stats(values):
+            """Return dict with list, mean, std, count for JSON."""
+            d = {
+                "values": [float(v) for v in values],
+                "count": int(len(values)),
+            }
+            if len(values) >= 1:
+                d["mean"] = float(mean(values))
+            if len(values) >= 2:
+                d["std"] = float(stdev(values))
+            else:
+                d["std"] = None
+            return d
+
+        # Baseline summary
+        console.print("\n[bold]Baseline (averaged over runs):[/bold]")
+        base_table = Table(show_header=True, header_style="bold magenta")
+        base_table.add_column("Metric", style="cyan")
+        base_table.add_column("Mean ± Std", style="magenta")
+        base_table.add_row("Train Accuracy", fmt(base_train))
+        base_table.add_row("Validation Accuracy", fmt(base_val))
+        base_table.add_row("Test Accuracy", fmt(base_test))
+        base_table.add_row("Time (s)", fmt(base_time, prec=2))
+        console.print(base_table)
+
+        # Save baseline multi-run summary
+        results["multi_run"] = {
+            "runs": int(runs),
+            "baseline": {
+                "train": stats(base_train),
+                "val": stats(base_val),
+                "test": stats(base_test),
+                "time_s": stats(base_time),
+            }
+        }
+
+        if data_C is None:
+            console.print("[bold yellow]LRMC core is empty; no core runs to average.[/bold yellow]")
+            maybe_save_results()
+            return
+
+        # Core summary
+        console.print("\n[bold]LRMC‑core (averaged over runs):[/bold]")
+        core_table = Table(show_header=True, header_style="bold magenta")
+        core_table.add_column("Metric", style="cyan")
+        core_table.add_column("Mean ± Std", style="magenta")
+        core_table.add_row("Train Accuracy", fmt(core_train))
+        core_table.add_row("Validation Accuracy", fmt(core_val))
+        core_table.add_row("Test Accuracy", fmt(core_test))
+        core_table.add_row("Core Training Time (s)", fmt(core_time, prec=2))
+        core_table.add_row("Speedup vs. Baseline", fmt(speedups, prec=2))
+        console.print(core_table)
+
+        # Save core multi-run summary
+        results["multi_run"]["core_model"] = {
+            "mode": str(args.core_mode),
+            "train": stats(core_train),
+            "val": stats(core_val),
+            "test": stats(core_test),
+            "core_train_time_s": stats(core_time),
+            "speedup_vs_baseline": stats(speedups),
+        }
+
+        # Comparison summary
+        console.print("\n[bold]Model Comparison (averaged): Baseline vs. L-RMC-core[/bold]")
+        comparison_table = Table(title="Performance Comparison (Mean ± Std)", show_header=True, header_style="bold magenta")
+        comparison_table.add_column("Metric", style="cyan")
+        comparison_table.add_column("Baseline", style="magenta")
+        comparison_table.add_column("L-RMC-core", style="green")
+        comparison_table.add_column("Speedup", style="yellow")
+
+        for metric, b_vals, c_vals in [
+            ("Train Accuracy", base_train, core_train),
+            ("Validation Accuracy", base_val, core_val),
+            ("Test Accuracy", base_test, core_test),
+        ]:
+            comparison_table.add_row(metric, fmt(b_vals), fmt(c_vals), "")
+
+        comparison_table.add_row("Training Time (s)", fmt(base_time, prec=2), fmt(core_time, prec=2), fmt(speedups, prec=2))
+        comparison_table.add_row("Speedup", "1x", fmt(speedups, prec=2), "")
+        console.print(comparison_table)
+
+        # Optional warm summary
+        if args.warm_ft_epochs > 0 and warm_time:
+            console.print("\n[bold]Warm‑start finetune (averaged over runs):[/bold]")
+            warm_table = Table(show_header=True, header_style="bold magenta")
+            warm_table.add_column("Metric", style="cyan")
+            warm_table.add_column("Mean ± Std", style="magenta")
+            warm_table.add_row("Train Accuracy", fmt(warm_train))
+            warm_table.add_row("Validation Accuracy", fmt(warm_val))
+            warm_table.add_row("Test Accuracy", fmt(warm_test))
+            warm_table.add_row("Finetune Time (s)", fmt(warm_time, prec=2))
+            warm_table.add_row("Total (core train + warm)", fmt(warm_total_time, prec=2))
+            console.print(warm_table)
+
+            # Save warm multi-run summary
+            results["multi_run"]["warm_finetune"] = {
+                "train": stats(warm_train),
+                "val": stats(warm_val),
+                "test": stats(warm_test),
+                "finetune_time_s": stats(warm_time),
+                "total_time_s": stats(warm_total_time),
+            }
+
+        # Emit results for multi-run
+        maybe_save_results()
+
+if __name__ == "__main__":
+    main()

src/2_epsilon_seed_sweep.py

@@ -0,0 +1,215 @@
+import json
+import os
+import shutil
+from typing import List, Set, Dict, Tuple
+from decimal import Decimal, getcontext
+from rich import print
+
+from generate_lrmc_seeds import build_lrmc_single_graph
+
+def get_seed_nodes(seeds_path: str) -> Set[int]:
+    """Extract all seed nodes from a seeds JSON file.
+
+    Handles either 'seed_nodes' or 'members' fields.
+    """
+    try:
+        with open(seeds_path, 'r') as f:
+            data = json.load(f)
+
+        seed_nodes: Set[int] = set()
+        clusters = data.get('clusters', [])
+        for cluster in clusters:
+            nodes = cluster.get('seed_nodes')
+            if nodes is None:
+                nodes = cluster.get('members', [])
+            seed_nodes.update(nodes)
+
+        return seed_nodes
+    except (FileNotFoundError, json.JSONDecodeError, KeyError) as e:
+        print(f"[red]Error reading {seeds_path}: {e}[/red]")
+        return set()
+
+def _format_eps_label(val: Decimal) -> str:
+    """Return a stable, unique string label for epsilon values.
+
+    - Use integer string for integral values (e.g., '50000').
+    - Otherwise, use a compact decimal without trailing zeros.
+    This avoids duplicates like many '1e+04' when step is small.
+    """
+    # Normalize to remove exponent if integral
+    if val == val.to_integral_value():
+        return str(val.to_integral_value())
+    # Use 'f' then strip trailing zeros/decimal point for uniqueness and readability
+    s = format(val, 'f')
+    if '.' in s:
+        s = s.rstrip('0').rstrip('.')
+    return s
+
+def generate_epsilon_range(start: float, end: float, step: float) -> List[str]:
+    """Generate epsilon values as unique, stable strings.
+
+    Uses Decimal to avoid float accumulation and label collisions.
+    """
+    if step <= 0:
+        raise ValueError("epsilon_step must be > 0")
+
+    getcontext().prec = 28
+    s = Decimal(str(start))
+    e = Decimal(str(end))
+    t = Decimal(str(step))
+
+    vals: List[str] = []
+    cur = s
+    # Safety margin to include end due to decimal rounding
+    while cur <= e + Decimal('1e-18'):
+        label = _format_eps_label(cur)
+        if not vals or vals[-1] != label:
+            vals.append(label)
+        cur += t
+    return vals
+
+def run_epsilon_sweep(input_edgelist: str, out_dir: str, levels: int,
+                     epsilon_start: float = 1e4, epsilon_end: float = 5e5,
+                     epsilon_step: float = 1e4, cleanup_duplicates: bool = True):
+    """
+    Run LRMC for multiple epsilon values and remove duplicate results.
+
+    Args:
+        input_edgelist: Path to input edgelist file
+        out_dir: Output directory
+        levels: Number of levels to build
+        epsilon_start: Starting epsilon value (default: 1e4)
+        epsilon_end: Ending epsilon value (default: 5e5)
+        epsilon_step: Step size for epsilon (default: 1e4)
+        cleanup_duplicates: Whether to remove duplicate seed sets (default: True)
+    """
+    print(f"[blue]Starting epsilon sweep from {epsilon_start} to {epsilon_end} with step {epsilon_step}[/blue]")
+
+    # Preflight: check input edgelist path and fix a common typo
+    if not os.path.isfile(input_edgelist):
+        fixed_path = None
+        if input_edgelist.endswith('.tx') and os.path.isfile(input_edgelist + 't'):
+            fixed_path = input_edgelist + 't'
+        elif input_edgelist.endswith('.txt'):
+            # Try relative to CWD if a bare filename was intended
+            alt = os.path.join(os.getcwd(), input_edgelist)
+            if os.path.isfile(alt):
+                fixed_path = alt
+        if fixed_path:
+            print(f"[yellow]Input edgelist not found at '{input_edgelist}'. Using '{fixed_path}' instead.[/yellow]")
+            input_edgelist = fixed_path
+        else:
+            raise FileNotFoundError(f"Input edgelist not found: '{input_edgelist}'. Did you mean '.txt'?")
+
+    # Generate epsilon values
+    epsilons = generate_epsilon_range(epsilon_start, epsilon_end, epsilon_step)
+    print(f"[blue]Will test {len(epsilons)} epsilon values: {epsilons}[/blue]")
+
+    # Track seen seed sets and their corresponding epsilon values
+    seen_seed_sets: Dict[Tuple[int, ...], str] = {}
+
+    # Run for each epsilon
+    for epsilon in epsilons:
+        print(f"[yellow]Processing epsilon: {epsilon}[/yellow]")
+
+        try:
+            # Create temporary output directory for this epsilon
+            temp_out_dir = f"{out_dir}_temp_{epsilon}"
+
+            # Run LRMC
+            seeds_path = build_lrmc_single_graph(
+                input_edgelist=input_edgelist,
+                out_dir=temp_out_dir,
+                levels=levels,
+                epsilon=epsilon
+            )
+
+            # Get seed nodes
+            seed_nodes = get_seed_nodes(seeds_path)
+            seed_nodes_tuple = tuple(sorted(seed_nodes))
+
+            print(f"[green]Epsilon {epsilon}: Found {len(seed_nodes)} unique seed nodes[/green]")
+
+            # Check if this seed set has been seen before
+            if seed_nodes_tuple in seen_seed_sets:
+                existing_epsilon = seen_seed_sets[seed_nodes_tuple]
+                print(f"[yellow]Duplicate seed set found! Epsilon {epsilon} has same seeds as {existing_epsilon}[/yellow]")
+                print(f"[yellow]Removing duplicate results for epsilon {epsilon}[/yellow]")
+
+                # Clean up temporary directory
+                if os.path.exists(temp_out_dir):
+                    shutil.rmtree(temp_out_dir)
+                continue
+
+            # If we get here, this is a unique seed set
+            seen_seed_sets[seed_nodes_tuple] = epsilon
+
+            # Move results to final location
+            final_out_dir = f"{out_dir}_epsilon_{epsilon}"
+            if os.path.exists(final_out_dir):
+                shutil.rmtree(final_out_dir)
+            shutil.move(temp_out_dir, final_out_dir)
+
+            # Move seeds_XXXXX.json to the stage0 directory
+            seeds_file = os.path.join(final_out_dir, "stage0", f"seeds_{epsilon}.json")
+            if os.path.exists(seeds_file):
+                stage0_dir = os.path.join(out_dir, "stage0")
+                if not os.path.exists(stage0_dir):
+                    os.makedirs(stage0_dir)
+                shutil.move(seeds_file, os.path.join(stage0_dir, f"seeds_{epsilon}.json"))
+
+            print(f"[green]Unique results saved to {os.path.join(stage0_dir, f"seeds_{epsilon}.json")}[/green]")
+
+        except Exception as e:
+            print(f"[red]Error processing epsilon {epsilon}: {e}[/red]")
+            # Clean up temporary directory if it exists
+            temp_out_dir = f"{out_dir}_temp_{epsilon}"
+            if os.path.exists(temp_out_dir):
+                shutil.rmtree(temp_out_dir)
+
+    # Print summary
+    print("\n[blue]--- Summary ---[/blue]")
+    print(f"[blue]Total epsilon values tested: {len(epsilons)}[/blue]")
+    print(f"[blue]Unique seed sets found: {len(seen_seed_sets)}[/blue]")
+    print(f"[blue]Duplicates removed: {len(epsilons) - len(seen_seed_sets)}[/blue]")
+
+    if seen_seed_sets:
+        print("\n[green]Unique epsilon values kept:[/green]")
+        for seed_tuple, epsilon in sorted(seen_seed_sets.items()):
+            seed_count = len(seed_tuple)
+            print(f"  {epsilon}: {seed_count} seed nodes")
+
+def main():
+    """Main function with command line interface."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run LRMC epsilon sweep with duplicate removal")
+    parser.add_argument('--input_edgelist', type=str, required=True,
+                       help='Path to input edgelist file')
+    parser.add_argument('--out_dir', type=str, required=True,
+                       help='Base output directory (results will be saved as out_dir_epsilon_X)')
+    parser.add_argument('--levels', type=int, required=True,
+                       help='Number of levels to build')
+    parser.add_argument('--epsilon_start', type=float, default=1e4,
+                       help='Starting epsilon value (default: 1e4)')
+    parser.add_argument('--epsilon_end', type=float, default=5e5,
+                       help='Ending epsilon value (default: 5e5)')
+    parser.add_argument('--epsilon_step', type=float, default=1e4,
+                       help='Epsilon step size (default: 1e4)')
+    parser.add_argument('--no_cleanup', action='store_true',
+                       help='Do not remove duplicates (keep all results)')
+
+    args = parser.parse_args()
+
+    run_epsilon_sweep(
+        input_edgelist=args.input_edgelist,
+        out_dir=args.out_dir,
+        levels=args.levels,
+        epsilon_start=args.epsilon_start,
+        epsilon_end=args.epsilon_end,
+        epsilon_step=args.epsilon_step,
+        cleanup_duplicates=not args.no_cleanup
+    )
+
+if __name__ == '__main__':
+    main()

src/2_lrmc_bilevel.py

@@ -1,6 +1,4 @@
-# lrmc_bilevel.py
 # Bi-level Node↔Cluster message passing with fixed LRMC seeds
-# Requires: torch, torch_geometric, torch_sparse
 
 import argparse, json, os
 from pathlib import Path
@@ -19,6 +17,8 @@ from torch_geometric.loader import DataLoader
 from torch_geometric.datasets import Planetoid, TUDataset
 from torch_geometric.nn import GCNConv, global_mean_pool
 
+from rich import print
+
 
 # ---------------------------
 # Utilities: edges and seeds

src/2_random_seed_sweep.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------
+#  generate_lrmc_seeds.sh
+#
+#  Run generate_lrmc_seeds.py for num_nodes = 1, 101, 201, …,
+#  2601, 2701 and then once more for 2708.
+#
+#  Usage: ./generate_lrmc_seeds.sh
+# ------------------------------------------------------------
+
+set -euo pipefail          # Exit on error, undefined var, or pipe failure
+
+# ------------------------------------------------------------------
+# Configuration – adjust these if your file layout changes
+# ------------------------------------------------------------------
+INPUT_EDGELIST="cora_seeds/edgelist.txt"
+OUT_DIR="cora_seeds"
+LEVELS=1
+BASELINE="random"
+
+# ------------------------------------------------------------------
+# Main loop – 1 … 2708 stepping by 100
+# ------------------------------------------------------------------
+echo "Starting loop over num_nodes = 1, 101, 201, …, 2601, 2701 …"
+
+for NUM_NODES in $(seq 1 100 2708); do
+    echo ">>> num_nodes=$NUM_NODES"
+    python3 generate_lrmc_seeds.py \
+        --input_edgelist "$INPUT_EDGELIST" \
+        --out_dir "$OUT_DIR" \
+        --levels "$LEVELS" \
+        --baseline "$BASELINE" \
+        --num_nodes "$NUM_NODES"
+done
+
+# ------------------------------------------------------------------
+# Explicitly run for 2708 (not hit by the 100‑step sequence)
+# ------------------------------------------------------------------
+echo ">>> num_nodes=2708 (explicitly added)"
+python3 generate_lrmc_seeds.py \
+    --input_edgelist "$INPUT_EDGELIST" \
+    --out_dir "$OUT_DIR" \
+    --levels "$LEVELS" \
+    --baseline "$BASELINE" \
+    --num_nodes 2708
+
+echo "All done!"