Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Jack B. Coughlin, Archis Joglekar, Jonathan Brodrick, Alexander Lavin

TL;DR

The paper tackles the challenge of building end-to-end differentiable simulations for a coupled multiphysics system in nuclear fusion, specifically a Z Pinch scenario with a macroscale pulsed-power circuit and a microscale VPFP-based plasma closure. It introduces Tesseract as a differentiable boundary layer that preserves gradient flow across heterogeneous solvers, including a high-fidelity C/CUDA code (Gkeyll), neural surrogates, and analytic closures, within a JAX-based AD workflow. The key contributions are (i) demonstration of end-to-end differentiability across solver boundaries, (ii) flexible V-to-I closures implemented via Gkeyll, symbolic regression with AD, and a closed-form surrogate, and (iii) rapid design-space exploration using gradient-based optimization (L-BFGS) over circuit and plasma parameters. The work highlights the practicality of modular, differentiable multiphysics workflows for engineering applications, enabling efficient tradeoffs between fidelity and computational budget while retaining the ability to prototype quickly.

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.

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

TL;DR

The paper tackles the challenge of building end-to-end differentiable simulations for a coupled multiphysics system in nuclear fusion, specifically a Z Pinch scenario with a macroscale pulsed-power circuit and a microscale VPFP-based plasma closure. It introduces Tesseract as a differentiable boundary layer that preserves gradient flow across heterogeneous solvers, including a high-fidelity C/CUDA code (Gkeyll), neural surrogates, and analytic closures, within a JAX-based AD workflow. The key contributions are (i) demonstration of end-to-end differentiability across solver boundaries, (ii) flexible V-to-I closures implemented via Gkeyll, symbolic regression with AD, and a closed-form surrogate, and (iii) rapid design-space exploration using gradient-based optimization (L-BFGS) over circuit and plasma parameters. The work highlights the practicality of modular, differentiable multiphysics workflows for engineering applications, enabling efficient tradeoffs between fidelity and computational budget while retaining the ability to prototype quickly.

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.

Paper Structure

This paper contains 11 sections, 4 equations, 2 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Fusion Z Pinch Compression Modeling
  3. Background
  4. Macroscale circuit equation
  5. Microscale plasma physics closure
  6. Software architecture
  7. Results
  8. Conclusion
  9. Technical Appendices
  10. Simulation parameters
  11. Computing resources

Figures (2)

  • Figure 1: Dataflow diagram of the software components used in solving Equations \ref{['eqn:Q_I_s']} and \ref{['eqn:vlasov_poisson']}. Different implementations of the $V \mapsto I$ map are packaged in Tesseracts. Implementation details are hidden behind a tesseract_jax call enabling rapid experimentation with different models.
  • Figure 2: Density-temperature trajectories of three rollouts of the circuit ODE, each with a different plasma model. Straight lines are lines of constant specific entropy.