add envelope computation

Author: oceandie

pydomcfg/domzgr/sco.py

@@ -7,6 +7,7 @@
 from typing import Optional  # , Tuple
 
 import numpy as np
+import xarray as xr
 from xarray import DataArray, Dataset
 
 from pydomcfg.utils import _smooth_MB06  # , _calc_rmax
@@ -121,15 +122,14 @@ def __call__(
 
         bathy = self._bathy["Bathymetry"]
 
-        # Land-Sea mask of the domain:
+        # compute land-sea mask of the domain:
         # 0 = land
         # 1 = ocean
-        self._ocean = bathy.where(bathy == 0, 1)
+        self._lsm = xr.where(bathy > 0, 1, 0)
 
         # 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
+        bathy = np.minimum(bathy, self._max_dep)
+        bathy = np.maximum(bathy, self._min_dep) * self._lsm
 
         # compute envelope bathymetry DataArray
         self._envlp = self._compute_env(bathy)
@@ -193,46 +193,79 @@ def _compute_env(self, depth: DataArray) -> DataArray:
         ----------
         depth: DataArray
             xarray DataArray of the 2D bottom topography
+            it MUST have only two dimensions named "x" and "y"
         Returns
         -------
         DataArray
             xarray DataArray of the 2D envelope bathymetry
         """
 
-        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]
+            lsm = self._lsm
 
             # 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
+
+            # ------------------------------------------------------------
+            # This is the original NEMO Fortran90 code: translated
+            # in python it is very inefficient
+            # ------------------------------------------------------------
+            # zenv = depth.copy()
+            # env  = zenv.data
+            #
+            # nj = env.shape[0]
+            # ni = env.shape[1]
+            #
+            # for j in range(nj - 1):
+            #    for i in range(ni - 1):
+            #        if not lsm[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:
+            #               env[j, i] = min_dep
+            #
+            # zenv.data = env
+            # ------------------------------------------------------------
+
+            # ------------------------------------------------------------
+            # This is my translation into xarray. I think it does what
+            # it should, a part on the boundaries: I tested with AMM7,
+            # and zenv computed with NEMO-like code (above) and this
+            # one are perfectly identical apart for two single different
+            # points just on the border ... I don't think it will make
+            # a huge difference but if there is a better way to manage
+            # the borders with xarray and obtain exactly the same results
+            # of the original NEMO-like code, happy to use it.
+            # ------------------------------------------------------------
+            cst_lsm = lsm * 0.0
+            ngb_pnt = [-1, 0, 1]
+            for j in ngb_pnt:
+                for i in ngb_pnt:
+                    if j != 0 or i != 0:
+                        lsm_sft = lsm.shift({lsm.dims[1]: i, lsm.dims[0]: j})
+                        cst_lsm += lsm_sft
+
+            cst_lsm = cst_lsm.where(lsm == 0, 0)
+            cst_lsm = cst_lsm.where(cst_lsm == 0, 1)
+            zenv = depth.where(cst_lsm == 0, self._min_dep)
+            for dim in lsm.dims:
+                for indx in [0, -1]:
+                    zenv[{dim: indx}] = depth[{dim: indx}]
+            # ------------------------------------------------------------
+
+            zenv = _smooth_MB06(zenv, self._rmax)
+            zenv = zenv.where(zenv > self._min_dep, self._min_dep)
+
+        return zenv

pydomcfg/utils.py

@@ -23,13 +23,12 @@ def _are_nemo_none(var: Iterable) -> Iterator[bool]:
         yield _is_nemo_none(v)
 
 
-# -----------------------------------------------------------------------------
 def _calc_rmax(depth: DataArray) -> DataArray:
     """
     Calculate rmax: measure of steepness
-    This function returns the maximum slope paramater
+    This function returns the slope paramater field
 
-    rmax = abs(Hb - Ha) / (Ha + Hb)
+    r = abs(Hb - Ha) / (Ha + Hb)
 
     where Ha and Hb are the depths of adjacent grid cells (Mellor et al 1998).
 
@@ -44,33 +43,43 @@ def _calc_rmax(depth: DataArray) -> DataArray:
     Returns
     -------
     DataArray
-        2D maximum slope parameter (units: None)
-    """
+        2D slope parameter (units: None)
 
-    depth = DataArray(depth.reset_index(list(depth.dims)))
+    Notes
+    -----
+    This function uses a "conservative approach" and rmax is overestimated.
+    rmax at T points is the maximum rmax estimated at any adjacent U/V point.
+    """
+    # Mask land
+    depth = depth.where(depth > 0)
 
+    # Loop over x and y
     both_rmax = []
     for dim in depth.dims:
 
-        # (Hb - Ha) / (Ha + Hb)
-        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)
+        # Compute rmax
+        rolled = depth.rolling({dim: 2}).construct("window_dim")
+        diff = rolled.diff("window_dim").squeeze("window_dim")
+        rmax = np.abs(diff) / rolled.sum("window_dim")
+
+        # Construct dimension with velocity points adjacent to any T point
+        # We need to shift as we rolled twice
+        rmax = rmax.rolling({dim: 2}).construct("vel_points")
+        rmax = rmax.shift({dim: -1})
 
-        # (rmax_a + rmax_b) / 2
-        rmax = rmax.rolling({dim: 2}).mean().dropna(dim)
+        both_rmax.append(rmax)
 
-        # Fill first row and column
-        rmax = rmax.pad({dim: (1, 1)}, constant_values=0)
+    # Find maximum rmax at adjacent U/V points
+    rmax = xr.concat(both_rmax, "vel_points")
+    rmax = rmax.max("vel_points", skipna=True)
 
-        both_rmax.append(np.abs(rmax))
+    # Mask halo points
+    for dim in rmax.dims:
+        rmax[{dim: [0, -1]}] = 0
 
-    return np.maximum(*both_rmax)
+    return rmax.fillna(0)
 
 
-# -----------------------------------------------------------------------------
 def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray:
     """
     This is NEMO implementation of the direct iterative method
@@ -120,9 +129,9 @@ def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray:
     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
+    zrmax = 1.0  # np.nanmax(_calc_rmax(depth))
+    zri = np.ones(zenv.shape)  # * zrmax
+    zrj = np.ones(zenv.shape)  # * zrmax
 
     tol = 1.0e-8
     itr = 0
@@ -144,6 +153,7 @@ def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray:
 
         zri *= 0.0
         zrj *= 0.0
+
         for j in range(nj - 1):
             for i in range(ni - 1):
                 ip1 = np.minimum(i + 1, ni)
@@ -173,7 +183,6 @@ def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray:
     return da_zenv
 
 
-# -----------------------------------------------------------------------------
 def generate_cartesian_grid(
     ppe1_m,
     ppe2_m,