diff --git a/pydomcfg/domzgr/zco.py b/pydomcfg/domzgr/zco.py index 4c92dcd..b143662 100644 --- a/pydomcfg/domzgr/zco.py +++ b/pydomcfg/domzgr/zco.py @@ -8,7 +8,7 @@ import numpy as np from xarray import DataArray, Dataset -from ..utils import _check_parameters +from ..utils import VerticalGridTuple, _check_parameters from .zgr import Zgr @@ -202,12 +202,12 @@ def _stretch_zco(self, sigma: DataArray, use_tanh2: bool = False) -> DataArray: def _zco_z3(self) -> Tuple[DataArray, ...]: """Compute and return z3{t,w} for z-coordinates grids""" - grids = ("T", "W") + top_are_zero = VerticalGridTuple(t=False, w=True) sigmas = self._sigmas sigmas_p1 = self._compute_sigma(self._z + 1) - both_z3 = [] - for grid, sigma, sigma_p1 in zip(grids, sigmas, sigmas_p1): + z3_staggered = [] + for top_is_zero, sigma, sigma_p1 in zip(top_are_zero, sigmas, sigmas_p1): if self._is_uniform: # Uniform zco grid @@ -230,13 +230,13 @@ def _zco_z3(self) -> Tuple[DataArray, ...]: a4 = self._ppa2 * self._ppacr2 z3 += ss2 * a4 - if grid == "W": + if top_is_zero: # Force first w-level to be exactly at zero z3[{"z": 0}] = 0.0 - both_z3 += [z3] + z3_staggered.append(z3) - return tuple(both_z3) + return VerticalGridTuple(*z3_staggered) # -------------------------------------------------------------------------- @property @@ -249,9 +249,9 @@ def _analyt_e3(self) -> Tuple[DataArray, ...]: if self._is_uniform: # Uniform: Return 0d DataArrays e3 = DataArray((self._pphmax / (self._jpk - 1.0))) - return tuple([e3, e3]) + return VerticalGridTuple(t=e3, w=e3) - both_e3 = [] + e3_staggered = [] for sigma in self._sigmas: # Stretched zco grid a0 = self._ppa0 @@ -266,9 +266,9 @@ def _analyt_e3(self) -> Tuple[DataArray, ...]: tanh2 = np.tanh((kk - self._ppkth2) / self._ppacr2) e3 += a2 * tanh2 - both_e3 += [e3] + e3_staggered.append(e3) - return tuple(both_e3) + return VerticalGridTuple(*e3_staggered) def _set_add_tanh2_and_pp2( self, pp2: Tuple[Optional[float], ...] diff --git a/pydomcfg/domzgr/zgr.py b/pydomcfg/domzgr/zgr.py index 771147f..a516797 100644 --- a/pydomcfg/domzgr/zgr.py +++ b/pydomcfg/domzgr/zgr.py @@ -8,6 +8,8 @@ import xarray as xr from xarray import DataArray, Dataset +from ..utils import VerticalGridTuple + class Zgr: """ @@ -59,10 +61,9 @@ def _compute_sigma(self, kindx: DataArray) -> Tuple[DataArray, ...]: sigma-coordinate (-1 <= sigma <= 0) for T and W grids """ - k_t = kindx + 0.5 - k_w = kindx + k_staggered = VerticalGridTuple(t=kindx + 0.5, w=kindx) - return tuple(-k / (self._jpk - 1.0) for k in (k_t, k_w)) + return VerticalGridTuple(*(-k / (self._jpk - 1.0) for k in k_staggered)) # ------------------------------------------------------------------------- @staticmethod @@ -109,13 +110,15 @@ def _compute_e3(z3t: DataArray, z3w: DataArray) -> Tuple[DataArray, ...]: (central-difference) of levels' depth (z3{t,w}). """ - both_e3 = [] - for z3, k_to_fill in zip((z3t, z3w), (-1, 0)): + z3_staggered = VerticalGridTuple(t=z3t, w=z3w) + k_to_fill_staggered = VerticalGridTuple(t=-1, w=0) + e3_staggered = [] + for z3, k_to_fill in zip(z3_staggered, k_to_fill_staggered): diff = z3.diff("z") fill = 2.0 * (z3t[{"z": k_to_fill}] - z3w[{"z": k_to_fill}]) - both_e3 += [xr.concat([diff, fill], "z")] + e3_staggered.append(xr.concat([diff, fill], "z")) - return tuple(both_e3) + return VerticalGridTuple(*e3_staggered) # -------------------------------------------------------------------------- def _merge_z3_and_e3( diff --git a/pydomcfg/utils.py b/pydomcfg/utils.py index caba33a..f8a23d4 100644 --- a/pydomcfg/utils.py +++ b/pydomcfg/utils.py @@ -3,7 +3,7 @@ """ import inspect -from collections import ChainMap +from collections import ChainMap, namedtuple from functools import wraps from typing import Any, Callable, Mapping, Optional, Tuple, TypeVar, Union, cast @@ -13,6 +13,8 @@ F = TypeVar("F", bound=Callable[..., Any]) +VerticalGridTuple = namedtuple("VerticalGridTuple", ["t", "w"]) + def generate_cartesian_grid( ppe1_m,