torax.transport_model package

Subpackages

Submodules

torax.transport_model.base_qlknn_model module

Base class for QLKNN Models.

class torax.transport_model.base_qlknn_model.BaseQLKNNModel(path, name)[source]

Bases: ABC

Base class for QLKNN Models.

Parameters:

torax.transport_model.bohm_gyrobohm module

The BohmGyroBohmModel class.

class torax.transport_model.bohm_gyrobohm.BohmGyroBohmTransportModel[source]

Bases: TransportModel

Calculates various coefficients related to particle transport according to the Bohm + gyro-Bohm Model.

class torax.transport_model.bohm_gyrobohm.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere, chi_e_bohm_coeff, chi_e_gyrobohm_coeff, chi_i_bohm_coeff, chi_i_gyrobohm_coeff, d_face_c1, d_face_c2, v_face_coeff, chi_e_bohm_multiplier, chi_e_gyrobohm_multiplier, chi_i_bohm_multiplier, chi_i_gyrobohm_multiplier)[source]

Bases: DynamicRuntimeParams, Mapping

Dynamic runtime params for the BgB transport model.

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

torax.transport_model.constant module

The ConstantTransportModel class.

A simple model assuming constant transport.

class torax.transport_model.constant.ConstantTransportModel[source]

Bases: TransportModel

Calculates various coefficients related to particle transport.

class torax.transport_model.constant.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere, chii_const, chie_const, De_const, Ve_const)[source]

Bases: DynamicRuntimeParams, Mapping

Extends the base runtime params with additional params for this model.

See base class runtime_params.DynamicRuntimeParams docstring for more info.

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

torax.transport_model.critical_gradient module

The CriticalGradientModel class.

class torax.transport_model.critical_gradient.CriticalGradientTransportModel[source]

Bases: TransportModel

Calculates various coefficients related to particle transport.

class torax.transport_model.critical_gradient.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere, alpha, chistiff, chiei_ratio, chi_D_ratio, VR_D_ratio)[source]

Bases: DynamicRuntimeParams, Mapping

Dynamic runtime params for the CGM transport model.

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

torax.transport_model.pydantic_model module

Pydantic config for Transport models.

class torax.transport_model.pydantic_model.BohmGyroBohmTransportModel(**data)[source]

Bases: TransportBase

Model for the Bohm + Gyro-Bohm transport model.

transport_model

The transport model to use. Hardcoded to ‘bohm-gyrobohm’.

chi_e_bohm_coeff

Prefactor for Bohm term for electron heat conductivity.

chi_e_gyrobohm_coeff

Prefactor for GyroBohm term for electron heat conductivity.

chi_i_bohm_coeff

Prefactor for Bohm term for ion heat conductivity.

chi_i_gyrobohm_coeff

Prefactor for GyroBohm term for ion heat conductivity.

chi_e_bohm_multiplier

Multiplier for chi_e_bohm_coeff. Intended for user-friendly default modification.

chi_e_gyrobohm_multiplier

Multiplier for chi_e_gyrobohm_coeff. Intended for user-friendly default modification.

chi_i_bohm_multiplier

Multiplier for chi_i_bohm_coeff. Intended for user-friendly default modification.

chi_i_gyrobohm_multiplier

Multiplier for chi_i_gyrobohm_coeff. Intended for user-friendly default modification.

d_face_c1

Constant for the electron diffusivity weighting factor.

d_face_c2

Constant for the electron diffusivity weighting factor.

v_face_coeff

Proportionality factor between convectivity and diffusivity.

Parameters:

data (Any)

build_transport_model()[source]

Builds a transport model from the config.

Return type:

BohmGyroBohmTransportModel

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.transport_model.pydantic_model.ConstantTransportModel(**data)[source]

Bases: TransportBase

Model for the Constant transport model.

transport_model

The transport model to use. Hardcoded to ‘constant’.

chii_const

coefficient in ion heat equation diffusion term in m^2/s.

chie_const

coefficient in electron heat equation diffusion term in m^2/s.

De_const

diffusion coefficient in electron density equation in m^2/s.

Ve_const

convection coefficient in electron density equation in m^2/s.

Parameters:

data (Any)

build_transport_model()[source]

Builds a transport model from the config.

Return type:

ConstantTransportModel

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.transport_model.pydantic_model.CriticalGradientTransportModel(**data)[source]

Bases: TransportBase

Model for the Critical Gradient transport model.

transport_model

The transport model to use. Hardcoded to ‘CGM’.

alpha

Exponent of chi power law: chi ∝ (R/LTi - R/LTi_crit)^alpha.

chistiff

Stiffness parameter.

chiei_ratio

Ratio of electron to ion heat transport coefficient (ion higher for ITG).

chi_D_ratio

Ratio of electron particle to ion heat transport coefficient.

VR_D_ratio

Ratio of major radius * electron particle convection, to electron diffusion. Sets the value of electron particle convection in the model.

Parameters:

data (Any)

build_transport_model()[source]

Builds a transport model from the config.

Return type:

CriticalGradientTransportModel

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.transport_model.pydantic_model.QLKNNTransportModel(**data)[source]

Bases: TransportBase

Model for the QLKNN transport model.

To determine which model to load, TORAX uses the following logic:

  • If model_path is provided, then we load the model from this path.

  • Otherwise, if the TORAX_QLKNN_MODEL_PATH environment variable is set, then we load the model from this path.

  • Otherwise, if model_name is provided, we load that model from registered models in the fusion_surrogates library.

  • If model_name is not set either, we load the default QLKNN model from fusion_surrogates (currently QLKNN_7_11).

It is recommended to not set model_name, TORAX_QLKNN_MODEL_PATH or model_path to use the default QLKNN model.

transport_model

The transport model to use. Hardcoded to ‘qlknn’.

model_path

Path to the model. Takes precedence over model_name and TORAX_QLKNN_MODEL_PATH.

model_name

Name of the model to use. Used to select a model from the fusion_surrogates library.

include_ITG

Whether to include ITG modes.

include_TEM

Whether to include TEM modes.

include_ETG

Whether to include ETG modes.

ITG_flux_ratio_correction

Correction factor for ITG electron heat flux.

ETG_correction_factor

Correction factor for ETG electron heat flux. https://gitlab.com/qualikiz-group/QuaLiKiz/-/commit/5bcd3161c1b08e0272ab3c9412fec7f9345a2eef

clip_inputs

Whether to clip inputs within desired margin of the QLKNN training set boundaries.

clip_margin

Margin to clip inputs within desired margin of the QLKNN training set boundaries.

coll_mult

Collisionality multiplier.

avoid_big_negative_s

Ensure that smag - alpha > -0.2 always, to compensate for no slab modes.

smag_alpha_correction

Reduce magnetic shear by 0.5*alpha to capture main impact of alpha.

q_sawtooth_proxy

If q < 1, modify input q and smag as if q~1 as if there are sawteeth.

DVeff

Effective D / effective V approach for particle transport.

An_min

Minimum |R/Lne| below which effective V is used instead of effective D.

Parameters:

data (Any)

build_transport_model()[source]

Builds a transport model from the config.

Return type:

QLKNNTransportModel

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.transport_model.pydantic_model_base module

Base pydantic config for Transport models.

class torax.transport_model.pydantic_model_base.TransportBase(**data)[source]

Bases: BaseModelFrozen, ABC

Base model holding parameters common to all transport models.

chimin

Lower bound on heat conductivity.

chimax

Upper bound on heat conductivity (can be helpful for stability).

Demin

minimum electron density diffusivity.

Demax

maximum electron density diffusivity.

Vemin

minimum electron density convection.

Vemax

minimum electron density convection.

apply_inner_patch

set inner core transport coefficients (ad-hoc MHD/EM transport).

De_inner

inner core electron density diffusivity.

Ve_inner

inner core electron density convection.

chii_inner

inner core ion heat equation diffusion term.

chie_inner

inner core electron heat equation diffusion term.

rho_inner

normalized radius below which inner patch is applied.

apply_outer_patch

set outer core transport coefficients (ad-hoc MHD/EM transport). Only used when pedestal.set_pedestal = False Useful for L-mode near-edge region where QLKNN10D is not applicable.

De_outer

outer core electron density diffusivity.

Ve_outer

outer core electron density convection.

chii_outer

outer core ion heat equation diffusion term.

chie_outer

outer core electron heat equation diffusion term.

rho_outer

normalized radius above which outer patch is applied.

smoothing_sigma

Width of HWHM Gaussian smoothing kernel operating on transport model outputs.

smooth_everywhere

Smooth across entire radial domain regardless of inner and outer patches.

Parameters:

data (Any)

abstract build_transport_model()[source]

Builds a transport model from the config.

Return type:

TransportModel

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.transport_model.qlknn_10d module

Class for handling QLKNN10D models.

class torax.transport_model.qlknn_10d.MLP(hidden_sizes, activations, parent=<flax.linen.module._Sentinel object>, name=None)[source]

Bases: Module

Parameters:
class torax.transport_model.qlknn_10d.QLKNN10D(path, name)[source]

Bases: BaseQLKNNModel

Class holding QLKNN10D networks.

path

Path to qlknn-hyper

net_itgleading

ITG Qi net

net_itgqediv

ITG Qe/Qi net

net_temleading

TEM Qe net

net_temqediv

TEM Qi/Qe net

net_etgleading

ETG Qe net

net_temqidiv

TEM Qi/Qe net

net_tempfediv

Tem pfe/Qe net

net_etgleading

ITG Qe/Qi net

net_itgpfediv

ITG pfe/Qi net

Parameters:
get_model_inputs_from_qualikiz_inputs(qualikiz_inputs)[source]

Converts QualikizInputs to model inputs.

Parameters:

qualikiz_inputs (QualikizInputs)

Return type:

Array

predict(inputs)[source]

Feed forward through the network and compute fluxes.

Parameters:

inputs (Array)

Return type:

dict[str, Array]

class torax.transport_model.qlknn_10d.QuaLiKizNDNN(model_config)[source]

Bases: object

Jax implementation for QLKNN10D inference.

Parameters:

model_config (dict[str, Any])

torax.transport_model.qlknn_model_wrapper module

A wrapper for QLKNN transport surrogate models.

class torax.transport_model.qlknn_model_wrapper.QLKNNModelWrapper(path, name='', flux_name_map=None)[source]

Bases: BaseQLKNNModel

A TORAX wrapper for a QLKNN Model from the fusion_surrogates library.

Parameters:
get_model_inputs_from_qualikiz_inputs(qualikiz_inputs)[source]

Converts QualikizInputs to model inputs.

Parameters:

qualikiz_inputs (QualikizInputs)

Return type:

Array

predict(inputs)[source]

Predicts the fluxes given the inputs.

Parameters:

inputs (Array)

Return type:

dict[str, Array]

torax.transport_model.qlknn_transport_model module

A transport model that uses a QLKNN model.

class torax.transport_model.qlknn_transport_model.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere, DVeff, An_min, coll_mult, avoid_big_negative_s, smag_alpha_correction, q_sawtooth_proxy, include_ITG, include_TEM, include_ETG, ITG_flux_ratio_correction, ETG_correction_factor, clip_inputs, clip_margin)[source]

Bases: DynamicRuntimeParams, 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.transport_model.qlknn_transport_model.QLKNNRuntimeConfigInputs(nref, Zeff_face, transport, Ped_top, set_pedestal)[source]

Bases: Mapping

Runtime config inputs for QLKNN.

The runtime DynamicRuntimeParamsSlice contains global runtime parameters, not all of which are cacheable. This set of inputs IS cacheable, and using this added layer allows the global config to change without affecting how QLKNNTransportModel works.

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.transport_model.qlknn_transport_model.QLKNNTransportModel(path, name)[source]

Bases: QualikizBasedTransportModel

Calculates turbulent transport coefficients.

Parameters:
torax.transport_model.qlknn_transport_model.clip_inputs(feature_scan, clip_margin, inputs_and_ranges)[source]

Clip input values according to the training set limits + optional user-defined margin for qlknn.

Parameters:
Return type:

Array

torax.transport_model.qlknn_transport_model.get_model(path, name)[source]

Load the model.

Parameters:
Return type:

BaseQLKNNModel

torax.transport_model.qualikiz_based_transport_model module

Base class and utils for Qualikiz-based models.

class torax.transport_model.qualikiz_based_transport_model.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere, DVeff, An_min, coll_mult, avoid_big_negative_s, smag_alpha_correction, q_sawtooth_proxy)[source]

Bases: DynamicRuntimeParams, Mapping

Shared parameters for Qualikiz-based models.

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.transport_model.qualikiz_based_transport_model.QualikizBasedTransportModel[source]

Bases: QuasilinearTransportModel

Base class for Qualikiz-based transport models.

class torax.transport_model.qualikiz_based_transport_model.QualikizInputs(chiGB, Rmin, Rmaj, lref_over_lti, lref_over_lte, lref_over_lne, lref_over_lni0, lref_over_lni1, Zeff_face, q, smag, x, Ti_Te, log_nu_star_face, normni, alpha, epsilon_lcfs)[source]

Bases: QuasilinearInputs, Mapping

Inputs to Qualikiz-based models.

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

torax.transport_model.qualikiz_transport_model module

torax.transport_model.quasilinear_transport_model module

Base class for quasilinear models.

class torax.transport_model.quasilinear_transport_model.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere, DVeff, An_min)[source]

Bases: DynamicRuntimeParams, Mapping

Shared parameters for Quasilinear models.

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.transport_model.quasilinear_transport_model.NormalizedLogarithmicGradients(lref_over_lti, lref_over_lte, lref_over_lne, lref_over_lni0, lref_over_lni1)[source]

Bases: Mapping

Normalized logarithmic gradients of plasma profiles.

Defined as Lref/Lprofile. Lref is an arbitrary reference length [m]. lprofile is each profile gradient length [m] defined as -1/grad(log(profile)), e.g. lti = -1/grad(log(ti)), i.e. lti = - ti / (dti/dr). The specific radial coordinate r used for the gradient is a user input.

Parameters:
classmethod from_profiles(core_profiles, radial_coordinate, reference_length)[source]

Calculates the normalized logarithmic gradients.

Parameters:

  • core_profiles (CoreProfiles)

  • radial_coordinate (Array)

  • reference_length (Array)

Return type:

Self

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.transport_model.quasilinear_transport_model.QuasilinearInputs(chiGB, Rmin, Rmaj, lref_over_lti, lref_over_lte, lref_over_lne, lref_over_lni0, lref_over_lni1)[source]

Bases: Mapping

Variables required to convert outputs to TORAX CoreTransport outputs.

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.transport_model.quasilinear_transport_model.QuasilinearTransportModel[source]

Bases: TransportModel

Base class for quasilinear models.

torax.transport_model.quasilinear_transport_model.calculate_alpha(core_profiles, nref, q, reference_magnetic_field, normalized_logarithmic_gradients)[source]

Calculates the alpha_MHD parameter.

alpha_MHD = Lref q^2 beta’ , where beta’ is the radial gradient of beta, the ratio of plasma pressure to magnetic pressure, Lref a reference length, and q is the safety factor. Lref is included within the NormalizedLogarithmicGradients.

Parameters:
Return type:

Union[Array, ndarray, bool, number]

Returns:

Alpha value as a chex.Array.

torax.transport_model.quasilinear_transport_model.calculate_chiGB(reference_temperature, reference_magnetic_field, reference_mass, reference_length)[source]

Calculates the gyrobohm diffusivity.

Different transport models make different choices for the reference temperature, magnetic field, and mass used for gyrobohm normalization.

Parameters:
Return type:

Union[Array, ndarray, bool, number]

Returns:

Gyrobohm diffusivity as a chex.Array [dimensionless].

torax.transport_model.quasilinear_transport_model.calculate_normalized_logarithmic_gradient(var, radial_coordinate, reference_length)[source]

Calculates the normalized logarithmic gradient of a CellVariable on the face grid.

Parameters:

  • var (CellVariable)

  • radial_coordinate (Array)

  • reference_length (Array)

Return type:

Array

torax.transport_model.runtime_params module

Dataclass representing runtime parameter inputs to the transport models.

This is the dataclass runtime config exposed to the user. The actual model gets a time-interpolated version of this config via the DynamicRuntimeParams.

class torax.transport_model.runtime_params.DynamicRuntimeParams(chimin, chimax, Demin, Demax, Vemin, Vemax, apply_inner_patch, De_inner, Ve_inner, chii_inner, chie_inner, rho_inner, apply_outer_patch, De_outer, Ve_outer, chii_outer, chie_outer, rho_outer, smoothing_sigma, smooth_everywhere)[source]

Bases: Mapping

Input params for the transport model which can be used as compiled args.

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

torax.transport_model.transport_model module

The TransportModel abstract base class.

The transport model calculates heat and particle turbulent transport coefficients.

class torax.transport_model.transport_model.TransportModel[source]

Bases: ABC

Calculates various coefficients related to heat and particle transport.

Subclass responsbilities: - Must implement __hash__, __eq__, and be immutable, so that the class can

be used as a static argument (or a subcomponent of a larger static argument) to jax.jit

  • Must set _frozen = True at the end of the subclass __init__ method to activate immutability.

torax.transport_model.transport_model.build_smoothing_matrix(geo, dynamic_runtime_params_slice, pedestal_model_outputs)[source]

Builds a smoothing matrix for the turbulent transport model.

Uses a Gaussian kernel of HWHM defined in the transport config.

Parameters:

  • geo (Geometry) – Geometry of the torus.

  • dynamic_runtime_params_slice (DynamicRuntimeParamsSlice) – Input runtime parameters that can change without triggering a JAX recompilation.

  • pedestal_model_outputs (PedestalModelOutput) – Output of the pedestal model.

Returns:

A smoothing matrix for convolution with the transport outputs.

Return type:

kernel

Module contents

Transport model.

This module contains functions calculating turbulent transport coefficients.