test 999999 == None and final cleanup

Author: malmans2

pydomcfg/domzgr/zco.py

@@ -124,12 +124,12 @@ def __call__(
         self._sigmas = self._compute_sigma(self._z)
 
         # Compute z3 depths of zco vertical levels
-        zco_z3 = self._zco_z3
+        z3t, z3w = self._zco_z3
 
         # Compute e3 scale factors (cell thicknesses)
-        zco_e3 = self._compute_e3(*zco_z3) if self._ln_e3_dep else self._analyt_e3
+        e3t, e3w = self._compute_e3(z3t, z3w) if self._ln_e3_dep else self._analyt_e3
 
-        return self._merge_z3_and_e3(*zco_z3, *zco_e3)
+        return self._merge_z3_and_e3(z3t, z3w, e3t, e3w)
 
     # --------------------------------------------------------------------------
     def _compute_pp(
@@ -211,7 +211,7 @@ def _zco_z3(self) -> Tuple[DataArray, ...]:
         sigmas = self._sigmas
         sigmas_p1 = self._compute_sigma(self._z + 1)
 
-        zco_z3 = []
+        both_z3 = []
         for grid, sigma, sigma_p1 in zip(grids, sigmas, sigmas_p1):
 
             if self._is_uniform:
@@ -227,25 +227,21 @@ def _zco_z3(self) -> Tuple[DataArray, ...]:
                 a1 = self._ppsur
                 a2 = self._ppa0 * (self._jpk - 1.0)
                 a3 = self._ppa1 * self._ppacr
-
-            # Compute z3
             z3 = self._compute_z3(su, s1, a1, a2, a3)
 
-            # Double tanh
             if self._ldbletanh:
                 # Add double tanh term
                 ss2 = self._stretch_zco(-sigma, self._ldbletanh)
                 a4 = self._ppa2 * self._ppacr2
                 z3 += ss2 * a4
 
-            # Force first w-level to be exactly at zero
             if grid == "W":
+                # Force first w-level to be exactly at zero
                 z3[{"z": 0}] = 0.0
 
-            # Append
-            zco_z3 += [z3]
+            both_z3 += [z3]
 
-        return tuple(zco_z3)
+        return tuple(both_z3)
 
     # --------------------------------------------------------------------------
     @property
@@ -255,33 +251,29 @@ def _analyt_e3(self) -> Tuple[DataArray, ...]:
         Return e3{t,w} as analytical derivative of depth function z3{t,w}.
         """
 
-        # Return here if is uniform
         if self._is_uniform:
-            # 0d DataArray
+            # Uniform: Return 0d DataArrays
             e3 = DataArray((self._pphmax / (self._jpk - 1.0)))
             return tuple([e3, e3])
 
-        # Stretched zco grid
-        zco_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)
-
-            # Compute e3
             e3 = a0 + a1 * tanh1
 
-            # Add double tanh term
             if self._ldbletanh:
+                # Add double tanh term
                 a2 = self._ppa2
                 tanh2 = np.tanh((kk - self._ppkth2) / self._ppacr2)
                 e3 += a2 * tanh2
 
-            # Append
-            zco_e3 += [e3]
+            both_e3 += [e3]
 
-        return tuple(zco_e3)
+        return tuple(both_e3)
 
     def _get_ldbletanh_and_pp2(
         self, ldbletanh: Optional[bool], pp2: Tuple[Optional[float], ...]
@@ -306,16 +298,18 @@ def _get_ldbletanh_and_pp2(
             # ldbletanh False and double tanh coefficients specified
             warning_msg = f"{prefix_msg} are ignored when ldbletanh is False"
         else:
+            # All good
             warning_msg = ""
 
         if warning_msg:
+            # Warn and return dummy values
             warnings.warn(warning_msg)
             return (False, (0, 0, 0))
 
-        # Errors: Wrong types
+        # 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 float")
+            raise ValueError(f"{prefix_msg} MUST be all None or all float")
 
         return (ldbletanh_out, tuple(pp or 0 for pp in pp2))

pydomcfg/domzgr/zgr.py

@@ -107,13 +107,13 @@ def _compute_e3(z3t: DataArray, z3w: DataArray) -> Tuple[DataArray, ...]:
         (central-difference) of levels' depth (z3{t,w}).
         """
 
-        e3s = []
-        for z3, k in zip((z3t, z3w), (-1, 0)):
+        both_e3 = []
+        for z3, k_to_fill in zip((z3t, z3w), (-1, 0)):
             diff = z3.diff("z")
-            fill = 2.0 * (z3t[{"z": k}] - z3w[{"z": k}])
-            e3s += [xr.concat([diff, fill], "z")]
+            fill = 2.0 * (z3t[{"z": k_to_fill}] - z3w[{"z": k_to_fill}])
+            both_e3 += [xr.concat([diff, fill], "z")]
 
-        return tuple(e3s)
+        return tuple(both_e3)
 
     # --------------------------------------------------------------------------
     def _merge_z3_and_e3(

pydomcfg/tests/test_zco.py

@@ -105,7 +105,7 @@ def test_zco_errors():
 
     # 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 float"
+        ValueError, match="ppa2, ppkth2 and ppacr2 MUST be all None or all float"
     ):
         zco(**kwargs, ppa2=1, ppkth2=1, ppacr2=None)
 
@@ -121,16 +121,16 @@ def test_zco_warnings():
     ds_bathy = Bathymetry(1.0e3, 1.2e3, 1, 1).flat(5.0e3)
     zco = Zco(ds_bathy, jpk=10)
 
-    # Uniform: Ignore coefficients controlling stretching
+    # Uniform: Ignore stretching
     kwargs = dict(ppdzmin=10, pphmax=5.0e3, ppkth=0, ppacr=0)
-    expected = zco(**kwargs, ppsur=None, ppa0=None, ppa1=None)
+    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 = False : Ignore parameters controlling double tanh
+    # 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(