Efficient dataset construction using active learning and uncertainty-aware neural networks for plasma turbulent transport surrogate models

This work demonstrates a proof-of-principle for using uncertainty-aware architectures, in combination with active learning techniques and an in-the-loop physics simulation code as a data labeler, to construct efficient datasets for data-driven surrogate model generation. This was applied to the tokamak plasma turbulent transport problem. Specifically, the QuaLiKiz quasilinear electrostatic gyrokinetic turbulent transport code was chosen as the base simulator, building off of a previous proof-of-principle that successfully demonstrating training set reduction on static pre-labeled datasets using the ADEPT framework. While QuaLiKiz provides relatively fast evaluations, this study specifically targeted small datasets to serve as a proxy for more expensive codes, such as CGYRO or GENE. The newly implemented algorithm uses the spectral-normalized Gaussian process architecture for the classification component of the problem and the Bayesian neural network with noise contrastive prior architecture for the regression component, training models for all turbulent modes (ion temperature gradient, trapped electron mode, and electron temperature gradient) and all transport fluxes (Qe, Qi, Γe, Γi, and Πi) described by the general QuaLiKiz output. With 45 active learning iterations, moving from a small initial training set of 102 to a final set of 104, the resulting models reached a F1 classification performance of ∼ 0.8 and a R2 regression performance of ∼ 0.75 on an independent test set across all outputs. This extrapolates to reaching the same performance and efficiency as the previous pipeline, although with 1 extra input dimension. While the improvement rate achieved in this implementation diminishes faster than expected, the overall technique is formulated with components that can be upgraded and generalized to many surrogate modeling applications beyond plasma turbulent transport predictions.