Implement VerticalGridTuple

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], ...]

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(

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,