diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e6dc3a3..af628b0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -41,7 +41,7 @@ repos: hooks: - id: mypy exclude: docs - additional_dependencies: [types-pkg_resources] + additional_dependencies: [xarray, types-pkg_resources] - repo: https://github.com/PyCQA/doc8 rev: 0.9.0a1 diff --git a/docs/conf.py b/docs/conf.py index 0b9382e..e6ae9e0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -61,6 +61,9 @@ # Darglint is already forcing us to have consistent types in the docstring. autodoc_typehints = "none" +# Do not warn about missing "Methods" in class docstring +numpydoc_show_class_members = False + # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for diff --git a/docs/developers/whats-new.rst b/docs/developers/whats-new.rst index 65ec2ee..84541b4 100644 --- a/docs/developers/whats-new.rst +++ b/docs/developers/whats-new.rst @@ -8,6 +8,8 @@ Unreleased - Added :py:func:`~pydomcfg.bathymetry.sea_mount` useful to generate classic sea mount test case. (:pr:`17`) +- Added :py:class:`~pydomcfg.domzgr.zco.Zco` + to generate geopotential z-coordinates. (:pr:`15`) - Added :py:func:`~pydomcfg.utils.generate_cartesian_grid` useful to generate test datasets. (:pr:`8`) - pyDOMCFG is on GitHub! (:pr:`1`) diff --git a/docs/users/api.rst b/docs/users/api.rst index 71d6cf0..79d0b59 100644 --- a/docs/users/api.rst +++ b/docs/users/api.rst @@ -5,6 +5,17 @@ API reference This page provides an auto-generated summary of pyDOMCFG's API. +Domzgr +====== + +Zco +---- +.. autoclass:: pydomcfg.domzgr.zco.Zco + :members: + + .. automethod:: __init__ + .. automethod:: __call__ + Utils ===== .. autosummary:: diff --git a/pydomcfg/domzgr/zco.py b/pydomcfg/domzgr/zco.py index 0a06924..297128a 100644 --- a/pydomcfg/domzgr/zco.py +++ b/pydomcfg/domzgr/zco.py @@ -1,15 +1,14 @@ -#!/usr/bin/env python - """ Class to generate NEMO v4.0 standard geopotential z-coordinates """ +import warnings from typing import Optional, Tuple import numpy as np -from xarray import Dataset +from xarray import DataArray, Dataset -from pydomcfg.utils import is_nemo_none +from pydomcfg.utils import _are_nemo_none, _is_nemo_none from .zgr import Zgr @@ -17,21 +16,6 @@ class Zco(Zgr): """ Class to generate geopotential z-coordinates dataset objects. - - 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_z subroutine - *) Marti, Madec & Delecluse, 1992, JGR, 97, No8, 12,763-12,766. """ # -------------------------------------------------------------------------- @@ -39,15 +23,15 @@ def __call__( self, ppdzmin: float, pphmax: float, - ppkth: float = 0.0, + ppkth: float = 0, ppacr: int = 0, ppsur: Optional[float] = None, ppa0: Optional[float] = None, ppa1: Optional[float] = None, - ldbletanh: bool = False, - ppa2: float = 0.0, - ppkth2: float = 0.0, - ppacr2: float = 0.0, + ppa2: Optional[float] = None, + ppkth2: Optional[float] = None, + ppacr2: Optional[float] = None, + ldbletanh: Optional[bool] = None, ln_e3_dep: bool = True, ) -> Dataset: """ @@ -59,109 +43,141 @@ def __call__( Minimum thickness for the top layer (units: m) pphmax: float Depth of the last W-level (total depth of the ocean basin) (units: m) - ppacr: int - Stretching factor, nondimensional > 0. The larger ppacr, - the smaller the stretching. Values from 3 to 10 are usual. - (default = 0., no stretching, uniform grid) - ppkth: float + ppkth: float, default = 0 Model level at which approximately max stretching occurs. - Nondimensional > 0, usually of order 1/2 or 2/3 of jpk. - (default = 0., i.e. no stretching, uniform grid) - ppsur: float, optional - Coeff. controlling distibution of vertical levels - (default = None, i.e. computed from ppdzmin, pphmax, ppkth and ppacr) - ppa0: float, optional - Coeff. controlling distibution of vertical levels - (default = None, i.e. computed from ppdzmin, pphmax, ppkth and ppacr) - ppa1: float, optional - Coeff. controlling distibution of vertical levels - (default = None, i.e. computed from ppdzmin, pphmax, ppkth and ppacr) - ldbletanh: bool - Logical flag to use or not double tanh stretching function - (default = False) - ppa2: float - Double tanh stretching function parameter - ppkth2: float - Double tanh stretching function parameter - ppacr2: float - Double tanh stretching function parameter - ln_e3_dep: bool - Logical flag to comp. e3 as fin. diff. (True) or - analyt. (False) (default = True) + Nondimensional > 0, usually of order 1/2 or 2/3 of ``jpk``. + (0 means no stretching, uniform grid) + ppacr: int, default = 0 + Stretching factor, nondimensional > 0. The larger ``ppacr``, + the smaller the stretching. Values from 3 to 10 are usual. + (0 means no stretching, uniform grid) + ppsur, ppa0, ppa1: float, optional + Coefficients controlling distibution of vertical levels + (None: compute from ``ppdzmin``, ``pphmax``, ``ppkth`` and ``ppacr``) + ppa2, ppkth2, ppacr2: float, optional + Double tanh stretching function parameters. + (None: Double tanh OFF) + ldbletanh: bool, optional + Logical flag to switch ON/OFF the double tanh stretching function. + This flag is only needed for compatibility with NEMO DOMAINcfg tools. + Just set ``ppa2``, ``ppkth2``, and ``ppacr2`` to switch ON + the double tanh stretching function. + ln_e3_dep: bool, default = True + Logical flag to comp. e3 as fin. diff. (True) oranalyt. (False) + (default = True) Returns ------- Dataset Describing the 3D geometry of the model - Raises - ------- - ValueError - If ldbletanh = True and parametrs are equal to 0. + Notes + ----- + * Model levels' depths ``z3{T,W}`` are defined from analytical function. + + * Model vertical scale factors ``e3{T,W}`` (i.e., grid cell thickness) + can be computed as: + + - analytical derivative of depth function (``ln_e3_dep=False``); + for backward compatibility with v3.6. + + - discrete derivative (central-difference) of levels' depth + (``ln_e3_dep=True``). The only possibility from v4.0. + + * See [1]_ and equations E.1-4 in the NEMO4 v4 ocean engine manual [2]_. + + See Also + -------- + NEMOv4.0: ``domzgr/zgr_z`` subroutine + References + ---------- + .. [1] Marti, O., Madec, G., and Delecluse, P. (1992), Comment [on “Net + diffusivity in ocean general circulation models with nonuniform grids” + by F. L. Yin and I. Y. Fung], J. Geophys. Res., 97( C8), 12763– 12766, + http://doi.org/10.1029/92JC00306. + + .. [2] Madec Gurvan, Romain Bourdallé-Badie, Jérôme Chanut, Emanuela Clementi, + Andrew Coward, Christian Ethé, … Guillaume Samson. (2019, October 24). + NEMO ocean engine (Version v4.0). + Notes Du Pôle De Modélisation De L'institut Pierre-simon Laplace (IPSL). + Zenodo. http://doi.org/10.5281/zenodo.3878122 """ + + # Init self._ppdzmin = ppdzmin self._pphmax = pphmax self._ppkth = ppkth self._ppacr = ppacr - self._ldbletanh = ldbletanh - self._ppa2 = ppa2 - self._ppkth2 = ppkth2 - self._ppacr2 = ppacr2 - self._is_uniform = not ppkth or not ppacr self._ln_e3_dep = ln_e3_dep + self._is_uniform = not (ppkth or ppacr) - # Checking consistency of input parameters - if self._ldbletanh: - if self._ppa2 * self._ppkth2 * self._ppacr2 == 0.0: - raise ValueError( - "ppa2, ppkth2 and ppacr2 MUST be > 0. " "when ldbletanh = True" - ) - - ds = self._init_ds() + # Set double tanh flag and coefficients (dummy floats when double tanh is OFF) + pp2_in = (ppa2, ppkth2, ppacr2) + self._ldbletanh, pp2_out = self._get_ldbletanh_and_pp2(ldbletanh, pp2_in) + self._ppa2, self._ppkth2, self._ppacr2 = pp2_out - # computing coeff. if needed - if not self._is_uniform: - self._ppsur, self._ppa0, self._ppa1 = self._compute_pp(ppsur, ppa0, ppa1) + # Initialize stretching coefficients (dummy floats for uniform case) + self._ppsur, self._ppa0, self._ppa1 = self._compute_pp(ppsur, ppa0, ppa1) - # compute sigma-coordinates for z3 computation - kindx = ds["z"] - sigma, sigma_p1 = (self._compute_sigma(kk) for kk in (kindx, kindx + 1.0)) - self._sigT, self._sigW = sigma - self._sigTp1, self._sigWp1 = sigma_p1 + # Compute and store sigma + self._sigmas = self._compute_sigma(self._z) - # compute z3 depths of zco vertical levels - dsz = self._zco_z3(ds) + # Compute z3 depths of zco vertical levels + z3t, z3w = self._zco_z3 - # compute e3 scale factors - dse = self._compute_e3(dsz) if self._ln_e3_dep else self._analyt_e3(dsz) + # Compute e3 scale factors (cell thicknesses) + e3t, e3w = self._compute_e3(z3t, z3w) if self._ln_e3_dep else self._analyt_e3 - return dse + return self._merge_z3_and_e3(z3t, z3w, e3t, e3w) # -------------------------------------------------------------------------- - def _compute_pp(self, ppsur, ppa0, ppa1) -> Tuple[float, ...]: - """Compute the coefficients for zco grid if needed.""" + def _compute_pp( + self, + ppsur: Optional[float], + ppa0: Optional[float], + ppa1: Optional[float], + ) -> Tuple[float, ...]: + """ + Compute and return the coefficients for zco grid. + Only None or 999999. are replaced. + Return 0s for uniform case. + """ - aa = self._ppdzmin - self._pphmax / float(self._jpk - 1) + # Uniform grid, return dummy zeros + if self._is_uniform: + if not all(_are_nemo_none((ppsur, ppa0, ppa1))): + warnings.warn( + "Uniform grid case (no stretching):" + " ppsur, ppa0 and ppa1 are ignored when ppacr == ppkth == 0" + ) + + return (0, 0, 0) + + # Strecthing parameters + aa = self._ppdzmin - self._pphmax / (self._jpk - 1) bb = np.tanh((1 - self._ppkth) / self._ppacr) - cc = self._ppacr / float(self._jpk - 1) + cc = self._ppacr / (self._jpk - 1) dd = np.log(np.cosh((self._jpk - self._ppkth) / self._ppacr)) ee = np.log(np.cosh((1.0 - self._ppkth) / self._ppacr)) - ppa1 = aa / (bb - cc * (dd - ee)) if is_nemo_none(ppa1) else ppa1 - ppa0 = self._ppdzmin - ppa1 * bb if is_nemo_none(ppa0) else ppa0 - ppsur = -(ppa0 + ppa1 * self._ppacr * ee) if is_nemo_none(ppsur) else ppsur + # Substitute only if is None or 999999 + ppa1_out = (aa / (bb - cc * (dd - ee))) if _is_nemo_none(ppa1) else ppa1 + ppa0_out = (self._ppdzmin - ppa1_out * bb) if _is_nemo_none(ppa0) else ppa0 + ppsur_out = ( + -(ppa0_out + ppa1_out * self._ppacr * ee) if _is_nemo_none(ppsur) else ppsur + ) - return (ppsur, ppa0, ppa1) + return (ppsur_out, ppa0_out, ppa1_out) # -------------------------------------------------------------------------- - def _stretch_zco(self, sigma: float, ldbletanh: bool = False): + def _stretch_zco(self, sigma: DataArray, ldbletanh: bool = False) -> DataArray: """ Provide the generalised analytical stretching function for NEMO z-coordinates. Parameters ---------- - sigma: float + sigma: DataArray Uniform non-dimensional sigma-coordinate: MUST BE positive, i.e. 0 <= sigma <= 1 ldbletanh: bool @@ -169,106 +185,131 @@ def _stretch_zco(self, sigma: float, ldbletanh: bool = False): Returns ------- - ss: float + DataArray Stretched coordinate """ kk = sigma * (self._jpk - 1.0) + 1.0 - if not ldbletanh: - kth = self._ppkth - acr = float(self._ppacr) - else: + if ldbletanh: + # Double tanh kth = self._ppkth2 - acr = float(self._ppacr2) + acr = self._ppacr2 + else: + # Single tanh + kth = self._ppkth + acr = self._ppacr - ss = np.log(np.cosh((kk - kth) / acr)) - return ss + return np.log(np.cosh((kk - kth) / acr)) # -------------------------------------------------------------------------- - def _zco_z3(self, ds: Dataset): - """ - Compute z3T/W for z-coordinates grids + @property + def _zco_z3(self) -> Tuple[DataArray, ...]: + """Compute and return z3{t,w} for z-coordinates grids""" - Parameters - ---------- - ds: Dataset - xarray dataset with empty ``z3{T,W}`` - Returns - ------- - ds: Dataset - xarray dataset with ``z3{T,W}`` correctly computed - """ + grids = ("T", "W") + sigmas = self._sigmas + sigmas_p1 = self._compute_sigma(self._z + 1) + + both_z3 = [] + for grid, sigma, sigma_p1 in zip(grids, sigmas, sigmas_p1): + + if self._is_uniform: + # Uniform zco grid + su = -sigma + s1 = DataArray((0.0)) + a1 = a3 = 0.0 + a2 = self._pphmax + else: + # Stretched zco grid + su = -sigma_p1 + s1 = 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) + + if self._ldbletanh: + # Add double tanh term + ss2 = self._stretch_zco(-sigma, self._ldbletanh) + a4 = self._ppa2 * self._ppacr2 + z3 += ss2 * a4 + + if grid == "W": + # Force first w-level to be exactly at zero + z3[{"z": 0}] = 0.0 - sigmas = [self._sigT, self._sigW] - sigmas_p1 = [self._sigTp1, self._sigWp1] - - for varname, sig, sig_p1 in zip(["z3T", "z3W"], sigmas, sigmas_p1): - for k in range(self._jpk): - if self._is_uniform: - # uniform zco grid - su = -sig[{"z": k}] - s1 = s2 = 0.0 - a1 = a3 = a4 = 0.0 - a2 = self._pphmax - else: - # stretched zco grid - su = -sig_p1[{"z": k}] - s1 = self._stretch_zco(-sig[{"z": k}]) - a1 = self._ppsur - a2 = self._ppa0 * (self._jpk - 1.0) - a3 = self._ppa1 * self._ppacr - # double tahh - if self._ldbletanh: - s2 = self._stretch_zco(-sig[{"z": k}], self._ldbletanh) - a4 = self._ppa2 * self._ppacr2 - else: - s2 = a4 = 0.0 - ds[varname][{"z": k}] = self._compute_z3(su, s1, a1, a2, a3, s2, a4) - - # force first w-level to be exactly at zero - ds["z3W"][{"z": 0}] = 0.0 - - return ds + both_z3 += [z3] + + return tuple(both_z3) # -------------------------------------------------------------------------- - def _analyt_e3(self, ds: Dataset): + @property + def _analyt_e3(self) -> Tuple[DataArray, ...]: + """ + Backward compatibility with v3.6: + Return e3{t,w} as analytical derivative of depth function z3{t,w}. """ - Provide e3T/W as analytical derivative of depth function - for backward compatibility with v3.6. - Parameters - ---------- - ds: Dataset - xarray dataset with empty ``e3{T,W}`` and ``z3{T,W}`` - correctly computed - Returns - ------- - ds: Dataset - xarray dataset with ``e3{T,W}`` correctly computed + if self._is_uniform: + # Uniform: Return 0d DataArrays + e3 = DataArray((self._pphmax / (self._jpk - 1.0))) + return tuple([e3, e3]) + + both_e3 = [] + for sigma in self._sigmas: + # Stretched zco grid + a0 = self._ppa0 + a1 = self._ppa1 + kk = -sigma * (self._jpk - 1.0) + 1.0 + tanh1 = np.tanh((kk - self._ppkth) / self._ppacr) + e3 = a0 + a1 * tanh1 + + if self._ldbletanh: + # Add double tanh term + a2 = self._ppa2 + tanh2 = np.tanh((kk - self._ppkth2) / self._ppacr2) + e3 += a2 * tanh2 + + both_e3 += [e3] + + return tuple(both_e3) + + def _get_ldbletanh_and_pp2( + self, ldbletanh: Optional[bool], pp2: Tuple[Optional[float], ...] + ) -> Tuple[bool, Tuple[float, ...]]: + """ + If ldbletanh is None, its bool value is inferred from pp2. + Return pp2=(0, 0, 0) when double tanh is switched off. """ - sigmas = [self._sigT, self._sigW] + pp_are_none = tuple(pp is None for pp in pp2) + prefix_msg = "ppa2, ppkth2 and ppacr2" + ldbletanh_out = ldbletanh if (ldbletanh is not None) else not any(pp_are_none) + + # Warnings: Ignore double tanh coeffiecients + if ldbletanh_out and self._is_uniform: + # Uniform and double tanh + warning_msg = ( + "Uniform grid case (no stretching):" + f" {prefix_msg} are ignored when ppacr == ppkth == 0" + ) + elif ldbletanh is False and not all(pp_are_none): + # ldbletanh False and double tanh coefficients specified + warning_msg = f"{prefix_msg} are ignored when ldbletanh is False" + else: + # All good + warning_msg = "" - for varname, sig in zip(["e3T", "e3W"], sigmas): - if self._is_uniform: - # uniform zco grid - ds[varname][{"z": slice(self._jpk)}] = self._pphmax / (self._jpk - 1.0) - else: - # stretched zco grid - a0 = self._ppa0 - a1 = self._ppa1 - # now faster to use loop, to be optimised - # using xarray features in the future - for k in range(self._jpk): - kk = -sig[{"z": k}] * (self._jpk - 1.0) + 1.0 - tanh1 = np.tanh((kk - self._ppkth) / self._ppacr) - if self._ldbletanh: - a2 = self._ppa2 - tanh2 = np.tanh((kk - self._ppkth2) / self._ppacr2) - else: - a2 = tanh2 = 0.0 - - ds[varname][{"z": k}] = a0 + a1 * tanh1 + a2 * tanh2 - - return ds + if warning_msg: + # Warn and return dummy values + warnings.warn(warning_msg) + return (False, (0, 0, 0)) + + # Errors: pp have inconsistent types + if ldbletanh is True and any(pp_are_none): + raise ValueError(f"{prefix_msg} MUST be all float when ldbletanh is True") + if ldbletanh is None and (any(pp_are_none) and not all(pp_are_none)): + raise ValueError(f"{prefix_msg} MUST be all None or all float") + + return (ldbletanh_out, tuple(pp or 0 for pp in pp2)) diff --git a/pydomcfg/domzgr/zgr.py b/pydomcfg/domzgr/zgr.py index 00c404a..5910107 100644 --- a/pydomcfg/domzgr/zgr.py +++ b/pydomcfg/domzgr/zgr.py @@ -2,7 +2,7 @@ Base class to generate NEMO v4.0 vertical grids. """ -from itertools import product +from typing import Tuple, Union import xarray as xr from xarray import DataArray, Dataset @@ -25,131 +25,121 @@ def __init__(self, ds_bathy: Dataset, jpk: int): number of computational levels """ - self._bathy = ds_bathy + # Input arguments + self._bathy = ds_bathy.copy() self._jpk = jpk - # ------------------------------------------------------------------------- - def _init_ds(self): - """ - Initialise the xarray dataset with empty - ``z3{T,W}`` and ``e3{T,W}`` - - Returns - ------- - ds: Dataset - A copy of the dataset used to initialise the class with new - coordinates ``z3{T,W}`` and ``e3{T,W}`` empty dataarrays - - """ - ds = self._bathy.copy() - - var = ["z3", "e3"] - grd = ["T", "W"] - - # Initialise a dataset with empty z3 and e3 dataarrays - da = xr.full_like(ds["Bathymetry"], None).expand_dims(z=range(self._jpk)) - for v, g in product(var, grd): - ds[v + g] = da.copy() - ds = ds.set_coords(v + g) - return ds + # Z dimension + self._z = DataArray(range(jpk), dims="z") # ------------------------------------------------------------------------- - def _compute_sigma(self, kindx: DataArray): + def _compute_sigma(self, kindx: DataArray) -> Tuple[DataArray, ...]: """ Provide the analytical function for sigma-coordinate, a uniform non-dimensional vertical coordinate describing the non-dimensional position of model levels. - Consider that -1 <= sigma <= 0, with + Consider that -1 <= sigma <= 0, with: - sigma = 0 at the shallower boundary - sigma = -1 at the deeper boundary + * sigma = 0 at the shallower boundary + + * sigma = -1 at the deeper boundary Parameters ---------- kindx: DataArray + z-axis indexes Returns ------- - ps: tuple + tuple (sigmaT, sigmaW) of Datarrays of uniform non-dimensional sigma-coordinate (-1 <= sigma <= 0) for T and W grids """ - k_T = kindx + 0.5 - k_W = kindx - ps = (-k / (self._jpk - 1.0) for k in (k_T, k_W)) + k_t = kindx + 0.5 + k_w = kindx - return ps + return tuple(-k / (self._jpk - 1.0) for k in (k_t, k_w)) # ------------------------------------------------------------------------- @staticmethod def _compute_z3( su: DataArray, ss1: DataArray, - a1: float, - a2: float, - a3: float, - ss2: DataArray = 0.0, - a4: float = 0.0, - ): + a1: Union[float, DataArray], + a2: Union[float, DataArray], + a3: Union[float, DataArray], + ) -> DataArray: """ - Generalised function providing the analytical - transformation from computational space to - physical space. It takes advantage of the fact that - z-coordinates can be considered a special case of - s-coordinate. - - Note that z is downward positive. + Generalised function providing the analytical transformation from computational + space to physical space. It takes advantage of the fact that z-coordinates can + be considered a special case of s-coordinate. Parameters ---------- su: DataArray - uniform non-dimensional vertical coordinate s, - aka sigma-coordinates. 0 <= s <= 1 + uniform non-dimensional vertical coordinate s, aka sigma-coordinates. + 0 <= s <= 1 ss1: DataArray stretched non-dimensional vertical coordinate s, 0 <= s <= 1 - a1: float - parameter of the transformation - a2: float - parameter of the transformation - a3: float - parameter of the transformation - ss2: DataArray - second stretched non-dimensional vertical coordinate s, - 0 <= s <= 1 (only used for zco with ldbletanh = True) - a4: float - parameter of the transformation (only used for zco with ldbletanh = True) + a1, a2, a3: float, DataArray + parameters of the transformation Returns ------- - z: DataArray + DataArray Depths of model levels + + Notes + ----- + z is downward positive. """ - z = a1 + a2 * su + a3 * ss1 + a4 * ss2 - return z + return a1 + a2 * su + a3 * ss1 # ------------------------------------------------------------------------- - def _compute_e3(self, ds: Dataset): + @staticmethod + def _compute_e3(z3t: DataArray, z3w: DataArray) -> Tuple[DataArray, ...]: + """ + Return grid cell thickness (e3{t,w}) computed as discrete derivative + (central-difference) of levels' depth (z3{t,w}). """ - Grid cell thickness computed as discrete derivative - (central-difference) of levels' depth - Parameters - ---------- - ds: Dataset - xarray dataset with empty ``e3{T,W}`` and ``z3{T,W}`` - correctly computed - Returns - ------- - ds: Dataset - xarray dataset with ``e3{T,W}`` correctly computed + both_e3 = [] + for z3, k_to_fill in zip((z3t, z3w), (-1, 0)): + diff = z3.diff("z") + fill = 2.0 * (z3t[{"z": k_to_fill}] - z3w[{"z": k_to_fill}]) + both_e3 += [xr.concat([diff, fill], "z")] + + return tuple(both_e3) + + # -------------------------------------------------------------------------- + def _merge_z3_and_e3( + self, z3t: DataArray, z3w: DataArray, e3t: DataArray, e3w: DataArray + ) -> Dataset: """ - ds["e3T"][{"z": slice(None, -1)}] = ds["z3W"].diff("z", label="lower") - ds["e3W"][{"z": slice(1, None)}] = ds["z3T"].diff("z", label="upper") - # Bottom and surface: - for varname, k in zip(["e3T", "e3W"], [-1, 0]): - ds[varname][{"z": k}] = 2.0 * (ds["z3T"][{"z": k}] - ds["z3W"][{"z": k}]) + Merge {z,e}3{t,w} with ds_bathy, broadcasting and adding CF attributes. + {z,e}3{t,w} are coordinates of the returned dataset. + """ + + # Merge computed variables with bathymetry + ds = xr.Dataset({"z3T": z3t, "z3W": z3w, "e3T": e3t, "e3W": e3w}) + ds = ds.broadcast_like(self._bathy["Bathymetry"]) + ds = ds.set_coords(ds.variables) + ds = ds.merge(self._bathy) + + # Add CF attributes + ds["z"] = ds["z"] + ds["z"].attrs = dict(axis="Z", long_name="z-dimension index") + for var in ["z3T", "z3W"]: + ds[var].attrs = dict( + standard_name="depth", long_name="Depth", units="m", positive="down" + ) + for var in ["e3T", "e3W"]: + ds[var].attrs = dict( + standard_name="cell_thickness", long_name="Thickness", units="m" + ) + return ds diff --git a/pydomcfg/tests/bathymetry.py b/pydomcfg/tests/bathymetry.py index 10bb2b3..a135a6a 100644 --- a/pydomcfg/tests/bathymetry.py +++ b/pydomcfg/tests/bathymetry.py @@ -45,7 +45,7 @@ def flat(self, depth: float) -> Dataset: ds["Bathymetry"] = xr.full_like(ds["glamt"], depth) return _add_attributes(_add_mask(ds)) - def sea_mount(self, depth: float, stiff=1.0) -> Dataset: + def sea_mount(self, depth: float, stiff: float = 1) -> Dataset: """ Channel with seamount case. @@ -123,7 +123,7 @@ def _add_attributes(ds: Dataset) -> Dataset: Dataset """ - attrs_dict = { + attrs_dict: dict = { "Bathymetry": dict(standard_name="sea_floor_depth_below_geoid", units="m"), "mask": dict(standard_name="sea_binary_mask", units="1"), } diff --git a/pydomcfg/tests/data.py b/pydomcfg/tests/data.py new file mode 100644 index 0000000..1c595d1 --- /dev/null +++ b/pydomcfg/tests/data.py @@ -0,0 +1,45 @@ +""" +Test data +""" +import numpy as np + +# Results to replicate: +# ORCA2 zco model levels depth and vertical +# scale factors as computed by NEMO v3.6. +# ------------------------------------------------------------- +# gdept_1d gdepw_1d e3t_1d e3w_1d +ORCA2_VGRID = np.array( + [ + [4.9999, 0.0000, 10.0000, 9.9998], + [15.0003, 10.0000, 10.0008, 10.0003], + [25.0018, 20.0008, 10.0023, 10.0014], + [35.0054, 30.0032, 10.0053, 10.0036], + [45.0133, 40.0086, 10.0111, 10.0077], + [55.0295, 50.0200, 10.0225, 10.0159], + [65.0618, 60.0429, 10.0446, 10.0317], + [75.1255, 70.0883, 10.0876, 10.0625], + [85.2504, 80.1776, 10.1714, 10.1226], + [95.4943, 90.3521, 10.3344, 10.2394], + [105.9699, 100.6928, 10.6518, 10.4670], + [116.8962, 111.3567, 11.2687, 10.9095], + [128.6979, 122.6488, 12.4657, 11.7691], + [142.1952, 135.1597, 14.7807, 13.4347], + [158.9606, 150.0268, 19.2271, 16.6467], + [181.9628, 169.4160, 27.6583, 22.7828], + [216.6479, 197.3678, 43.2610, 34.2988], + [272.4767, 241.1259, 70.8772, 55.2086], + [364.3030, 312.7447, 116.1088, 90.9899], + [511.5348, 429.7234, 181.5485, 146.4270], + [732.2009, 611.8891, 261.0346, 220.3500], + [1033.2173, 872.8738, 339.3937, 301.4219], + [1405.6975, 1211.5880, 402.2568, 373.3136], + [1830.8850, 1612.9757, 444.8663, 426.0031], + [2289.7679, 2057.1314, 470.5516, 459.4697], + [2768.2423, 2527.2169, 484.9545, 478.8342], + [3257.4789, 3011.8994, 492.7049, 489.4391], + [3752.4422, 3504.4551, 496.7832, 495.0725], + [4250.4012, 4001.1590, 498.9040, 498.0165], + [4749.9133, 4500.0215, 500.0000, 499.5419], + [5250.2266, 5000.0000, 500.5646, 500.3288], + ] +) diff --git a/pydomcfg/tests/test_zco.py b/pydomcfg/tests/test_zco.py index b13f530..9d0f795 100644 --- a/pydomcfg/tests/test_zco.py +++ b/pydomcfg/tests/test_zco.py @@ -3,11 +3,13 @@ """ import numpy as np +import pytest import xarray as xr from pydomcfg.domzgr.zco import Zco from .bathymetry import Bathymetry +from .data import ORCA2_VGRID def test_zco_orca2(): @@ -19,79 +21,33 @@ def test_zco_orca2(): double tanh. """ - # Results to replicate: - # ORCA2 zco model levels depth and vertical - # scale factors as computed by NEMO v3.6. - # ------------------------------------------------------------- - # gdept_1d gdepw_1d e3t_1d e3w_1d - expected = np.array( - [ - [4.9999, 0.0000, 10.0000, 9.9998], - [15.0003, 10.0000, 10.0008, 10.0003], - [25.0018, 20.0008, 10.0023, 10.0014], - [35.0054, 30.0032, 10.0053, 10.0036], - [45.0133, 40.0086, 10.0111, 10.0077], - [55.0295, 50.0200, 10.0225, 10.0159], - [65.0618, 60.0429, 10.0446, 10.0317], - [75.1255, 70.0883, 10.0876, 10.0625], - [85.2504, 80.1776, 10.1714, 10.1226], - [95.4943, 90.3521, 10.3344, 10.2394], - [105.9699, 100.6928, 10.6518, 10.4670], - [116.8962, 111.3567, 11.2687, 10.9095], - [128.6979, 122.6488, 12.4657, 11.7691], - [142.1952, 135.1597, 14.7807, 13.4347], - [158.9606, 150.0268, 19.2271, 16.6467], - [181.9628, 169.4160, 27.6583, 22.7828], - [216.6479, 197.3678, 43.2610, 34.2988], - [272.4767, 241.1259, 70.8772, 55.2086], - [364.3030, 312.7447, 116.1088, 90.9899], - [511.5348, 429.7234, 181.5485, 146.4270], - [732.2009, 611.8891, 261.0346, 220.3500], - [1033.2173, 872.8738, 339.3937, 301.4219], - [1405.6975, 1211.5880, 402.2568, 373.3136], - [1830.8850, 1612.9757, 444.8663, 426.0031], - [2289.7679, 2057.1314, 470.5516, 459.4697], - [2768.2423, 2527.2169, 484.9545, 478.8342], - [3257.4789, 3011.8994, 492.7049, 489.4391], - [3752.4422, 3504.4551, 496.7832, 495.0725], - [4250.4012, 4001.1590, 498.9040, 498.0165], - [4749.9133, 4500.0215, 500.0000, 499.5419], - [5250.2266, 5000.0000, 500.5646, 500.3288], - ] - ) - - # Testing pydomcfg.domzgr.zco - # ----------------------------------- - # ORCA2 input parameters - ppdzmin = 10.0 - pphmax = 5000.0 - ppkth = 21.43336197938 - ppacr = 3 - ppsur = -4762.96143546300 - ppa0 = 255.58049070440 - ppa1 = 245.58132232490 - ldbletanh = False - - jpk = 31 - # Bathymetry dataset - ds_bathy = Bathymetry(1000.0, 1200.0, 1, 1).flat(5000.0) + ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3) # zco grid generator - zco = Zco(ds_bathy, jpk) + zco = Zco(ds_bathy, jpk=31) - # zco mesh with analytical e3 + # zco mesh with analytical e3 using ORCA2 input parameters + # See pag 62 of v3.6 manual for the input parameters dsz_an = zco( - ppdzmin, pphmax, ppkth, ppacr, ppsur, ppa0, ppa1, ldbletanh, ln_e3_dep=False + 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, ) - varname = ["z3T", "z3W", "e3T", "e3W"] # reference ocean.output values are # given with 4 digits precision eps = 1.0e-5 - for n, var in enumerate(varname): - actual = dsz_an[var].squeeze().values - np.testing.assert_allclose(expected[:, n], actual, rtol=eps, atol=0) + 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) def test_zco_uniform(): @@ -102,38 +58,93 @@ def test_zco_uniform(): """ # Input parameters - ppdzmin = 10.0 - pphmax = 5000.0 - ppkth = 0.0 - ppacr = 0.0 - - jpk = 51 + kwargs = dict( + ppdzmin=10, + pphmax=5.0e3, + ppkth=0, + ppacr=0, + ) # Bathymetry dataset - ds_bathy = Bathymetry(1000.0, 1200.0, 1, 1).flat(5000.0) + ds_bathy = Bathymetry(10.0e3, 1.2e3, 1, 1).flat(5.0e3) # zco grid generator - zco = Zco(ds_bathy, jpk) + zco = Zco(ds_bathy, jpk=51) - # zco mesh with analytical e3 - dsz_an = zco(ppdzmin, pphmax, ppkth, ppacr, ln_e3_dep=False) - - # zco mesh with finite difference e3 - dsz_fd = zco(ppdzmin, pphmax, ppkth, ppacr, ln_e3_dep=True) - - # truncation errors - eps = 1.0e-14 - xr.testing.assert_allclose(dsz_fd, dsz_an, rtol=eps, atol=0) + # Compare zco mesh with analytical VS finite difference e3 + expected = zco(**kwargs, ln_e3_dep=True) + actual = zco(**kwargs, ln_e3_dep=False) + eps = 1.0e-14 # truncation errors + xr.testing.assert_allclose(expected, actual, rtol=eps, atol=0) def test_zco_x_y_invariant(): """Make sure all vertical columns are identical""" # Generate 2x2 flat bathymetry dataset - ds_bathy = Bathymetry(1000.0, 1200.0, 2, 2).flat(5000.0) - zco = Zco(ds_bathy, 10) - ds = zco(10.0, 5.0e3) + ds_bathy = Bathymetry(10.0e3, 1.2e3, 2, 2).flat(5.0e3) + zco = Zco(ds_bathy, jpk=10) + ds = zco(ppdzmin=10, pphmax=5.0e3) # Check z3 and e3 - for var in ["z3T", "z3W", "e3T", "e3W"]: - assert (ds[var] == ds[var].mean(["x", "y"])).all() + for varname in ["z3T", "z3W", "e3T", "e3W"]: + expected = ds[varname].isel(x=0, y=0) + actual = ds[varname] + assert (expected == actual).all() + + +def test_zco_errors(): + """Make sure we raise informative errors""" + + # Input parameters + kwargs = dict(ppdzmin=10, pphmax=5.0e3, ppkth=1, ppacr=1) + + # Generate test data + ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3) + zco = Zco(ds_bathy, jpk=10) + + # Without ldbletanh flag, only allow all pps set or none of them + with pytest.raises( + ValueError, match="ppa2, ppkth2 and ppacr2 MUST be all None or all float" + ): + zco(**kwargs, ppa2=1, ppkth2=1, ppacr2=None) + + # When ldbletanh flag is True, all coefficients must be specified + with pytest.raises(ValueError, match="ppa2, ppkth2 and ppacr2 MUST be all float"): + zco(**kwargs, ldbletanh=True, ppa2=1, ppkth2=1, ppacr2=None) + + +def test_zco_warnings(): + """Make sure we warn when arguments are ignored""" + + # Initialize test class + ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3) + zco = Zco(ds_bathy, jpk=10) + + # Uniform: Ignore stretching + kwargs = dict(ppdzmin=10, pphmax=5.0e3, ppkth=0, ppacr=0) + expected = zco(**kwargs, ppsur=None, ppa0=999_999, ppa1=None) + with pytest.warns( + UserWarning, match="ppsur, ppa0 and ppa1 are ignored when ppacr == ppkth == 0" + ): + actual = zco(**kwargs, ppsur=2, ppa0=2, ppa1=2) + xr.testing.assert_identical(expected, actual) + + # ldbletanh OFF: Ignore double tanh + kwargs = dict(ppdzmin=10, pphmax=5.0e3, ldbletanh=False) + expected = zco(**kwargs, ppa2=None, ppkth2=None, ppacr2=None) + with pytest.warns( + UserWarning, match="ppa2, ppkth2 and ppacr2 are ignored when ldbletanh is False" + ): + actual = zco(**kwargs, ppa2=2, ppkth2=2, ppacr2=2) + xr.testing.assert_identical(expected, actual) + + # Uniform case: Ignore double tanh + kwargs = dict(ppdzmin=10, pphmax=5.0e3, ppkth=0, ppacr=0) + expected = zco(**kwargs, ppa2=None, ppkth2=None, ppacr2=None) + with pytest.warns( + UserWarning, + match="ppa2, ppkth2 and ppacr2 are ignored when ppacr == ppkth == 0", + ): + actual = zco(**kwargs, ppa2=2, ppkth2=2, ppacr2=2) + xr.testing.assert_identical(expected, actual) diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index e272acc..3b2dd53 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -2,28 +2,24 @@ Utilities """ -from typing import Dict, Optional +from typing import Hashable, Iterable, Iterator, Optional import numpy as np import xarray as xr from xarray import DataArray, Dataset +NEMO_NONE = 999_999 -def is_nemo_none(var: Optional[float] = None) -> bool: - """ - Assess if a namelist parameter is None - Parameters - ---------- - var: float, optional - Variable to be tested +def _is_nemo_none(var: Hashable) -> bool: + """Assess if a NEMO parameter is None""" + return (var or NEMO_NONE) == NEMO_NONE - Returns - ------- - bool - True if the namelist parameter has to be considered as None - """ - return var in [None, 999999.0] + +def _are_nemo_none(var: Iterable) -> Iterator[bool]: + """Iterate over namelist parameters and assess if they are None""" + for v in var: + yield _is_nemo_none(v) def generate_cartesian_grid( @@ -33,7 +29,7 @@ def generate_cartesian_grid( jpjglo: Optional[int] = None, ppglam0: float = 0, ppgphi0: float = 0, - chunks: Optional[Dict[str, int]] = None, + chunks: Optional[dict] = None, ) -> Dataset: """ Generate coordinates and spacing of a NEMO Cartesian grid. @@ -80,7 +76,7 @@ def generate_cartesian_grid( ) # c: center f:face - delta_c = DataArray(ppe if ppe.shape else np.full(jp, ppe), dims=dim) + delta_c = DataArray(ppe if ppe.shape else ppe.repeat(jp), dims=dim) coord_f = delta_c.cumsum(dim) + (ppg - 0.5 * delta_c[0]) coord_c = coord_f.rolling({dim: 2}).mean().fillna(ppg) delta_f = coord_c.diff(dim).pad({dim: (0, 1)}, constant_values=delta_c[-1]) diff --git a/setup.cfg b/setup.cfg index 4f86405..bcf4d12 100644 --- a/setup.cfg +++ b/setup.cfg @@ -38,11 +38,15 @@ ignore = max-line-length = 88 [darglint] -ignore_regex=^_(.*) +# Ignore private and tests, retain __init__ and __call__ +ignore_regex = ^(_|test_)(.*)(?