torax.config package

Submodules

torax.config.build_runtime_params module

Methods for building simulation parameters.

For the static_runtime_params_slice this is a method build_static_runtime_params__from_config. For the dynamic_runtime_params_slice this is a class DynamicRuntimeParamsSliceProvider which provides a slice of the DynamicRuntimeParamsSlice to use during time t of the sim. This module also provides a method get_consistent_dynamic_runtime_params_slice_and_geometry which returns a DynamicRuntimeParamsSlice and a corresponding geometry with consistent Ip.

class torax.config.build_runtime_params.DynamicRuntimeParamsSliceProvider(torax_config)[source]

Bases: object

Provides a DynamicRuntimeParamsSlice to use during time t of the sim.

The DynamicRuntimeParamsSlice may change from time step to time step, so this class interpolates any time-dependent params in the input config to the values they should be at time t.

NOTE: In order to maintain consistency between the DynamicRuntimeParamsSlice and the geometry, sim.get_consistent_dynamic_runtime_params_slice_and_geometry should be used to get a slice of the DynamicRuntimeParamsSlice and a corresponding geometry.

See run_simulation() for how this callable is used.

After this object has been constructed changes any runtime params may not be picked up if they are updated and it is safest to construct a new provider object (if for example updating the simulation). ```

Parameters:

torax_config (ToraxConfig)

classmethod from_config(config)[source]

Constructs a DynamicRuntimeParamsSliceProvider from a ToraxConfig.

Parameters:

config (ToraxConfig)

Return type:

Self

torax.config.build_runtime_params.build_static_params_from_config(config)[source]

Builds a StaticRuntimeParamsSlice from a ToraxConfig.

Parameters:

config (ToraxConfig)

Return type:

StaticRuntimeParamsSlice

torax.config.build_runtime_params.get_consistent_dynamic_runtime_params_slice_and_geometry(*, t, dynamic_runtime_params_slice_provider, geometry_provider)[source]

Returns the dynamic runtime params and geometry for a given time.

Parameters:
Return type:

tuple[DynamicRuntimeParamsSlice, Geometry]

torax.config.config_args module

Functions for building arguments for configs and runtime input params.

torax.config.config_args.recursive_replace(obj, ignore_extra_kwargs=False, **changes)[source]

Recursive version of dataclasses.replace.

This allows updating of nested dataclasses. Assumes all dict-valued keys in changes are themselves changes to apply to fields of obj.

Parameters:

  • obj (TypeVar(_T)) – Any dataclass instance.

  • ignore_extra_kwargs (bool) – If True, any kwargs from changes are ignored if they do not apply to obj.

  • **changes – Dict of updates to apply to fields of obj.

Return type:

TypeVar(_T)

Returns:

A copy of obj with the changes applied.

torax.config.config_loader module

Functions to load a config from config file or directory.

torax.config.config_loader.build_torax_config_from_config_module(config_module_str, config_package=None)[source]

Returns a Sim and RuntimeParams from the config module.

Parameters:

  • config_module_str (str) – Python package path to config module. E.g. torax.examples.iterhybrid_predictor_corrector.

  • config_package (Optional[str]) – Optional, base package config is imported from. See config_package flag docs.

Return type:

ToraxConfig

torax.config.config_loader.example_config_paths()[source]

Returns a tuple of example config paths.

Return type:

dict[Literal['basic_config', 'iterhybrid_predictor_corrector', 'iterhybrid_rampup'], Path]

torax.config.config_loader.import_config_dict(path)[source]

Import a Torax config dictionary from a file.

Parameters:

path (str | Path) – The path to the config file. The path can be represented as a string or a pathlib.Path object.

Return type:

dict[str, Any]

Returns:

The config dictionary.

torax.config.config_loader.import_module(module_name, config_package=None)[source]

Imports a module.

Parameters:

torax.config.numerics module

Numerics parameters used throughout TORAX simulations.

class torax.config.numerics.DynamicNumerics(t_initial, t_final, exact_t_final, maxdt, mindt, dtmult, fixed_dt, dt_reduction_factor, resistivity_mult, nref, largeValue_T, largeValue_n, calcphibdot)[source]

Bases: Mapping

Generic numeric parameters for the simulation.

Parameters:
items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
values() → an object providing a view on D's values
class torax.config.numerics.Numerics(**data)[source]

Bases: BaseModelFrozen

Generic numeric parameters for the simulation.

The from_dict(…) method can accept a dictionary defined by https://torax.readthedocs.io/en/latest/configuration.html#numerics.

t_initial

Simulation start time, in units of seconds.

t_final

Simulation end time, in units of seconds.

exact_t_final

If True, ensures that the simulation end time is exactly t_final, by adapting the final dt to match.

maxdt

Maximum timesteps allowed in the simulation. This is only used with the chi_time_step_calculator time_step_calculator.

mindt

Minimum timestep allowed in simulation.

dtmult

Prefactor in front of chi_timestep_calculator base timestep dt=dx^2/(2*chi). In most use-cases can be increased further above this.

fixed_dt

Timestep used for fixed_time_step_calculator.

adaptive_dt

Iterative reduction of dt if nonlinear step does not converge, if nonlinear step does not converge, then the step is redone iteratively at successively lower dt until convergence is reached.

dt_reduction_factor

Factor by which to reduce dt if adaptive_dt is True.

ion_heat_eq

Solve the ion heat equation (ion temperature evolves over time).

el_heat_eq

Solve the electron heat equation (electron temperature evolves over time)

current_eq

Solve the current equation (current evolves over time).

dens_eq

Solve the density equation (n evolves over time).

calcphibdot

Calculate Phibdot in the geometry dataclasses. This is used in calc_coeffs to calculate terms related to time-dependent geometry. Can set to false to zero out for testing purposes.

resistivity_mult

1/multiplication factor for sigma (conductivity) to reduce current diffusion timescale to be closer to heat diffusion timescale

nref

Reference density value for normalizations.

largeValue_T

Prefactor for adaptive source term for setting temperature internal boundary conditions.

largeValue_n

Prefactor for adaptive source term for setting density internal boundary conditions.

Parameters:

data (Any)

build_dynamic_params(t)[source]

Builds a DynamicNumerics.

Parameters:

t (Union[Array, ndarray, bool, number, float, int])

Return type:

DynamicNumerics

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

torax.config.plasma_composition module

Plasma composition parameters used throughout TORAX simulations.

class torax.config.plasma_composition.DynamicIonMixture(fractions, avg_A, Z_override=None)[source]

Bases: Mapping

Represents a fixed mixture of ion species at a specific time.

Information on ion names are not stored here, but rather in StaticRuntimeParamsSlice, to simplify JAX logic and performance in source functions for fusion power and radiation which are species-dependent.

fractions

Ion fractions for a time slice.

avg_A

Average A of the mixture.

Z_override

Typically, the average Z is calculated according to the temperature dependent charge-state-distribution, or for low-Z cases by the atomic numbers of the ions assuming full ionization. If Z_override is provided, it is used instead for the average Z.

Parameters:

  • fractions (Union[Float[Array, 'rhon'], Float[ndarray, 'rhon']])

  • avg_A (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • Z_override (Union[Float[Array, ''], Float[ndarray, ''], number, float, None])

items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
values() → an object providing a view on D's values
class torax.config.plasma_composition.DynamicPlasmaComposition(main_ion, impurity, Zeff, Zeff_face)[source]

Bases: Mapping

Parameters:
items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
values() → an object providing a view on D's values
class torax.config.plasma_composition.IonMixture(**data)[source]

Bases: BaseModelFrozen

Represents a mixture of ion species. The mixture can depend on time.

Main use cases: 1. Represent a bundled mixture of hydrogenic main ions (e.g. D and T) 2. Represent a bundled impurity species where the avg charge state, mass,

and radiation is consistent with each fractional concentration, and these quantities are then averaged over the mixture to represent a single impurity species in the transport equations for efficiency.

species

A dict mapping ion symbols (from ION_SYMBOLS) to their fractional concentration in the mixture. The fractions must sum to 1.

Z_override

An optional override for the average charge (Z) of the mixture.

A_override

An optional override for the average mass (A) of the mixture.

Parameters:

data (Any)

build_dynamic_params(t)[source]

Creates a DynamicIonMixture object at a given time.

Optional overrides for Z and A can be provided.

Parameters:

t (Union[Array, ndarray, bool, number, float, int]) – The time at which to build the DynamicIonMixture.

Return type:

DynamicIonMixture

Returns:

A DynamicIonMixture object.

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class torax.config.plasma_composition.PlasmaComposition(**data)[source]

Bases: BaseModelFrozen

Configuration for the plasma composition.

The from_dict(…) method can accept a dictionary defined by https://torax.readthedocs.io/en/latest/configuration.html#plasma-composition.

main_ion

Main ion species. Can be single ion or a mixture of ions (e.g. D and T). Either a single ion, and constant mixture, or a time-dependent mixture. For single ions the input is one of the allowed strings in ION_SYMBOLS. For mixtures the input is an IonMixture object, constructed from a dict mapping ion symbols to their fractional concentration in the mixture.

impurity

Impurity ion species. Same format as main_ion.

Zeff

Constraint for impurity densities.

Zi_override

Optional arbitrary masses and charges which can be used to override the data for the average Z and A of each IonMixture for main_ions or impurities. Useful for testing or testing physical sensitivities, outside the constraint of allowed impurity species.

Ai_override

Optional arbitrary masses and charges which can be used to override the data for the average Z and A of each IonMixture for main_ions or impurities. Useful for testing or testing physical sensitivities, outside the constraint of allowed impurity species.

Zimp_override

Optional arbitrary masses and charges which can

Parameters:

data (Any)

get_impurity_names()[source]

Returns the impurity symbol strings from the input.

Return type:

tuple[str, ...]

get_main_ion_names()[source]

Returns the main ion symbol strings from the input.

Return type:

tuple[str, ...]

property impurity_mixture: IonMixture

Returns the IonMixture object for the impurity ions.

property main_ion_mixture: IonMixture

Returns the IonMixture object for the main ions.

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

torax.config.profile_conditions module

Profile condition parameters used throughout TORAX simulations.

class torax.config.profile_conditions.DynamicProfileConditions(Ip_tot, vloop_lcfs, Ti_bound_right, Te_bound_right, Te, Ti, psi, ne, normalize_to_nbar, nbar, ne_is_fGW, ne_bound_right, ne_bound_right_is_fGW, ne_bound_right_is_absolute, nu, initial_j_is_total_current, initial_psi_from_j)[source]

Bases: Mapping

Prescribed values and boundary conditions for the core profiles.

Parameters:

  • Ip_tot (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • vloop_lcfs (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • Ti_bound_right (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • Te_bound_right (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • Te (Union[Float[Array, 'rhon'], Float[ndarray, 'rhon']])

  • Ti (Union[Float[Array, 'rhon'], Float[ndarray, 'rhon']])

  • psi (Union[Float[Array, 'rhon'], Float[ndarray, 'rhon'], None])

  • ne (Union[Float[Array, 'rhon'], Float[ndarray, 'rhon']])

  • normalize_to_nbar (bool)

  • nbar (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • ne_is_fGW (bool)

  • ne_bound_right (Union[Float[Array, ''], Float[ndarray, ''], number, float])

  • ne_bound_right_is_fGW (bool)

  • ne_bound_right_is_absolute (bool)

  • nu (float)

  • initial_j_is_total_current (bool)

  • initial_psi_from_j (bool)

items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
values() → an object providing a view on D's values
class torax.config.profile_conditions.ProfileConditions(**data)[source]

Bases: BaseModelFrozen

Generic numeric parameters for the simulation.

The from_dict(…) method can accept a dictionary defined by https://torax.readthedocs.io/en/latest/configuration.html#profile-conditions.

Ip_tot

Total plasma current in MA. Note that if Ip_from_parameters=False in geometry, then this Ip will be overwritten by values from the geometry data. If use_vloop_lcfs_boundary_condition, only used as an initial condition.

use_vloop_lcfs_boundary_condition

Boundary condition at LCFS for Vloop ( = dspsi_lcfs/dt ). If use_vloop_lcfs_boundary_condition is True, then the specfied Vloop at the LCFS is used as the boundary condition for the psi equation; otherwise, Ip is used as the boundary condition.

vloop_lcfs

Boundary condition at LCFS for Vloop ( = dpsi_lcfs/dt ).

Ti_bound_right

Temperature boundary conditions at r=Rmin. If this is None the boundary condition will instead be taken from Ti and Te at rhon=1.

Te_bound_right

Temperature boundary conditions at r=Rmin. If this is None the boundary condition will instead be taken from Ti and Te at rhon=1.

Ti

Prescribed or evolving values for temperature at different times.

Te

Prescribed or evolving values for temperature at different times.

psi

Initial values for psi. If provided, the initial psi will be taken from here. Otherwise, the initial psi will be calculated from either the geometry or the “nu formula” dependant on the initial_psi_from_j field.

ne

Prescribed or evolving values for electron density at different times.

normalize_to_nbar

Whether to renormalize the density profile to have the desired line averaged density nbar.

nbar

Line averaged density. In units of reference density if ne_is_fGW = False. In Greenwald fraction if ne_is_fGW = True. nGW = Ip/(pi*a^2) with a in m, nGW in 10^20 m-3, Ip in MA

ne_is_fGW

Toggle units of nbar

ne_bound_right

Density boundary condition for r=Rmin. In units of reference density if ne_bound_right_is_fGW = False. In Greenwald fraction if ne_bound_right_is_fGW = True. If ne_bound_right is None then the boundary condition will instead be taken from ne at rhon=1. In this case, ne_bound_right_is_absolute will be set to False and ne_bound_right_is_fGW will be set to ne_is_fGW. If ne_bound_right is not None then ne_bound_right_is_absolute will be set to True.

ne_bound_right_is_fGW

Toggle units of ne_bound_right.

ne_bound_right_is_absolute

Toggle units of ne_bound_right

nu

Peaking factor of “Ohmic” current: johm = j0*(1 - r^2/a^2)^nu

initial_j_is_total_current

Toggles if “Ohmic” current is treated as total current upon initialization, or if non-inductive current should be included in initial jtot calculation.

initial_psi_from_j

Toggles if the initial psi calculation is based on the “nu” current formula, or from the psi available in the numerical geometry file. This setting is ignored for the ad-hoc circular geometry, which has no numerical geometry.

Parameters:

data (Any)

build_dynamic_params(t)[source]

Builds a DynamicProfileConditions.

Parameters:

t (Union[Array, ndarray, bool, number, float, int])

Return type:

DynamicProfileConditions

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

torax.config.runtime_params module

General runtime input parameters used throughout TORAX simulations.

class torax.config.runtime_params.DynamicGeneralRuntimeParams(plasma_composition, profile_conditions, numerics)[source]

Bases: Mapping

General runtime input parameters for the torax module.

Parameters:
items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
values() → an object providing a view on D's values
class torax.config.runtime_params.GeneralRuntimeParams(**data)[source]

Bases: BaseModelFrozen

Pydantic model for runtime parameters.

The from_dict(…) method can accept a dictionary defined by https://torax.readthedocs.io/en/latest/configuration.html#runtime-params.

profile_conditions

Pydantic model for the profile conditions.

numerics

Pydantic model for the numerics.

plasma_composition

Pydantic model for the plasma composition.

output_dir

File directory where the simulation outputs will be saved. If not provided, this will default to /tmp/torax_results_<YYYYMMDD_HHMMSS>/.

Parameters:

data (Any)

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

torax.config.runtime_params_slice module

Inputs to TORAX steppers and functions based on the input runtime parameters.

When running a TORAX simulation, the stepper is (by default) a JAX-compiled function, meaning it has two types of arguments: “dynamic” and “static”.

The “dynamic” arguments can change from call to call. These arguments must be arrays, scalars, or standard (possibly nested) Python containers. See the JAX docs for more info on allowed types. They cannot influence the logical branches the SimulationStepFn may take (again, see the sharp bits in the JAX docs to learn more about the how these “dynamic” args can be used within the function).

Note that the “dynamic” arguments are NOT necessarily time-dependent. They do not need to vary from time step to time step (though they can). They can change from time step to time step, or from simulation run to simulation run, without triggering a recompile. Changing these params without needing to recompile the stepper is the defining quality of the dynamic arguments.

The “static” arguments are compile-time constant. Any changes to them would trigger a recompilation of the stepper. These arguments don’t have the same restrictions as the dynamic arguments both in terms of types and how they are used.

class torax.config.runtime_params_slice.DynamicRuntimeParamsSlice(transport, stepper, plasma_composition, profile_conditions, numerics, sources, pedestal, mhd)[source]

Bases: Mapping

Input params that are ok to use as inputs to a JAX-compiled function.

This PyTree of params is input to the sim.SimulationStepFn, which updates the joint state and evolves the mesh state. This config includes various “dynamic” parameters which can change from step to step, or from simulation run to simulation run, without requiring components in the SimulationStepFn to recompile.

Note that “dynamic” does NOT mean time dependent necessarily (though these params can be time dependent). Here “dynamic” means these params can change without trigerring or requiring a recompile.

While the parameters are not necessarily time-dependent, that is how the class gets its name: a config “slice” refers to a subset of the overall TORAX config at a specific time t.

This class contains “slices” of various RuntimeParams attributes defined throughout TORAX:

  • from the profile_conditions runtime params

  • from the numerics runtime params

  • from the plasma_composition runtime params

  • from the transport model’s runtime params

  • from the stepper’s runtime params

  • from each of the sources’ runtime params

  • from the pedestal model’s runtime params

  • from each of the mhd models’ runtime params

This class packages all these together for convenience, as it simplifies many of the internal APIs within TORAX.

Parameters:
items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
values() → an object providing a view on D's values
class torax.config.runtime_params_slice.StaticRuntimeParamsSlice(stepper, sources, torax_mesh, ion_heat_eq, el_heat_eq, current_eq, dens_eq, main_ion_names, impurity_names, use_vloop_lcfs_boundary_condition, adaptive_dt)[source]

Bases: Mapping

Static arguments to SimulationStepFn which cannot be changed.

If any changes are made to these arguments, then components in SimulationStepFn must be recompiled.

NOTE: These are not the only parameters which can trigger recompilations! For instance, if the geometry changes its shape (i.e. nr or hires_fac change), that can also trigger a recompile. This is just to note that this list is not an exhaustive list of what can cause recompilations.

TODO(b/335596447): Add function to help users detect whether their change in config will trigger a recompile.

Parameters:
items() → a set-like object providing a view on D's items
keys() → a set-like object providing a view on D's keys
validate_new(new_params)[source]

Validates that the new static runtime params slice is compatible.

Parameters:

new_params (Self)

values() → an object providing a view on D's values
torax.config.runtime_params_slice.make_ip_consistent(dynamic_runtime_params_slice, geo)[source]

Fixes Ip to be the same across dynamic_runtime_params_slice and geo.

Parameters:
Return type:

tuple[DynamicRuntimeParamsSlice, Geometry]

torax.config.runtime_validation_utils module

Utilities for validating the config inputs.

torax.config.runtime_validation_utils.time_varying_array_bounded(time_varying_array, lower_bound=-inf, upper_bound=inf)[source]

Validates the input for the TimeVaryingArray.

Parameters:

  • time_varying_array (TimeVaryingArray)

  • lower_bound (float)

  • upper_bound (float)

Return type:

TimeVaryingArray

torax.config.runtime_validation_utils.time_varying_array_defined_at_1(time_varying_array)[source]

Validates the input for the TimeVaryingArray.

Parameters:

time_varying_array (TimeVaryingArray)

Return type:

TimeVaryingArray

Module contents

Config.

This module contains functionality related to parsing user-facing config files and setting up TORAX dynamic and static runtime parameters.