🤖 machine learning

Towards Fast GNN Surrogates for CO2 Migration in Complex Geological Formations

This paper proposes an end-to-end graph neural network surrogate that leverages anisotropic message passing and autoregressive latent dynamics to accurately forecast CO $_2$ plume migration and fluid properties in complex geological formations, demonstrating competitive performance on the challenging SPE11A benchmark.

Original authors: Rodrigo S. Luna, Thiago H. N. Coelho, Luiz S. L. Neto, Roberto M. Velho, Adriano M. A. Cortes, Renato N. Elias, Alexandre G. Evsukoff, Fernando A. Rochinha, Mauricio Araya-Polo, Herve Gross, Alvaro L.

Published 2026-06-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Rodrigo S. Luna, Thiago H. N. Coelho, Luiz S. L. Neto, Roberto M. Velho, Adriano M. A. Cortes, Renato N. Elias, Alexandre G. Evsukoff, Fernando A. Rochinha, Mauricio Araya-Polo, Herve Gross, Alvaro L. G. A. Coutinho

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict how a drop of ink spreads through a sponge. In the real world, that sponge isn't uniform; it has some parts that are tight and hard to squeeze through, and other parts that are loose and easy. If you want to know exactly where the ink will be in an hour, you could run a super-accurate physics simulation. But that's like trying to calculate the path of every single water molecule—it takes forever and requires a massive computer.

This paper introduces a new "smart shortcut" (a machine learning model) that learns the rules of how that ink spreads, but does it thousands of times faster. Here is how they did it, explained simply:

1. The Problem: The "Sponge" is Complicated

The researchers are studying how Carbon Dioxide (CO2) moves underground to be stored safely. The ground isn't a flat, smooth sheet; it's a messy, complex 3D structure made of different rock types (like a giant, irregular sponge).

The Old Way: Traditional computer simulations treat the ground like a grid of tiny boxes. They calculate the physics step-by-step. It's accurate, but it's so slow that you can't run it enough times to test different scenarios or predict the future with high confidence.
The Challenge: The CO2 doesn't just spread evenly. It gets "stuck" in some rocks, rushes through others, and forms weird, finger-like shapes (called "fingering") because of gravity and pressure. Standard AI models often miss these tricky details because they treat all directions the same.

2. The Solution: A "Smart Map" with Directional Sense

The authors built a new type of AI called AnisoMeshGraph-LSTM. Let's break down what that means using an analogy:

The Graph (The Map): Instead of a rigid grid, they turned the underground rock formation into a "social network" or a map of connected cities. Each rock cell is a "node" (a city), and the connections between them are "edges" (roads).
The "Anisotropic" Secret (The Traffic Rules): Most AI models treat traffic on these roads as if it flows equally in all directions. But in reality, if a road is steep or the ground is slippery, the flow goes one way much faster than another.
- The authors taught their AI to look at the geometry of the rocks. It learned that if a road (connection) is aligned with a steep slope or a high-permeability rock, the "message" (the CO2) should travel faster that way.
- Analogy: Imagine a messenger running through a city. A standard AI tells the messenger to run to all neighbors equally. This new AI tells the messenger: "Hey, the road to the North is a highway, but the road to the East is a muddy swamp. Run North fast, and ignore the East for now." This allows the AI to predict the "finger-like" shapes of the CO2 much better.

3. The Time Machine (The LSTM)

Predicting the future isn't just about looking at the map; it's about remembering the past.

The AI uses a special memory component (called GraphConv-LSTM) that acts like a diary. It doesn't just guess the next second; it remembers the history of how the CO2 moved.
Instead of guessing the entire new state from scratch, it only guesses the change (the difference) from the current state. This is like a driver who doesn't need to recalculate the whole route every second, but just adjusts slightly based on the last turn. This keeps the prediction stable even when looking far into the future.

4. The Test: The "SPE11A" Challenge

To see if this worked, they tested it on a famous industry benchmark called SPE11A.

The Setup: They simulated a lab experiment where CO2 is injected into a rock formation. They created 10 different "versions" of this rock formation by randomly changing how porous or permeable the rocks were, just to make sure the AI wasn't just memorizing one specific map.
The Result:
- Speed: The traditional simulation took about 2 hours to run. The new AI model took about 160 seconds (less than 3 minutes) to predict 50 steps into the future. That's a massive speedup.
- Accuracy: The AI was very good at predicting where the CO2 would be and how dense the liquid would be. It handled the tricky "finger" shapes well.
- The Weakness: The AI struggled a bit at the very beginning of the injection (the first few minutes). It's like a driver who is great at cruising on the highway but gets a little nervous when first merging onto the road. Once the CO2 started moving, the predictions became very accurate.

5. What They Found (and Didn't Find)

Success: The model successfully learned to predict the movement of CO2 in complex, messy underground environments without needing to solve heavy physics equations every time. It is fast and accurate enough for long-term monitoring.
Limitations:
- The test was done on a 2D slice of rock (like a flat map), not a full 3D building.
- It didn't account for the rocks physically cracking or shifting (geomechanics).
- It still gets a little "foggy" at the very start of the process.

In a nutshell: The authors built a "smart navigator" for underground CO2 storage. By teaching the AI to respect the specific direction and shape of the underground rocks, they created a tool that predicts the future of the gas plume almost instantly, with enough accuracy to be useful for safety monitoring, though it still needs work on the very first moments of injection and in full 3D environments.

Technical Summary: Towards Fast GNN Surrogates for CO2 Migration in Complex Geological Formations

Problem Statement
Accurate forecasting of CO2 plume migration in complex geological formations is essential for the safe deployment of Carbon Capture and Storage (CCS) technologies. Traditional numerical reservoir simulators, while physically reliable, incur prohibitive computational costs for long-term simulations and uncertainty quantification. While recent scientific machine learning approaches (e.g., Fourier Neural Operators, DeepONets, Graph Neural Networks) offer speedups, standard Graph Neural Networks (GNNs) often rely on isotropic message passing. This uniform propagation of information fails to capture the strongly directional transport inherent in reservoir flows, which is driven by grid geometry, permeability contrasts, gravity, and geological heterogeneity.

Methodology
The authors propose AnisoMeshGraph-LSTM, an end-to-end graph neural surrogate tailored for CO2 plume migration. The model is built upon the MeshGraphNet architecture but introduces three critical modifications to address the specific physics of subsurface flow:

Anisotropic Message Passing: Unlike standard isotropic mechanisms, the proposed model employs a geometry-conditioned edge-weighting mechanism. Interaction weights are computed via learnable MLPs based on geometric features (displacement vectors and Euclidean distances) and edge embeddings. This biases message aggregation toward physically relevant transport directions, explicitly modulated by geometric and physical edge attributes rather than relying solely on unconstrained node representations.
Latent Temporal Dynamics: To mitigate error accumulation during long-horizon rollouts, the model integrates a GraphConv-LSTM module after each message-passing block. This component captures temporal dependencies in the latent node embeddings using a spectral graph convolution operator (implemented via Chebyshev polynomials) rather than standard convolutions.
Residual Autoregressive Formulation: Instead of predicting absolute state values directly, the decoder predicts increments (residuals) of the physical state. The dynamic state is advanced using an explicit Euler update. The model is trained using multi-step supervision in a closed-loop autoregressive manner, where the model receives its own previous predictions as input, rather than ground truth, to ensure stability over extended forecasting horizons.

Experimental Setup
The approach is evaluated on the SPE11A benchmark, a two-dimensional laboratory-scale test case characterized by sharp gas-water interfaces, strong advective transport, and rapid convective mixing with fingering development.

Data Generation: 10 trajectories were generated by randomly varying the permeability and porosity of seven facies within the bounds of the benchmark.
Resolution: A 2 cm grid resolution was selected as a compromise between accuracy and computational cost, reducing storage requirements by ~50% compared to the 1 cm benchmark resolution while maintaining quantitative consistency.
Graph Representation: The mesh is converted to a static graph with 25,992 nodes and 51,557 edges. Node features include physical state variables (gas saturation, liquid density), cell volume, permeability, porosity, and cell type. Edge features include displacement vectors and Euclidean distances.
Training: The model was trained on 90% of the data (10 trajectories) with a 10-step multi-step supervision horizon. Three training regimes were tested: injection-only, post-injection-only, and a combined regime.

Key Results

Accuracy: The model produces competitive forecasts for gas saturation and liquid-phase density. For the combined training regime, the 50-step cumulative RMSE for gas saturation remains below 2.5%, and for liquid-phase density, the relative error is on the order of $10^{-5}$ .
Long-Horizon Stability: The error growth remains bounded over extended forecasting horizons (up to 50 steps). The anisotropic mechanism successfully captures directional transport effects, including sharp interfaces and density-driven fingering patterns.
Performance: Training on 10 trajectories takes approximately 22 hours on an NVIDIA H100. Inference for a 50-step forecast takes only 160 seconds, compared to 111–135 minutes required for a single high-fidelity numerical simulation.
Comparison: The proposed model significantly outperforms a baseline (standard MeshGraphNet without anisotropic passing or residual formulation), particularly in reducing diffuse errors in the fingering regions.
Limitations: The authors note that prediction accuracy is lower during the early injection phase where CO2 spreads rapidly. The current work is limited to a 2D laboratory-scale problem without geomechanics coupling.

Significance and Claims
The paper claims that the AnisoMeshGraph-LSTM architecture demonstrates that geometry-conditioned anisotropic message passing is beneficial for graph-based simulators in subsurface flow applications. By explicitly capturing directional transport mechanisms, the model achieves stable long-horizon predictions and generalizes to unseen geological configurations. The authors conclude that this approach opens promising directions for physically informed surrogate models, though they modestly acknowledge that extending these methods to full 3D problems with coupled flow and geomechanics remains an open challenge.