def _plot_segfile():
try:
if 'seg' in w.FIGS:
w.FIGS['seg'][0].remove()
w.FIGS['seg'][1][0].remove()
except:
pass
if hasattr(w, 'segfile'):
ptarray = np.hstack([[w.segfile[seg][:, 0], w.segfile[seg][:, 1]]
for seg in w.segfile]).T
tmparray = [
list(zip(w.segfile[seg][:, 0], w.segfile[seg][:, 1]))
for seg in w.segfile
]
w.linecollection = LC(tmparray,
color=w.config['seg']['linecolor'],
linewidth=float(w.config['seg']['linewidth']),
zorder=int(w.config['seg']['zorder']))
w.FIGS['seg'] = [
w.linecollection,
w.ax.plot(ptarray[:, 0],
ptarray[:, 1],
w.config['seg']['linecolor'] + '.',
markersize=float(w.config['seg']['linedotsize']),
zorder=int(w.config['seg']['zorder']))
]
w.ax.add_collection(w.linecollection)
w.figure.canvas.draw()
return
def speed_plot(i):
print(i)
f = plt.figure()
ax = plt.axes([.125, .1, .775, .8])
if coastflag == True:
plotcoast(ax,
filename='mid_nwatl6c_sjh_lr.nc',
color='k',
fcolor='0.75',
fill=True,
zorder=100)
triax = ax.tripcolor(data['trigrid'],
np.sqrt(data['ua'][i, :]**2 + data['va'][i, :]**2),
vmin=cmin,
vmax=cmax)
if np.shape(cages) != ():
lseg_t = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax.add_collection(lseg_t)
if vectorflag == True:
Q1 = ax.quiver(data['uvnodell'][vidx, 0],
data['uvnodell'][vidx, 1],
data['ua'][i, vidx],
data['va'][i, vidx],
angles='xy',
scale_units='xy',
scale=vector_scale,
zorder=100,
width=.0025)
if uniformvectorflag == True:
norm = np.sqrt(data['ua'][i, vidx]**2 + data['va'][i, vidx]**2)
Q1 = ax.quiver(data['uvnodell'][vidx, 0],
data['uvnodell'][vidx, 1],
np.divide(data['ua'][i, vidx], norm),
np.divide(data['va'][i, vidx], norm),
angles='xy',
scale_units='xy',
scale=vector_scale,
zorder=100,
width=.002,
color='k')
sand = np.argwhere(data['wet_cells'][i, :] == 0)
tmparray = [
list(
zip(data['nodell'][data['nv'][j, [0, 1, 2]], 0],
data['nodell'][data['nv'][j, [0, 1, 2]], 1])) for j in sand
]
lseg_sand = PC(tmparray, facecolor='0.75', edgecolor='0.75')
ax.add_collection(lseg_sand)
prettyplot_ll(ax, setregion=region, cblabel=r'Speed (ms$^{-1}$)', cb=triax)
f.savefig(savepath + grid + '_' + region['regionname'] + '_speed_' +
("%04d" % (i)) + '.png',
dpi=150)
plt.close(f)
def plot_example_trajectories(env,reps_to_plot, pcts_to_plot):
run_colors = [LINCLAB_COLS['red'],LINCLAB_COLS['blue'],LINCLAB_COLS['green']]
run_labels = {'onehot': 'Unstructured','analytic successor':'Structured'}
start_locations = [105,114,285]
fig, ax = plt.subplots(len(reps_to_plot),len(pcts_to_plot))
for r,rep in enumerate(reps_to_plot):
for p, pct in enumerate(pcts_to_plot):
run_id = list(gb.get_group((env_name,rep,cache_limits[env_name][pct])))[-1]
print(rep, pct, run_id)
with open(f'../../Data/results/{run_id}_data.p', 'rb') as f:
data = pickle.load(f)
rep_dict = {'analytic successor': sr, 'onehot':onehot}
state_reps, _, __, ___ = rep_dict[rep](env)
if rep == 'analytic successor':
for s1 in env.obstacle:
state_reps.pop(s1)
all_the_eps = data['ec_dicts']
rwd_colrow=(14,14)
rect = plt.Rectangle(rwd_colrow, 1, 1, facecolor='gray', alpha=0.3,edgecolor='k')
ax[r,p].pcolor(env.grid,cmap='bone_r',edgecolors='k', linewidths=0.1)
ax[r,p].axis(xmin=0, xmax=20, ymin=0,ymax=20)
ax[r,p].set_aspect('equal')
ax[r,p].add_patch(rect)
ax[r,p].set_yticks([])
ax[r,p].get_xaxis().set_visible(False)
#ax[r,p].get_yaxis().set_visible(False)
ax[r,p].invert_yaxis()
for xx in range(len(run_colors)):
trajectory = sample_from_ec_pol(state_reps,all_the_eps[-1],start_state=start_locations[xx])
print(trajectory)
lines = []
for i in range(len(trajectory)-1):
# need to reverse r-c from state2d to be x-y for line collection
state = trajectory[i]
next_state = trajectory[i+1]
lines.append([(state[1]+0.5,state[0]+0.5),(next_state[1]+0.5,next_state[0]+0.5)])
ax[r,p].add_patch(plt.Circle(lines[0][0],radius=0.3,color=run_colors[xx]))
lc = LC(lines, colors=run_colors[xx], linestyle="--",linewidths=0.85,alpha=0.75)
ax[r,p].add_collection(lc)
ax[r,0].set_ylabel(f"{run_labels[rep]}")
plt.savefig('../figures/CH2/example_trajectories1.svg')
plt.show()
def res_plot(i):
print(i)
f = plt.figure()
ax = plt.axes([.125, .1, .775, .8])
triax = ax.tripcolor(data['trigrid'],
np.sqrt(resu[:, i]**2 + resv[:, i]**2),
vmin=cmin,
vmax=cmax)
if coastflag == True:
plotcoast(ax,
filename='pacific_harbour.nc',
color='None',
fcolor='darkgreen',
fill=True)
if np.shape(cages) != ():
lseg_t = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax.add_collection(lseg_t)
if vectorflag == True:
Q1 = ax.quiver(data['uvnodell'][vidx, 0],
data['uvnodell'][vidx, 1],
resu[vidx, i],
resv[vidx, i],
angles='xy',
scale_units='xy',
scale=vector_scale,
zorder=100,
width=.0025)
if uniformvectorflag == True:
norm = np.sqrt(resu[vidx, i]**2 + resv[vidx, i]**2)
Q1 = ax.quiver(data['uvnodell'][vidx, 0],
data['uvnodell'][vidx, 1],
np.divide(resu[vidx, i], norm),
np.divide(resv[vidx, i], norm),
angles='xy',
scale_units='xy',
scale=vector_scale,
zorder=100,
width=.002,
color='k')
prettyplot_ll(ax,
setregion=region,
cblabel=r'Residual (ms$^{-1}$)',
cb=triax)
f.savefig(savepath + grid + '_' + region['regionname'] + '_residual_' +
("%04d" % (i)) + '.png',
dpi=150)
plt.close(f)
def myplot(i):
print i
f = plt.figure()
ax = plt.axes([.125, .1, .775, .8])
speed1 = np.sqrt(data1['ua'][i, :]**2 + data1['va'][i, :]**2) + .01
speed2 = np.sqrt(data2['ua'][i, :]**2 + data2['va'][i, :]**2) + .01
divco = np.divide((speed2 - speed1), speed1)
divco[np.isnan(divco)] = 0
triax = ax.tripcolor(data1['trigrid'],
divco,
vmin=cmin,
vmax=cmax,
cmap=mpl.cm.seismic)
#plotcoast(ax,color='k',fill=True)
if cages != None:
lseg_t = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax.add_collection(lseg_t)
vidx = equal_vectors(data2, region, 300)
Q1 = ax.quiver(data2['uvnodell'][vidx, 0],
data2['uvnodell'][vidx, 1],
data2['ua'][i, vidx],
data2['va'][i, vidx],
angles='xy',
scale_units='xy',
scale=50,
zorder=10)
aqk1 = ax.quiverkey(Q1,
.1,
.95,
0.5,
r'0.5 m/s',
labelpos='E',
fontproperties={'size': 10})
aqk1.set_zorder(30)
prettyplot_ll(ax,
setregion=region,
cblabel=r'Speed Rel Diff. (ms$^{-1}$)',
cb=triax)
f.savefig(savepath + grid + '_' + regionname + '_speed_reldiff_' +
("%04d" % (i)) + '.png',
dpi=150)
plt.close(f)
def load_coastline():
"""
Load and plot a coastline.
"""
filename = ()
filename = askopenfilename(initialdir=w.init_dir)
if filename != ():
w.coastline = ut.load_coastline(filename=filename)
w.TF['coast'] = True
w.CBVar['coast'].set(1)
#Reduce number of segments - helps but not enough
#axis=[-68.5,-64,44,46.5]
#yep=[]
#def isin(axis,line):
#return ((line[:,0]>=axis[0]) & (line[:,0]<=axis[1]) & (line[:,1]>=axis[2]) & (line[:,1]<=axis[3])).sum()
#for i,line in enumerate(w.coastline):
#if isin(axis,np.array(line)):
#yep+=[i]
if w.config['coast']['fill'] == 'True':
w.FIGS['coast'] = PC(w.coastline,
facecolor=w.config['coast']['facecolor'],
edgecolor=w.config['coast']['edgecolor'],
linewidths=float(
w.config['coast']['linewidth']),
zorder=int(w.config['coast']['zorder']))
else:
w.FIGS['coast'] = LC(w.coastline,
color=w.config['coast']['edgecolor'],
linewidths=float(
w.config['coast']['linewidth']),
zorder=int(w.config['coast']['zorder']))
w.ax.add_collection(w.FIGS['coast'])
w.figure.canvas.draw()
return
ax1.set_ylabel('')
for label in cb2.ax.get_xticklabels():
label.set_rotation(90)
else:
ax0ca = f.add_axes(
[ax0bb[0] + ax0bb[2] + .025, ax0bb[1], .025, ax0bb[3]])
ax1ca = f.add_axes(
[ax1bb[0] + ax1bb[2] + .025, ax1bb[1], .025, ax1bb[3]])
cb = plt.colorbar(axtri0, cax=ax0ca)
cb.set_label(r'Asymmetry', fontsize=8)
cb2 = plt.colorbar(axtri1, cax=ax1ca)
cb2.set_label(r'Asymmetry', fontsize=8)
plotcoast(ax0, filename='pacific.nc', color='k')
lseg0 = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax0.add_collection(lseg0)
plotcoast(ax1, filename='pacific.nc', color='k')
lseg1 = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax1.add_collection(lseg1)
ax0.text(
ABC[0], ABC[1], "A", transform=ax0.transAxes
) #,bbox={'facecolor':'white','edgecolor':'None', 'alpha':1, 'pad':3},zorder=31)
ax1.text(
ABC[0], ABC[1], "B", transform=ax1.transAxes
) #,bbox={'facecolor':'white','edgecolor':'None', 'alpha':1, 'pad':3},zorder=31)
f.savefig(savepath + grid + '_' + name + '_' + name2 + '_' + regionname +
'_ebbfld_asymmetry.png',
region = regions(regionname)
region = expand_region(region, dist=[5000, 6000], shift=[0, 6000])
region = regionll2xy(data, region)
nidx = get_nodes(data, region)
f = plt.figure()
ax = f.add_axes([.125, .1, .775, .8])
ax.triplot(data['trigrid'], lw=.5)
#clims=np.percentile(data['h'][nidx],[1,99])
#trip=ax.tripcolor(data['trigrid'],data['h'],vmin=clims[0],vmax=clims[1])
#prettyplot_ll(ax,setregion=region,cb=trip,cblabel='Depth (m)',grid=True)
plotcoast(ax, color='k', fill=True)
if np.shape(cages) != ():
lseg_t = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax.add_collection(lseg_t)
ax.axis(region['region'])
#f.savefig(savepath + grid + '_' + regionname+'_current_variance_magnitude_ratio.png',dpi=600)
vec = f.ginput(n=-1, timeout=-1)
plt.close(f)
print(len(vec))
if len(vec) == 0:
sys.exit("No points selected.")
locs = np.zeros((len(vec), runtime + 1, 2)) + np.nan
proj = gridproj(grid)
if forward:
def cross_shore_transect_2d(grid, name, region, vec, npt):
data = dt.loadnc('runs/' + grid + '/' + name + '/output/',
singlename=grid + '_0001.nc')
print('done load')
data = dt.ncdatasort(data, trifinder=True)
print('done sort')
cages = gt.loadcage('runs/' + grid + '/' + name + '/input/' + grid +
'_cage.dat')
if np.shape(cages) != ():
tmparray = [
list(
zip(data['nodell'][data['nv'][i, [0, 1, 2, 0]], 0],
data['nodell'][data['nv'][i, [0, 1, 2, 0]], 1]))
for i in cages
]
color = 'g'
lw = .2
ls = 'solid'
vectorstart = np.array(vec[0])
vectorend = np.array(vec[1])
vectorx = np.array([vectorstart[0], vectorend[0]])
vectory = np.array([vectorstart[1], vectorend[1]])
snv = (vectorend - vectorstart) / np.linalg.norm(vectorend - vectorstart)
xi = np.linspace(vectorstart[0], vectorend[0], npt)
yi = np.linspace(vectorstart[1], vectorend[1], npt)
us = data['ua'].shape
savepath = 'data/cross_shore_transect/'
if not os.path.exists(savepath): os.makedirs(savepath)
plotpath = 'figures/png/' + grid + '_2d/cross_shore_transect/'
if not os.path.exists(plotpath): os.makedirs(plotpath)
nidx = dt.get_nodes(data, region)
f = plt.figure()
ax = f.add_axes([.125, .1, .775, .8])
triax = ax.tripcolor(data['trigrid'],
data['h'],
vmin=data['h'][nidx].min(),
vmax=data['h'][nidx].max())
ax.plot(xi, yi, 'k', lw=3)
if np.shape(cages) != ():
lseg_t = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
coast = ax.add_collection(lseg_t)
coast.set_zorder(30)
pt.prettyplot_ll(ax, setregion=region, cb=triax, cblabel=r'Depth (m)')
f.savefig(plotpath + name + '_' + ('%f' % vectorx[0]) + '_' +
('%f' % vectorx[1]) + '_' + ('%f' % vectory[0]) + '_' +
('%f' % vectory[1]) + '_' + ('%d' % len(xi)) +
'_line_location.png',
dpi=600)
plt.close(f)
fillarray_u = np.empty((us[0], npt))
fillarray_v = np.empty((us[0], npt))
fillalong = np.empty((us[0], npt))
fillcross = np.empty((us[0], npt))
dist = np.empty((npt, ))
h = np.empty((npt, ))
print('interp uvw on path')
for i in range(0, len(xi)):
print(i)
fillarray_u[:, i] = interpE_at_loc(data, 'ua', [xi[i], yi[i]])
fillarray_v[:, i] = interpE_at_loc(data, 'va', [xi[i], yi[i]])
h[i] = interpN_at_loc(data, 'h', [xi[i], yi[i]])
print('Calc along path current')
for i in range(0, len(xi)):
print(i)
inner = np.inner(
np.vstack([fillarray_u[:, i], fillarray_v[:, i]]).T, snv)
along = np.vstack([inner * snv[0], inner * snv[1]]).T
tmpa = np.multiply(np.sign(np.arctan2(along[:, 1], along[:, 0])),
np.linalg.norm(along, axis=1))
fillalong[:, i] = tmpa
cross = np.vstack([fillarray_u[:, i], fillarray_v[:, i]]).T - along
tmpc = np.multiply(np.sign(np.arctan2(cross[:, 1], cross[:, 0])),
np.linalg.norm(cross, axis=1))
fillcross[:, i] = tmpc
dist[i] = (sw.dist([vectorstart[1], yi[i]], [vectorstart[0], xi[i]],
'km'))[0] * 1000
if np.shape(cages) != ():
incage = np.zeros((len(xi), ))
host = data['trigrid'].get_trifinder().__call__(xi, yi)
incage[np.in1d(host, cages)] = 1
savedic = {}
savedic['u'] = fillarray_u
savedic['v'] = fillarray_v
savedic['along'] = fillalong
savedic['cross'] = fillcross
savedic['distance'] = dist
savedic['h'] = h
savedic['lon'] = xi
savedic['lat'] = yi
if np.shape(cages) != ():
savedic['incage'] = incage
np.save(
savepath + grid + '_' + name + '_' + ('%f' % vectorx[0]) + '_' +
('%f' % vectorx[1]) + '_' + ('%f' % vectory[0]) + '_' +
('%f' % vectory[1]) + '_' + ('%d' % len(xi)) + '_2d.npy', savedic)
sio.savemat(savepath + 'matfiles/' + grid + '_' + name + '_' +
('%f' % vectorx[0]) + '_' + ('%f' % vectorx[1]) + '_' +
('%f' % vectory[0]) + '_' + ('%f' % vectory[1]) + '_' +
('%d' % len(xi)) + '_2d.mat',
mdict=savedic)
def res_plot(i):
print(i)
f = plt.figure()
ax = plt.axes([.125, .1, .775, .8])
triax = ax.tripcolor(data['trigrid'],
np.sqrt(resu[:, i]**2 + resv[:, i]**2),
vmin=cmin,
vmax=cmax)
if coastflag == True:
plotcoast(ax,
filename='pacific_harbour.nc',
color='None',
fcolor='darkgreen',
fill=True)
if np.shape(cages) != ():
lseg_t = LC(tmparray, linewidths=lw, linestyles=ls, color=color)
ax.add_collection(lseg_t)
if vectorflag == True:
Q1 = ax.quiver(data['uvnodell'][vidx, 0],
data['uvnodell'][vidx, 1],
resu[vidx, i],
resv[vidx, i],
angles='xy',
scale_units='xy',
scale=vector_scale,
zorder=100,
width=.0025)
if uniformvectorflag == True:
norm = np.sqrt(resu[vidx, i]**2 + resv[vidx, i]**2)
Q1 = ax.quiver(data['uvnodell'][vidx, 0],
data['uvnodell'][vidx, 1],
np.divide(resu[vidx, i], norm),
np.divide(resv[vidx, i], norm),
angles='xy',
scale_units='xy',
scale=vector_scale,
zorder=100,
width=.002,
color='k')
prettyplot_ll(ax,
setregion=region,
cblabel=r'Residual (ms$^{-1}$)',
cb=triax)
#ax1=plt.axes([.125,.675,.675,.2])
ax1 = plt.axes([.125, .545, .675, .2])
ax1.plot(data['time'][starttime:i] - data['time'][starttime],
data['zeta'][starttime:i, nidx[zetanode]])
ax1.set_ylabel(r'Elevation (m)')
ax1.set_xlabel(r'Time (days)')
ax1.xaxis.set_tick_params(labeltop='on', labelbottom='off')
ax1.axis([
data['time'][starttime] - data['time'][starttime],
data['time'][endtime] - data['time'][starttime], ymin, ymax
])
_formatter = mpl.ticker.ScalarFormatter(useOffset=False)
ax1.yaxis.set_major_formatter(_formatter)
ax1.xaxis.set_major_formatter(_formatter)
ax1.xaxis.set_label_coords(0.5, 1.4)
f.savefig(savepath + grid + '_' + region['regionname'] + '_residual_' +
("%04d" % (i)) + '.png',
dpi=150)
plt.close(f)