Bayesian Optimization With GPU Acceleration for Ocean Models

• Published in: Journal of geophysical research. Machine learning and computation Vol. 2; no. 3
• Main Authors: Mrozowska, Marta Agnieszka, Avery, James, Stoustrup, Aster, Nuterman, Roman, Johnsen, Carl‐Johannes, Thormann, Asger, Jochum, Markus
• Format: Journal Article
• Language: English
• Published: 01.09.2025
• ISSN: 2993-5210; 2993-5210
• DOI: 10.1029/2024JH000517

Ocean general circulation models (OGCMs) contain numerous parameterizations of subgrid scale processes. The parameter tuning procedure is rarely reported and often done by hand. We present an automated alternative: VerOpt, a Python package for the ocean model Veros, adapts Bayesian optimization to climate model tuning. We use VerOpt to identify a set of model parameter values that minimize mixed layer depth (MLD) bias in Veros. We present the results of three optimization procedures: TWIN, OBS_TKE, and OBS_9D. The goal is to minimize the MLD error relative to a target. In TWIN, the target is MLD simulated using Veros with a known vertical mixing parameterization. VerOpt identifies a set of parameter values that simulate the target MLD with 99% accuracy. In OBS_TKE and OBS_9D, the target is MLD climatology. In the OBS_TKE optimization, only the vertical turbulence closure parameters are calibrated, which is insufficient to reduce MLD biases in Veros. In the OBS_9D experiment, the mesoscale eddy closure and the scaling of the wind stress and salinity forcing masks are tuned in addition to the vertical mixing scheme parameters. MLD biases in Veros are reduced by up to 23% with the largest impact in the Southern Ocean. In the ensemble of the simulations for which MLD biases are minimized, the mesoscale turbulence diffusion coefficients and are halved compared to the control run, and the zonal wind stress is 10% weaker. Numerous schemes in ocean models approximate the impact of unresolved processes on the mean state of the system. It is not always possible to determine the values of the model parameters based on empirical arguments. Often, the schemes are instead calibrated by hand to reproduce the most realistic ocean state. We present an alternative which automates the tuning process. We use Bayesian optimization (BO) to determine which parameter values minimize the mixed layer depth (MLD) biases in the ocean model Veros. The versatile optimizer (VerOpt) adapts BO to climate modeling problems. We show the results of three optimization sequences: TWIN, OBS_TKE, and OBS_9D. In the TWIN experiment, the optimizer is tasked with finding a set of parameters which reproduce MLD simulated with Veros. VerOpt identifies a range of parameter values which reproduce the target MLD up to 1% accuracy. In the OBS_TKE experiment, the target is the observed MLD. We find that tuning the vertical mixing scheme parameters in isolation does not yield a reduction in MLD biases. We achieve this instead by including the horizontal mixing scheme and atmospheric forcing parameters in the OBS_9D optimization, where the model biases are reduced by up to 23%. Bayesian optimization is used to reduce mixed layer depth biases in a primitive equation 1 ocean model The JAX library allows us to use GPUs and reduce energy consumption by more than 90% The tool can be easily used for optimization of other models