Graph Neural ODE Digital Twins for Control-Oriented Reactor Thermal-Hydraulic Forecasting Under Partial Observability

Akzhol Almukhametov; Doyeong Lim; Rui Hu; Yang Liu

cs.LGApr 8, 2026

Graph Neural ODE Digital Twins for Control-Oriented Reactor Thermal-Hydraulic Forecasting Under Partial Observability

Authors: Akzhol Almukhametov, Doyeong Lim, Rui Hu, Yang Liu

Organizations: Department of Nuclear Engineering, Texas A&M University, College Station, TX 77843, USA · Argonne National Laboratory, Nuclear Science and Engineering Division, USA

Abstract

Real-time supervisory control of advanced reactors requires accurate forecasting of plant-wide thermal-hydraulic states, including locations where physical sensors are unavailable. Meeting this need calls for surrogate models that combine predictive fidelity, millisecond-scale inference, and robustness to partial observability. In this work, we present a physics-informed message-passing Graph Neural Network coupled with a Neural Ordinary Differential Equation (GNN-ODE) to addresses all three requirements simultaneously. We represent the whole system as a directed sensor graph whose edges encode hydraulic connectivity through flow/heat transfer-aware message passing, and we advance the latent dynamics in continuous time via a controlled Neural ODE. A topology-guided missing-node initializer reconstructs uninstrumented states at rollout start; prediction then proceeds fully autoregressively. The GNN-ODE surrogate achieves satisfactory results for the system dynamics prediction. On held-out simulation transients, the surrogate achieves an average MAE of 0.91 K at 60 s and 2.18 K at 300 s for uninstrumented nodes, with $R^2$ up to 0.995 for missing-node state reconstruction. Inference runs at approximately 105 times faster than simulated time on a single GPU, enabling 64-member ensemble rollouts for uncertainty quantification. To assess sim-to-real transfer, we adapt the pretrained surrogate to experimental facility data using layerwise discriminative fine-tuning with only 30 training sequences. The learned flow-dependent heat-transfer scaling recovers a Reynolds-number exponent consistent with established correlations, indicating constitutive learning beyond trajectory fitting. The model tracks a steep power change transient and produces accurate trajectories at uninstrumented locations.

Figures & tables

Panel of Figure 1 — Figure 1: Architecture of the proposed physics-informed GNN-ODE surrogate for thermal-hydraulic forecasting under partial observability. Plant-node features are represented as $[T_{\mathrm{norm}},mask,\dot{T}_{\mathrm{norm}}]$ , and a topology-guided missing-node initializer (TGMI) provides initial estimates for uninstrumented nodes. Type-specific MLP encoders map plant, actuator, and ambient inputs to a shared hidden space $H$ , followed by a three-layer GNN encoder in which plant states are updated via physics-informed message passing while actuator and ambient states are propagated unchanged. A controlled latent Neural ODE, integrated with Runge Kutta 4-th order (RK4), advances the latent state, and the decoder predicts node-wise temperature increments $\widehat{\Delta T}\in\mathbb{R}^{|\mathcal{V}_{p}|}$ , where $|\mathcal{V}_{p}|$ is a number of plant nodes. The control input $u$ is assembled from actuator-node and flow-node features.

Component / Sensor	Measurement Type	Uncertainty
T-Type Thermocouples	Loop Temperatures	$\pm 1.0~^{\circ}\mathrm{C}$
J-Type Thermocouples	Heater Cartridge Temperatures	$\pm 2.2~^{\circ}\mathrm{C}$
Variable-Area Transmitters	Loop Flow Rates (up to $0.15~\mathrm{kg/s}$ )	$\pm 5\%$
Gauge Pressure Transducers	System Pressures (up to $206.8~\mathrm{kPa}$ )	$\pm 1\%$
Heater Controller	Actuated Thermal Power	$\pm 5\%$

Table 1: Summary of experimental hardware measurement and actuation uncertainties.

Idx	Sensor	Mask	GNN–ODE node name
0	$\mathrm{T_{H,avg}}$	$m_{i}=1$	Heater rods average temperature
1	$\mathrm{T_{ch}}$	$m_{i}=0$	Mixing chamber average temperature
2	TF12	$m_{i}=1$	Leg 1 (hot), segment 1
3	TF13	$m_{i}=1$	Leg 1 (hot), segment 2
4	TF14	$m_{i}=1$	Leg 1 (hot), segment 3
5	HX1_L1	$m_{i}=0$	Heat Exchanger 1, primary side

Table 2: Nodal discretization and signal mapping scheme (“Uninstrumented” nodes are latent states with

m_{i}=0

Panel of Figure 4 — Figure 4: Horizontal, compact representation of the 3-loop thermal system graph topology.

Scenario	Description
Stepwise heat-load ramp	Heat load is increased from 1 $\rightarrow$ 5 kW in 1 kW increments every 60 s, with all three loop flow rates fixed at 0.14 kg/s.
Cascaded loop-flow changes	Cascaded flow-rate changes at 7.5 kW: Loop 1, 0.14 $\rightarrow$ 0.1 kg/s at $t=30$ s; Loop 2, 0.1 $\rightarrow$ 0.12 kg/s at $t=40$ s; and Loop 3, 0.07 $\rightarrow$ 0.14 kg/s at $t=50$ s.

Table 3: Transient scenarios used in this study. Each scenario is 300 s. long and is being generated by SAM model.

Node	$R^{2}$ (w/o $\tilde{\mathbf{u}}$ )	MAE (w/o $\tilde{\mathbf{u}})$	$R^{2}$	MAE	$\Delta$ MAE
Ch	0.99989	0.076	0.99996	0.064	-0.012
HX1_L1	0.93401	2.322	0.99328	0.686	-1.637
HX1_L2	0.92213	1.980	0.99537	0.448	-1.532
HX2_L2	0.93471	1.911	0.98953	0.632	-1.279
HX2_L3	0.88419	2.559	0.99196	0.559	-2.000

Table 4: Reconstruction performance for missing nodes on validation dataset.

Device/Mode	$M=1$	$M=16$	$M=32$	$M=64$
CPU 1-thread	146.02	35.31	16.26	8.47
CPU 128-thread	63.27	24.52	19.41	14.50
GPU (RTX 4090)	104.49	66.43	52.07	36.57

Table 5: No-plot rollout speedup factors for single-model and ensemble inference, computed as simulated (predicted) time divided by wall-clock runtime.

Figure 5: Forecasting comparison across the two transient scenarios listed in Table 3 . The red dashed line corresponds to non-perturbed GNN-ODE predictions, while the red background indicates the predictive uncertainty band between the 5th and 95th percentiles. The green background denotes the imposed power profile on the secondary y-axis. GT (SAM) denotes SAM ground truth, which is used to compare against GNN-ODE predictions and compute MAE.

Node	10s	30s	60s	120s	180s	300s
Heater_avg	0.14	0.44	0.76	1.52	2.31	0.55
TF12	0.23	0.94	1.78	2.71	2.96	6.59
TF13	0.26	1.11	2.03	2.83	2.83	4.68
TF14	0.28	1.24	2.29	3.34	3.46	4.86
TF15	0.10	0.31	0.58	0.95	1.32	2.53
TF11	0.16	0.55	1.09	1.96	2.51	2.97

Table 6: Forecasting rollout summary for observed nodes (GNN-ODE forecast).

Node	10s	30s	60s	120s	180s	300s
Chamber_avg	0.18	0.49	0.85	1.34	1.87	1.17
HX1_L1	0.74	0.87	1.22	1.72	2.10	3.12
HX1_L2	0.53	0.65	0.87	1.29	1.53	2.82
HX2_L2	0.72	0.76	0.93	1.22	1.52	2.11
HX2_L3	0.52	0.51	0.66	0.88	1.06	1.71
AVERAGE	0.54	0.66	0.91	1.29	1.62	2.18

Table 7: Forecasting rollout summary for missing nodes (virtual sensors + GNN-ODE forecast).

Figure 6: Long-horizon rollout Mean Absolute Error (MAE) versus forecast horizon on held-out SAM test transients for (a) observed nodes and (b) unobserved (missing) nodes.

Figure 7: Surrogate predictions versus experimental measurements for observable facility nodes during steep power changes ( $0\rightarrow 10\rightarrow 5\rightarrow 7$ kW). Shaded bands denote 95th-percentile confidence intervals from $M=64$ parallel ensemble rollouts.

Figure 8: Inferred thermal trajectories for permanently uninstrumented (hidden) nodes during steep power changes (

0\rightarrow 10\rightarrow 5\rightarrow 7

kW). Initial states are seeded by TGMI; shaded bands show 95th-percentile confidence intervals from

M=64

ensemble rollouts.

Surrogate Modeling

Surrogate modeling builds faster or more interpretable approximations of expensive simulations and black-box models.

See all 222 papers

HyperODE: Zero-Shot Surrogate for Simulation and Inference of Dynamical Systems
Aug 1, 2026Ajitesh SrivastavaNeural Surrogate Modeling Zero-Shot Learning
Northeastern University
Physics-based Digital Twins for Integrated Thermal Energy Systems Using Active Learning
May 7, 2026Umme Mahbuba Nabila, Paul Seurin, Linyu Lin +1Surrogate Modeling Sparse Identification of Nonlinear Dynamics
Department of Nuclear Engineering and Radiological Sciences, University of Michigan, Ann Arbor, MI 48109, United States · Nuclear Science & Technology Division, Idaho National Laboratory, Idaho Falls, ID 83415, United States · Department of Computer Science and Engineering, University of Michigan, Ann Arbor, MI 48109, United States
INDEQS: Informed Neural controlled Differential EQuationS
Jun 17, 2026Michael Detzel, Gabriel Nobis, Kristiyan Blagov +3Neural Controlled Differential Equations Spatiotemporal Forecasting
Fraunhofer Heinrich Hertz Institute · Technische Universität Berlin

Aug 1, 2026J/K move · Enter open · S save

HyperODE: Zero-Shot Surrogate for Simulation and Inference of Dynamical Systems

Ajitesh Srivastava

Northeastern University

Neural Surrogate Modeling Zero-Shot Learning Simulation-Based Inference Surrogate Modeling Ordinary Differential Equations

Open paper AskPDF arXiv

Understanding and controlling complex dynamical systems often requires executing thousands of numerical simulations across vast parametric landscapes, which is time-consuming. Machine learning surrogates significantly accelerate simulation by predicting state trajectories across different initializations and parameter values. However, surrogate models are specialized to one simulation model. Modifying the underlying differential equations - e.g., adding a physiological state or altering an epidemiological contact network - renders trained models obsolete and forces computationally expensive retraining from scratch. We introduce HyperODE, a surrogate capable of operating across an entire class of approximately mass-conserving compartmental models without retraining. By mapping the structure of ordinary differential equations (ODEs) into directed hypergraphs, HyperODE decouples the functional form of system interactions from the neural network architecture. HyperODE takes a compartmental model in the form of an ODE with an arbitrary parameter distribution defined through quantiles and transforms it into a hypergraph. It outputs the distribution of the trajectories for all the states in the original ODE in the form of quantiles. We then use this surrogate to build an encoder that takes a noisy trajectory and outputs a distribution over the parameters of the original ODE, thus calibrating the model in a single pass. On families and system sizes never seen in training, HyperODE produces calibrated quantile bands in a single forward pass, with weighted-interval score and coverage on par with specialized surrogates for each structure. For inverse inference, HyperODE produces calibration from noisy state trajectories in a few milliseconds with a single shared encoder, competitive with existing methods. HyperODE extends zero-shot to ODEs that break mass conservation and to external forcing.

Abstract

Surrogate Modeling

Explore similar work

HyperODE: Zero-Shot Surrogate for Simulation and Inference of Dynamical Systems

Physics-based Digital Twins for Integrated Thermal Energy Systems Using Active Learning

INDEQS: Informed Neural controlled Differential EQuationS