torax.fvm package

Submodules

torax.fvm.block_1d_coeffs module

The Block1DCoeffs dataclass.

This is the key interface between the fvm module, which is abstracted to the level of a coupled 1D fluid dynamics PDE, and the rest of torax, which includes calculations specific to plasma physics to provide these coefficients.

class torax.fvm.block_1d_coeffs.Block1DCoeffs(transient_in_cell, transient_out_cell=None, d_face=None, v_face=None, source_mat_cell=None, source_cell=None, auxiliary_outputs=None)[source]

Bases: Mapping

The coefficients of coupled 1D fluid dynamics PDEs.

The differential equation is: transient_out_coeff partial x transient_in_coeff / partial t = F where F = divergence(diffusion_coeff * grad(x)) - divergence(convection_coeff * x) + source_mat_coeffs * u + sources.

source_mat_coeffs exists for specific classes of sources where this decomposition is valid, allowing x to be treated implicitly in linear solvers, even if source_mat_coeffs contains state-dependent terms

This class captures a snapshot of the coefficients of the equation at one instant in time, discretized spatially across a mesh.

This class imposes the following structure on the problem: - It assumes the variables are arranged on a 1-D, evenly spaced grid. - It assumes the x variable is broken up into “channels,” so the resulting matrix equation has one block per channel.

transient_out_cell

Tuple with one entry per channel, transient_out_cell[i] gives the transient coefficients outside the time derivative for channel i on the cell grid.

transient_in_cell

Tuple with one entry per channel, transient_in_cell[i] gives the transient coefficients inside the time derivative for channel i on the cell grid.

d_face

Tuple, with d_face[i] containing diffusion term coefficients for channel i on the face grid.

v_face

Tuple, with v_face[i] containing convection term coefficients for channel i on the face grid.

source_mat_cell

2-D matrix of Tuples, with source_mat_cell[i][j] adding to block-row i a term of the form source_cell[j] * u[channel j]. Depending on the source runtime_params, may be constant values for a timestep, or updated iteratively with new states in a nonlinear solver

source_cell

Additional source terms on the cell grid for each channel. Depending on the source runtime_params, may be constant values for a timestep, or updated iteratively with new states in a nonlinear solver

auxiliary_outputs

Optional extra output which can include auxiliary state or information useful for inspecting the computation inside the callback which calculated these coeffs.

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.fvm.calc_coeffs module

Calculates Block1DCoeffs for a time step.

class torax.fvm.calc_coeffs.CoeffsCallback(static_runtime_params_slice, transport_model, explicit_source_profiles, source_models, evolving_names, pedestal_model)[source]

Bases: object

Calculates Block1DCoeffs for a state.

static_runtime_params_slice

See the docstring for stepper.Stepper.

transport_model

See the docstring for stepper.Stepper.

explicit_source_profiles

See the docstring for stepper.Stepper.

source_models

See the docstring for stepper.Stepper.

evolving_names

The names of the evolving variables.

pedestal_model

See the docstring for stepper.Stepper.

Parameters:
torax.fvm.calc_coeffs.calc_coeffs(static_runtime_params_slice, dynamic_runtime_params_slice, geo, core_profiles, transport_model, explicit_source_profiles, source_models, pedestal_model, evolving_names, use_pereverzev=False, explicit_call=False)[source]

Calculates Block1DCoeffs for the time step described by core_profiles.

Parameters:

  • 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 (DynamicRuntimeParamsSlice) – General input parameters that can change from time step to time step or simulation run to run, and do so without triggering a recompile.

  • geo (Geometry) – Geometry describing the torus.

  • core_profiles (CoreProfiles) – Core plasma profiles for this time step during this iteration of the solver. Depending on the type of stepper being used, this may or may not be equal to the original plasma profiles at the beginning of the time step.

  • transport_model (TransportModel) – A TransportModel subclass, calculates transport coeffs.

  • explicit_source_profiles (SourceProfiles) – Precomputed explicit source profiles. These profiles either do not depend on the core profiles or depend on the original core profiles at the start of the time step, not the “live” updating core profiles. For sources that are implicit, their explicit profiles are set to all zeros.

  • source_models (SourceModels) – All TORAX source/sink functions that generate the explicit and implicit source profiles used as terms for the core profiles equations.

  • pedestal_model (PedestalModel) – A PedestalModel subclass, calculates pedestal values.

  • evolving_names (tuple[str, ...]) – The names of the evolving variables in the order that their coefficients should be written to coeffs.

  • use_pereverzev (bool) – Toggle whether to calculate Pereverzev terms

  • explicit_call (bool) – If True, indicates that calc_coeffs is being called for the explicit component of the PDE. Then calculates a reduced Block1DCoeffs if theta_imp=1. This saves computation for the default fully implicit implementation.

Returns:

Block1DCoeffs containing the coefficients at this time step.

Return type:

coeffs

torax.fvm.cell_variable module

The CellVariable class.

A chex dataclass used to represent variables on meshes for the 1D fvm solver. Naming conventions and API are similar to those developed in the FiPy fvm solver [https://www.ctcms.nist.gov/fipy/]

class torax.fvm.cell_variable.CellVariable(value, dr, left_face_constraint=None, right_face_constraint=None, left_face_grad_constraint=<factory>, right_face_grad_constraint=<factory>)[source]

Bases: Mapping

A variable representing values of the cells along the radius.

value

A jax.Array containing the value of this variable at each cell.

dr

Distance between cell centers.

left_face_constraint

An optional jax scalar specifying the value of the leftmost face. Defaults to None, signifying no constraint. The user can modify this field at any time, but when face_grad is called exactly one of left_face_constraint and left_face_grad_constraint must be None.

left_face_grad_constraint

An optional jax scalar specifying the (otherwise underdetermined) value of the leftmost face. See left_face_constraint.

right_face_constraint

Analogous to left_face_constraint but for the right face, see left_face_constraint.

right_face_grad_constraint

A jax scalar specifying the undetermined value of the gradient on the rightmost face variable.

Parameters:

  • value (Union[Float[Array, 't* cell'], Float[ndarray, 't* cell']])

  • dr (Union[Float[Array, 't*'], Float[ndarray, 't*']])

  • left_face_constraint (Union[Float[Array, 't*'], Float[ndarray, 't*'], None])

  • right_face_constraint (Union[Float[Array, 't*'], Float[ndarray, 't*'], None])

  • left_face_grad_constraint (Union[Float[Array, 't*'], Float[ndarray, 't*'], None])

  • right_face_grad_constraint (Union[Float[Array, 't*'], Float[ndarray, 't*'], None])

cell_plus_boundaries()[source]

Returns the value of this variable plus left and right boundaries.

Return type:

Float[Array, 't* cell+2']

face_grad(x=None)[source]

Returns the gradient of this value with respect to the faces.

Implemented using forward differencing of cells. Leftmost and rightmost gradient entries are determined by user specify constraints, see CellVariable class docstring.

Parameters:

x (Union[Float[Array, 'cell'], Float[ndarray, 'cell'], None]) – (optional) coordinates over which differentiation is carried out

Return type:

Union[Float[Array, 'face'], Float[ndarray, 'face']]

Returns:

A jax.Array of shape (num_faces,) containing the gradient.

face_value()[source]

Calculates values of this variable on the face grid.

Return type:

Float[Array, 't* face']

grad()[source]

Returns the gradient of this variable wrt cell centers.

Return type:

Float[Array, 't* face']

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.fvm.convection_terms module

The make_convection_terms function.

Builds the convection terms of the discrete matrix equation.

torax.fvm.convection_terms.make_convection_terms(v_face, d_face, var, dirichlet_mode='ghost', neumann_mode='ghost')[source]

Makes the terms of the matrix equation derived from the convection term.

The convection term of the differential equation is of the form - (partial / partial r) v u

Parameters:

  • v_face (Array) – Convection coefficient on faces.

  • d_face (Array) – Diffusion coefficient on faces. The relative strength of convection to diffusion is used to weight the contribution of neighboring cells when calculating face values of u.

  • var (CellVariable) – CellVariable to define mesh and boundary conditions.

  • dirichlet_mode (str) – The strategy to use to handle Dirichlet boundary conditions. The default is ‘ghost’, which has superior stability. ‘ghost’ -> Boundary face values are inferred by constructing a ghost cell then alpha weighting cells ‘direct’ -> Boundary face values are read directly from constraints ‘semi-implicit’ -> Matches FiPy. Boundary face values are alpha weighted with the constraint value specifying the value of the “other” cell: x_{boundary_face} = alpha x_{last_cell} + (1 - alpha) BC

  • neumann_mode (str) – Which strategy to use to handle Neumann boundary conditions. The default is ghost, which has superior stability. ‘ghost’ -> Boundary face values are inferred by constructing a ghost cell then alpha weighting cells. ‘semi-implicit’ -> Matches FiPy. Boundary face values are alpha weighted, with the (1 - alpha) weight applied to the external face value rather than to a ghost cell.

Returns:

Tridiagonal matrix of coefficients on u c: Vector of terms not dependent on u

Return type:

mat

torax.fvm.diffusion_terms module

The make_diffusion_terms function.

Builds the diffusion terms of the discrete matrix equation.

torax.fvm.diffusion_terms.make_diffusion_terms(d_face, var)[source]

Makes the terms of the matrix equation derived from the diffusion term.

The diffusion term is of the form (partial / partial x) D partial x / partial x

Parameters:
Returns:

Tridiagonal matrix of coefficients on u c: Vector of terms not dependent on u

Return type:

mat

torax.fvm.discrete_system module

Functionality for building discrete linear systems.

This file is expected to be used mostly internally by fvm itself.

The functionality here is for constructing a description of one discrete time step of a PDE in terms of a linear equation. In practice, the actual expressive power of the resulting Jax expression may still be nonlinear because the coefficients of this linear equation are Jax expressions, not just numeric values, so nonlinear solvers like newton_raphson_solve_block can capture nonlinear dynamics even when each step is expressed using a matrix multiply.

torax.fvm.discrete_system.calc_c(x, coeffs, convection_dirichlet_mode='ghost', convection_neumann_mode='ghost')[source]

Calculate C and c such that F = C x + c.

See docstrings for Block1DCoeff and implicit_solve_block for more detail.

Parameters:

  • x (tuple[CellVariable, ...]) – Tuple containing CellVariables for each channel. This function uses only their shape and their boundary conditions, not their values.

  • coeffs (Block1DCoeffs) – Coefficients defining the differential equation.

  • convection_dirichlet_mode (str) – See docstring of the convection_terms function, dirichlet_mode argument.

  • convection_neumann_mode (str) – See docstring of the convection_terms function, neumann_mode argument.

Returns:

matrix C, such that F = C x + c c: the vector c

Return type:

c_mat

torax.fvm.enums module

The enums module.

Enums shared through the fvm package.

class torax.fvm.enums.InitialGuessMode(value)[source]

Bases: Enum

Modes for initial guess of x_new for iterative solvers.

torax.fvm.fvm_conversions module

Conversions utilities for fvm objects.

torax.fvm.fvm_conversions.cell_variable_tuple_to_vec(x_tuple)[source]

Converts a tuple of CellVariables to a flat array.

Parameters:

x_tuple (tuple[CellVariable, ...]) – A tuple of CellVariables.

Return type:

Array

Returns:

A flat array of evolving state variables.

torax.fvm.fvm_conversions.vec_to_cell_variable_tuple(x_vec, core_profiles, evolving_names)[source]

Converts a flat array of core profile state vars to CellVariable tuple.

Parameters:

  • x_vec (Array) – A flat array of evolving core profile state variables. The order of the variables in the array must match the order of the evolving_names.

  • core_profiles (CoreProfiles) – CoreProfiles containing all CellVariables with appropriate boundary conditions.

  • evolving_names (tuple[str, ...]) – The names of the evolving cell variables.

Return type:

tuple[CellVariable, ...]

Returns:

A tuple of updated CellVariables.

torax.fvm.implicit_solve_block module

The implicit_solve_block function.

See function docstring for details.

torax.fvm.implicit_solve_block.implicit_solve_block(dt, x_old, x_new_guess, coeffs_old, coeffs_new, theta_imp=1.0, convection_dirichlet_mode='ghost', convection_neumann_mode='ghost')[source]

Runs one time step of an implicit solver on the equation defined by coeffs.

This solver is relatively generic in that it models diffusion, convection, etc. abstractly. The caller must do the problem-specific physics calculations to obtain the coefficients for a particular problem.

Parameters:

  • dt (Array) – Discrete time step.

  • x_old (tuple[CellVariable, ...]) – Tuple containing CellVariables for each channel with their values at

  • x_new_guess (tuple[CellVariable, ...]) – Tuple containing initial guess for x_new.

  • coeffs_old (Block1DCoeffs) – Coefficients defining the equation, computed for time t.

  • coeffs_new (Block1DCoeffs) – Coefficients defining the equation, computed for time t+dt.

  • theta_imp (float) – Coefficient in [0, 1] determining which solution method to use. We solve transient_coeff (x_new - x_old) / dt = theta_imp F(t_new) + (1 - theta_imp) F(t_old). Three values of theta_imp correspond to named solution methods: theta_imp = 1: Backward Euler implicit method (default). theta_imp = 0.5: Crank-Nicolson. theta_imp = 0: Forward Euler explicit method

  • convection_dirichlet_mode (str) – See docstring of the convection_terms function, dirichlet_mode argument.

  • convection_neumann_mode (str) – See docstring of the convection_terms function, neumann_mode argument.

Returns:

Tuple, with x_new[i] giving channel i of x at the next time step

Return type:

x_new

torax.fvm.newton_raphson_solve_block module

The newton_raphson_solve_block function.

See function docstring for details.

torax.fvm.newton_raphson_solve_block.body(input_state, jacobian_fun, delta_cond_fun, delta_reduction_factor, log_iterations)[source]

Calculates next guess in Newton-Raphson iteration.

Parameters:

input_state (dict[str, Array])

Return type:

dict[str, Array]

torax.fvm.newton_raphson_solve_block.cond(state, tau_min, maxiter, tol)[source]

Check if exit condition reached for Newton-Raphson iterations.

Parameters:
Return type:

bool

torax.fvm.newton_raphson_solve_block.delta_body(input_delta_state, delta_reduction_factor)[source]

Reduces step size for this Newton iteration.

Parameters:

  • input_delta_state (dict[str, Array])

  • delta_reduction_factor (float)

Return type:

dict[str, Array]

torax.fvm.newton_raphson_solve_block.delta_cond(delta_state, residual_fun)[source]

Check if delta obtained from Newton step is valid.

Parameters:

  • delta_state (dict[str, Array]) – see delta_body.

  • residual_fun (Callable[[Array], Array]) – Residual function.

Return type:

bool

Returns:

True if the new value of x causes any NaNs or has increased the residual relative to the old value of x.

torax.fvm.newton_raphson_solve_block.newton_raphson_solve_block(dt, static_runtime_params_slice, dynamic_runtime_params_slice_t, dynamic_runtime_params_slice_t_plus_dt, geo_t, geo_t_plus_dt, x_old, core_profiles_t, core_profiles_t_plus_dt, transport_model, explicit_source_profiles, source_models, pedestal_model, coeffs_callback, evolving_names, initial_guess_mode, maxiter, tol, coarse_tol, delta_reduction_factor, tau_min, log_iterations=False)[source]

Runs one time step of a Newton-Raphson based root-finding on the equation defined by coeffs.

This solver is relatively generic in that it models diffusion, convection, etc. abstractly. The caller must do the problem-specific physics calculations to obtain the coefficients for a particular problem.

This solver uses iterative root finding on the linearized residual between two sides of the equation describing a theta method update.

The linearized residual for a trial x_new is:

R(x_old) + jacobian(R(x_old))*(x_new - x_old)

Setting delta = x_new - x_old, we solve the linear system:

A*x_new = b, with A = jacobian(R(x_old)), b = A*x_old - R(x_old)

Each successive iteration sets x_new = x_old - delta, until the residual or delta is under a tolerance (tol). If either the delta step leads to an unphysical state, represented by NaNs in the residual, or if the residual doesn’t shrink following the delta step, then delta is successively reduced by a delta_reduction_factor. If tau = delta_now / delta_original is below a tolerance, then the iterations stop. If residual > tol then the function exits with an error flag, producing either a warning or recalculation with a lower dt.

Parameters:

  • dt (Array) – Discrete time step.

  • static_runtime_params_slice (StaticRuntimeParamsSlice) – Static runtime parameters. Changes to these runtime params will trigger recompilation.

  • dynamic_runtime_params_slice_t (DynamicRuntimeParamsSlice) – Runtime parameters for time t (the start time of the step). These config params can change from step to step without triggering a recompilation.

  • dynamic_runtime_params_slice_t_plus_dt (DynamicRuntimeParamsSlice) – Runtime parameters for time t + dt.

  • geo_t (Geometry) – Geometry at time t.

  • geo_t_plus_dt (Geometry) – Geometry at time t + dt.

  • x_old (tuple[CellVariable, ...]) – Tuple containing CellVariables for each channel with their values at the start of the time step.

  • core_profiles_t (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the start of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • core_profiles_t_plus_dt (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the end of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • transport_model (TransportModel) – Turbulent transport model callable.

  • explicit_source_profiles (SourceProfiles) – Pre-calculated sources implemented as explicit sources in the PDE.

  • source_models (SourceModels) – Collection of source callables to generate source PDE coefficients.

  • pedestal_model (PedestalModel) – Model of the pedestal’s behavior.

  • coeffs_callback (CoeffsCallback) – Calculates diffusion, convection etc. coefficients given a core_profiles. Repeatedly called by the iterative optimizer.

  • evolving_names (tuple[str, ...]) – The names of variables within the core profiles that should evolve.

  • initial_guess_mode (InitialGuessMode) – chooses the initial_guess for the iterative method, either x_old or linear step. When taking the linear step, it is also recommended to use Pereverzev-Corrigan terms if the transport coefficients are stiff, e.g. from QLKNN. This can be set by setting use_pereverzev = True in the solver config.

  • maxiter (int) – Quit iterating after this many iterations reached.

  • tol (float) – Quit iterating after the average absolute value of the residual is <= tol.

  • coarse_tol (float) – Coarser allowed tolerance for cases when solver develops small steps in the vicinity of the solution.

  • delta_reduction_factor (float) – Multiply by delta_reduction_factor after each failed line search step.

  • tau_min (float) – Minimum delta/delta_original allowed before the newton raphson routine resets at a lower timestep.

  • log_iterations (bool) – If true, output diagnostic information from within iteration loop.

Returns:

Tuple, with x_new[i] giving channel i of x at the next time step stepper_numeric_outputs: state_module.StepperNumericOutputs. Iteration and

error info. For the error, 0 signifies residual < tol at exit, 1 signifies residual > tol, steps became small.

aux_output: Extra auxiliary output from calc_coeffs.

Return type:

x_new

torax.fvm.optimizer_solve_block module

The optimizer_solve_block function.

See function docstring for details.

torax.fvm.optimizer_solve_block.optimizer_solve_block(dt, static_runtime_params_slice, dynamic_runtime_params_slice_t, dynamic_runtime_params_slice_t_plus_dt, geo_t, geo_t_plus_dt, x_old, core_profiles_t, core_profiles_t_plus_dt, transport_model, explicit_source_profiles, source_models, pedestal_model, coeffs_callback, evolving_names, initial_guess_mode, maxiter, tol)[source]

Runs one time step of an optimization-based solver on the equation defined by coeffs.

This solver is relatively generic in that it models diffusion, convection, etc. abstractly. The caller must do the problem-specific physics calculations to obtain the coefficients for a particular problem.

This solver uses iterative optimization to minimize the norm of the residual between two sides of the equation describing a theta method update.

Parameters:

  • dt (Array) – Discrete time step.

  • static_runtime_params_slice (StaticRuntimeParamsSlice) – Static runtime parameters. Changes to these runtime params will trigger recompilation. A key parameter in this params slice is theta_imp, a coefficient in [0, 1] determining which solution method to use. We solve transient_coeff (x_new - x_old) / dt = theta_imp F(t_new) + (1 - theta_imp) F(t_old). Three values of theta_imp correspond to named solution methods: theta_imp = 1: Backward Euler implicit method (default). theta_imp = 0.5: Crank-Nicolson. theta_imp = 0: Forward Euler explicit method.

  • dynamic_runtime_params_slice_t (DynamicRuntimeParamsSlice) – Runtime params for time t (the start time of the step). These runtime params can change from step to step without triggering a recompilation.

  • dynamic_runtime_params_slice_t_plus_dt (DynamicRuntimeParamsSlice) – Runtime params for time t + dt.

  • geo_t (Geometry) – Geometry object used to initialize auxiliary outputs at time t.

  • geo_t_plus_dt (Geometry) – Geometry object used to initialize auxiliary outputs at time t + dt.

  • x_old (tuple[CellVariable, ...]) – Tuple containing CellVariables for each channel with their values at the start of the time step.

  • core_profiles_t (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the start of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • core_profiles_t_plus_dt (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the end of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • transport_model (TransportModel) – Turbulent transport model callable.

  • explicit_source_profiles (SourceProfiles) – Pre-calculated sources implemented as explicit sources in the PDE.

  • source_models (SourceModels) – Collection of source callables to generate source PDE coefficients.

  • pedestal_model (PedestalModel) – Model of the pedestal’s behavior.

  • coeffs_callback (CoeffsCallback) – Calculates diffusion, convection etc. coefficients given a core_profiles. Repeatedly called by the iterative optimizer.

  • evolving_names (tuple[str, ...]) – The names of variables within the core profiles that should evolve.

  • initial_guess_mode (InitialGuessMode) – Chooses the initial_guess for the iterative method, either x_old or linear step. When taking the linear step, it is also recommended to use Pereverzev-Corrigan terms if the transport use pereverzev terms for linear solver. Is only applied in the nonlinear solver for the optional initial guess from the linear solver.

  • maxiter (int) – See docstring of jaxopt.LBFGS.

  • tol (float) – See docstring of jaxopt.LBFGS.

Returns:

Tuple, with x_new[i] giving channel i of x at the next time step stepper_numeric_outputs: StepperNumericOutputs. Info about iterations and

errors

aux_output: Extra auxiliary output from the calc_coeffs.

Return type:

x_new

torax.fvm.residual_and_loss module

Residual functions and loss functions.

Residual functions define a full differential equation and give a vector measuring (left hand side) - (right hand side). Loss functions collapse these to scalar functions, for example using mean squared error. Residual functions are for use with e.g. the Newton-Raphson method while loss functions can be minimized using any optimization method.

torax.fvm.residual_and_loss.jaxopt_solver(dt, static_runtime_params_slice, dynamic_runtime_params_slice_t_plus_dt, geo_t_plus_dt, x_old, init_x_new_vec, core_profiles_t_plus_dt, transport_model, pedestal_model, explicit_source_profiles, source_models, coeffs_old, evolving_names, maxiter, tol)[source]

Advances jaxopt solver by one timestep.

Parameters:

  • dt (Array) – Time step duration.

  • static_runtime_params_slice (StaticRuntimeParamsSlice) – Static runtime parameters. Changes to these runtime params will trigger recompilation. A key parameter in this params slice is theta_imp, a coefficient in [0, 1] determining which solution method to use. We solve transient_coeff (x_new - x_old) / dt = theta_imp F(t_new) + (1 - theta_imp) F(t_old). Three values of theta_imp correspond to named solution methods: theta_imp = 1: Backward Euler implicit method (default). theta_imp = 0.5: Crank-Nicolson. theta_imp = 0: Forward Euler explicit method.

  • dynamic_runtime_params_slice_t_plus_dt (DynamicRuntimeParamsSlice) – Runtime parameters for time t + dt.

  • geo_t_plus_dt (Geometry) – geometry object for time t + dt.

  • x_old (tuple[CellVariable, ...]) – The starting x defined as a tuple of CellVariables.

  • init_x_new_vec (Array) – Flattened array of initial guess of x_new for all evolving core profiles.

  • core_profiles_t_plus_dt (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the end of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • transport_model (TransportModel) – turbulent transport model callable.

  • pedestal_model (PedestalModel) – Model of the pedestal’s behavior.

  • explicit_source_profiles (SourceProfiles) – pre-calculated sources implemented as explicit sources in the PDE.

  • source_models (SourceModels) – Collection of source callables to generate source PDE coefficients.

  • coeffs_old (Block1DCoeffs) – The coefficients calculated at x_old.

  • evolving_names (tuple[str, ...]) – The names of variables within the core profiles that should evolve.

  • maxiter (int) – maximum number of iterations of jaxopt solver.

  • tol (float) – tolerance for jaxopt solver convergence.

Returns:

Flattened evolving profile array after jaxopt evolution. final_loss: loss after jaxopt evolution aux_output: auxiliary outputs from calc_coeffs. num_iterations: number of iterations ran in jaxopt

Return type:

x_new_vec

torax.fvm.residual_and_loss.theta_method_block_loss(x_new_guess_vec, dt, static_runtime_params_slice, dynamic_runtime_params_slice_t_plus_dt, geo_t_plus_dt, x_old, core_profiles_t_plus_dt, transport_model, explicit_source_profiles, source_models, coeffs_old, evolving_names, pedestal_model)[source]

Loss for the optimizer method of nonlinear solution.

Parameters:

  • x_new_guess_vec (Array) – Flattened array of current guess of x_new for all evolving core profiles.

  • dt (Array) – Time step duration.

  • static_runtime_params_slice (StaticRuntimeParamsSlice) – Static runtime parameters. Changes to these runtime params will trigger recompilation. A key parameter in this params slice is theta_imp, a coefficient in [0, 1] determining which solution method to use. We solve transient_coeff (x_new - x_old) / dt = theta_imp F(t_new) + (1 - theta_imp) F(t_old). Three values of theta_imp correspond to named solution methods: theta_imp = 1: Backward Euler implicit method (default). theta_imp = 0.5: Crank-Nicolson. theta_imp = 0: Forward Euler explicit method.

  • dynamic_runtime_params_slice_t_plus_dt (DynamicRuntimeParamsSlice) – Runtime parameters for time t + dt.

  • geo_t_plus_dt (Geometry) – geometry object at time t + dt.

  • x_old (tuple[CellVariable, ...]) – The starting x defined as a tuple of CellVariables.

  • core_profiles_t_plus_dt (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the end of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • transport_model (TransportModel) – turbulent transport model callable

  • explicit_source_profiles (SourceProfiles) – pre-calculated sources implemented as explicit sources in the PDE

  • source_models (SourceModels) – Collection of source callables to generate source PDE coefficients.

  • coeffs_old (Block1DCoeffs) – The coefficients calculated at x_old.

  • evolving_names (tuple[str, ...]) – The names of variables within the core profiles that should evolve.

  • pedestal_model (PedestalModel) – Model of the pedestal’s behavior.

Returns:

mean squared loss of theta method residual.

Return type:

loss

torax.fvm.residual_and_loss.theta_method_block_residual(x_new_guess_vec, dt, static_runtime_params_slice, dynamic_runtime_params_slice_t_plus_dt, geo_t_plus_dt, x_old, core_profiles_t_plus_dt, transport_model, explicit_source_profiles, source_models, coeffs_old, evolving_names, pedestal_model)[source]

Residual of theta-method equation for core profiles at next time-step.

Parameters:

  • x_new_guess_vec (Array) – Flattened array of current guess of x_new for all evolving core profiles.

  • dt (Array) – Time step duration.

  • static_runtime_params_slice (StaticRuntimeParamsSlice) – Static runtime parameters. Changes to these runtime params will trigger recompilation. A key parameter in this params slice is theta_imp, a coefficient in [0, 1] determining which solution method to use. We solve transient_coeff (x_new - x_old) / dt = theta_imp F(t_new) + (1 - theta_imp) F(t_old). Three values of theta_imp correspond to named solution methods: theta_imp = 1: Backward Euler implicit method (default). theta_imp = 0.5: Crank-Nicolson. theta_imp = 0: Forward Euler explicit method.

  • dynamic_runtime_params_slice_t_plus_dt (DynamicRuntimeParamsSlice) – Runtime parameters for time t + dt.

  • geo_t_plus_dt (Geometry) – The geometry at time t + dt.

  • x_old (tuple[CellVariable, ...]) – The starting x defined as a tuple of CellVariables.

  • core_profiles_t_plus_dt (CoreProfiles) – Core plasma profiles which contain all available prescribed quantities at the end of the time step. This includes evolving boundary conditions and prescribed time-dependent profiles that are not being evolved by the PDE system.

  • transport_model (TransportModel) – Turbulent transport model callable.

  • explicit_source_profiles (SourceProfiles) – Pre-calculated sources implemented as explicit sources in the PDE.

  • source_models (SourceModels) – Collection of source callables to generate source PDE coefficients.

  • coeffs_old (Block1DCoeffs) – The coefficients calculated at x_old.

  • evolving_names (tuple[str, ...]) – The names of variables within the core profiles that should evolve.

  • pedestal_model (PedestalModel) – Model of the pedestal’s behavior.

Returns:

Vector residual between LHS and RHS of the theta method equation.

Return type:

residual

torax.fvm.residual_and_loss.theta_method_matrix_equation(dt, x_old, x_new_guess, coeffs_old, coeffs_new, theta_imp=1.0, convection_dirichlet_mode='ghost', convection_neumann_mode='ghost')[source]

Returns the left-hand and right-hand sides of the theta method equation.

The theta method solves a differential equation

tc_out partial (tc_in x) / partial t = F

where tc is the transient coefficient, with tc_out being outside the partial derivative and tc_in inside it.

We rearrange this to

partial tc_in x / partial t = F / tc_out

The theta method calculates one discrete time step by solving:

(tc_in_new x_new - tc_in_old x_old) / dt =
theta_imp F_new / tc_out_new + theta_exp F_old / tc_out_old

The equation is on the cell grid where tc is never zero. Therefore it’s safe to multiply equation by dt/tc_in_new and scale the residual to x, which has O(1) values and thus the residual is scaled appropriately.

We thus rearrange to:

x_new - tc_in_old/tc_in_new x_old =
dt theta_imp F_new / (tc_out_new tc_in_new) +
dt theta_exp F_old / (tc_out_old tc_in_new)

Rearranging we obtain

x_new - dt theta_imp F_new / (tc_out_new tc_in_new) =
tc_in_old/tc_in_new x_old + dt theta_exp F_old / (tc_out_old tc_in_new)

We now substitute in F = Cu + c:

(I - dt theta_imp diag(1/(tc_out_new tc_in_new)) C_new) x_new
- dt theta_imp diag(1/(tc_out_new tc_in_new)) c_new
=
(diag(tc_in_old/tc_in_new)
+ dt theta_exp diag(1/(tc_out_old tc_in_new)) C_old) x_old
+ dt theta_exp diag(1/(tc_out_old tc_in_new)) c_old
Parameters:

  • dt (Array) – Time step duration.

  • x_old (tuple[CellVariable, ...]) – The starting x defined as a tuple of CellVariables.

  • x_new_guess (tuple[CellVariable, ...]) – Current guess of x_new defined as a tuple of CellVariables.

  • coeffs_old (Block1DCoeffs) – The coefficients calculated at x_old.

  • coeffs_new (Block1DCoeffs) – The coefficients calculated at x_new.

  • theta_imp (float) – Coefficient on implicit term of theta method.

  • convection_dirichlet_mode (str) – See docstring of the convection_terms function, dirichlet_mode argument.

  • convection_neumann_mode (str) – See docstring of the convection_terms function, neumann_mode argument.

Return type:

tuple[Array, Array, Array, Array]

Returns:

For the equation A x_new + a_vec = B x_old + b_vec. This function returns

  • left-hand side matrix, A

  • left-hand side vector, a

  • right-hand side matrix B

  • right-hand side vector, b

Module contents

Finite volume method.

This module contains functionality related to solving differential equations built using the finite volume method. This module is meant to be somewhat problem-independent, with other modules providing the coefficients on relatively generic differential equations.