将可微分的等离子体求解器集成到核聚变仿真中,实现参数优化与快速迭代。
Case study of a differentiable heterogeneous multiphysics solver for a nuclear fusion application
- 用JAX构建可微分常微分方程求解器,驱动脉冲功率电路与等离子体演化
- 通过梯度下降牛顿法高效求解每步的等离子体阻抗,保持全程可微
- 借助Tesseract框架兼容高保真求解器与神经网络代理模型,支持快速原型
本工作展示了一个面向核聚变领域的异构多物理场求解器的案例研究。在宏观尺度上,基于JAX的可微分常微分方程求解器用于计算压缩型Z箍缩中的脉冲功率电路与整体等离子体参数演化。该求解器需要在每个时间步通过根查找获得等离子体负载的阻抗,我们采用基于梯度的牛顿迭代法高效求解。然而,将非可微的生产级等离子体求解器(如Gkeyll,一个C/CUDA等离子体模拟套件)集成到基于梯度的工作流中存在挑战。'Tesseract'软件通过提供一个完全兼容JAX(通过`tesseract_jax`适配器)的多物理场可微抽象层,解决了这一问题。该架构确保了端到端可微性,同时允许在高保真求解器(Gkeyll)、神经网络代理模型和解析近似之间无缝切换,支持快速、渐进式原型开发。
原文摘要 · Abstract (English)
This work presents a case study of a heterogeneous multiphysics solver from the nuclear fusion domain. At the macroscopic scale, an auto-differentiable ODE solver in JAX computes the evolution of the pulsed power circuit and bulk plasma parameters for a compressing Z Pinch. The ODE solver requires a closure for the impedance of the plasma load obtained via root-finding at every timestep, which we solve efficiently using gradient-based Newton iteration. However, incorporating non-differentiable production-grade plasma solvers like Gkeyll (a C/CUDA plasma simulation suite) into a gradient-based workflow is non-trivial. The ''Tesseract'' software addresses this challenge by providing a multi-physics differentiable abstraction layer made fully compatible with JAX (through the `tesseract_jax` adapter). This architecture ensures end-to-end differentiability while allowing seamless interchange between high-fidelity solvers (Gkeyll), neural surrogates, and analytical approximations for rapid, progressive prototyping.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。