Source code for salishsea_tools.LiveOcean_grid

# 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