NTVTOK-ML: Fast surrogate model for neoclassical toroidal viscosity torque calculation in tokamaks based on machine learning methods

The Neoclassical Toroidal Viscosity (NTV) torque is a crucial source of toroidal momentum in tokamaks, exerting significant influence on plasma instability and performance. Accurate numerical modeling of NTV torque is essential for experimental design and operation, as well as for gaining insight into the relevant physical processes. However, the time-consuming nature of NTV torque calculation poses challenges for its practical application in experiment analysis and physical investigations. In this study, we have developed NTVTOK-ML, a surrogate model for NTV torque calculation that combines the expressive power and fast inference of machine learning methods to achieve simultaneous accuracy and time efficiency. To obtain datasets for NTV torque, extensive numerical calculations using NTVTOK and MARS-F codes were performed under various plasma conditions of Experimental Advanced Superconducting Tokamak (EAST), covering a wide range of experimentally relevant parameter regimes and incorporating rich physical effects such as pitch angle scattering, full toroidal geometry, resonances, etc. For fixed magnetic perturbation case, NTVTOK-ML employs Multi-Layer Perceptron (MLP) deep neural network and eXtreme Gradient Boosting (XGBoost) ensemble learning techniques respectively. Furthermore, when considering linear plasma response effect, Convolutional Neural Network (CNN) is utilized to process two-dimensional magnetic perturbation data. The prediction accuracy of NTVTOK-ML is evaluated based on statistical metrics including coefficient of determination ($R^2$), mean squared error (MSE), and relative error; single sample prediction ability; and generalization ability - demonstrating its reliability in NTV torque prediction tasks. Importantly, the computational time required for predicting NTV torque using our proposed approach is significantly reduced compared to the original numerical code by several orders of magnitude. Additionally, the flexibility offered by the NTVTOK-ML framework allows users to optimize model performance under specific circumstances. Overall, our developed method provides an accessible solution for rapid yet accurate prediction of NTV torque while incorporating essential physical effects - thereby facilitating real-time or inter-shot analysis in experiments as well as comprehensive multi-scale nonlinear time evolution modeling.

Program summary

Program Title: NTVTOK-ML

CPC Library link to program files: https://doi.org/10.17632/thcd9fbjd5.1

Licensing provisions: Apache-2.0

Programming language: Python

Nature of problem: The traditional numerical calculation of NTV torque in tokamaks is time-consuming, which hinders real-time or inter-shot experimental analysis and multi-scale nonlinear time evolution modeling. Simplifications to physical models have been employed to accelerate the calculations; however, they often result in quantitative or qualitative deviations within certain parameter regimes. Therefore, achieving both accuracy and high time efficiency simultaneously for NTV modeling is essential for experiment design and operation, as well as a comprehensive understanding of relevant physical processes.

Solution method: 1. Generate a comprehensive plasma parameter space that encompasses a wide range of experimentally relevant tokamak plasma conditions; 2. Conduct numerous NTVTOK and MARS-F calculations to construct NTV torque datasets that incorporate various physical effects, such as pitch angle scattering, full toroidal geometry, and resonances; 3. Pre-process the datasets using dimensionless methods and logarithm transformation. For two-dimensional magnetic perturbation data, employ CNN for pre-processing; 4. Develop and train machine learning models based on deep neural networks or ensemble learning methods; 5. Evaluate the model performance in terms of statistical metrics, single sample prediction capability, generalization ability, and computational efficiency. The results indicate that the NTVTOK-ML surrogate model can be applied to diverse physical tasks in future studies.