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

arXiv: https://arxiv.org/abs/2511.13262 · HTML: https://mlsanity.com/paper/2511.13262/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 16 references that are in the map.
- First posted on arXiv 17 Nov 2025.

## Builds on (most cited first, 6 of 6)

- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [Simulation Intelligence: Towards a New Generation of Scientific Methods](https://mlsanity.com/paper/2112.03235/) (2021)
- [JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science](https://mlsanity.com/paper/2212.00964/) (2022)
- [Optimistix: modular optimisation in JAX and Equinox](https://mlsanity.com/paper/2402.09983/) (2024)
- [JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows](https://mlsanity.com/paper/2402.05193/) (2024)
- [PICT -- A Differentiable, GPU-Accelerated Multi-Block PISO Solver for Simulation-Coupled Learning Tasks in Fluid Dynamics](https://mlsanity.com/paper/2505.16992/) (2025)

## Abstract (arXiv)

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.