Source code for opendrift.readers.roppy.depth

# -*- coding: utf-8 -*-

"""Vertical structure functions for ROMS

:func:`sdepth`
  Depth of s-levels
:func:`zslice`
  Slice a 3D field in s-coordinates to fixed depth
:func:`multi_zslice`
  Slice a 3D field to several depth levels
:func:`z_average`
  Vertical average of a 3D field
:func:`s_stretch`
  Compute vertical stretching arrays Cs_r or Cs_w

"""

# -----------------------------------
# Bjørn Ådlandsvik <bjorn@imr.no>
# Institute of Marine Research
# Bergen, Norway
# 2010-09-30
# -----------------------------------

import numpy as np

[docs] 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 *cs_r* : 1D array s-level stretching curve *stagger* : [ 'rho' | 'w' ] *Vtransform* : [ 1 | 2 ] defines the transform used, defaults 1 = Song-Haidvogel Returns an array with ndim = H.ndim + 1 and shape = cs_r.shape + H.shape with the depths of the mid-points in the s-levels. Typical usage:: .. code:: 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, zeta, Hc, C) """ H = np.asarray(H) zeta = np.asarray(zeta) Hshape = H.shape # Save the shape of H Hflat = H.ravel() # and make H 1D for easy shape manipulation C = np.asarray(C) N = len(C) outshape = (N,) + Hshape # Shape of output if stagger == 'rho': S = -1.0 + (0.5+np.arange(N))/N # Unstretched coordinates elif stagger == 'w': S = np.linspace(-1.0, 0.0, N) 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) 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, Hflat) D = (Hflat + Hc) z_rho = zeta + (zeta + H) * (N/D).reshape(outshape) return z_rho else: raise ValueError("Unknown Vtransform")
# ------------------------------------
[docs] def sdepth_w(H, Hc, cs_w): """Return depth of w-points in s-levels Kept for backwards compatibility use *sdepth(H, Hc, cs_w, stagger='w')* instead """ return sdepth(H, Hc, cs_w, stagger='w')
# ------------------------------------------ # Vertical slicing e.t.c. # ------------------------------------------
[docs] def zslice(F, S, z): """Vertical slice of a 3D ROMS field Vertical interpolation of a field in s-coordinates to (possibly varying) depth level *F* : array with vertical profiles, first dimension is vertical *S* : array with depths of the F-values, *z* : Depth level(s) for output, scalar or ``shape = F.shape[1:]`` The z values should be negative Return value : array, `shape = F.shape[1:]`, the vertical slice Example: H is an array of depths (positive values) Hc is the critical depth C is 1D containing the s-coordinate stretching at rho-points returns F50, interpolated values at 50 meter with F50.shape = H.shape .. code:: z_rho = sdepth(H, Hc, C) F50 = zslice(F, z_rho, -50.0) """ # TODO: # Option to Save A, D, Dm # => faster interpolate more fields to same depth F = np.asarray(F) S = np.asarray(S) z = np.asarray(z, dtype='float') Fshape = F.shape # Save original shape if S.shape != Fshape: raise ValueError("F and z_r must have same shape") if z.shape and z.shape != Fshape[1:]: raise ValueError("z must be scalar or have shape = F.shape[1:]") # Flatten all non-vertical dimensions N = F.shape[0] # Length of vertical dimension M = F.size // N # Combined length of horizontal dimension(s) F = F.reshape((N, M)) S = S.reshape((N, M)) if z.shape: z = z.reshape((M,)) # Find integer array C with shape (M,) # with S[C[i]-1, i] < z <= S[C[i], i] # C = np.apply_along_axis(np.searchsorted, 0, S, z) # but the following is much faster C = np.sum(S < z, axis=0) C = C.clip(1, N-1) # For vectorisation # construct index array tuples D and Dm such that # F[D][i] = F[C[i], i] # F[Dm][i] = F[C[i]-1, i] I = np.arange(M, dtype='int') D = (C, I) Dm = (C-1, I) # Compute interpolation weights A = (z - S[Dm]) / (S[D]-S[Dm]) A = A.clip(0.0, 1.0) # Control the extrapolation # Do the linear interpolation R = (1-A)*F[Dm]+A*F[D] # Give the result the correct s R = R.reshape(Fshape[1:]) return R
# -----------------------------------------------
[docs] def multi_zslice(F, S, Z): """Slice a 3D ROMS field to fixed depth Vertical interpolation of a field in s-coordinates to fixed vertical level *F* : array of with vertical profiles, first dimension is vertical *S* : array with depth of s-levels (at rho-points) 1D (constant depth) or S.shape = F.shape *Z* : single depth value, negative Returns : array, ``shape = F.shape[1:]`` the vertical slice """ # TODO: # Option to Save A, D, Dm # => faster interpolate more fields to same depth F = np.asarray(F) S = np.asarray(S) Fshape = F.shape # Save original shape # Flat all dimensions after first N = F.shape[0] M = F.size // N F = F.reshape((N, M)) S = S.reshape((N, M)) # Make z.shape = (M,) Z = np.asarray(Z, dtype='float') # Valid possibilities # 1) Z = single scalar (shape = ()), one constant value # 2) Z = 1D array, shape=(kmax), a set of constant depths # 3) Z = 2D or more, reshapeable to (kmax, M) if Z.ndim == 0: Z = Z + np.zeros((1, M)) kmax = 1 elif Z.ndim == 1: kmax = Z.size Z = Z[:, np.newaxis] + np.zeros((kmax, M)) else: kmax = Z.size // M Z = Z.reshape((kmax, M)) # Find C, C.shape = (kmax, M) such that # z_r[C[k,i]-1, i] < Z[k] <= z_r[C[k,i], i] # shape: kmax, N, M => kmax, M C = np.sum(S[np.newaxis, :, :] < Z[:, np.newaxis, :], axis=1) C = C.clip(1, N-1) # Horizontal index I = np.arange(M, dtype=int) # Compute interpolation weights A = (Z - S[(C-1, I)])/(S[(C, I)]-S[(C-1, I)]) A = A.clip(0.0, 1.0) # Control the extrapolation # Do the interpolation R = (1-A)*F[(C-1, I)]+A*F[(C, I)] # Give the result the correct shape R = R.reshape((kmax,) + Fshape[1:]) return R, (A, C, I, kmax)
# ------------------------------------------------------
[docs] def z_average(F, z_r, z0, z1): """Slice a 3D ROMS field to fixed depth Vertical interpolation of a field in s-coordinates to fixed vertical level *F* : array Vertical profiles, first dimension is vertical *z_r* : array Depth of s-levels (at rho-points), requires `z_r.shape = F.shape` *z0*, *z1* : floats Single depth values with z0 <= z1 <= 0 return value : array `shape = F.shape[1:]`, the vertical average """ F = np.asarray(F) z_r = np.asarray(z_r) Fshape = F.shape # Save original shape # Flatten all dimensions after first N = F.shape[0] M = F.size // N F = F.reshape((N, M)) z_r = z_r.reshape((N, M)) # z0, z1 are scalars or horizontal arrays z0 = np.asarray(z0) if z0.shape: # Array, must be 2D z0 = z0.reshape((M,)) z1 = np.asarray(z1) if z1.shape: z1 = z1.reshape((M,)) # Bracket z0, i.e. # Find integer array C0 with shape (M,) # with z_r[C0[i]-1, i] < z0 <= z_r[C0[i], i] # Can be done with: # C0 = np.apply_along_axis(np.searchsorted, 0, z_r, z0) # but the following is much faster C0 = np.sum(z_r < z0, axis=0) C0 = C0.clip(1, N-1) # Clip to avoid illegal indices # Bracket z1 C1 = np.sum(z_r < z1, axis=0) C1 = C1.clip(1, N-1) # Use advanced indexing for vectorisation # F[(C0,I)][i] = F[C0[i], i] I = np.arange(M, dtype='int') # Interpolate F to the two levels A0 = (z0 - z_r[(C0-1, I)]) / (z_r[(C0, I)]-z_r[(C0-1, I)]) A0 = A0.clip(0.0, 1.0) # Control the extrapolation F0 = (1-A0)*F[(C0-1, I)]+A0*F[(C0, I)] A1 = (z1 - z_r[(C1-1, I)])/(z_r[(C1, I)]-z_r[(C1-1, I)]) A1 = A1.clip(0.0, 1.0) F1 = (1-A1)*F[(C1-1, I)] + A1*F[(C1, I)] # Find indices again (unclipped) C0 = np.sum(z_r < z0, axis=0) C1 = np.sum(z_r < z1, axis=0) R = np.zeros(M, dtype=np.float64) X = np.zeros(N+2, dtype=np.float64) Y = np.zeros(N+2, dtype=np.float64) z0 = z0 + R # Make sure they are spatial arrays z1 = z1 + R # For indexing below for i in I: X[:] = 0.0 Y[:] = 0.0 nz = C1[i] - C0[i] # Number of rho-points between z0 and z1 # Set up arrays for trapezoidal integration X[0] = z0[i] X[1:nz+1] = z_r[C0[i]:C1[i], i] X[nz+1] = z1[i] Y[0] = F0[i] Y[1:nz+1] = F[C0[i]:C1[i], i] Y[nz+1] = F1[i] # Perform the integration R[i] = 0.5 * np.dot(X[1:nz+2]-X[0:nz+1], Y[1:nz+2]+Y[0:nz+1]) # Compute average and revert to correct shape R = R / (z1-z0) R = R.reshape(Fshape[1:]) return R
# ----------------------------------
[docs] def s_stretch(N, theta_s, theta_b, stagger='rho', Vstretching=1): """Compute a s-level stretching array *N* : Number of vertical levels *theta_s* : Surface stretching factor *theta_b* : Bottom stretching factor *stagger* : "rho"|"w" *Vstretching* : 1|2|4 """ if stagger == 'rho': S = -1.0 + (0.5+np.arange(N))/N elif stagger == "w": S = np.linspace(-1.0, 0.0, N+1) else: raise ValueError("stagger must be 'rho' or 'w'") if Vstretching == 1: cff1 = 1.0 / np.sinh(theta_s) cff2 = 0.5 / np.tanh(0.5*theta_s) return ((1.0-theta_b)*cff1*np.sinh(theta_s*S) + theta_b*(cff2*np.tanh(theta_s*(S+0.5))-0.5)) elif Vstretching == 2: a, b = 1.0, 1.0 Csur = (1 - np.cosh(theta_s * S))/(np.cosh(theta_s) - 1) Cbot = np.sinh(theta_b * (S+1)) / np.sinh(theta_b) - 1 mu = (S+1)**a * (1 + (a/b)*(1-(S+1)**b)) return mu*Csur + (1-mu)*Cbot elif Vstretching == 4: C = (1 - np.cosh(theta_s * S)) / (np.cosh(theta_s) - 1) C = (np.exp(theta_b * C) - 1) / (1 - np.exp(-theta_b)) return C else: raise ValueError("Unknown Vstretching")
# wrapper for backwards compatibility
[docs] def s_stretch_w(N, theta_s, theta_b, Vstretching=1): """Obsolete use *s_stretch* instead""" return s_stretch(N, theta_s, theta_b, stagger='w', Vstretching=Vstretching)