add zgr and zco classes

pydomcfg/domzgr/zco.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python
+
+"""
+Class to generate NEMO v4.0 standard geopotential z-coordinates
+"""
+
+
+import numpy as np
+from xarray import Dataset
+
+from .zgr import Zgr
+
+
+class Zco(Zgr):
+
+    # In future we may want to get rid of this
+    # and use a more pythonic way. We could do
+    # it when  we read the namelist.
+    pp_to_be_computed = 999999.0
+
+    """
+      Class to generate geopotential z-coordinates dataset objects.
+
+      Method
+      ------
+      *) Model levels' depths depT/W are defined from analytical function.
+      *) Model vertical scale factors e3 (i.e., grid cell thickness) can
+         be computed as
+
+             1) analytical derivative of depth function
+                (ln_e3_dep=False); for backward compatibility with v3.6.
+             2) discrete derivative (central-difference) of levels' depth
+                (ln_e3_dep=True). The only possibility from v4.0.
+
+      References:
+       *) NEMO v4.0 domzgr/zgr_z subroutine
+       *) Marti, Madec & Delecluse, 1992, JGR, 97, No8, 12,763-12,766.
+      """
+
+    def __init__(self):
+        """
+        Initialise class.
+        """
+
+    # --------------------------------------------------------------------------
+    def generate(
+        self,
+        ppdzmin,
+        pphmax,
+        ppkth: float = 0.0,
+        ppacr: int = 0,
+        ppsur: float = pp_to_be_computed,
+        ppa0: float = pp_to_be_computed,
+        ppa1: float = pp_to_be_computed,
+        ldbletanh: bool = False,
+        ppa2: float = 0.0,
+        ppkth2: float = 0.0,
+        ppacr2: float = 0.0,
+    ) -> Dataset:
+        """
+        Generate NEMO geopotential z-coordinates model levels.
+
+        Parameters
+        ----------
+        ppdzmin: float
+            Minimum thickness for the top layer (units: m)
+        pphmax: float
+            Depth of the last W-level (total depth of the ocean basin) (units: m)
+        ppacr: int
+            Stretching factor, nondimensional > 0. The larger ppacr,
+            the smaller the stretching. Values from 3 to 10 are usual.
+            (default = 0., no stretching, uniform grid)
+        ppkth: float
+            Model level at which approximately max stretching occurs.
+            Nondimensional > 0, usually of order 1/2 or 2/3 of jpk.
+            (default = 0., i.e. no stretching, uniform grid)
+        ppsur: float
+            Coeff. controlling distibution of vertical levels
+            (default = pp_to_be_computed, i.e. computed from
+            ppdzmin, pphmax, ppkth and ppacr)
+        ppa0: float
+            Coeff. controlling distibution of vertical levels
+            (default = pp_to_be_computed, i.e. computed from
+            ppdzmin, pphmax, ppkth and ppacr)
+        ppa1: float
+            Coeff. controlling distibution of vertical levels
+            (default = pp_to_be_computed, i.e. computed from
+            ppdzmin, pphmax, ppkth and ppacr)
+        ldbletanh: bool
+            Logical flag to use or not double tanh stretching function (default = False)
+        ppa2: float
+            Double tanh stretching function parameter
+        ppkth2: float
+            Double tanh stretching function parameter
+        ppacr2: float
+            Double tanh stretching function parameter
+
+        Returns
+        -------
+        Dataset
+           Describing the 3D geometry of the model
+
+        """
+        self._ppdzmin = ppdzmin
+        self._pphmax = pphmax
+        self._ppkth = ppkth
+        self._ppacr = ppacr
+        self._ppsur = ppsur
+        self._ppa0 = ppa0
+        self._ppa1 = ppa1
+        self._ldbletanh = ldbletanh
+        self._ppa2 = ppa2
+        self._ppkth2 = ppkth2
+        self._ppacr2 = ppacr2
+
+        ds = self.init_ds()
+
+        # computing coeff. if needed
+        if self._ppkth * self._ppacr > 0.0:
+            self.compute_pp()
+
+        # compute z3 depths of vertical levels
+        for k in range(self._jpk):
+
+            if self._ppkth * self._ppacr == 0.0:
+                # uniform zco grid
+                suT = -self.sigma(k, "T")
+                suW = -self.sigma(k, "W")
+                s1T = s1W = s2T = s2W = 0.0
+                a1 = a3 = a4 = 0.0
+                a2 = self._pphmax
+            else:
+                # stretched zco grid
+                suT = -self.sigma(k, "T") * (self._jpk - 1) + 1
+                suW = -self.sigma(k, "W") * (self._jpk - 1) + 1
+                s1T = self.stretch_zco1(suT)
+                s1W = self.stretch_zco1(suW)
+                a1 = self._ppsur
+                a2 = self._ppa0
+                a3 = self._ppa1 * self._ppacr
+                if self._ldbletanh:
+                    s2T = self.stretch_zco2(suT)
+                    s2W = self.stretch_zco2(suW)
+                    a4 = self._ppa2 * self._ppacr2
+                else:
+                    s2T = s2W = a4 = 0.0
+
+            ds["z3T"][k, :, :] = self.compute_z3(suT, s1T, a1, a2, a3, s2T, a4)
+            ds["z3W"][k, :, :] = self.compute_z3(suW, s1W, a1, a2, a3, s2W, a4)
+
+        # force first w-level to be exactly at zero
+        ds["z3W"][0, :, :] = 0.0
+
+        # compute e3 scale factors
+        dsz = self.compute_e3(ds)
+
+        return dsz
+
+    # --------------------------------------------------------------------------
+    def compute_pp(self):
+        """
+        Compute the coefficients for zco grid if requested.
+
+        Raises
+        ------
+        ValueError
+            If ldbletanh = True and parametrs are equal to 0.
+
+        """
+        if (
+            self._ppsur == self.pp_to_be_computed
+            or self._ppa0 == self.pp_to_be_computed
+            or self._ppa1 == self.pp_to_be_computed
+        ):
+
+            aa = self._ppdzmin - self._pphmax / float(self._jpk - 1)
+            bb = np.tanh((1 - self._ppkth) / self._ppacr)
+            cc = self._ppacr / float(self._jpk - 1)
+            dd = np.log(np.cosh((self._jpk - self._ppkth) / self._ppacr)) - np.log(
+                np.cosh((1 - self._ppkth) / self._ppacr)
+            )
+
+            self._ppa1 = aa / (bb - cc * dd)
+            self._ppa0 = self._ppdzmin - self._ppa1 * bb
+            self._ppsur = -self._ppa0 - self._ppa1 * self._ppacr * np.log(
+                np.cosh((1 - self._ppkth) / self._ppacr)
+            )
+
+            if self._ldbletanh and self._ppa2 * self._ppkth2 * self._ppacr2 == 0.0:
+                raise ValueError(
+                    "ppa2, ppkth2 and ppacr2 must > 0. " "when ldbletanh = True"
+                )
+
+    # --------------------------------------------------------------------------
+    def stretch_zco1(self, k: float) -> float:
+        """
+        Provide the standard analytical stretching function for NEMO z-coordinates.
+
+        Parameters
+        ----------
+        k : float
+
+        Returns
+        -------
+        float
+
+        """
+
+        ss = np.log(np.cosh((k - self._ppkth) / self._ppacr))
+        return ss
+
+    # --------------------------------------------------------------------------
+    def stretch_zco2(self, k: float):
+        """
+        Provide the double tanh analytical stretching function for NEMO z-coordinates.
+
+        Parameters
+        ----------
+        k: float
+
+        Returns
+        -------
+        ss: float
+
+        """
+
+        ss = np.log(np.cosh((k - self._ppkth2) / self._ppacr2))
+        return ss

pydomcfg/domzgr/zgr.py

@@ -0,0 +1,154 @@
+"""
+Base class to generate NEMO v4.0 vertical grids.
+"""
+
+from itertools import product
+
+import numpy as np
+from xarray import DataArray, Dataset
+
+
+class Zgr:
+    """
+    Base class to generate NEMO vertical grids.
+    """
+
+    def __init__(self, ds_bathy: Dataset, jpk: int):
+        """
+        Initialize class.
+
+        Parameters
+        ----------
+        ds_bathy: Dataset
+             xarray dataset including grid coordinates and bathymetry
+        jpk: int
+             number of computational levels
+        """
+
+        self._bathy = ds_bathy
+        self._jpk = jpk
+
+    # -------------------------------------------------------------------------
+    def init_ds(self) -> Dataset:
+        ds = self._bathy.copy()
+        jpi = ds.dims["x"]
+        jpj = ds.dims["y"]
+        jpk = self._jpk
+
+        var = ["z3", "e3"]
+        grd = ["T", "W"]
+        crd = [("z", range(jpk)), ("y", range(jpj)), ("x", range(jpi))]
+
+        # Initialise a dataset with z3 and e3 dataarray filled with nan
+        da = DataArray(np.zeros(shape=(jpk, jpj, jpi)) * np.nan, coords=crd)
+        for v, g in product(var, grd):
+            ds[v + g] = da.copy()
+            ds = ds.set_coords(v + g)
+        return ds
+
+    # -------------------------------------------------------------------------
+    def sigma(self, k: int, grd: str) -> float:
+        """
+        Provide the analytical function for sigma-coordinate,
+        a uniform non-dimensional vertical coordinate describing
+        the non-dimensional position of model levels.
+
+        Consider that 0. <= sigma <= -1, with
+
+                 sigma =  0 at the shallower boundary
+                 sigma = -1 at the deeper boundary
+
+        Parameters
+        ----------
+        k: int
+          Model level index. Note that
+             *) T-points are at integer values (between 1 and jpk)
+             *) W-points are at integer values - 1/2 (between 0.5 and jpk-0.5)
+        grd: str
+          If we are dealing with "T" or "W" model levels.
+
+        Returns
+        -------
+        float
+        """
+
+        kindx = float(k + 1)  # to deal with python convention
+        if grd == "W":
+            kindx -= 0.5
+
+        ps = -(kindx - 0.5) / float(self._jpk - 1)
+        return ps
+
+    # -------------------------------------------------------------------------
+    @staticmethod
+    def compute_z3(
+        su: float,
+        ss1: float,
+        a1: float,
+        a2: float,
+        a3: float,
+        ss2: float = 0.0,
+        a4: float = 0.0,
+    ) -> float:
+        """
+        Generalised function providing the analytical
+        transformation from computational space to
+        physical space. It takes advantage of the fact that
+        z-coordinates can be considered a special case of
+        s-coordinate.
+
+        Note that z is downward positive.
+
+        Parameters
+        ----------
+        su: float
+            uniform non-dimensional vertical coordinate s,
+            aka sigma-coordinates. 0 <= s <= 1
+        ss1: float
+            stretched non-dimensional vertical coordinate s,
+            0 <= s <= 1
+        a1: float
+            parameter of the transformation
+        a2: float
+            parameter of the transformation
+        a3: float
+            parameter of the transformation
+        ss2: float
+             second stretched non-dimensional vertical coordinate s,
+             0 <= s <= 1 (only used for zco with ldbletanh = True)
+        a4: float
+            parameter of the transformation (only used for zco with ldbletanh = True)
+
+        Returns
+        -------
+        float
+
+        """
+
+        z = a1 + a2 * su + a3 * ss1 + a4 * ss2
+        return z
+
+    # -------------------------------------------------------------------------
+    def compute_e3(self, ds: Dataset) -> Dataset:
+        """
+        Grid cell thickness computed as discrete derivative
+        (central-difference) of levels' depth
+
+        Parameters
+        ----------
+        ds: Dataset
+
+        Returns
+        -------
+        Dataset
+        """
+        for k in range(self._jpk - 1):
+            ds["e3T"][k, :, :] = ds["z3W"][k + 1, :, :] - ds["z3W"][k, :, :]
+            ds["e3W"][k + 1, :, :] = ds["z3T"][k + 1, :, :] - ds["z3T"][k, :, :]
+        # Bottom:
+        k = -1
+        ds["e3T"][k, :, :] = 2.0 * (ds["z3T"][k, :, :] - ds["z3W"][k, :, :])
+        # Surface:
+        k = 0
+        ds["e3W"][k, :, :] = 2.0 * (ds["z3T"][k, :, :] - ds["z3W"][k, :, :])
+        return ds