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

Jack Coughlin, Jonathan Brodrick, Archis Joglekar, Alexander Lavin

Keywords: nuclear fusion, differentiable programming, computational physics, tesseract

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 ab-
straction 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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 47