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

[1]

M. Perrot and F. Lemarié. Energetically consistent Eddy-Diffusivity Mass-Flux convective schemes. Part I: Theory and Models (2024). url : hal.science/hal-04439113.

[2]

A. Zhou, L. Hawkins and P. Gentine. Proof-of-concept: Using ChatGPT to Translate and Modernize an Earth System Model from Fortran to Python/JAX (2024). url : arxiv.org/abs/2405.00018.

class SingleColumnModel(time_frame, dt, out_dt, init_state, case, closure_name)

Single column forward model core.

This forward model of tunax is the combination of 4 things : the physical case case, an initial state of the water column init_state, the time information with time_frame, dt and out_dt 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 compute_trajectory_with().

Parameters:

• time_frame (float) – Total time of the simulation \([\text h]\).

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

• out_dt (float) – Time-step for the output writing \([\text s]\).

• init_state (State) – cf. attribute.

• case (Case) – cf. attribute.

• closure_name (str) – Name of the chosen closure, must be a key of CLOSURES_REGISTRY, see its documentation for the available closures.

Return type:

SingleColumnModel

nt

Number of integration interations.

Type:

int

dt

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

Type:

float

n_out

Number of of time-steps between every output.

Type:

int

init_state

Initial physical state of the water column.

Type:

State

case

Physical case and forcings of the experiment.

Type:

Case

closure

Abstraction representing the chosen closure.

Type:

Closure

Warning

• If time_frame is not proportional to the time-step dt.

• If time_frame is not proportional to the out time-step out_dt.

Raises:

• ValueError – If out_dt is not proportional to the time step dt.

• ValueError – If closure_name is not registerd in CLOSURES_REGISTRY.

Parameters:

• time_frame (float)

• dt (float)

• out_dt (float)

• init_state (State)

• case (Case)

• closure_name (str)

Return type:

SingleColumnModel

Note

To make this forward model compatible with the fitter part of Tunax, the parameters of the closure are only given during of the call of the run with compute_trajectory_with().

compute_trajectory_with(closure_parameters)

Run the model with a specific set of closure parameters.

This method is the main one for runing the model. It calls nt times the function step() and regulary writes the output to build the Trajectory output.

Parameters:

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

Returns:

trajectory – Timeseries of the evolution of the variables of the model every out_dt.

Return type:

Trajectory

step(dt, case, closure, state, closure_state, closure_parameters, swr_frac)

Run 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:

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

• case (Case) – Physical case and forcings of the experiment.

• closure (Closure) – Abstraction representing the chosen closure.

• state (State) – Curent state of the water column.

• closure_state (ClosureStateAbstract) – Curent state of the water column for the closure variables.

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

• swr_frac (Float[Array, 'nz+1']) – Fraction of solar penetration throught the water column \([\text{dimensionless}]\).

Return type:

Tuple[State, ClosureStateAbstract]

Returns:

• state (State) – State of the water column at next time-step.

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

Note

This function is jitted with JAX, it should make it faster, but the jit() decorator can be removed.

lmd_swfrac(hz)

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, '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, ‘nz+1’]

advance_tra_ed(t, s, akt, swr_frac, hz, dt, case)

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.

Parameters:

• t (Float[Array, 'nz']) – Current temperature on the center of the cells \([° \text C]\).

• s (Float[Array, 'nz']) – Current salinity on the center of the cells \([\text{psu}]\).

• akt (Float[Array, 'nz+1']) – Current eddy-diffusivity \(K_m\) on the interfaces of the cells \(\left[\text m ^2 \cdot \text s ^{-1}\right]\).

• swr_frac (Float[Array, 'nz+1']) – Fraction of solar penetration throught the water column \([\text{dimensionless}]\).

• hz (Float[Array, 'nz']) – Thickness of cells from deepest to shallowest \([\text m]\).

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

• case (Case) – Physical case and forcings of the experiment.

Return type:

Tuple[Float[Array, 'nz'], Float[Array, 'nz']]

Returns:

• t (Float[Array, ‘nz’]) – Temperature on the center of the cells at next step \([° \text C]\).

• s (Float[Array, ‘nz’]) – Salinity on the center of the cells at next step \([\text{psu}]\).

advance_dyn_cor_ed(u, v, hz, akv, dt, case)

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.

Parameters:

• u (Float[Array, 'nz']) – Current zonal velocity on the center of the cells \(\left[\text m \cdot \text s^{-1}\right]\).

• v (Float[Array, 'nz']) – Current meridional velocity on the center of the cells \(\left[\text m \cdot \text s^{-1}\right]\).

• akv (Float[Array, 'nz+1']) – Current eddy-viscosity \(K_v\) on the interfaces of the cells \(\left[\text m ^2 \cdot \text s ^{-1}\right]\).

• hz (Float[Array, 'nz']) – Thickness of cells from deepest to shallowest \([\text m]\).

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

• case (Case) – Physical case and forcings of the experiment.

Return type:

Tuple[Float[Array, 'nz'], Float[Array, 'nz']]

Returns:

• u (Float[Array, ‘nz’]) – Zonal velocity on the center of the cells at the next time step \(\left[\text m \cdot \text s^{-1}\right]\).

• v (Float[Array, ‘nz’]) – Meridional velocity on the center of the cells at the next time step \(\left[\text m \cdot \text s^{-1}\right]\).

diffusion_solver(ak, hz, f, dt)

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, 'nz+1']) – Diffusion at the cell interfaces \(K\) in \(\left[\text m ^2 \cdot \text s ^{-1}\right]\).

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

• f (Float[Array, '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, ‘nz’]