Source-linked AI summary
Graph neural networks for exoplanet atmospheres
Antonia Vojtekova, Kai Hou Yip, Ingo P. Waldmann, Nikolaos Nikolaou, Olivia Venot, Ahmed Faris Al-Refaie, Bruno Merín
TL;DR
Disequilibrium chemistry is a major computational challenge in atmospheric retrievals, and earlier convolutional surrogates introduce spatial bias inconsistent with reaction-network topology. This study uses a graph neural network to model chemical connections and reports accurate abundance and spectral predictions with faster-than-kinetic-code forward modeling.
Problem
Disequilibrium steady-state calculations remain a major computational challenge, while the earlier U-Net introduced an artificial spatial bias inconsistent with chemical reaction-network topology.
Source: §1 · p. 1, §6 · p. 15Method
A graph neural network encodes chemical connections between species and reactions, using atmospheres generated with Guillot temperature–pressure profiles and the Venot+2020 scheme.
Source: §2 · p. 2, §5 · p. 11, §6 · p. 15Results
More than an order of magnitude speed-up relative to kinetic codes accompanies accurate abundance reconstruction across the sampled parameter space and spectral errors consistent with JWST and Ariel expectations.
Source: §5 · p. 15, §6 · p. 16Takeaways & Limitations
The speed-up enables inclusion of disequilibrium chemistry in retrieval analyses, while performance varies near C/O = 1 and at low temperatures.
Source: §5 · p. 15, §6 · p. 16Solver-precision outliers can cause abrupt log(VMR) drops and jitter, while decreasing absolute tolerance substantially increases computational cost.
Source: §2.2.1 · p. 3
Abstract
from arXiv · showhide
Calculating disequilibrium chemistry in exoplanet atmospheres remains a significant computational bottleneck in atmospheric retrievals. The increasing observational precision from facilities such as JWST and the Ariel mission requires including disequilibrium chemistry in these analyses. Previous studies have demonstrated that neural networks can emulate kinetic chemistry, although their spatial inductive bias does not align with the topology of chemical reaction networks. This study introduces a graph neural network surrogate that represents chemical species as nodes and temperature-dependent reaction rates as edges, thereby enabling information propagation along physically meaningful chemical pathways. The model is trained on atmospheres generated using the Venot+2020 chemical scheme and Guillot temperature-pressure profiles. The GNN accurately reconstructs disequilibrium abundances across the sampled parameter space and reduces the mean abundance error by a factor of approximately 3 compared to the previous U-Net model. When applied to transmission spectra, most predictions fall within the observational precision expected for JWST and Ariel, with only about 7% of test atmospheres exceeding a 20 ppm mean spectral error. Performance variations are primarily observed in chemically transitional regimes near a carbon-to-oxygen ratio of one and at low temperatures. An evaluation of the boundary-case planet WASP-39b demonstrates effective performance under a moderate domain shift. Perturbation analysis indicates that disturbances propagate along chemical connectivity rather than spatial adjacency, confirming that the architecture captures the structure of reaction networks. These results suggest that GNN surrogates provide accurate, computationally efficient predictions of disequilibrium chemistry, facilitating integration into the atmospheric retrieval pipeline.
1 INTRODUCTION
Earlier U-Net work accelerated disequilibrium-chemistry calculations but introduced a spatial bias that may limit physicality and interpretability. Because chemical schemes are naturally graphs, graph neural networks offer a representation aligned with species and their interactions.
1 INTRODUCTION
- Explainability analysis found that the U-Net’s data-matrix representation and 2D convolution introduce a bias that may limit performance, physicality, and interpretability.The earlier model’s overall performance was described as satisfactory.Source: §1 · p. 1
- Chemical schemes naturally form graphs, with species as nodes and chemical interactions as edges.This structure motivates adopting a graph-based representation for disequilibrium steady-state abundances.Source: §1 · p. 1
- GNN message passing updates a target node using messages from neighboring nodes connected by edges.The passage identifies MPNNs, GCNs, and GATs as architectures relying on this core principle.Source: §1 · p. 1
- GNNs have been applied in chemistry, medicine, biology, engineering, astrophysics, and planetary science.Planetary-science examples include exploring Solar System orbital mechanics.Source: §1 · pp. 1–2
2 DATA
The dataset combines synthetic atmospheres with Guillot temperature–pressure profiles and the Venot+2020 chemical scheme. It is filtered, formatted as graph inputs, and normalised for training a surrogate of disequilibrium abundances.
2 DATA
- Synthetic atmospheres use non-isothermal Guillot profiles and the Venot+2020 chemical scheme, which contains 108 C/H/O/N-bearing species and 1906 reactions.The scheme also includes three condensed species; it was retained for consistency with the previous study.Source: §2 · p. 2, §2.1 · p. 2
2.1.1 Global parameters
- Samples span planetary, atmospheric, and stellar parameters, with 128 pressure layers from 1 × 10^-6 to 1 × 10^1 bar.Planet mass and radius selection is restricted, including surface-gravity bounds of 4–40 m s^-2.Source: §2.1.1 · p. 2, Table 1 · p. 2
2.1.2 Temperature–pressure profile
- The Guillot temperature–pressure profile is an analytical radiative-transfer solution for irradiated, close-orbiting giant planets.The model assumes a plane-parallel semi-grey atmosphere in radiative equilibrium, parameterised by mean visible and infrared opacities.Source: §2.1.2 · p. 2
2.1.3 Equilibrium chemistry
- Equilibrium profiles are generated with ACE chemistry, then FRECKLL calculates disequilibrium steady-state abundances with vertical mixing.The Eddy diffusion coefficient K_zz is held constant at 10^9 cm^2 · s^-1 across samples.Source: §2.1.3 · p. 2, §2.1.4 · p. 2
2.1.5 Reaction Rates
- Reaction-rate vectors encode chemical connections; rates are extracted per atmosphere and depend mainly on temperature, sometimes also on pressure or collisional bodies.The method covers modified Arrhenius/Kooij rates and low-pressure, high-pressure, and fall-off behavior for collisional reactions.Source: §2.1.5 · p. 3
2.2.1 Data Filtering
- Filtering selects 40 species, removes or interpolates solver-jitter outliers, and restricts profile temperatures to 700–2500 K.The retained dataset has 22,973 samples split 80/10/10 into training, validation, and test sets.Source: §2.2.1 · p. 3
2.3 Data Formatting
- The graph input represents normalized abundances as node features, reaction rates as directed edges, and atmospheric and planetary properties as global parameters.The network has 40 species nodes and 3,286 edges; each node and edge carries a 128-layer vector.Source: §2.3 · p. 3
2.3.1 Data Normalisation
- Abundances and reaction rates are log-transformed, while global parameters remain in linear space and are scaled independently.Abundances are clipped at a lower bound of −30; reaction rates are clipped and normalised to [0, 1].Source: §2.3.1 · pp. 3–4
3 METHODS
The study uses a message-passing GNN surrogate to predict disequilibrium steady-state abundances from equilibrium abundances. Its graph layers update species representations using reaction-connected neighbors, followed by a shared decoder.
3 METHODS
- The GNN surrogate predicts disequilibrium steady-state abundances from input equilibrium abundances.Source: §3 · p. 4
3.1 Model Architecture
- The architecture comprises graph layers built from edge-encoding, message-passing, global-encoder, and node-encoder modules, followed by a decoder.Each graph layer has its own set of module parameters.Source: §3.1 · p. 4, §3.1.2 · p. 4
3.1.1 Decoder MLP
- A shared decoder processes each final node representation independently to produce a 128-value output for each species.Source: §3.1.1 · p. 4
3.1.2 Graph Neural Network Layer
- Each graph layer updates species representations from reaction-connected neighbors, using encoded reaction rates and a residual node update.The node update combines aggregated neighbor messages, global information, and target-node features; the residual connection mitigates oversmoothing.Source: §3.1.2 · pp. 4, 6
3.2 Training Strategy
Training used MAE loss with Adam, learning-rate scheduling, and early stopping; convergence took 138 epochs and around 10 hours on an NVIDIA L4 GPU.
3.2 Training Strategy
- 138 epochs and around 10 hours on an NVIDIA L4 GPU were required for training convergence.Training used MAE loss, Adam, a learning-rate scheduler triggered by validation-loss plateaus, and early stopping; batch size was 8, with a maximum of 150 epochs.Source: §3.2 · p. 6
3.3 Hyperparameter Search
The authors searched key architectural hyperparameters, then selected a final model from five validation-loss leaders based on overall evaluation performance.
3.3 Hyperparameter Search
- The final model was selected from five top validation-loss candidates based on overall performance across evaluation metrics.The search varied architectural choices including MLP depth and width, graph-layer count, message aggregation, and global-parameter involvement; the U-Net comparator was retrained on the same dataset.Source: §3.3 · p. 6
3.4 Evaluation
Evaluation combines numerical and visual techniques. The numerical metrics include MAE and MAPE, which respectively measure absolute error and relative prediction error as a percentage of ground truth.
3.4 Evaluation
- Evaluation combines numerical metrics with visual techniques to assess predictions against ground truth.MAE is defined over atmospheres, pressure layers, and species.Source: §3.4 · p. 6
- MAPE expresses relative prediction error as a percentage of the ground truth.Source: §3.4 · p. 7
3.5 Full forward modelling
The forward-model evaluation converts predicted abundances into transmission spectra and compares them with spectra from ground-truth abundances. WASP-39b provides an edge-of-training-distribution test using a Guillot profile that approximates the retrieved profile within the training range.
3.5 Full forward modelling
- Predicted abundances are combined with planetary and stellar parameters and a Guillot-type profile to generate TauREx transmission spectra, including molecular absorption, collisionally induced absorption, and Rayleigh scattering.Clouds are excluded from the simulations.Source: §3.5 · p. 7
- Spectral MAE compares GNN-generated spectra with FRECKLL ground truth, emphasizing pressure layers that contribute more strongly to spectra than the highest and lowest layers.Spectra are displayed binned to JWST/NIRSpec G395H resolution against an adopted 20 ppm noise floor representative of best-case JWST precision and Ariel's target precision.Source: §3.5 · p. 7
- Each GNN layer independently encodes reaction-rate edges, jointly aggregates neighboring-node messages, and updates nodes using local features and shared global atmospheric context.The aggregated message is normalized before activation.Source: Figure 6 · p. 8, Figure 7 · p. 8, Figure 8 · p. 8
3.5.1 WASP-39b
- WASP-39b is tested with a Guillot profile approximating its retrieved structure while staying within training bounds, placing it near the distribution boundary in atmospheric and planetary properties.The retrieved profile contains parameters outside the sampled range.Source: §3.5.1 · p. 8
3.6 Perturbation test
The perturbation test probes how changing one molecule affects network outputs. It introduces noise to a selected molecule and examines whether effects follow chemical connections, while leaving an in-depth study of perturbation effects outside this work's scope.
3.6 Perturbation test
- The test hypothesizes that perturbing a molecule will most affect directly connected species, with some influence spreading further because the graph becomes fully connected after two steps.Source: §3.6 · p. 9
- A single-molecule perturbation adds Gaussian noise with mean 0 and standard deviation 0.1 to the selected molecule across atmospheric layers, leaving other species unchanged.The test is described as a post-hoc, model-agnostic local approach.Source: §3.6 · p. 9
- The study demonstrates improvement over the U-Net's spatial bias but does not investigate perturbation effects in depth.Source: §3.6 · p. 9
4 RESULTS
The GNN outperforms the U-Net on test-set abundance error, while most test spectra remain within the cited observational precision. Errors vary with atmospheric conditions, and the WASP-39b and perturbation analyses probe boundary behavior and chemical connectivity.
4 RESULTS
- Errors are higher for samples with C/O ratios around one and at lower temperatures.The highest-error sample has C/O ratio 1, while the 95th-percentile case performs substantially better, making the worst case a statistical outlier.Source: §4 · p. 9
- Abundance errors can have minor spectral impact because spectra primarily reflect mid-pressure layers, while larger errors often occur in low-pressure regions affected by jitter.The highest species errors include O2, C2H5, and C2H6, whose jittered profiles were interpolated during preprocessing.Source: §4 · pp. 9–10, Figure 11 · p. 11
- About 7% of test samples (168) have spectral MAE greater than 20 ppm, while roughly 2% exceed 50 ppm.Source: §4 · p. 10
4.1 WASP-39b
- For boundary-case WASP-39b, abundance MAE is 7.42 · 10^-3 versus 1.43 · 10^-3 at the test set's 95th percentile; its unbinned spectral MAE is 26 ppm.The performance degradation is described as consistent with expected extrapolative behavior and within the test distribution's upper tail.Source: §4.1 · p. 10
4.2 Perturbation
- The propagated error lacks the diagonal clustering seen with a CNN, suggesting that it follows molecular edges rather than spatial adjacency.Source: §4.2 · p. 10
- Perturbation hotspots reflect edge strength and interaction count; asymmetry may also track less-buffered species, though abundance, temperature, initial conditions, and baseline errors remain unmodeled.Four of five main hotspots are among the ten strongest reactant–product pairs.Source: §4.2 · p. 10
- Perturbing NH2 does not strongly affect H despite 30 shared connections, while 3CH2's largest error response is in C2H2 despite fewer shared edges.The authors suggest H and CO receive information from substantially more species than C2H2.Source: §4.2 · p. 10
4 RESULTS
- The GNN's mean test-set MAE is approximately 3 times lower than the U-Net's.Source: §4 · p. 9, Table 5 · p. 11
5 DISCUSSION
The discussion reports that GNN-generated spectra have errors consistent with JWST and Ariel expectations, while perturbation responses follow chemical connectivity rather than spatial adjacency. Errors remain higher in transitional and low-temperature regimes, and the study notes numerical and training-scope limitations.
5 DISCUSSION
- The GNN addresses the U-Net’s spatial inductive bias by encoding chemical connections and passing information along reaction pathways rather than grid-adjacent locations.The discussion also describes the GNN as more flexible for complex Guillot temperature–pressure profiles.Source: §5 · pp. 11–12
- Higher errors persist near C/O = 1 and at low temperatures, patterns the authors interpret as physically meaningful challenges rather than random failures or architectural limitations.Figure 12 clips MAE values at the 99th percentile to suppress extreme outliers for visualization.Source: §5 · p. 12, Figure 12 · p. 12
- Elevated errors for C2H5, C2H6, and O2 are attributed primarily to FRECKLL numerical artefacts and preprocessing interpolation, not inherent machine-learning architecture limits.The dataset discussion links solver-precision jitter to abrupt abundance drops and notes that tighter tolerance would substantially increase computational cost.Source: §5 · p. 12
- Spectral errors remain moderate for JWST and anticipated Ariel data, though occasional outliers occur; the authors suggest using surrogates during broad retrieval exploration and kinetic codes near convergence.The forward-modeling test also reports more than an order-of-magnitude speed-up relative to kinetic codes.Source: §5 · p. 15
- WASP-39b spectrum error remains below 30 ppm despite higher-tail abundance error, near the 95th percentile of evaluation errors and below cited JWST standard errors of 50–100 ppm.Its temperature–pressure profile lies near the training-distribution boundary; degradation is described as expected outside densely sampled regions, not instability.Source: §5 · p. 15
- Perturbations predominantly follow chemically meaningful connections, but a species’ reaction involvement also affects its robustness or susceptibility, unlike convolutional propagation through spatial kernel overlap.The authors interpret these tests as evidence that the GNN encodes chemical topology.Source: §5 · p. 15
- The study’s scope is limited to Guillot temperature–pressure profiles, fixed K_z, and training-set boundaries; newer chemical schemes require retraining and revalidation.Photochemistry and complex cloud-condensation models are also excluded from the present model.Source: §5 · p. 15
6 CONCLUSIONS
The paper presents a graph neural network surrogate for disequilibrium steady-state abundances, trained with Guillot profiles and the Venot+2020 scheme. It reports accurate abundance and spectral reconstruction, with perturbations following reaction connectivity, and concludes that the surrogate could support retrieval analyses.
6 CONCLUSIONS
- The surrogate encodes chemical connections between species and reactions and is trained on atmospheres using Guillot temperature–pressure profiles and the Venot+2020 chemical scheme.The paper frames this graph representation as addressing the U-Net’s artificial spatial bias.Source: §6 · p. 15
- The GNN accurately reconstructs disequilibrium abundances across the sampled parameter space; spectral errors remain within expected JWST and Ariel observational precision on average.Performance variations appear near C/O = 1 and at low temperatures; WASP-39b indicates robustness under moderate domain shift.Source: §6 · p. 16
- Accurate predictions make retrieval-pipeline use feasible, with orders-of-magnitude computational acceleration relative to kinetic codes.Source: §6 · p. 16
APPENDIX A: HYPERPARAMETER SEARCH
The appendix describes an iterative hyperparameter-selection process for the GNN, using stability and training and validation losses to narrow candidates before test-set selection.
APPENDIX A: HYPERPARAMETER SEARCH
- Unstable configurations and those with high training and validation losses were excluded during iterative model selection.Source: Appendix A · p. 16
- Stable, well-performing candidates were retrained across multiple fixed random seeds, and the final model was selected using test-set performance.Source: Appendix A · p. 16
- The search explored principal architectural and training hyperparameters, with the final configuration recorded in Table A1.Source: Table A1 · p. 17