Data-driven preconditioning for parallel-in-time PDE solvers

Project Description

In our group we are developing solvers for parallel-in-time simulation based on a technique called ParaDiag. The goal of parallel-in-time methods is to solve time-dependent PDEs using a parallel computer with computation across several time steps at once in addition to the usual parallelism across spatial subdomains, to yield additional parallelism so that we can get the solution to bigger problems more quickly. ParaDiag is a computational linear algebra trick for solving the implicit system coupling together several time steps at once. For problems with waves and/or transport (such as those arising in weather and climate, or fusion plasma simulation) the big challenge is that ParaDiag requires the solution of systems that look like backward Euler applied to the original system but with a complex valued timestep! This complex valued aspect causes trouble for standard iterative solvers, and here we plan to use an iteration consisting of solving the linear system restricted independently on each of a set of overlapping patches, inspired by Graham et al (2020) which solved similar problems in the context of frequency domain wave equations. From that work it is evident that the patch systems must have an (approximate) absorbing boundary condition for the solver to converge well. For nondispersive wave equations and pure advection problems it is clear what these boundary conditions should be, but for nonlinear problems the theory is incomplete. In this project we propose to use a data-driven approach, taking solutions of the backward Euler problem, and restricting them to patches as input data, before fitting parameters using an appropriate model (such as a neural network, random features model, etc.).

• Graham, Ivan G., Euan A. Spence, and Jun Zou. “Domain decomposition with local impedance conditions for the Helmholtz equation with absorption.” SIAM Journal on Numerical Analysis 58, no. 5 (2020): 2515-2543.

Existing background work

The history of the ParaDiag framework is surveyed in Gander et al (2020); we are pursuing the ParaDiag II approach as it is described in that article, particularly focussing on achieving practical performance for nonlinear PDEs. Funded by the UK ExCALIBUR project led by UKRI, the Met Office, UKAEA and STFC, we have developed a software library, asQ (Hope-Collins et al (2024)) for investigating ParaDiag algorithms and benchmarking them on high performance supercomputers. This has involved specific test cases relevant to weather and climate, fusion simulation and computational geology, but the scope is broad. asQ is built using Firedrake (firedrakeproject.org) which has a PyTorch interface for seamless integration of neural networks (Bouziani, Ham and Farsi (2024)).

• Gander, Martin J., Jun Liu, Shu-Lin Wu, Xiaoqiang Yue, and Tao Zhou. “Paradiag: Parallel-in-time algorithms based on the diagonalization technique.” arXiv preprint arXiv:2005.09158 (2020).
• Hope-Collins, Joshua, Abdalaziz Hamdan, Werner Bauer, Lawrence Mitchell, and Colin Cotter. “asQ: parallel-in-time finite element simulations using ParaDiag for geoscientific models and beyond.” arXiv preprint arXiv:2409.18792 (2024).
• Bouziani, Nacime, David A. Ham, and Ado Farsi. “Differentiable programming across the PDE and Machine Learning barrier.” arXiv preprint arXiv:2409.06085 (2024).

Main objectives of the project

The objectives are: * To develop a data-driven approach to absorbing boundary conditions for ParaDiag problems * To build up the capability through dispersive wave equations, coupled wave-transport problems, and the problems arising from ParaDiag applied to specific nonlinear PDEs

Details of Software/Data Deliverables

• implemented preconditioners for asQ using petsc4py integrated with Firedrake
• tuned parameters for example problems
• everything developed open source on Github