torax.time_step_calculator package

Subpackages

Submodules

torax.time_step_calculator.chi_time_step_calculator module

The ChiTimeStepCalculator class.

Steps through time using a heuristic based on chi_max.

class torax.time_step_calculator.chi_time_step_calculator.ChiTimeStepCalculator(tolerance=1e-07)[source]

Bases: TimeStepCalculator

TimeStepCalculator based on chi_max heuristic.

config

General configuration parameters.

Parameters:

tolerance (float)

next_dt(dynamic_runtime_params_slice, geo, core_profiles, core_transport)[source]

Calculates the next time step duration.

This calculation is a heuristic scaling of the maximum stable step size for the explicit method, and is therefore a function of chi_max.

Parameters:

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

  • geo (Geometry) – Geometry for the tokamak being simulated.

  • core_profiles (CoreProfiles) – Current core plasma profiles.

  • core_transport (CoreTransport) – Used to calculate maximum step size.

Returns:

Scalar time step duration.

Return type:

dt

torax.time_step_calculator.fixed_time_step_calculator module

The FixedTimeStepCalculator class.

Steps through time using a constant time step.

class torax.time_step_calculator.fixed_time_step_calculator.FixedTimeStepCalculator(tolerance=1e-07)[source]

Bases: TimeStepCalculator

TimeStepCalculator based on constant time steps.

config

General configuration parameters.

Parameters:

tolerance (float)

next_dt(dynamic_runtime_params_slice, geo, core_profiles, core_transport)[source]

Calculates the next time step duration.

Parameters:

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

  • geo (Geometry) – Geometry for the tokamak being simulated.

  • core_profiles (CoreProfiles) – Current core plasma profiles.

  • core_transport (CoreTransport) – Used to calculate chi, which determines maximum step size.

Returns:

Scalar time step duration.

Return type:

dt

torax.time_step_calculator.pydantic_model module

Pydantic config for time step calculators.

class torax.time_step_calculator.pydantic_model.TimeStepCalculator(**data)[source]

Bases: BaseModelFrozen

Config for a time step calculator.

calculator_type

The type of time step calculator to use.

tolerance

The tolerance within the final time for which the simulation will be considered done.

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].

class torax.time_step_calculator.pydantic_model.TimeStepCalculatorType(value)[source]

Bases: Enum

Types of time step calculators.

torax.time_step_calculator.time_step_calculator module

The TimeStepCalculator class.

Abstract base class defining time stepping interface.

class torax.time_step_calculator.time_step_calculator.TimeStepCalculator(tolerance=1e-07)[source]

Bases: ABC

Iterates over time during simulation.

Usage follows this pattern:

Parameters:

tolerance (float)

abstract next_dt(dynamic_runtime_params_slice, geo, core_profiles, core_transport)[source]

Returns the next time step duration and internal time stepper state.

Parameters:

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

  • geo (Geometry) – Geometry for the Tokamak.

  • core_profiles (CoreProfiles) – Core plasma profiles in the tokamak.

  • core_transport (CoreTransport) – Transport coefficients.

Return type:

Array

Module contents

Time step calculator.

This module contains functionality related to calculating dt before each stepper step.