From 4fb1bdd1b802b812ac2d6a09d5d9e7e97f97a587 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 15 Jun 2021 21:19:04 +0100 Subject: [PATCH 01/19] including calc_rmax in utils to be used with sco class --- pydomcfg/utils.py | 45 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index e272acc..cee7c8d 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -8,7 +8,7 @@ import xarray as xr from xarray import DataArray, Dataset - +# ----------------------------------------------------------------------------- def is_nemo_none(var: Optional[float] = None) -> bool: """ Assess if a namelist parameter is None @@ -25,7 +25,50 @@ def is_nemo_none(var: Optional[float] = None) -> bool: """ return var in [None, 999999.0] +# ----------------------------------------------------------------------------- +def calc_rmax(depth: DataArray) -> DataArray: + """ + 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 + ------- + rmax: DataArray + Maximum slope parameter (units: None) + """ + + depth = depth.reset_index(list(depth.dims)) + + 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 + + # (rmax_a + rmax_b) / 2 + rmax = rmax.rolling({dim: 2}).mean().dropna(dim) + + # Fill first row and column + rmax = rmax.pad({dim: (1, 1)}, constant_values=0) + + both_rmax.append(np.abs(rmax)) + + return np.maximum(*both_rmax) +# ----------------------------------------------------------------------------- def generate_cartesian_grid( ppe1_m, ppe2_m, From 6d3163611d2e9cf14151d78aecaeb0b5a257a823 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 15 Jun 2021 21:21:44 +0100 Subject: [PATCH 02/19] including calc_rmax in utils to be used with sco class --- pydomcfg/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index cee7c8d..4ce74d1 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -8,6 +8,7 @@ import xarray as xr from xarray import DataArray, Dataset + # ----------------------------------------------------------------------------- def is_nemo_none(var: Optional[float] = None) -> bool: """ @@ -25,11 +26,12 @@ def is_nemo_none(var: Optional[float] = None) -> bool: """ return var in [None, 999999.0] + # ----------------------------------------------------------------------------- def calc_rmax(depth: DataArray) -> DataArray: """ Calculate rmax: measure of steepness - This function returns the maximum slope paramater + This function returns the maximum slope paramater rmax = abs(Hb - Ha) / (Ha + Hb) @@ -68,6 +70,7 @@ def calc_rmax(depth: DataArray) -> DataArray: return np.maximum(*both_rmax) + # ----------------------------------------------------------------------------- def generate_cartesian_grid( ppe1_m, From 5d0926c188fa5dbf19eef04ebb4d51c948f803de Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 15 Jun 2021 21:23:35 +0100 Subject: [PATCH 03/19] including calc_rmax in utils to be used with sco class --- pydomcfg/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index 4ce74d1..9b81926 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -46,7 +46,7 @@ def calc_rmax(depth: DataArray) -> DataArray: Bottom depth (units: m). Returns ------- - rmax: DataArray + DataArray Maximum slope parameter (units: None) """ From 80b4b58645fa3f3fc219eef0fb6ee05b094a8a5d Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 15 Jun 2021 21:39:01 +0100 Subject: [PATCH 04/19] start work on sco --- pydomcfg/domzgr/sco.py | 199 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 199 insertions(+) create mode 100644 pydomcfg/domzgr/sco.py diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py new file mode 100644 index 0000000..a349fe3 --- /dev/null +++ b/pydomcfg/domzgr/sco.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python + +""" +Class to generate NEMO v4.0 s-coordinates +""" + +from typing import Optional # , Tuple + +# import numpy as np +from xarray import Dataset + +from .zgr import Zgr + +# from pydomcfg.utils import is_nemo_none + + +class Sco(Zgr): + """ + Class to generate terrain-following coordinates dataset objects. + Currently, four types of terrain-following grids can be genrated: + *) uniform sigma-coordinates (Phillips 1957) + *) stretched s-coordinates with Song & Haidvogel 1994 stretching + *) stretched s-coordinates with Siddorn & Furner 2013 stretching + *) stretched s-coordinates with Madec et al. 1996 stretching + + 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_sco,s_sh94,s_sf12,s_tanh} subroutine + *) Phillips, J. Meteorol., 14, 184-185, 1957. + *) Song & Haidvogel, J. Comp. Phy., 115, 228-244, 1994. + *) Siddorn & Furner, Oce. Mod. 66:1-13, 2013. + *) Madec, Delecluse, Crepon & Lott, JPO 26(8):1393-1408, 1996. + """ + + # -------------------------------------------------------------------------- + def __call__( + self, + bot_min: float, + bot_max: float, + hc: float = 0.0, + rmax: Optional[float] = None, + stretch: Optional[str] = None, + psurf: Optional[float] = None, + pbott: Optional[float] = None, + alpha: Optional[float] = None, + efold: Optional[float] = None, + pbot2: Optional[float] = None, + ln_e3_dep: bool = True, + ) -> Dataset: + """ + Generate NEMO terrain-following model levels. + + Parameters + ---------- + bot_min: float + Minimum depth of bottom topography surface (>0) (m) + bot_max: float + Maximum depth of bottom topography surface (>0) (m) + hc: float + critical depth for transition from uniform sigma + to stretched s-coordinates (>0) (m) + rmax: float, optional + maximum slope parameter value allowed + stretch: str, optional + Type of stretching applied: + *) None = no stretching, i.e. uniform sigma-coord. + *) "sh94" = Song & Haidvogel 1994 stretching + *) "sf12" = Siddorn & Furner 2013 stretching + *) "md96" = Madec et al. 1996 stretching + psurf: float, optional + sh94: surface control parameter (0<= psurf <=20) + md96: surface control parameter (0<= psurf <=20) + sf12: thickness of first model layer (m) + pbott: float, optional + sh94: bottom control parameter (0<= pbott <=1) + md96: bottom control parameter (0<= pbott <=1) + sf12: scaling factor for computing thickness + of bottom level Zb + alpha: float, optional + sf12: stretching parameter + efold: float, optional + sf12: efold length scale for transition from sigma + to stretched coord + pbot2: float, optional + sf12 offset for calculating Zb = H*pbott + pbot2 + ln_e3_dep: bool + Logical flag to comp. e3 as fin. diff. (True) or + analyt. (False) (default = True) + + Returns + ------- + Dataset + Describing the 3D geometry of the model + """ + + self._bot_min = bot_min + self._bot_max = bot_max + self._hc = hc + self.rmax = rmax + self._stretch = stretch + self._ln_e3_dep = ln_e3_dep + + # set stretching parameters after checking their consistency + if self._stretch: + self._set_stretch_par(psurf, pbott, alpha, efold, pbot2) + + ds = self._init_ds() + + # compute envelope bathymetry + ds_env = self._compute_env(ds) + + # compute sigma-coordinates for z3 computation + kindx = ds_env["z"] + sigma = (self._compute_sigma(kk) for kk in kindx) + self._sigT, self._sigW = sigma + + # compute z3 depths of zco vertical levels + # dsz = self._sco_z3(ds_env) + + # compute e3 scale factors + # dse = self._compute_e3(dsz) if self._ln_e3_dep else self._analyt_e3(dsz) + + # return dse + + # -------------------------------------------------------------------------- + def _set_stretch_par(self, psurf, pbott, alpha, efold, pbot2): + """ + Set stretching parameters after checking + consistency of input parameters + """ + if not (psurf and pbott): + if self._stretch == "sh94": + srf = "rn_theta" + bot = "rn_bb" + elif self._stretch == "md96": + srf = "rn_theta" + bot = "rn_thetb" + elif self._stretch == "sf12": + srf = "rn_zs" + bot = "rn_zb_a" + msg = ( + srf + + " and " + + bot + + "MUST be set when using " + + self._stretch + + " stretching." + ) + raise ValueError(msg) + + if self._stretch == "sf12": + if not (alpha and efold and pbot2): + msg = "rn_alpha, rn_efold and rn_zb_b MUST be set when \ + using sf12 stretching." + raise ValueError(msg) + + # setting stretching parameters + self._psurf = psurf if psurf else 0.0 + self._pbott = pbott if pbott else 0.0 + self._alpha = alpha if alpha else 0.0 + self._efold = efold if efold else 0.0 + self._pbot2 = pbot2 if pbot2 else 0.0 + + # -------------------------------------------------------------------------- + def _compute_env(self, ds: Dataset) -> Dataset: + """ + Compute the envelope bathymetry surface by applying the + Martinho & Batteen (2006) smoothing algorithm to the + actual topography to reduce the maximum value of the slope + parameter + r = abs(Hb-Ha) / (Ha+Hb) + + where Ha and Hb are the depths of adjacent grid cells. + The maximum slope parameter is reduced to be <= rmax. + + Reference: + *) Martinho & Batteen, Oce. Mod. 13(2):166-175, 2006. + + Parameters + ---------- + ds: Dataset + xarray dataset with the 2D bottom topography DataArray + Returns + ------- + ds: Dataset + xarray dataset with the 2D envelope bathymetry DataArray + """ + + return ds From c8a28f598cb44b61a1e4b79aa2b94049fecdf77a Mon Sep 17 00:00:00 2001 From: oceandie Date: Sat, 19 Jun 2021 18:52:32 +0100 Subject: [PATCH 05/19] merging main to sco_dev and adding _calc_rmax to utils --- pydomcfg/domzgr/sco.py | 65 ++++++++++++++++++------------------------ pydomcfg/utils.py | 16 +++++++---- 2 files changed, 39 insertions(+), 42 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index a349fe3..bfb897b 100644 --- a/pydomcfg/domzgr/sco.py +++ b/pydomcfg/domzgr/sco.py @@ -7,7 +7,7 @@ from typing import Optional # , Tuple # import numpy as np -from xarray import Dataset +from xarray import DataArray, Dataset from .zgr import Zgr @@ -100,7 +100,7 @@ def __call__( Returns ------- Dataset - Describing the 3D geometry of the model + Describing the 3D geometry of the model """ self._bot_min = bot_min @@ -112,17 +112,20 @@ def __call__( # set stretching parameters after checking their consistency if self._stretch: - self._set_stretch_par(psurf, pbott, alpha, efold, pbot2) + self._check_stretch_par(psurf, pbott, alpha, efold, pbot2) + self._psurf = psurf or 0.0 + self._pbott = pbott or 0.0 + self._alpha = alpha or 0.0 + self._efold = efold or 0.0 + self._pbot2 = pbot2 or 0.0 - ds = self._init_ds() + # ds = self._init_ds() - # compute envelope bathymetry - ds_env = self._compute_env(ds) + # compute envelope bathymetry DataArray + self._envlp = self._compute_env(self._bathy["Bathymetry"]) # compute sigma-coordinates for z3 computation - kindx = ds_env["z"] - sigma = (self._compute_sigma(kk) for kk in kindx) - self._sigT, self._sigW = sigma + self._sigmas = self._compute_sigma(self._z) # compute z3 depths of zco vertical levels # dsz = self._sco_z3(ds_env) @@ -130,13 +133,13 @@ def __call__( # compute e3 scale factors # dse = self._compute_e3(dsz) if self._ln_e3_dep else self._analyt_e3(dsz) - # return dse + # addind this only to not make darglint complying + return self._merge_z3_and_e3(self._envlp, self._envlp, self._envlp, self._envlp) # -------------------------------------------------------------------------- - def _set_stretch_par(self, psurf, pbott, alpha, efold, pbot2): + def _check_stretch_par(self, psurf, pbott, alpha, efold, pbot2): """ - Set stretching parameters after checking - consistency of input parameters + Check consistency of stretching parameters """ if not (psurf and pbott): if self._stretch == "sh94": @@ -148,31 +151,19 @@ def _set_stretch_par(self, psurf, pbott, alpha, efold, pbot2): elif self._stretch == "sf12": srf = "rn_zs" bot = "rn_zb_a" - msg = ( - srf - + " and " - + bot - + "MUST be set when using " - + self._stretch - + " stretching." + raise ValueError( + f"{srf} and {bot} MUST be set when using {self._stretch} stretching." ) - raise ValueError(msg) if self._stretch == "sf12": if not (alpha and efold and pbot2): - msg = "rn_alpha, rn_efold and rn_zb_b MUST be set when \ - using sf12 stretching." - raise ValueError(msg) - - # setting stretching parameters - self._psurf = psurf if psurf else 0.0 - self._pbott = pbott if pbott else 0.0 - self._alpha = alpha if alpha else 0.0 - self._efold = efold if efold else 0.0 - self._pbot2 = pbot2 if pbot2 else 0.0 + raise ValueError( + "rn_alpha, rn_efold and rn_zb_b MUST be set when using \ + sf12 stretching." + ) # -------------------------------------------------------------------------- - def _compute_env(self, ds: Dataset) -> Dataset: + def _compute_env(self, da: DataArray) -> DataArray: """ Compute the envelope bathymetry surface by applying the Martinho & Batteen (2006) smoothing algorithm to the @@ -188,12 +179,12 @@ def _compute_env(self, ds: Dataset) -> Dataset: Parameters ---------- - ds: Dataset - xarray dataset with the 2D bottom topography DataArray + da: DataArray + xarray DataArray of the 2D bottom topography Returns ------- - ds: Dataset - xarray dataset with the 2D envelope bathymetry DataArray + DataArray + xarray DataArray of the 2D envelope bathymetry """ - return ds + return da diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index 8c8d65a..bcbf94e 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -10,6 +10,7 @@ NEMO_NONE = 999_999 + def _is_nemo_none(var: Hashable) -> bool: """Assess if a NEMO parameter is None""" return (var or NEMO_NONE) == NEMO_NONE @@ -22,29 +23,34 @@ def _are_nemo_none(var: Iterable) -> Iterator[bool]: # ----------------------------------------------------------------------------- -def calc_rmax(depth: DataArray) -> DataArray: +def _calc_rmax(depth: DataArray) -> float: """ 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). Reference: - *) Mellor, Oey & Ezer, J Atm. Oce. Tech. 15(5):1122-1131, 1998. + *) Mellor, Oey & Ezer, J Atm. Oce. Tech. 15(5):1122-1131, 1998. Parameters ---------- depth: DataArray Bottom depth (units: m). + Returns ------- - DataArray + float Maximum slope parameter (units: None) """ - depth = depth.reset_index(list(depth.dims)) + # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + # Do we actually need this? mypy complains since + # DataArray.reset_indexe() returns Optional["DataArray"] + # + # depth = depth.reset_index(list(depth.dims)) both_rmax = [] for dim in depth.dims: From be52a5cf940b2387954f23c9f69553a20cc884f0 Mon Sep 17 00:00:00 2001 From: oceandie Date: Mon, 21 Jun 2021 10:07:25 +0100 Subject: [PATCH 06/19] add computation of envelope bathymetry --- pydomcfg/domzgr/sco.py | 80 ++++++++++++++++++----- pydomcfg/utils.py | 142 +++++++++++++++++++++++++++++++++++------ 2 files changed, 186 insertions(+), 36 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index bfb897b..67e652e 100644 --- a/pydomcfg/domzgr/sco.py +++ b/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,7 +191,7 @@ def _compute_env(self, da: DataArray) -> DataArray: Parameters ---------- - da: DataArray + depth: DataArray xarray DataArray of the 2D bottom topography Returns ------- @@ -187,4 +199,40 @@ def _compute_env(self, da: DataArray) -> 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 diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index bcbf94e..b75d86b 100644 --- a/pydomcfg/utils.py +++ b/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, From 966bf236476259843bf67774dc66ded8b812441a Mon Sep 17 00:00:00 2001 From: oceandie Date: Fri, 25 Jun 2021 21:23:41 +0100 Subject: [PATCH 07/19] amending doc --- pydomcfg/utils.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index b75d86b..dda2880 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -26,25 +26,25 @@ def _are_nemo_none(var: Iterable) -> Iterator[bool]: # ----------------------------------------------------------------------------- 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 - ------- + Returns + ------- DataArray - 2D maximum slope parameter (units: None) + 2D maximum slope parameter (units: None) """ depth = DataArray(depth.reset_index(list(depth.dims))) From 4567fde3d96ce284f9dc945753390783ec1377cc Mon Sep 17 00:00:00 2001 From: oceandie Date: Fri, 25 Jun 2021 21:29:18 +0100 Subject: [PATCH 08/19] little test just for me --- pydomcfg/tests/test_sco.py | 66 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 pydomcfg/tests/test_sco.py diff --git a/pydomcfg/tests/test_sco.py b/pydomcfg/tests/test_sco.py new file mode 100644 index 0000000..14e542b --- /dev/null +++ b/pydomcfg/tests/test_sco.py @@ -0,0 +1,66 @@ +""" +Tests for sco +""" + +import numpy as np +import pytest +import xarray as xr + +from .bathymetry import Bathymetry +from pydomcfg.utils import _calc_rmax, _smooth_MB06 +from pydomcfg.domzgr.sco import Sco + +ds_bathy = Bathymetry(1000,1000, 200, 200).sea_mount(5000., 1) +ds_bathy["Bathymetry"]=ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"]>550, 0.) +sco = Sco(ds_bathy, jpk=51) + +ocean = ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"]==0,1) + +np.nanmax(_calc_rmax(ds_bathy["Bathymetry"])*ocean) +ds_s = sco(min_dep=500.,max_dep=3500.,rmax=0.01) +np.nanmax(_calc_rmax(ds_s["hbatt"])*ocean) + +ds_s["hbatt"].isel({'y':100}).plot() +ds_bathy["Bathymetry"].isel({'y':100}).plot() +plt.gca().invert_yaxis() +plt.show() + + + +#def test_zco_orca2(): +# """ +# The test consists in reproducing ORCA2 grid +# z3T/W and e3T/W as computed by NEMO v3.6 +# (see pag 62 of v3.6 manual for the input parameters). +# This test validates stretched grids with analytical e3 and no +# double tanh. +# """ + +# # Bathymetry dataset +# ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3) +# +# # zco grid generator +# zco = Zco(ds_bathy, jpk=31) + +# # zco mesh with analytical e3 using ORCA2 input parameters +# # See pag 62 of v3.6 manual for the input parameters +# dsz_an = zco( +# ppdzmin=10.0, +# pphmax=5000.0, +# ppkth=21.43336197938, +# ppacr=3, +# ppsur=-4762.96143546300, +# ppa0=255.58049070440, +# ppa1=245.58132232490, +# ldbletanh=False, +# ln_e3_dep=False, +# ) + +# # reference ocean.output values are +# # given with 4 digits precision +# eps = 1.0e-5 +# for n, varname in enumerate(["z3T", "z3W", "e3T", "e3W"]): +# expected = ORCA2_VGRID[:, n] +# actual = dsz_an[varname].squeeze().values +# np.testing.assert_allclose(expected, actual, rtol=eps, atol=0) + From 8b90ea97a772ffd1de76c56a7ff8555f3d599ce9 Mon Sep 17 00:00:00 2001 From: oceandie Date: Fri, 25 Jun 2021 21:30:12 +0100 Subject: [PATCH 09/19] little test just for me --- pydomcfg/tests/test_sco.py | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/pydomcfg/tests/test_sco.py b/pydomcfg/tests/test_sco.py index 14e542b..e3efe60 100644 --- a/pydomcfg/tests/test_sco.py +++ b/pydomcfg/tests/test_sco.py @@ -6,28 +6,28 @@ import pytest import xarray as xr -from .bathymetry import Bathymetry -from pydomcfg.utils import _calc_rmax, _smooth_MB06 from pydomcfg.domzgr.sco import Sco +from pydomcfg.utils import _calc_rmax, _smooth_MB06 -ds_bathy = Bathymetry(1000,1000, 200, 200).sea_mount(5000., 1) -ds_bathy["Bathymetry"]=ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"]>550, 0.) +from .bathymetry import Bathymetry + +ds_bathy = Bathymetry(1000, 1000, 200, 200).sea_mount(5000.0, 1) +ds_bathy["Bathymetry"] = ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"] > 550, 0.0) sco = Sco(ds_bathy, jpk=51) -ocean = ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"]==0,1) +ocean = ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"] == 0, 1) -np.nanmax(_calc_rmax(ds_bathy["Bathymetry"])*ocean) -ds_s = sco(min_dep=500.,max_dep=3500.,rmax=0.01) -np.nanmax(_calc_rmax(ds_s["hbatt"])*ocean) +np.nanmax(_calc_rmax(ds_bathy["Bathymetry"]) * ocean) +ds_s = sco(min_dep=500.0, max_dep=3500.0, rmax=0.01) +np.nanmax(_calc_rmax(ds_s["hbatt"]) * ocean) -ds_s["hbatt"].isel({'y':100}).plot() -ds_bathy["Bathymetry"].isel({'y':100}).plot() +ds_s["hbatt"].isel({"y": 100}).plot() +ds_bathy["Bathymetry"].isel({"y": 100}).plot() plt.gca().invert_yaxis() plt.show() - -#def test_zco_orca2(): +# def test_zco_orca2(): # """ # The test consists in reproducing ORCA2 grid # z3T/W and e3T/W as computed by NEMO v3.6 @@ -63,4 +63,3 @@ # expected = ORCA2_VGRID[:, n] # actual = dsz_an[varname].squeeze().values # np.testing.assert_allclose(expected, actual, rtol=eps, atol=0) - From 1dd355e338a6ef77f14043118482b67085c408a0 Mon Sep 17 00:00:00 2001 From: oceandie Date: Sun, 27 Jun 2021 16:43:07 +0100 Subject: [PATCH 10/19] cleaning --- pydomcfg/tests/test_sco.py | 65 -------------------------------------- 1 file changed, 65 deletions(-) delete mode 100644 pydomcfg/tests/test_sco.py diff --git a/pydomcfg/tests/test_sco.py b/pydomcfg/tests/test_sco.py deleted file mode 100644 index e3efe60..0000000 --- a/pydomcfg/tests/test_sco.py +++ /dev/null @@ -1,65 +0,0 @@ -""" -Tests for sco -""" - -import numpy as np -import pytest -import xarray as xr - -from pydomcfg.domzgr.sco import Sco -from pydomcfg.utils import _calc_rmax, _smooth_MB06 - -from .bathymetry import Bathymetry - -ds_bathy = Bathymetry(1000, 1000, 200, 200).sea_mount(5000.0, 1) -ds_bathy["Bathymetry"] = ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"] > 550, 0.0) -sco = Sco(ds_bathy, jpk=51) - -ocean = ds_bathy["Bathymetry"].where(ds_bathy["Bathymetry"] == 0, 1) - -np.nanmax(_calc_rmax(ds_bathy["Bathymetry"]) * ocean) -ds_s = sco(min_dep=500.0, max_dep=3500.0, rmax=0.01) -np.nanmax(_calc_rmax(ds_s["hbatt"]) * ocean) - -ds_s["hbatt"].isel({"y": 100}).plot() -ds_bathy["Bathymetry"].isel({"y": 100}).plot() -plt.gca().invert_yaxis() -plt.show() - - -# def test_zco_orca2(): -# """ -# The test consists in reproducing ORCA2 grid -# z3T/W and e3T/W as computed by NEMO v3.6 -# (see pag 62 of v3.6 manual for the input parameters). -# This test validates stretched grids with analytical e3 and no -# double tanh. -# """ - -# # Bathymetry dataset -# ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3) -# -# # zco grid generator -# zco = Zco(ds_bathy, jpk=31) - -# # zco mesh with analytical e3 using ORCA2 input parameters -# # See pag 62 of v3.6 manual for the input parameters -# dsz_an = zco( -# ppdzmin=10.0, -# pphmax=5000.0, -# ppkth=21.43336197938, -# ppacr=3, -# ppsur=-4762.96143546300, -# ppa0=255.58049070440, -# ppa1=245.58132232490, -# ldbletanh=False, -# ln_e3_dep=False, -# ) - -# # reference ocean.output values are -# # given with 4 digits precision -# eps = 1.0e-5 -# for n, varname in enumerate(["z3T", "z3W", "e3T", "e3W"]): -# expected = ORCA2_VGRID[:, n] -# actual = dsz_an[varname].squeeze().values -# np.testing.assert_allclose(expected, actual, rtol=eps, atol=0) From eac503baa8c84b6e028cee45b1f2bf1a586edc43 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 29 Jun 2021 09:24:26 +0100 Subject: [PATCH 11/19] add envelope computation --- pydomcfg/domzgr/sco.py | 107 +++++++++++++++++++++++++++-------------- pydomcfg/utils.py | 55 ++++++++++++--------- 2 files changed, 102 insertions(+), 60 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index 67e652e..71ea3e5 100644 --- a/pydomcfg/domzgr/sco.py +++ b/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 diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index dda2880..b670690 100644 --- a/pydomcfg/utils.py +++ b/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, From 95896795b7781dd2529f464e3dfa1da5ffd5d774 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 29 Jun 2021 14:06:49 +0100 Subject: [PATCH 12/19] cleaning syntax --- pydomcfg/domzgr/sco.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index 71ea3e5..390ea94 100644 --- a/pydomcfg/domzgr/sco.py +++ b/pydomcfg/domzgr/sco.py @@ -4,13 +4,14 @@ Class to generate NEMO v4.0 s-coordinates """ -from typing import Optional # , Tuple +from itertools import product +from typing import Optional import numpy as np import xarray as xr from xarray import DataArray, Dataset -from pydomcfg.utils import _smooth_MB06 # , _calc_rmax +from pydomcfg.utils import _smooth_MB06 from .zgr import Zgr @@ -193,7 +194,7 @@ 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" + it MUST have only two dimensions Returns ------- DataArray @@ -251,11 +252,10 @@ def _compute_env(self, depth: DataArray) -> DataArray: # ------------------------------------------------------------ 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 + for j, i in product(ngb_pnt, repeat=2): + if not (j == 0 and 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) From 3119aada462da5a665d7dbc72fd511107cb9560e Mon Sep 17 00:00:00 2001 From: diegobruciaferri Date: Tue, 29 Jun 2021 14:57:24 +0100 Subject: [PATCH 13/19] Fix little inconsistencies at the boundaries Co-authored-by: Mattia Almansi --- pydomcfg/domzgr/sco.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index 390ea94..aa29bc1 100644 --- a/pydomcfg/domzgr/sco.py +++ b/pydomcfg/domzgr/sco.py @@ -250,19 +250,10 @@ def _compute_env(self, depth: DataArray) -> DataArray: # 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, i in product(ngb_pnt, repeat=2): - if not (j == 0 and 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) + cst_lsm = lsm.rolling({dim: 3 for dim in lsm.dims}, min_periods=2).sum() + cst_lsm = cst_lsm.shift({dim: -1 for dim in lsm.dims}) + cst_lsm = (cst_lsm > 0) & (lsm == 0) 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) From ef02eec66e5fdd19a3f4064753478bcce1a5df68 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 29 Jun 2021 15:03:03 +0100 Subject: [PATCH 14/19] cleaning --- pydomcfg/domzgr/sco.py | 45 ------------------------------------------ 1 file changed, 45 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index aa29bc1..7b69ab6 100644 --- a/pydomcfg/domzgr/sco.py +++ b/pydomcfg/domzgr/sco.py @@ -4,7 +4,6 @@ Class to generate NEMO v4.0 s-coordinates """ -from itertools import product from typing import Optional import numpy as np @@ -207,54 +206,10 @@ def _compute_env(self, depth: DataArray) -> DataArray: # set first land point adjacent to a wet cell to # min_dep as this needs to be included in smoothing - - # ------------------------------------------------------------ - # 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.rolling({dim: 3 for dim in lsm.dims}, min_periods=2).sum() cst_lsm = cst_lsm.shift({dim: -1 for dim in lsm.dims}) cst_lsm = (cst_lsm > 0) & (lsm == 0) zenv = depth.where(cst_lsm == 0, self._min_dep) - # ------------------------------------------------------------ zenv = _smooth_MB06(zenv, self._rmax) zenv = zenv.where(zenv > self._min_dep, self._min_dep) From 02842f7e63b1957dfcc3ebf0d8bbcf1425bd9145 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 29 Jun 2021 16:49:22 +0100 Subject: [PATCH 15/19] updating name of stretched var for consistency --- pydomcfg/domzgr/zgr.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pydomcfg/domzgr/zgr.py b/pydomcfg/domzgr/zgr.py index d24b2a9..e713a7f 100644 --- a/pydomcfg/domzgr/zgr.py +++ b/pydomcfg/domzgr/zgr.py @@ -67,7 +67,7 @@ def _compute_sigma(self, kindx: DataArray) -> Tuple[DataArray, ...]: @staticmethod def _compute_z3( su: DataArray, - ss1: DataArray, + ss: DataArray, a1: Union[float, DataArray], a2: Union[float, DataArray], a3: Union[float, DataArray], @@ -82,7 +82,7 @@ def _compute_z3( su: DataArray uniform non-dimensional vertical coordinate s, aka sigma-coordinates. 0 <= s <= 1 - ss1: DataArray + ss: DataArray stretched non-dimensional vertical coordinate s, 0 <= s <= 1 a1, a2, a3: float, DataArray @@ -98,7 +98,7 @@ def _compute_z3( z is downward positive. """ - return a1 + a2 * su + a3 * ss1 + return a1 + a2 * su + a3 * ss # ------------------------------------------------------------------------- @staticmethod From bb206e22b03efcb3f1364ac97161b3115a203ed8 Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 29 Jun 2021 16:55:56 +0100 Subject: [PATCH 16/19] updating name of stretched var for consistency --- pydomcfg/domzgr/zco.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pydomcfg/domzgr/zco.py b/pydomcfg/domzgr/zco.py index 6ec0c14..0f500ee 100644 --- a/pydomcfg/domzgr/zco.py +++ b/pydomcfg/domzgr/zco.py @@ -218,17 +218,17 @@ def _zco_z3(self) -> Tuple[DataArray, ...]: if self._is_uniform: # Uniform zco grid su = -sigma - s1 = DataArray((0.0)) + ss = DataArray((0.0)) a1 = a3 = 0.0 a2 = self._pphmax else: # Stretched zco grid su = -sigma_p1 - s1 = self._stretch_zco(-sigma) + ss = self._stretch_zco(-sigma) a1 = self._ppsur a2 = self._ppa0 * (self._jpk - 1.0) a3 = self._ppa1 * self._ppacr - z3 = self._compute_z3(su, s1, a1, a2, a3) + z3 = self._compute_z3(su, ss, a1, a2, a3) if self._ldbletanh: # Add double tanh term From 675053358d3446b76978e54cc23b221e0676b07f Mon Sep 17 00:00:00 2001 From: oceandie Date: Tue, 29 Jun 2021 18:38:00 +0100 Subject: [PATCH 17/19] add sh94 and md96 stretching and computation of gdep and e3 --- pydomcfg/domzgr/sco.py | 157 ++++++++++++++++++++++++++++++++++++++--- 1 file changed, 148 insertions(+), 9 deletions(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index 7b69ab6..d98561a 100644 --- a/pydomcfg/domzgr/sco.py +++ b/pydomcfg/domzgr/sco.py @@ -4,7 +4,7 @@ Class to generate NEMO v4.0 s-coordinates """ -from typing import Optional +from typing import Optional, Tuple import numpy as np import xarray as xr @@ -43,7 +43,6 @@ class Sco(Zgr): *) Madec, Delecluse, Crepon & Lott, JPO 26(8):1393-1408, 1996. """ - # -------------------------------------------------------------------------- def __call__( self, min_dep: float, @@ -138,17 +137,16 @@ def __call__( self._sigmas = self._compute_sigma(self._z) # compute z3 depths of zco vertical levels - # dsz = self._sco_z3(ds_env) + z3t, z3w = self._sco_z3 # compute e3 scale factors - # dse = self._compute_e3(dsz) if self._ln_e3_dep else self._analyt_e3(dsz) + e3t, e3w = self._compute_e3(z3t, z3w) # addind this only to not make darglint complying - ds = self._bathy.copy() - ds["hbatt"] = self._envlp - return ds + # ds = self._bathy.copy() + # ds["hbatt"] = self._envlp + return self._merge_z3_and_e3(z3t, z3w, e3t, e3w) - # -------------------------------------------------------------------------- def _check_stretch_par(self, psurf, pbott, alpha, efold, pbot2): """ Check consistency of stretching parameters @@ -174,7 +172,6 @@ def _check_stretch_par(self, psurf, pbott, alpha, efold, pbot2): sf12 stretching." ) - # -------------------------------------------------------------------------- def _compute_env(self, depth: DataArray) -> DataArray: """ Compute the envelope bathymetry surface by applying the @@ -215,3 +212,145 @@ def _compute_env(self, depth: DataArray) -> DataArray: zenv = zenv.where(zenv > self._min_dep, self._min_dep) return zenv + + @property + def _sco_z3(self) -> Tuple[DataArray, ...]: + """Compute and return z3{t,w} for s-coordinates grids""" + + grids = ("T", "W") + sigmas = self._sigmas + scosrf = self._envlp * 0.0 # unperturbed free-surface + + both_z3 = [] + for grid, sigma in zip(grids, sigmas): + + if self._stretch: + # Stretched sco grid + su = -sigma + ss = self._stretch_sco(-sigma) + a1 = scosrf + a2 = 0.0 + a3 = self._envlp + if self._stretch != "sf12": + a2 += self._hc + a3 -= self._hc + else: + # Uniform sco grid + su = -sigma + ss = DataArray((0.0)) + a1 = a3 = scosrf + a2 = self._envlp + + z3 = self._compute_z3(su, ss, a1, a2, a3) + + both_z3 += [z3] + + return tuple(both_z3) + + def _stretch_sco(self, sigma: DataArray) -> DataArray: + """ + Wrapping method for calling generalised analytical + stretching function for terrain-following s-coordinates. + + Parameters + ---------- + sigma: DataArray + Uniform non-dimensional sigma-coordinate: + MUST BE positive, i.e. 0 <= sigma <= 1 + + Returns + ------- + DataArray + Stretched coordinate + """ + if self._stretch == "sh94": + ss = self._sh94(sigma) + elif self._stretch == "md96": + ss = self._md96(sigma) + # elif self._stretch == "sf12": + # ss = self._sf12(sigma) + + return ss + + def _sh94(self, sigma: DataArray) -> DataArray: + """ + Song and Haidvogel 1994 analytical stretching + function for terrain-following s-coordinates. + + Reference: + Song & Haidvogel, J. Comp. Phy., 115, 228-244, 1994. + + Parameters + ---------- + sigma: DataArray + Uniform non-dimensional sigma-coordinate: + MUST BE positive, i.e. 0 <= sigma <= 1 + + Returns + ------- + DataArray + Stretched coordinate + """ + ca = self._psurf + cb = self._pbott + + if ca == 0.0: + ss = sigma + else: + ss = (1.0 - cb) * np.sinh(ca * sigma) / np.sinh(ca) + cb * ( + (np.tanh(ca * (sigma + 0.0)) - np.tanh(0.0 * ca)) + / (2.0 * np.tanh(0.0 * ca)) + ) + + return ss + + def _md96(self, sigma: DataArray) -> DataArray: + """ + Madec et al. 1996 analytical stretching + function for terrain-following s-coordinates. + + Reference: + pag 65 of NEMO Manual + Madec, Lott, Delecluse and Crepon, 1996. JPO, 26, 1393-1408 + + Parameters + ---------- + sigma: DataArray + Uniform non-dimensional sigma-coordinate: + MUST BE positive, i.e. 0 <= sigma <= 1 + + Returns + ------- + DataArray + Stretched coordinate + """ + ca = self._psurf + cb = self._pbott + + ss = ( + (np.tanh(ca * (sigma + cb)) - np.tanh(cb * ca)) + * (np.cosh(ca) + np.cosh(ca * (2.0e0 * cb - 1.0e0))) + / (2.0 * np.sinh(ca)) + ) + + return ss + + # def _sf12(self, sigma: DataArray) -> DataArray: + # """ + # Siddorn and Furner 2012 analytical stretching + # function for terrain-following s-coordinates. + # + # Reference: + # Siddorn & Furner, Oce. Mod. 66:1-13, 2013. + + # Parameters + # ---------- + # sigma: DataArray + # Uniform non-dimensional sigma-coordinate: + # MUST BE positive, i.e. 0 <= sigma <= 1 + # + # Returns + # ------- + # DataArray + # Stretched coordinate + # """ From 8e08398a0ea94fe2fb6fc7f444f8c41aece81e47 Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Wed, 30 Jun 2021 11:41:10 +0100 Subject: [PATCH 18/19] Refactor MB06 (#51) * refactor to avoid loops * apply Diego's suggestions * remove useless return after error --- pydomcfg/utils.py | 127 +++++++++++++++++++--------------------------- 1 file changed, 53 insertions(+), 74 deletions(-) diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index b670690..27f13b8 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -2,7 +2,6 @@ Utilities """ -# from itertools import product from typing import Hashable, Iterable, Iterator, Optional import numpy as np @@ -80,10 +79,15 @@ def _calc_rmax(depth: DataArray) -> DataArray: return rmax.fillna(0) -def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray: +def _smooth_MB06( + depth: DataArray, + rmax: float, + tol: float = 1.0e-8, + max_iter: int = 10_000, +) -> DataArray: """ - This is NEMO implementation of the direct iterative method - of Martinho and Batteen (2006). + Direct iterative method of Martinho and Batteen (2006) consistent + with NEMO implementation. The algorithm ensures that @@ -100,87 +104,62 @@ def _smooth_MB06(depth: DataArray, rmax: float) -> DataArray: Parameters ---------- depth: DataArray - Bottom depth (units: m). + Bottom depth. rmax: float Maximum slope parameter allowed + tol: float, default = 1.0e-8 + Tolerance for the iterative method + max_iter: int, default = 10000 + Maximum number of iterations Returns ------- DataArray Smooth version of the bottom topography with - a maximum slope parameter < rmax (units: m). - + a maximum slope parameter < rmax. """ - # set scaling factor used for smoothing + # 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 = 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 - 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 + # Initialize envelope bathymetry + zenv = depth + + for _ in range(max_iter): + + # Initialize lists of DataArrays to concatenate + all_ztmp = [] + all_zr = [] + for dim in zenv.dims: + + # Shifted arrays + zenv_m1 = zenv.shift({dim: -1}) + zenv_p1 = zenv.shift({dim: +1}) + + # Compute zr + zr = (zenv_m1 - zenv) / (zenv_m1 + zenv) + zr = zr.where((zenv > 0) & (zenv_m1 > 0), 0) + for dim_name in zenv.dims: + zr[{dim_name: -1}] = 0 + all_zr += [zr] + + # Compute ztmp + zr_p1 = zr.shift({dim: +1}) + all_ztmp += [zenv.where(zr <= rmax, zenv_m1 * zrfact)] + all_ztmp += [zenv.where(zr_p1 >= -rmax, zenv_p1 * zrfact)] + + # Update envelope bathymetry + zenv = xr.concat([zenv] + all_ztmp, "dummy_dim").max("dummy_dim") + + # Check target rmax + zr = xr.concat(all_zr, "dummy_dim") + if ((np.abs(zr) - rmax) <= tol).all(): + return zenv + + raise ValueError( + "Iterative method did NOT converge." + " You might want to increase the number of iterations and/or the tolerance." + ) def generate_cartesian_grid( From 1d80a84d2795c4445cdeff94c0ec429d1ea52b2f Mon Sep 17 00:00:00 2001 From: oceandie Date: Wed, 30 Jun 2021 11:49:07 +0100 Subject: [PATCH 19/19] implementing Mattia's mypy workaround --- pydomcfg/domzgr/sco.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydomcfg/domzgr/sco.py b/pydomcfg/domzgr/sco.py index d98561a..d4990d5 100644 --- a/pydomcfg/domzgr/sco.py +++ b/pydomcfg/domzgr/sco.py @@ -229,7 +229,7 @@ def _sco_z3(self) -> Tuple[DataArray, ...]: su = -sigma ss = self._stretch_sco(-sigma) a1 = scosrf - a2 = 0.0 + a2 = DataArray((0.0)) # TODO: Why can't use float here? a3 = self._envlp if self._stretch != "sf12": a2 += self._hc