tunax 0.1.0__py3-none-any.whl

tunax/__init__.py

@@ -0,0 +1,18 @@
+"""
+Importation of Tunax classes and functions for shortcuts.
+"""
+
+
+from .space import Grid, State, Trajectory
+from .case import Case
+from .closure import (
+    ClosureParametersAbstract, ClosureStateAbstract, Closure
+)
+from .database import Obs, Database
+from .functions import tridiag_solve, add_boundaries
+from .closures_registry import CLOSURES_REGISTRY
+from .model import (
+    SingleColumnModel, step, lmd_swfrac, advance_tra_ed, advance_dyn_cor_ed,
+    diffusion_solver
+)
+from .fitter import FittableParameter, FittableParametersSet, Fitter

tunax/case.py

@@ -0,0 +1,282 @@
+"""
+Physical parameters and forcings.
+
+This module comes down to Case class. This class can be obtained by the prefix
+:code:`tunax.case.` or directly by :code:`tunax.`.
+
+"""
+
+from __future__ import annotations
+
+import equinox as eqx
+import jax.numpy as jnp
+from jax import device_get
+
+
+_OMEGA = 7.292116e-05
+"""float: Rotation rate of the Earth [rad.s-1]."""
+_RAD_DEG = jnp.pi/180.
+"""float: Measure of one degree in radiant [rad.°-1]."""
+
+
+class Case(eqx.Module):
+    r"""
+    Physical parameters and forcings.
+
+    This class contains all the physical constants, and the constant forcings
+    that definine an experience for the model.
+
+    Parameters
+    ----------
+    rho0 : float, default=1024.
+        cf. attribute.
+    grav : float, default=9.81
+        cf. attribute.
+    cp : float, default=3985.
+        cf. attribute.
+    alpha : float, default=2e-4
+        cf. attribute.
+    beta : float, default=8e-4
+        cf. attribute.
+    t_rho_ref : float, default=0.
+        cf. attribute.
+    s_rho_ref : float, default=35.
+        cf. attribute.
+    vkarmn : float, default=0.384
+        cf. attribute.
+    fcor : float, default=0.
+        cf. attribute.
+    ustr_sfc : float, default=0.
+        cf. attribute.
+    ustr_btm : float, default=0.
+        cf. attribute.
+    vstr_sfc : float, default=0.
+        cf. attribute.
+    vstr_btm : float, default=0.
+        cf. attribute.
+    tflx_sfc : float, default=0.
+        cf. attribute.
+    tflx_btm : float, default=0.
+        cf. attribute.
+    sflx_sfc : float, default=0.
+        cf. attribute.
+    sflx_btm : float, default=0.
+        cf. attribute.
+    rflx_sfc_max : float, default=0.
+        cf. attribute.
+
+    Attributes
+    ----------
+    rho0 : float, default=1024.
+        Default density of saltwater :math:`[\text{kg} \cdot \text{m}^{-3}]`.
+    grav : float, default=9.81
+        Gravity acceleration :math:`[\text{m} \cdot \text{s}^{-2}]`.
+    cp : float, default=3985.
+        Specific heat capacity of saltwater
+        :math:`[\text{J} \cdot \text{kg}^{-1} \cdot \text{K}^{-1}]`.
+    alpha : float, default=2e-4
+        Thermal expansion coefficient :math:`[\text{K}^{-1}]`.
+    beta : float, default=8e-4
+        Salinity expansion coefficient :math:`[\text{psu}^{-1}]`.
+    t_rho_ref : float, default=0.
+        Reference temperature for the density computation :math:`[° \text C]`.
+    s_rho_ref : float, default=35.
+        Reference salinity for the density computation :math:`[\text{psu}]`.
+    vkarmn : float, default=0.384
+        Von Kármán constant :math:`[\text{dimensionless}]`.
+    fcor : float, default=0.
+        Coriolis frequency at the water column
+        :math:`[\text{rad} \cdot \text{s}^{-1}]`.
+    ustr_sfc : float, default=0.
+        Zonal wind stress :math:`[\text{m}^{2} \cdot \text{s}^{-2}]`.
+    ustr_btm : float, default=0.
+        Zonal current stress at the bottom
+        :math:`[\text{m}^{2} \cdot \text{s}^{-2}]`.
+    vstr_sfc : float, default=0.
+        Meridional wind stress :math:`[\text{m}^{2} \cdot \text{s}^{-2}]`.
+    vstr_btm : float, default=0.
+        Meridional current stress at the bottom
+        :math:`[\text{m}^{2} \cdot \text{s}^{-2}]`.
+    tflx_sfc : float, default=0.
+        Non-penetrative heat flux at the surface
+        :math:`[\text{K} \cdot \text{m} \cdot \text{s}^{-1}]`.
+    tflx_btm : float, default=0.
+        Non-penetrative heat flux at the bottom
+        :math:`[\text{K} \cdot \text{m} \cdot \text{s}^{-1}]`.
+    sflx_sfc : float, default=0.
+        Fresh water flux at the surface
+        :math:`[\text{psu} \cdot \text{m} \cdot \text{s}^{-1}]`.
+    sflx_btm : float, default=0.
+        Fresh water flux at the bottom
+        :math:`[\text{psu} \cdot \text{m} \cdot \text{s}^{-1}]`.
+    rflx_sfc_max : float, default=0.
+        Maximum solar radiation flux at the surface
+        :math:`[\text{K} \cdot \text{m} \cdot \text{s}^{-1}]`.
+
+    """
+
+    # physcal constants
+    rho0: float = 1024.
+    grav: float = 9.81
+    cp: float = 3985.
+    alpha: float = 2e-4
+    beta: float = 8e-4
+    t_rho_ref: float = 0.
+    s_rho_ref: float = 35.
+    vkarmn: float = 0.384
+    #forcings
+    fcor: float = 0.
+    ustr_sfc: float = 0.
+    ustr_btm: float = 0.
+    vstr_sfc: float = 0.
+    vstr_btm: float = 0.
+    tflx_sfc: float = 0.
+    tflx_btm: float = 0.
+    sflx_sfc: float = 0.
+    sflx_btm: float = 0.
+    rflx_sfc_max: float = 0.
+
+    def set_lat(self, lat: float) -> Case:
+        """
+        Set the Coriolis frequency from the latitude.
+
+        Parameters
+        ----------
+        lat : float
+            Latitude of the water column :math:`[°]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the the new value of :attr:`fcor`.
+        """
+        fcor = float(device_get(2.*_OMEGA*jnp.sin(_RAD_DEG*lat)))
+        print(type(fcor))
+        print(fcor)
+        case = eqx.tree_at(lambda t: t.fcor, self, fcor)
+        return case
+
+    def set_u_wind(self, u_wind: float) -> Case:
+        r"""
+        Set the zonal wind stress with the zonal wind speed.
+
+        Parameters
+        ----------
+        u_wind : float
+            Zonal wind speed :math:`[\text m \cdot \text s ^{-1}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`ustr_sfc`.
+        """
+        case = eqx.tree_at(lambda t: t.ustr_sfc, self, u_wind**2)
+        return case
+
+    def set_u_cur(self, u_cur: float) -> Case:
+        r"""
+        Set the zonal current stress with the zonal current.
+
+        Parameters
+        ----------
+        u_cur : float
+            Zonal current speed :math:`[\text m \cdot \text s ^{-1}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`ustr_btm`.
+        """
+        case = eqx.tree_at(lambda t: t.ustr_btm, self, u_cur**2)
+        return case
+
+    def set_v_wind(self, v_wind: float) -> Case:
+        r"""
+        Set the meridional wind stress with the meridional wind.
+
+        Parameters
+        ----------
+        v_wind : float
+            Meridional wind speed :math:`[\text m \cdot \text s ^{-1}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`vstr_sfc`.
+        """
+        case = eqx.tree_at(lambda t: t.vstr_sfc, self, v_wind**2)
+        return case
+
+    def set_v_cur(self, v_cur: float) -> Case:
+        r"""
+        Set the meridional current stress with the meridional current.
+
+        Parameters
+        ----------
+        v_cur : float
+            Meridional current speed :math:`[\text m \cdot \text s ^{-1}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`vstr_btm`.
+        """
+        case = eqx.tree_at(lambda t: t.vstr_btm, self, v_cur**2)
+        return case
+
+    def set_tpw_sfc(self, tpw_sfc: float) -> Case:
+        r"""
+        Set the heat flux at surface from the heat power.
+
+        Parameters
+        ----------
+        tpw_sfc : float
+            Non-penetrative heat power at the surface
+            :math:`[\text W \cdot \text m ^{-2}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`tflx_sfc`.
+        """
+        tflx_sfc = tpw_sfc/(self.rho0*self.cp)
+        case = eqx.tree_at(lambda t: t.tflx_sfc, self, tflx_sfc)
+        return case
+
+    def set_tpw_btm(self, tpw_btm: float) -> Case:
+        r"""
+        Set the heat flux at bottom from the heat power.
+
+        Parameters
+        ----------
+        tpw_btm : float
+            Non-penetrative heat power at the bottom
+            :math:`[\text W \cdot \text m ^{-2}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`tflx_btm`.
+        """
+        tflx_btm = tpw_btm/(self.rho0*self.cp)
+        case = eqx.tree_at(lambda t: t.tflx_btm, self, tflx_btm)
+        return case
+
+    def set_rpw_sfc_max(self, rpw_sfc_max: float) -> Case:
+        r"""
+        Set the maximum solar radiation from the solar power
+
+        Parameters
+        ----------
+        rpw_sfc_max : float
+            Maximum solar radiation power at the surface (penetrative)
+            :math:`[\text W \cdot \text m ^{-2}]`.
+        
+        Returns
+        -------
+        case : Case
+            The :code:`self` object with the new value of :attr:`rflx_sfc_max`.
+        """
+        rflx_sfc_max = rpw_sfc_max/(self.rho0*self.cp)
+        case = eqx.tree_at(lambda t: t.rflx_sfc_max, self, rflx_sfc_max)
+        return case

tunax/closure.py

@@ -0,0 +1,165 @@
+"""
+Abstractions for defining closures.
+
+A closure is a a set of physical equations which determines the sub-mesh
+turbulence of the water column at a specific time. The purpose of Tunax is to
+calibrate the parameters of these equations. This module contains the abstract
+classes required to define a closure, which is done in the folder
+:code:`closures/`. These classes can be obtained by the prefix
+:code:`tunax.closure.` or directly by :code:`tunax.`.
+
+"""
+
+from abc import ABC
+from typing import Callable, Type, TypeVar, Generic
+
+import equinox as eqx
+from jaxtyping import Float, Array
+
+from tunax.case import Case
+from tunax.space import Grid, State
+
+
+class ClosureParametersAbstract(eqx.Module, ABC):
+    """
+    Abstraction for the parameters of a closure.
+    
+    To define a closure, a class that inherits from this one must be created.
+    This parent class does not impose anything on the child one. The child
+    class should be fill of the parameters involved in the closure which may be
+    calibrated. The child class must be the only place where these parameters
+    are defined for running a forward model : the instance of this class is
+    given to the closure functions to recover the parameters and do the
+    computation. For a run of a calibration, the description of the parameters
+    to calibration is done from the child class. It can also includes
+    parameters that are not dedicated to be calibrated such as physical or
+    mathematical constants. The attributes of this class can be used to avoid
+    the systematic computation of some values which are independent of the time
+    (eg. k-epsilon with the method :code:`__post_init__`). 
+    
+    """
+
+
+# variable that represent a type which contains ClosureParametersAbstract and
+# all its subclasses
+CloParT = TypeVar('CloParT', bound=ClosureParametersAbstract)
+
+
+class ClosureStateAbstract(eqx.Module, ABC):
+    r"""
+    Abstraction for the water column state linked to the closure.
+
+    To define a closure, a class that inherits rom this one must be created.
+    This parent class imposes a grid and three variables as attributes : these
+    variables are used at each step of the forward model to compute the next
+    step of tracers and momentum. The child class should be fill with other
+    arrays and scalars that may evolve throught time and necessary for the
+    computation of the closure at the next step. The child class is similary
+    to the :class:`~space.State` class but for the diffusivity part.
+    The :code:`__init__` method can be overwrited to fill the variables with
+    initial values.
+
+    Parameters
+    ----------
+    grid : Grid
+        cf. attribute.
+    akt : Float[~jax.Array, 'nz+1']
+        cf. attribute.
+    akv : Float[~jax.Array, 'nz+1']
+        cf. attribute.
+
+    Attributes
+    ----------
+    grid : Grid
+        Geometry of the water column, should be the same than for the
+        :class:`~space.State` instance used in the model.
+    akt : Float[~jax.Array, 'nz+1']
+        Eddy-diffusivity on the interfaces of the cells
+        :math:`\left[\text m ^2 \cdot \text s ^{-1}\right]`.
+    akv : Float[~jax.Array, 'nz+1']
+        Eddy-viscosity on the interfaces of the cells
+        :math:`\left[\text m ^2 \cdot \text s ^{-1}\right]`.
+
+    """
+
+    grid: Grid
+    akt: Float[Array, 'nz+1']
+    akv: Float[Array, 'nz+1']
+
+
+# variable that represent a type which contains ClosureStateAbstract and
+# all its subclasses
+CloStateT = TypeVar('CloStateT', bound=ClosureStateAbstract)
+
+
+class Closure(eqx.Module, Generic[CloStateT, CloParT]):
+    r"""
+    Implementation of a physical closure for computing eddy-diffusivity.
+
+    This class contains the three parts that defines a closure of a vertical
+    physics. These parts will be used in the model to compute the eddy-
+    diffusivity and eddy-viscosity at each time-step from the current water
+    column :class:`~space.State`, the current water column state of the
+    closure (child class of :class:`~ClosureStateAbstract`), the physical
+    :class:`~case.Case` and the closure parameters (child class of
+    :class:`~ClosureParametersAbstract`).
+
+    Parameters
+    ----------
+    name : str
+        cf. attribute.
+    parameters_class : Type[ClosureParametersAbstract]
+        cf. attribute.
+    state_class : Type[ClosureStateAbstract]
+        cf. attribute.
+    step_fun : Callable[[State, CloStateT, float, CloParT, Case], CloStateT]
+        cf. attribute.
+
+    Attributes
+    ----------
+    name : str
+        The name of the closure.
+    parameters_class : Type[ClosureParametersAbstract]
+        A child class of :class:`~ClosureParametersAbstract` that defines the
+        constant parameters used in the computation done by the closure, it
+        includes the parameters that may be calibrated.
+    state_class : Type[ClosureStateAbstract]
+        A child class of :class:`~ClosureStateAbstract` that defines the state
+        of the water column for the variables used by the closure computation.
+    step_fun : Callable[[State, CloStateT, float, CloParT, Case], CloStateT]
+        Ths function is called at every step of the forward model to compute
+        the eddy-diffusivity before resolving the equation of the tracers and
+        of the momentum.
+        
+        Parameters
+        ----------
+        state : State
+            Current state of the water column.
+        closure_state : CloStateT
+            Current state of the water column for the variables used by the
+            closure.
+        dt : float
+            Time-step of the forward model :math:`[\text s]`.
+        closure_params : CloParT
+            Values of the parameters used by the closure (time-independant).
+        case : Case
+            Physical parameters and forcings of the model run.
+        
+        Returns
+        -------
+        closure_state : CloStateT
+            State of the water column for the variables used by the
+            closure at the next time-step.
+        
+        Notes
+        -----
+        :code:`CloStateT` is the type reprensenting the instances of the child
+        classes of :class:`~ClosureStateAbstract` and :code:`CloParT` is the
+        same for :class:`~ClosureParametersAbstract`.
+
+    """
+
+    name : str
+    parameters_class: Type[ClosureParametersAbstract]
+    state_class: Type[ClosureStateAbstract]
+    step_fun: Callable[[State, CloStateT, float, CloParT, Case], CloStateT]

tunax/closures/__init__.py

@@ -0,0 +1,6 @@
+"""
+Importation of Tunax classes and functions of closures for shortcuts.
+"""
+
+
+from .k_epsilon import KepsParameters, KepsState, keps_step