SpeedyWeather.jl: A 16-bit weather model with machine learning

We present SpeedyWeather.jl, a global atmospheric model currently developed as a prototype for a 16-bit climate model incorporating machine learning for accuracy and computational efficiency on different hardware. SpeedyWeather.jl is designed for type flexibility with low precision, and automatic differentiation to replace parts of the model with neural networks for a more accurate representation of climate processes and computational efficiency.

Computational resources are a major limitation to improve reliability in numerical predictions of weather and climate. Most simulations run on conventional CPUs in 64-bit floats, although some weather forecast centres now use 32 bits operationally for higher performance. Successful 16-bit simulations have been previously demonstrated with projects like ShallowWaters.jl, increasing performance by 4x with respect to a 64-bit simulation on Fujitsu’s A64FX CPU. However, it remains to be seen whether these results can also be achieved for global atmospheric models, like those used for weather and climate simulation. A new model, SpeedyWeather.jl aims to address this question. As with ShallowWaters.jl, SpeedyWeather.jl aims to support hardware-accelerated low precision arithmetic, yet will be substantially more complex. Much like state-of-the-art numerical weather prediction models, SpeedyWeather.jl includes a “dynamical core” for advancing forward the basic equations describing fluid flow in the Earth’s atmosphere and “parametrizations” for representing physical processes that take place below the scale of the model’s spatial grid, such as the development of clouds from convective updrafts. As such, it is intended to be a simple model for exploring weather and climate simulation in the Julia ecosystem. SpeedyWeather.jl is, like ShallowWaters.jl, fully type-flexible to support arbitrary number formats for performance and analysis (like Sherlogs.jl) simultaneously. This means the model development is precision-agnostic, which allows us to address the common problems of dynamic range and critical precision loss often incurred from using low-precision number formats. The aim of this project is to develop a prototype towards the first global 16-bit weather and climate models.

Beyond numerical weather prediction, low-precision arithmetic is now routinely used in deep learning and neural networks. SpeedyWeather.jl is developed so that entire parts of the model may be replaced by artificial neural networks, thereby complementing conventional physics-based climate modelling with a data-driven approach. Such “hybrid” climate models promise to improve the representation of climate processes that are conventionally poorly resolved, either by training against higher resolution simulations or simulations based on more sophisticated, yet expensive, algorithms. In addition, hybrid models offer the prospect of fitting climate models to observational data. In order to train the neural network components of the model, SpeedyWeather.jl aims to be fully differentiable using automatic differentiation. Implementing parts of weather and climate models with artificial neural networks can also improve computational efficiency and facilitate low precision linear algebra. This talk presents the concept, implementation details, challenges and first results in the development of SpeedyWeather.jl towards a hybrid model incorporating both differential equation solvers and machine learning.

Co-Authors:

• Tom Kimpson (University of Oxford, UK)
• Alistair White and Maximilian Gelbrecht (Potsdam Institute for Climate Impact Research and Technical University of Munich, Germany)
• Sam Hatfield (European Centre for Medium-Range Weather Forecasts, Reading, UK)