v1.3.1

Release with new API symbols for downstream users to experiment with.

This should be treated as a pre-release of new features - API and bug fixes may come.

New release with fast ions and adaptive pedestal finalised coming soon.

v1.3.0

What's Changed

In this version, we introduce significant improvements to the geometry grid handling, enabling support for non-uniform grids. We also expand the physics capabilities with the Redl bootstrap current model, extended radiation modeling, and increased robustness in edge coupling. IMAS integration has improved with imas-core integration, and improved handling of core_profiles.

Physics features

• Neoclassical and Bootstrap Current
  • Added the Redl model for bootstrap current calculation.
  • Added Shaing corrections to the Angioni-Sauter neoclassical model, including smooth transitions between collision regimes.
• Radiation
  • Refactored Mavrin radiative cooling models to support smooth interpolation between coronal and non-coronal models.
  • Added extrapolation logic to better model radiative collapse scenarios (e.g., smoothly decaying to 0 below the polynomial fit range).
• Edge Modeling (Extended Lengyel)
  • Improved robustness
  • Added configuration for initial guesses and the ability to use previous time-step outputs as initial guesses to improve convergence speed and stability.
  • Parallelized multi-solution exploration in the Extended Lengyel forward mode solver, due to the possible existence of multiple roots. All unique roots are saved in output.
  • Fixed handling of impurity concentrations to prevent negative seeded impurities in inverse solver modes.
• Transport Models
  • Combined Transport Model: Introduced ADD and OVERWRITE merge modes. OVERWRITE allows specific models to lock transport coefficients in their domain, preventing subsequent models from modifying them.
  • Output Masking: Added flags (enable_chi_i, enable_D_e, etc.) to selectively disable specific transport channels within transport models.
  • Rotation: Changed the definition of toroidal_velocity from linear velocity (m/s) to angular velocity (rad/s). Introduced two different ExB suppression rules: a simple "Waltz rule" and a rule from van de Plassche 2020 which is more appropriate for toroidal-velocity limited cases.
• Heating and Source Terms
  • Added dW/dt smoothing (exponential moving average) to $P_{SOL}$ and $P_{loss}$ calculations to reduce noise in confinement time estimates.

Technical features

• Grid and Geometry
  • Non-uniform Grid Support: Major refactor of Grid1D and finite difference operators (face_grad, diffusion terms) to support non-uniformly spaced grids defined by face centers.
  • Updated CoreProfiles and Geometry to allow configuration via face_centers.

IMAS Integration

• IMAS Core
  TORAX now depends on IMAS Core as part of the imas Python v2.1.0 upgrade.
• Versioning
  • Upgraded support for imas Python library to v2.1.0.
  • Added options to explicitly convert equilibrium IDSs to the current Data Dictionary (DD) version during validation.
• Core Profiles
  • Added explicit validators for core_profiles IDS.
  • CoreProfiles now stores full ChargeStateInfo, optimizing impurity density scaling calculations.

UX/QoL improvements

• Outputs
  • Added new post-processing outputs for non-inductive currents: total non-inductive current density, total external current, and non-inductive current fraction.
• Documentation
  • Added CONTRIBUTING.rst to the documentation.
  • Updated documentation for rotation physics, sawtooth crash durations, and transport model inputs.

Bugfixes

• Lengyel Model: Fixed a bug where fixed impurities were incorrectly included in seeded impurity outputs, causing NaNs.
• Transport: Fixed a bug where optional transport sub-channels (like chi_bohm) were leaking outside their configured radial domains.
• Transport: Fixed a numerical instability for the QLKNN transport model when rotation was enabled.
• Transport: Fixed tracer leak in tglf_based_transport_model
• Numerics: Fixed an issue where deterministic time-stepping (with adaptive steps done internally as substeps) could hang if a minimum time-step (min_dt) was reached; it now exits gracefully.
• Geometry: Fixed boundary gradient calculations (face_grad) to correctly handle extrapolation on non-uniform grids.

v1.2.0

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.

v1.1.1

What's Changed

• Bugfixes for new input impurity API
  • Fixed input documentation examples
  • Included documentation for new per-species impurity outputs
  • Bugfix for sawtooth redistribution for current-diffusion-only
• New cumulative energy outputs
  • E_aux_total - total cumulative energy absorbed by plasma from external heating sources. Replaces E_aux which previously confusingly contained both external heating and Ohmic heating
  • E_external_total - total cumulative energy absorbed by plasma by both external heating sources and Ohmic source. Same as previous E_aux
  • E_external_injected - total cumulative injected energy from external heating sources. Differs from E_aux_total due to absorption factors.
  • E_ohmic_e - total cumulative ohmic heating
  • P_fusion - total fusion power (including neutron contribution)

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
• 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 and post_processed_outputs into three distinct categories. These are xarray.DataTrees contained in the top-level xarray.DataTree:
  • profiles: containing 1D profile values
  • scalars: containing scalar quantities
  • numerics: 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.