google-deepmind/torax v1.2.0

v1.2.0

Repository: google-deepmind/torax

Tag: v1.2.0

Published: 2025-12-09T19:22:52Z

Prerelease: no

Release notes:

What's Changed

In this version we introduce an experimental API - torax.experimental.

This exposes several elements that are useful for running a simulation under a JAX JITted context and for loosely coupling explicit geometry overrides to TORAX. See https://torax.readthedocs.io/en/v1.2.0/developer_guides.html for example usages for both. For more detailed questions on use please feel free to reach out.

The experimental API is subject to change and experimental symbols will be moved to the stable public API in TORAX v2.0.0.

Physics features

• Edge Modeling (Extended Lengyel Model)
• Implementation of the Extended Lengyel model for determining separatrix boundary conditions and impurity concentrations. Note: the coupling behaviour is currently undergoing verification so is currently classified as experimental.
• Includes Kallenbach 2024 divertor enrichment model scaling.
• Rotation and Electric Field
• Calculation of Neoclassical poloidal velocity.
• Calculation of Radial Electric Field ($E_r$) and $E \times B$ velocity.
• Added $v_{E \times B}$ shear as an input to TGLF-based and QLK-based models.
• Note that these terms and their impact are currently undergoing verification so are currently classified as experimental.
• Heating and Current Drive
• Corrected Electron Cyclotron (EC) current drive efficiency formulae (Lin-Liu model).
• Conversion routines between parallel and toroidal currents
• Transport Models
• Added support for the TGLFNN-UKAEA transport model.

Technical features

• Differentiability and Optimization
• Implementation of whilei_loop and while_loop_bounded to enable backward-mode differentiation (gradients) through the adaptive time-stepping loop.
• Enabled use_jax_custom_root by default for Newton-Raphson solvers to facilitate differentiation.
• JAX and Compilation
• StaticDataclass: Introduced and migrated to a base class for configuration and model objects to ensure correct hashing and equality semantics within jax.jit, preventing spurious recompilations.
• Promoted experimental JIT compilation of the step function to default (removed TORAX_EXPERIMENTAL_COMPILE flag).
• Bumped JAX version to 0.8.1 and Chex to 0.1.90.

IMAS Integration

• Core Profiles I/O
• Added capability to populate profile_conditions directly from an IMAS core_profiles IDS.
• Added function to output TORAX state to an IMAS core_profiles IDS.
• Loader improvements
• Added support for explicit data dictionary conversion in the IMAS loader.

UX/QoL improvements

• Examples and Tutorials
• Added OpenSTEP (STEP SPP-001) flat-top example configuration.
• Added a tutorial notebook demonstrating vmap and grad over a simulation.

Bugfixes

• Corrected Phi_b_dot calculation cache misses by setting it at object creation time.