v1.4.0

What's Changed

In this version, we introduce expansions to core physics modeling, including new dynamic pedestal formation models, a new interface to the TGLF transport model, and support for prescribed fast ion profiles. Performance and numerics have been upgraded with a new Thomas Algorithm Block Tridiagonal solver. IMAS integration has been further extended to handle time-dependent geometries and input core sources.

Physics features

• Dynamic Pedestal Modelling
  • Formation Modelling: Introduced dynamic pedestal modelling with two distinct modes:
    • ADAPTIVE_SOURCE: Ramps pedestal boundary conditions to target values over a specified time window during transitions, using an adaptive artificial source to set kinetic profile values in the pedestal region. The transition time is set by the user.
    • ADAPTIVE_TRANSPORT: Dynamically adapts transport coefficients based on trigger pedestal top values and includes a continuous ELM model.
  • Transition Hysteresis: Introduced H-L back-transition hysteresis for both transport and source adaptive modes. Back-transitions now occur when $P_{SOL}$ drops below a user-configurable $P_{LH} \times$ hysteresis_factor (defaulting to 0.8).
  • Dithering Support: A new pedestal state machine is introduced with modes L_MODE, H_MODE, TRANSITIONING_TO_H_MODE and TRANSITIONING_TO_L_MODE. These states are used to control dynamic pedestal behaviour including the possibility of dithering.
• Core Transport
  • TGLF Interface: Added a new interface to execute TGLF as a core transport model.
• Heating and Fast Ions
  • Prescribed Fast Ion Profiles: Added global support for prescribing fast ion populations (density and temperature) via the profile_conditions module. This allows overriding specific (source, species) pairs independent of localized heating models.
  • Scaled ICRH Profile Model: Added a new scaled_profile Ion Cyclotron Resonance Heating (ICRH) model. This approximates RF heating adaptations by radially shifting prescribed heating profiles based on changes to the on-axis toroidal magnetic field ($B_0$) and rescaling to match a specified total absorbed power.
  • Fast Ion Clipping: Added input clipping capabilities to the fast ion stabilization model to ensure inputs remain within the neural network's expected validity range.

Technical features

• Solvers and Numerics
  • Block Tridiagonal Solver: Introduced a new Thomas Algorithm Block Tridiagonal solver class to more efficiently represent and solve discrete system convection and diffusion matrices.
  • Bounded Execution: Added a jitted run_simulation wrapper that accepts a max_steps argument to bound simulation execution and prevent infinite hangs.
  • Adaptive Timestepping: Improved efficiency with adaptive timestepping. Added a from_previous_dt timestep calculator which can reduce the number of total solver iterations.
  • Bumped TORAX JAX dependency to 0.10.0 and switched to a new FAST_COMPILE flag (superceding xla_cpu_flatten_after_fusion) which includes specific optimizations for TORAX while_loops.
• Architecture
  • Refactored Source Configurations to extract a base class, preparing the codebase for multiple ion cyclotron source models.
  • Exposed source registration and pedestal model registration to the public API.

IMAS Integration

• Geometry
  • Added the ability to load all IMAS time slices directly from the equilibrium IDS, generating a time-dependent geometry automatically without needing to manually specify multiple configs.
• Core Profiles and Sources
  • Added parsing support for IMAS input core sources, automatically mapping them to corresponding existing TORAX sources.
  • Added mapping for several new quantities to the output equilibrium IDS.

UX/QoL improvements

• Internal Boundary Conditions (IBCs)
  • Set IBC profiles: Added the ability to set Internal Boundary Conditions directly via the profile_conditions config. This UX upgrade allows users to enforce internal boundary constraints at specific points or across spatial ranges. This can be used for imposing kinetic profiles in the L-mode edge, or as an alternative way to impose a full pedestal profile in H-mode.
  • SparseTimeVaryingArray: Introduced a new sparse array structure to drive these IBCs. It can set uniform values or linear interpolations across spatial ranges, and smoothly interpolates between user-defined times. This enables setting partial profiles as internal boundary conditions instead of a value at a single radial location.
• Plotting and Visualization
  • Plotly Migration: Upgraded the visualization backend to use Plotly, providing highly interactive, browser-friendly plots, enabling slider features also when working in notebooks
  • Slider Modes Spatial plots now feature UI buttons to toggle the slider between "Time" (linearly spaced time points) and "Timesteps" (actual solver timesteps).
  • Enhanced Plotruns Upgraded plotruns to accept a Mapping of plot data, allowing arbitrary labeled lines (removing the previous two-line restriction) while preventing unwanted in-place mutations of the data tree.
• Documentation and Warnings
  • Geometry Clarity: Clarified the definition of $B_0$ throughout the codebase and docs as the vacuum toroidal magnetic field at $R_{major}$ (the geometric center of the LCFS), preventing ambiguity.
  • Grid Warnings: Added a logging warning if the maximum spacing between rho_norm points in an input equilibrium grid is > 0.02, as coarse grids can lead to artifacts in derived quantities like $j_{total}$.
  • Added comprehensive documentation for fast ion physics models detailing Bimaxwellian Split, Stix equations, and pressure dilution impacts.

Bugfixes

• Fast Ions: Fixed a critical numerical bug in safe_divide used in the new bimaxwellian_split $n_{tail}$ computation. A default epsilon of 1e-7 was dominating the W/particle loss rate (~1e-13), suppressing fast ion density by 5 orders of magnitude. The epsilon was adjusted to 1e-30, restoring physically correct ~1e18 m^-3 density values.
• Numerics: Modified safe_divide to make eps a mandatory argument, forcing developers to explicitly consider edge-case numerics.
• Numerics: Guarded against division-by-zero JAX edge cases in jnp.where expressions involving impurity properties when $Z_{eff} = 1.0$.
• Dilution: Fixed an operator precedence bug in the dilution factor calculation for n_e_ratios that caused only the last term of the sum to be divided by the main ion charge. This has no impact on hydrogenic main ions.
• Plotting: Fixed a bug causing slider failures when data was defined on different time grids.