def get_itp_dict(sta_dict, G):
# get interpolants
itp_dict = dict()
for sta_name in sta_dict.keys():
# target position
Lon = np.array(sta_dict[sta_name][0])
Lat = np.array(sta_dict[sta_name][1])
# get interpolants for this point
Xi0 = dict()
Yi0 = dict()
Xi1 = dict()
Yi1 = dict()
Aix = dict()
Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1, :]
yy = G['lat_' + grd][:, 1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx, extrap_nan=True)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy, extrap_nan=True)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1 - xfr, xfr]).reshape((1, 1, 2))
Aiy[grd] = np.array([1 - yfr, yfr]).reshape((1, 2))
itp_dict[sta_name] = (Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
return itp_dict
def get_itp_dict(sta_dict, G):
# get interpolants
itp_dict = dict()
for sta_name in sta_dict.keys():
# target position
Lon = np.array(sta_dict[sta_name][0])
Lat = np.array(sta_dict[sta_name][1])
# get interpolants for this point
Xi0 = dict(); Yi0 = dict()
Xi1 = dict(); Yi1 = dict()
Aix = dict(); Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1,:]
yy = G['lat_' + grd][:,1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx, extrap_nan=True)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy, extrap_nan=True)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1-xfr, xfr]).reshape((1,1,2))
Aiy[grd] = np.array([1-yfr, yfr]).reshape((1,2))
itp_dict[sta_name] = (Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
return itp_dict
def get_interpolated(G, S, b, lon, lat, z, N):
# start input dict
c = dict()
# 2D fields
for vn in ['ssh', 'ubar', 'vbar']:
xr, yr = get_xyr(G, vn)
Mr, Lr = xr.shape
u = b[vn]
uu = zfun.interp_scattered_on_plaid(xr, yr, lon, lat, u)
uu = np.reshape(uu, xr.shape)
if np.isnan(uu).any():
print('Warning: nans in output array for ' + vn)
else:
c[vn] = uu
# 3D fields
for vn in ['theta', 's3d', 'u3d', 'v3d']:
xr, yr = get_xyr(G, vn)
zr = get_zr(G, S, vn)
Nr, Mr, Lr = zr.shape
v = b[vn]
# create an intermediate array, which is on the ROMS lon, lat grid
# but has the HyCOM vertical grid
# get interpolants
xi0, xi1, xf = zfun.get_interpolant(xr, lon, extrap_nan=True)
yi0, yi1, yf = zfun.get_interpolant(yr, lat, extrap_nan=True)
# bi linear interpolation
u00 = v[:, yi0, xi0]
u10 = v[:, yi1, xi0]
u01 = v[:, yi0, xi1]
u11 = v[:, yi1, xi1]
vi = (1 - yf) * ((1 - xf) * u00 + xf * u01) + yf * (
(1 - xf) * u10 + xf * u11)
vi = vi.reshape((N, Mr, Lr))
# make interpolants to go from the HyCOM vertical grid to the
# ROMS vertical grid
I0, I1, FR = zfun.get_interpolant(zr, z, extrap_nan=True)
zrs = zr.shape
vif = vi.flatten()
LL = np.tile(np.arange(Lr), Mr * Nr)
MM = np.tile(np.repeat(np.arange(Mr), Lr), Nr)
f0 = LL + MM * Lr + I0 * Mr * Lr
f1 = LL + MM * Lr + I1 * Mr * Lr
vv = (1 - FR) * vif[f0] + FR * vif[f1]
vv = vv.reshape(zrs)
if np.isnan(vv).any():
print('Warning: nans in output array for ' + vn)
else:
c[vn] = vv
return c
def get_interpolated(G, S, b, lon, lat, z, N):
# start input dict
c = dict()
# 2D fields
for vn in ['ssh', 'ubar', 'vbar']:
xr, yr = get_xyr(G, vn)
Mr, Lr = xr.shape
u = b[vn]
uu = zfun.interp_scattered_on_plaid(xr, yr, lon, lat, u)
uu = np.reshape(uu, xr.shape)
if np.isnan(uu).any():
print('Warning: nans in output array for ' + vn)
else:
c[vn] = uu
# 3D fields
for vn in ['theta', 's3d', 'u3d', 'v3d']:
xr, yr = get_xyr(G, vn)
zr = get_zr(G, S, vn)
Nr, Mr, Lr = zr.shape
v = b[vn]
# create an intermediate array, which is on the ROMS lon, lat grid
# but has the HyCOM vertical grid
# get interpolants
xi0, xi1, xf = zfun.get_interpolant(xr, lon, extrap_nan=True)
yi0, yi1, yf = zfun.get_interpolant(yr, lat, extrap_nan=True)
# bi linear interpolation
u00 = v[:,yi0,xi0]
u10 = v[:,yi1,xi0]
u01 = v[:,yi0,xi1]
u11 = v[:,yi1,xi1]
vi = (1-yf)*((1-xf)*u00 + xf*u01) + yf*((1-xf)*u10 + xf*u11)
vi = vi.reshape((N, Mr, Lr))
# make interpolants to go from the HyCOM vertical grid to the
# ROMS vertical grid
I0, I1, FR = zfun.get_interpolant(zr, z, extrap_nan=True)
zrs = zr.shape
vif = vi.flatten()
LL = np.tile(np.arange(Lr), Mr*Nr)
MM = np.tile(np.repeat(np.arange(Mr), Lr), Nr)
f0 = LL + MM*Lr + I0*Mr*Lr
f1 = LL + MM*Lr + I1*Mr*Lr
vv = (1-FR)*vif[f0] + FR*vif[f1]
vv = vv.reshape(zrs)
if np.isnan(vv).any():
print('Warning: nans in output array for ' + vn)
else:
c[vn] = vv
return c
# get coordinates
indir = indir0 + ocn + '/Data/'
coord_dict = pickle.load(open(indir + 'coord_dict.p', 'rb'))
lon = coord_dict['lon']
lat = coord_dict['lat']
# get fields
# note that when we use the "xfh" fields there should be no nans
xfh = pickle.load(open(indir + '/xfh' + mds + '.p', 'rb'))
zeta = xfh['ssh']
# get indices around certain regions, and their data
#
# (a) Mouth of Strait of Juan de Fuca
i0, i1, fr = zfun.get_interpolant(np.array([-124.7, -124.5]), lon, extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([48.4, 48.6]), lat, extrap_nan=False)
zeta_jdf = zeta[j0[0]:j1[1]+1, i0[0]:i1[1]+1]
#
# (b) Northern Strait of Georgia
i0, i1, fr = zfun.get_interpolant(np.array([-125.5, -124.5]), lon, extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([50.25, 50.45]), lat, extrap_nan=False)
zeta_sog = zeta[j0[0]:j1[1]+1, i0[0]:i1[1]+1]
#
# (c) Offshore
i0, i1, fr = zfun.get_interpolant(np.array([-127, -126]), lon, extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([46, 47]), lat, extrap_nan=False)
zeta_off = zeta[j0[0]:j1[1]+1, i0[0]:i1[1]+1]
mdt = datetime.strptime(mds,'%Y.%m.%d')
zdf.loc[mdt, 'z_jdf'] = zeta_jdf.mean()
def get_V(vn_list, ds, plon, plat, pcs, R, surface, bound='stop'):
from warnings import filterwarnings
filterwarnings('ignore') # skip some warning messages
# get interpolant arrays
i0lon_d = dict()
i1lon_d = dict()
frlon_d = dict()
i0lat_d = dict()
i1lat_d = dict()
frlat_d = dict()
for gg in ['r', 'u', 'v']:
exn = False
i0lon_d[gg], i1lon_d[gg], frlon_d[gg] = zfun.get_interpolant(
plon, R['rlon' + gg], extrap_nan=exn)
i0lat_d[gg], i1lat_d[gg], frlat_d[gg] = zfun.get_interpolant(
plat, R['rlat' + gg], extrap_nan=exn)
i0csr, i1csr, frcsr = zfun.get_interpolant(pcs, R['rcsr'], extrap_nan=exn)
i0csw, i1csw, frcsw = zfun.get_interpolant(pcs, R['rcsw'], extrap_nan=exn)
NV = len(vn_list)
NP = len(plon)
# get interpolated values
V = np.nan * np.ones((NP, NV))
vcount = 0
for vn in vn_list:
if vn in ['w']:
i0cs = i0csw
i1cs = i1csw
frcs = frcsw
else:
i0cs = i0csr
i1cs = i1csr
frcs = frcsr
# Bartos - adding 'AKs_turb' and 'dAKs_dz' to list
if vn in [
'salt', 'temp', 'zeta', 'h', 'Uwind', 'Vwind', 'w', 'AKs',
'dAKs_dz'
]:
gg = 'r'
elif vn in ['u']:
gg = 'u'
elif vn in ['v']:
gg = 'v'
i0lat = i0lat_d[gg]
i1lat = i1lat_d[gg]
frlat = frlat_d[gg]
i0lon = i0lon_d[gg]
i1lon = i1lon_d[gg]
frlon = frlon_d[gg]
# get the data field and put nan's in masked points
if vn in ['salt', 'temp', 'u', 'v', 'w'] and surface == True:
v0 = ds[vn][0, -1, :, :].squeeze()
else:
v0 = ds[vn][:].squeeze()
try:
vv = v0.data
vv[v0.mask] = np.nan
except AttributeError:
# it is not a masked array
vv = v0
# Bartos - adding 'AKs' to list
if vn in ['salt', 'temp', 'AKs', 'u', 'v', 'w'] and surface == False:
# For variables with depth axis
# Get just the values around our particle positions.
# each row in VV corresponds to a "box" around a point
VV = np.nan * np.ones((NP, 8))
VV[:, 0] = vv[i0cs, i0lat, i0lon]
VV[:, 1] = vv[i0cs, i0lat, i1lon]
VV[:, 2] = vv[i0cs, i1lat, i0lon]
VV[:, 3] = vv[i0cs, i1lat, i1lon]
VV[:, 4] = vv[i1cs, i0lat, i0lon]
VV[:, 5] = vv[i1cs, i0lat, i1lon]
VV[:, 6] = vv[i1cs, i1lat, i0lon]
VV[:, 7] = vv[i1cs, i1lat, i1lon]
# Work on edge values. If all in a box are masked
# then that row will be nan's, and also:
# Bartos - adding reflective boundary condition and 'AKs_turb' to velocity list
if bound == 'stop':
if vn in ['u', 'v', 'w', 'AKs']:
# set all velocities to zero if any in the box are masked
VV[np.isnan(VV).any(axis=1), :] = 0
elif bound == 'reflect':
# get magnitude of average velocity
newval_mag = abs(np.nanmean(VV, axis=1).reshape(NP, 1))
if vn in ['w', 'AKs']:
# set vertical velocity and AKs to zero if any in the box are masked
VV[np.isnan(VV).any(axis=1), :] = 0
elif vn in ['u']:
# mask0 for masked points on i0lon
mask0 = np.isnan(VV[:, [0, 2, 4, 6]]).any(axis=1)
# mask1 for masked points on i1lon
mask1 = np.isnan(VV[:, [1, 3, 5, 7]]).any(axis=1)
# set u-velocity away from mask
VV[mask0, :] = newval_mag[mask0]
VV[mask1, :] = -newval_mag[mask1]
elif vn in ['v']:
# mask0 for masked points on i0lat
mask0 = np.isnan(VV[:, [0, 1, 4, 5]]).any(axis=1)
# mask1 for masked points on i1lat
mask1 = np.isnan(VV[:, [2, 3, 6, 7]]).any(axis=1)
# set v-velocity away from mask
VV[mask0, :] = newval_mag[mask0]
VV[mask1, :] = -newval_mag[mask1]
elif vn in ['salt', 'temp']:
# set all tracers to their average if any in the box are masked
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones(
(1, 8))
mask = np.isnan(VV)
VV[mask] = newval[mask]
# now do the interpolation in each box
vl = ((1 - frlat) * ((1 - frlon) * VV[:, 0] + frlon * VV[:, 1]) +
frlat * ((1 - frlon) * VV[:, 2] + frlon * VV[:, 3]))
vu = ((1 - frlat) * ((1 - frlon) * VV[:, 4] + frlon * VV[:, 5]) +
frlat * ((1 - frlon) * VV[:, 6] + frlon * VV[:, 7]))
v = (1 - frcs) * vl + frcs * vu
elif vn in ['salt', 'temp', 'u', 'v', 'w'] and surface == True:
VV = np.nan * np.ones((NP, 4))
VV[:, 0] = vv[i0lat, i0lon]
VV[:, 1] = vv[i0lat, i1lon]
VV[:, 2] = vv[i1lat, i0lon]
VV[:, 3] = vv[i1lat, i1lon]
# Work on edge values. If all in a box are masked
# then that row will be nan's, and also:
# Bartos - adding reflective boundary condition and 'AKs_turb' to velocity list
if bound == 'stop':
if vn in ['u', 'v', 'w', 'AKs']:
# set all velocities to zero if any in the box are masked
VV[np.isnan(VV).any(axis=1), :] = 0
elif bound == 'reflect':
# get magnitude of average velocity
newval_mag = abs(np.nanmean(VV, axis=1).reshape(NP, 1))
if vn in ['w', 'AKs']:
# set vertical velocity and AKs to zero if any in the box are masked
VV[np.isnan(VV).any(axis=1), :] = 0
elif vn in ['u']:
# mask0 for masked points on i0lon
mask0 = np.isnan(VV[:, [0, 2]]).any(axis=1)
# mask1 for masked points on i1lon
mask1 = np.isnan(VV[:, [1, 3]]).any(axis=1)
# set u-velocity away from mask
VV[mask0, :] = newval_mag[mask0]
VV[mask1, :] = -newval_mag[mask1]
elif vn in ['v']:
# mask0 for masked points on i0lat
mask0 = np.isnan(VV[:, [0, 1]]).any(axis=1)
# mask1 for masked points on i1lat
mask1 = np.isnan(VV[:, [2, 3]]).any(axis=1)
# set v-velocity away from mask
VV[mask0, :] = newval_mag[mask0]
VV[mask1, :] = -newval_mag[mask1]
elif vn in ['salt', 'temp']:
# set all tracers to their average if any in the box are masked
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones(
(1, 4))
mask = np.isnan(VV)
VV[mask] = newval[mask]
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones((1, 4))
mask = np.isnan(VV)
VV[mask] = newval[mask]
v = ((1 - frlat) * ((1 - frlon) * VV[:, 0] + frlon * VV[:, 1]) +
frlat * ((1 - frlon) * VV[:, 2] + frlon * VV[:, 3]))
elif vn in ['zeta', 'Uwind', 'Vwind', 'h']:
# For variables without depth axis
VV = np.nan * np.ones((NP, 4))
VV[:, 0] = vv[i0lat, i0lon]
VV[:, 1] = vv[i0lat, i1lon]
VV[:, 2] = vv[i1lat, i0lon]
VV[:, 3] = vv[i1lat, i1lon]
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones((1, 4))
mask = np.isnan(VV)
VV[mask] = newval[mask]
v = ((1 - frlat) * ((1 - frlon) * VV[:, 0] + frlon * VV[:, 1]) +
frlat * ((1 - frlon) * VV[:, 2] + frlon * VV[:, 3]))
V[:, vcount] = v
vcount += 1
return V
h_list = os.listdir(dir0 + f_dir)
h_list.sort()
h_list = [x for x in h_list if x[:9] == 'ocean_his']
if tt == 0:
pass
else:
h_list = h_list[1:] # drop the zero hour for all but the first day
#
for hi in h_list:
fn = dir0 + f_dir + '/' + hi
ds = nc.Dataset(fn)
if tt == 0:
G, S, T = zrfun.get_basic_info(fn)
lon_vec = G['lon_rho'][0, :].squeeze()
lat_vec = G['lat_rho'][:, 0].squeeze()
ii0, ii1, ifr = zfun.get_interpolant(np.array([0.02, 1.5]),
lon_vec)
jj0, jj1, jfr = zfun.get_interpolant(np.array([44.8, 45.2]),
lat_vec)
i0 = ii0[0]
i1 = ii1[1]
j0 = jj0[0]
j1 = jj1[1]
h = G['h'][j0:j1, i0:i1]
dx = G['DX'][j0:j1, i0:i1]
dy = G['DY'][j0:j1, i0:i1]
da = dx * dy
ny, nx = da.shape
zeta = ds['zeta'][0, j0:j1, i0:i1].squeeze()
salt = ds['salt'][0, :, j0:j1, i0:i1].squeeze()
zr, zw = zrfun.get_z(h, zeta, S)
dista = np.tile(distr, [N, 1]) # array # pack fields to process in dicts d2 = dict() d2['zbot'] = zbot d2['zeta'] = zeta d2['lon'] = lon d2['lat'] = lat d3 = dict() d3['zr'] = zr d3['salt'] = salt # get vectors describing the (plaid) grid xx = lon[1, :] yy = lat[:, 1] xit = zfun.get_interpolant(x, xx) yit = zfun.get_interpolant(y, yy) # and prepare them to do the bilinear interpolation xita = np.array(xit) yita = np.array(yit) col0 = xita[:, 0].astype(int) col1 = xita[:, 1].astype(int) colf = xita[:, 2] colff = 1 - colf row0 = yita[:, 0].astype(int) row1 = yita[:, 1].astype(int) rowf = yita[:, 2] rowff = 1 - rowf # now actually do the interpolation
def get_V(vn_list, ds, plon, plat, pcs, R, surface):
"""
The all-purpose tool for getting properties at
particle locations using interpolation.
"""
from warnings import filterwarnings
filterwarnings('ignore') # skip some warning messages
# get interpolant arrays
i0lon_d = dict()
i1lon_d = dict()
frlon_d = dict()
i0lat_d = dict()
i1lat_d = dict()
frlat_d = dict()
for gg in ['r', 'u', 'v']:
exn = True # nan-out particles that leave domain
i0lon_d[gg], i1lon_d[gg], frlon_d[gg] = zfun.get_interpolant(
plon, R['rlon' + gg], extrap_nan=exn)
i0lat_d[gg], i1lat_d[gg], frlat_d[gg] = zfun.get_interpolant(
plat, R['rlat' + gg], extrap_nan=exn)
i0csr, i1csr, frcsr = zfun.get_interpolant(pcs, R['rcsr'], extrap_nan=exn)
i0csw, i1csw, frcsw = zfun.get_interpolant(pcs, R['rcsw'], extrap_nan=exn)
NV = len(vn_list)
NP = len(plon)
# get interpolated values
V = np.nan * np.ones((NP, NV))
vcount = 0
for vn in vn_list:
if vn in ['w', 'AKs']:
i0cs = i0csw
i1cs = i1csw
frcs = frcsw
else:
i0cs = i0csr
i1cs = i1csr
frcs = frcsr
if vn == 'u':
gg = 'u'
elif vn == 'v':
gg = 'v'
else:
gg = 'r'
i0lat = i0lat_d[gg]
i1lat = i1lat_d[gg]
frlat = frlat_d[gg]
i0lon = i0lon_d[gg]
i1lon = i1lon_d[gg]
frlon = frlon_d[gg]
# get the data field and put nan's in masked points
# (later we do more massaging)
if vn in ['salt', 'temp', 'u', 'v', 'w'] and surface == True:
v0 = ds[vn][0, -1, :, :].squeeze()
else:
v0 = ds[vn][:].squeeze()
try:
vv = v0.data
vv[v0.mask] = np.nan
except AttributeError:
# it is not a masked array
vv = v0
if vn in ['salt', 'temp', 'AKs', 'u', 'v', 'w'] and surface == False:
# Get just the values around our particle positions.
# each row in VV corresponds to a "box" around a point
#
VV = np.nan * np.ones((NP, 8))
# using "fancy indexing" in which the three indexing arrays
# are the same length and the resulting vector has the same length
# (works for numpy arrays, not for NetCDF)
VV[:, 0] = vv[i0cs, i0lat, i0lon]
VV[:, 1] = vv[i0cs, i0lat, i1lon]
VV[:, 2] = vv[i0cs, i1lat, i0lon]
VV[:, 3] = vv[i0cs, i1lat, i1lon]
VV[:, 4] = vv[i1cs, i0lat, i0lon]
VV[:, 5] = vv[i1cs, i0lat, i1lon]
VV[:, 6] = vv[i1cs, i1lat, i0lon]
VV[:, 7] = vv[i1cs, i1lat, i1lon]
# Work on edge values. If all in a box are masked
# then that row will be 8 nan's, and also:
if vn in ['u', 'v', 'w', 'AKs']:
# set all velocities and AKs to zero if any in the box are nan
mask = np.isnan(VV)
VV[mask] = 0
elif vn in ['salt', 'temp']:
# set all tracers to their average if any in the box are masked
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones(
(1, 8))
mask = np.isnan(VV)
VV[mask] = newval[mask]
# now do the interpolation in each box
vl = ((1 - frlat) * ((1 - frlon) * VV[:, 0] + frlon * VV[:, 1]) +
frlat * ((1 - frlon) * VV[:, 2] + frlon * VV[:, 3]))
vu = ((1 - frlat) * ((1 - frlon) * VV[:, 4] + frlon * VV[:, 5]) +
frlat * ((1 - frlon) * VV[:, 6] + frlon * VV[:, 7]))
v = (1 - frcs) * vl + frcs * vu
elif vn in ['salt', 'temp', 'u', 'v', 'w'] and surface == True:
#from time import time
#tt0 = time()
if vn == 'w':
v = np.zeros(NP)
else:
VV = np.nan * np.ones((NP, 4))
VV[:, 0] = vv[i0lat, i0lon]
VV[:, 1] = vv[i0lat, i1lon]
VV[:, 2] = vv[i1lat, i0lon]
VV[:, 3] = vv[i1lat, i1lon]
# Work on edge values. If all in a box are masked
# then that row will be nan's, and also:
if vn in ['u', 'v', 'w']:
# set all velocities to zero if any in the box are masked
# wait... this looks like it just sets masked velocities to
# zero, not all of them... 2019.11.08
mask = np.isnan(VV)
VV[mask] = 0
elif vn in ['salt', 'temp']:
# set all tracers to their average if any in the box are masked
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones(
(1, 4))
mask = np.isnan(VV)
VV[mask] = newval[mask]
v = ((1 - frlat) *
((1 - frlon) * VV[:, 0] + frlon * VV[:, 1]) + frlat *
((1 - frlon) * VV[:, 2] + frlon * VV[:, 3]))
#print('%s took %0.3f seconds' % (vn, time() - tt0))
elif vn in ['zeta', 'Uwind', 'Vwind', 'h']:
VV = np.nan * np.ones((NP, 4))
VV[:, 0] = vv[i0lat, i0lon]
VV[:, 1] = vv[i0lat, i1lon]
VV[:, 2] = vv[i1lat, i0lon]
VV[:, 3] = vv[i1lat, i1lon]
newval = np.nanmean(VV, axis=1).reshape(NP, 1) * np.ones((1, 4))
mask = np.isnan(VV)
VV[mask] = newval[mask]
v = ((1 - frlat) * ((1 - frlon) * VV[:, 0] + frlon * VV[:, 1]) +
frlat * ((1 - frlon) * VV[:, 2] + frlon * VV[:, 3]))
V[:, vcount] = v
vcount += 1
return V
def get_cast(gridname, tag, ex_name, date_string, station, lon_str, lat_str):
# get the dict Ldir
Ldir = Lfun.Lstart(gridname, tag)
Ldir['gtagex'] = Ldir['gtag'] + '_' + ex_name
Ldir['date_string'] = date_string
Ldir['station'] = station
Ldir['lon_str'] = lon_str
Ldir['lat_str'] = lat_str
# make sure the output directory exists
outdir0 = Ldir['LOo'] + 'cast/'
Lfun.make_dir(outdir0)
outdir = outdir0 + Ldir['gtagex'] + '/'
Lfun.make_dir(outdir)
dt = Ldir['date_string']
#%% function definitions
def get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy):
dims = ds.variables[vv].dimensions
if 'eta_rho' in dims:
grd = 'rho'
elif 'eta_u' in dims:
grd = 'u'
elif 'eta_v' in dims:
grd = 'v'
else:
print('grid error!')
xi0 = Xi0[grd]
yi0 = Yi0[grd]
xi1 = Xi1[grd]
yi1 = Yi1[grd]
aix = Aix[grd]
aiy = Aiy[grd]
xi01 = np.array([xi0, xi1]).flatten()
yi01 = np.array([yi0, yi1]).flatten()
return xi01, yi01, aix, aiy
#%% set up for the extraction
# target position
Lon = np.array(float(Ldir['lon_str']))
Lat = np.array(float(Ldir['lat_str']))
# get grid info
indir = Ldir['roms'] + 'output/' + Ldir['gtagex'] + '/f' + dt + '/'
fn = indir + 'ocean_his_0021.nc' # approx. noon local standard time
if os.path.isfile(fn):
pass
else:
print('Not found: ' + fn)
return
G = zrfun.get_basic_info(fn, only_G=True)
S = zrfun.get_basic_info(fn, only_S=True)
lon = G['lon_rho']
lat = G['lat_rho']
mask = G['mask_rho']
xvec = lon[0, :].flatten()
yvec = lat[:, 0].flatten()
i0, i1, frx = zfun.get_interpolant(np.array([float(Ldir['lon_str'])]),
xvec)
j0, j1, fry = zfun.get_interpolant(np.array([float(Ldir['lat_str'])]),
yvec)
i0 = int(i0)
j0 = int(j0)
# find indices of nearest good point
if mask[j0, i0] == 1:
print('- ' + station + ': point OK')
elif mask[j0, i0] == 0:
print('- ' + station + ':point masked')
i0, j0 = get_ij_good(lon, lat, xvec, yvec, i0, j0, mask)
new_lon = xvec[i0]
new_lat = yvec[j0]
Lon = np.array(new_lon)
Lat = np.array(new_lat)
# get interpolants for this point
Xi0 = dict()
Yi0 = dict()
Xi1 = dict()
Yi1 = dict()
Aix = dict()
Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1, :]
yy = G['lat_' + grd][:, 1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx, extrap_nan=True)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy, extrap_nan=True)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1 - xfr, xfr]).reshape((1, 1, 2))
Aiy[grd] = np.array([1 - yfr, yfr]).reshape((1, 2))
# generating some lists
v0_list = ['h', 'lon_rho', 'lat_rho', 'lon_u', 'lat_u', 'lon_v', 'lat_v']
v1_list = ['ocean_time']
v2_list = []
v3_list_rho = []
v3_list_w = []
ds = nc.Dataset(fn)
for vv in ds.variables:
vdim = ds.variables[vv].dimensions
if (('ocean_time' in vdim) and ('s_rho' not in vdim)
and ('s_w' not in vdim) and (vv != 'ocean_time')):
v2_list.append(vv)
elif (('ocean_time' in vdim) and ('s_rho' in vdim)):
v3_list_rho.append(vv)
elif (('ocean_time' in vdim) and ('s_w' in vdim)):
v3_list_w.append(vv)
V = dict()
V_long_name = dict()
V_units = dict()
v_all_list = v0_list + v1_list + v2_list + v3_list_rho + v3_list_w
for vv in v_all_list:
V[vv] = np.array([])
try:
V_long_name[vv] = ds.variables[vv].long_name
except:
V_long_name[vv] = ''
try:
V_units[vv] = ds.variables[vv].units
except:
V_units[vv] = ''
# get static variables
for vv in v0_list:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][yi01, xi01].squeeze()
V[vv] = (aiy * ((aix * vvtemp).sum(-1))).sum(-1)
ds.close()
#%% extract time-dependent fields
print('-- Working on date: ' + dt)
sys.stdout.flush()
ds = nc.Dataset(fn)
for vv in v1_list:
vtemp = ds.variables[vv][:].squeeze()
V[vv] = np.append(V[vv], vtemp)
for vv in v2_list:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][:, yi01, xi01].squeeze()
vtemp = (aiy * ((aix * vvtemp).sum(-1))).sum(-1)
V[vv] = np.append(V[vv], vtemp)
for vv in v3_list_rho:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][:, :, yi01, xi01].squeeze()
vtemp = (aiy * ((aix * vvtemp).sum(-1))).sum(-1)
V[vv] = vtemp.reshape((S['N'], 1))
for vv in v3_list_w:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][:, :, yi01, xi01].squeeze()
vtemp = (aiy * ((aix * vvtemp).sum(-1))).sum(-1)
V[vv] = vtemp.reshape((S['N'] + 1, 1))
ds.close()
# create z_rho and z_w (has to be done after we have V['zeta'])
hh = V['h'][:] * np.ones_like(V['zeta'])
z_rho, z_w = zrfun.get_z(hh, V['zeta'][:], S)
V['hh'] = hh
V_long_name['hh'] = 'bottom depth (positive down) as a vector'
V_units['hh'] = 'm'
V['z_rho'] = z_rho
V_long_name['z_rho'] = 'z on rho points (positive up)'
V_units['z_rho'] = 'm'
V['z_w'] = z_w
V_long_name['z_w'] = 'z on w points (positive up)'
V_units['z_w'] = 'm'
v2_list.append('hh')
v3_list_rho.append('z_rho')
v3_list_w.append('z_w')
#%% save the output to NetCDF
out_fn = (outdir + Ldir['station'] + '_' + Ldir['date_string'] + '.nc')
# get rid of the old version, if it exists
try:
os.remove(out_fn)
except OSError:
pass # assume error was because the file did not exist
foo = nc.Dataset(out_fn, 'w')
N = S['N']
NT = len(V['ocean_time'][:])
foo.createDimension('scalar', 1)
foo.createDimension('s_rho', N)
foo.createDimension('s_w', N + 1)
foo.createDimension('ocean_time', NT)
for vv in v0_list:
v_var = foo.createVariable(vv, float, ('scalar'))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v1_list:
v_var = foo.createVariable(vv, float, ('ocean_time', ))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v2_list:
v_var = foo.createVariable(vv, float, ('ocean_time', ))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v3_list_rho:
v_var = foo.createVariable(vv, float, ('s_rho', 'ocean_time'))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v3_list_w:
v_var = foo.createVariable(vv, float, ('s_w', 'ocean_time'))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
foo.close()
def get_extrapolated(in_fn,
L,
M,
N,
X,
Y,
lon,
lat,
z,
Ldir,
add_CTD=False,
fix_NSoG=False):
b = pickle.load(open(in_fn, 'rb'))
vn_list = list(b.keys())
# check that things are the expected shape
def check_coords(shape_tuple, arr_shape):
if arr_shape != shape_tuple:
print('WARNING: array shape mismatch')
for vn in vn_list:
if vn == 'dt':
pass
elif vn == 'ssh':
check_coords((M, L), b[vn].shape)
else:
check_coords((N, M, L), b[vn].shape)
# creat output array and add dt to it.
vn_list.remove('dt')
V = dict()
for vn in vn_list:
V[vn] = np.nan + np.ones(b[vn].shape)
V['dt'] = b['dt']
# extrapolate ssh
vn = 'ssh'
v = b[vn]
if fix_NSoG:
print(' ^^ Fixing NSoG ^^')
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# (a) Mouth of Strait of Juan de Fuca
i0, i1, fr = zfun.get_interpolant(np.array([-124.7, -124.5]),
lon,
extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([48.4, 48.6]),
lat,
extrap_nan=False)
zeta_jdf = v[j0[0]:j1[1] + 1, i0[0]:i1[1] + 1]
#
# (b) Northern Strait of Georgia
i0, i1, fr = zfun.get_interpolant(np.array([-125.5, -124.5]),
lon,
extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([50.15, 50.45]),
lat,
extrap_nan=False)
v[j0[0]:j1[1] + 1, i0[0]:i1[1] + 1] = np.nanmean(zeta_jdf) + 0.1
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
vv = extrap_nearest_to_masked(X, Y, v)
V[vn] = vv
vn_list.remove('ssh')
# extrapolate 3D fields
for vn in vn_list:
v = b[vn]
if vn == 't3d':
v0 = np.nanmin(v)
elif vn == 's3d':
v0 = np.nanmax(v)
if vn in ['t3d', 's3d']:
print(' -- extrapolating ' + vn)
if add_CTD == False:
for k in range(N):
fld = v[k, :, :]
fldf = extrap_nearest_to_masked(X, Y, fld, fld0=v0)
V[vn][k, :, :] = fldf
elif add_CTD == True:
print(vn + ' Adding CTD data before extrapolating')
Cast_dict, sta_df = Ofun_CTD.get_casts(Ldir)
for k in range(N):
fld = v[k, :, :]
zz = z[k]
xyorig, fldorig = Ofun_CTD.get_orig(
Cast_dict, sta_df, X, Y, fld, lon, lat, zz, vn)
fldf = Ofun_CTD.extrap_nearest_to_masked_CTD(
X, Y, fld, xyorig=xyorig, fldorig=fldorig, fld0=v0)
V[vn][k, :, :] = fldf
elif vn in ['u3d', 'v3d']:
print(' -- extrapolating ' + vn)
vv = v.copy()
vv = np.ma.masked_where(np.isnan(vv), vv)
vv[vv.mask] = 0
V[vn] = vv.data
# Create ubar and vbar.
# Note: this is slightly imperfect because the z levels are at the same
# position as the velocity levels.
dz = np.nan * np.ones((N, 1, 1))
dz[1:, 0, 0] = np.diff(z)
dz[0, 0, 0] = dz[1, 0, 0]
# account for the fact that the new hycom fields do not show up masked
u3d = np.ma.masked_where(np.isnan(b['u3d']), b['u3d'])
v3d = np.ma.masked_where(np.isnan(b['v3d']), b['v3d'])
dz3 = dz * np.ones_like(u3d) # make dz a masked array
b['ubar'] = np.sum(u3d * dz3, axis=0) / np.sum(dz3, axis=0)
b['vbar'] = np.sum(v3d * dz3, axis=0) / np.sum(dz3, axis=0)
for vn in ['ubar', 'vbar']:
v = b[vn]
vv = v.copy()
vv = np.ma.masked_where(np.isnan(vv), vv)
vv[vv.mask] = 0
V[vn] = vv.data
# calculate potential temperature
press_db = -z.reshape((N, 1, 1))
V['theta'] = seawater.ptmp(V['s3d'], V['t3d'], press_db)
return V
dista = np.tile(distr, [N, 1]) # array # pack fields to process in dicts d2 = dict() d2['zbot'] = zbot d2['zeta'] = zeta d2['lon'] = lon d2['lat'] = lat d3 = dict() d3['zr'] = zr d3['salt'] = salt # get vectors describing the (plaid) grid xx = lon[1,:] yy = lat[:,1] xit = zfun.get_interpolant(x, xx) yit = zfun.get_interpolant(y, yy) # and prepare them to do the bilinear interpolation xita = np.array(xit) yita = np.array(yit) col0 = xita[:, 0].astype(int) col1 = xita[:, 1].astype(int) colf = xita[:, 2] colff = 1 - colf row0 = yita[:, 0].astype(int) row1 = yita[:, 1].astype(int) rowf = yita[:, 2] rowff = 1 - rowf # now actually do the interpolation
qnet_up_b = np.cumsum(q_up)
qnet_down_b = np.cumsum(q_down)
qnet_r_b = np.cumsum(q_r)
qnet_b = qnet_up_b - qnet_down_b
dQ_a = qnet_a[0] - q_s[0]
dQ_b = qnet_b[-1] - q_s[-1]
qnet_aa = np.append(qnet_a, 0)
qnet_bb = np.append(0, qnet_b)
qnet_aa -= dQ_a
qnet_bb -= dQ_b
# now let's find a way to interpolate to get the amount to split
i0, i1, fr = zfun.get_interpolant(np.array([0]), qnet_bb)
i0 = i0[0]
i1 = i1[0]
print('i0=%d, i1=%d, fr=%0.2f' % (i0, i1, fr))
# qnet[0:i0+1] = qnet_aa[0:i0+1]
# qnet[i1:] = qnet_bb[i1:]
qnet[0:i0 + 1] = qnet_aa[0:i0 + 1] - (qnet_aa[i0] + fr *
(qnet_aa[i1] - qnet_aa[i0]))
qnet[i1:] = qnet_bb[i1:] - (qnet_bb[i0] + fr *
(qnet_bb[i1] - qnet_bb[i0]))
# Now apply the same split to the other vectors
qnet_up_aa = np.append(qnet_up_a, 0)
qnet_up_bb = np.append(0, qnet_up_b)
qnet_down_aa = np.append(qnet_down_a, 0)
def get_section(ds, vn, x, y, in_dict):
# PLOT CODE
from warnings import filterwarnings
filterwarnings('ignore') # skip a warning message
# GET DATA
G, S, T = zrfun.get_basic_info(in_dict['fn'])
h = G['h']
zeta = ds['zeta'][:].squeeze()
zr = zrfun.get_z(h, zeta, S, only_rho=True)
sectvar = ds[vn][:].squeeze()
L = G['L']
M = G['M']
N = S['N']
lon = G['lon_rho']
lat = G['lat_rho']
mask = G['mask_rho']
maskr = mask.reshape(1, M, L).copy()
mask3 = np.tile(maskr, [N, 1, 1])
zbot = -h # don't need .copy() because of the minus operation
# make sure fields are masked
zeta[mask == False] = np.nan
zbot[mask == False] = np.nan
sectvar[mask3 == False] = np.nan
# create dist
earth_rad = zfun.earth_rad(np.mean(lat[:, 0])) # m
xrad = np.pi * x / 180
yrad = np.pi * y / 180
dx = earth_rad * np.cos(yrad[1:]) * np.diff(xrad)
dy = earth_rad * np.diff(yrad)
ddist = np.sqrt(dx**2 + dy**2)
dist = np.zeros(len(x))
dist[1:] = ddist.cumsum() / 1000 # km
# find the index of zero
i0, i1, fr = zfun.get_interpolant(np.zeros(1), dist)
idist0 = i0
distr = dist.reshape(1, len(dist)).copy()
dista = np.tile(distr, [N, 1]) # array
# pack fields to process in dicts
d2 = dict()
d2['zbot'] = zbot
d2['zeta'] = zeta
d2['lon'] = lon
d2['lat'] = lat
d3 = dict()
d3['zr'] = zr
d3['sectvar'] = sectvar
# get vectors describing the (plaid) grid
xx = lon[1, :]
yy = lat[:, 1]
col0, col1, colf = zfun.get_interpolant(x, xx)
row0, row1, rowf = zfun.get_interpolant(y, yy)
# and prepare them to do the bilinear interpolation
colff = 1 - colf
rowff = 1 - rowf
# now actually do the interpolation
# 2-D fields
v2 = dict()
for fname in d2.keys():
fld = d2[fname]
fldi = (rowff * (colff * fld[row0, col0] + colf * fld[row0, col1]) +
rowf * (colff * fld[row1, col0] + colf * fld[row1, col1]))
if type(fldi) == np.ma.core.MaskedArray:
fldi = fldi.data # just the data, not the mask
v2[fname] = fldi
# 3-D fields
v3 = dict()
for fname in d3.keys():
fld = d3[fname]
fldi = (rowff *
(colff * fld[:, row0, col0] + colf * fld[:, row0, col1]) +
rowf *
(colff * fld[:, row1, col0] + colf * fld[:, row1, col1]))
if type(fldi) == np.ma.core.MaskedArray:
fldid = fldi.data # just the data, not the mask
fldid[fldi.mask == True] = np.nan
v3[fname] = fldid
v3['dist'] = dista # distance in km
# make "full" fields by padding top and bottom
nana = np.nan * np.ones((N + 2, len(dist))) # blank array
v3['zrf'] = nana.copy()
v3['zrf'][0, :] = v2['zbot']
v3['zrf'][1:-1, :] = v3['zr']
v3['zrf'][-1, :] = v2['zeta']
#
v3['sectvarf'] = nana.copy()
v3['sectvarf'][0, :] = v3['sectvar'][0, :]
v3['sectvarf'][1:-1, :] = v3['sectvar']
v3['sectvarf'][-1, :] = v3['sectvar'][-1, :]
#
v3['distf'] = nana.copy()
v3['distf'][0, :] = v3['dist'][0, :]
v3['distf'][1:-1, :] = v3['dist']
v3['distf'][-1, :] = v3['dist'][-1, :]
# attempt to skip over nan's
v3.pop('zr')
v3.pop('sectvar')
v3.pop('dist')
mask3 = ~np.isnan(v3['sectvarf'][:])
#print(mask3.shape)
mask2 = mask3[-1, :]
dist = dist[mask2]
NC = len(dist)
NR = mask3.shape[0]
for k in v2.keys():
#print('v2 key: ' + k)
v2[k] = v2[k][mask2]
for k in v3.keys():
#print('v3 key: ' + k)
v3[k] = v3[k][mask3]
v3[k] = v3[k].reshape((NR, NC))
#print(v3[k].shape)
return v2, v3, dist, idist0
def get_inds(x0, x1, y0, y1, G):
# determine the direction of the section
# and make sure indices are *increasing*
if (x0==x1) and (y0!=y1):
sdir = 'NS'
a = [y0, y1]; a.sort()
y0 = a[0]; y1 = a[1]
elif (x0!=x1) and (y0==y1):
sdir = 'EW'
a = [x0, x1]; a.sort()
x0 = a[0]; x1 = a[1]
else:
print('Input points do not form a proper section')
sdir='bad'
sys.exit()
# we assume a plaid grid, as usual
if sdir == 'NS':
lon = G['lon_u'][0,:].squeeze()
lat = G['lat_u'][:,0].squeeze()
elif sdir == 'EW':
lon = G['lon_v'][0,:].squeeze()
lat = G['lat_v'][:,0].squeeze()
# we get all 4 i's or j's but only 3 are used
i0, i1, fr = zfun.get_interpolant(np.array([x0]), lon, extrap_nan=True)
if np.isnan(fr):
print('Bad x point')
sys.exit()
else:
ii0 = int(i0)
i0, i1, fr = zfun.get_interpolant(np.array([x1]), lon, extrap_nan=True)
if np.isnan(fr):
print('Bad x point')
sys.exit()
else:
ii1 = int(i1)
j0, j1, fr = zfun.get_interpolant(np.array([y0]), lat, extrap_nan=True)
if np.isnan(fr):
print('Bad y0 point')
sys.exit()
else:
jj0 = int(j0)
j0, j1, fr = zfun.get_interpolant(np.array([y1]), lat, extrap_nan=True)
if np.isnan(fr):
print('Bad y1 point')
sys.exit()
else:
jj1 = int(j1)
# get mask and trim indices
# Note: the mask in G is True on water points
if sdir == 'NS':
mask = G['mask_u'][jj0:jj1+1, ii0]
# Note: argmax finds the index of the first True in this case
igood0 = np.argmax(mask)
igood1 = np.argmax(mask[::-1])
# check to see if section is "closed"
if (igood0==0) | (igood1==0):
print('Warning: not closed one or both ends')
# keep only to end of water points, to allow for ocean sections
Mask = mask[igood0:-igood1]
# and change the indices to match. These will be the indices
# of the start and end points.
jj0 = jj0 + igood0
jj1 = jj1 - igood1
print(' sdir=%2s: jj0=%4d, jj1=%4d, ii0=%4d' % (sdir, jj0, jj1, ii0))
Lat = lat[jj0:jj1+1]
Lon = lon[ii0] * np.ones_like(Mask)
elif sdir == 'EW':
mask = G['mask_v'][jj0, ii0:ii1+1]
igood0 = np.argmax(mask)
igood1 = np.argmax(mask[::-1])
# check to see if section is "closed"
if (igood0==0) | (igood1==0):
print('Warning: not closed one or both ends')
Mask = mask[igood0:-igood1]
ii0 = ii0 + igood0
ii1 = ii1 - igood1
print(' sdir=%2s: jj0=%4d, ii0=%4d, ii1=%4d' % (sdir, jj0, ii0, ii1))
Lon = lon[ii0:ii1+1]
Lat = lat[jj0] * np.ones_like(Mask)
return ii0, ii1, jj0, jj1, sdir, Lon, Lat, Mask
lon_vec = G['lon_rho'][0, :]
lat_vec = G['lat_rho'][:, 0]
# test values for cascadia1
# XBS = [55, 80] # x-limits
# YBS = [295, 325] # y-limits
# Bounds of subdomain (rho grid) from Susan Allen 2018_11
lon0 = -125.016452048434
lon1 = -124.494612925929
lat0 = 48.3169463809796
lat1 = 48.7515055163539
# find indices containing these bonds
i0, i1, ifr = zfun.get_interpolant(np.array([lon0, lon1]),
lon_vec,
extrap_nan=True)
j0, j1, jfr = zfun.get_interpolant(np.array([lat0, lat1]),
lat_vec,
extrap_nan=True)
if np.nan in ifr:
print('Warning: nan in ifr')
if np.nan in jfr:
print('Warning: nan in jfr')
XBS = [i0[0], i1[1]]
YBS = [j0[0], j1[1]]
print(XBS)
print(YBS)
from datetime import datetime, timedelta
start_time = datetime.now()
else:
inname = 'cascadia1_base_lo1_47N124.5W_low_pass_2013.01.02_2015.12.31.p'
import pickle
V, v1_list, v2_list, v3_list, G, S, Lon, Lat, sta_name, h = pickle.load( open( indir + inname, 'rb' ) )
#%% messing with z
z_rho, z_w = zfun.get_z(h, 0*h, S)
z_rho_ext = zfun.make_full((-h, z_rho, np.array([0.,])))
zees = np.array([-70, -35, -20, -10, -0])
zcolor = ['k', 'b', 'g', 'orange', 'r']
nz = len(zees)
int_arr = zfun.get_interpolant(zees, z_rho_ext)
v3_list_alt = ['temp', 'salt', 'v']
v3_names = ['Temperature $(^{\circ}C)$', 'Salinity',
'NS Velocity $(m s^{-1})$']
v3_lims = [(6,19), (34,26), (-.6,.6)]
v3_limdict = dict(zip(v3_list_alt, v3_lims))
v3_namedict = dict(zip(v3_list_alt, v3_names))
v2_list_alt = ['pe_mix_v', 'svstr', 'shflux']
v2_names = ['Energy to Mix $(Jm^{-3})$', 'NS wind Stress (Pa)',
'Surface Net Heat Flux $(Wm^{-2})$']
v2_lims = [(0,160), (-.15,.25), (-300,300)]
v2_limdict = dict(zip(v2_list_alt, v2_lims))
v2_namedict = dict(zip(v2_list_alt, v2_names))
def get_cast(gridname, tag, ex_name, date_string, station, lon_str, lat_str):
# get the dict Ldir
Ldir = Lfun.Lstart(gridname, tag)
Ldir['gtagex'] = Ldir['gtag'] + '_' + ex_name
Ldir['date_string'] = date_string
Ldir['station'] = station
Ldir['lon_str'] = lon_str
Ldir['lat_str'] = lat_str
# make sure the output directory exists
outdir0 = Ldir['LOo'] + 'cast/'
Lfun.make_dir(outdir0)
outdir = outdir0 + Ldir['gtagex'] + '/'
Lfun.make_dir(outdir)
dt = Ldir['date_string']
#%% function definitions
def get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy):
dims = ds.variables[vv].dimensions
if 'eta_rho' in dims:
grd = 'rho'
elif 'eta_u' in dims:
grd = 'u'
elif 'eta_v' in dims:
grd = 'v'
else:
print('grid error!')
xi0 = Xi0[grd]; yi0 = Yi0[grd]
xi1 = Xi1[grd]; yi1 = Yi1[grd]
aix = Aix[grd]; aiy = Aiy[grd]
xi01 = np.array([xi0, xi1]).flatten()
yi01 = np.array([yi0, yi1]).flatten()
return xi01, yi01, aix, aiy
#%% set up for the extraction
# target position
Lon = np.array(float(Ldir['lon_str']))
Lat = np.array(float(Ldir['lat_str']))
# get grid info
indir = Ldir['roms'] + 'output/' + Ldir['gtagex'] + '/f' + dt + '/'
fn = indir + 'ocean_his_0021.nc' # approx. noon local standard time
if os.path.isfile(fn):
pass
else:
print('Not found: ' + fn)
return
G = zrfun.get_basic_info(fn, only_G=True)
S = zrfun.get_basic_info(fn, only_S=True)
lon = G['lon_rho']
lat = G['lat_rho']
mask = G['mask_rho']
xvec = lon[0,:].flatten()
yvec = lat[:,0].flatten()
i0, i1, frx = zfun.get_interpolant(np.array([float(Ldir['lon_str'])]), xvec)
j0, j1, fry = zfun.get_interpolant(np.array([float(Ldir['lat_str'])]), yvec)
i0 = int(i0)
j0 = int(j0)
# find indices of nearest good point
if mask[j0,i0] == 1:
print('- ' + station + ': point OK')
elif mask[j0,i0] == 0:
print('- ' + station + ':point masked')
i0, j0 = get_ij_good(lon, lat, xvec, yvec, i0, j0, mask)
new_lon = xvec[i0]
new_lat = yvec[j0]
Lon = np.array(new_lon)
Lat = np.array(new_lat)
# get interpolants for this point
Xi0 = dict(); Yi0 = dict()
Xi1 = dict(); Yi1 = dict()
Aix = dict(); Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1,:]
yy = G['lat_' + grd][:,1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx, extrap_nan=True)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy, extrap_nan=True)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1-xfr, xfr]).reshape((1,1,2))
Aiy[grd] = np.array([1-yfr, yfr]).reshape((1,2))
# generating some lists
v0_list = ['h', 'lon_rho', 'lat_rho', 'lon_u', 'lat_u', 'lon_v', 'lat_v']
v1_list = ['ocean_time']
v2_list = []
v3_list_rho = []
v3_list_w = []
ds = nc.Dataset(fn)
for vv in ds.variables:
vdim = ds.variables[vv].dimensions
if ( ('ocean_time' in vdim)
and ('s_rho' not in vdim)
and ('s_w' not in vdim)
and (vv != 'ocean_time') ):
v2_list.append(vv)
elif ( ('ocean_time' in vdim) and ('s_rho' in vdim) ):
v3_list_rho.append(vv)
elif ( ('ocean_time' in vdim) and ('s_w' in vdim) ):
v3_list_w.append(vv)
V = dict()
V_long_name = dict()
V_units = dict()
v_all_list = v0_list + v1_list + v2_list + v3_list_rho + v3_list_w
for vv in v_all_list:
V[vv] = np.array([])
try:
V_long_name[vv] = ds.variables[vv].long_name
except:
V_long_name[vv] = ''
try:
V_units[vv] = ds.variables[vv].units
except:
V_units[vv] = ''
# get static variables
for vv in v0_list:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][yi01, xi01].squeeze()
V[vv] = ( aiy*((aix*vvtemp).sum(-1)) ).sum(-1)
ds.close()
#%% extract time-dependent fields
print('-- Working on date: ' + dt)
sys.stdout.flush()
ds = nc.Dataset(fn)
for vv in v1_list:
vtemp = ds.variables[vv][:].squeeze()
V[vv] = np.append(V[vv], vtemp)
for vv in v2_list:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][:, yi01, xi01].squeeze()
vtemp = ( aiy*((aix*vvtemp).sum(-1)) ).sum(-1)
V[vv] = np.append(V[vv], vtemp)
for vv in v3_list_rho:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][:, :, yi01, xi01].squeeze()
vtemp = ( aiy*((aix*vvtemp).sum(-1)) ).sum(-1)
V[vv] = vtemp.reshape((S['N'],1))
for vv in v3_list_w:
xi01, yi01, aix, aiy = get_its(ds, vv, Xi0, Yi0, Xi1, Yi1, Aix, Aiy)
vvtemp = ds.variables[vv][:, :, yi01, xi01].squeeze()
vtemp = ( aiy*((aix*vvtemp).sum(-1)) ).sum(-1)
V[vv] = vtemp.reshape((S['N']+1,1))
ds.close()
# create z_rho and z_w (has to be done after we have V['zeta'])
hh = V['h'][:] * np.ones_like(V['zeta'])
z_rho, z_w = zrfun.get_z(hh, V['zeta'][:], S)
V['hh'] = hh
V_long_name['hh'] = 'bottom depth (positive down) as a vector'
V_units['hh'] = 'm'
V['z_rho'] = z_rho
V_long_name['z_rho'] = 'z on rho points (positive up)'
V_units['z_rho'] = 'm'
V['z_w'] = z_w
V_long_name['z_w'] = 'z on w points (positive up)'
V_units['z_w'] = 'm'
v2_list.append('hh')
v3_list_rho.append('z_rho')
v3_list_w.append('z_w')
#%% save the output to NetCDF
out_fn = (outdir +
Ldir['station'] + '_' +
Ldir['date_string'] +
'.nc')
# get rid of the old version, if it exists
try:
os.remove(out_fn)
except OSError:
pass # assume error was because the file did not exist
foo = nc.Dataset(out_fn, 'w')
N = S['N']
NT = len(V['ocean_time'][:])
foo.createDimension('scalar', 1)
foo.createDimension('s_rho', N)
foo.createDimension('s_w', N+1)
foo.createDimension('ocean_time', NT)
for vv in v0_list:
v_var = foo.createVariable(vv, float, ('scalar'))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v1_list:
v_var = foo.createVariable(vv, float, ('ocean_time',))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v2_list:
v_var = foo.createVariable(vv, float, ('ocean_time',))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v3_list_rho:
v_var = foo.createVariable(vv, float, ('s_rho', 'ocean_time'))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
for vv in v3_list_w:
v_var = foo.createVariable(vv, float, ('s_w', 'ocean_time'))
v_var[:] = V[vv][:]
v_var.long_name = V_long_name[vv]
v_var.units = V_units[vv]
foo.close()
Lat = np.array(float(Ldir['lat_str']))
# get grid info
indir = Ldir['roms'] + 'output/' + Ldir['gtagex'] + '/f' + dt + '/'
fn = indir + 'ocean_his_0021.nc' # noon local standart time
G = zrfun.get_basic_info(fn, only_G=True)
S = zrfun.get_basic_info(fn, only_S=True)
# get interpolants for this point
Xi0 = dict(); Yi0 = dict()
Xi1 = dict(); Yi1 = dict()
Aix = dict(); Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1,:]
yy = G['lat_' + grd][:,1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx, extrap_nan=True)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy, extrap_nan=True)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1-xfr, xfr]).reshape((1,1,2))
Aiy[grd] = np.array([1-yfr, yfr]).reshape((1,2))
# generating some lists
v0_list = ['h', 'lon_rho', 'lat_rho', 'lon_u', 'lat_u', 'lon_v', 'lat_v']
v1_list = ['ocean_time']
v2_list = []
v3_list_rho = []
v3_list_w = []
def get_section(ds, vn, x, y, in_dict):
# PLOT CODE
from warnings import filterwarnings
filterwarnings('ignore') # skip a warning message
# GET DATA
G, S, T = zrfun.get_basic_info(in_dict['fn'])
h = G['h']
zeta = ds['zeta'][:].squeeze()
zr = zrfun.get_z(h, zeta, S, only_rho=True)
sectvar = ds[vn][:].squeeze()
L = G['L']
M = G['M']
N = S['N']
lon = G['lon_rho']
lat = G['lat_rho']
mask = G['mask_rho']
maskr = mask.reshape(1, M, L).copy()
mask3 = np.tile(maskr, [N, 1, 1])
zbot = -h # don't need .copy() because of the minus operation
# make sure fields are masked
zeta[mask==False] = np.nan
zbot[mask==False] = np.nan
sectvar[mask3==False] = np.nan
# create dist
earth_rad = zfun.earth_rad(np.mean(lat[:,0])) # m
xrad = np.pi * x /180
yrad = np.pi * y / 180
dx = earth_rad * np.cos(yrad[1:]) * np.diff(xrad)
dy = earth_rad * np.diff(yrad)
ddist = np.sqrt(dx**2 + dy**2)
dist = np.zeros(len(x))
dist[1:] = ddist.cumsum()/1000 # km
# find the index of zero
i0, i1, fr = zfun.get_interpolant(np.zeros(1), dist)
idist0 = i0
distr = dist.reshape(1, len(dist)).copy()
dista = np.tile(distr, [N, 1]) # array
# pack fields to process in dicts
d2 = dict()
d2['zbot'] = zbot
d2['zeta'] = zeta
d2['lon'] = lon
d2['lat'] = lat
d3 = dict()
d3['zr'] = zr
d3['sectvar'] = sectvar
# get vectors describing the (plaid) grid
xx = lon[1,:]
yy = lat[:,1]
col0, col1, colf = zfun.get_interpolant(x, xx)
row0, row1, rowf = zfun.get_interpolant(y, yy)
# and prepare them to do the bilinear interpolation
colff = 1 - colf
rowff = 1 - rowf
# now actually do the interpolation
# 2-D fields
v2 = dict()
for fname in d2.keys():
fld = d2[fname]
fldi = (rowff*(colff*fld[row0, col0] + colf*fld[row0, col1])
+ rowf*(colff*fld[row1, col0] + colf*fld[row1, col1]))
if type(fldi) == np.ma.core.MaskedArray:
fldi = fldi.data # just the data, not the mask
v2[fname] = fldi
# 3-D fields
v3 = dict()
for fname in d3.keys():
fld = d3[fname]
fldi = (rowff*(colff*fld[:, row0, col0] + colf*fld[:, row0, col1])
+ rowf*(colff*fld[:, row1, col0] + colf*fld[:, row1, col1]))
if type(fldi) == np.ma.core.MaskedArray:
fldid = fldi.data # just the data, not the mask
fldid[fldi.mask == True] = np.nan
v3[fname] = fldid
v3['dist'] = dista # distance in km
# make "full" fields by padding top and bottom
nana = np.nan * np.ones((N + 2, len(dist))) # blank array
v3['zrf'] = nana.copy()
v3['zrf'][0,:] = v2['zbot']
v3['zrf'][1:-1,:] = v3['zr']
v3['zrf'][-1,:] = v2['zeta']
#
v3['sectvarf'] = nana.copy()
v3['sectvarf'][0,:] = v3['sectvar'][0,:]
v3['sectvarf'][1:-1,:] = v3['sectvar']
v3['sectvarf'][-1,:] = v3['sectvar'][-1,:]
#
v3['distf'] = nana.copy()
v3['distf'][0,:] = v3['dist'][0,:]
v3['distf'][1:-1,:] = v3['dist']
v3['distf'][-1,:] = v3['dist'][-1,:]
# attempt to skip over nan's
v3.pop('zr')
v3.pop('sectvar')
v3.pop('dist')
mask3 = ~np.isnan(v3['sectvarf'][:])
#print(mask3.shape)
mask2 = mask3[-1,:]
dist = dist[mask2]
NC = len(dist)
NR = mask3.shape[0]
for k in v2.keys():
#print('v2 key: ' + k)
v2[k] = v2[k][mask2]
for k in v3.keys():
#print('v3 key: ' + k)
v3[k] = v3[k][mask3]
v3[k] = v3[k].reshape((NR, NC))
#print(v3[k].shape)
return v2, v3, dist, idist0
x = G['lon_rho']
y = G['lat_rho']
mask = G['mask_rho']
F = ds0[name][:].squeeze()
if name == 'zeta':
fld = F.copy()
fldf = fld.copy() # initialize the "filled" field
if tt == 0:
# fill masked values using nearest neighbor
xyorig = np.array((X[~fld.mask], Y[~fld.mask])).T
xynew = np.array((X[fld.mask], Y[fld.mask])).T
a = cKDTree(xyorig).query(xynew)
aa_rho = a[1]
# INTERPOLATION
# get interpolants
xi0_rho, xi1_rho, xf_rho = zfun.get_interpolant(
x, X[0, :], extrap_nan=True)
yi0_rho, yi1_rho, yf_rho = zfun.get_interpolant(
y, Y[:, 0], extrap_nan=True)
# create the filled field
fldf[fld.mask] = fld[~fld.mask][aa_rho]
fx = fldf.data
# do the bi-linear interpolation
u00 = fx[yi0_rho, xi0_rho]
u10 = fx[yi1_rho, xi0_rho]
u01 = fx[yi0_rho, xi1_rho]
u11 = fx[yi1_rho, xi1_rho]
fi = (1 - yf_rho) * (
(1 - xf_rho) * u00 + xf_rho * u01) + yf_rho * (
(1 - xf_rho) * u10 + xf_rho * u11)
ff = np.reshape(fi, x.shape)
fm = np.ma.masked_where(mask == False, ff)
G, S, T = zrfun.get_basic_info(fn) ds.close() # get grid and indices (rho-grid) for extraction domain # x0 = -122.85 # x1 = -122.6 # y0 = 47.6 # y1 = 47.9 x0 = -122.5 x1 = -122.3 y0 = 47.4 y1 = 47.55 Lon = G['lon_rho'] Lat = G['lat_rho'] Lon_vec = Lon[0,:] Lat_vec = Lat[:,0] i0, junk, junk = zfun.get_interpolant(np.array(x0), Lon_vec, extrap_nan=True) junk, i1, junk = zfun.get_interpolant(np.array(x1), Lon_vec, extrap_nan=True) j0, junk, junk = zfun.get_interpolant(np.array(y0), Lat_vec, extrap_nan=True) junk, j1, junk = zfun.get_interpolant(np.array(y1), Lat_vec, extrap_nan=True) i0 = i0[0] i1 = i1[0] + 1 j0 = j0[0] j1 = j1[0] + 1 lon_rho = Lon[j0:j1, i0:i1] lat_rho = Lat[j0:j1, i0:i1] h = G['h'][j0:j1, i0:i1] NR, NC = lon_rho.shape N = S['N']
for iz in range(F.shape[0]):
# EXTRAPOLATION
fld = F[iz, :, :].squeeze()
# do the extrapolation using nearest neighbor
fldf = fld.copy() # initialize the "filled" field
if iz==0:
xyorig = np.array((X[~fld.mask],Y[~fld.mask])).T
xynew = np.array((X[fld.mask],Y[fld.mask])).T
a = cKDTree(xyorig).query(xynew)
aa = a[1]
fldf[fld.mask] = fld[~fld.mask][aa]
fx = fldf.data
# INTERPOLATION
if iz==0:
# get interpolants
xi0, xi1, xf = zfun.get_interpolant(x,X[0,:], extrap_nan=True)
yi0, yi1, yf = zfun.get_interpolant(y,Y[:,0], extrap_nan=True)
# bi linear interpolation
u00 = fx[yi0,xi0]
u10 = fx[yi1,xi0]
u01 = fx[yi0,xi1]
u11 = fx[yi1,xi1]
fi = (1-yf)*((1-xf)*u00 + xf*u01) + yf*((1-xf)*u10 + xf*u11)
ff = np.reshape(fi, x.shape)
fm = np.ma.masked_where(mask==False, ff)
ds1[name][tt,iz,:,:] = fm
print(' -- %s took %0.1f seconds' % (name, time.time() - tt00))
sys.stdout.flush()
print('Hour %d took %0.1f seconds' % (tt, time.time() - tt0))
G, S, T = zrfun.get_basic_info(fn) ds.close() # get grid and indices (rho-grid) for extraction domain # x0 = -122.85 # x1 = -122.6 # y0 = 47.6 # y1 = 47.9 x0 = -122.5 x1 = -122.3 y0 = 47.4 y1 = 47.55 Lon = G['lon_rho'] Lat = G['lat_rho'] Lon_vec = Lon[0, :] Lat_vec = Lat[:, 0] i0, junk, junk = zfun.get_interpolant(np.array(x0), Lon_vec, extrap_nan=True) junk, i1, junk = zfun.get_interpolant(np.array(x1), Lon_vec, extrap_nan=True) j0, junk, junk = zfun.get_interpolant(np.array(y0), Lat_vec, extrap_nan=True) junk, j1, junk = zfun.get_interpolant(np.array(y1), Lat_vec, extrap_nan=True) i0 = i0[0] i1 = i1[0] + 1 j0 = j0[0] j1 = j1[0] + 1 lon_rho = Lon[j0:j1, i0:i1] lat_rho = Lat[j0:j1, i0:i1] h = G['h'][j0:j1, i0:i1] NR, NC = lon_rho.shape N = S['N']
def get_tracks(fn_list, plon0, plat0, pcs0, TR, trim_loc=False):
"""
This is the main function doing the particle tracking.
"""
# unpack items needed from TR
if TR['rev']:
dir_tag = 'reverse'
else:
dir_tag = 'forward'
surface = not TR['3d']
turb = TR['turb']
ndiv = TR['ndiv']
windage = TR['windage']
# get basic info
G = zrfun.get_basic_info(fn_list[0], only_G=True)
maskr = np.ones_like(G['mask_rho']) # G['mask_rho'] = True in water
maskr[G['mask_rho'] == False] = 0 # maskr = 1 in water
S = zrfun.get_basic_info(fn_list[0], only_S=True)
# get time vector of history files
NT = len(fn_list)
rot = np.nan * np.ones(NT)
counter = 0
for fn in fn_list:
ds = nc4.Dataset(fn)
rot[counter] = ds.variables['ocean_time'][:].squeeze()
counter += 1
ds.close
# This is an attempt to fix a bug that happens avery couple of weeks,
# one which I can't reliably reproduce! In theory fn_list is sorted
# by the calling function, but maybe ...??
# Now (2019.11.08) I think this glitch was caused when the carbon code was
# modifying the history files at the same time as the particle tracking was
# going.
rot.sort()
delta_t = rot[1] - rot[0] # second between saves
# this is how we track backwards in time
if dir_tag == 'reverse':
delta_t = -delta_t
fn_list = fn_list[::-1]
# make vectors to feed to interpolant maker
R = dict()
R['rlonr'] = G['lon_rho'][0, :].squeeze()
R['rlatr'] = G['lat_rho'][:, 0].squeeze()
R['rlonu'] = G['lon_u'][0, :].squeeze()
R['rlatu'] = G['lat_u'][:, 0].squeeze()
R['rlonv'] = G['lon_v'][0, :].squeeze()
R['rlatv'] = G['lat_v'][:, 0].squeeze()
R['rcsr'] = S['Cs_r'][:]
R['rcsw'] = S['Cs_w'][:]
# these lists are used internally to get other variables as needed
# (they must all be present in the history files)
vn_list_vel = ['u', 'v', 'w']
vn_list_zh = ['zeta', 'h']
vn_list_wind = ['Uwind', 'Vwind']
vn_list_other = ['salt', 'temp'] + vn_list_zh + vn_list_vel
if windage > 0:
vn_list_other = vn_list_other + vn_list_wind
# plist_main is what ends up written to output
plist_main = ['lon', 'lat', 'cs', 'ot', 'z'] + vn_list_other
if save_dia:
# diagnostic info
vn_list_dia = ['hit_sidewall', 'hit_bottom', 'hit_top', 'bad_pcs']
# Step through times.
#
counter = 0
# rot is a list of all the ocean times (sec) in the current file list(e.g. 25)
# and pot is a single one of these
for pot in rot[:-1]:
# get time indices of surrounding (in time) history files
it0, it1, frt = zfun.get_interpolant(np.array(pot),
rot,
extrap_nan=False)
# get Datasets
ds0 = nc4.Dataset(fn_list[it0[0]])
ds1 = nc4.Dataset(fn_list[it1[0]])
if counter == 0:
if trim_loc == True:
# remove points on land
pmask = zfun.interp_scattered_on_plaid(plon0,
plat0,
R['rlonr'],
R['rlatr'],
maskr,
exnan=True)
# keep only points with pmask >= maskr_crit
pcond = pmask >= maskr_crit
plon = plon0[pcond]
plat = plat0[pcond]
pcs = pcs0[pcond]
else:
plon = plon0.copy()
plat = plat0.copy()
pcs = pcs0.copy()
# create result arrays
NP = len(plon)
P = dict()
for vn in plist_main:
# NOTE: info is packed in a dict P of arrays
# packed in order (time, particle)
P[vn] = np.nan * np.ones((NT, NP))
# initialize diagnostic arrays
if save_dia:
for vn in vn_list_dia:
P[vn] = np.zeros((NT, NP)) # 0=OK, 1=violation
# write positions to the results arrays
P['lon'][it0, :] = plon
P['lat'][it0, :] = plat
if surface == True:
pcs[:] = S['Cs_r'][-1]
P['cs'][it0, :] = pcs
P = get_properties(vn_list_other, ds0, it0, P, plon, plat, pcs, R,
surface)
delt = delta_t / ndiv
# do the particle tracking for a single pair of history files in ndiv steps
for nd in range(ndiv):
fr0 = nd / ndiv
fr1 = (nd + 1) / ndiv
frmid = (fr0 + fr1) / 2
# RK4 integration
V0, ZH0 = get_vel(vn_list_vel, vn_list_zh, ds0, ds1, plon, plat,
pcs, R, fr0, surface)
plon1, plat1, pcs1, dia_dict = update_position(
R, maskr, V0, ZH0, S, delt / 2, plon, plat, pcs, surface)
V1, ZH1 = get_vel(vn_list_vel, vn_list_zh, ds0, ds1, plon1, plat1,
pcs1, R, frmid, surface)
plon2, plat2, pcs2, dia_dict = update_position(
R, maskr, V1, ZH1, S, delt / 2, plon, plat, pcs, surface)
V2, ZH2 = get_vel(vn_list_vel, vn_list_zh, ds0, ds1, plon2, plat2,
pcs2, R, frmid, surface)
plon3, plat3, pcs3, dia_dict = update_position(
R, maskr, V2, ZH2, S, delt, plon, plat, pcs, surface)
V3, ZH3 = get_vel(vn_list_vel, vn_list_zh, ds0, ds1, plon3, plat3,
pcs3, R, fr1, surface)
# add windage, calculated from the middle time
if (surface == True) and (windage > 0):
Vwind = get_wind(vn_list_wind, ds0, ds1, plon, plat, pcs, R,
frmid, surface)
Vwind3 = np.concatenate((windage * Vwind, np.zeros((NP, 1))),
axis=1)
else:
Vwind3 = np.zeros((NP, 3))
plon, plat, pcs, dia_dict = update_position(
R, maskr, (V0 + 2 * V1 + 2 * V2 + V3) / 6 + Vwind3,
(ZH0 + 2 * ZH1 + 2 * ZH2 + ZH3) / 6, S, delt, plon, plat, pcs,
surface)
if save_dia:
# diagnostic info for horizontal advection
P['hit_sidewall'][it1, :] = dia_dict['hit_sidewall']
# add turbulence to vertical position change (advection already added above)
if turb == True:
# pull values of VdAKs and add up to 3-dimensions
VdAKs = get_dAKs(vn_list_zh, ds0, ds1, plon, plat, pcs, R, S,
frmid, surface)
# print('VdAKs has %d nans out of %d' % (np.isnan(VdAKs).sum(), len(VdAKs)))
VdAKs3 = np.concatenate((np.zeros(
(NP, 2)), VdAKs[:, np.newaxis]),
axis=1)
# update position advecting vertically with 1/2 of AKs gradient
ZH = get_zh(vn_list_zh, ds0, ds1, plon, plat, pcs, R, frmid,
surface)
plon_junk, plat_junk, pcs_half, dia_dict = update_position(
R, maskr, VdAKs3 / 2, ZH, S, delt, plon, plat, pcs,
surface)
# get AKs at this height, and thence the turbulent velocity
Vturb = get_turb(ds0, ds1, VdAKs, delt, plon, plat, pcs_half,
R, frmid, surface)
# update position vertically for real
Vturb3 = np.concatenate((np.zeros(
(NP, 2)), Vturb[:, np.newaxis]),
axis=1)
plon_junk, plat_junk, pcs, dia_dict = update_position(
R, maskr, Vturb3, ZH, S, delt, plon, plat, pcs, surface)
if save_dia:
# diagnostic info for vertical advection
P['bad_pcs'][it1, :] = dia_dict['bad_pcs']
P['hit_top'][it1, :] = dia_dict['hit_top']
P['hit_bottom'][it1, :] = dia_dict['hit_bottom']
# write positions to the results arrays
P['lon'][it1, :] = plon
P['lat'][it1, :] = plat
if surface == True:
pcs[:] = S['Cs_r'][-1]
P['cs'][it1, :] = pcs
P = get_properties(vn_list_other, ds1, it1, P, plon, plat, pcs, R,
surface)
ds0.close()
ds1.close()
counter += 1
# by doing this the points are going forward in time
if dir_tag == 'reverse':
for vn in plist_main:
P[vn] = P[vn][::-1, :]
# and save the time vector (seconds in whatever the model reports)
P['ot'] = rot
return P
lon_vec = G['lon_rho'][0,:]
lat_vec = G['lat_rho'][:,0]
# test values for cascadia1
# XBS = [55, 80] # x-limits
# YBS = [295, 325] # y-limits
# Bounds of subdomain (rho grid) from Susan Allen 2018_11
lon0 = -125.016452048434
lon1 = -124.494612925929
lat0 = 48.3169463809796
lat1 = 48.7515055163539
# find indices containing these bonds
i0, i1, ifr = zfun.get_interpolant(np.array([lon0,lon1]), lon_vec, extrap_nan=True)
j0, j1, jfr = zfun.get_interpolant(np.array([lat0,lat1]), lat_vec, extrap_nan=True)
if np.nan in ifr:
print('Warning: nan in ifr')
if np.nan in jfr:
print('Warning: nan in jfr')
XBS = [i0[0], i1[1]]
YBS = [j0[0], j1[1]]
print(XBS)
print(YBS)
from datetime import datetime, timedelta
start_time = datetime.now()
# the output file name
out_fn = (Ldir['roms'] + 'output/' + Ldir['gtagex'] +
Lon = np.array(float(Ldir['lon_str']))
Lat = np.array(float(Ldir['lat_str']))
# get grid info
fn = fn_list[0]
G = zrfun.get_basic_info(fn, only_G=True)
S = zrfun.get_basic_info(fn, only_S=True)
# get interpolants for this point
Xi0 = dict(); Yi0 = dict()
Xi1 = dict(); Yi1 = dict()
Aix = dict(); Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1,:]
yy = G['lat_' + grd][:,1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx, extrap_nan=True)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy, extrap_nan=True)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1-xfr, xfr]).reshape((1,1,2))
Aiy[grd] = np.array([1-yfr, yfr]).reshape((1,2))
# generating some lists
v0_list = ['h', 'lon_rho', 'lat_rho', 'lon_u', 'lat_u', 'lon_v', 'lat_v']
v1_list = ['ocean_time']
v2_list = []
v3_list_rho = []
v3_list_w = []
Lon = np.array(float(Ldir['lon_str']))
Lat = np.array(float(Ldir['lat_str']))
# get grid info
indir = Ldir['roms'] + 'output/' + Ldir['gtagex'] + '/f' + date_list[0] + '/'
fn = indir + 'ocean_his_0002.nc'
[G, S] = zfun.get_basic_info(fn, getS=True, getT=False)
# get interpolants for this point
Xi0 = dict(); Yi0 = dict()
Xi1 = dict(); Yi1 = dict()
Aix = dict(); Aiy = dict()
for grd in ['rho', 'u', 'v']:
xx = G['lon_' + grd][1,:]
yy = G['lat_' + grd][:,1]
xi0, xi1, xfr = zfun.get_interpolant(Lon, xx)
yi0, yi1, yfr = zfun.get_interpolant(Lat, yy)
Xi0[grd] = xi0
Yi0[grd] = yi0
Xi1[grd] = xi1
Yi1[grd] = yi1
# create little arrays that are used in the actual interpolation
Aix[grd] = np.array([1-xfr, xfr]).reshape((1,1,2))
Aiy[grd] = np.array([1-yfr, yfr]).reshape((1,2))
# generating some lists
v0_list = ['h', 'lon_rho', 'lat_rho', 'lon_u', 'lat_u', 'lon_v', 'lat_v']
v1_list = ['ocean_time']
v2_list = []
v3_list_rho = []
v3_list_w = []
def get_inds(x0, x1, y0, y1, G, verbose=False):
# determine the direction of the section
# and make sure indices are *increasing*
if (x0 == x1) and (y0 != y1):
sdir = 'NS'
a = [y0, y1]
a.sort()
y0 = a[0]
y1 = a[1]
elif (x0 != x1) and (y0 == y1):
sdir = 'EW'
a = [x0, x1]
a.sort()
x0 = a[0]
x1 = a[1]
else:
print('Input points do not form a proper section')
sdir = 'bad'
sys.exit()
# we assume a plaid grid, as usual
if sdir == 'NS':
lon = G['lon_u'][0, :].squeeze()
lat = G['lat_u'][:, 0].squeeze()
elif sdir == 'EW':
lon = G['lon_v'][0, :].squeeze()
lat = G['lat_v'][:, 0].squeeze()
# we get all 4 i's or j's but only 3 are used
i0, i1, fr = zfun.get_interpolant(np.array([x0]), lon, extrap_nan=True)
if np.isnan(fr):
print('Bad x point')
sys.exit()
else:
ii0 = int(i0)
i0, i1, fr = zfun.get_interpolant(np.array([x1]), lon, extrap_nan=True)
if np.isnan(fr):
print('Bad x point')
sys.exit()
else:
ii1 = int(i1)
j0, j1, fr = zfun.get_interpolant(np.array([y0]), lat, extrap_nan=True)
if np.isnan(fr):
print('Bad y0 point')
sys.exit()
else:
jj0 = int(j0)
j0, j1, fr = zfun.get_interpolant(np.array([y1]), lat, extrap_nan=True)
if np.isnan(fr):
print('Bad y1 point')
sys.exit()
else:
jj1 = int(j1)
# get mask and trim indices
# Note: the mask in G is True on water points
if sdir == 'NS':
mask = G['mask_u'][jj0:jj1 + 1, ii0]
# Note: argmax finds the index of the first True in this case
igood0 = np.argmax(mask)
igood1 = np.argmax(mask[::-1])
# check to see if section is "closed"
if (igood0 == 0) | (igood1 == 0):
print('Warning: not closed one or both ends')
# keep only to end of water points, to allow for ocean sections
Mask = mask[igood0:-igood1]
# and change the indices to match. These will be the indices
# of the start and end points.
jj0 = jj0 + igood0
jj1 = jj1 - igood1
if verbose:
print(' sdir=%2s: jj0=%4d, jj1=%4d, ii0=%4d' %
(sdir, jj0, jj1, ii0))
Lat = lat[jj0:jj1 + 1]
Lon = lon[ii0] * np.ones_like(Mask)
elif sdir == 'EW':
mask = G['mask_v'][jj0, ii0:ii1 + 1]
igood0 = np.argmax(mask)
igood1 = np.argmax(mask[::-1])
# check to see if section is "closed"
if (igood0 == 0) | (igood1 == 0):
print('Warning: not closed one or both ends')
Mask = mask[igood0:-igood1]
ii0 = ii0 + igood0
ii1 = ii1 - igood1
if verbose:
print(' sdir=%2s: jj0=%4d, ii0=%4d, ii1=%4d' %
(sdir, jj0, ii0, ii1))
Lon = lon[ii0:ii1 + 1]
Lat = lat[jj0] * np.ones_like(Mask)
return ii0, ii1, jj0, jj1, sdir, Lon, Lat, Mask
def get_extrapolated(in_fn, L, M, N, X, Y, lon, lat, z, Ldir, add_CTD=False, fix_NSoG=False):
b = pickle.load(open(in_fn, 'rb'))
vn_list = list(b.keys())
# check that things are the expected shape
def check_coords(shape_tuple, arr_shape):
if arr_shape != shape_tuple:
print('WARNING: array shape mismatch')
for vn in vn_list:
if vn == 'dt':
pass
elif vn == 'ssh':
check_coords((M, L), b[vn].shape)
else:
check_coords((N, M, L), b[vn].shape)
# creat output array and add dt to it.
vn_list.remove('dt')
V = dict()
for vn in vn_list:
V[vn] = np.nan + np.ones(b[vn].shape)
V['dt'] = b['dt']
# extrapolate ssh
vn = 'ssh'
v = b[vn]
if fix_NSoG:
print(' ^^ Fixing NSoG ^^')
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# (a) Mouth of Strait of Juan de Fuca
i0, i1, fr = zfun.get_interpolant(np.array([-124.7, -124.5]), lon, extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([48.4, 48.6]), lat, extrap_nan=False)
zeta_jdf = v[j0[0]:j1[1]+1, i0[0]:i1[1]+1]
#
# (b) Northern Strait of Georgia
i0, i1, fr = zfun.get_interpolant(np.array([-125.5, -124.5]), lon, extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([50.15, 50.45]), lat, extrap_nan=False)
v[j0[0]:j1[1]+1, i0[0]:i1[1]+1] = np.nanmean(zeta_jdf) + 0.1
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
vv = extrap_nearest_to_masked(X, Y, v)
V[vn] = vv
vn_list.remove('ssh')
# extrapolate 3D fields
for vn in vn_list:
v = b[vn]
if vn == 't3d':
v0 = np.nanmin(v)
elif vn == 's3d':
v0 = np.nanmax(v)
if vn in ['t3d', 's3d']:
print(' -- extrapolating ' + vn)
if add_CTD==False:
for k in range(N):
fld = v[k, :, :]
fldf = extrap_nearest_to_masked(X, Y, fld, fld0=v0)
V[vn][k, :, :] = fldf
elif add_CTD==True:
print(vn + ' Adding CTD data before extrapolating')
Cast_dict, sta_df = Ofun_CTD.get_casts(Ldir)
for k in range(N):
fld = v[k, :, :]
zz = z[k]
xyorig, fldorig = Ofun_CTD.get_orig(Cast_dict, sta_df,
X, Y, fld, lon, lat, zz, vn)
fldf = Ofun_CTD.extrap_nearest_to_masked_CTD(X,Y,fld,
xyorig=xyorig,fldorig=fldorig,fld0=v0)
V[vn][k, :, :] = fldf
elif vn in ['u3d', 'v3d']:
print(' -- extrapolating ' + vn)
vv = v.copy()
vv = np.ma.masked_where(np.isnan(vv), vv)
vv[vv.mask] = 0
V[vn] = vv.data
# Create ubar and vbar.
# Note: this is slightly imperfect because the z levels are at the same
# position as the velocity levels.
dz = np.nan * np.ones((N, 1, 1))
dz[1:, 0, 0]= np.diff(z)
dz[0, 0, 0] = dz[1, 0, 0]
# account for the fact that the new hycom fields do not show up masked
u3d = np.ma.masked_where(np.isnan(b['u3d']),b['u3d'])
v3d = np.ma.masked_where(np.isnan(b['v3d']),b['v3d'])
dz3 = dz * np.ones_like(u3d) # make dz a masked array
b['ubar'] = np.sum(u3d*dz3, axis=0) / np.sum(dz3, axis=0)
b['vbar'] = np.sum(v3d*dz3, axis=0) / np.sum(dz3, axis=0)
for vn in ['ubar', 'vbar']:
v = b[vn]
vv = v.copy()
vv = np.ma.masked_where(np.isnan(vv), vv)
vv[vv.mask] = 0
V[vn] = vv.data
# calculate potential temperature
press_db = -z.reshape((N,1,1))
V['theta'] = seawater.ptmp(V['s3d'], V['t3d'], press_db)
return V
def get_tracks(fn_list,
plon0,
plat0,
pcs0,
dir_tag,
method,
surface,
turb,
ndiv,
windage,
bound='stop',
dep_range=None):
'''
Create tracks of particles starting from the locations at (plon0, plat0, pcs0)
Inputs:
fn_list - array of dataset file locations
plon0/plat0/pcs0 - arrays of position (longitude/latitude/proportion of total depth)
dir_tag - string determining direction of tracking ('forward' or 'reverse')
method - string determining interpolation level ('rk2' or 'rk4')
surface - Boolean, True traps particles to the surface
turb - Boolean, True adds random vertical walk
ndiv - int determining number of divisions to make between saves for the integration
windage - float (0 to 1) determining proportion of velocity to add from surface wind
bound - string determing boundary velocity conditions
'stop' (default) prevents particles from moving when a gridcell contains land
'reflect' causes particles to move away from gridcells that contain land
dep_range - tuple containing arrays (minimum depth, maximum depth) for all particles
Outputs:
P - dictionary containing variables in plist_main
each variable has rows of time and columns of particles
except the array 'ot', seconds since 1970
G - dictionary containing grid and bathymetry information
S - dictionary containing water surface information
'''
# Bartos - removed on land particle removal, so set these to plon, not plonA
plon = plon0.copy()
plat = plat0.copy()
pcs = pcs0.copy()
# get basic info
G = zrfun.get_basic_info(fn_list[0], only_G=True)
S = zrfun.get_basic_info(fn_list[0], only_S=True)
# get time vector of history files
NT = len(fn_list)
rot = np.nan * np.ones(NT)
counter = 0
for fn in fn_list:
ds = nc4.Dataset(fn)
rot[counter] = ds.variables['ocean_time'][:].squeeze()
counter += 1
ds.close
delta_t = rot[1] - rot[0] # seconds between saves
# this is how we track backwards in time
if dir_tag == 'reverse':
delta_t = -delta_t
fn_list = fn_list[::-1]
# make vectors to feed to interpolant maker
R = dict()
R['rlonr'] = G['lon_rho'][0, :].squeeze()
R['rlatr'] = G['lat_rho'][:, 0].squeeze()
R['rlonu'] = G['lon_u'][0, :].squeeze()
R['rlatu'] = G['lat_u'][:, 0].squeeze()
R['rlonv'] = G['lon_v'][0, :].squeeze()
R['rlatv'] = G['lat_v'][:, 0].squeeze()
R['rcsr'] = S['Cs_r'][:]
R['rcsw'] = S['Cs_w'][:]
# these lists are used internally to get other variables as needed
vn_list_vel = ['u', 'v', 'w']
vn_list_zh = ['zeta', 'h']
vn_list_wind = ['Uwind', 'Vwind']
vn_list_other = ['salt', 'temp', 'zeta', 'h', 'u', 'v', 'w']
# Step through times.
# Bartos - removed on land particle removal
counter = 0
nrot = len(rot)
for pot in rot[:-1]:
# if np.mod(counter,24) == 0:
# print(' - time %d out of %d' % (counter, nrot))
# sys.stdout.flush()
# get time indices
it0, it1, frt = zfun.get_interpolant(np.array(pot),
rot,
extrap_nan=False)
# get the velocity zeta, and h at all points
ds0 = nc4.Dataset(fn_list[it0[0]])
ds1 = nc4.Dataset(fn_list[it1[0]])
if counter == 0:
# create result arrays
NP = len(plon)
P = dict()
# plist main is what ends up written to output
plist_main = [
'lon', 'lat', 'cs', 'ot', 'z', 'zeta', 'zbot', 'salt', 'temp',
'u', 'v', 'w', 'Uwind', 'Vwind', 'h'
]
for vn in plist_main:
P[vn] = np.nan * np.ones((NT, NP))
# write positions to the results arrays
P['lon'][it0, :] = plon
P['lat'][it0, :] = plat
if surface == True:
pcs[:] = S['Cs_r'][-1]
P['cs'][it0, :] = pcs
P = get_properties(vn_list_other, ds0, it0, P, plon, plat, pcs, R,
surface)
delt = delta_t / ndiv
for nd in range(ndiv):
fr0 = nd / ndiv
fr1 = (nd + 1) / ndiv
frmid = (fr0 + fr1) / 2
if method == 'rk4':
# RK4 integration
V0, ZH0 = get_vel(vn_list_vel,
vn_list_zh,
ds0,
ds1,
plon,
plat,
pcs,
R,
fr0,
surface,
bound=bound)
plon1, plat1, pcs1 = update_position(V0, ZH0, S, delt / 2,
plon, plat, pcs, surface)
V1, ZH1 = get_vel(vn_list_vel,
vn_list_zh,
ds0,
ds1,
plon1,
plat1,
pcs1,
R,
frmid,
surface,
bound=bound)
plon2, plat2, pcs2 = update_position(V1, ZH1, S, delt / 2,
plon, plat, pcs, surface)
V2, ZH2 = get_vel(vn_list_vel,
vn_list_zh,
ds0,
ds1,
plon2,
plat2,
pcs2,
R,
frmid,
surface,
bound=bound)
plon3, plat3, pcs3 = update_position(V2, ZH2, S, delt, plon,
plat, pcs, surface)
V3, ZH3 = get_vel(vn_list_vel,
vn_list_zh,
ds0,
ds1,
plon3,
plat3,
pcs3,
R,
fr1,
surface,
bound=bound)
# add windage, calculated from the middle time
if (surface == True) and (windage > 0):
Vwind = get_wind(vn_list_wind, ds0, ds1, plon, plat, pcs,
R, frmid, surface)
Vwind3 = np.concatenate((windage * Vwind, np.zeros(
(NP, 1))),
axis=1)
else:
Vwind3 = np.zeros((NP, 3))
plon, plat, pcs = update_position(
(V0 + 2 * V1 + 2 * V2 + V3) / 6 + Vwind3,
(ZH0 + 2 * ZH1 + 2 * ZH2 + ZH3) / 6, S, delt, plon, plat,
pcs, surface)
# Bartos - begin turbulence edit
# add turbulence in two distinct timesteps
if turb == True:
# pull values of VdAKs and add up to 3-dimensions
VdAKs = get_dAKs(vn_list_zh, ds0, ds1, plon, plat, pcs, R,
S, frmid, surface)
VdAKs3 = np.concatenate((np.zeros(
(NP, 2)), VdAKs[:, np.newaxis]),
axis=1)
# update position with 1/2 of AKs gradient
plon, plat, pcs = update_position(
VdAKs3 / 2, (ZH0 + 2 * ZH1 + 2 * ZH2 + ZH3) / 6, S,
delt, plon, plat, pcs, surface)
# update position with rest of turbulence
Vturb = get_turb(ds0, ds1, VdAKs, delta_t, plon, plat, pcs,
R, frmid, surface)
Vturb3 = np.concatenate((np.zeros(
(NP, 2)), Vturb[:, np.newaxis]),
axis=1)
plon, plat, pcs = update_position(
Vturb3, (ZH0 + 2 * ZH1 + 2 * ZH2 + ZH3) / 6, S, delt,
plon, plat, pcs, surface)
# Bartos - end edit
elif method == 'rk2':
# RK2 integration
V0, ZH0 = get_vel(vn_list_vel,
vn_list_zh,
ds0,
ds1,
plon,
plat,
pcs,
R,
fr0,
surface,
bound=bound)
plon1, plat1, pcs1 = update_position(V0, ZH0, S, delt / 2,
plon, plat, pcs, surface)
V1, ZH1 = get_vel(vn_list_vel,
vn_list_zh,
ds0,
ds1,
plon1,
plat1,
pcs1,
R,
frmid,
surface,
bound=bound)
# add windage, calculated from the middle time
if (surface == True) and (windage > 0):
Vwind = get_wind(vn_list_wind, ds0, ds1, plon, plat, pcs,
R, frmid, surface)
Vwind3 = np.concatenate((windage * Vwind, np.zeros(
(NP, 1))),
axis=1)
else:
Vwind3 = np.zeros((NP, 3))
plon, plat, pcs = update_position(V1 + Vwind3, ZH1, S, delt,
plon, plat, pcs, surface)
# Bartos - begin turbulence edit
# add turbulence in two distinct timesteps
if turb == True:
# pull values of VdAKs and add up to 3-dimensions
VdAKs = get_dAKs(vn_list_zh, ds0, ds1, plon, plat, pcs, R,
S, frmid, surface)
VdAKs3 = np.concatenate((np.zeros(
(NP, 2)), VdAKs[:, np.newaxis]),
axis=1)
# update position with 1/2 of AKs gradient
plon, plat, pcs = update_position(VdAKs3 / 2, ZH1, S, delt,
plon, plat, pcs, surface)
# update position with rest of turbulence
Vturb = get_turb(ds0, ds1, VdAKs, delta_t, plon, plat, pcs,
R, frmid, surface)
Vturb3 = np.concatenate((np.zeros(
(NP, 2)), Vturb[:, np.newaxis]),
axis=1)
plon, plat, pcs = update_position(Vturb3, ZH1, S, delt,
plon, plat, pcs, surface)
# write positions to the results arrays
P['lon'][it1, :] = plon
P['lat'][it1, :] = plat
if surface == True:
pcs[:] = S['Cs_r'][-1]
P['cs'][it1, :] = pcs
P = get_properties(vn_list_other, ds1, it1, P, plon, plat, pcs, R,
surface)
# Bartos - begin maximum and minimum depth edit
if dep_range != None:
dep_max = dep_range[1]
# mask of depths below dep_max
dep_mask = P['z'][it1, :] < dep_max
dep_mask = dep_mask.squeeze()
# total depth
dep_tot = P['zeta'][it1, dep_mask] + P['h'][it1, dep_mask]
# replace masked depths with dep_max
P['z'][it1, dep_mask] = dep_max[dep_mask]
P['cs'][it1, dep_mask] = dep_max[dep_mask] / dep_tot
pcs[dep_mask] = dep_max[dep_mask] / dep_tot
dep_min = dep_range[0]
# mask of depths above dep_min
dep_mask = P['z'][it1, :] > dep_min
dep_mask = dep_mask.squeeze()
# total depth
dep_tot = P['zeta'][it1, dep_mask] + P['h'][it1, dep_mask]
# replace masked depths with dep_min
P['z'][it1, dep_mask] = dep_min[dep_mask]
P['cs'][it1, dep_mask] = dep_min[dep_mask] / dep_tot
pcs[dep_mask] = dep_max[dep_mask] / dep_tot
# Bartos - end edits
ds0.close()
ds1.close()
counter += 1
# by doing this the points are going forward in time
if dir_tag == 'reverse':
for vn in plist_main:
P[vn] = P[vn][::-1, :]
# and save the time vector (seconds in whatever the model reports)
P['ot'] = rot
return P, G, S
lon = dsr['lon'][:]
lat = dsr['lat'][:]
ot_vec = dsr['ot'][:]
days = (ot_vec - ot_vec[0]) / 86400
dt_list = [Lfun.modtime_to_datetime(ot) for ot in ot_vec]
NT, NP = lon.shape
day_list = []
frac_list = []
for tt in range(0, NT, 96):
x = lon[tt, :]
y = lat[tt, :]
#tt1 = time()
i0, i1, frx = zfun.get_interpolant(x, xvec)
j0, j1, fry = zfun.get_interpolant(y, yvec)
ii = i0.data + np.round(frx.data)
jj = j0.data + np.round(fry.data)
iii = ii.astype(int)
jjj = jj.astype(int)
this_ji_list = list(zip(jjj, iii))
#print('Get ji list %0.3f seconds' % (time()-tt1))
#tt3 = time()
this_ji_cvec = jjj + iii * 1j
isin_vec = np.isin(this_ji_cvec, ji_cvec)
incount = isin_vec.sum()
#print('Search ji list alt %0.3f seconds' % (time()-tt3))
day_list.append(days[tt])
for iz in range(F.shape[0]):
# EXTRAPOLATION
fld = F[iz, :, :].squeeze()
# do the extrapolation using nearest neighbor
fldf = fld.copy() # initialize the "filled" field
if iz==0:
xyorig = np.array((X[~fld.mask],Y[~fld.mask])).T
xynew = np.array((X[fld.mask],Y[fld.mask])).T
a = cKDTree(xyorig).query(xynew)
aa = a[1]
fldf[fld.mask] = fld[~fld.mask][aa]
fx = fldf.data
# INTERPOLATION
if iz==0:
# get interpolants
xi0, xi1, xf = zfun.get_interpolant(x,X[0,:], extrap_nan=True)
yi0, yi1, yf = zfun.get_interpolant(y,Y[:,0], extrap_nan=True)
# bi linear interpolation
u00 = fx[yi0,xi0]
u10 = fx[yi1,xi0]
u01 = fx[yi0,xi1]
u11 = fx[yi1,xi1]
fi = (1-yf)*((1-xf)*u00 + xf*u01) + yf*((1-xf)*u10 + xf*u11)
ff = np.reshape(fi, x.shape)
fm = np.ma.masked_where(mask==False, ff)
ds1[name][tt,iz,:,:] = fm
print('Hour %d took %0.1f seconds' % (tt, time.time() - tt0))
tt += 1
ds0.close()
ds1.close()
# get coordinates
indir = indir0 + ocn + '/Data/'
coord_dict = pickle.load(open(indir + 'coord_dict.p', 'rb'))
lon = coord_dict['lon']
lat = coord_dict['lat']
# get fields
# note that when we use the "xfh" fields there should be no nans
xfh = pickle.load(open(indir + '/xfh' + mds + '.p', 'rb'))
zeta = xfh['ssh']
# get indices around certain regions, and their data
#
# (a) Mouth of Strait of Juan de Fuca
i0, i1, fr = zfun.get_interpolant(np.array([-124.7, -124.5]),
lon,
extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([48.4, 48.6]),
lat,
extrap_nan=False)
zeta_jdf = zeta[j0[0]:j1[1] + 1, i0[0]:i1[1] + 1]
#
# (b) Northern Strait of Georgia
i0, i1, fr = zfun.get_interpolant(np.array([-125.5, -124.5]),
lon,
extrap_nan=False)
j0, j1, fr = zfun.get_interpolant(np.array([50.25, 50.45]),
lat,
extrap_nan=False)
zeta_sog = zeta[j0[0]:j1[1] + 1, i0[0]:i1[1] + 1]
#