functions

Usefull calculation functions.

Thefunctions in this module are supposed to be used in various other modules. They can be called by the prefix tunax.functions. or directly by tunax..

tridiag_solve(a, b, c, f)

Solve a trigiagonal problem.

The tridiagonal problem can be written \(\mathbb MX = F\) where \(\mathbb M = \begin{pmatrix} b_1 & c_1 & & \\ a_2 & \ddots & \ddots & \\ & \ddots & \ddots & c_{n-1} \\ & & a_n & b_n \end{pmatrix}\) and \(F = \begin{pmatrix} f_1 \\ \vdots \\ f_n \end{pmatrix}\). The problem is solved by recurrence using jax.lax

Parameters:

• a (Float[Array, 'n']) – Left diagonal of \(\mathbb M\), the first element is not used.

• b (Float[Array, 'n']) – Middle diagonal of \(\mathbb M\).

• c (Float[Array, 'n']) – Right diagonal of \(\mathbb M\), the last element is not used.

• f (Float[Array, 'n']) – Right hand of the equation \(F\).

Returns:

x – Solution \(X\) of tridiagonal problem.

Return type:

Float[Array, ‘nz’]

add_boundaries(vec_btm, vec_in, vec_sfc)

Concatenate the three parts of a vector : surface, bottom and inside.

This functions is made to avoid loops and make JAX more efficient by writing the calculations with vectorization when possible.

Parameters:

• vec_btm (float) – Bottom value of the vector.

• vec_in (Float[Array, 'n-2']) – Middle values of the vector.

• vec_sfc (float) – Surface value of the vector.

Return type:

Float[Array, 'n']

Returns:

• vec (Float[~jax.Array, ‘n’]) – Concatenated vector.