The Challenge

Bayesian calibration is a powerful statistical method for tuning the parameters of complex computer models, but it is notoriously slow and computationally expensive. This has limited its application for the next generation of high-resolution climate and weather models.

My Contribution

To address this challenge, I developed KOH-GPJax, an open-source Python package that leverages the power of JAX and GPU acceleration to dramatically speed up the calibration process.

Key features of KOH-GPJax:

• High-performance: By using JAX for automatic differentiation and GPU computation, KOH-GPJax can be orders of magnitude faster than traditional implementations.
• Scalable: The package is designed to handle the large datasets and complex models used in modern scientific research.
• User-friendly: KOH-GPJax provides a simple and intuitive API, making it easier for researchers to apply Bayesian calibration to their own models.

Impact

KOH-GPJax has been successfully applied to a number of real-world projects, including:

• Weather and climate model calibration: In partnership with the UK Met Office, we have used KOH-GPJax to calibrate key parameters in their weather and climate models, leading to improved forecasts.
• Collaboration with Stanford University: I have also collaborated with researchers at Stanford to apply KOH-GPJax to their own climate modeling problems.

This project demonstrates my ability to develop high-quality scientific software that solves real-world problems.

I presented this topic as part of the UCL Department of Statistical Science PhD Seminar Series back in February 2024. A recording of this presentation from YouTube is embedded below.