Note: Last update on 2026.07.20
Note
⭐ Now available: This repository additionally hosts the extended code for our follow-up paper
"Mechanism-Dependent Antagonism of Auxiliary Information in Substation-Level Load Disaggregation for Distribution Network Planning" (arXiv:2605.24491) — see the SpatialAllocation/ library and the multi-country case studies under StudyCase/.
The exact code that accompanied the first (EPSR) paper is preserved at the git tag paper1-epsr.
AllocateGNN proposes a Graph Neural Network (GNN)-based approach to improve the spatial allocation of electricity demand in energy systems. Traditional methods such as Voronoi tessellation assign demand to the nearest substation using simple geometric proximity, ignoring structural and contextual information. This work formulates the spatial allocation task as an edge weight prediction problem on a heterogeneous graph and uses self-supervised learning to produce more accurate, context-aware allocations.
The framework models the spatial allocation problem as a heterogeneous graph with three node types:
- Source nodes: Represent regional administrative areas (e.g., ITL3 regions) with known aggregate demand.
- Agent nodes: Represent grid cells with land-use features that serve as intermediaries for demand distribution.
- Target nodes: Represent substations where demand is physically consumed.
A GNN encoder learns node embeddings via message passing, and a differentiable edge weighting layer predicts allocation weights, which are optimized using a combination of self-supervised and weakly-supervised loss functions including entropy regularization, feature similarity loss, and land-use prediction loss. The real substation demand
The follow-up paper studies substation-level load disaggregation when auxiliary information — night-time lights (NTL) and proximity priors — is injected into the learned weighter. A central finding is that auxiliary signals are not unconditionally helpful: depending on the correction mechanism (multiplicative reweighting, additive correction, conservation renormalization) two otherwise useful cues can antagonize each other. The effect is examined across three independent grids (Great Britain, Australia, and the German Börde region). The corresponding modules (FeatureExtractor, the Weighter/Allocator registries, and the NTL/proximity prior losses and correctors) and the full experiment suite live in this repository.
AllocateGNN/
├── SpatialAllocation/ # Core GNN-based spatial allocation library
│ ├── GNN/
│ │ ├── core/ # Training/inference orchestration + model config
│ │ ├── Layer/ # Graph encoder, edge-weight layer, gating, losses
│ │ │ └── LossFunction/ # Modular losses incl. NTL / Proximity prior losses
│ │ ├── Allocation/ # End-to-end allocation graph/solver/losses
│ │ └── utils/ # Graph construction and feature preprocessing
│ ├── Allocator/ # Allocation methods (Voronoi, CIVD, …) via a registry
│ ├── Weighter/ # Weighter methods (heterogeneous GNN, GPM, uniform, …)
│ ├── FeatureExtractor/ # Feature pipeline: fetchers / extractors / correctors
│ │ ├── fetchers/ # Sentinel-2, WorldCover, OSM, NTL data fetchers
│ │ ├── extractors/ # Spectral / land-use / WorldCover / NTL extractors
│ │ └── correctors/ # NTL and proximity feature correctors
│ └── utils/ # Imagery, network distance, spectral indices, CNN features
│
├── ClusterBasedVoronoi/ # Baseline: cluster-based Voronoi approach
│
├── StudyCase/ # Multi-country case studies (follow-up paper)
│ ├── British/ # Great Britain case study (notebooks + training scripts)
│ ├── British_weighter_experiments/# Antagonism experiment suite (main experiments)
│ ├── Australia/ # Australia (Ausgrid) case study
│ └── Germany/ # German Börde case study
│
├── requirements.txt
├── README.md
├── license.md
└── icon_kit.png
Note on data: raw and intermediate datasets, cached artifacts, and result files are not distributed with this repository (they are ignored via
.gitignore). The case-study scripts expect processed inputs to be present locally; the data-ingestion / feature-pipeline steps are intentionally excluded from the release.
Supports multiple GNN convolution types (GCN, GraphSAGE, GAT, GIN, HGT) wrapped in HeteroConv for heterogeneous graph learning. Includes residual connections, layer normalization, and L2-normalized embeddings with learnable scaling.
Predicts edge weights using embedding distances gated by a learned MLP, followed by temperature-scaled grouped softmax via scatter_softmax. Ensures that weights from each source node sum to 1.
Modular loss system with learnable uncertainty-based weighting for multi-task optimization. Self-supervised / weakly-supervised objectives include entropy regularization, feature-similarity/consistency losses, and a land-use (KL) prediction loss. The follow-up study adds NTL and proximity prior losses that inject auxiliary spatial information into the weighter.
The Allocator, Weighter, and FeatureExtractor subpackages expose registry-based interfaces so that allocation methods (Voronoi, CIVD, …), weighters (heterogeneous GNN, GPM, uniform, …), and feature fetchers/extractors/correctors can be selected and composed by name.
An alternative approach using clustering (DBSCAN, HDBSCAN, K-Means, etc.) combined with Voronoi tessellation. Supports optimization-based allocation via Pyomo with CIVD/IVD influence methods.
The follow-up paper is reproduced through four case-study folders. Each contains the analysis notebooks and the training/experiment scripts (numbered by execution order); raw-data ingestion and feature construction are excluded.
| Folder | Grid | Content |
|---|---|---|
British/ |
Great Britain | Base case study: static allocation, degradation analysis, CIVD, GNN training with NTL/proximity priors. |
British_weighter_experiments/ |
Great Britain | Main experiment suite for the antagonism study (main results, significance, strength sweeps, mechanism isolation, robustness, r-series ablations). |
Australia/ |
Ausgrid (AU) | Static baselines, GNN training, statistical evaluation, feature-fusion training. |
Germany/ |
Börde (DE) | Börde training, results tables, LOOCV, additive-correction matrix, pandapower downstream. |
- Python 3.12+
- CUDA-compatible GPU (recommended)
pip install -r requirements.txtFor PyTorch Geometric and its compiled extensions (torch-scatter, torch-sparse), follow the official installation guide to match your PyTorch and CUDA version. Pyomo additionally requires an external solver (e.g. CBC or GLPK) for the cluster-based Voronoi baseline.
from SpatialAllocation.GNN.core.ModelConfig import ModelConfig
from SpatialAllocation.GNN.core.EdgeWeightSolver import EdgeWeightSolver
from torch_geometric.loader import DataLoader
# Configure model
config = ModelConfig(
hidden_dim=128,
embedding_dim=64,
num_layers=3,
conv_type='sage',
epochs=300,
learning_rate=0.001,
learnable=True
)
# Create data loaders
train_loader = DataLoader(train_graphs, batch_size=1, shuffle=True)
test_loader = DataLoader(test_graphs, batch_size=1, shuffle=False)
# Train
solver = EdgeWeightSolver(config)
solver.train_multi_graph(
train_loader,
test_loader,
objective_weights={
'entropy_regularization': 1.0,
'landuse_prediction_loss': 1.0
}
)# Predict edge weights for a new graph
result_df = solver.predict_edge_weights(test_data)
# result_df contains: source_node_idx, agent_node_idx, agent_original_idx, predicted_weightEach StudyCase/<grid>/ folder holds numbered scripts and notebooks. The training scripts (e.g. 005_kfold_prior_training.py in British_weighter_experiments/) and the downstream experiment scripts expect processed inputs to be available locally.
If you use this framework in your research, please consider citing our papers 📝 and giving the repository a star ⭐:
@article{Mu2026Improving,
author={Mu, Xuanhao and Geiges, Jakob and Liu, Nan and Schlachter, Thorsten and Hagenmeyer, Veit},
title={Improving spatial allocation for energy system coupling with graph neural networks},
journal={Electric Power Systems Research},
volume={262},
pages={113519},
year={2027},
issn={0378-7796},
doi={10.1016/j.epsr.2026.113519},
url={https://www.sciencedirect.com/science/article/pii/S0378779626008126}
}
@article{Mu2026Antagonism,
author={Mu, Xuanhao and Thota, Kundan and Liu, Nan and Schlachter, Thorsten and Hagenmeyer, Veit},
title={Mechanism-Dependent Antagonism of Auxiliary Information in Substation-Level Load Disaggregation for Distribution Network Planning},
journal={arXiv preprint arXiv:2605.24491},
year={2026},
url={https://arxiv.org/abs/2605.24491}
}This code is licensed under the MIT License. For any issues or any intention of cooperation, please feel free to contact me at xuanhao.mu@kit.edu.