Releases: google-deepmind/torax
v1.1.0
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
- 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
- Extended and improved geometry inputs. Now includes <1/R>, <B^2>, <B^-2>. 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.

v1.0.3
What's Changed
Bugfix for scenarios with non-zero Phi_b_dot
(rate of change of toroidal flux at LCFS).
An erroneous implementation of the Phi_b_dot
terms led to non-physical evolution for scenarios with large Phi_b_dot
, e.g. rapid ramp-down including reduction of plasma volume. The current diffusion equation and Ohmic power calculations were particularly affected.
Previous tests did not have large enough Phi_b_dot
to notice the effect. The present correct implementation was verified against RAPTOR for significantly larger Phi_b_dot
values than previously used.
It is highly recommended to rerun any previous simulations which included time-varying geometry with Phi_b
variations, and investigate the impact.
v1.0.2
What's Changed
Features
- New 'combined' transport model for additive coefficients from an arbitrary number of transport models each with configurable radial domain
- 'constant' transport model can now accept generally prescribed profile shapes (TimeVaryingArrays) at input
- More outputs:
- A_i, A_impurity
- beta_tor, beta_pol, beta_N
- integrated particle sources
- Z_eff
- Improved Pydantic validations of profile_conditions
- Order-of-magnitude sanity checking for plasma current, densities, and temperatures
- More informative logging if NaNs are discovered during the simulation
Cleanups / refactors
- Internal CoreProfiles densities now in m^-3
- Convertor utilities between scaled state vector x used in solver, and CoreProfiles attributes
- Various docs cleanups and fixes
Bugfixes
- Corrected spurious filtering of some geometry attributes from the output structure
ToraxConfig.update_fields
now validates updated Pydantic objects after all keys are updated, instead of one at a time- W_thermal calculation now uses consistent integration method as sawtooth energy conservation, improving dW_thermal/dt consistency. NOTE: as part of this fix,
pressure_thermal_e
,pressure_thermal_i
,pressure_thermal_total
have internally become CellVariables and are now outputted on the extended cell grid, instead of the face grid. - Z_eff_face used in physics formulas are now self-consistently calculated from CoreProfiles, instead of using a less consistent interpolation from the input Z_eff profile
v1.0.1
Bugfixes and I/O consistency improvements.
- Bugfix plot_config loading for
plotrun.py
- I/O naming consistency improvement for v_loop
- ICRH minority concentration made a fraction as opposed to a percentage, to align with the main_ion and impurity ion mixture API
v1.0.0
Key Highlights
This release focuses heavily on internal refactoring, code simplification, and improving the structure and consistency of outputs and configurations, as well as releasing new physics features. Users will notice significant changes in how simulation data is organized and accessed, as well as improvements to documentation and configuration handling. A TORAX API is defined and breaking changes (e.g. in existing config structure, input and output variable names) will be avoided throughout TORAX v1, ensuring stability for user applications even as new features roll out.
TORAX API
The following objects are released as part of the TORAX API:
Classes and functions for loading configs and running a simulation:
build_torax_config_from_file
import_module
ToraxConfig
run_simulation
Classes returned by the simulation:
StateHistory
CoreProfiles
SourceProfiles
CoreTransport
PostProcessedOutputs
SolverNumericOutputs
SimError
If you would like further classes or methods exposed please reach out to us.
Physics features
Sawtooth modeling
- Sawtooth trigger and redistribution model. A sawtooth model inherits from the
Solver
class and implements an alternative simulation step if a sawtooth is triggered. If a state-dependent condition is triggered, then the kinetic profiles and plasma current are redistributed over a short user-defined timestep. Currently simple trigger and redistribution models are implemented, setting a user-defined critical magnetic shear at the q=1 surface as the trigger, flattening the profiles inside q=1 and redistributing them up to a user-defined mixing radius while conserving energy and total current.
Neoclassical API
- Bootstrap current and conductivity have moved to a new "neoclassical" section of the config
- This paves the way for the addition of future neoclassical models e.g. for transport
UX improvements
PyPI availability
Both TORAX v0.3.3 and v1.0.0 are now available on PyPI. Can pip install torax.
Output Restructure and Renaming:
A major effort has been undertaken to standardize and improve the organization and naming of output variables. This output structure will be maintained for TORAX v1. This includes:
- Separating outputs from the existing
core_profiles
,core_sources
,geometry
,core_transport
andpost_processed_outputs
into three distinct categories. These arexarray.DataTrees
contained in the top-levelxarray.DataTree
:profiles
: containing 1D profile valuesscalars
: containing scalar quantitiesnumerics
: containing quantities relating to solver numerics
- Renaming variables to be clearer and more consistent
- Removing internal TORAX fvm information from output names, e.g. remove suffixes β_faceβ from names. Instead, where appropriate, the grid that a variable is defined on can be accessed by its
Dataset
coordinates - Variables which previously had both face-grid and cell-grid versions are now merged onto a single grid, which consists of the cell grid extended to include the rho_norm=0 and rho_norm=1 boundaries
- Consistent units. MW β W, MA βA, and output density now in m^-3 instead of normalized to a reference density
- P_external has been removed and replaced by P_aux, the auxiliary (external heating) power. For total input power it is necessary to add P_aux to P_ohmic
Configuration Refactoring and Renaming:
The configuration system has been significantly updated for better clarity and consistency. This configuration structure will be supported throughout TORAX v1. This involves:
- Widespread renaming of configuration fields and model names across various modules (sources, transport models, profiles, geometry, etc.)
- Improvements to how configurations are loaded and validated, including resolving file paths as opposed to module loading
- Simplifies auto-complete when loading config files from scripts
- Avoids complications on resolving the root package when loading modules
- Density inputs now in m^-3 as opposed to normalized by a reference density. Reference density input is removed
- Better considered default values
See the documentation at https://torax.readthedocs.io/en/v1.0.0/configuration.html
Improved Documentation:
Significant updates have been made throughout the documentation. Streamlined launching of the notebook tutorials.
Removal of env vars
All environment variables used for setting various paths in TORAX have been removed. These are now configured in the config dict and Pydantic ToraxConfig.
Additional output variables
- Number of solver iteration steps
- Location of q-profile rational surfaces, and q_min, over time
Bugfixes
- Ohmic power now always positive regardless of sign of v_loop and current
- Prescribed sources now supports sources with multiple outputs (e.g. electron heat and current), as expected
v0.3.3
Upcoming API Changes: The next major release of TORAX will include significant API updates including: changes to the config structure/nesting, improved output naming and structure and improved module loading.
This release (v0.3.3) provides a minor bug fix, along with some non-breaking API changes and backend improvements.
- Small bug fix to be more robust to QLKNN10D paths.
- Removes the runtime_params nesting in the config moving them to the top level. For now configs with the nesting are supported but this will be going away.
v0.3.2
Upcoming API Changes: The next major release of TORAX will include significant API updates including: changes to the config structure/nesting, improved output naming and structure and improved module loading. This release (v0.3.2) provides primarily UX improvements as well as a sawteeth model. A few minor breaking changes to the config and scripts are part of this release. This version will be packaged and made available on PyPI.
- UX improvements
- We now use fusion_surrogates model registry, so specifying a transport model path using environment variables is now optional
- Geo data default path is set relative to the installed TORAX library, so setting the environment variable at install is not required anymore.
- Added binary scripts
run_torax
andplot_torax
with pip installation.
- Phyisics features
- Sawteeth model added
- Breaking changes
- Default pedestal model has changed, it is now
no_pedestal
set_pedestal
has moved to the pedestal config and is now by defaultfalse
nottrue
- The
run_simulation_main.py
script has been moved inside the library code, so it can be packaged in the PyPI version.
- Default pedestal model has changed, it is now
- Bug fixes
- Fixed some JAX compile cache misses related to ion cyclotron source and time dependent geometry. This may speed up subsequent simulation runs.
v0.3.1
Upcoming API Changes: The next major release of TORAX will include significant API updates including: changes to the config structure/nesting, improved output naming and structure and improved module loading. This release (v0.3.1) provides primarily bugfixes, UX improvements, and new tutorials compared to v0.3.0, and allows users to pick up improvements with a stable API.
-
Physics features
- QLKNN_7_11 now the default turbulent transport model
- Absorbed fraction input field for ICRH and generic heating
- vloop_lcfs output available for Ip boundary condition cases
-
UX improvements
- Allow lists to be input in user-facing config dict
- Remove ToraxSimOutput, StateHistory is the new public output class (before xarray conversion)
- Create utility for registering new source model configs.
- Add simulation error for negative temperatures and densities
- Add serialized ToraxConfig to TORAX outputs
- Tutorials
-
Bugfixes + miscellanous
- Fix missed jit cache hits
- Fix stale documentation
- Various refactors for internal API simplification
v0.3.0
Upcoming API Changes: The next major release of TORAX will include significant API updates including: changes to the config structure/nesting, improved output naming and structure and improved module loading. This release (v0.3.0) provides a stable version with significant changes from v0.2.0 for users to pick up improvements with a stable API.
TORAX improvements since v0.2.0 (highlights)
-
Changes to interacting with TORAX programatically
- Moved the config to Pydantic improving validation and readability
- Replaced sim.Sim with
torax.run_simulation
which takes atorax.ToraxConfig
-
Physics features
- Average charge state equilibria added
- QLKNN_7_11 transport model
- qualikiz transport model
- Improved simple models for pedestal boundary condition
- Can split main ion and impurity into mixtures of ion species. DT ratio now impacts fusion power. Radiation for each separate impurity can be calculated
- Improved radiation models including Mavrin polynomial fits to ADAS data
- Cyclotron radiation
- New Vloop boundary condition option for current diffusion equation
-
**Miscellaneous and bugfixes
- Fixes for q_face and s_face being stale
- Fix in He3 tail temperature calculation
- P_LH_minimum bugfix
- Extra post-processed outputs added
- Improved on-axis treatment of pprime and plasma current
- Minor performance optimizations (~5% in compilation time and runtime)
- Refactors: various restructuring and cleanups
New Contributors
- @hassec made their first contribution in #599
- @anteplu made their first contribution in #617
- @emmanuel-ferdman made their first contribution in #740
- @shreyans413 made their first contribution in #773
- @harshtech123 made their first contribution in #827
- @demoncoder-crypto made their first contribution in #869
Full Changelog: v0.2.0...v0.3.0
v0.2.0
Upcoming API Changes: The next major release of TORAX will include significant API updates, including changes to the configuration structure, a transition to a hierarchical xarray DataTree output format, and variable renaming for improved consistency. This release (v0.2.0) provides a stable version for users who wish to continue development with the current API.
TORAX improvements since v0.1.0 (highlights)
-
Physics Model Improvements:
- Added Gaussian smoothing option for transport models.
- Implemented support for spatially varying Zeff profiles.
- Integrated a neural network wrapper for the Toric ICRH ML-surrogate model (using Flax).
- Added Bremsstrahlung radiation sink.
- Added BohmGyroBohm transport model.
- Implemented ECCD with Lin-Liu current drive formula.
-
Geometry Input Extensions:
- Added support for MEQ (FBT) geometry input.
- Enabled EQDSK geometry input with flux-surface averaging.
- Implemented support for time-dependent geometries:
- Support for multiple geometry files or bundled FBT files (via liuqe meqlpack).
- Calculation of time-dependent
phibdot
term in governing equations.
-
Configuration and Sim object construction improvements
- Enabled flexible input of prescribed data using Python primitives, NumPy, and xarray.
- Expanded support for prescribing time-dependent profiles, including core profiles and sources.
- Significantly improved simulation setup speed by pre-interpolating input arrays and preparing time interpolation providers.
- Enhanced the methodology for Ip scaling, supporting both parameter-based and geometry file-based definitions.
-
Developer Quality of Life:
- Redesigned source management for improved modularity and easier integration of new sources.
- Restructured transport_models to increase modularity, simplify new model integration, and reduce redundant code.
- Enhanced type safety with array typing and typeguard annotations.
- Implemented helper scripts for managing and updating simulation integration tests.
- Divided
SimulationStepFn
into distinct phases, enabling external modification of dynamic runtime parameters at each simulation step.
-
Bugfixes:
- Corrected QLKNN magnetic shear calculation.
- Fixed Coulomb constant calculation.
- Resolved issues with impurity density and boundary condition updates.
- Fixed a bug in external current profile updates.
-
Extended Outputs:
- Added stepper numeric outputs, including error codes and iteration counts.
- Expanded post-processed outputs, such as integrated currents, sources, fusion gain, and pressure.
- Added output of the plasma current profile calculated from the poloidal flux.
-
Documentation:
- Significantly expanded documentation, covering physics models, solver details, code structure, and developer guides.
- Improved docstrings and comments for enhanced code readability.
-
Miscellaneous:
- Enabled nightly and pull request-triggered unit tests in CI.
- Added a JAX profiler (optional, flag-activated) for performance analysis.
- Implemented a configurable, data-driven plotting library.
- Implemented various performance optimizations, including solver improvements and JIT compilation of simulation setup.
- Enabled JAX persistent cache for faster execution.
- Added support for restarting simulations from arbitrary time points using existing TORAX output files.
- Modified governing equations for consistency with normalized toroidal flux as the radial coordinate.
- Improved various equation calculations, including q-profile and current calculations.