Neural-Parareal: Self-improving acceleration of fusion MHD simulations using time-parallelisation and neural operators

Abstract

The fusion research facility ITER is currently being assembled to demonstrate that fusion can be used for industrial energy production, while several other programmes across the world are also moving forward, such as EU-DEMO, CFETR, SPARC and STEP. The high engineering complexity of a tokamak makes it an extremely challenging device to optimise, and test-based optimisation would be too slow and too costly. Instead, digital design and optimisation must be favoured, which requires strongly-coupled suites of multi-physics, multi-scale High-Performance Computing calculations. Safety regulation, uncertainty quantification, and optimisation of fusion digital twins is undoubtedly an Exascale grand challenge. In this context, having surrogate models to provide quick estimates with uncertainty quantification is essential to explore and optimise new design options. But there lies the dilemma: accurate surrogate training first requires simulation data. Extensive work has explored solver-in-the-loop solutions to maximise the training of such surrogates. Likewise, innovative methods have been proposed to accelerate conventional HPC solvers using surrogates. Here, a novel approach is designed to do both. By bootstrapping neural operators and HPC methods together, a self-improving framework is achieved. As more simulations are being run within the framework, the surrogate improves, while the HPC simulations get accelerated. This idea is demonstrated on fusion-relevant MHD simulations, where Fourier Neural Operator based surrogates are used to create neural coarse-solver for the Parareal (time-parallelisation) method. Parareal is particularly relevant for large HPC simulations where conventional spatial parallelisation has saturated, and the temporal dimension is thus parallelised as well. This Neural-Parareal framework is a step towards exploiting the convergence of HPC and AI, where scientists and engineers can benefit from automated, self-improving, ever faster simulations. All data/codes developed here are made available to the community.