Firedrake - automated, differentiable building blocks for geoscientific simulation

Computer simulations of continuous processes described by partial differential equations are a bedrock of geoscientific simulation. Each simulation is a complex composition of equations, discretisations, solvers and parameterisations. Realistic geoscientific simulation also depends on the integration of observed data either as facing functions or through data assimilation. The result of this complexity is that creating new models, or even extending existing ones, can often be exceptionally resource intensive, even for large and highly capable institutions.

 

Firedrake (https://www.firedrakeproject.org/) offers a revolutionary different approach to model creation. Rather than coding the implementation of a model in low level code in a compiled language, Firedrake users write the mathematical formulation of their model in high-level Python. The high performance, parallel implementation of that code is then automatically generated and executed. Users have access to:

 

• A huge range of finite element discretisations for any PDE they choose, including generalisations of the various variables staggerings that are typically used across the geosciences.
• Programmable, composable solvers and preconditioners, including algebraic and geometric multigrid approaches, and physics-based preconditions based on the characteristics of the system being solved.
• Seemless coupling to external processes, including the ML frameworks JAX and PyTorch.
• Fully automated adjoint computations: the adjoint to a Firedrake simulation is available with no additional coding required.
• Integration with optimisation algorithms for data assimilation.

 

Firedrake already provides the basis for:

• The GUSTO toolkit, used for dynamical core development research at the Met Office and University of Exeter (https://www.firedrakeproject.org/gusto/).
• The Thetis coastal ocean model (Kärna et al. 2018)
• G-ADOPT: The Geoscientific ADjoint Optimisation PlaTform for mantle convection and glacial isostatic adjustment from the Australian National University (Ghelichkhan et al 2024). 

As well as hundreds of bespoke simulations by users around the world.

 

This PICO will present the key features of Firedrake and illustrate the applications to which it is put.

 

References

 

Ghelichkhan, Sia, et al. "Automatic adjoint-based inversion schemes for geodynamics: reconstructing the evolution of Earth's mantle in space and time." Geoscientific Model Development 17.13 (2024): 5057-5086.

Kärnä, Tuomas, et al. "Thetis coastal ocean model: discontinuous Galerkin discretization for the three-dimensional hydrostatic equations." Geoscientific Model Development11.11 (2018): 4359-4382.