torax.stepper package

Subpackages

Submodules

torax.stepper.linear_theta_method module

The LinearThetaMethodStepper class.

class torax.stepper.linear_theta_method.LinearThetaMethod(transport_model, source_models, pedestal_model)[source]

Bases: Stepper

Time step update using theta method, linearized on coefficients at t.

Parameters:

torax.stepper.nonlinear_theta_method module

The NonLinearThetaMethod class.

class torax.stepper.nonlinear_theta_method.DynamicNewtonRaphsonRuntimeParams(chi_per, d_per, corrector_steps, log_iterations, initial_guess_mode, maxiter, tol, coarse_tol, delta_reduction_factor, tau_min)[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.stepper.nonlinear_theta_method.DynamicOptimizerRuntimeParams(chi_per, d_per, corrector_steps, initial_guess_mode, maxiter, tol)[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.stepper.nonlinear_theta_method.NewtonRaphsonThetaMethod(transport_model, source_models, pedestal_model)[source]

Bases: NonlinearThetaMethod

Nonlinear theta method using Newton Raphson.

transport_model

A TransportModel subclass, calculates transport coeffs.

callback_class

Which class should be used to calculate the coefficients.

Parameters:
class torax.stepper.nonlinear_theta_method.NonlinearThetaMethod(transport_model, source_models, pedestal_model)[source]

Bases: Stepper

Time step update using theta method.

transport_model

A TransportModel subclass, calculates transport coeffs.

source_models

All TORAX sources used to compute both the explicit and implicit source profiles used for each time step as terms in the state evolution equations. Though the explicit profiles are computed outside the call to Stepper, the same sources should be used to compute those. The Sources are exposed here to provide a single source of truth for which sources are used during a run.

Parameters:
class torax.stepper.nonlinear_theta_method.OptimizerThetaMethod(transport_model, source_models, pedestal_model)[source]

Bases: NonlinearThetaMethod

Minimize the squared norm of the residual of the theta method equation.

transport_model

A TransportModel subclass, calculates transport coeffs.

callback_class

Which class should be used to calculate the coefficients.

initial_guess_mode

Passed through to fvm.optimizer_solve_block.

maxiter

Passed through to jaxopt.LBFGS.

tol

Passed through to jaxopt.LBFGS.

Parameters:

torax.stepper.predictor_corrector_method module

Carries out the predictor corrector method for the PDE solution.

Picard iterations to approximate the nonlinear solution. If static_runtime_params_slice.stepper.predictor_corrector is False, reverts to a standard linear solution.

torax.stepper.predictor_corrector_method.predictor_corrector_method(dt, static_runtime_params_slice, dynamic_runtime_params_slice_t_plus_dt, geo_t_plus_dt, x_old, x_new_guess, core_profiles_t_plus_dt, coeffs_exp, coeffs_callback)[source]

Predictor-corrector method.

Parameters:

  • dt (Array) – current timestep

  • static_runtime_params_slice (StaticRuntimeParamsSlice) – General input parameters which are fixed through a simulation run, and if changed, would trigger a recompile.

  • dynamic_runtime_params_slice_t_plus_dt (DynamicRuntimeParamsSlice) – Dynamic runtime parameters corresponding to the next time step, needed for the implicit PDE coefficients.

  • geo_t_plus_dt (Geometry) – Geometry at the next time step.

  • x_old (tuple[CellVariable, ...]) – Tuple of CellVariables correspond to the evolving core profiles at time t.

  • x_new_guess (tuple[CellVariable, ...]) – Tuple of CellVariables corresponding to the initial guess for the next time step.

  • core_profiles_t_plus_dt (CoreProfiles) – Core profiles at the next time step.

  • coeffs_exp (Block1DCoeffs) – Block1DCoeffs PDE coefficients at beginning of timestep

  • coeffs_callback (CoeffsCallback) – coefficient callback function

Returns:

Solution of evolving core profile state variables auxiliary_outputs: Extra outputs containing auxiliary information from the coeffs_callback computed based on x_new.

Return type:

x_new

torax.stepper.pydantic_model module

Pydantic config for Stepper.

class torax.stepper.pydantic_model.BaseStepper(**data)[source]

Bases: BaseModelFrozen, ABC

Base class for stepper configs.

theta_imp

The theta value in the theta method 0 = explicit, 1 = fully implicit, 0.5 = Crank-Nicolson.

predictor_corrector

Enables predictor_corrector iterations with the linear solver. If False, compilation is faster.

corrector_steps

The number of corrector steps for the predictor-corrector linear solver. 0 means a pure linear solve with no corrector steps.

convection_dirichlet_mode

See fvm.convection_terms docstring, dirichlet_mode argument.

convection_neumann_mode

See fvm.convection_terms docstring, neumann_mode argument.

use_pereverzev

Use pereverzev terms for linear solver. Is only applied in the nonlinear solver for the optional initial guess from the linear solver

chi_per

(deliberately) large heat conductivity for Pereverzev rule.

d_per

(deliberately) large particle diffusion for Pereverzev rule.

Parameters:

data (Any)

abstract property build_dynamic_params: DynamicRuntimeParams

Builds dynamic runtime params from the config.

build_static_params()[source]

Builds static runtime params from the config.

Return type:

StaticRuntimeParams

abstract build_stepper(transport_model, source_models, pedestal_model)[source]

Builds a stepper from the config.

Parameters:
Return type:

Stepper

abstract property linear_solver: bool

Returns True if the stepper is a linear solver.

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

Bases: BaseStepper

Model for the linear stepper.

stepper_type

The type of stepper to use, hardcoded to ‘linear’.

Parameters:

data (Any)

property build_dynamic_params: DynamicRuntimeParams

Builds dynamic runtime params from the config.

build_stepper(transport_model, source_models, pedestal_model)[source]

Builds a stepper from the config.

Parameters:
Return type:

Stepper

property linear_solver: bool

Returns True if the stepper is a linear solver.

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

Bases: BaseStepper

Model for nonlinear Newton-Raphson stepper.

stepper_type

The type of stepper to use, hardcoded to ‘newton_raphson’.

log_iterations

If True, log internal iterations in Newton-Raphson solver.

initial_guess_mode

The initial guess mode for the Newton-Raphson solver.

maxiter

The maximum number of iterations for the Newton-Raphson solver.

tol

The tolerance for the Newton-Raphson solver.

coarse_tol

The coarse tolerance for the Newton-Raphson solver.

delta_reduction_factor

The delta reduction factor for the Newton-Raphson solver.

tau_min

The minimum value of tau for the Newton-Raphson solver.

Parameters:

data (Any)

property build_dynamic_params: DynamicNewtonRaphsonRuntimeParams

Builds dynamic runtime params from the config.

build_stepper(transport_model, source_models, pedestal_model)[source]

Builds a stepper from the config.

Parameters:
Return type:

NewtonRaphsonThetaMethod

property linear_solver: bool

Returns True if the stepper is a linear solver.

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

Bases: BaseStepper

Model for nonlinear OptimizerThetaMethod stepper.

stepper_type

The type of stepper to use, hardcoded to ‘optimizer’.

initial_guess_mode

The initial guess mode for the optimizer.

maxiter

The maximum number of iterations for the optimizer.

tol

The tolerance for the optimizer.

Parameters:

data (Any)

property build_dynamic_params: DynamicOptimizerRuntimeParams

Builds dynamic runtime params from the config.

build_stepper(transport_model, source_models, pedestal_model)[source]

Builds a stepper from the config.

Parameters:
Return type:

OptimizerThetaMethod

property linear_solver: bool

Returns True if the stepper is a linear solver.

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.stepper.runtime_params module

Runtime params for the stepper.

class torax.stepper.runtime_params.DynamicRuntimeParams(chi_per, d_per, corrector_steps)[source]

Bases: Mapping

Input params for the stepper 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
class torax.stepper.runtime_params.StaticRuntimeParams(theta_imp, convection_dirichlet_mode, convection_neumann_mode, use_pereverzev, predictor_corrector)[source]

Bases: Mapping

Static params for the stepper.

Parameters:

  • theta_imp (float)

  • convection_dirichlet_mode (str)

  • convection_neumann_mode (str)

  • use_pereverzev (bool)

  • predictor_corrector (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

torax.stepper.stepper module

The Stepper class.

Abstract base class defining updates to State.

class torax.stepper.stepper.Stepper(transport_model, source_models, pedestal_model)[source]

Bases: ABC

Calculates a single time step’s update to State.

transport_model

A TransportModel subclass, calculates transport coeffs.

source_models

All TORAX sources used to compute both the explicit and implicit source profiles used for each time step as terms in the state evolution equations. Though the explicit profiles are computed outside the call to Stepper, the same sources should be used to compute those. The Sources are exposed here to provide a single source of truth for which sources are used during a run.

pedestal_model

A PedestalModel subclass, calculates pedestal values.

Parameters:

Module contents

Stepper.

This module is the PDE time stepper. It discretizes the PDE in time and solves for the next time step with linear or nonlinear methods.