In search of long-time stability with neural network surrogates of Fokker-Planck collision operator solvers

Collisional transport is a critical physical process in the edge of tokamak plasmas, yet for kinetic simulations with multiple ion species numerically solving for collisional transport are expensive. Previous work [1] showed the promise of using machine learning to train a surrogate for the numerical Fokker-Planck-Landau collision solver used in the edge turbulence code XGC. An encoder-decoder neural network was trained on a large dataset from XGC, with soft-constraints in the optimization loss to enforce conservation properties. This resulted in a surrogate which for single-time predictions on a test set gave acceptable conservation errors (1e-4). However, when using the ML model in the full XGC simulation over multiple timesteps, instabilities invariably occurred. Here we probe the ability of neural networks to produce stable, asymptotically correct solutions to a simpler collisional system, that of a bi-Maxwellian distribution function collisionally relaxing to an isotropic Maxwellian. We show the benefits of new operator learning models, such as the Fourier Neural Operator (FNO), in generalizing. We also show the benefits of imposing hard constraints to enforce conservation properties, and explore more directly the use of entropy functions to further constrain solutions.