v1.0.0: Merge pull request #804 from google-deepmind:onnx-docs

[jcitrin]

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

---

This discussion was created from the release v1.0.0: Merge pull request #804 from google-deepmind:onnx-docs.