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.
The deep learning revolution has borne many fruits, but for scientific computing one of the main benefits has been the development of software frameworks for writing automatically differentiable (AD) numerical programs. End-to-end differentiable solvers have been created for fluid dynamics [1,6], finite element analysis [17], molecular dynamics [12], and plasma transport in tokamaks [3], to name just a few examples. Differentiable simulations unlock sensitivity analysis and gradient-based optimization workflows end-to-end through large PDE solutions.
However, many physical systems of interest in engineering applications are “multiphysics” in nature. By a multiphysics system we mean a system involving a coupling between two or more physical domains or scales, each governed by its own set of ordinary or partial differential equations. The coupling between equation is often constrained by an implicit relation that requires the solution of a nonlinear system. Moreover, spatial and temporal scales may vary widely, necessitating the use of implicit time stepping methods or convergence to a steady state solution on the fast timescale. Examples of multiphysics systems include fluid-structure interactions, combusting and reacting flows, and radiation hydrodynamics [10].
Applying AD to multiphysics problems presents a particular challenge because their complexity may make it infeasible to rewrite all physics solvers in an AD-native framework. In this work we present one such example. One component of our system is an ODE for which it is quite simple to write a new solver in JAX. However, this is coupled to a high-dimensional PDE for which an excellent, high-accuracy and efficient solver is available in C. To leverage this existing battle-tested code we use 1st Workshop on Differentiable Systems and Scientific Machine Learning @ EurIPS 2025 arXiv:2511.13262v1 [physics.comp-ph] 17 Nov 2025
Tesseracts [8] to provide an abstraction layer at the boundary between solvers. Crucially, as Gkeyll itself is non-differentiable, Jacobian-vector products and vector-Jacobian products of the C solver are implemented using finite differences, and the resulting differentiable function is incorporated into the outer JAX program using tesseract_jax1 . Furthermore, Tesseracts enable a modular software architecture, enabling us to swap out different implementations for the PDE component, such as neural surrogates or analytic approximations. End-to-end differentiability of the program is maintained regardless of the underlying solver.
2 Fusion Z Pinch Compression Modeling
The Sheared Flow-Stabilized Z Pinch [13] is a promising approach to magnetically confined nuclear fusion (MCF). It has the benefits of a compact size, requires no external magnetic field coils, and relatively simple engineering compared to other MCF concepts. The Z Pinch functions on the basis of the plasma pinch principle: a plasma carrying a current with density vector j produces a magnetic field B and corresponding force vector F = j × B which acts to compress the plasma in on itself. By ramping up the plasma current, one can increase F and thus the equal and opposite pressure gradient force -∇p, driving the plasma density and temperature up to fusion conditions.
The pulsed power driver of the Z Pinch device can be modeled as a series RLC circuit as done in Ref. [5]. Let Q and I denote the time-dependent capacitor charge and total circuit current respectively. The voltage balance law around the circuit, including all three RLC circuit components and the plasma load, is
where L is the circuit inductance, R the circuit resistance and C the capacitance. The plasma load impedance is characterized by the plasma inductance L p and the resistive plasma voltage V Rp . The plasma inductance is a geometric property of the pinch profile and is derived in Ref. [4]. The resistive plasma voltage V Rp is the voltage required to drive a current I through the plasma at steady state, and is a function of the current I and the plasma temperature T and density n. Temperature T is related to current I, linear density N , and charge state Z by the well-known Bennett relation for pinches.
The relation between T and n is given by the specific entropy s = ln(T /n γ-1 ), where γ = 5/3 is the adiabatic constant for hydrogenic plasmas. The complete ODE system for the plasma and circuit state is
The terms P η and P Br are the contributions from Ohmic heating of the plasma and bremsstrahlung radiative cooling, a key energy loss mechanism in fusion plasmas. Expressions for them may be found in [7].
The key remaining challenge is to specify a closure for V Rp (I, T, n). While the circuit dynamics can be described by the lumped element model eq. ( 2), the plasma impedance depends closely on the distribution of current-carrying plasma particles near the electrodes. Plasma-electrode boundaries form a structure known as a plasma sheath, in which the mobility difference between light electrons a
This content is AI-processed based on open access ArXiv data.