google-deepmind/torax v1.4.0

v1.4.0

Repository: google-deepmind/torax

Tag: v1.4.0

Published: 2026-05-15T17:53:42Z

Prerelease: no

Release notes:

What's Changed

In this version, we introduce expansions to core physics modeling, including new dynamic pedestal formation models, a new interface to the TGLF transport model, and support for prescribed fast ion profiles. Performance and numerics have been upgraded with a new Thomas Algorithm Block Tridiagonal solver. IMAS integration has been further extended to handle time-dependent geometries and input core sources.

Physics features

• Dynamic Pedestal Modelling
• Formation Modelling: Introduced dynamic pedestal modelling with two distinct modes:
• ADAPTIVE_SOURCE: Ramps pedestal boundary conditions to target values over a specified time window during transitions, using an adaptive artificial source to set kinetic profile values in the pedestal region. The transition time is set by the user.
• ADAPTIVE_TRANSPORT: Dynamically adapts transport coefficients based on trigger pedestal top values and includes a continuous ELM model.
• Transition Hysteresis: Introduced H-L back-transition hysteresis for both transport and source adaptive modes. Back-transitions now occur when $P_{SOL}$ drops below a user-configurable $P_{LH} \times$ hysteresis_factor (defaulting to 0.8).
• Dithering Support: A new pedestal state machine is introduced with modes L_MODE, H_MODE, TRANSITIONING_TO_H_MODE and TRANSITIONING_TO_L_MODE. These states are used to control dynamic pedestal behaviour including the possibility of dithering.
• Core Transport
• TGLF Interface: Added a new interface to execute TGLF as a core transport model.
• Heating and Fast Ions
• Prescribed Fast Ion Profiles: Added global support for prescribing fast ion populations (density and temperature) via the profile_conditions module. This allows overriding specific (source, species) pairs independent of localized heating models.
• Scaled ICRH Profile Model: Added a new scaled_profile Ion Cyclotron Resonance Heating (ICRH) model. This approximates RF heating adaptations by radially shifting prescribed heating profiles based on changes to the on-axis toroidal magnetic field ($B_0$) and rescaling to match a specified total absorbed power.
• Fast Ion Clipping: Added input clipping capabilities to the fast ion stabilization model to ensure inputs remain within the neural network's expected validity range.

Technical features

• Solvers and Numerics
• Block Tridiagonal Solver: Introduced a new Thomas Algorithm Block Tridiagonal solver class to more efficiently represent and solve discrete system convection and diffusion matrices.
• Bounded Execution: Added a jitted run_simulation wrapper that accepts a max_steps argument to bound simulation execution and prevent infinite hangs.
• Adaptive Timestepping: Improved efficiency with adaptive timestepping. Added a from_previous_dt timestep calculator which can reduce the number of total solver iterations.
• Bumped TORAX JAX dependency to 0.10.0 and switched to a new FAST_COMPILE flag (superceding xla_cpu_flatten_after_fusion) which includes specific optimizations for TORAX while_loops.
• Architecture
• Refactored Source Configurations to extract a base class, preparing the codebase for multiple ion cyclotron source models.
• Exposed source registration and pedestal model registration to the public API.

IMAS Integration

• Geometry
• Added the ability to load *all* IMAS time slices directly from the equilibrium IDS, generating a time-dependent geometry automatically without needing to manually specify multiple configs.
• Core Profiles and Sources
• Added parsing support for IMAS input core sources, automatically mapping them to corresponding existing TORAX sources.
• Added mapping for several new quantities to the output equilibrium IDS.

UX/QoL improvements

• Internal Boundary Conditions (IBCs)
• Set IBC profiles: Added the ability to set Internal Boundary Conditions directly via the profile_conditions config. This UX upgrade allows users to enforce internal boundary constraints at specific points or across spatial ranges. This can be used for imposing kinetic profiles in the L-mode edge, or as an alternative way to impose a full pedestal profile in H-mode.
• SparseTimeVaryingArray: Introduced a new sparse array structure to drive these IBCs. It can set uniform values or linear interpolations across spatial ranges, and smoothly interpolates between user-defined times. This enables setting partial profiles as internal boundary conditions instead of a value at a single radial location.
• Plotting and Visualization
• Plotly Migration: Upgraded the visualization backend to use Plotly, providing highly interactive, browser-friendly plots, enabling slider features also when working in notebooks
• Slider Modes Spatial plots now feature UI buttons to toggle the slider between "Time" (linearly spaced time points) and "Timesteps" (actual solver timesteps).
• Enhanced Plotruns Upgraded plotruns to accept a Mapping of plot data, allowing arbitrary labeled lines (removing the previous two-line restriction) while preventing unwanted in-place mutations of the data tree.
• Documentation and Warnings
• Geometry Clarity: Clarified the definition of $B_0$ throughout the codebase and docs as the vacuum toroidal magnetic field at $R_{major}$ (the geometric center of the LCFS), preventing ambiguity.
• Grid Warnings: Added a logging warning if the maximum spacing between rho_norm points in an input equilibrium grid is > 0.02, as coarse grids can lead to artifacts in derived quantities like $j_{total}$.
• Added comprehensive documentation for fast ion physics models detailing Bimaxwellian Split, Stix equations, and pressure dilution impacts.

Bugfixes

• Fast Ions: Fixed a critical numerical bug in safe_divide used in the new bimaxwellian_split $n_{tail}$ computation. A default epsilon of 1e-7 was dominating the W/particle loss rate (~1e-13), suppressing fast ion density by 5 orders of magnitude. The epsilon was adjusted to 1e-30, restoring physically correct ~1e18 m^-3 density values.
• Numerics: Modified safe_divide to make eps a mandatory argument, forcing developers to explicitly consider edge-case numerics.
• Numerics: Guarded against division-by-zero JAX edge cases in jnp.where expressions involving impurity…