diff --git a/opendrift/readers/reader_ROMS_native.py b/opendrift/readers/reader_ROMS_native.py index 5313c4042..ced5a33cb 100644 --- a/opendrift/readers/reader_ROMS_native.py +++ b/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) diff --git a/opendrift/readers/roppy/depth.py b/opendrift/readers/roppy/depth.py index f58a9beee..5f220e6c0 100644 --- a/opendrift/readers/roppy/depth.py +++ b/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") diff --git a/tests/models/test_stranding.py b/tests/models/test_stranding.py index 22224fc5e..032a29000 100644 --- a/tests/models/test_stranding.py +++ b/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) diff --git a/tests/readers/test_depth.py b/tests/readers/test_depth.py new file mode 100644 index 000000000..ce6765920 --- /dev/null +++ b/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()