Metadata-Version: 2.4
Name: gnn-augment
Version: 0.4.1
Summary: Graph-augmented semi-supervised node classification: boost graph neural networks (GCN, GraphSAGE, GAT, GIN) by replacing the adjacency with node-similarity representatives (Katz, Rooted PageRank, Graph Gravity, LLM text-embedding kNN) -- single view, multi-view training, sparse mode for large graphs
Author: Akanda Wahid-Ul-Ashraf
License: MIT License
        
        Copyright (c) 2026 Akanda Wahid-Ul-Ashraf
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://github.com/AkandaAshraf/gnn-augment
Project-URL: Documentation, https://github.com/AkandaAshraf/gnn-augment/blob/main/docs/USAGE.md
Project-URL: Repository, https://github.com/AkandaAshraf/gnn-augment
Project-URL: Changelog, https://github.com/AkandaAshraf/gnn-augment/blob/main/CHANGELOG.md
Project-URL: Issues, https://github.com/AkandaAshraf/gnn-augment/issues
Project-URL: Paper (arXiv:1905.09087), https://arxiv.org/abs/1905.09087
Project-URL: PhD thesis (Bournemouth), https://eprints.bournemouth.ac.uk/34163/
Project-URL: Gravity method (akanda-method), https://github.com/AkandaAshraf/akanda-method
Keywords: graph-neural-networks,gnn,graph-convolutional-networks,gcn,graphsage,graph-attention-networks,graph-augmentation,semi-supervised-learning,node-classification,network-science,link-prediction,pytorch,scikit-learn
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Information Analysis
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: networkx>=3.0
Requires-Dist: numpy>=1.24
Requires-Dist: scipy>=1.10
Requires-Dist: scikit-learn>=1.3
Provides-Extra: semantic
Requires-Dist: sentence-transformers>=2.2; extra == "semantic"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: build>=1.0; extra == "dev"
Requires-Dist: twine>=4.0; extra == "dev"
Dynamic: license-file

# gnn-augment

[![CI](https://github.com/AkandaAshraf/gnn-augment/actions/workflows/ci.yml/badge.svg)](https://github.com/AkandaAshraf/gnn-augment/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/gnn-augment)](https://pypi.org/project/gnn-augment/)
[![Python ≥ 3.10](https://img.shields.io/badge/python-%E2%89%A5%203.10-blue)](https://github.com/AkandaAshraf/gnn-augment/blob/main/pyproject.toml)
[![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/AkandaAshraf/gnn-augment/blob/main/LICENSE)

**Boost a graph neural network for semi-supervised node classification by
augmenting the graph with node-similarity representatives.**

A graph neural network (GNN) propagates node features through the graph's
adjacency matrix, so at layer *l* only *l*-hop neighbours influence a
node. `gnn-augment` replaces the adjacency with a **node-similarity
representative** — Katz, Rooted PageRank, or **Graph Gravity** (a
gravitational link-prediction score: degree as mass, shortest path as
distance) — letting distant-but-similar nodes influence each other even
in a shallow network, which improves node-classification accuracy.

The trick is architecture-agnostic — it changes *which graph* the network
sees — so it works across four GNN architectures: the Kipf–Welling **GCN**,
**GraphSAGE**, **GAT** (graph attention), and **GIN**.

This is a packaged, reusable implementation of the augmentation method from:

> A. Wahid-Ul-Ashraf, M. Budka, K. Musial.
> **Simulation and Augmentation of Social Networks for Building Deep
> Learning Models.** [arXiv:1905.09087](https://arxiv.org/abs/1905.09087)

The Graph Gravity representative is the same gravitational method whose
link-prediction reference implementation is
[akanda-method](https://github.com/AkandaAshraf/akanda-method).

## Installation

```bash
pip install gnn-augment                   # release
pip install "gnn-augment[semantic]"       # + sentence-transformers, for the semantic (LLM-embedding) view
pip install git+https://github.com/AkandaAshraf/gnn-augment   # latest main
```

Requires Python ≥ 3.10; depends on `torch`, `networkx`, `numpy`, `scipy`
and `scikit-learn`. A CUDA build of PyTorch is used automatically when a
GPU is available (`device="auto"`) — install the torch wheel you want (CPU
or a specific CUDA) first if pip's default is not it. From a clone:
`pip install -e ".[dev,semantic]"` (see *Development and tests* below).

## Documentation

- [docs/USAGE.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/docs/USAGE.md) — user guide: every option, semi-supervised conventions, multi-view, the semantic view, sparse mode, datasets, reproducing the benchmarks
- [docs/API.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/docs/API.md) — API reference generated from the docstrings
- [benchmarks/](https://github.com/AkandaAshraf/gnn-augment/tree/main/benchmarks) — every result table with the script that produced it; [CHANGELOG.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/CHANGELOG.md)

## Quick start

```python
from gnn_augment import GNNClassifier, load_cora, make_semi_supervised_split

G, X, y = load_cora()                                    # networkx graph, features, labels
y_semi, train, val, test = make_semi_supervised_split(y, seed=0)  # -1 = unlabeled

clf = GNNClassifier(model="gcn", representative="rpr", threshold="auto")
clf.fit(G, X, y_semi, train_mask=train, val_mask=val)
print("test accuracy:", clf.score(G, X, y, mask=test))
```

Pick any architecture with `model=` — `"gcn"`, `"sage"` (GraphSAGE),
`"gat"` (graph attention), or `"gin"`. The augmentation (`representative`,
`threshold`) is independent of the architecture, so any combination works.

The estimator follows the scikit-learn `fit` / `predict` / `predict_proba`
/ `score` convention. Semi-supervised learning is expressed the
scikit-learn way — set `-1` for unlabeled nodes in `y` — or pass explicit
`train_mask` / `val_mask` arrays. The graph may be a networkx graph, a
dense adjacency matrix, or a SciPy sparse matrix.

## The augmented representatives

| `representative` | what replaces the adjacency | paper |
|---|---|---|
| `"adjacency"` | none — the original Kipf–Welling GCN | Eq. 10–12 |
| `"katz"` | Katz index Σ βˡ Aˡ (truncated at A⁵) | Eq. 17 |
| `"rpr"` | Rooted PageRank | Eq. 20 |
| `"gravity"` | Graph Gravity: DC(i)·DC(j) / SP(i,j)² | Eq. 22 |

Similarity matrices are L2 row-normalised and **thresholded** (paper
§9.1): pass `threshold=(t1, t2)` (entries ≤ t1 → 0, > t2 → 1), `"auto"`
(binarise at the mean), `None` (keep weights), or `"select"` to pick the
best threshold on the validation set automatically. `weighted_features=True`
adds the learnable per-feature weight vector (Eq. 25).

Build a representative on its own for use in another model:

```python
from gnn_augment import build_representative
G_matrix = build_representative(A, kind="gravity", threshold="auto")  # normalised Laplacian
```

## Multi-view training: a different augmented graph every epoch

The paper's conclusions suggest using *several* augmented matrices of the
same graph as a regulariser. `views=` turns that on: each epoch the
network trains on a different view of the graph (graph data augmentation),
and at inference the class probabilities can be averaged over the views
(test-time augmentation):

```python
clf = GNNClassifier(
    model="gin",
    views=["adjacency", ("rpr", (0.0, 0.5)), ("gravity", (0.0, 1.0))],
    view_schedule="random",   # "random" | "cycle" | "mix" (Dirichlet mixture)
    predict_views="mean",     # "mean" (TTA) | "first"
)
```

On Cora (3 splits; [benchmarks/multiview-cora.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/multiview-cora.md)),
multi-view training with TTA beats the raw adjacency on every architecture
without any view selection, and on the GCN it is the **best configuration
of all** (0.819 vs 0.816 for the best single view). When one view clearly
dominates (RPR on GraphSAGE/GIN), committing to it still wins — multi-view
is the robust "don't know the best view" strategy. Most of the gain comes
from averaging over views at inference, so `predict_views="mean"` is the
recommended mode.

Is it a real regulariser? Tested properly — high-capacity models (up to
256 hidden × 4 layers), **no dropout, no weight decay, no early stopping**,
so they memorise the training set ([benchmarks/overfitting-cora.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/overfitting-cora.md)):
multi-view training beats the raw adjacency at 3 of 4 capacities on the GCN
(up to +2.3 points at 256 × 4) and the single RPR view at all 4 — a genuine
but modest effect that complements early stopping (the dominant
regulariser) rather than replacing it. A fixed augmented view, by contrast,
does *not* regularise: its gains depend on early stopping. One limitation:
deep unregularised **GIN** fails to train under multi-view (its
un-normalised sum aggregation cannot handle the message-scale jump between
views) — use a single view, `"cycle"`, or scaled views for deep GIN.

## The semantic view: LLM embeddings of node text as a graph (Graph + LLM)

Nodes of citation graphs are papers with real text. `"semantic"` builds a
graph representative from **text embeddings**: cosine-similarity kNN over
sentence-transformer (or any LLM) embeddings of each node's text, with the
original edges kept — a *text-similarity* view next to the topological ones.
Embeddings are computed locally (MiniLM, 22M parameters; no API):

```python
from gnn_augment import GNNClassifier, embed_texts, load_cora, load_cora_text

G, X, y = load_cora()
E = embed_texts(load_cora_text())          # n x 384, local sentence-transformer

clf = GNNClassifier(representative="semantic", threshold=None, semantic_k=10)
clf.fit(G, X, y_semi, train_mask=train, val_mask=val, node_embeddings=E)

# or as one view among several in multi-view training:
clf = GNNClassifier(views=["adjacency", ("rpr", (0.0, 0.5)), ("semantic", None)])
```

Raw texts: `load_cora_text()` (aligned with `load_cora`) and
`load_pubmed_raw()` (graph + features + labels + texts), both reading the
TAPE raw-text releases (`data/raw_text/{cora_orig,PubMed_orig}.zip`, from
the Google-Drive links in the [TAPE](https://github.com/XiaoxinHe/TAPE)
README). `embed_texts(model_name=...)` also accepts a local model
directory — handy offline or behind TLS-intercepting proxies: download the
`sentence-transformers/all-MiniLM-L6-v2` files once (e.g. with curl) into
`data/models/all-MiniLM-L6-v2` and the example uses it automatically.
`cache_representatives=True` caches augmented adjacencies in-process so
repeated fits on the same graph (seeds, architectures, views) reuse the
expensive Katz/RPR/gravity builds.

On Cora the semantic view, added to multi-view training, gives the **best
accuracy yet on every architecture** (3 splits;
[benchmarks/semantic.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/semantic.md)):

| architecture | adjacency | best topological | multi adj+RPR+gravity+semantic |
|---|---|---|---|
| GCN | 0.797 | 0.819 | **0.829** |
| GraphSAGE | 0.750 | 0.794 | **0.805** |
| GIN | 0.729 | 0.784 | **0.792** |
| GAT | 0.814 | 0.818 | **0.828** |

Text *complements* topology rather than replacing it: the semantic graph
alone is weaker than the citation graph, but as an extra view it adds what
the topological views lack — and it is the first augmentation that moves GAT.

At 7× the scale (PubMed, 19,717 nodes, run in sparse mode, 2 splits) the
semantic view transfers: semantic kNN + adjacency is the most consistent
single view (+2.4 / +3.3 / +3.4 points on GCN / GraphSAGE / GIN; the best
single view on GCN and GraphSAGE), and multi-view training that includes it
is best or tied-best on GraphSAGE (+3.4) and GIN (+4.0, with gravity as a
fourth view) — every architecture gains at least 2.4 points where the
topological views alone gave at most +2.3 at this scale.

## The augmentation helps every architecture

The paper applied augmentation only to the GCN, but the trick generalises.
On Cora, the RPR representative improves all four architectures (mean test
accuracy over 3 splits; see [benchmarks/architectures-cora.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/architectures-cora.md)):

| architecture | adjacency | + RPR |
|---|---|---|
| GCN | 0.797 | **0.816** (+0.018) |
| GraphSAGE | 0.750 | **0.794** (+0.044) |
| GAT | 0.814 | **0.818** (+0.003) |
| GIN | 0.729 | **0.784** (+0.054) |

```bash
python examples/cora_demo.py            # vanilla vs augmented GCN
python examples/architectures_demo.py   # augmentation across GCN/SAGE/GAT/GIN
python examples/pubmed_demo.py          # scale up to ~20k nodes (GCN/SAGE/GIN)
```

Datasets: `load_cora()`, `load_planetoid("cora" | "citeseer" | "pubmed")`,
plus `load_cora_text()` / `load_pubmed_raw()` for node texts. Dense views
fit to roughly 20k nodes on an 8 GB GPU (see
[benchmarks/pubmed.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/pubmed.md)); GAT's O(N²·heads) attention
caps it at ~8–10k nodes; beyond that, sparse mode (next paragraph).

**Large graphs — sparse mode.** `GNNClassifier(sparse="auto")` (the
default) switches above 8k nodes to *chunked, sparse* views: every
representative is built in row chunks keeping only its `view_top_k`
strongest entries per row (RPR via a sparse LU solve, gravity via chunked
BFS, Katz via chunked powers, semantic via chunked kNN) and stored as a
SciPy CSR matrix — O(n·k) memory instead of O(n²) — with sparse propagation
on the GPU for GCN / GraphSAGE / GIN. A 20k-node view drops from ~1.5 GB to
~10 MB, and much larger graphs fit (the PubMed table in
[benchmarks/semantic.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/semantic.md) ran this way in a few GB
of host RAM). Force it with `sparse=True`; GAT needs the dense path.

The best *representative* is dataset-dependent (as the paper notes): RPR is
the most robust across architectures on Cora, while Graph Gravity suits
friendship/social networks rather than citation graphs. At PubMed scale the
purely topological gains shrink and become architecture-specific (GIN
benefits most; [benchmarks/pubmed.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/pubmed.md)), whereas the
semantic view keeps paying off there (+2.4 to +4.0 points;
[benchmarks/semantic.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/benchmarks/semantic.md)).

## Development and tests

```bash
git clone https://github.com/AkandaAshraf/gnn-augment && cd gnn-augment
pip install -e ".[dev,semantic]"
pytest                      # unit + integration (synthetic graphs, offline, ~1 min on CPU)
pytest -m unit              # fast, isolated tests of each function / class
pytest -m integration       # end-to-end through the public API, dense and sparse paths
pytest -m "not slow"        # skip the real-data tests (Cora / local MiniLM; auto-skipped when absent)
python -m build && twine check dist/*
```

Layout: `src/gnn_augment/` (`classifier`, `representatives`, `sparse`,
`models`, `data`), `tests/unit/`, `tests/integration/`, `examples/`, and
`benchmarks/` (every result table in this README, with the script that
produced it). CI ([ci.yml](https://github.com/AkandaAshraf/gnn-augment/blob/main/.github/workflows/ci.yml)) runs the unit and
integration suites on Python 3.10 and 3.12 with the CPU build of torch,
then builds the sdist/wheel, installs the wheel into a clean venv and runs
the integration tests against it; tagged releases are published to PyPI by
[release.yml](https://github.com/AkandaAshraf/gnn-augment/blob/main/.github/workflows/release.yml) (see [RELEASING.md](https://github.com/AkandaAshraf/gnn-augment/blob/main/RELEASING.md)).

## Relation to other work

The `"gravity"` representative is the *node-topological* gravitational
score (centrality as mass, path length as distance). It is unrelated to
gravity-inspired graph autoencoders, which apply a gravity term to learned
embeddings — see the akanda-method README for that distinction.

## Papers and citing

This package implements, and extends to more architectures, multi-view
training, a semantic view and large graphs, the augmentation method of

> Akanda Wahid-Ul-Ashraf, Marcin Budka, Katarzyna Musial.
> **Simulation and Augmentation of Social Networks for Building Deep
> Learning Models.** arXiv:1905.09087 (2019).
> [https://arxiv.org/abs/1905.09087](https://arxiv.org/abs/1905.09087)

The Graph Gravity representative is the gravitational link-prediction
method of

> Akanda Wahid-Ul-Ashraf, Marcin Budka, Katarzyna Musial.
> **How to predict social relationships — Physics-inspired approach to link
> prediction.** *Physica A: Statistical Mechanics and its Applications* 523
> (2019) 1110–1129. DOI [10.1016/j.physa.2019.04.246](https://doi.org/10.1016/j.physa.2019.04.246)
>
> Earlier version: Wahid-Ul-Ashraf, Budka, Musial-Gabrys. **Newton's
> gravitational law for link prediction in social networks.** *Complex
> Networks & Their Applications VI* (COMPLEX NETWORKS 2017), Springer, 2018,
> pp. 93–104. DOI [10.1007/978-3-319-72150-7_8](https://doi.org/10.1007/978-3-319-72150-7_8)

Both lines of work are part of the author's PhD thesis:

> Akanda Wahid-Ul-Ashraf. **Prediction and modelling of complex social
> networks and their evolution.** PhD thesis, Bournemouth University, 2020.
> [eprints.bournemouth.ac.uk/34163](https://eprints.bournemouth.ac.uk/34163/)

If you use `gnn-augment`, please cite the arXiv paper (and the Physica A
paper when you use the gravity representative); the thesis collects both:

```bibtex
@article{wahidulashraf2019augmentation,
  title   = {Simulation and Augmentation of Social Networks for Building Deep Learning Models},
  author  = {Wahid-Ul-Ashraf, Akanda and Budka, Marcin and Musial, Katarzyna},
  journal = {arXiv preprint arXiv:1905.09087},
  year    = {2019},
  url     = {https://arxiv.org/abs/1905.09087}
}

@article{wahidulashraf2019gravity,
  title   = {How to predict social relationships --- Physics-inspired approach to link prediction},
  author  = {Wahid-Ul-Ashraf, Akanda and Budka, Marcin and Musial, Katarzyna},
  journal = {Physica A: Statistical Mechanics and its Applications},
  volume  = {523},
  pages   = {1110--1129},
  year    = {2019},
  doi     = {10.1016/j.physa.2019.04.246}
}

@phdthesis{wahidulashraf2020thesis,
  title  = {Prediction and modelling of complex social networks and their evolution},
  author = {Wahid-Ul-Ashraf, Akanda},
  school = {Bournemouth University},
  year   = {2020},
  url    = {https://eprints.bournemouth.ac.uk/34163/}
}
```

## Acknowledgements

The method and the papers behind this package are the product of the
author's PhD research, **funded by Bournemouth University** and carried out
there under the supervision of **Marcin Budka** and **Katarzyna Musial**,
co-authors of the papers above.

## License

MIT — see [LICENSE](https://github.com/AkandaAshraf/gnn-augment/blob/main/LICENSE).
