Calibrating subgrid parametrizations of single-column ocean models via simulation-based inference
Authors: Luben M. C. Cabezas, Sacha Wendling, Aurèle Gallard, Gabriel Mouttapa, Julien Le Sommer, Pedro L. C. Rodrigues
Organizations: DEs, ICMC Federal Univ. of São Carlos, Univ. of São Paulo, Brazil · Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LJK, France · Univ. Grenoble Alpes, Inria, CNRS, IRD, Grenoble INP, INRAE, IGE, France
Subgrid parametrizations of vertical mixing in ocean models depend on free coefficients that cannot be measured directly and must be calibrated against high-fidelity references such as large-eddy simulations (LES). Existing approaches return point estimates and leave the associated uncertainty unquantified, a limitation when the inverse problem is ill-posed or when distinct parameter configurations fit the data comparably well. Simulation-based inference (SBI) addresses exactly this: given a prior and access to the simulator, it approximates the full posterior over parameters without requiring a tractable likelihood, at a cost set by the number of simulator evaluations. We apply it to tunax, a JAX-based single-column ocean model, to calibrate the coefficients of its k--ε closure. A blockwise PCA summary statistic compresses the simulator output along the depth axis while preserving its forcing--horizon--variable structure, making inference tractable at modest budgets. We compare neural posterior estimation and its sequential variants against a recent training-free approach built on a tabular foundation model. The latter recovers informative posteriors from a few hundred simulator calls, outperforming the trained estimators at every budget considered.
Figures & tables
Figure 1: RMSE between LES profiles and the tunax output evaluated at each posterior mean. Markers and error bars show the mean and SEM across 5 realizations, and the dashed line denotes the default tunax configuration. The gray marker denotes full-profile gradient calibration; its position counts 1000 simulator calls, although differentiation makes these calls more costly than ordinary forward simulations. NPE-PFN-MEAN reaches lower RMSE rapidly at small budgets and performs best among the posterior methods from moderate budgets onward. All methods plateau at budgets of around 1000 simulations, yet none approach the gradient descent optimum.
Figure 2: Prior and posterior predictive final profiles. Colored curves and shaded regions show the predictive medians and central 90% intervals, while black curves denote the LES profiles. SNPE adjusts the predictive medians but retains uncertainty bands broadly similar to the prior. In contrast, NPE-PFN produces substantially narrower intervals, particularly for u and v , indicating a stronger information gain in output space and close tracking of LES profiles.
Figure 3: Prior–posterior comparison in physical parameter space for cϵ1 , cϵ2 , cϵ3+ , and c2 , selected from the strongest aggregate Spearman correlations across methods detailed in Appendix B.2 . Diagonal panels show marginal KDEs and lower panels show bivariate density contours; dashed lines and diamonds denote θGD from the full-profile gradient calibration, which generally lies away from the main posterior concentration. SNPE remains broadly similar to the prior, whereas NPE-PFN produces clearer shifts and more concentrated densities, indicating a stronger posterior update.
Figure 4: Prior predictive final profiles for buoyancy b . Columns show forcing regimes and rows show 6, 24, and 72 hours. At 24 and 72 hours, the central 90% prior predictive interval generally encompasses the LES profiles. At 6 hours, mild departures occur over some depths and forcing regimes.
Figure 5: Prior predictive final profiles for the along-wind velocity u . The central 90% prior predictive interval encompasses the LES profiles over most retained forcing and depth combinations, with localized departures in some cases. Free convection is not retained because it has no imposed wind and therefore no informative u response.
Figure 6: Prior predictive final profiles for the transverse velocity v . Across the retained rotating and wind-forced configurations, the central 90% prior predictive interval encompasses most of the LES profiles. Free convection and the two zero-Coriolis configurations are not retained because their transverse profiles provide negligible information for calibration.
Figure 7: Prior predictive final profiles for the passive tracer pt . The central 90% prior predictive interval generally encompasses a larger portion of the LES profiles at 24 and 72 hours. More visible departures occur during the 6-hour transient, showing that the prior predictive coverage of the tracer varies throughout the simulated evolution.
Figure 8: Blockwise PCA dimensions selected using a 90% explained-variance threshold. Panels denote physical variables, rows denote time horizons, and columns denote forcing regimes; cells report the selected dimension and retained variance, while gray cells indicate masked combinations. Overall, b generally requires a larger PCA basis than pt , supporting an adaptive block-specific truncation.
Figure 9: Empirical covariance structure in the 159-dimensional PCA-summary space: (a) covariance magnitudes, showing nearly zero values for b and the largest values for pt ; (b) standardized within- and cross-variable correlations; and (c) eigenspectra before and after adding the 10−4 diagonal nugget. The regularization lifts the low-variance directions dominated by b while preserving the dominant covariance structure.
Figure 10: Physical-space RMSE between the posterior mean and θGD , reported as mean and SEM across seeds. The dashed line denotes the customary tunax parameters; the starred TSNPE point at B=50 uses a 1%-trimmed mean, and values above 2.2 are clipped. NPE-PFN yields lower RMSE across budgets, while NPE, SNPE, and TSNPE generally fluctuate slightly above the default-reference distance.
Figure 11: Posterior parameter differences relative to θGD for each method’s lowest-profile-RMSE seed at budget 3000 . NPE, SNPE, and TSNPE largely overlap, whereas NPE-PFN shows the clearest shifts for cϵ2 and c3 .
Figure 12: Posterior Spearman correlation matrices in physical parameter space for the selected runs at a simulation budget of 3000 . All panels share the same color scale. Pairwise correlations are generally weak, with the positive association between cε1 and cε2 providing the clearest relationship consistently recovered across methods. These matrices motivate the four-parameter subset examined in the posterior pairplots.
Parameter pair
NPE
SNPE
TSNPE
NPE-PFN
Rij
(cε1,cε2)
0.215
0.105
0.183
0.236
0.185
(cε2,cε3+)
0.063
0.060
0.106
0.040
0.067
(c2,cμ0)
-0.074
-0.091
-0.091
-0.007
0.066
(c2,cε1)
-0.065
-0.087
-0.071
-0.027
0.062
(cε1,cε3+)
0.092
0.057
0.062
0.021
0.058
(σε,cε2)
-0.120
-0.068
-0.035
-0.006
0.057
Appendix
Table 2: Parameter pairs with the largest mean absolute posterior Spearman correlation across methods. Method-specific coefficients retain their signs, while Rij summarizes the average association strength.
Figure 13: Prior and posterior predictive final profiles under strong wind and weak cooling at 6 hours. Colored curves and bands denote predictive medians and central 90% intervals, and black curves denote the LES profiles. SNPE largely retains the prior predictive spread, whereas NPE-PFN produces sharper intervals, particularly for u and v , together with clearer median adjustments for b , v , and pt .
Figure 14: Prior and posterior predictive final profiles under strong wind at 24 hours. NPE-PFN preserves the behavior observed at 6 hours, producing the clearest uncertainty contraction for the velocity profiles and predictive medians generally closer to the LES reference. SNPE improves the median profiles but retains uncertainty bands similar to the prior.
Global ocean models exhibit biases in the mean state and variability, particularly at coarse resolution, where mesoscale eddies are unresolved. To address these biases, parameterization coefficients are typically tuned ad hoc. Here, we formulate parameter tuning as a calibration problem using Ensemble Kalman Inversion (EKI). We optimize parameters of a neural network parameterization of mesoscale eddies in two idealized ocean models at coarse resolution. The calibrated parameterization reduces errors by factors of 1.7-3.3 in the time-averaged fluid interfaces and their variability compared to the unparameterized model, depending on the metric and configuration. The EKI method is robust to noise in time-averaged statistics arising from chaotic ocean dynamics. Furthermore, we propose an efficient calibration protocol that bypasses integration to statistical equilibrium by carefully choosing an initial condition. These results demonstrate that systematic calibration can substantially improve coarse-resolution ocean simulations and provide a practical pathway for reducing biases in global ocean models.
Pavel Perezhogin, Alistair Adcroft, Laure Zanna
Courant Institute School of Mathematics, Computing, and Data Science, New York University, New York, NY, USA · Program in Atmospheric and Oceanic Sciences, Princeton University, Princeton, NJ, USA
Neural Posterior Estimation (NPE) enables rapid parameter inference for complex simulators with intractable likelihoods. NPE trains an inference network to estimate a probability density over parameters given data, typically assumed to be continuous. However, many scientific models involve parameter spaces that are mixed, that is, they contain both discrete and continuous dimensions. We address this limitation by extending NPE to mixed parameter spaces through an inference network that jointly handles discrete and continuous parameters. The inference network factorizes the joint posterior into discrete and continuous components, combining an autoregressive classifier for the discrete parameters with a generative model for the continuous parameters, trained jointly under a single simulation-based objective. In addition, we propose a diagnostic tool to assess the calibration of the mixed posterior approximation. Across tractable toy examples and real-world scientific simulators, our joint inference approach yields accurate and calibrated posteriors. The inference framework is available in the sbi Python package.
Jan Boelts, Cornelius Schröder, Jonas Beck +3
appliedAI Institute for Europe · Machine Learning in Science, University of Tübingen · Tübingen AI Center +3
Accurate calibration of particle-based simulators is crucial for robotic earthwork simulation, but analytical calibration is challenging due to this task's highly nonlinear particle dynamics and the black-box nature of conventional simulators. Although simulation-based inference (SBI) can estimate posterior distributions over simulation parameters solely from forward simulations, applying SBI directly to high-fidelity (HF) particle simulators is often computationally prohibitive. Low-fidelity (LF) simulators with coarser particles can reduce this cost, but changes in particle size and particle count shift the parameter values needed to reproduce the same observation, producing biased LF posteriors. We propose Bridged SBI, which leverages a biased but informative LF posterior to guide HF inference. This method first uses inexpensive LF simulations to identify a coarse high-density parameter region, and then it learns a local residual bridge to transport LF posterior samples toward HF-consistent regions by correcting the LF--HF discrepancy. We analyze how sequential multi-fidelity SBI (Naive-MF) can suffer from LF-induced posterior miscoverage when it directly relies on the LF posterior without discrepancy correction. We then show that Bridged SBI is designed to alleviate this issue by explicitly modeling the LF--HF discrepancy through residual correction. Experiments on both sim-to-sim particle-parameter calibration and real-to-sim calibration with real soil observation show that Bridged SBI produces more accurate and reliable HF posteriors than HF-only SBI or the Naive-MF baseline, especially under limited HF simulation costs.
Gahee Kim, Yuki Kadokawa, Sandro M. Alcantara Tacora +5
NARA INSTITUTE OF SCIENCE AND TECHNOLOGY, Nara, Japan · PUBLIC WORKS RESEARCH INSTITUTE, Ibaraki, Japan