Merge pull request #1238 from kthyng/add_zeta_to_roms

Add zeta to sdepth in depths.py (part of effort to add zeta to ROMS reader)

opendrift/readers/reader_ROMS_native.py

@@ -81,6 +81,7 @@ def __init__(self, filename=None, name=None, gridfile=None, standard_name_mappin
         self._mask_rho = None
         self._mask_u = None
         self._mask_v = None
+        self._zeta = None
         self.land_binary_mask = None
         self.sea_floor_depth_below_sea_level = None
         self.z_rho_tot = None
@@ -359,6 +360,18 @@ def mask_v(self):
                     self._mask_v = self._mask_v.compute()
                 logger.info("Using mask_v for mask_v")
         return self._mask_v
+
+    @property
+    def zeta(self):
+        """Sea surface height."""
+        if self._zeta is None:
+            if 'zeta' in self.Dataset.data_vars:
+                self._zeta = self.Dataset.variables['zeta']
+                logger.info("Using zeta for sea surface height")
+            else:
+                self._zeta = np.zeros(self.mask_rho.shape)
+                logger.info("No zeta found, using 0 array for sea surface height")
+        return self._zeta        
 
     def get_variables(self, requested_variables, time=None,
                       x=None, y=None, z=None):
@@ -410,6 +423,17 @@ def get_variables(self, requested_variables, time=None,
         indy = np.arange(np.max([0, indy.min()-buffer]),
                             np.min([indy.max()+buffer, self.lon.shape[0]-1]))
 
+        # define indices
+        ixy = (indy,indx)
+        itxy = (indxTime,indy,indx)
+
+        # use these same indices for all mask subsetting since if one is
+        # 3D they all should be
+        if self.mask_rho.ndim == 2:
+            imask = ixy
+        elif self.mask_rho.ndim == 3:
+            imask = itxy
+
         # Find depth levels covering all elements
         if z.min() == 0 or self.hc is None:
             indz = self.num_layers - 1  # surface layer
@@ -424,11 +448,13 @@ def get_variables(self, requested_variables, time=None,
 
             if self.z_rho_tot is None:
                 Htot = self.sea_floor_depth_below_sea_level
-                self.z_rho_tot = depth.sdepth(Htot, self.hc, self.Cs_r,
+                zeta = self.zeta[indxTime]
+                self.z_rho_tot = depth.sdepth(Htot, zeta, self.hc, self.Cs_r,
                                               Vtransform=self.Vtransform)
 
             H = self.sea_floor_depth_below_sea_level[indy, indx]
-            z_rho = depth.sdepth(H, self.hc, self.Cs_r,
+            zeta = self.zeta[itxy]
+            z_rho = depth.sdepth(H, zeta, self.hc, self.Cs_r,
                                  Vtransform=self.Vtransform)
             # Element indices must be relative to extracted subset
             indx_el = np.clip(indx_el - indx.min(), 0, z_rho.shape[2]-1)
@@ -454,12 +480,8 @@ def get_variables(self, requested_variables, time=None,
                                           -z.min()) + self.verticalbuffer)
             variables['z'] = np.array(self.zlevels[zi1:zi2])
 
-        # use these same indices for all mask subsetting since if one is
-        # 3D they all should be
-        if self.mask_rho.ndim == 2:
-            imask = (indy,indx)
-        elif self.mask_rho.ndim == 3:
-            imask = (indxTime,indy,indx)        
+        # define another set of indices
+        itzxy = (indxTime, indz, indy, indx)
 
         def get_mask(mask_name, imask):
             if mask_name in masks_for_loop:
@@ -477,13 +499,13 @@ def get_mask(mask_name, imask):
             if par == 'land_binary_mask':
                 variables[par] = self.land_binary_mask[imask]
             elif par == 'sea_floor_depth_below_sea_level':
-                variables[par] = self.sea_floor_depth_below_sea_level[indy, indx]
+                variables[par] = self.sea_floor_depth_below_sea_level[ixy]
             elif var.ndim == 2:
-                variables[par] = var[indy, indx]
+                variables[par] = var[ixy]
             elif var.ndim == 3:
-                variables[par] = var[indxTime, indy, indx]
+                variables[par] = var[itxy]
             elif var.ndim == 4:
-                variables[par] = var[indxTime, indz, indy, indx]
+                variables[par] = var[itzxy]
             else:
                 raise Exception('Wrong dimension of variable: ' +
                                 self.variable_mapping[par])
@@ -495,11 +517,11 @@ def get_mask(mask_name, imask):
 
                 # make sure that var has matching horizontal dimensions with the mask
                 # make sure coord names also match
-                if self.mask_rho.shape[-2:] == var.shape[-2:] and set(self.mask_rho.dims).issubset(var.dims):
+                if self.mask_rho.shape[-2:] == var.shape[-2:] and self.mask_rho.dims[-2:] == var.dims[-2:]:
                     mask, mask_name = get_mask("mask_rho", imask)
-                elif self.mask_u.shape[-2:] == var.shape[-2:] and set(self.mask_u.dims).issubset(var.dims):
+                elif self.mask_u.shape[-2:] == var.shape[-2:] and self.mask_u.dims[-2:] == var.dims[-2:]:
                     mask, mask_name = get_mask("mask_u", imask)
-                elif self.mask_v.shape[-2:] == var.shape[-2:] and set(self.mask_v.dims).issubset(var.dims):
+                elif self.mask_v.shape[-2:] == var.shape[-2:] and self.mask_v.dims[-2:] == var.dims[-2:]:
                     mask, mask_name = get_mask("mask_v", imask)
                 else:
                     raise Exception('No mask found for ' + par)

opendrift/readers/roppy/depth.py

@@ -24,11 +24,14 @@
 
 import numpy as np
 
-def sdepth(H, Hc, C, stagger="rho", Vtransform=1):
+def sdepth(H, zeta, Hc, C, stagger="rho", Vtransform=1):
     """Depth of s-levels
 
     *H* : arraylike
       Bottom depths [meter, positive]
+    
+    *zeta* : scalar, arraylike
+      Surface elevation [meter]
 
     *Hc* : scalar
        Critical depth
@@ -51,14 +54,16 @@ def sdepth(H, Hc, C, stagger="rho", Vtransform=1):
 
         fid = Dataset(roms_file)
         H = fid.variables['h'][:, :]
+        zeta = fid.variables['zeta'][:, :]
         C = fid.variables['Cs_r'][:]
         Hc = fid.variables['hc'].getValue()
-        z_rho = sdepth(H, Hc, C)
+        z_rho = sdepth(H, zeta, Hc, C)
 
     """
     H = np.asarray(H)
+    zeta = np.asarray(zeta)
     Hshape = H.shape      # Save the shape of H
-    H = H.ravel()         # and make H 1D for easy shape maniplation
+    Hflat = H.ravel()         # and make H 1D for easy shape manipulation
     C = np.asarray(C)
     N = len(C)
     outshape = (N,) + Hshape       # Shape of output
@@ -69,15 +74,20 @@ def sdepth(H, Hc, C, stagger="rho", Vtransform=1):
     else:
         raise ValueError("stagger must be 'rho' or 'w'")
 
+    # https://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#_ocean_s_coordinate_generic_form_1
     if Vtransform == 1:         # Default transform by Song and Haidvogel
         A = Hc * (S - C)[:, None]
         B = np.outer(C, H)
-        return (A + B).reshape(outshape)
+        Zo_rho = (A + B).reshape(outshape)
+        z_rho = Zo_rho + zeta * (1 + Zo_rho / H)
+        return z_rho
 
+    # https://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#_ocean_s_coordinate_generic_form_2
     elif Vtransform == 2:       # New transform by Shchepetkin
-        N = Hc*S[:, None] + np.outer(C, H)
-        D = (1.0 + Hc/H)
-        return (N/D).reshape(outshape)
+        N = Hc*S[:, None] + np.outer(C, Hflat)
+        D = (Hflat + Hc)
+        z_rho = zeta + (zeta + H) * (N/D).reshape(outshape)
+        return z_rho
 
     else:
         raise ValueError("Unknown Vtransform")

tests/models/test_stranding.py

@@ -57,10 +57,10 @@ def test_stranding_3d(self):
 
         # Check calculated trajectory lengths and speeds
         total_length, distances, speeds = o.get_trajectory_lengths()
-        self.assertAlmostEqual(total_length.max(), 14978.4, 1)
-        self.assertAlmostEqual(total_length.min(), 1225.2, 1)
-        self.assertAlmostEqual(speeds.max(), 0.127, 1)
-        self.assertAlmostEqual(distances.max(), 2859.0, 1)
+        self.assertAlmostEqual(total_length.max(), 14974.8, 1)
+        self.assertAlmostEqual(total_length.min(), 1222.3, 1)
+        self.assertAlmostEqual(speeds.max(), 0.132, 1)
+        self.assertAlmostEqual(distances.max(), 2859.5, 1)
 
     def test_stranding_roms(self):
         o = PelagicEggDrift(loglevel=20)

tests/readers/test_depth.py

@@ -0,0 +1,100 @@
+"""test depth.py"""
+
+from opendrift.readers.roppy import depth
+import numpy as np
+
+import unittest
+
+class TestMyFunction(unittest.TestCase):
+    def setUp(self):
+        # don't need to test time because always first select time index or
+        # is interpolated to time
+        self.NZ, self.NY, self.NX = 3, 4, 5
+        self.Htot = np.ones((self.NY, self.NX))
+        self.hc = 0.1
+        self.Cs_r = np.linspace(-1, 0, self.NZ)
+
+    def test_sdepth_zeta_zero_Vtransform1(self):
+        """test zero zeta in sdepth, Vtransform 1."""
+
+        Vtransform = 1
+
+        # zeta 0 scalar
+        zeta = 0
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+
+        expected_output = np.array([-0.98333333, -0.5, -0.01666667])        
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+        # zeta 0s array
+        zeta = np.zeros((self.NY, self.NX))
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+
+        expected_output = np.array([-0.98333333, -0.5, -0.01666667])
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+
+    def test_sdepth_zeta_zero_Vtransform2(self):
+        """test zero zeta in sdepth, Vtransform 2."""
+
+        Vtransform = 2
+
+        # zeta 0 scalar
+        zeta = 0
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+        expected_output = np.array([-0.98484848, -0.5, -0.01515152])
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+        # zeta 0s array
+        zeta = np.zeros((self.NY, self.NX))
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+        expected_output = np.array([-0.98484848, -0.5, -0.01515152])
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+    def test_sdepth_zeta_nonzero_Vtransform1(self):
+        """test nonzero zeta in sdepth, Vtransform 1."""
+
+        Vtransform = 1
+
+        # zeta scalar
+        zeta = 1
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+        expected_output = np.array([-0.96666667,  0.,  0.96666667])        
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+        # zeta 1s array
+        zeta = np.ones((self.NY, self.NX))
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+        expected_output = np.array([-0.96666667,  0.,  0.96666667])
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+
+    def test_sdepth_zeta_nonzero_Vtransform2(self):
+        """test nonzero zeta in sdepth, Vtransform 2."""
+
+        Vtransform = 2
+
+        # zeta 1 scalar
+        zeta = 1
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+        expected_output = np.array([-0.96969697,  0.,  0.96969697])
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+        # zeta 1s array
+        zeta = np.ones((self.NY, self.NX))
+        z_rho = depth.sdepth(self.Htot, zeta, self.hc, self.Cs_r, Vtransform=Vtransform)
+        expected_output = np.array([-0.96969697,  0.,  0.96969697])
+        assert z_rho.shape == (self.NZ, self.NY, self.NX)
+        np.allclose(z_rho[:,0,0], expected_output)
+
+
+
+if __name__ == '__main__':
+    unittest.main()