functions

Usefull calculation functions.

The functions 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..

FloatJax

Type that represent a float in a Array, used only for the code linter.

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

Parameters:

• a (float Array of shape (n)) – Left diagonal of \(\mathbb M\), the first element is not used.

• b (float Array of shape (n)) – Middle diagonal of \(\mathbb M\).

• c (float Array of shape (n)) – Right diagonal of \(\mathbb M\), the last element is not used.

• f (float Array of shape (n)) – Right hand of the equation \(F\).

Returns:

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

Return type:

float Array of shape (n)

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 of shape (n-2)) – Middle values of the vector.

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

Returns:

vec – Concatenated vector.

Return type:

float Array of shape (n)