Case study of a differentiable heterogeneous multiphysics solver for a nuclear fusion application
This addresses the problem of efficient and differentiable multiphysics modeling for nuclear fusion researchers, though it appears incremental as it builds on existing tools like JAX and Gkeyll.
This work tackled the challenge of integrating non-differentiable, production-grade plasma solvers into gradient-based workflows for nuclear fusion applications, resulting in the development of the 'Tesseract' software that enables end-to-end differentiability and seamless interchange between high-fidelity solvers, neural surrogates, and analytical approximations.
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.