Pith. sign in

REVIEW 7 cited by

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

Not yet reviewed by Pith; the record is open.

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2406.06718 v4 pith:UDNVJRAO submitted 2024-06-10 physics.plasm-ph

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

classification physics.plasm-ph
keywords toraxtransporttokamakdifferentiablefastheatmodelssimulator
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved
0 comments
Share X Bluesky LinkedIn Reddit HN
read the original abstract

We present TORAX, a new, open-source, differentiable tokamak core transport simulator implemented in Python using the JAX framework. TORAX solves the coupled equations for ion heat transport, electron heat transport, particle transport, and current diffusion, incorporating modular physics-based and ML models. JAX's just-in-time compilation ensures fast runtimes, while its automatic differentiation capability enables gradient-based optimization workflows and simplifies the use of Jacobian-based PDE solvers. Coupling to ML-surrogates of physics models is greatly facilitated by JAX's intrinsic support for neural network development and inference. TORAX is verified against the established RAPTOR code, demonstrating agreement in simulated plasma profiles. TORAX provides a powerful and versatile tool for accelerating research in tokamak scenario modeling, pulse design, and control.

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 7 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Enhanced alpha channeling with spin-polarized fuel

    physics.plasm-ph 2026-07 conditional novelty 7.0

    Vector-aligned spin-polarized D-T fuel keeps its perpendicular alpha birth bias through slowing-down, raises alpha-channeling efficiency ~1.5x, and projects 2-4x fusion-power gains if high-efficiency channeling waves ...

  2. iGENE: A Differentiable Flux-Tube Gyrokinetic Code in TensorFlow

    physics.plasm-ph 2026-05 unverdicted novelty 7.0

    A fully differentiable TensorFlow gyrokinetic code allows approximate gradients of nonlinear turbulence quantities to be used for outer-loop tasks such as profile prediction despite stochasticity.

  3. Predictive capabilities of the integrated modeling TRANSP code for tokamak plasmas

    physics.plasm-ph 2026-05 conditional novelty 6.0

    TRANSP's PT_SOLVER module is described and verified via analytic, manufactured, stiff-diffusion, and TGYRO benchmarks, with a new T3D/GX coupling shown.

  4. Accelerating integrated modeling with surrogate-based optimization: the MAESTRO workflow

    physics.plasm-ph 2026-05 unverdicted novelty 6.0

    MAESTRO couples surrogate optimization transport modeling with external solvers to enable efficient full-physics steady-state plasma predictions in fusion devices.

  5. Kinetic Equilibrium Prediction at TCV using RAPTOR and FBT

    physics.plasm-ph 2026-03 conditional novelty 6.0

    A pre-shot workflow coupling RAPTOR transport and FBT equilibrium simulations predicts TCV discharge profiles and improves feedforward coil-current preparation for plasma shape control.

  6. Using Neural Emulators and Hamiltonian Monte Carlo to constrain the Epoch of Reionization's History with the Ly$\alpha$ Forest Power Spectrum

    astro-ph.CO 2025-09 conditional novelty 6.0

    A JAX pipeline that emulates both the Lyα forest power spectrum and its model-dependent covariance with neural networks, then runs Hamiltonian Monte Carlo, recovers reionization-history parameters from mock observations.

  7. Predictive capabilities of the integrated modeling TRANSP code for tokamak plasmas

    physics.plasm-ph 2026-05 accept novelty 5.0

    TRANSP now includes a new implicit solver PT_SOLVER and T3D/GX integration that provides robust numerical handling for predictive transport modeling in tokamaks.