# Tools for getting grid information from Live Ocean results
# Adapted from Parker MacCready's code
# Nancy Soontiens 2016
# updated based on Parker's 2019 version by Susan Allen and Doug Latornell
import netCDF4 as nc
import numpy as np
[docs]
def get_basic_info(fn, only_G=False, only_S=False, only_T=False):
"""
Gets grid, vertical coordinate, and time info from a ROMS NetCDF
history file with full name 'fn'
Input: the filename (with path if needed)
Output: dicts G, S, and T
Example calls:
G, S, T = zfun.get_basic_info(fn)
T = zfun.get_basic_info(fn, only_T=True)
"""
ds = nc.Dataset(fn, 'r')
def make_G(ds):
# get grid and bathymetry info
g_varlist = ['h', 'lon_rho', 'lat_rho', 'lon_u', 'lat_u', 'lon_v',
'lat_v', 'mask_rho', 'mask_u', 'mask_v', 'pm', 'pn', ]
G = dict()
for vv in g_varlist:
G[vv] = ds.variables[vv][:]
G['DX'] = 1/G['pm']
G['DY'] = 1/G['pn']
G['M'], G['L'] = np.shape(G['lon_rho']) # M = rows, L = columns
# make the masks boolean (True = water, False = land, opposite of masked arrays!)
G['mask_rho'] = G['mask_rho'] == 1
G['mask_u'] = G['mask_u'] == 1
G['mask_v'] = G['mask_v'] == 1
return G
def make_S(ds):
# get vertical sigma-coordinate info (vectors are bottom to top)
s_varlist = ['s_rho', 'hc', 'Cs_r', 'Vtransform']
S = dict()
for vv in s_varlist:
S[vv] = ds.variables[vv][:]
S['N'] = len(S['s_rho']) # number of vertical levels
return S
def make_T(ds):
# get time info
t_varlist = ['ocean_time']
T = dict()
for vv in t_varlist:
T[vv] = ds.variables[vv][:]
T_units = ds.variables['ocean_time'].units
tt = nc.num2date(T['ocean_time'][:], T_units)
T['time'] = tt
return T
# return results
if only_G:
return make_G(ds)
elif only_S:
return make_S(ds)
elif only_T:
return make_T(ds)
else:
return make_G(ds), make_S(ds), make_T(ds)
[docs]
def get_z(h, zeta, S):
"""
Used to calculate the z position of fields in a ROMS history file. Only for
rho levels
Input: arrays h (bathymetry depth) and zeta (sea surface height)
which must be the same size, and dict S created by get_basic_info()
Output: 3-D arrays of z_rho
NOTE: one foible is that if you input arrays of h and zeta that are
vectors of length VL, the output array (e.g. z_rho) will have size (N, VL)
(i.e. it will never return an array with size (N, VL, 1), even if (VL, 1)
was the input shape). This is a result of the initial and final squeeze
calls.
"""
# input error checking
if ( (not isinstance(h, np.ndarray))
or (not isinstance(zeta, (np.ndarray, np.ma.core.MaskedArray))) ):
print('WARNING from get_z(): Inputs must be numpy arrays')
if not isinstance(S, dict):
print('WARNING from get_z(): S must be a dict')
# number of vertical levels
N = S['N']
# remove singleton dimensions
h = h.squeeze()
zeta = zeta.squeeze()
# ensure that we have enough dimensions
h = np.atleast_2d(h)
zeta = np.atleast_2d(zeta)
# check that the dimensions are the same
if h.shape != zeta.shape:
print('WARNING from get_z(): h and zeta must be the same shape')
M, L = h.shape
# rho
# create some useful arrays
csr = S['Cs_r']
csrr = csr.reshape(N, 1, 1).copy()
Cs_r = np.tile(csrr, [1, M, L])
H_r = np.tile(h.reshape(1, M, L).copy(), [N, 1, 1])
Zeta_r = np.tile(zeta.reshape(1, M, L).copy(), [N, 1, 1])
if S['hc'] == 0: # if hc = 0 the transform is simpler (and faster)
z_rho = H_r*Cs_r + Zeta_r + Zeta_r*Cs_r
elif S['hc'] != 0: # need to calculate a few more useful arrays
sr = S['s_rho']
srr = sr.reshape(N, 1, 1).copy()
S_rho = np.tile(srr, [1, M, L])
Hc_r = np.tile(S['hc'], [N, M, L])
if S['Vtransform'] == 1:
zr0 = (S_rho - Cs_r) * Hc_r + Cs_r*H_r
z_rho = zr0 + Zeta_r * (1 + zr0/H_r)
elif S['Vtransform'] == 2:
zr0 = (S_rho*Hc_r + Cs_r*H_r) / (Hc_r + H_r)
z_rho = Zeta_r + (Zeta_r + H_r)*zr0
z_rho = z_rho.squeeze()
return z_rho