def get_casts(Ldir):
year = 2017
month = 1
# +++ load ecology CTD cast data +++
dir0 = Ldir['parent'] + 'ptools_data/ecology/'
# load processed station info and data
sta_df = pd.read_pickle(dir0 + 'sta_df.p')
# add Canadian data
dir1 = Ldir['parent'] + 'ptools_data/canada/'
# load processed station info and data
sta_df_ca = pd.read_pickle(dir1 + 'sta_df.p')
sta_df = pd.concat((sta_df, sta_df_ca), sort=False)
Casts = pd.read_pickle(dir0 + 'Casts_' + str(year) + '.p')
Casts_ca = pd.read_pickle(dir1 + 'Casts_' + str(year) + '.p')
Casts = pd.concat((Casts, Casts_ca), sort=False)
# limit the stations used, if desired
sta_list = [s for s in sta_df.index
] # if ('WPA' not in s) and ('GYS' not in s)]
# keep only certain columns
sta_df = sta_df.loc[sta_list, ['Max_Depth', 'Latitude', 'Longitude']]
#
# start a dict to store one cast per station (if it has data in the year)
Cast_dict = dict()
for station in sta_list:
casts = Casts[Casts['Station'] == station]
casts = casts.set_index('Date')
casts = casts.loc[:, ['Salinity', 'Temperature',
'Z']] # keep only selected columns
# identify a single cast by its date
alldates = casts.index
castdates = alldates.unique(
) # a short list of unique dates (1 per cast)
# get the CTD cast data for this station, in the nearest month
cdv = castdates.month.values # all the months with casts
if len(cdv) > 0:
# get the cast closest to the selected month
imo = zfun.find_nearest_ind(cdv, month)
new_mo = cdv[imo]
cast = casts[casts.index == castdates[imo]]
Cast = cast.set_index('Z') # reorganize so that the index is Z
Cast = Cast.dropna() # clean up
# store cast in a dict
Cast_dict[station] = Cast
# save the month, just so we know
sta_df.loc[station, 'Month'] = new_mo
print(' - Ofun_CTD.get_casts: including : ' + station +
' month=' + str(new_mo))
else:
print(' - Ofun_CTD.get_casts:' + station + ': no data')
# Cast_dict.keys() is the "official" list of stations to loop over
return Cast_dict, sta_df
def box_inds(lon0, lon1, lat0, lat1, grid_fn):
# label output dict
D = dict()
# get Datasets
ds = nc4.Dataset(grid_fn)
# gather some fields, for convenience
lon_rho = ds['lon_rho'][0, :].squeeze()
lat_rho = ds['lat_rho'][:, 0].squeeze()
lon_u = ds['lon_u'][0, :].squeeze()
lat_u = ds['lat_u'][:, 0].squeeze()
lon_v = ds['lon_v'][0, :].squeeze()
lat_v = ds['lat_v'][:, 0].squeeze()
# find rho index rel. to reference point, find u and v indeces rel. to rho point
# relation of u and v pts to rho pt determined from the definitions on the
# ROMS grid here: https://www.myroms.org/wiki/Numerical_Solution_Technique
lonr0 = zfun.find_nearest_ind(lon_rho, lon0)
D['lonr0'] = lonr0
D['lonu0'] = lonr0 - 1
D['lonv0'] = lonr0
lonr1 = zfun.find_nearest_ind(lon_rho, lon1)
D['lonr1'] = lonr1
D['lonu1'] = lonr1
D['lonv1'] = lonr1
latr0 = zfun.find_nearest_ind(lat_rho, lat0)
D['latr0'] = latr0
D['latu0'] = latr0
D['latv0'] = latr0 - 1
latr1 = zfun.find_nearest_ind(lat_rho, lat1)
D['latr1'] = latr1
D['latu1'] = latr1
D['latv1'] = latr1
ds.close()
return D
def get_coordinates(ds):
"""
Get lat, lon vectors, related indices, and z from the
Dataset for a hycom archive.
"""
# create z from the depth
depth = ds.variables['depth'][:]
# and pack bottom to top
z = -depth[::-1]
# get full coordinates (vectors for the plaid grid)
llon = ds.variables['lon'][:]
llat = ds.variables['lat'][:]
# find indices of a sub region
aa = get_extraction_limits()
llon = llon - 360 # convert from 0:360 to -360:0 format
i0 = zfun.find_nearest_ind(llon, aa[0])
i1 = zfun.find_nearest_ind(llon, aa[1])
j0 = zfun.find_nearest_ind(llat, aa[2])
j1 = zfun.find_nearest_ind(llat, aa[3])
# and just get the desired region (these are vectors, not matrices)
lon = llon[i0:i1]
lat = llat[j0:j1]
coords = dict()
coords['z'] = z
coords['lon'] = lon
coords['lat'] =lat
coords['i0'] = i0
coords['i1'] = i1
coords['j0'] = j0
coords['j1'] = j1
return coords
def get_coordinates(ds):
"""
Get lat, lon vectors, related indices, and z from the
Dataset for a hycom archive.
"""
# create z from the depth
depth = ds.variables['depth'][:]
# and pack bottom to top
z = -depth[::-1]
# get full coordinates (vectors for the plaid grid)
llon = ds.variables['lon'][:]
llat = ds.variables['lat'][:]
# find indices of a sub region
aa = get_extraction_limits()
llon = llon - 360 # convert from 0:360 to -360:0 format
i0 = zfun.find_nearest_ind(llon, aa[0])
i1 = zfun.find_nearest_ind(llon, aa[1])
j0 = zfun.find_nearest_ind(llat, aa[2])
j1 = zfun.find_nearest_ind(llat, aa[3])
# and just get the desired region (these are vectors, not matrices)
lon = llon[i0:i1]
lat = llat[j0:j1]
coords = dict()
coords['z'] = z
coords['lon'] = lon
coords['lat'] = lat
coords['i0'] = i0
coords['i1'] = i1
coords['j0'] = j0
coords['j1'] = j1
return coords
foo.createDimension('xi_rho', NC)
# add time data
vv = foo.createVariable(timename, float, (timename, ))
vv.units = 'seconds since 1970.01.01 UTC'
vv[:] = mod_time_vec
# add variable definition
vv = foo.createVariable(vn, float, (timename, 'eta_rho', 'xi_rho'))
vv.long_name = afun.longname_dict[vn]
vv.units = afun.units_dict[vn]
foo.close()
print('Initializing NetCDF files took %0.1f seconds' % (time.time() - tt0))
tt0 = time.time()
# find index to trim Eastern part of WRF fields
lon_max = lon[0, -1]
imax2 = zfun.find_nearest_ind(lon2[0, :], lon_max + .5)
lon2 = lon2[:, :imax2]
lat2 = lat2[:, :imax2]
if do_d3:
imax3 = zfun.find_nearest_ind(lon3[0, :], lon_max + .5)
lon3 = lon3[:, :imax3]
lat3 = lat3[:, :imax3]
if do_d4:
imax4 = zfun.find_nearest_ind(lon4[0, :], lon_max + .5)
lon4 = lon4[:, :imax4]
lat4 = lat4[:, :imax4]
# prepare coordinate arrays for interpolation
XY = np.array((lon.flatten(), lat.flatten())).T
XY2 = np.array((lon2.flatten(), lat2.flatten())).T
if do_d3:
XY3 = np.array((lon3.flatten(), lat3.flatten())).T
zeta = ds_avg['zeta'][:]
zr = zrfun.get_z(h, zeta, S, only_rho=True)
#read in all of the momentum terms of interest...
for key in sorted(list_choice):
if i == 0:
Du[key] = np.zeros((NT))
Dv[key] = np.zeros((NT))
u_key = 'u_' + key
v_key = 'v_' + key
#for u and v at the particle's position for that time step
indx_u = zfun.find_nearest_ind(lonu[0, :], tlon[i])
indy_u = zfun.find_nearest_ind(latu[:, 0], tlat[i])
indz_u = zfun.find_nearest_ind(
zr[:, indy_u,
indx_u + 1], tz[i]) #indx+1 because zr is on rho grid
Du[key][i] = ds_dia[u_key][0, indz_u, indy_u, indx_u]
indx_v = zfun.find_nearest_ind(lonv[0, :], tlon[i])
indy_v = zfun.find_nearest_ind(latv[:, 0], tlat[i])
indz_v = zfun.find_nearest_ind(
zr[:, indy_v + 1,
indx_v], tz[i]) #indy+1 because zr is on rho grid
Dv[key][i] = ds_dia[v_key][0, indz_v, indy_v, indx_v]
if np.mod(i, 100) == 0:
print('time step %i' % i)
hrs = len(d_list)
for hr in range(hrs):
ds_dia = nc.Dataset(d_list[hr])
ds_avg = nc.Dataset(a_list[hr])
ds_h0 = nc.Dataset(h_list[hr])
ds_h1 = nc.Dataset(h_list[hr + 1])
zeta = ds_avg['zeta'][:]
zr = zrfun.get_z(h, zeta, S, only_rho=True)
for particle in particlelist:
p = pdict[particle]
#calcualte direction based on coriolis
indx_u = zfun.find_nearest_ind(lonu[0, :], p['tlon'][tt])
indy_u = zfun.find_nearest_ind(latu[:, 0], p['tlat'][tt])
indz_u = zfun.find_nearest_ind(
zr[:, indy_u,
indx_u + 1], p['tz'][tt]) #indx+1 because zr is on rho grid
indx_v = zfun.find_nearest_ind(lonv[0, :], p['tlon'][tt])
indy_v = zfun.find_nearest_ind(latv[:, 0], p['tlat'][tt])
indz_v = zfun.find_nearest_ind(
zr[:, indy_v + 1,
indx_v], p['tz'][tt]) #indy+1 because zr is on rho grid
#build dist_from_mouth
x, y = zfun.ll2xy(p['tlon'][tt], p['tlat'][tt], reflon, reflat)
p['dist_from_mouth'][tt] = np.sqrt(x**2 + y**2)
p['dist_from_mouth'][tt] = np.sign(p['tlon'][tt]) * np.abs(
NT = len(fn_list)
lon = G['lon_rho']
lat = G['lat_rho']
mask = G['mask_rho']
xvec = lon[0, :].flatten()
yvec = lat[:, 0].flatten()
# automatically correct for moorings that are on land
# Note that this only checks rho points - what if a u or v point is masked?
ji_dict = {}
for sta in sta_dict.keys():
xy = sta_dict[sta]
slon = xy[0]
slat = xy[1]
i0 = zfun.find_nearest_ind(xvec, slon)
j0 = zfun.find_nearest_ind(yvec, slat)
# find indices of nearest good point
# Note: tha mask in G is Boolean, not 0, 1, but testing with 1 or 0 works.
if mask[j0, i0] == 1:
print(sta + ': point ok')
elif mask[j0, i0] == 0:
print(sta + ':point masked')
i0, j0 = mfun.get_ij_good(lon, lat, mask, xvec, yvec, i0, j0)
newlon = xvec[i0]
newlat = yvec[j0]
print(' Replacing (%0.3f,%0.3f) with (%0.3f,%0.3f)' %
(slon, slat, newlon, newlat))
sta_dict[sta] = (newlon, newlat)
ji_dict[sta] = (j0, i0)
def get_orig(Cast_dict, sta_df, X, Y, fld, lon, lat, zz, vn):
verbose = False
# make vectors or 1- or 2-column arrays (*) of the good points to feed to cKDTree
xyorig = np.array((X[~fld.mask],Y[~fld.mask])).T
fldorig = fld[~fld.mask]
#========================================================================
# +++ append good points from CTD data to our arrays (*) +++
goodcount = 0
for station in Cast_dict.keys():
Cast = Cast_dict[station]
cz = Cast.index.values
izc = zfun.find_nearest_ind(cz, zz)
# only take data from this cast if its bottom depth is at or above
# the chosen hycom level
czbot = -sta_df.loc[station,'Max_Depth']
if czbot <= zz:
# becasue we used find_nearest above we should always
# get data in the steps below
if vn == 't3d':
this_fld = Cast.iloc[izc]['Temperature']
elif vn == 's3d':
this_fld = Cast.iloc[izc]['Salinity']
# and store in sta_df (to align with lat, lon)
sta_df.loc[station,'fld'] = this_fld
goodcount += 1
else:
pass
if goodcount >= 1:
# drop stations that don't have T and s values at this depth
sta_df = sta_df.dropna()
# and for later convenience make a new list of stations
sta_list = list(sta_df.index)
# if we got any good points then append them
if verbose:
print(' - Ofun_CTD.get_orig: goodcount = %d, len(sta_df) = %d'
% (goodcount, len(sta_df)))
# append CTD values to the good points from HYCOM
x_sta = sta_df['Longitude'].values
y_sta = sta_df['Latitude'].values
xx_sta, yy_sta = zfun.ll2xy(x_sta, y_sta, lon.mean(), lat.mean())
xy_sta = np.stack((xx_sta,yy_sta), axis=1)
xyorig = np.concatenate((xyorig, xy_sta))
fld_arr = sta_df['fld'].values
fldorig = np.concatenate((fldorig, np.array(fld_arr,ndmin=1)))
else:
if verbose:
print(' - Ofun_CTD.get_orig: No points added')
return xyorig, fldorig
def get_layer(fn, NZ=-1, aa=[], print_info=False):
# function to extract and process fields from a history file
# returning a dict of arrays that can be passed to CO2SYS.m
# default is to get full surface field
ds = nc.Dataset(fn)
G, S, T = zrfun.get_basic_info(fn)
if len(aa) == 4:
# find indices that encompass region aa
i0 = zfun.find_nearest_ind(G['lon_rho'][0,:], aa[0]) - 1
i1 = zfun.find_nearest_ind(G['lon_rho'][0,:], aa[1]) + 2
j0 = zfun.find_nearest_ind(G['lat_rho'][:,0], aa[2]) - 1
j1 = zfun.find_nearest_ind(G['lat_rho'][:,0], aa[3]) + 2
else:
# full region
i0 = 0; j0 = 0
j1, i1 = G['lon_rho'].shape
i1 += 1; j1 += 1
plon = G['lon_psi'][j0:j1-1, i0:i1-1]
plat = G['lat_psi'][j0:j1-1, i0:i1-1]
# extract needed info from history file
v_dict = dict()
if print_info:
print('\nINPUT Variable Info:')
for vn in ['alkalinity', 'TIC', 'salt', 'temp','rho']:
v = ds[vn][0,NZ, j0:j1, i0:i1]
v = fillit(v)
v_dict[vn] = v
if print_info:
name = ds[vn].long_name
try:
units = ds[vn].units
except AttributeError:
units = ''
vmax = np.nanmax(v)
vmin = np.nanmin(v)
v_dict[vn] = v
print('%25s (%25s) max = %6.1f min = %6.1f' %
(name, units, vmax, vmin))
# create depth, pressure, and in situ temperature
h = ds['h'][j0:j1, i0:i1]
h = fillit(h)
lat = G['lat_rho'][j0:j1, i0:i1]
z_rho = zrfun.get_z(h, 0*h, S, only_rho=True)
depth = -z_rho[NZ, :, :].squeeze()
pres = sw.pres(depth, lat)
v_dict['pres'] = pres
# assume potential temperature is close enough
# temp = sw.ptmp(v_dict['salt'], v_dict['temp'], 0, v_dict['pres'])
# v_dict['temp'] = temp
# convert from umol/L to umol/kg using in situ dentity
v_dict['alkalinity'] = 1000 * v_dict['alkalinity'] / (v_dict['rho'] + 1000)
v_dict['TIC'] = 1000 * v_dict['TIC'] / (v_dict['rho'] + 1000)
# clean up
v_dict.pop('rho') # no longer needed, so don't pass to worker
ds.close()
return v_dict, plon, plat
fig=plt.figure(figsize=(9,7))
ax = plt.gca()
ax.axis(a_ll)
#plot background
sl = np.ma.masked_where(salt[45+i,:,:] > 40,salt[45+i,:,:])
ax.pcolormesh(lonr,latr,sl,vmax=31,vmin=25,cmap='rainbow',alpha=0.35)
#get and plot momentum forcings using quiver
Du = dict()
Dv = dict()
for key in sorted(DIA):
if key[0]=='u':
#all keys are 3D, helpfully
indx = zfun.find_nearest_ind(lonu[0,:],tlon_sm[i])
indy = zfun.find_nearest_ind(latu[:,0],tlat_sm[i])
indz = zfun.find_nearest_ind(z[45+i,:,indy,indx+1],tz_sm[i]) #indx+1 because zr is on rho grid
Du[key[2:]] = DIA[key][45+i,indz,indy,indx]
if key[0]=='v':
indx = zfun.find_nearest_ind(lonv[0,:],tlon_sm[i])
indy = zfun.find_nearest_ind(latv[:,0],tlat_sm[i])
indz = zfun.find_nearest_ind(z[45+i,:,indy+1,indx],tz_sm[i]) #indy+1 because zr is on rho grid
Dv[key[2:]] = DIA[key][45+i,indz,indy,indx]
Du['ageo'] = Du['prsgrd']+Du['cor']
Dv['ageo'] = Dv['prsgrd']+Dv['cor']
Du['adv'] = Du['hadv']+Du['vadv']
Dv['adv'] = Dv['hadv']+Dv['vadv']
Du['cap'] = Du['ageo']+Du['adv']
Dv['cap'] = Dv['ageo']+Dv['adv']
Du['balance'] = Du['cap'] + Du['vvisc'] - Du['accel']
G, S, T = zrfun.get_basic_info(oceanfna)
lonr = G['lon_rho'][:]
latr = G['lat_rho'][:]
maskr0 = G['mask_rho'][:]
H = G['h'][:]
maskr = np.tile(maskr0, (NZ, 1, 1))
# create cross sectional axes
latvec = latr[:, 0]
l_list = [45.6, 45.3, 45.0, 44.4]
llen = len(l_list)
l = np.zeros(len(l_list), dtype='int')
lonvec = np.zeros([len(l_list), np.shape(lonr)[1]])
for lat in range(len(l_list)):
l0 = zfun.find_nearest_ind(latvec, l_list[lat])
l[lat] = l0
lonvec[lat, :] = lonr[l0, :]
a_ll = [-1.0, 0.02, latr.min(), latr.max()]
vmin = 22.7
vmax = 23.9
for t in range(ndays):
#identify file names for that day
fn_list = os.listdir(dir0 + f_list[t])
fna = [x for x in fn_list if x[-11:-8] == 'avg']
#load in data for each hour and find the average
fna_list = [x for x in fn_list if x[6:9] == 'avg']
hrs = len(fna_list)
for hr in range(hrs):
oceanfna = dir0 + f_list[t] + '/' + fna_list[hr]
dsa = nc4.Dataset(oceanfna)
#read in variable slices from data
u0 = dsa['u'][:]
if tt == 0:
NT, NZ, NY, NX = np.shape(u0)
lonu = dsa['lon_u'][:]
l0 = zfun.find_nearest_ind(lonu[0, :], 0.3)
#calculate A
maskr = dsa['mask_rho']
masku = dsa['mask_u']
mu = np.tile(masku[:, l0], (NZ, 1))
h = dsa['h'][:, l0]
h[maskr[:, l0] == 0] = np.nan
pn = dsa['pn'][:]
dy = 1 / pn[:, l0]
dA = h * dy
A = np.nansum(dA)
H = np.nanmean(h)
Cd = dsa['rdrg2'][:]
#initialize matrix
oceanfnh1 = dir0 + f_list[i] + '/' + fnh[h + 1]
oceanfnd = dir0 + f_list[i] + '/' + fnd[h]
oceanfna = dir0 + f_list[i] + '/' + fna[h]
ds_his0 = nc.Dataset(oceanfnh0)
ds_his1 = nc.Dataset(oceanfnh1)
ds_dia = nc.Dataset(oceanfnd)
ds_avg = nc.Dataset(oceanfna)
if t == 0:
dt = ds_his1['ocean_time'][:] - ds_his0['ocean_time'][:] #dt
lon_var = ds_his0[lon_str]
lon_vec = lon_var[0, :]
lat_var = ds_his0[lat_str]
lat_vec = lat_var[:, 0]
x = zfun.find_nearest_ind(lon_vec, lonchoice)
y = zfun.find_nearest_ind(lat_vec, latchoice)
z = -1
# variable to reproduce
ot[t] = ds_avg['ocean_time'][0]
var_accel_old[t] = ds_dia[accel_str][0, z, y, x]
dvar = ds_his1[varchoice][0, z, y, x] - ds_his0[varchoice][0, z, y,
x] #du
delta0 = get_delta(oceanfnh0, pmpn, z, y, x)
delta1 = get_delta(oceanfnh1, pmpn, z, y, x)
ddelta = delta1 - delta0
#load in averaged variables
var = ds_avg[varchoice][0, z, y, x] #<u>
oceanfn = dir0 + fn
dso = nc4.Dataset(oceanfn)
if t == 0:
#get grid info from ocean
lonr = dso['lon_rho'][:]
latr = dso['lat_rho'][:]
h = dso['h'][:]
zeta = dso['zeta'][:]
maskr = dso['mask_rho'][:]
S = zrfun.get_basic_info(oceanfn, only_S=True)
zr = zrfun.get_z(h, zeta, S, only_rho=True)
NZ, NY, NX = np.shape(zr)
midlat = zfun.find_nearest_ind(latr[:, 0], 45.0)
bgsalt = dso['salt'][:]
surfsalt = bgsalt[0, -1, :, :].squeeze()
crosssalt = bgsalt[0, :, midlat, :].squeeze()
nt = t * 24
lon1 = lon0[nt, tma]
lat1 = lat0[nt, tma]
z1 = z0[nt, tma]
fig = plt.figure(figsize=(12, 4))
ax1 = fig.add_subplot(1, 3, 2)
a_ll = [lonr.min(), 1.2, latr.min(), latr.max()]
bgs = np.ma.masked_where(maskr == 0, surfsalt)
