```# -*- 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, Hc, C, stagger="rho", Vtransform=1):
"""Depth of s-levels

*H* : arraylike
Bottom depths [meter, positive]

*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'][:, :]
C = fid.variables['Cs_r'][:]
Hc = fid.variables['hc'].getValue()
z_rho = sdepth(H, Hc, C)

"""
H = np.asarray(H)
Hshape = H.shape      # Save the shape of H
H = H.ravel()         # and make H 1D for easy shape maniplation
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'")

if Vtransform == 1:         # Default transform by Song and Haidvogel
A = Hc * (S - C)[:, None]
B = np.outer(C, H)
return (A + B).reshape(outshape)

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)

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        # 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
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
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 = z0[i]
X[1:nz+1] = z_r[C0[i]:C1[i], i]
X[nz+1] = z1[i]
Y = 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):