Case study of a differentiable heterogeneous multiphysics solver for a nuclear fusion application

TL;DR

This paper demonstrates a differentiable multiphysics solver framework for nuclear fusion applications, integrating high-fidelity plasma simulations with gradient-based methods to enable efficient, end-to-end differentiability and rapid prototyping.

Contribution

It introduces Tesseract, a software layer that makes complex plasma solvers compatible with JAX, facilitating seamless integration and differentiation in multiphysics workflows.

Findings

End-to-end differentiability achieved with high-fidelity plasma solvers.

Efficient gradient-based root-finding for plasma impedance.

Flexible interchangeability between solvers and surrogates.

Abstract

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.