Automating the construction of time-parallel 4DVar data assimilation systems for finite element models.

Data assimilation (DA) has made a significant contribution to the increase in forecast skill in recent decades by leveraging real-world observations to improve the states used in the forecasting system. However, the DA stage comprises a large part of the computational work of the system. Therefore, improving DA efficiency is crucial for continued improvements in forecast skill, which in turn requires high productivity software to be available for DA researchers.

 

4D variational assimilation (4DVar) is a DA method in common use operationally. 4DVar optimises an objective function (agreement with observations and prior forecasts) by updating a control (initial conditions) using the adjoint method. This requires running the linearised system forwards in time, and the system adjoint backwards in time, at every optimisation step. Weak constraint 4DVar is a modification that allows an "inexact model" and splitting the time-series into chunks. This enables reformulating each optimisation step as a saddle point problem, where the forward and adjoint models on each time-chunk can be run in parallel. This time-parallelism can potentially greatly decrease the time required for the optimisation.

 

The adjoint method is highly effective but requires differentiating every operation in the system. Manually differentiation requires high developer effort for new system components, while automatic numerical differentiation is often computationally inefficient. This is particularly problematic in the development stage, where researchers want to run a variety of methods on a variety of equations. This causes a gap between simple equations (e.g. Lorenz and heat) often used in research, and the fluids models in operational systems.

Symbolic differentiation aims to achieve the efficiency of manual differentiation with the automation of numerical differentiation, and provides a route to closing this gap.

 

We present a library for constructing the 4DVar system using symbolic automatic differentiation. This is achieved using: Firedrake, a finite element library which provides symbolic differentiation using UFL, a high level DSL; and pyadjoint, a library for symbolically recording code execution and automatically constructing the forward and adjoint models.

To construct the 4DVar system, the user need only run the forward model (i.e. dynamical core) and the observation operators once. Firedrake and pyadjoint record this run and calculate the linearised forward model and the adjoint model, from which all components of 4DVar can be constructed.

The implementation is both space and time parallel, enabling the real performance of the methods to be tested on HPC machines. Because pyadjoint is model agnostic, this library can be used for any equation simply by changing the finite element model. Minimising the code modification required for different models is key for improving researcher productivity. This not only allows a smoother transition up the model hierarchy, but also allows straightforward testing of different dynamical core discretisations. We will demonstrate the software API and discuss the design choices taken, and present preliminary results on different equations.