model

Single column forward model core.

This module contains the main class SingleColumnModel which is the core of Tunax for implementing computing the evolution of a water column of the ocean. It also contains functions that are used in this computation. The model was traduced from Frotran to JAX with the work of Florian Lemarié and Manolis Perrot [1], the translation was done in part using the work of Anthony Zhou, Linnia Hawkins and Pierre Gentine [2]. this class and these functions can be obtained by the prefix tunax.model. or directly by tunax..

References

StatesTime

Type that represent the values that are transformed in an integration step of the model.

alias of Tuple[State, ClosureStateAbstract, float]

class SingleColumnModel[source]

Single column forward model core.

This forward model of tunax is the combination of 4 things : the physical case case_tracable, an initial state of the water column init_state, the time information with nt, dt and p_out and the abstraction of the chosen closure for eddy-diffusivity closure. Adding a set of parameters for the closure one can run the model with the method run(). The builder of this class takes the case as arguments all the attributes, except for the case : the parameter is an instance of Case which is transformed in a CaseTracable instance for JAX purpose and for the closure : the parameter is only the name of the closure.

Parameters:
nt

Number of integration interations.

Type:

int

dt

Time-step of integration for every iteration \([\text s]\).

Type:

float

p_out

Number of of time-steps between every output step.

Type:

int

init_state

Initial physical state of the water column.

Type:

State

case_tracable

Physical case and forcings of the experiment, made tracable for JAX purposes.

Type:

CaseTracable

closure

Abstraction representing the chosen closure.

Type:

Closure

start_time

Value of the starting time in \([\text s]\).

Type:

float, default=0.

checkpoint

Use the checkpoint() on the partial run method. Used for economize the memory when computing the gradient, especially on GPUs.

Type:

bool, default=False

Note

The closure parameters are given only at the end when we compute the model so that the calibrations of the parameters are easier.

tra_promote(promotions)[source]

Increase the dimension type of the tracers.

It is usefull to use vmap() on several instances of SingleColumnModel to do batch computing (run the model in parallel).

Parameters:

promotions (Dict[str, str]) – The keys of the dictionnary are the name of the tracer variables to modify the dimensions of the forcing, one of {'t', 's', 'b', 'pt’}, the values are the new dimensions of the forcings, possible values {'constant', 'variable'}.

Returns:

model – The self object with the several promoted forcings.

Return type:

SingleColumnModel

step(state, closure_state, time, closure_parameters)[source]

Runs one time-step of the model.

This functions first call the closure to compute the eddy-diffusivity and viscosity, and then integrate the equations of tracers and momentum. It modifies the state with these new values and then returns the new state and closure_state.

Parameters:

  • state (State) – State of the water column at the current time.

  • closure_state (ClosureStateAbstract) – State of the water column for the closure variables at the current time.

  • time (float) – Time of the current iteration (the mehtod integrates from this time to the next one).

  • closure_parameters (ClosureParametersAbstract) – A set of parameters of the used closure.

Return type:

Tuple[State, ClosureStateAbstract, float]

Returns:

  • state (State) – State of the water column after the integration.

  • closure_state (ClosureStateAbstract) – State of the water column for the closure variables after the integration.

  • time (float) – Value of the time after the integration.

step_check(state, closure_state, time, closure_parameters)[source]

Checkpointed version of run() to save memory during the gradient computation.

Parameters:

  • state (State) – State of the water column at the current time.

  • closure_state (ClosureStateAbstract) – State of the water column for the closure variables at the current time.

  • time (float) – Time of the current iteration (the mehtod integrates from this time to the next one).

  • closure_parameters (ClosureParametersAbstract) – A set of parameters of the used closure.

Return type:

Tuple[State, ClosureStateAbstract, float]

Returns:

  • state (State) – State of the water column after the integration.

  • closure_state (ClosureStateAbstract) – State of the water column for the closure variables after the integration.

  • time (float) – Value of the time after the integration.

run_partial(state0, closure_state0, time0, n_steps, closure_parameters)[source]

Runs a certain number of time steps.

Computes a loop of integration for a number of time steps of n_steps, and return the last states of the loop.

Parameters:

  • state0 (State) – State of the water column at the beginning of the integration loop.

  • closure_state (ClosureStateAbstract) – State of the water column for the closure variables at the beginning of the integration loop.

  • time0 (float) – Begining time of the integration loop.

  • n_steps (int) – Number of integration steps.

  • closure_parameters (ClosureParametersAbstract) – A set of parameters of the used closure.

  • closure_state0 (ClosureStateAbstract)

Return type:

Tuple[State, ClosureStateAbstract, float]

Returns:

  • state (State) – State of the water column after a number of n_steps integration steps.

  • closure_state (ClosureStateAbstract) – State of the water column for the closure variables after a number of n_steps integration steps.

  • time (float) – Value of the time after a number of n_steps integration steps.

run(closure_parameters)[source]

Main run the model.

Computes the nt integration steps of lenght dt from the initial state init_state with and doing an output of the state every p_out steps with the physical case and forcings corresponding to case_tracable. The closure for eddy- diffusivity used is closure with the values of the parameters described by the parameter closure_parameters.

Parameters:

closure_parameters (ClosureParametersAbstract) – The set of parameters to use for the computation of the closure of eddy-diffusivity.

Returns:

trajectory – The trajectory with all the output steps of the integration loop.

Return type:

Trajectory

jit_run(closure_parameters)[source]

Jitted version of run().

This method does exacly like run() but jit() is applied on it, which means that the first call of this method will be the compilation of the function, and the next ones will be the compiled execution of the function which are supposed to be faster. There will be a compilation each time that this method will be call for a new “shape” of the SingleColumnModel instance (which means that all the leafs of the pytree have the same shape as Array), but even if this method is call on different instances, if they have the same shape, the compilation will be done only one time. Moreover, this method should use only for direct methods, if one wants to apply a grad() over it’s better to put the jit() outside.

Parameters:

closure_parameters (ClosureParametersAbstract) – cf. run().

Returns:

trajectory – cf. run().

Return type:

Trajectory

lmd_swfrac(hz)[source]

Compute solar forcing.

Compute fraction of solar shortwave flux penetrating to specified depth due to exponential decay in Jerlov water type. This function is called once before running the model.

Parameters:

hz (float Array of shape (nz)) – Thickness of cells from deepest to shallowest \([\text m]\).

Returns:

swr_frac – Fraction of solar penetration throught the water column \([\text{dimensionless}]\).

Return type:

float Array of shape (nz+1)

tracer_flux(tracer, case_tracable, grid, i_time)[source]

Computes flux of the tracer forcing.

This function get the flux of the forcing at a certain time depending on the type of the forcing, the flux being the derivative of the forcing along the depth.

Parameters:

  • tracer (str) – Name of the tracer variable of the concerned forcing. One of {'t', 's', 'b', 'pt’}.

  • case_tracable (CaseTracable) – Physical case which contains the forcings type and values.

  • grid (Grid) – Vertical grid of the water column.

  • i_time (int) – Index of the time iteration corresponding to the index in the forcing.

Returns:

df – Flux of the forcing of the tracer. At each cell it represents the difference between the input and the ouput flux.

Return type:

float Array of shape (nz)

Raises:

ValueError – If the forcing type is not one of {‘borders’, ‘constant’, ‘variable’}.

advance_tra_ed(state, akt, dt, case_tracable, i_time)[source]

Integrate vertical diffusion term for tracers.

First the flux divergences are computed taking in account the forcings. Then the diffusion equation of the tracers system is solved, and the tracers at next time-step are returned. The solved equation is for \(C\) a tracer :

\(\partial _z ( K_m \partial _z C) + \partial _t C + F = 0\)

where \(F\) is the representation of the forcings. This equation is solved for every tracer indicated in eos_tracers and the passive tracer if do_pt is set.

Parameters:

  • state (State) – State of the water column at the current iteration.

  • akt (float Array of shape (nz)+1) – Eddy-diffusivity on the interfaces of the cells \(\left[\text m ^2 \cdot \text s ^{-1}\right]\).

  • dt (float) – Time-step of the integration step.

  • case_tracable (CaseTracable) – Physical case which contains the forcings type and values.

  • i_time (int) – Index of the time iteration corresponding to the index in the forcing.

Returns:

state – The state of the water column with the values of the tracers after the integration and the diffusion by the eddy-diffusivity.

Return type:

State

advance_dyn_cor_ed(state, akv, dt, case_tracable)[source]

Integrate vertical diffusion and Coriolis terms for momentum.

First the Coriolis term is computed, then the momentum forcings are applied and finally, the diffusion equation is solved. The momentum at next time-step is returned. The equation which is solved is :

\(\partial_z (K_v \partial_z U) + F_{\text{cor}}(U) + F = 0\)

where \(F_{\text{cor}}\) represent the Coriolis effect, and \(F\) represent the effect of the forcings on the momentum.

Parameters:

  • state (State) – State of the water column at the current iteration.

  • akv (float Array of shape (nz+1)) – Eddy-viscosity on the interfaces of the cells \(\left[\text m ^2 \cdot \text s ^{-1}\right]\).

  • dt (float) – Time-step of the integration step.

  • case_tracable (CaseTracable) – Physical case which contains the forcings type and values.

Returns:

state – The state of the water column with the values of the momentum after the integration and the diffusion by the eddy-viscosity.

Return type:

State

diffusion_solver(ak, hz, f, dt)[source]

Solve a diffusion problem with finite volumes.

The diffusion problems can be written

\(\partial _z (K \partial _z X) + \dfrac f {\Delta t \Delta x} = 0\)

where we are searching for \(X\) and where \(f\) represents the temporal derivative and forcings. This function transforms this problem in a tridiagonal system and then solve it.

Parameters:

  • ak (float Array of shape (nz+1)) – Diffusion at the cell interfaces \(K\) in \(\left[\text m ^2 \cdot \text s ^{-1}\right]\).

  • hz (float Array of shape (nz)) – Thickness of cells from deepest to shallowest \(\left[\text m\right]\).

  • f (float Array of shape (nz)) – Right-hand flux of the equation \(f\) in \([[X] \cdot \text m ]\).

  • dt (float) – Time-step of discretisation \([\text s]\).

Returns:

x – Solution of the diffusion problem \(X\) in \(\left[[X]\right]\).

Return type:

float Array of shape (nz)