Encoder–decoder neural network for solving the nonlinear Fokker–Planck–Landau collision operator in XGC

Abstract

An encoder–decoder neural network has been used to examine the possibility for acceleration of a partial integro-differential equation, the Fokker–Planck–Landau collision operator. This is part of the governing equation in the massively parallel particle-in-cell code XGC, which is used to study turbulence in fusion energy devices. The neural network emphasizes physics-inspired learning, where it is taught to respect physical conservation constraints of the collision operator by including them in the training loss, along with the $\ell _2$ loss. In particular, network architectures used for the computer vision task of semantic segmentation have been used for training. A penalization method is used to enforce the ‘soft’ constraints of the system and integrate error in the conservation properties into the loss function. During training, quantities representing the particle density, momentum and energy for all species of the system are calculated at each configuration vertex, mirroring the procedure in XGC. This simple training has produced a median relative loss, across configuration space, of the order of $10^{-4}$ , which is low enough if the error is of random nature, but not if it is of drift nature in time steps. The run time for the current Picard iterative solver of the operator is $O(n^2)$ , where $n$ is the number of plasma species. As the XGC1 code begins to attack problems including a larger number of species, the collision operator will become expensive computationally, making the neural network solver even more important, especially since its training only scales as $O(n)$ . A wide enough range of collisionality has been considered in the training data to ensure the full domain of collision physics is captured. An advanced technique to decrease the losses further will be subject of a subsequent report. Eventual work will include expansion of the network to include multiple plasma species.

2022 Review of Data-Driven Plasma Science

“…As an example, the Fokker-Planck-Landau collision operator has a computational cost that grows at a quadratic rate as the number of species increases and needs to be evaluated many times when forming the right-hand side of the system of equations. Machine learning surrogates have been successfully developed for this operator [44]. In other settings, data-driven machine learning models have been developed to estimate closures for plasma fluid models [45,46] and fluid turbulence models [47], and for coupled simulations.…”

“…Data driven AI/ML can accelerate the extreme scale simulations by replacing some compute intensive kernels with AI/ML inference routines. Fokker-Planck collision operation is an example [44]. Preconditioner and PDE solvers can be good candidates.…”

“…The machine learning moment closures has been used for accurate and efficient simulation of polydisperse evaporating sprays [532,533], for the radiative transfer equation [534], and for the moment system of the Boltzmann equation [535]. In [44], the authors pursue encoder-decoder neural network for solving the nonlinear Fokker-Planck-Landau collision operator in XGC. In [536], the surrogate models have been trained for integrated simulations for the calculation of the core turbulent transport fluxes and the pedestal structure.…”

Towards fast and accurate predictions of radio frequency power deposition and current profile via data-driven modelling: applications to lower hybrid current drive

“…2018) and also employed ML techniques to reduce computation time for codes such as XGC (Miller et al. 2021). The work of Boyer and Morosohk, in which they create a fast neural-network=based surrogate model for NUBEAM, is particularly relevant to the work presented in this paper.…”

Error-Bounded Learned Scientific Data Compression with Preservation of Derived Quantities

“…There are approaches for incorporating QoI using constrained optimization [25,26] that are closely related to our work. While augmented Lagrangians have been used to incorporate physical constraints into neural networks [27,28], some methods merely use soft penalties which violate QoI error tolerances [29,30]. These simulation approaches have not yet been used for compression, which is the focus of our present work.…”