Graph Neural Network Surrogates for Energy System Modeling

Graph Neural Network Surrogates for Energy System Modeling Gerild Pjetri, Anurag Mohapatra, Beneharo Reveron Baecker, Maximilian Hock, Thomas Hamacher Technical University of Munich, Germany {gerild.pjetri, anurag.mohapatra, beneharo.reveron-baecker, max.hock, thomas.hamacher}@tum.de Abstract—This paper introduces a surrogate modeling framework through Graph Neural Network (GNN), approximating the output of a Mixed Integer Linear Programming (MILP) based energy grid planning model. The surrogate is able to predict the net demand of an active distribution grid at the point of coupling with the upstream transmission grid, while taking into account different DER penetrations and paradigms of sectorcoupled flexibility and coordination. Additionally, it predicts the node voltages of the grid using the temporal GNN architectures and the power flow equations as a physical law in the training loss. The GNN surrogates are able to predict the net demand with under 1% accuracy and the best surrogate among them has a voltage prediction MAE of 0.002 pu for the grid. The proposed framework extends our previous results for LSTM and Transformer-based surrogates by using only one-third of the training data and providing a multi-output surrogate that is used to predict the node voltages. Once trained, the GNN surrogates can evaluate new scenarios in under 5 seconds. The voltage predictor can be used to identify critical buses and quickly eliminate scenarios which would lead to voltage violations without going for a full run of the optimization model. Index Terms—Machine Learning, Graph Neural Network, Energy System Modeling, Active Distribution Grids, PhysicsInformed Training I. INTRODUCTION Prosumer operation and sector coupling paradigms, such as flexibility and coordination, are needed to accommodate the renewable assets coming up in the active distribution grids (ADG). Energy system modeling (ESM) through convex optimization frameworks is the prevailing tool of choice among practitioners to understand and analyze these changes in ADG and consequently their impact on the net aggregated demand seen by the upstream transmission grid. Detailed ESM for ADG involves solving two optimization problems - capacity planning and economic dispatch, to be solved together or separately. In its combined form, such an ESM structure could determine the optimal DER installation capacity at every node of the ADG, along with a costoptimal dispatch schedule of required granularity, for a given demand and weather time series, alongside technology and operation costs. For an ADG, there are additional operation and quality of service constraints, sector coupling, flexibility and coordination mechanisms. The net demand seen by the upstream grid is then calculated by aggregating the net demand at every ADG node. This is a data of interest for upstream grid planning, as explained in [1]. A quick and reasonable estimate of the net aggregated demand could help narrow down the vast decision tree of operating and expansion decisions for the grid planner down to a few feasible candidates. Thereafter, the computational resources for detailed scenario analysis could be efficiently utilized. It should be noted that ESM frameworks require a large quantity of data to set up. They also use Mixed-Integer Linear Program (MILP) solvers which, depending on problem size, can lead to infeasibility on standard hardware. They also generate intermediate time-series, at a great computational cost, which are of no particular interest to the grid planner concerned with the aggregated behavior of ADG. A. Research Gap Surrogate modeling has emerged as a viable strategy to accelerate such analyzes. A surrogate model is a learned approximation of the input–output relationship of a complex simulation or optimization. We have shown in [2] that neural network surrogates with Long Short Term Memory (LSTM)/Transformer architecture can be successfully used to perform quick and reasonably accurate scenario analysis in place of a detailed ESM scenario run. However, the surrogate models in [2] typically predict only a single aggregate quantity, such as net demand or total cost, and often ignore network variables like bus voltages. This limits their usefulness for distribution grid applications, where voltage regulation and node-specific constraints are also a critical aspect for the grid planner. Additionally, we had used a large amount of data from 27 distinct scenarios of our reference ADG to generate the training data in our previous work. Therefore, a reduced training set to produce a similar quality of result would be considered a quality improvement for our surrogates. Moreover, existing purely data-driven surrogates without consideration of important physical laws can produce physically implausible results, especially when extrapolating to scenarios with limited training data. B. Contribution In this paper, we introduce a Graph neural network (GNN) based framework for building surrogate models of ADG to predict aggregated net demand from ESM and node-level voltage magnitudes. Incorporating spatial and temporal dependencies via a hybrid GNN architecture, our surrogate provides physically consistent predictions while significantly lowering the training data requirement.979-8-3315-2503-3/25/$31.00 ©2025 IEEE Fig. 1. Schematic overview of the GNN surrogate creation process. •GNN based surrogates are able to achieve similar performance as the best surrogates in [2], while using one-third of the training data. •GNN based surrogates forecasts voltage magnitudes at all nodes along with the net aggregated demand. This enables better scenario analysis and constraint verification. •Maintains a MAE of 0.002 pu in voltage predictions, ensuring accuracy within the tight bus voltage min-max band. II. METHODOLOGY The methodology visualized in Fig. 1 is similar to the structure presented in [2] except the differences emerge in the ”Surrogate Creation” step, due to the change from LSTM/Transformers architecture to GNN. Our change is motivated by two key factors. Firstly, GNNs have been shown to work with less training data without sacrificing performance and they have embedded physics from the graph structure which is missing in other neural network counterparts [3]. Secondly, power grids are by definition a graph network. The nodes can be denoted by tuples containing power injections and complex voltages, while the edge features can be line flows and impedance. Thus, the very nature of our physical system lends directly to a graph-based representation for any modeling or analysis. Recent works, such as [4] and [5], have leveraged these properties and proposed power system applications using GNN. A. Data preparation The data preparation step follows the sequence shown in the top row of Fig. 1. We adopt the same ADG as described in [2], based on the Kerber benchmark low-voltage grid [6] enhanced with distributed energy resources (DERs) such as photovoltaics (PV), heat pumps (HP), and electric vehicle (EV) chargers. Changing the installed DER capacities: The approach for generating data through scenario padding with varying DER capacities is followed as shown in [2], where 27 scenarios per grid are constructed using permutations of PV, BESS, and HP capacities, as seen in second row of Fig. 1. We then pick 10 randomly chosen scenarios from this list to create our limited training dataset for the GNN based surrogate, as compared to the full 27 scenario dataset for the LSTM and Transformer based surrogates in [2]. This results in 87,600 input-output TABLE I TRAINING DATASET CREATED WITH 8760 (1 YEAR)INPUT-OUTPUT PAIRS. Type with units Spatial resolution Temporal resolution Variability across scenarios Inputs System Demands, Solar Irradiation, Relative COP Electrical demand, kWel For each bus in the grid Hourly Constant Domestic hot water demand, kWth Space heating demand, kWth Mobility demand, kWel Solar irradiation capacity factor For each bus in the grid Heat pump COP Charging station availability System Capacities Battery capacity, kWel For each bus in the grid As installed 3*3*3 variation across scenarios Rooftop PV, kWp Heat Pump, kWel System Capacities Charging stations, kWel For each bus in the grid As installed Constant Output Net demand of the grid, kWel Grid aggregated Hourly Variable Active power, kWel For each transmission line Hourly Variable pairs for a given ADG, enabling the surrogate to learn how variations in DER penetration affect net demand. We include active power flows on edges as additional features to incorporate constraints from LinDistFlow as a physics-based loss in the GNN training step. It should be noted here that the reactive power flows were ignored for computation simplification and also due to a lack of an appropriate household reactive power demand time-series. When available, these can be included in the ESM scenario and the surrogate creation steps without any loss of generality. B. Data Pre-processing Three key pre-processing steps were followed to link the MILP simulation data with the GNN model input as a structured graph data suitable for learning. They are also shown in the GNN model definition schematic shown in Fig. 2. 1) Node Feature Extraction: We extract the time-series of node-level features (e.g., electrical demand, DER capacities, efficiency factors, weather inputs) for every bus from each scenario in the ESM output. This is in contrast to the grid level Fig. 2. Specialized pre-processing of data for GNN surrogate model definition. aggregation done in [2]. These features are organized into a temporal node feature matrix of dimensions (time steps × nodes ×features), effectively capturing the evolution of each node’s attributes over the year. 2) Graph Connectivity and Edge Features: Each bus in the network is treated as a node, and each line connecting two buses is represented as an edge. We construct an edge index (adjacency matrix) that lists all node pairs (source and target bus indices) that are connected by a line. This edge index, along with an edge mapping, allows the model to understand which nodes exchange power with each other. This is visualised in Fig. 3. In addition, we compile an edge feature matrix that stores physical line attributes – specifically the line resistance and reactance for each edge. These attributes encode the electrical characteristics of the connections, providing the GNN with relevant physical parameters for each link in the network. Together, the node feature matrices and the graph connectivity (edge index + attributes) form the structured input that the GNN will process, ensuring that both temporal demand patterns and network topology are reflected in the model’s input space. Furthermore, to incorporate domain knowledge of power flow physics, we extract the active power flow timeseries along each line from the simulation results and align them with the corresponding edges in our graph, for every hour in each scenario. While these flow values are not direct inputs to the GNN, as shown in Fig. 3, , they are used to compute a physics-informed loss term during training for node voltage prediction. In particular, we employ a LinDistFlow-based loss that penalizes the surrogate if its predictions violate physical power flow relationships, using the extracted line flow data as ground truth. 3) Feature Scaling and Normalization: Before training, all input features and target values are scaled to normalized ranges to prevent features with larger magnitudes from dominating the learning process. We apply a MinMax scaler to each feature set (node features, edge attributes, and target output) to bring them to a consistent range. Our optimization model, by definition, balances supply and demand within each ADG. This results in a highly skewed net-demand distribution, which spikes at the first bin next to zero, as seen in Fig. 4. Although masking and feature scaling have been applied, this imbalance complicates our net-demand prediction approach: the network tends to overfit in regions with sparse data and under-represent rare but relevant patterns. Future work could mitigate this by techniques such as, downsampling near-zero values or oversampling others, and adding more domain specific regularization. We mention this in the current manuscript as an important learning for surrogate creation from time-series output of ESM frameworks. C. Creating the Surrogates We develop several surrogate models for ADG based on Temporal Graph Neural Network (TGNN) architectures [7]. Fig. 3 shows our primary architecture which combines a Graph Attention Network (GAT) for spatial learning and aGated Recurrent Unit (GRU) for temporal modeling. The GAT mechanism enables nodes to aggregate information from their neighbors by learning attention scores, effectively identifying which nodes in the network have the most significant influence on a given node’s behavior. The GRU component captures sequential dependencies, processing time-series data to learn temporal trends in power demand, generation, and voltage dynamics. In addition to the GAT-GRU architecture, we evaluate several other TGNN-based models, each with distinct design characteristics and advantages: •Temporal Graph Convolutional Network (TGCN): Utilizes local node aggregation with convolution operations to capture immediate neighborhood dependencies over time. •Temporal Graph Attention Network (TGAT): Builds on neighborhood attention, incorporating temporal encoding to enhance node-level predictions. •Temporal Graph Transformer (TGT): Employs global attention mechanisms with positional and structural encoding, allowing the model to capture long-range dependencies in both space and time. •Temporal Graph Isomorphism Network with Edge Features (TGINE): Integrates edge-aware message passing and emphasizes isomorphism properties, providing better generalization across different graph structures. Physics-Informed Learning: To enhance physical consistency in voltage predictions, we integrate a physics-informed loss function derived from the LinDistFlow equations: vj=vi−2(rijPij +xijQij)(1) Here, viand vjrepresent squared voltage magnitudes at nodes iand j, while rij and xij are the line resistance and reactance, respectively. Pij and Qij denote the active and reactive power flows across the line. This loss penalizes violations of Kirchhoff’s Voltage Law, ensuring the surrogate’s predictions respect fundamental power flow constraints. Here we would again point out that the reactive power flows were set to zero in our work as the underlying ADG model did not have sensible reactive power demand data. However, this can be rolled into the surrogate for other ADG models, without any loss of generality and expected better performance. Fig. 3. TGNN illustration. Kerber grid; Radial structure; Nodes as grid bus; Edges as transmission lines; Message passing = Neighborhood averaging (AVG) + Linear transformation Fig. 4. Target output distribution is skewed for net grid demand. Implementation and Training: All models were implemented in Python using the PyTorch Geometric framework for efficient graph processing and spatio-temporal batching. Training was conducted using mini-batch stochastic gradient descent with early stopping and regularization to prevent overfitting. We employed a 10-fold cross-validation approach at the scenario level to assess the models’ ability to generalize across unseen grid configurations. Hyperparameter optimization was performed using Optuna. Most models were trained on a local MacBook Pro - 2021, Chip M1 Pro, Memory 32GB. Part of the computations were also ported to the compute cluster, using a Nvidia A100 GPU, at the Leibniz-Rechenzentrum (LRZ)-AI System provided by the Technical University of Munich. Since the code did not support parallel computation in its current state, the local Macbook had superior performance than the compute cluster. Upgrades of this work could leverage code parallelization to use multiple GPUs available at the compute clusters for faster performance. III. RESULTS This section demonstrates results of TGNN models on net demand and bus voltages predictions. For net demand forecasting we trained all TGNN architectures listed in Table II, while only the best performing TGNN architecture was used for bus voltages. A. Net demand prediction The summary of results is shown in Tab. II. We list the detailed performance metrics, both statistical and physical, and training time for the best performing models. The last row of the table corresponds to the best performing model from [2], to compare the GNN based surrogates against the LSTM/Transformer based surrogates at a glance. The first four TGNN models in Tab. II do not have the explicit physicsloss from II-C, but they still include the inherent physical information of the graph connectivity and edge features such as resistance and reactance of the transmission lines. These GNN surrogates achieved a R2of 0.85–0.87 in net demand prediction across all scenarios, closely matching the performance shown by LSTM/Transformer based surrogate, while using only one-third of the training data. The PiTGAT model uses the additional physics-loss term from II-C in its training routine. Although this addition should improve the overall surrogate performance, we see that the R2and some other metrics show a drop compared to the pure TGNN architectures. An exhaustive weight tuning of the data loss versus the physics-loss could be used to better the current performance. We put this aside for future work as the topic of balancing physics and data loss in the training function merits its own manuscript. However, we still include the results for the sake of completeness of the current work. In terms of computation time, compared to the LSTM/Transformer surrogates from [2], the GNN-based surrogates take around eight times longer for hyperparameter tuning and four times longer for training. This is to be expected due to the lack of parallelization and also the inherent speed of GNN tuning compared to recurrent NN. However, the GNN surrogates needed one-third less data and still provide the inference in less than five seconds. Since most ESM analysis ends up being scenario runs at different demand or DER penetration levels, the GNN surrogates can be useful when comparatively less training data is available. Again, as highlighted in [2], we have to remark that the computation time comparison is not for benchmarking as the code structure and the hardware used were not exactly the same. Yet they do serve a representative purpose as standard, if not the same, hardware was used for both analysis. B. Voltage Magnitude Prediction The overall best performing TGAT model from Tab. II, was tested for the node voltage predictions. The surrogate successfully predicted node voltage magnitudes across the grid, with a MAE of 0.002 pu. Fig. 5 shows a random week voltage predictions compared to the ground truth, leading to additional observations - •The maximum and minimum voltage squares were 1.03 pu and 0.97 pu, respectively and the error remains well within TABLE II NET DEMAND PREDICTION USING TGNN MODELS Label R2MSE MAE [kW] Net Demand [%] Peak Demand [%] Peak Feed-in [%] Hypertuning Time [h] Training Time [h] Evaluation Time [s] TGCN 0.85 0.02 38.834 0.19 9.31 2.77 22:15:20 03:52:42 4.081 TGAT 0.81 0.06 30.551 0.58 7.39 3.01 24:20:29 04:04:12 4.311 TGT 0.87 0.02 32.172 0.46 14.20 1.79 25:01:34 05:20:12 4.823 TGINE 0.79 0.05 32.320 0.80 13.06 1.32 21:28:40 04:02:50 4.452 PiTGAT 0.65 0.03 61.402 7.50 4.15 57.50 27:07:32 05:16:27 3.470 TrainDeltaLoss Transformer24 TS-VMD5 - 0.048 19.441 0.882 0.826 5.455 03:15:09 01:05:36 4.509 Legend - TGCN: Temporal Graph Convolution Network, (no physics loss); PiTGAT: Physics-informed Temporal Graph Attention (LinDistFlow physics-informed loss included); TrainDeltaLoss Transformer24 TS-VMD5(TrainDeltaLoss: Distance-weighted sum of MSE and MAE on training data as loss function; Transformer24: Transformer with look-back window of size 24; VMD5: 5 mode variable mode decomposition); Fig. 5. Random week bus voltage squares predictions. regulatory voltage thresholds. This is visualized in Fig. 6. In other words, if the voltage constraint is limited to ± 5% (it could also be ±10% based on the grid directive), the sensible region of voltage prediction is limited to this. Therefore, the voltage prediction should have a tight spread within this range to be useful for any analysis. Our MAE when projected on the ±5% limited active voltage range from grid codes, have an average error of approximately 3.3%. •The tight spread of voltage predictions could identify voltage constraint violations, ensuring that infeasible scenarios can be flagged early. One can imagine a rapid scenario elimination workflow where a critical bus voltage is inferred across a variety of demand profiles and DER penetration rate, without having to re-run the ESM framework. •The TGAT model achieved R2= 0.70 with MAE = 0.002 pu for the voltage prediction. When switching to the PiTGAT model, the R2dropped to 0.48, suggesting once again that further tuning is needed to balance physics-based learning with data-driven accuracy. Fig. 6. MAE illustration with respect to the observed min-max range of grid code mandated voltage magnitudes. IV. CONCLUSION This work extended our previous surrogate modeling framework for ADG, through a GNN based approach, offering an alternative to traditional ESM scenario analysis. Leveraging the inherent graph nature of a power grid, the GNN archictectures capture both the spatial and temporal interlinks between nodes to provide similar performance as the previous work, while using only one-third of the training data. Furthermore, we extracted active power flows from the ESM output timeseries to setup a voltage magnitude predictor. We were able to achieve a very tight variation in our voltage predictions (±0.002 pu, or 3% of our permissible voltage range), showing that the predictions can be used with reasonable confidence for flagging critical scenarios. We also used simplified LinDistFlow equations as an additional physics-based loss to get physically consistent voltage predictions. This addition did not bring the expected improvements, highlighting the need to further tune the balance between the data-driven loss and physics-loss component in the training routine. Inspite of this, the overall results from the GNN-based surrogates can still be used in its current state, for rapid scenario analysis and identifying critical voltages on buses in ADG. REFERENCES [1] B. Reveron Baecker and S. Candas, “Co-optimizing transmission and active distribution grids to assess demand-side flexibilities of a carbon-neutral German energy system,” Renewable and Sustainable Energy Reviews, vol. 163, p. 112 422, Jul. 2022. DOI: 10.1016/j.rser.2022.112422. [2] A. Mohapatra, M. Hock, B. Reveron Baecker, and T. Hamacher, Surrogate Framework for Energy System Modeling, accepted in IEEE PowerTech 2025, 2025. [3] B. Sanchez-Lengeling, E. Reif, A. Pearce, and A. Wiltschko, “A gentle introduction to graph neural networks,” Distill, vol. 6, no. 8, Aug. 2021. DOI: 10.23915/distill.00033. [4] B. Donon, R. Couillet, and D. Guenard, “Neural Networks for Power Flow: Graph Neural Solver,” arXiv preprint arXiv:2010.00186, 2020. [5] A. Varbella et al.,Physics-Informed GNN for non-linear constrained optimization: PINCO a solver for the AC-optimal power flow, 2024. DOI: 10.48550/ARXIV.2410.04818. [6] G. Kerber, “Aufnahmef¨ ahigkeit von Niederspannungsverteilnetzen f¨ ur die Einspeisung aus Photovoltaikkleinanlagen,” Ph.D. dissertation, TU Munich, 2011. [7] E. S¨ ohner, M. Helm, and B. Jaeger, “Temporal graph neural networks,” in School of Computation, Information and Technology, Technical University of Munich, Chair of Network Architectures and Services, 2020.