Accelerating UKCA using data-driven timestep adaptation
Main supervisor: Joe Wallwork, ICCS. Co-supervisor: Luke Abraham, Department of Chemistry/NCAS
Project description
The chemistry solver within the United Kingdom Chemistry and Aerosols (UKCA) model – called ASAD – has been established as its most computational intensive component. As such, it is a key target for computational optimisation. Reducing its cost benefits not just UKCA's performance but also that of coupled models that make use of UKCA, such as the Met Office's weather forecasting suites and the UK Earth System Model.
ASAD uses the implicit Backward Euler timestepping scheme, applying a Newton solver at each timestep. It supports adaptive timestepping in the sense that a different timestep lengths can be used at different points in time and space. However, the model currently lacks an intelligent way to select the timestep length. A large default timestep is selected, which is sufficient in most cases. In cases where the chemistry is more complex, the default timestep is insufficient to achieve convergence, so the timestep length is repeatedly halved until convergence is attained.
In this project, we propose to predict the required number of timestep halvings in advance, using information available to the solver at the beginning of the timestep. In doing so, the redundant calculation involved with attempting to solve with larger timestep lengths will be avoided. A PyTorch-based machine learning (ML) architecture will be designed to achieve the mapping and tested 'offline' on data from a wide variety of model runs.
Sensitivity analysis experiments will be conducted using the ML approach, assessing the dependencies of the number of halving steps on the various input variables. If sufficient progress is made, we will also deploy the ML approach will be deployed and tested 'online' by coupling it into UKCA using FTorch.
The main outputs of the project will take the form of open source software contributions. There will also be the opportunity to contribute towards a research paper on the proposed approach.
Background information
UKCA is a community atmospheric chemistry and aerosol micro-physics model developed in partnership between the Met Office, University of Cambridge, University of Leeds and others. It is written in Fortran and has benefited from several performance-oriented development efforts (involving both supervisors) in recent years. However, the redundant calculation in the timestepping scheme remains a computational bottleneck stopping greater performance enhancements being unlocked. Co-supervisor Luke Abraham is the main code owner of UKCA.
The project will build on an extensive Python-based data analysis framework designed by Joe Wallwork. A core part of the project will involve contributions to the open source mlstep package for predicting the number of halving steps using machine learning.
Working environment
The student will be included in weekly development meetings for FTorch and can also meet regularly with research software engineers in the group for additional guidance. They can work remotely for part of the time.
The student will have the option of attending the ICCS Summer School.
Essential knowledge, skills and attributes
- Basic knowledge of numerical methods for solving ODEs (i.e. timestepping schemes)
- Basic knowledge of machine learning techniques (preferably neural networks)
- Python programming
- Basic knowledge of PyTorch and/or Fortran would be useful but not essential.
Contact us
If you have questions about this project, email Joe Wallwork