google-deepmind/torax v1.1.0

v1.1.0

Repository: google-deepmind/torax

Tag: v1.1.0

Published: 2025-09-12T15:01:33Z

Prerelease: no

Release notes:

What's Changed

Physics features

• Angioni-Sauter neoclassical transport model
• Note: not yet with near-axis potato orbit corrections (coming soon)
• Extended impurity input API and more detailed impurity outputs
• Impurities can now be set as ratios of electron density, with emergent Z_eff. Optionally one species can be set as unknown, and its density then constrained by input Z_eff. Legacy fractions impurity input mode is still available.
• Support for radially varying input of prescribed impurity species
• More detailed per-species impurity outputs for density, average charge, and radiation
• Extended and improved geometry inputs. Now includes , , . Important for strong shaping.

Technical features

• Full sim JAX compatibility
• Enables full simulation differentiability (coming soon)
• Added custom run loop that fully runs under jax.jit
• Expanded JIT scope to step function (usable with TORAX_EXPERIMENTAL_COMPILE=True)
• Significant runtime speed-ups. Only minor impact on compile time due to XLA/CPU improvements and usage of `xla_cpu_flatten_after_fusion' flag.
• Reduced compiler inlining of JAX functions due to a new XLA feature turned on by xla_cpu_flatten_after_fusion
• Introduction of runtime shape checking with jaxtyping
• This feature is not rolled out yet everywhere and will be completed in a later release

Cleanups / refactors

• Restructure of internal solver API and state scalings
• Improved modularity of initialization and simulation preparation routines
• New "neoclassical_models" container for transport, conductivity, bootstrap current
• Unification of standard and adaptive step routines
• Dataclasses are now registered as standard JAX registered dataclasses. Allowed unification of "dynamic" and "static" parameters, with static parameters (that retrigger compilation) set through Annotation.
• Runtime parameters jargon simplification. Now we just have RuntimeParamsProvider and RuntimeParams (which are the params at a given time). No more concept of dynamic, static, slice.
• Improved structure for IMAS I/O code
• Extension to more IDS I/O beyond equilibrium is ongoing and will be part of future releases

UX/QoL improvements

• Upgraded Python requirement to 3.11
• Improved linting and VSCode workspace settings
• Minor plotting improvements
• Various docstring and documentation fixes/improvements

Bugfixes

• Tightening of various torax_pydantic types for validation
• Correct geo_t_plus_dt now used in adaptive stepping
• Fixed n_rho update in ToraxConfig.update_fields
• Corrected calculations for single effective Z_impurity and n_impurity when bundling several impurities together
• Minor corrections to physics constants

To illustrate the impact of TORAX technical improvements, see the plot below showing the typical runtimes for the iterhybrid_rampup config over time.