From fe4df9c515ad2a6d5060c3d2d5abe1e23c66f7d8 Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Thu, 1 Dec 2022 15:08:41 +0100
Subject: [PATCH 1/7] Initial commit for the wind diagnostic

---
 qgs/diagnostics/wind.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/qgs/diagnostics/wind.py b/qgs/diagnostics/wind.py
index c5e5ca6..0986e58 100644
--- a/qgs/diagnostics/wind.py
+++ b/qgs/diagnostics/wind.py
@@ -27,6 +27,7 @@
 
 from qgs.diagnostics.differential import DifferentialFieldDiagnostic
 from qgs.diagnostics.util import create_grid_basis
+from qgs.functions.tendencies import create_tendencies
 
 
 class AtmosphericWindDiagnostic(DifferentialFieldDiagnostic):
@@ -632,6 +633,46 @@ def _get_diagnostic(self, dimensional):
         return self._diagnostic_data
 
 
+class VerticalVelocity(AtmosphericWindDiagnostic):
+
+    def __init__(self, model_params, f=None, atmospheric_inner_products=None, delta_x=None, delta_y=None, dimensional=True):
+
+        self.type = "W"
+        AtmosphericWindDiagnostic.__init__(self, model_params, delta_x, delta_y, dimensional)
+        self._plot_title = r'Atmospheric vertical wind in the middle layer'
+
+        if f is None or atmospheric_inner_products is None:
+            self._f, _, self._aips = create_tendencies(model_params, return_inner_products=True)
+        else:
+            self._f = f
+            self._aips = atmospheric_inner_products
+
+    def set_data(self, time, data):
+
+        AtmosphericWindDiagnostic.set_data(self, time, data)
+        tendencies = self._f(self._time[0], self._data[:, 0])[:, np.newaxis]
+        for i in range(1, self._data.shape[-1]):
+            new_tendency = self._f(self._time[i], self._data[:, i])[:, np.newaxis]
+            tendencies = np.concatenate((tendencies, new_tendency), axis=-1)
+
+        self._tendencies = tendencies
+
+
+    def _get_diagnostic(self,dimensional):
+
+        vr = self._model_params.variables_range
+
+
+        sigma = self._model_params.atmospheric_params.sig0
+        hd=self._model_params.atemperature_params.hd
+        thetas = self._model_params.atemperature_params.thetas
+        if dimensional:
+            self._diagnostic_data=*self._model_params.scale_params.deltap*self._model_params.scale_params.f0
+        else:
+            self._diagnostic_data=
+        return self._diagnostic_data
+
+
 if __name__ == '__main__':
     from qgs.params.params import QgParams
     from qgs.integrators.integrator import RungeKuttaIntegrator

From ea6c0302a6cfeb3c41fc3e2698752142908bded3 Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Thu, 1 Dec 2022 22:17:49 +0100
Subject: [PATCH 2/7] First working diagnostic class computing the middle layer
 vertical wind intensity ! To be validated !

---
 qgs/diagnostics/base.py           |   2 +-
 qgs/diagnostics/wind.py           |  86 ++--
 qgs/functions/tendencies.py       |  82 ++++
 qgs/params/params.py              |   2 +
 qgs/tensors/atmo_thermo_tensor.py | 627 ++++++++++++++++++++++++++++++
 5 files changed, 773 insertions(+), 26 deletions(-)
 create mode 100644 qgs/tensors/atmo_thermo_tensor.py

diff --git a/qgs/diagnostics/base.py b/qgs/diagnostics/base.py
index 3a93f46..4ba6242 100644
--- a/qgs/diagnostics/base.py
+++ b/qgs/diagnostics/base.py
@@ -321,7 +321,7 @@ def plot(self, time_index, style="image", ax=None, figsize=(16, 9),
         elif style == "contour":
             if self._orography and oro_kwargs is not False and self._oro_basis is not None:
                 hk = np.array(self._model_params.ground_params.hk, dtype=np.float)
-                oro = hk @ np.swapaxes(self._oro_basis, 0, 1) * 8500  # average height in meters of the 500 hPa pressure level at midlatitude
+                oro = hk @ np.swapaxes(self._oro_basis, 0, 1) * self._model_params.scale_params.Ha
                 im = ax.imshow(oro, origin='lower',
                                extent=[0, 2 * np.pi / self._model_params.scale_params.n, 0, np.pi], **oro_kwargs)
                 if color_bar:
diff --git a/qgs/diagnostics/wind.py b/qgs/diagnostics/wind.py
index 0986e58..9127d80 100644
--- a/qgs/diagnostics/wind.py
+++ b/qgs/diagnostics/wind.py
@@ -23,11 +23,12 @@
 import warnings
 
 import numpy as np
+from numba import njit
 import matplotlib.pyplot as plt
 
 from qgs.diagnostics.differential import DifferentialFieldDiagnostic
 from qgs.diagnostics.util import create_grid_basis
-from qgs.functions.tendencies import create_tendencies
+from qgs.functions.tendencies import create_tendencies, create_atmo_thermo_tendencies
 
 
 class AtmosphericWindDiagnostic(DifferentialFieldDiagnostic):
@@ -110,6 +111,11 @@ def _configure(self, delta_x=None, delta_y=None):
         elif self.type == "U":
             self._configure_differential_grid(basis, "dy", 1, delta_x, delta_y)
 
+        elif self.type == "W":
+            self._compute_grid(delta_x, delta_y)
+            self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
+
+
         elif self.type is None:
             warnings.warn("AtmosphericWindDiagnostic: Basis type note specified." +
                           " Unable to configure the diagnostic properly.")
@@ -633,50 +639,77 @@ def _get_diagnostic(self, dimensional):
         return self._diagnostic_data
 
 
-class VerticalVelocity(AtmosphericWindDiagnostic):
+class MiddleLayerVerticalVelocity(AtmosphericWindDiagnostic):
+    """Diagnostic giving the middle atmospheric layer vertical wind intensity fields.
+
+    Parameters
+    ----------
+
+    model_params: QgParams
+        An instance of the model parameters.
+    delta_x: float, optional
+        Spatial step in the zonal direction `x` for the gridded representation of the field.
+        If not provided, take an optimal guess based on the provided model's parameters.
+    delta_y: float, optional
+        Spatial step in the meridional direction `y` for the gridded representation of the field.
+        If not provided, take an optimal guess based on the provided model's parameters.
+    dimensional: bool, optional
+        Indicate if the output diagnostic must be dimensionalized or not.
+        Default to `True`.
+
+    Attributes
+    ----------
+
+    dimensional: bool
+        Indicate if the output diagnostic must be dimensionalized or not.
+
+    """
 
-    def __init__(self, model_params, f=None, atmospheric_inner_products=None, delta_x=None, delta_y=None, dimensional=True):
+    def __init__(self, model_params, delta_x=None, delta_y=None, dimensional=True):
 
         self.type = "W"
+
         AtmosphericWindDiagnostic.__init__(self, model_params, delta_x, delta_y, dimensional)
+
         self._plot_title = r'Atmospheric vertical wind in the middle layer'
+        self._plot_units = r" (in " + r'Pa s$^{-1}$' + r")"
 
-        if f is None or atmospheric_inner_products is None:
-            self._f, _, self._aips = create_tendencies(model_params, return_inner_products=True)
-        else:
-            self._f = f
-            self._aips = atmospheric_inner_products
+        self._f, _ = create_tendencies(model_params)
+        self._f_thermo = create_atmo_thermo_tendencies(model_params)
 
     def set_data(self, time, data):
 
-        AtmosphericWindDiagnostic.set_data(self, time, data)
-        tendencies = self._f(self._time[0], self._data[:, 0])[:, np.newaxis]
-        for i in range(1, self._data.shape[-1]):
-            new_tendency = self._f(self._time[i], self._data[:, i])[:, np.newaxis]
-            tendencies = np.concatenate((tendencies, new_tendency), axis=-1)
+        self._time = time
+        self._data = _compute_omega_term(time, data, self._f, self._f_thermo)
+        self._data = 2 * self._data / self._model_params.atmospheric_params.sig0
 
-        self._tendencies = tendencies
-
-
-    def _get_diagnostic(self,dimensional):
+    def _get_diagnostic(self, dimensional):
 
         vr = self._model_params.variables_range
+        omega = np.swapaxes(self._data[vr[0]:vr[1], ...].T @ np.swapaxes(self._grid_basis, 0, 1), 0, 1)
 
-
-        sigma = self._model_params.atmospheric_params.sig0
-        hd=self._model_params.atemperature_params.hd
-        thetas = self._model_params.atemperature_params.thetas
         if dimensional:
-            self._diagnostic_data=*self._model_params.scale_params.deltap*self._model_params.scale_params.f0
+            self._diagnostic_data = omega * self._model_params.scale_params.deltap * self._model_params.scale_params.f0
         else:
-            self._diagnostic_data=
+            self._diagnostic_data = omega
+
         return self._diagnostic_data
 
 
+@njit
+def _compute_omega_term(time, data, func, thermo_func):
+    tendencies = np.zeros_like(data)
+    thermo_tendencies = np.zeros_like(tendencies)
+    for i in range(data.shape[-1]):
+        tendencies[:, i] = func(time[i], data[:, i])
+        thermo_tendencies[:, i] = thermo_func(time[i], data[:, i])
+
+    return tendencies - thermo_tendencies
+
+
 if __name__ == '__main__':
     from qgs.params.params import QgParams
     from qgs.integrators.integrator import RungeKuttaIntegrator
-    from qgs.functions.tendencies import create_tendencies
 
     pars = QgParams()
     pars.set_atmospheric_channel_fourier_modes(2, 2)
@@ -684,7 +717,7 @@ def _get_diagnostic(self,dimensional):
     integrator = RungeKuttaIntegrator()
     integrator.set_func(f)
     ic = np.random.rand(pars.ndim) * 0.1
-    integrator.integrate(0., 200000., 0.1, ic=ic, write_steps=5)
+    integrator.integrate(0., 10000., 0.1, ic=ic, write_steps=5)
     time, traj = integrator.get_trajectories()
     integrator.terminate()
 
@@ -714,3 +747,6 @@ def _get_diagnostic(self,dimensional):
 
     psi1_wind = UpperLayerAtmosphericWindIntensityDiagnostic(pars)
     psi1_wind(time, traj)
+
+    vert_wind = MiddleLayerVerticalVelocity(pars)
+    vert_wind(time, traj)
diff --git a/qgs/functions/tendencies.py b/qgs/functions/tendencies.py
index 4a29165..8ecd054 100644
--- a/qgs/functions/tendencies.py
+++ b/qgs/functions/tendencies.py
@@ -13,6 +13,7 @@
 from qgs.inner_products.analytic import AtmosphericAnalyticInnerProducts, OceanicAnalyticInnerProducts, GroundAnalyticInnerProducts
 from qgs.inner_products.symbolic import AtmosphericSymbolicInnerProducts, OceanicSymbolicInnerProducts, GroundSymbolicInnerProducts
 from qgs.tensors.qgtensor import QgsTensor, QgsTensorDynamicT, QgsTensorT4
+from qgs.tensors.atmo_thermo_tensor import AtmoThermoTensor, AtmoThermoTensorDynamicT, AtmoThermoTensorT4
 from qgs.functions.sparse_mul import sparse_mul5, sparse_mul4, sparse_mul3, sparse_mul2
 
 
@@ -129,6 +130,87 @@ def Df(t, x):
     return ret
 
 
+def create_atmo_thermo_tendencies(params, return_atmo_thermo_tensor=False):
+    """Returns a function which return a part of the atmospheric thermodynamic tendencies useful for computing the vertical wind.
+
+    Parameters
+    ----------
+    params: QgParams
+        The parameters fully specifying the model configuration.
+    return_atmo_thermo_tensor: bool
+        If True, return the tendencies tensor of these tencencies. Default to False.
+
+
+    Returns
+    -------
+    f: callable
+        The numba-jitted tendencies function.
+    atmo_thermo_tensor: AtmoThermoTensor
+        If `return_atmo_thermo_tensor` is True, the tendencies tensor of the system.
+    """
+
+    if params.ablocks is not None:
+        aip = AtmosphericAnalyticInnerProducts(params)
+    elif params.atmospheric_basis is not None:
+        aip = AtmosphericSymbolicInnerProducts(params)
+    else:
+        aip = None
+
+    if params.oblocks is not None:
+        oip = OceanicAnalyticInnerProducts(params)
+    elif params.oceanic_basis is not None:
+        oip = OceanicSymbolicInnerProducts(params)
+    else:
+        oip = None
+
+    if params.gblocks is not None:
+        gip = GroundAnalyticInnerProducts(params)
+    elif params.ground_basis is not None:
+        gip = GroundSymbolicInnerProducts(params)
+    else:
+        gip = None
+
+    if aip is not None and oip is not None:
+        if not aip.connected_to_ocean:
+            aip.connect_to_ocean(oip)
+    elif aip is not None and gip is not None:
+        if not aip.connected_to_ground:
+            aip.connect_to_ground(gip)
+
+    if params.T4:
+        agotensor = AtmoThermoTensorT4(params, aip, oip, gip)
+    elif params.dynamic_T:
+        agotensor = AtmoThermoTensorDynamicT(params, aip, oip, gip)
+    else:
+        agotensor = AtmoThermoTensor(params, aip, oip, gip)
+
+    coo = agotensor.tensor.coords.T
+    val = agotensor.tensor.data
+
+    if params.T4 or params.dynamic_T:
+        @njit
+        def f(t, x):
+            xx = np.concatenate((np.full((1,), 1.), x))
+            xr = sparse_mul5(coo, val, xx, xx, xx, xx)
+            return xr[1:]
+
+    else:
+        @njit
+        def f(t, x):
+            xx = np.concatenate((np.full((1,), 1.), x))
+            xr = sparse_mul3(coo, val, xx, xx)
+            return xr[1:]
+
+    if return_atmo_thermo_tensor:
+        ret = list()
+        ret.append(f)
+        ret.append(agotensor)
+    else:
+        ret = f
+
+    return ret
+
+
 if __name__ == '__main__':
     from qgs.params.params import QgParams
 
diff --git a/qgs/params/params.py b/qgs/params/params.py
index 44eae2a..1b63119 100644
--- a/qgs/params/params.py
+++ b/qgs/params/params.py
@@ -257,6 +257,8 @@ def __init__(self, dic=None):
         self.phi0_npi = Parameter(0.25e0, input_dimensional=False, description="latitude expressed in fraction of pi")
         self.deltap = Parameter(5.e4, units='[Pa]', description='pressure difference between the two atmospheric layers',
                                 return_dimensional=True)
+        self.Ha = Parameter(8500., units='[m]', description="Average height of the 500 hPa pressure level at midlatitude",
+                            return_dimensional=True)
         self.set_params(dic)
 
     # ----------------------------------------
diff --git a/qgs/tensors/atmo_thermo_tensor.py b/qgs/tensors/atmo_thermo_tensor.py
new file mode 100644
index 0000000..f0d0138
--- /dev/null
+++ b/qgs/tensors/atmo_thermo_tensor.py
@@ -0,0 +1,627 @@
+"""
+    qgs tensor module
+    =================
+
+    This module computes and holds the tensors representing the tendencies of the model's equations.
+
+    TODO: Add a list of the different tensor available
+
+"""
+
+import numpy as np
+import sparse as sp
+
+from qgs.tensors.qgtensor import QgsTensor
+
+real_eps = np.finfo(np.float64).eps
+
+
+class AtmoThermoTensor(QgsTensor):
+    """Atmospheric thermodynamic tendencies tensor class.
+
+    Parameters
+    ----------
+    params: None or QgParams, optional
+        The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`.
+    atmospheric_inner_products: None or AtmosphericInnerProducts, optional
+        The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected.
+        If None, disable the atmospheric tendencies. Default to `None`.
+    oceanic_inner_products: None or OceanicInnerProducts, optional
+        The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected.
+        If None, disable the oceanic tendencies. Default to `None`.
+    ground_inner_products: None or GroundInnerProducts, optional
+        The inner products of the ground basis functions on which the model's PDE ground equations are projected.
+        If None, disable the ground tendencies. Default to `None`.
+
+    Attributes
+    ----------
+    params: None or QgParams
+        The models parameters used to configure the tensor. `None` for an empty tensor.
+    atmospheric_inner_products: None or AtmosphericInnerProducts
+        The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected.
+        If None, disable the atmospheric tendencies. Default to `None`.
+    oceanic_inner_products: None or OceanicInnerProducts
+        The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected.
+        If None, disable the oceanic tendencies. Default to `None`.
+    ground_inner_products: None or GroundInnerProducts
+        The inner products of the ground basis functions on which the model's PDE ground equations are projected.
+        If None, disable the ground tendencies. Default to `None`.
+    tensor: sparse.COO(float)
+        The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components.
+    jacobian_tensor: sparse.COO(float)
+        The jacobian tensor :math:`\\mathcal{T}_{i,j,k} + \\mathcal{T}_{i,k,j}` :math:`i`-th components.
+    """
+
+    def __init__(self, params=None, atmospheric_inner_products=None, oceanic_inner_products=None, ground_inner_products=None):
+
+        QgsTensor.__init__(self, params, atmospheric_inner_products, oceanic_inner_products, ground_inner_products)
+
+    def _compute_tensor_dicts(self):
+
+        if self.params is None:
+            return None
+
+        if self.atmospheric_inner_products is None and self.oceanic_inner_products is None \
+                and self.ground_inner_products is None:
+            return None
+
+        aips = self.atmospheric_inner_products
+        par = self.params
+        atp = par.atemperature_params
+        ap = par.atmospheric_params
+        op = par.oceanic_params
+        scp = par.scale_params
+        gp = par.ground_params
+        nvar = par.number_of_variables
+        ndim = par.ndim
+
+        bips = None
+        if self.oceanic_inner_products is not None:
+            bips = self.oceanic_inner_products
+            ocean = True
+        else:
+            ocean = False
+
+        if self.ground_inner_products is not None:
+            bips = self.ground_inner_products
+            ground_temp = True
+        else:
+            ground_temp = False
+
+        if self.params.dynamic_T:
+            offset = 1
+        else:
+            offset = 0
+
+        # constructing some derived matrices
+        if aips is not None:
+            a_theta = np.zeros((nvar[1], nvar[1]))
+            for i in range(nvar[1]):
+                for j in range(nvar[1]):
+                    a_theta[i, j] = aips.u(i, j)
+            a_theta = np.linalg.inv(a_theta)
+            a_theta = sp.COO(a_theta)
+
+        #################
+
+        if bips is not None:
+            go = bips.stored
+        else:
+            go = True
+
+        sparse_arrays_dict = dict()
+
+        if aips.stored and go:
+
+            # theta_a part
+            for i in range(nvar[1]):
+                t = sp.zeros((ndim + 1, ndim + 1), dtype=np.float64, format='dok')
+
+                if par.Cpa is not None:
+                    t[0, 0] += par.Cpa[i]
+
+                if atp.hd is not None and atp.thetas is not None:
+                    t[0, 0] += atp.thetas[i] * atp.hd
+
+                for j in range(nvar[0]):
+
+                    jo = j + offset  # skipping the theta 0 variable if it exists
+
+                    for k in range(nvar[0]):
+                        ko = k + offset  # skipping the theta 0 variable if it exists
+
+                        val = a_theta[i, :] @ aips._g[:, jo, ko]
+                        t[self._psi_a(j), self._theta_a(ko)] -= val
+
+                for j in range(nvar[1]):
+                    val = a_theta[i, :] @ aips._u[:, j]
+                    if par.Lpa is not None:
+                        t[self._theta_a(j), 0] -= val * atp.sc * par.Lpa
+                    if par.LSBpa is not None:
+                        t[self._theta_a(j), 0] -= val * par.LSBpa
+
+                    if atp.hd is not None:
+                        t[self._theta_a(j), 0] -= val * atp.hd
+
+                if ocean:
+                    for j in range(nvar[2]):
+                        jo = j + offset  # skipping the theta 0 variable if it exists
+
+                    if par.Lpa is not None:
+                        for j in range(nvar[3]):
+                            val = a_theta[i, :] @ aips._s[:, j]
+                            t[self._deltaT_o(j), 0] += val * par.Lpa / 2
+                            if par.LSBpgo is not None:
+                                t[self._deltaT_o(j), 0] += val * par.LSBpgo
+
+                if ground_temp:
+                    if par.Lpa is not None:
+                        for j in range(nvar[2]):
+                            val = a_theta[i, :] @ aips._s[:, j]
+                            t[self._deltaT_g(j), 0] += val * par.Lpa / 2
+                            if par.LSBpgo is not None:
+                                t[self._deltaT_g(j), 0] += val * par.LSBpgo
+
+                sparse_arrays_dict[self._theta_a(i)] = t.to_coo()
+
+        else:
+
+            # theta_a part
+            for i in range(nvar[1]):
+                t = sp.zeros((ndim + 1, ndim + 1), dtype=np.float64, format='dok')
+
+                if par.Cpa is not None:
+                    t[0, 0] += par.Cpa[i]
+
+                if atp.hd is not None and atp.thetas is not None:
+                    t[0, 0] += atp.thetas[i] * atp.hd
+
+                for j in range(nvar[0]):
+
+                    jo = j + offset  # skipping the theta 0 variable if it exists
+
+                    for k in range(nvar[0]):
+                        ko = k + offset  # skipping the theta 0 variable if it exists
+
+                        val = 0
+                        for jj in range(nvar[1]):
+                            val += a_theta[i, jj] * aips.g(jj, jo, ko)
+
+                        t[self._psi_a(j), self._theta_a(ko)] -= val
+
+                for j in range(nvar[1]):
+                    val = 0
+                    for jj in range(nvar[1]):
+                        val += a_theta[i, jj] * aips.u(jj, j)
+
+                    if par.Lpa is not None:
+                        t[self._theta_a(j), 0] -= val * atp.sc * par.Lpa
+                    if par.LSBpa is not None:
+                        t[self._theta_a(j), 0] -= val * par.LSBpa
+
+                    if atp.hd is not None:
+                        t[self._theta_a(j), 0] -= val * atp.hd
+
+                if ocean:
+                    if par.Lpa is not None:
+                        for j in range(nvar[3]):
+                            val = 0
+                            for jj in range(nvar[1]):
+                                val += a_theta[i, jj] * aips.s(jj, j)
+                            t[self._deltaT_o(j), 0] += val * par.Lpa / 2
+                            if par.LSBpgo is not None:
+                                t[self._deltaT_o(j), 0] += val * par.LSBpgo
+
+                if ground_temp:
+                    if par.Lpa is not None:
+                        for j in range(nvar[2]):
+                            val = 0
+                            for jj in range(nvar[1]):
+                                val += a_theta[i, jj] * aips.s(jj, j)
+                            t[self._deltaT_g(j), 0] += val * par.Lpa / 2
+                            if par.LSBpgo is not None:
+                                t[self._deltaT_g(j), 0] += val * par.LSBpgo
+
+                sparse_arrays_dict[self._theta_a(i)] = t.to_coo()
+
+        return sparse_arrays_dict
+
+
+class AtmoThermoTensorDynamicT(AtmoThermoTensor):
+    """Atmospheric thermodynamic dynamical temperature first order (linear) tendencies tensor class.
+
+    Parameters
+    ----------
+    params: None or QgParams, optional
+        The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`.
+    atmospheric_inner_products: None or AtmosphericInnerProducts, optional
+        The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected.
+        If None, disable the atmospheric tendencies. Default to `None`.
+    oceanic_inner_products: None or OceanicInnerProducts, optional
+        The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected.
+        If None, disable the oceanic tendencies. Default to `None`.
+    ground_inner_products: None or GroundInnerProducts, optional
+        The inner products of the ground basis functions on which the model's PDE ground equations are projected.
+        If None, disable the ground tendencies. Default to `None`.
+
+    Attributes
+    ----------
+    params: None or QgParams
+        The models parameters used to configure the tensor. `None` for an empty tensor.
+    atmospheric_inner_products: None or AtmosphericInnerProducts
+        The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected.
+        If None, disable the atmospheric tendencies. Default to `None`.
+    oceanic_inner_products: None or OceanicInnerProducts
+        The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected.
+        If None, disable the oceanic tendencies. Default to `None`.
+    ground_inner_products: None or GroundInnerProducts
+        The inner products of the ground basis functions on which the model's PDE ground equations are projected.
+        If None, disable the ground tendencies. Default to `None`.
+    tensor: sparse.COO(float)
+        The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components.
+    jacobian_tensor: sparse.COO(float)
+        The jacobian tensor :math:`\\mathcal{T}_{i,j,k} + \\mathcal{T}_{i,k,j}` :math:`i`-th components.
+    """
+
+    def __init__(self, params=None, atmospheric_inner_products=None, oceanic_inner_products=None, ground_inner_products=None):
+
+        AtmoThermoTensor.__init__(self, params, atmospheric_inner_products, oceanic_inner_products, ground_inner_products)
+
+    def _compute_tensor_dicts(self):
+
+        if self.params is None:
+            return None
+
+        if self.atmospheric_inner_products is None and self.oceanic_inner_products is None \
+                and self.ground_inner_products is None:
+            return None
+
+        aips = self.atmospheric_inner_products
+
+        bips = None
+        if self.oceanic_inner_products is not None:
+            bips = self.oceanic_inner_products
+
+        elif self.ground_inner_products is not None:
+            bips = self.ground_inner_products
+
+        if bips is not None:
+            go = bips.stored
+        else:
+            go = True
+
+        if aips.stored and go:
+
+            sparse_arrays_full_dict = self._compute_stored_full_dict()
+
+        else:
+
+            sparse_arrays_full_dict = self._compute_non_stored_full_dict()
+
+        return sparse_arrays_full_dict
+
+    def _compute_stored_full_dict(self):
+        par = self.params
+        ap = par.atmospheric_params
+        nvar = par.number_of_variables
+        aips = self.atmospheric_inner_products
+
+        bips = None
+        if self.oceanic_inner_products is not None:
+            bips = self.oceanic_inner_products
+            ocean = True
+        else:
+            ocean = False
+
+        if self.ground_inner_products is not None:
+            bips = self.ground_inner_products
+            ground_temp = True
+        else:
+            ground_temp = False
+
+        # constructing identity matrices
+        if aips is not None:
+            if aips is not None:
+                a_theta = np.zeros((nvar[1], nvar[1]))
+                for i in range(nvar[1]):
+                    for j in range(nvar[1]):
+                        a_theta[i, j] = aips.u(i, j)
+                a_theta = np.linalg.inv(a_theta)
+                a_theta = sp.COO(a_theta)
+
+        #################
+
+        sparse_arrays_full_dict = dict()
+        # theta_a part
+        for i in range(nvar[1]):
+            sparse_arrays_full_dict[self._theta_a(i)] = list()
+
+            if par.T4LSBpa is not None:
+                val = sp.tensordot(a_theta[i], aips._z, axes=1)
+                if val.nnz > 0:
+                    sparse_arrays_full_dict[self._theta_a(i)].append(self._shift_tensor_coordinates(- par.T4LSBpa * val, self._theta_a(0)))
+
+            if ocean:
+                if par.T4LSBpgo is not None:
+                    val = sp.tensordot(a_theta[i], aips._v, axes=1)
+                    if val.nnz > 0:
+                        sparse_arrays_full_dict[self._theta_a(i)].append(self._shift_tensor_coordinates(par.T4LSBpgo * val, self._deltaT_o(0)))
+
+            if ground_temp:
+
+                if par.T4LSBpgo is not None:
+                    val = sp.tensordot(a_theta[i], aips._v, axes=1)
+                    if val.nnz > 0:
+                        sparse_arrays_full_dict[self._theta_a(i)].append(self._shift_tensor_coordinates(par.T4LSBpgo * val, self._deltaT_g(0)))
+
+        return sparse_arrays_full_dict
+
+    def _compute_non_stored_full_dict(self):
+        par = self.params
+        ap = par.atmospheric_params
+        nvar = par.number_of_variables
+        ndim = par.ndim
+        aips = self.atmospheric_inner_products
+
+        bips = None
+        if self.oceanic_inner_products is not None:
+            bips = self.oceanic_inner_products
+            ocean = True
+        else:
+            ocean = False
+
+        if self.ground_inner_products is not None:
+            bips = self.ground_inner_products
+            ground_temp = True
+        else:
+            ground_temp = False
+
+        # constructing some derived matrices
+        if aips is not None:
+
+            a_theta = np.zeros((nvar[1], nvar[1]))
+            for i in range(nvar[1]):
+                for j in range(nvar[1]):
+                    a_theta[i, j] = ap.sig0 * aips.a(i, j) - aips.u(i, j)
+            a_theta = np.linalg.inv(a_theta)
+            a_theta = sp.COO(a_theta)
+
+        #################
+
+        sparse_arrays_full_dict = dict()
+        # theta_a part
+        for i in range(nvar[1]):
+            t_full = sp.zeros((ndim + 1, ndim + 1, ndim + 1, ndim + 1), dtype=np.float64, format='dok')
+
+            if par.T4LSBpa is not None:
+                j = k = ell = 0
+                for m in range(nvar[1]):
+                    val = 0
+                    for jj in range(nvar[1]):
+                        val -= a_theta[i, jj] * aips.z(jj, j, k, ell, m)
+                    if m == 0:
+                        t_full[self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)] = par.T4LSBpa * val
+                    else:
+                        t_full[self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)] = 4 * par.T4LSBpa * val
+
+            if ocean:
+                if par.T4LSBpgo is not None:
+                    j = k = ell = 0
+                    for m in range(nvar[3]):
+                        val = 0
+                        for jj in range(nvar[1]):
+                            val += a_theta[i, jj] * aips.v(jj, j, k, ell, m)
+                        if m == 0:
+                            t_full[self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)] = par.T4LSBpgo * val
+                        else:
+                            t_full[self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)] = 4 * par.T4LSBpgo * val
+
+            if ground_temp:
+                if par.T4LSBpgo is not None:
+                    j = k = ell = 0
+                    for m in range(nvar[2]):
+                        val = 0
+                        for jj in range(nvar[1]):
+                            val += a_theta[i, jj] * aips.v(jj, j, k, ell, m)
+                        if m == 0:
+                            t_full[self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)] = par.T4LSBpgo * val
+                        else:
+                            t_full[self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)] = 4 * par.T4LSBpgo * val
+
+            sparse_arrays_full_dict[self._theta_a(i)] = t_full.to_coo()
+
+        return sparse_arrays_full_dict
+
+    def compute_tensor(self):
+        """Routine to compute the tensor."""
+        # gathering
+        par = self.params
+        ndim = par.ndim
+
+        sparse_arrays_dict = QgsTensor._compute_tensor_dicts(self)
+        sparse_arrays_full_dict = self._compute_tensor_dicts()
+
+        tensor = sp.zeros((ndim + 1, ndim + 1, ndim + 1, ndim + 1, ndim + 1), dtype=np.float64, format='coo')
+        if sparse_arrays_dict is not None:
+            tensor = self._add_dict_to_tensor(sparse_arrays_dict, tensor)
+        if sparse_arrays_full_dict is not None:
+            tensor = self._add_dict_to_tensor(sparse_arrays_full_dict, tensor)
+        self._set_tensor(tensor)
+
+
+class AtmoThermoTensorT4(AtmoThermoTensorDynamicT):
+    """Atmospheric thermodynamic :math:`T^4` tendencies tensor class. Implies dynamical zeroth-order temperature.
+
+    Parameters
+    ----------
+    params: None or QgParams, optional
+        The models parameters to configure the tensor. `None` to initialize an empty tensor. Default to `None`.
+    atmospheric_inner_products: None or AtmosphericInnerProducts, optional
+        The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected.
+        If None, disable the atmospheric tendencies. Default to `None`.
+    oceanic_inner_products: None or OceanicInnerProducts, optional
+        The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected.
+        If None, disable the oceanic tendencies. Default to `None`.
+    ground_inner_products: None or GroundInnerProducts, optional
+        The inner products of the ground basis functions on which the model's PDE ground equations are projected.
+        If None, disable the ground tendencies. Default to `None`.
+
+    Attributes
+    ----------
+    params: None or QgParams
+        The models parameters used to configure the tensor. `None` for an empty tensor.
+    atmospheric_inner_products: None or AtmosphericInnerProducts
+        The inner products of the atmospheric basis functions on which the model's PDE atmospheric equations are projected.
+        If None, disable the atmospheric tendencies. Default to `None`.
+    oceanic_inner_products: None or OceanicInnerProducts
+        The inner products of the oceanic basis functions on which the model's PDE oceanic equations are projected.
+        If None, disable the oceanic tendencies. Default to `None`.
+    ground_inner_products: None or GroundInnerProducts
+        The inner products of the ground basis functions on which the model's PDE ground equations are projected.
+        If None, disable the ground tendencies. Default to `None`.
+    tensor: sparse.COO(float)
+        The tensor :math:`\\mathcal{T}_{i,j,k}` :math:`i`-th components.
+    jacobian_tensor: sparse.COO(float)
+        The jacobian tensor :math:`\\mathcal{T}_{i,j,k} + \\mathcal{T}_{i,k,j}` :math:`i`-th components.
+    """
+
+    def __init__(self, params=None, atmospheric_inner_products=None, oceanic_inner_products=None, ground_inner_products=None):
+
+        AtmoThermoTensorDynamicT.__init__(self, params, atmospheric_inner_products, oceanic_inner_products, ground_inner_products)
+
+    def _compute_non_stored_full_dict(self):
+
+        if self.params is None:
+            return None
+
+        if self.atmospheric_inner_products is None and self.oceanic_inner_products is None \
+                and self.ground_inner_products is None:
+            return None
+
+        aips = self.atmospheric_inner_products
+        par = self.params
+        ap = par.atmospheric_params
+        nvar = par.number_of_variables
+        ndim = par.ndim
+
+        bips = None
+        if self.oceanic_inner_products is not None:
+            bips = self.oceanic_inner_products
+            ocean = True
+        else:
+            ocean = False
+
+        if self.ground_inner_products is not None:
+            bips = self.ground_inner_products
+            ground_temp = True
+        else:
+            ground_temp = False
+
+        # constructing some derived matrices
+        if aips is not None:
+
+            a_theta = np.zeros((nvar[1], nvar[1]))
+            for i in range(nvar[1]):
+                for j in range(nvar[1]):
+                    a_theta[i, j] = ap.sig0 * aips.a(i, j) - aips.u(i, j)
+            a_theta = np.linalg.inv(a_theta)
+            a_theta = sp.COO(a_theta)
+
+        #################
+
+        sparse_arrays_full_dict = dict()
+
+        # theta_a part
+        for i in range(nvar[1]):
+            t_full = sp.zeros((ndim + 1, ndim + 1, ndim + 1, ndim + 1), dtype=np.float64, format='dok')
+
+            if par.T4LSBpa is not None:
+                for j in range(nvar[1]):
+                    for k in range(nvar[1]):
+                        for ell in range(nvar[1]):
+                            for m in range(nvar[1]):
+                                val = 0
+                                for jj in range(nvar[1]):
+                                    val -= a_theta[i, jj] * aips.z(jj, j, k, ell, m)
+                                t_full[self._theta_a(j), self._theta_a(k), self._theta_a(ell), self._theta_a(m)] = par.T4LSBpa * val
+
+            if ocean:
+                if par.T4LSBpgo is not None:
+                    for j in range(nvar[3]):
+                        for k in range(nvar[3]):
+                            for ell in range(nvar[3]):
+                                for m in range(nvar[3]):
+                                    val = 0
+                                    for jj in range(nvar[1]):
+                                        val += a_theta[i, jj] * aips.v(jj, j, k, ell, m)
+                                    t_full[self._deltaT_o(j), self._deltaT_o(k), self._deltaT_o(ell), self._deltaT_o(m)] = par.T4LSBpgo * val
+
+            if ground_temp:
+                if par.T4LSBpgo is not None:
+                    for j in range(nvar[2]):
+                        for k in range(nvar[2]):
+                            for ell in range(nvar[2]):
+                                for m in range(nvar[2]):
+                                    val = 0
+                                    for jj in range(nvar[1]):
+                                        val += a_theta[i, jj] * aips.v(jj, j, k, ell, m)
+                                    t_full[self._deltaT_g(j), self._deltaT_g(k), self._deltaT_g(ell), self._deltaT_g(m)] = par.T4LSBpgo * val
+
+            sparse_arrays_full_dict[self._theta_a(i)] = t_full.to_coo()
+
+        return sparse_arrays_full_dict
+
+
+def _kronecker_delta(i, j):
+
+    if i == j:
+        return 1
+
+    else:
+        return 0
+
+
+if __name__ == '__main__':
+    from qgs.params.params import QgParams
+    from qgs.inner_products.analytic import AtmosphericAnalyticInnerProducts, OceanicAnalyticInnerProducts
+    from qgs.inner_products.symbolic import AtmosphericSymbolicInnerProducts, OceanicSymbolicInnerProducts
+
+    # Analytic test
+
+    params = QgParams({'rr': 287.e0, 'sb': 5.6e-8})
+    params.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
+    params.set_atmospheric_channel_fourier_modes(2, 2)
+    params.set_oceanic_basin_fourier_modes(2, 4)
+    aip = AtmosphericAnalyticInnerProducts(params)
+    oip = OceanicAnalyticInnerProducts(params)
+    aip.connect_to_ocean(oip)
+    agotensor = AtmoThermoTensor(params, aip, oip)
+
+    # Symbolic dynamic T test
+
+    params_t = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, dynamic_T=True)
+    params_t.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
+    params_t.set_atmospheric_channel_fourier_modes(2, 2, mode='symbolic')
+    params_t.set_oceanic_basin_fourier_modes(2, 4, mode='symbolic')
+
+    aip = AtmosphericSymbolicInnerProducts(params_t, quadrature=True)  # , stored=False)
+    oip = OceanicSymbolicInnerProducts(params_t, quadrature=True)  # , stored=False)
+    agotensor_t = AtmoThermoTensorDynamicT(params_t, aip, oip)
+
+    # Symbolic dynamic T4 test
+
+    params_t4 = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, T4=True)
+    params_t4.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
+    params_t4.set_atmospheric_channel_fourier_modes(2, 2, mode='symbolic')
+    params_t4.set_oceanic_basin_fourier_modes(2, 4, mode='symbolic')
+
+    aip = AtmosphericSymbolicInnerProducts(params_t4, quadrature=True)  # , stored=False)
+    oip = OceanicSymbolicInnerProducts(params_t4, quadrature=True)  # , stored=False)
+
+    aip.save_to_file("aip.ip")
+    oip.save_to_file("oip.ip")
+
+    aip.load_from_file("aip.ip")
+    oip.load_from_file("oip.ip")
+
+    agotensor_t4 = AtmoThermoTensorT4(params_t4, aip, oip)

From 00d27213916cbba762bb3d764fbab7e2dc055bbf Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Sat, 10 Dec 2022 14:16:00 +0100
Subject: [PATCH 3/7] Update of the documentation with the vertical velocity
 fields diagnostic

---
 documentation/source/files/user_guide.rst |  7 ++++---
 qgs/diagnostics/wind.py                   | 14 +++++++-------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/documentation/source/files/user_guide.rst b/documentation/source/files/user_guide.rst
index a6ad713..e7f58fd 100644
--- a/documentation/source/files/user_guide.rst
+++ b/documentation/source/files/user_guide.rst
@@ -510,9 +510,10 @@ Note that it is also possible to use other ordinary differential equations integ
     * :class:`.MiddleAtmosphericUWindDiagnostic`: Diagnostic giving the middle atmospheric U wind fields :math:`- \partial_y \psi_{\rm a}`.
     * :class:`.UpperLayerAtmosphericVWindDiagnostic`: Diagnostic giving the upper layer atmospheric V wind fields :math:`\partial_x \psi^1_{\rm a}`.
     * :class:`.UpperLayerAtmosphericUWindDiagnostic`: Diagnostic giving the upper layer atmospheric U wind fields :math:`- \partial_y \psi^1_{\rm a}`.
-    * :class:`.LowerLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the lower layer atmospheric wind intensity fields.
-    * :class:`.MiddleAtmosphericWindIntensityDiagnostic`: Diagnostic giving the middle atmospheric wind intensity fields.
-    * :class:`.UpperLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the upper layer atmospheric wind intensity fields.
+    * :class:`.LowerLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the lower layer atmospheric horizontal wind intensity fields.
+    * :class:`.MiddleAtmosphericWindIntensityDiagnostic`: Diagnostic giving the middle atmospheric horizontal wind intensity fields.
+    * :class:`.UpperLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the upper layer atmospheric horizontal wind intensity fields.
+    * :class:`.MiddleLayerVerticalVelocity`: Diagnostic giving the middle atmospheric layer vertical wind intensity fields.
     * :class:`.MiddleAtmosphericEddyHeatFluxDiagnostic`: Diagnostic giving the middle atmospheric eddy heat flux field.
     * :class:`.MiddleAtmosphericEddyHeatFluxProfileDiagnostic`: Diagnostic giving the middle atmospheric eddy heat flux zonally averaged profile.
 
diff --git a/qgs/diagnostics/wind.py b/qgs/diagnostics/wind.py
index 9127d80..bb3c292 100644
--- a/qgs/diagnostics/wind.py
+++ b/qgs/diagnostics/wind.py
@@ -14,9 +14,10 @@
     * :class:`MiddleAtmosphericUWindDiagnostic`: Diagnostic giving the middle atmospheric U wind fields :math:`- \\partial_y \\psi_{\\rm a}`.
     * :class:`UpperLayerAtmosphericVWindDiagnostic`: Diagnostic giving the upper layer atmospheric V wind fields :math:`\\partial_x \\psi^1_{\\rm a}`.
     * :class:`UpperLayerAtmosphericUWindDiagnostic`: Diagnostic giving the upper layer atmospheric U wind fields :math:`- \\partial_y \\psi^1_{\\rm a}`.
-    * :class:`LowerLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the lower layer atmospheric wind intensity fields.
-    * :class:`MiddleAtmosphericWindIntensityDiagnostic`: Diagnostic giving the middle atmospheric wind intensity fields.
-    * :class:`UpperLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the upper layer atmospheric wind intensity fields.
+    * :class:`LowerLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the lower layer atmospheric horizontal wind intensity fields.
+    * :class:`MiddleAtmosphericWindIntensityDiagnostic`: Diagnostic giving the middle atmospheric horizontal wind intensity fields.
+    * :class:`UpperLayerAtmosphericWindIntensityDiagnostic`: Diagnostic giving the upper layer atmospheric horizontal wind intensity fields.
+    * :class:`MiddleLayerVerticalVelocity`: Diagnostic giving the middle atmospheric layer vertical wind intensity fields.
 
 """
 
@@ -115,7 +116,6 @@ def _configure(self, delta_x=None, delta_y=None):
             self._compute_grid(delta_x, delta_y)
             self._grid_basis = create_grid_basis(basis, self._X, self._Y, self._subs)
 
-
         elif self.type is None:
             warnings.warn("AtmosphericWindDiagnostic: Basis type note specified." +
                           " Unable to configure the diagnostic properly.")
@@ -478,7 +478,7 @@ def _get_diagnostic(self, dimensional):
 
 
 class LowerLayerAtmosphericWindIntensityDiagnostic(AtmosphericWindDiagnostic):
-    """Diagnostic giving the lower layer atmospheric wind intensity fields.
+    """Diagnostic giving the lower layer atmospheric horizontal wind intensity fields.
 
     Parameters
     ----------
@@ -531,7 +531,7 @@ def _get_diagnostic(self, dimensional):
 
 
 class MiddleAtmosphericWindIntensityDiagnostic(AtmosphericWindDiagnostic):
-    """Diagnostic giving the middle atmospheric wind intensity fields.
+    """Diagnostic giving the middle atmospheric horizontal wind intensity fields.
 
     Parameters
     ----------
@@ -584,7 +584,7 @@ def _get_diagnostic(self, dimensional):
 
 
 class UpperLayerAtmosphericWindIntensityDiagnostic(AtmosphericWindDiagnostic):
-    """Diagnostic giving the lower layer atmospheric wind intensity fields.
+    """Diagnostic giving the lower layer atmospheric horizontal wind intensity fields.
 
     Parameters
     ----------

From 6944a39a2a602dedba10d3dd3428fad7d5eb5349 Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Sat, 10 Dec 2022 15:20:22 +0100
Subject: [PATCH 4/7] Corrected a bug + Added the references

---
 .../source/files/technical/tensors.rst        |  4 ++
 qgs/diagnostics/wind.py                       |  2 +-
 qgs/tensors/atmo_thermo_tensor.py             | 62 +++++++++----------
 3 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/documentation/source/files/technical/tensors.rst b/documentation/source/files/technical/tensors.rst
index 93f89ba..3a20dee 100644
--- a/documentation/source/files/technical/tensors.rst
+++ b/documentation/source/files/technical/tensors.rst
@@ -7,3 +7,7 @@ Module holding the model's tendencies tensor encoding each of their additive ter
 .. automodule:: qgs.tensors.qgtensor
     :show-inheritance:
     :members:
+
+.. automodule:: qgs.tensors.atmo_thermo_tensor
+    :show-inheritance:
+    :members:
diff --git a/qgs/diagnostics/wind.py b/qgs/diagnostics/wind.py
index bb3c292..db7eae6 100644
--- a/qgs/diagnostics/wind.py
+++ b/qgs/diagnostics/wind.py
@@ -681,7 +681,7 @@ def set_data(self, time, data):
 
         self._time = time
         self._data = _compute_omega_term(time, data, self._f, self._f_thermo)
-        self._data = 2 * self._data / self._model_params.atmospheric_params.sig0
+        self._data = self._data / self._model_params.atmospheric_params.sig0
 
     def _get_diagnostic(self, dimensional):
 
diff --git a/qgs/tensors/atmo_thermo_tensor.py b/qgs/tensors/atmo_thermo_tensor.py
index f0d0138..f7d0357 100644
--- a/qgs/tensors/atmo_thermo_tensor.py
+++ b/qgs/tensors/atmo_thermo_tensor.py
@@ -1,8 +1,8 @@
 """
-    qgs tensor module
-    =================
+    qgs atmospheric thermodynamic tensor module
+    ===========================================
 
-    This module computes and holds the tensors representing the tendencies of the model's equations.
+    This module computes and holds the tensors representing the atmospheric tendencies of the model's equations.
 
     TODO: Add a list of the different tensor available
 
@@ -68,10 +68,6 @@ def _compute_tensor_dicts(self):
         aips = self.atmospheric_inner_products
         par = self.params
         atp = par.atemperature_params
-        ap = par.atmospheric_params
-        op = par.oceanic_params
-        scp = par.scale_params
-        gp = par.ground_params
         nvar = par.number_of_variables
         ndim = par.ndim
 
@@ -599,29 +595,29 @@ def _kronecker_delta(i, j):
 
     # Symbolic dynamic T test
 
-    params_t = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, dynamic_T=True)
-    params_t.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
-    params_t.set_atmospheric_channel_fourier_modes(2, 2, mode='symbolic')
-    params_t.set_oceanic_basin_fourier_modes(2, 4, mode='symbolic')
-
-    aip = AtmosphericSymbolicInnerProducts(params_t, quadrature=True)  # , stored=False)
-    oip = OceanicSymbolicInnerProducts(params_t, quadrature=True)  # , stored=False)
-    agotensor_t = AtmoThermoTensorDynamicT(params_t, aip, oip)
-
-    # Symbolic dynamic T4 test
-
-    params_t4 = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, T4=True)
-    params_t4.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
-    params_t4.set_atmospheric_channel_fourier_modes(2, 2, mode='symbolic')
-    params_t4.set_oceanic_basin_fourier_modes(2, 4, mode='symbolic')
-
-    aip = AtmosphericSymbolicInnerProducts(params_t4, quadrature=True)  # , stored=False)
-    oip = OceanicSymbolicInnerProducts(params_t4, quadrature=True)  # , stored=False)
-
-    aip.save_to_file("aip.ip")
-    oip.save_to_file("oip.ip")
-
-    aip.load_from_file("aip.ip")
-    oip.load_from_file("oip.ip")
-
-    agotensor_t4 = AtmoThermoTensorT4(params_t4, aip, oip)
+    # params_t = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, dynamic_T=True)
+    # params_t.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
+    # params_t.set_atmospheric_channel_fourier_modes(2, 2, mode='symbolic')
+    # params_t.set_oceanic_basin_fourier_modes(2, 4, mode='symbolic')
+    #
+    # aip = AtmosphericSymbolicInnerProducts(params_t, quadrature=True)  # , stored=False)
+    # oip = OceanicSymbolicInnerProducts(params_t, quadrature=True)  # , stored=False)
+    # agotensor_t = AtmoThermoTensorDynamicT(params_t, aip, oip)
+    #
+    # # Symbolic dynamic T4 test
+    #
+    # params_t4 = QgParams({'rr': 287.e0, 'sb': 5.6e-8}, T4=True)
+    # params_t4.set_params({'kd': 0.04, 'kdp': 0.04, 'n': 1.5})
+    # params_t4.set_atmospheric_channel_fourier_modes(2, 2, mode='symbolic')
+    # params_t4.set_oceanic_basin_fourier_modes(2, 4, mode='symbolic')
+    #
+    # aip = AtmosphericSymbolicInnerProducts(params_t4, quadrature=True)  # , stored=False)
+    # oip = OceanicSymbolicInnerProducts(params_t4, quadrature=True)  # , stored=False)
+    #
+    # # aip.save_to_file("aip.ip")
+    # # oip.save_to_file("oip.ip")
+    # #
+    # # aip.load_from_file("aip.ip")
+    # # oip.load_from_file("oip.ip")
+    #
+    # agotensor_t4 = AtmoThermoTensorT4(params_t4, aip, oip)

From e1c7b0833c3582298070245142c061282c469d3d Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Sat, 10 Dec 2022 15:32:59 +0100
Subject: [PATCH 5/7] Adapted vertical wind diagnostic to dynamical temperature
 settings

---
 qgs/diagnostics/wind.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/qgs/diagnostics/wind.py b/qgs/diagnostics/wind.py
index db7eae6..675c048 100644
--- a/qgs/diagnostics/wind.py
+++ b/qgs/diagnostics/wind.py
@@ -685,8 +685,13 @@ def set_data(self, time, data):
 
     def _get_diagnostic(self, dimensional):
 
+        if self._model_params.dynamic_T:
+            offset = 1
+        else:
+            offset = 0
+
         vr = self._model_params.variables_range
-        omega = np.swapaxes(self._data[vr[0]:vr[1], ...].T @ np.swapaxes(self._grid_basis, 0, 1), 0, 1)
+        omega = np.swapaxes(self._data[vr[0]+offset:vr[1], ...].T @ np.swapaxes(self._grid_basis, 0, 1), 0, 1)
 
         if dimensional:
             self._diagnostic_data = omega * self._model_params.scale_params.deltap * self._model_params.scale_params.f0

From ae0852da9fbaefa14d12ff3dd2f12369c62e33d4 Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Sun, 11 Dec 2022 10:52:47 +0100
Subject: [PATCH 6/7] Added a new notebook to show the vertical wind velocity

---
 notebooks/diagnostics/eddy_heat_flux.ipynb    |   6 +-
 .../diagnostics/vertical_wind_velocity.ipynb  | 495 ++++++++++++++++++
 2 files changed, 496 insertions(+), 5 deletions(-)
 create mode 100644 notebooks/diagnostics/vertical_wind_velocity.ipynb

diff --git a/notebooks/diagnostics/eddy_heat_flux.ipynb b/notebooks/diagnostics/eddy_heat_flux.ipynb
index 40241dc..ed2f54e 100644
--- a/notebooks/diagnostics/eddy_heat_flux.ipynb
+++ b/notebooks/diagnostics/eddy_heat_flux.ipynb
@@ -401,16 +401,12 @@
     "                 diagnostic_kwargs={'show_time': False, 'background': background},\n",
     "                 plot_kwargs={'ms': 0.2})\n",
     "\n",
-    "# m.add_diagnostic(vwind)#,\n",
-    "                 #diagnostic_kwargs={'style': 'contour', 'contour_labels': False},\n",
-    "                 #plot_kwargs={'colors': 'k'})\n",
-    "        \n",
     "m.add_diagnostic(eddy, plot_kwargs={'cmap':plt.get_cmap('magma')})\n",
     "m.add_diagnostic(eddyp, plot_kwargs={'figsize':(10,8)})\n",
     "\n",
     "m.add_diagnostic(temp)\n",
     "                 \n",
-    "m.set_data(reference_time, reference_traj)"
+    "m.set_data(reference_time[:2000], reference_traj[:,:2000])"
    ]
   },
   {
diff --git a/notebooks/diagnostics/vertical_wind_velocity.ipynb b/notebooks/diagnostics/vertical_wind_velocity.ipynb
new file mode 100644
index 0000000..a30d71d
--- /dev/null
+++ b/notebooks/diagnostics/vertical_wind_velocity.ipynb
@@ -0,0 +1,495 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Diagnostic: Vertical wind velocity in the middle layer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We show here the usage of the vertical wind diagnostics with is a simple 2-layer channel QG atmosphere truncated at wavenumber 2 on a beta-plane with a simple orography (a montain and a valley). "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Modules import"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "First, setting the path and loading of some modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys, os"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sys.path.extend([os.path.abspath('../../')])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from mpl_toolkits.mplot3d import Axes3D"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from matplotlib import rc\n",
+    "rc('font',**{'family':'serif','sans-serif':['Times'],'size':14})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Initializing the random number generator (for reproducibility). -- Disable if needed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.random.seed(210217)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Importing the model's modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qgs.params.params import QgParams\n",
+    "from qgs.integrators.integrator import RungeKuttaIntegrator, RungeKuttaTglsIntegrator\n",
+    "from qgs.functions.tendencies import create_tendencies\n",
+    "from qgs.plotting.util import std_plot"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Systems definition"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "General parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Time parameters\n",
+    "dt = 0.1\n",
+    "# Saving the model state n steps\n",
+    "write_steps = 5\n",
+    "\n",
+    "number_of_trajectories = 1\n",
+    "number_of_perturbed_trajectories = 10"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setting some model parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Model parameters instantiation with some non-default specs\n",
+    "model_parameters = QgParams({'phi0_npi': np.deg2rad(50.)/np.pi, 'hd':0.045})\n",
+    "# Mode truncation at the wavenumber 2 in both x and y spatial coordinate\n",
+    "model_parameters.set_atmospheric_channel_fourier_modes(2, 2)\n",
+    "\n",
+    "# Changing (increasing) the orography depth and the meridional temperature gradient\n",
+    "model_parameters.ground_params.set_orography(0.2, 1)\n",
+    "model_parameters.atemperature_params.set_thetas(0.1, 0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Printing the model's parameters\n",
+    "model_parameters.print_params()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Creating the tendencies function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "f, Df = create_tendencies(model_parameters)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Time integration"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Defining an integrator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "integrator = RungeKuttaIntegrator()\n",
+    "integrator.set_func(f)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Start on a random initial condition and integrate over a transient time to obtain an initial condition on the attractors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "ic = np.random.rand(model_parameters.ndim)*0.1\n",
+    "integrator.integrate(0., 200000., dt, ic=ic, write_steps=0)\n",
+    "time, ic = integrator.get_trajectories()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now integrate to obtain a trajectory on the attractor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "integrator.integrate(0., 100000., dt, ic=ic, write_steps=write_steps)\n",
+    "reference_time, reference_traj = integrator.get_trajectories()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "varx = 0\n",
+    "vary = 1\n",
+    "varz = 2\n",
+    "\n",
+    "fig = plt.figure(figsize=(10, 8))\n",
+    "axi = fig.add_subplot(111, projection='3d')\n",
+    "\n",
+    "axi.scatter(reference_traj[varx], reference_traj[vary], reference_traj[varz], s=0.2);\n",
+    "\n",
+    "axi.set_xlabel('$'+model_parameters.latex_var_string[varx]+'$')\n",
+    "axi.set_ylabel('$'+model_parameters.latex_var_string[vary]+'$')\n",
+    "axi.set_zlabel('$'+model_parameters.latex_var_string[varz]+'$');"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "varx = 2\n",
+    "vary = 1\n",
+    "plt.figure(figsize=(10, 8))\n",
+    "\n",
+    "plt.plot(reference_traj[varx], reference_traj[vary], marker='o', ms=0.07, ls='')\n",
+    "\n",
+    "plt.xlabel('$'+model_parameters.latex_var_string[varx]+'$')\n",
+    "plt.ylabel('$'+model_parameters.latex_var_string[vary]+'$');"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "var = 1\n",
+    "plt.figure(figsize=(10, 8))\n",
+    "\n",
+    "plt.plot(model_parameters.dimensional_time*reference_time, reference_traj[var])\n",
+    "\n",
+    "plt.xlabel('time (days)')\n",
+    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Showing the vertical wind velocity"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Creating the diagnostics:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qgs.diagnostics.streamfunctions import MiddleAtmosphericStreamfunctionDiagnostic, UpperLayerAtmosphericStreamfunctionDiagnostic, LowerLayerAtmosphericStreamfunctionDiagnostic\n",
+    "from qgs.diagnostics.variables import VariablesDiagnostic\n",
+    "from qgs.diagnostics.multi import MultiDiagnostic\n",
+    "from qgs.diagnostics.wind import MiddleLayerVerticalVelocity, MiddleAtmosphericUWindDiagnostic, MiddleAtmosphericWindIntensityDiagnostic"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "* For the temperature, and the 250, 500 & 750 hPa streamfunctins :"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "psi = MiddleAtmosphericStreamfunctionDiagnostic(model_parameters)\n",
+    "psi1 = UpperLayerAtmosphericStreamfunctionDiagnostic(model_parameters)\n",
+    "psi3 = LowerLayerAtmosphericStreamfunctionDiagnostic(model_parameters)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "* For the vertical wind velocity :"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vertwind = MiddleLayerVerticalVelocity(model_parameters)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "* For the 500 hPa wind :"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "wind = MiddleAtmosphericWindIntensityDiagnostic(model_parameters)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "* For the nondimensional variables $\\psi_{{\\rm a}, 2}$ and $\\psi_{{\\rm a}, 3}$:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qgs.diagnostics.variables import VariablesDiagnostic\n",
+    "from qgs.diagnostics.multi import MultiDiagnostic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "variable_nondim = VariablesDiagnostic([2, 1], model_parameters, False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# setting also the background\n",
+    "background = VariablesDiagnostic([2, 1], model_parameters, False)\n",
+    "background.set_data(reference_time, reference_traj)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Creating a multi diagnostic with both:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m = MultiDiagnostic(2,3)\n",
+    "m.add_diagnostic(vertwind)\n",
+    "m.add_diagnostic(psi)\n",
+    "m.add_diagnostic(psi1)\n",
+    "m.add_diagnostic(variable_nondim,\n",
+    "                 diagnostic_kwargs={'show_time': False, 'background': background},\n",
+    "                 plot_kwargs={'ms': 0.2})\n",
+    "m.add_diagnostic(psi3)\n",
+    "m.add_diagnostic(wind)\n",
+    "\n",
+    "m.set_data(reference_time[:2000], reference_traj[:,:2000])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "and show an interactive animation:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "rc('font',**{'family':'serif','sans-serif':['Times'],'size':12})\n",
+    "m.animate(figsize=(22,8))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "or a movie:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "rc('font',**{'family': 'serif','sans-serif': ['Times'],'size': 12})\n",
+    "m.movie(figsize=(22,8), anim_kwargs={'interval': 100, 'frames': 1000})"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

From 03eaeb3e7a612382c03d0a6b67a6f6f47338344c Mon Sep 17 00:00:00 2001
From: Jonathan Demaeyer <jodemaey@meteo.be>
Date: Sun, 11 Dec 2022 10:53:44 +0100
Subject: [PATCH 7/7] Added some modifications of the docs

---
 qgs/diagnostics/wind.py           | 3 ++-
 qgs/tensors/atmo_thermo_tensor.py | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/qgs/diagnostics/wind.py b/qgs/diagnostics/wind.py
index 675c048..c68302b 100644
--- a/qgs/diagnostics/wind.py
+++ b/qgs/diagnostics/wind.py
@@ -640,7 +640,8 @@ def _get_diagnostic(self, dimensional):
 
 
 class MiddleLayerVerticalVelocity(AtmosphericWindDiagnostic):
-    """Diagnostic giving the middle atmospheric layer vertical wind intensity fields.
+    """Diagnostic giving the middle atmospheric layer vertical wind intensity fields :math:`\\omega`.
+    See also the :ref:`files/model/atmosphere:Atmospheric component` and :ref:`files/model/oro_model:Mid-layer equations and the thermal wind relation` sections.
 
     Parameters
     ----------
diff --git a/qgs/tensors/atmo_thermo_tensor.py b/qgs/tensors/atmo_thermo_tensor.py
index f7d0357..1a740a7 100644
--- a/qgs/tensors/atmo_thermo_tensor.py
+++ b/qgs/tensors/atmo_thermo_tensor.py
@@ -2,7 +2,8 @@
     qgs atmospheric thermodynamic tensor module
     ===========================================
 
-    This module computes and holds the tensors representing the atmospheric tendencies of the model's equations.
+    This module computes and holds the tensors representing the atmospheric thermodynamic tendencies of the model's equations
+    needed to compute the vertical wind velocity :math:`\\omega`.
 
     TODO: Add a list of the different tensor available