Towards Fast GNN Surrogates for CO2 Migration in Complex Geological Formations

The use of variable well perforation and injection strategies can improve the efficiency of geological carbon storage operations. We develop a new multimodal auto-regressive transformer surrogate to model these operations under geological uncertainty. A modified SEAM CO2 geomodel, which involves a faulted system with three stacked aquifers, is considered. The two injection wells are perforated in stages, from bottom to top, with the stage durations and individual well injection rates treated as control variables. The surrogate model processes three input modalities - the 3D geomodel, scalar parameters characterizing relative permeability functions, and control variables - through separate encoders. These are fused via self-attention in a transformer encoder, and a temporal decoder generates predictions auto-regressively through encoder-decoder cross-attention. The surrogate is trained, using 4000 GEOS flow simulations, to predict saturation and pressure at monitoring locations, total injected and mobile CO2 mass, and saturation footprints. For a new test set, involving randomly sampled geomodels and control variables, the surrogate achieves a median saturation MAE of 0.028 and median relative errors of 0.2-5% for the other quantities of interest. Importantly, it captures the switch from rate to bottom-hole-pressure control. The surrogate model is used within a hierarchical Markov chain Monte Carlo data assimilation procedure for a synthetic true model under three operational strategies. Substantial uncertainty reduction is achieved for key metaparameters, particularly the fault permeabilities. Posterior predictions for saturation footprints and total injected and mobile CO2 mass are also shown to be generally consistent with true model results.

Accurate prediction of bottom-hole pressure (BHP) and CO2 plume migration is essential for safe geological carbon storage, yet practical simulations often rely on truncated domains where artificial boundaries distort pressure diffusion and CO2 saturation footprints. In this study, we evaluate how boundary-condition fidelity affects BHP and CO2 plume prediction by comparing ten reduced-domain boundary treatments against full-domain reference simulations in homogeneous and heterogeneous reservoirs. We test uniform pore-volume multipliers, transmissibility modifiers, corner-adjusted pore-volume corrections, layered corrections, and gradual modifiers using BHP RMSE, NRMSE, peak pressure deviation, and plume Intersection over Union (IoU) as performance metrics. Our results show that conserving corner pore volume is the most important requirement for truncated-domain modeling. We find that uniform treatments which neglect corner storage generate large pressure errors, with BHP RMSE of 362 to 382 psi in the homogeneous model and 250 to 304 psi in the heterogeneous model, and yield plume IoU values near 0.80 to 0.84, indicating roughly 16 to 20% of the combined plume area is misrepresented. Corner-adjusted scenarios substantially reduce pressure errors and raise plume IoU above 0.94, but we observe that transmissibility correction is not universally beneficial. In homogeneous reservoirs, uniform transmissibility adjustment improves pressure fidelity; in heterogeneous reservoirs, it can over-restrict flow across variable-permeability boundary faces, increasing BHP error and contracting the predicted plume. We find the gradual modifier with transmissibility correction provides the most consistent performance, achieving BHP NRMSE below 3.7% and plume IoU above 0.97 in both reservoir types.

Compressible multiphase flows involving shocks and material interfaces arise in applications such as bubble collapse and droplet breakup, where strong nonlinear interactions produce complex interface deformation, mixing, and multiscale dynamics. Developing reliable machine learning surrogates for these flows remains challenging due to the simultaneous presence of compressibility, sharp discontinuities, and multiphase effects. In this work, we introduce the first large-scale benchmark specifically designed for shock-driven compressible multiphase flows, comprising 2.4 TB of high-fidelity 2D and 3D datasets featuring shock-induced bubble collapse and droplet breakup. We evaluate diverse surrogate model families on our benchmarking framework: Neptuna {https://github.com/tumaer/Neptuna}, including convolutional, spectral, transformer-based, and pre-trained PDE foundation models. Beyond standard MSE training, we investigate composite losses combining MSE with Sobolev, interface-aware, and structure-aware terms, together with adaptive loss balancing using SoftAdapt and GradNorm. Evaluation includes pointwise, spectral, feature-focused, structural, and physics-informed metrics. Results show that no single model performs best across all datasets and metrics, while composite losses significantly improve interface preservation and spectral fidelity. Among adaptive weighting strategies, SoftAdapt provides the most consistent improvements with almost no overhead compared to MSE-only training.