From 9002bb012e3e68cc434b47deda26fa3157b09dab Mon Sep 17 00:00:00 2001
From: Mattia Almansi <mattia.almansi@noc.ac.uk>
Date: Sun, 27 Jun 2021 16:24:58 +0100
Subject: [PATCH] Do not drop indexes when computing rmax (#40)

* Do not drop indexes when computing rmax

* implement #39

* replace land with 0s

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* handle land properly

* use max rather than mean and use xarray max instead of np maximum

* same as pyroms

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 pydomcfg/domzgr/zgr.py       |  3 ++-
 pydomcfg/tests/bathymetry.py | 43 ++++++++++++++++++++++++------------
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/pydomcfg/domzgr/zgr.py b/pydomcfg/domzgr/zgr.py
index 5910107..d24b2a9 100644
--- a/pydomcfg/domzgr/zgr.py
+++ b/pydomcfg/domzgr/zgr.py
@@ -20,7 +20,8 @@ def __init__(self, ds_bathy: Dataset, jpk: int):
         Parameters
         ----------
         ds_bathy: Dataset
-             xarray dataset including grid coordinates and bathymetry
+             xarray dataset including grid coordinates and bathymetry.
+             All bathymetry values MUST be >= 0, where 0 is land.
         jpk: int
              number of computational levels
         """
diff --git a/pydomcfg/tests/bathymetry.py b/pydomcfg/tests/bathymetry.py
index a135a6a..ea2bcd0 100644
--- a/pydomcfg/tests/bathymetry.py
+++ b/pydomcfg/tests/bathymetry.py
@@ -142,34 +142,49 @@ def _calc_rmax(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])
+    |H[0] - H[1]| / (H[0] + H[1])
 
     Parameters
     ----------
     depth: float
-            Bottom depth (units: m).
+        Bottom depth (units: m).
 
     Returns
     -------
     rmax: float
-            Slope steepness value (units: None)
+        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.
     """
-    depth = depth.reset_index(list(depth.dims))
 
+    # Mask land
+    depth = depth.where(depth > 0)
+
+    # Loop over x and y
     both_rmax = []
     for dim in depth.dims:
 
-        # (H[0] - H[1]) / (H[0] + H[1])
-        depth_diff = depth.diff(dim)
-        depth_rolling_sum = depth.rolling({dim: 2}).sum().dropna(dim)
-        rmax = depth_diff / depth_rolling_sum
+        # 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})
 
-        # (R[0] + R[1]) / 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)