Developing sco class

pydomcfg/domzgr/sco.py

@@ -6,12 +6,12 @@
 
 from typing import Optional  # , Tuple
 
-# import numpy as np
+import numpy as np
 from xarray import DataArray, Dataset
 
-from .zgr import Zgr
+from pydomcfg.utils import _smooth_MB06  # , _calc_rmax
 
-# from pydomcfg.utils import is_nemo_none
+from .zgr import Zgr
 
 
 class Sco(Zgr):
@@ -45,8 +45,8 @@ class Sco(Zgr):
     # --------------------------------------------------------------------------
     def __call__(
         self,
-        bot_min: float,
-        bot_max: float,
+        min_dep: float,
+        max_dep: float,
         hc: float = 0.0,
         rmax: Optional[float] = None,
         stretch: Optional[str] = None,
@@ -62,9 +62,9 @@ def __call__(
 
         Parameters
         ----------
-        bot_min: float
+        min_dep: float
             Minimum depth of bottom topography surface (>0) (m)
-        bot_max: float
+        max_dep: float
             Maximum depth of bottom topography surface (>0) (m)
         hc: float
             critical depth for transition from uniform sigma
@@ -103,10 +103,10 @@ def __call__(
             Describing the 3D geometry of the model
         """
 
-        self._bot_min = bot_min
-        self._bot_max = bot_max
+        self._min_dep = min_dep
+        self._max_dep = max_dep
         self._hc = hc
-        self.rmax = rmax
+        self._rmax = rmax
         self._stretch = stretch
         self._ln_e3_dep = ln_e3_dep
 
@@ -119,10 +119,20 @@ def __call__(
             self._efold = efold or 0.0
             self._pbot2 = pbot2 or 0.0
 
-        # ds = self._init_ds()
+        bathy = self._bathy["Bathymetry"]
+
+        # Land-Sea mask of the domain:
+        # 0 = land
+        # 1 = ocean
+        self._ocean = bathy.where(bathy == 0, 1)
+
+        # set maximum and minumum depths of model bathymetry
+        bathy = bathy.where(bathy < self._max_dep, self._max_dep)
+        bathy = bathy.where(bathy > self._min_dep, self._min_dep)
+        bathy *= self._ocean
 
         # compute envelope bathymetry DataArray
-        self._envlp = self._compute_env(self._bathy["Bathymetry"])
+        self._envlp = self._compute_env(bathy)
 
         # compute sigma-coordinates for z3 computation
         self._sigmas = self._compute_sigma(self._z)
@@ -134,7 +144,9 @@ def __call__(
         # dse = self._compute_e3(dsz) if self._ln_e3_dep else self._analyt_e3(dsz)
 
         # addind this only to not make darglint complying
-        return self._merge_z3_and_e3(self._envlp, self._envlp, self._envlp, self._envlp)
+        ds = self._bathy.copy()
+        ds["hbatt"] = self._envlp
+        return ds
 
     # --------------------------------------------------------------------------
     def _check_stretch_par(self, psurf, pbott, alpha, efold, pbot2):
@@ -163,7 +175,7 @@ def _check_stretch_par(self, psurf, pbott, alpha, efold, pbot2):
                 )
 
     # --------------------------------------------------------------------------
-    def _compute_env(self, da: DataArray) -> DataArray:
+    def _compute_env(self, depth: DataArray) -> DataArray:
         """
         Compute the envelope bathymetry surface by applying the
         Martinho & Batteen (2006) smoothing algorithm to the
@@ -179,12 +191,48 @@ def _compute_env(self, da: DataArray) -> DataArray:
 
         Parameters
         ----------
-        da: DataArray
+        depth: DataArray
             xarray DataArray of the 2D bottom topography
         Returns
         -------
         DataArray
             xarray DataArray of the 2D envelope bathymetry
         """
 
-        return da
+        da_zenv = depth.copy()
+
+        if self._rmax:
+
+            # getting the actual numpy array
+            # TO BE OPTIMISED
+            zenv = da_zenv.data
+            nj = zenv.shape[0]
+            ni = zenv.shape[1]
+
+            # set first land point adjacent to a wet cell to
+            # min_dep as this needs to be included in smoothing
+            for j in range(nj - 1):
+                for i in range(ni - 1):
+                    if not self._ocean[j, i]:
+                        ip1 = np.minimum(i + 1, ni)
+                        jp1 = np.minimum(j + 1, nj)
+                        im1 = np.maximum(i - 1, 0)
+                        jm1 = np.maximum(j - 1, 0)
+                        if (
+                            depth[jp1, im1]
+                            + depth[jp1, i]
+                            + depth[jp1, ip1]
+                            + depth[j, im1]
+                            + depth[j, ip1]
+                            + depth[jm1, im1]
+                            + depth[jm1, i]
+                            + depth[jm1, ip1]
+                        ) > 0.0:
+                            zenv[j, i] = self._min_dep
+
+            da_zenv.data = zenv
+            # print(np.nanmax(_calc_rmax(da_zenv)*self._ocean))
+            da_zenv = _smooth_MB06(da_zenv, self._rmax)
+            da_zenv = da_zenv.where(da_zenv > self._min_dep, self._min_dep)
+
+        return da_zenv

pydomcfg/utils.py

@@ -2,6 +2,7 @@
 Utilities
 """
 
+# from itertools import product
 from typing import Hashable, Iterable, Iterator, Optional
 
 import numpy as np
@@ -23,34 +24,30 @@ def _are_nemo_none(var: Iterable) -> Iterator[bool]:
 
 
 # -----------------------------------------------------------------------------
-def _calc_rmax(depth: DataArray) -> float:
+def _calc_rmax(depth: DataArray) -> DataArray:
     """
-    Calculate rmax: measure of steepness
-    This function returns the maximum slope paramater
+     Calculate rmax: measure of steepness
+     This function returns the maximum slope paramater
 
-    rmax = abs(Hb - Ha) / (Ha + Hb)
+     rmax = abs(Hb - Ha) / (Ha + Hb)
 
-    where Ha and Hb are the depths of adjacent grid cells (Mellor et al 1998).
+     where Ha and Hb are the depths of adjacent grid cells (Mellor et al 1998).
 
-    Reference:
-    *) Mellor, Oey & Ezer, J Atm. Oce. Tech. 15(5):1122-1131, 1998.
+     Reference:
+     *) Mellor, Oey & Ezer, J Atm. Oce. Tech. 15(5):1122-1131, 1998.
 
-    Parameters
-    ----------
-    depth: DataArray
-        Bottom depth (units: m).
+     Parameters
+     ----------
+     depth: DataArray
+         Bottom depth (units: m).
 
-    Returns
-    -------
-    float
-        Maximum slope parameter (units: None)
+     Returns
+     -------
+    DataArray
+         2D maximum slope parameter (units: None)
     """
 
-    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
-    # Do we actually need this? mypy complains since
-    # DataArray.reset_indexe() returns Optional["DataArray"]
-    #
-    # depth = depth.reset_index(list(depth.dims))
+    depth = DataArray(depth.reset_index(list(depth.dims)))
 
     both_rmax = []
     for dim in depth.dims:
@@ -59,6 +56,8 @@ def _calc_rmax(depth: DataArray) -> float:
         depth_diff = depth.diff(dim)
         depth_rolling_sum = depth.rolling({dim: 2}).sum().dropna(dim)
         rmax = depth_diff / depth_rolling_sum
+        # dealing with nans at land points
+        rmax = rmax.where(np.isfinite(rmax), 0)
 
         # (rmax_a + rmax_b) / 2
         rmax = rmax.rolling({dim: 2}).mean().dropna(dim)
@@ -71,6 +70,109 @@ def _calc_rmax(depth: DataArray) -> float:
     return np.maximum(*both_rmax)
 
 
+# -----------------------------------------------------------------------------
+def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray:
+    """
+    This is NEMO implementation of the direct iterative method
+    of Martinho and Batteen (2006).
+
+    The algorithm ensures that
+
+                H_ij - H_n
+                ---------- < rmax
+                H_ij + H_n
+
+    where H_ij is the depth at some point (i,j) and H_n is the
+    neighbouring depth in the east, west, south or north direction.
+
+    Reference:
+    *) Martinho & Batteen, Oce. Mod. 13(2):166-175, 2006.
+
+    Parameters
+    ----------
+    depth: DataArray
+        Bottom depth (units: m).
+    rmax: float
+        Maximum slope parameter allowed
+
+    Returns
+    -------
+    DataArray
+        Smooth version of the bottom topography with
+        a maximum slope parameter < rmax (units: m).
+
+    """
+
+    # set scaling factor used for smoothing
+    zrfact = (1.0 - rmax) / (1.0 + rmax)
+
+    # getting the actual numpy array
+    # TO BE OPTIMISED
+    da_zenv = depth.copy()
+    zenv = da_zenv.data
+    nj = zenv.shape[0]
+    ni = zenv.shape[1]
+
+    # initialise temporary evelope depth arrays
+    ztmpi1 = zenv.copy()
+    ztmpi2 = zenv.copy()
+    ztmpj1 = zenv.copy()
+    ztmpj2 = zenv.copy()
+
+    # Computing the initial maximum slope parameter
+    zrmax = np.nanmax(_calc_rmax(depth))
+    zri = np.ones(zenv.shape) * zrmax
+    zrj = np.ones(zenv.shape) * zrmax
+
+    tol = 1.0e-8
+    itr = 0
+    max_itr = 10000
+
+    while itr <= max_itr and (zrmax - rmax) > tol:
+
+        itr += 1
+        zrmax = 0.0
+        # we set zrmax from previous r-values (zri and zrj) first
+        # if set after current r-value calculation (as previously)
+        # we could exit DO WHILE prematurely before checking r-value
+        # of current zenv
+        max_zri = np.nanmax(np.absolute(zri))
+        max_zrj = np.nanmax(np.absolute(zrj))
+        zrmax = np.nanmax([zrmax, max_zrj, max_zri])
+
+        print("Iter:", itr, "rmax: ", zrmax)
+
+        zri *= 0.0
+        zrj *= 0.0
+        for j in range(nj - 1):
+            for i in range(ni - 1):
+                ip1 = np.minimum(i + 1, ni)
+                jp1 = np.minimum(j + 1, nj)
+                if zenv[j, i] > 0.0 and zenv[j, ip1] > 0.0:
+                    zri[j, i] = (zenv[j, ip1] - zenv[j, i]) / (
+                        zenv[j, ip1] + zenv[j, i]
+                    )
+                if zenv[j, i] > 0.0 and zenv[jp1, i] > 0.0:
+                    zrj[j, i] = (zenv[jp1, i] - zenv[j, i]) / (
+                        zenv[jp1, i] + zenv[j, i]
+                    )
+                if zri[j, i] > rmax:
+                    ztmpi1[j, i] = zenv[j, ip1] * zrfact
+                if zri[j, i] < -rmax:
+                    ztmpi2[j, ip1] = zenv[j, i] * zrfact
+                if zrj[j, i] > rmax:
+                    ztmpj1[j, i] = zenv[jp1, i] * zrfact
+                if zrj[j, i] < -rmax:
+                    ztmpj2[jp1, i] = zenv[j, i] * zrfact
+
+        ztmpi = np.maximum(ztmpi1, ztmpi2)
+        ztmpj = np.maximum(ztmpj1, ztmpj2)
+        zenv = np.maximum(zenv, np.maximum(ztmpi, ztmpj))
+
+    da_zenv.data = zenv
+    return da_zenv
+
+
 # -----------------------------------------------------------------------------
 def generate_cartesian_grid(
     ppe1_m,