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

G. H. Zheng, S. F. Liu+15 more|arXiv (Cornell University)|Jan 18, 2026

Magnetic confinement fusion research被引用 0

一句话总结

本论文介绍了 PTEFIT，一种基于 PyTorch/TensorRT 的 GPU 加速实时 Grad–Shafranov 平衡重建算法，用于托克马克，实现亚毫秒级每片切片性能并在 EXL-50U 球形托卡马克上实现实时反馈集成。

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 imes 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.

研究动机与目标

• 在不牺牲跨配置普适性的前提下，推动实时、高精度的托卡马克平衡重建。
• 利用机器学习框架来优化数值算法，同时保持基于物理的理论基础。
• 提供一个模块化、GPU 加速的重建流程，兼容 EFIT 原则。

提出的方法

• 在 PyTorch 中构建 PTEFIT，包含 21 个模块化计算单元，处理格林函数生成、电流密度参数化和 Grad–Shafranov (G-S) 求解。
• 使用格林函数矩阵将圆柱电流与圆极场和通量联系起来，利用 GPU 加速的矩阵运算来计算。
• 用归一化通量的多项式对 J_phi 进行参数化，并使用正规方程或 QR 分解（Gram-Schmidt）在无 SVD 的情况下求解最小二乘系统。
• 将 G-S 方程转化为一组并行的三对角问题，并为快速 GPU 求解进行预计算。
• 用二次多项式拟合处理通量面，以定位 O 点和 X 点，从而实现通量归一化的轴向/边界识别的高精度。
• 通过高效的 GPU 可处理积分和基于二分搜索的通量面顶点搜索来计算如安全因子等的通量面平均量。
• 通过将数据预计算、最小化 CPU-GPU 数据传输并以 TensorRT 的 CUDA 图实现，优化为可在实时性要求下运行。

Figure 1 : Schematic diagram of the bisection method for finding flux surface vertices. The flux outside the plasmas region need to be filtered first. For each thread, the initial endpoints are located at the magnetic axis and the reconstructed boundary, respectively, and the points on the reconstru

实验结果

研究问题

• RQ1 一个基于物理的平衡重建是否可以完全在机器学习框架内实现，以在不牺牲 EFIT 一致性的前提下达到实时性能？
• RQ2 基于 PyTorch/TensorRT 的方法在速度和准确性上与离线 EFIT 及传统的 GPU 加速重建在不同通量面配置下的对比表现如何？
• RQ3 在托卡马克中使用 PTEFIT 进行反馈控制时，能够达到的实时控制性能（例如 R_max 与等通量）的表现如何？
• RQ4 该方法是否能够适应不同的托卡马克配置和诊断，同时保持物理可解释性？
• RQ5 在此背景下求解最小二乘步骤时，正规方程与 QR 的计算权衡是什么？

主要发现

• PTEFIT 在 129×129 分辨率下的时间片平均推理时间为 0.268 ms。
• PTEFIT 大约比离线 EFIT 快 10,000×，且大约比传统的 GPU 加速重建快 2–10×。
• PTEFIT 在 LCFS 重建中与离线 EFIT 展现出较高的一致性，在 XPT 配置中的偏差约为 1 cm。
• PTEFIT 能在 ICRF 加热实验中实现最大等离子体半径 R_max 的实时 PID 反馈控制。
• PTEFIT 能以 MIMO 反馈方式实现等通量控制，保持通量差异接近参考点，尽管 VDE 限制了一个 X 点的可见性。
• PTEFIT 的设计在保持与 EFIT 原则的一致性的同时，提供了较低的开发成本和更广的可配置性。

Figure 2 : Comparison of LCFS results obtained by PTEFIT and offline EFIT for XPT configuration reconstruction of shot #13899. Here, $\Delta r_{\min}$ and $\Delta r_{\max}$ are the deviations between the two at the minimum and maximum radial positions of the LCFS, respectively.