A Novel Numerical Algorithms Optimization Method with Machine Learning Frameworks: Application on Real-time Plasmas Equilibrium Reconstruction in EXL-50U Spherical Torus

TL;DR

This paper introduces a novel machine learning-based optimization method for real-time plasma reconstruction in tokamaks, utilizing PyTorch and TensorRT for enhanced performance, flexibility, and usability.

Contribution

It presents a new GPU-accelerated algorithm, PTEFIT, that integrates machine learning frameworks for real-time plasma equilibrium reconstruction in fusion devices.

Findings

Achieved 0.268ms inference time per slice

Successfully controlled plasma position and flux in real-time

Demonstrated potential for accelerating numerical algorithm development

Abstract

This work proposes for the first time a novel optimization method for numerical algorithms, which takes advantages of machine learning frameworks PyTorch and TensorRT, leveraging their modularity, low development threshold, and automatic tuning characteristics to achieve a real-time plasmas reconstruction algorithm called PTEFIT as an application in tokamak-based controlled fusion that combines performance, flexibility, and usability. The algorithm has been deployed and routinely operated on the EXL-50U spherical tokamak, with an average inference time of only 0.268ms per time slice at $129\times 129$ resolution, and has successfully driven feedback control of the maximum radial position of plasmas and isoflux control. We believe that its design philosophy has sufficient potential to accelerate development and optimization in GPU parallel computing, and is expected to be extended to other numerical algorithms.