Test calc rmax closes #48

pydomcfg/tests/lib_for_rmax.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python
+
+import netCDF4 as nc4
+import numpy as np
+import xarray as xr
+from xarray import DataArray, Dataset
+
+
+# =======================================================================================
+def calc_rmax_np(depth):
+    """
+    Calculate rmax: measure of steepness
+
+    This function returns the slope steepness criteria rmax, which is simply
+    (H[0] - H[1]) / (H[0] + H[1])
+    Parameters
+    ----------
+    depth: float
+            Bottom depth (units: m).
+    Returns
+    -------
+    rmax: float
+            Slope steepness value (units: None)
+    """
+
+    rmax_x, rmax_y = np.zeros_like(depth), np.zeros_like(depth)
+
+    rmax_x[1:-1, 1:-1] = np.maximum (
+        np.abs( np.diff(depth[1:-1, :-1], axis=1) / (depth[1:-1,  :-2] + depth[1:-1, 1:-1]) ),
+        np.abs( np.diff(depth[1:-1,1:  ], axis=1) / (depth[1:-1, 1:-1] + depth[1:-1, 2:  ]) )
+    )
+    rmax_y[1:-1, 1:-1] = np.maximum (
+        np.abs( np.diff(depth[ :-1, 1:-1], axis=0) / (depth[ :-2, 1:-1] + depth[1:-1, 1:-1]) ) ,
+        np.abs( np.diff(depth[1:  , 1:-1], axis=0) / (depth[1:-1, 1:-1] + depth[2:,   1:-1]) )
+    )
+
+    return np.maximum(np.abs(rmax_x), np.abs(rmax_y))
+
+
+# =======================================================================================
+def calc_rmax_xr1(depth):
+    """
+    Calculate rmax: measure of steepness
+    This function returns the maximum slope paramater
+
+    rmax = abs(Hb - Ha) / (Ha + Hb)
+
+    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.
+
+    Parameters
+    ----------
+    depth: DataArray
+        Bottom depth (units: m).
+
+    Returns
+    -------
+    DataArray
+        2D maximum slope parameter (units: None)
+
+    """
+    both_rmax = []
+
+    for dim in depth.dims:
+
+        # |Ha - Hb| / (Ha + Hb)
+        rolled = depth.rolling({dim: 2}).construct("window_dim")
+        # Ha - Hb: diff is computed at U/V points
+        diff = rolled.diff("window_dim").squeeze("window_dim")
+        # rmax is computed at U/V points
+        rmax = np.abs(diff) / rolled.sum("window_dim")
+
+        # (rmax_a + rmax_b]) / 2 -> to compute rmax on T points
+        rolled = rmax.rolling({dim: 2}).construct("window_dim")
+        rmax = rolled.mean("window_dim", skipna=True)
+
+        # 1. Place on the correct index (shift -1 as we rolled twice)
+        # 2. Force first/last values = 0
+        # 3. Replace land values with 0
+        rmax = rmax.shift({dim: -1})
+        rmax[{dim: [0, -1]}] = 0
+        rmax = rmax.fillna(0)
+
+        both_rmax.append(rmax)
+
+    return np.maximum(*both_rmax)
+
+# =======================================================================================
+def calc_rmax_xr2(depth):
+    """
+    Calculate rmax: measure of steepness
+    This function returns the slope steepness criteria rmax, which is simply
+    |H[0] - H[1]| / (H[0] + H[1])
+    Parameters
+    ----------
+    depth: float
+        Bottom depth (units: m).
+    Returns
+    -------
+    rmax: float
+        Slope steepness value (units: None)
+    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.
+    """
+
+    # Loop over x and y
+    both_rmax = []
+    for dim in depth.dims:
+
+        # 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 and to force boundaries = NaN
+        rmax = rmax.rolling({dim: 2}).construct("vel_points")
+        rmax = rmax.shift({dim: -1})
+        rmax[{dim: [0, -1]}] = None
+
+        both_rmax.append(rmax)
+
+    rmax = xr.concat(both_rmax, "vel_points")
+    rmax = rmax.max("vel_points", skipna=True)
+
+    return rmax.fillna(0)
+
+# =======================================================================================
+def SlopeParam(raw_bathy, msk):
+
+    # This code is slightly modified from
+    # https://github.com/ESMG/pyroms/blob/master/bathy_smoother/bathy_smoother/bathy_tools.py
+    #
+    # This function computes the slope parameter defined as
+    #
+    #                Z_ij - Z_n
+    #                ----------
+    #                Z_ij + Z_n
+    #
+    # where Z_ij is the depth at some point i,j
+    # and Z_n is the neighbouring depth in the
+    # east,west,south or north sense.
+    #
+    # This code is adapted from the matlab code
+    # "LP Bathymetry" by Mathieu Dutour Sikiric
+    # http://drobilica.irb.hr/~mathieu/Bathymetry/index.html
+    # For a description of the method, see
+    # M. Dutour Sikiric, I. Janekovic, M. Kuzmic, A new approach to
+    # bathymetry smoothing in sigma-coordinate ocean models, Ocean
+    # Modelling 29 (2009) 128--136.
+
+    """
+    SloParMat = SlopeParam(raw_bathy)
+
+    raw_bathy: raw bathymetry interpolated on the model grid.
+               It must be a positive depths field.
+    msk      : is the mask of the grid
+
+    """
+
+    bathy = np.copy(raw_bathy)
+    # print bathy.shape
+    nj, ni = bathy.shape
+
+    # Masking land points: bathy is a positive depths field
+
+    bathy[msk == 0.0] = np.nan
+
+    nghb_pnts = np.array([[0, 1], [1, 0], [0, -1], [-1, 0]])
+
+    SloParMat = np.zeros(bathy.shape)
+
+    for j in range(1, nj - 1):
+        for i in range(1, ni - 1):
+            if msk[j, i] == 1:
+                slopar = 0.0
+
+                for n in range(4):
+                    j_nghb = j + nghb_pnts[n][0]
+                    i_nghb = i + nghb_pnts[n][1]
+                    if msk[j_nghb, i_nghb] == 1:
+                        dep1 = bathy[j, i]
+                        dep2 = bathy[j_nghb, i_nghb]
+                        delta = abs((dep1 - dep2) / (dep1 + dep2))
+                        slopar = np.maximum(slopar, delta)
+                SloParMat[j, i] = slopar
+
+    return SloParMat

pydomcfg/tests/test_rmax_calc.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python
+
+import netCDF4 as nc4
+import numpy as np
+import xarray as xr
+from lib_for_rmax import *
+from matplotlib import pyplot as plt
+
+rn_sbot_max = 7000.0
+rn_sbot_min = 10.0
+rn_rmax = 0.24
+
+nc = "/home/h01/dbruciaf/mod_dev/jmmp_MEs/amm7_mes_vgrid/MEs_2env_0.24_0.07_opt_v2/bathymeter_mo_ps43.nc"
+ds_bathy = xr.open_dataset(nc)
+bathy = ds_bathy["Bathymetry"].squeeze()
+
+# Computing LSM
+ocean = xr.where(bathy > 0, 1, 0)
+
+# setting max depth
+bathy = np.minimum(bathy, rn_sbot_max)
+
+# setting min depth
+bathy = np.maximum(bathy, rn_sbot_min) * ocean
+
+da_env = bathy.copy()
+zenv = da_env.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 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 (
+                bathy[jp1, im1]
+                + bathy[jp1, i]
+                + bathy[jp1, ip1]
+                + bathy[j, im1]
+                + bathy[j, ip1]
+                + bathy[jm1, im1]
+                + bathy[jm1, i]
+                + bathy[jm1, ip1]
+            ) > 0.0:
+                zenv[j, i] = rn_sbot_min
+
+# set scaling factor used for smoothing
+zrfact = (1.0 - rn_rmax) / (1.0 + rn_rmax)
+
+# initialise temporary evelope depth arrays
+ztmpi1 = zenv.copy()
+ztmpi2 = zenv.copy()
+ztmpj1 = zenv.copy()
+ztmpj2 = zenv.copy()
+
+# initial maximum slope parameter
+zrmax = 1.0
+zri = np.ones(zenv.shape)
+zrj = np.ones(zenv.shape)
+
+tol = 1.0e-8
+itr = 0
+max_itr = 10000
+
+while itr <= max_itr and (zrmax - rn_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])
+    for j in range(nj):
+        for i in range(ni):
+            zrmax = np.amax([zrmax, np.absolute(zri[j, i]), np.absolute(zrj[j, i])])
+
+    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] > rn_rmax:
+                ztmpi1[j, i] = zenv[j, ip1] * zrfact
+            if zri[j, i] < -rn_rmax:
+                ztmpi2[j, ip1] = zenv[j, i] * zrfact
+            if zrj[j, i] > rn_rmax:
+                ztmpj1[j, i] = zenv[jp1, i] * zrfact
+            if zrj[j, i] < -rn_rmax:
+                ztmpj2[jp1, i] = zenv[j, i] * zrfact
+
+    for j in range(nj):
+        for i in range(ni):
+            zenv[j, i] = np.amax(
+                [zenv[j, i], ztmpi1[j, i], ztmpi2[j, i], ztmpj1[j, i], ztmpj2[j, i]]
+            )
+
+    # ztmpi = np.maximum(ztmpi1, ztmpi2)
+    # ztmpj = np.maximum(ztmpj1, ztmpj2)
+    # zenv = np.maximum(zenv, np.maximum(ztmpi, ztmpj))
+
+# set all points to avoid undefined scale value warnings
+zenv = np.maximum(zenv, rn_sbot_min)
+
+da_env.data = zenv
+
+# calc rmax according to:
+
+# 1. calc_rmax_xr1: first implem. using the average
+rmax = np.amax(calc_rmax_xr1(da_env) * ocean).values
+print("1. calc_rmax_xr1: " + str(rmax))
+
+# 2. calc_rmax_xr2: last implem. using the maximum
+rmax = np.amax(calc_rmax_xr2(da_env) * ocean).values
+print("2. calc_rmax_xr2: " + str(rmax))
+
+# 2. calc_rmax_np: original numpy code
+rmax = np.amax(calc_rmax_np(zenv) * ocean.values)
+print("3. calc_rmax_np: " + str(rmax))
+
+# 3. SlopeParam: pyROMS very inefficient implem.
+rmax = np.amax(SlopeParam(zenv, ocean.values))
+print("4. SlopeParam: " + str(rmax))
+
+da_env.data = calc_rmax_np(zenv) * ocean.values - SlopeParam(zenv, ocean.values)
+da_env.plot()
+plt.show()