import numpy as np
import scipy.io as io
from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML, display
import matplotlib
import datetime
data = io.loadmat('./trajectory_100.mat')
lat, lon, z, t = data['lat_t'], data['lon_t'], data['psizt'], data['times']
app = data['app_solver']
nx,ny,n_depths = app['Np_vec'][0,0][0,:]
N = int(nx)*ny
n = nx
simstart = datetime.datetime(2018,5,25)
t_mask = ~(t[:,0] == 0)
t_mask[0] = True
lat,lon,z,t = lat[:,t_mask],lon[:,t_mask],z[:,t_mask],t[t_mask,:]
minlat = np.nanmin(lat)
maxlat = np.nanmax(lat)
minlon = np.nanmin(lon)
maxlon = np.nanmax(lon)
latspan = np.nanmax(lat) - np.nanmin(lat)
lonspan = np.nanmax(lon) - np.nanmin(lon)
scale = 0.1
nanmask = np.isnan((abs(lat[:,-1] - lat[:,0])+abs(lon[:,-1] - lon[:,0])+abs(z[:,-1] - z[:,0])))
mask = (abs(lat[:,-1] - lat[:,0])+abs(lon[:,-1] - lon[:,0])+abs(z[:,-1] - z[:,0]))>1e-3
mask_land = np.logical_or(mask,nanmask)
sparcity = 5
mask_sparce = np.zeros_like(lat[:,0],dtype=bool)
mask_sparce = mask_sparce.reshape([n_depths,nx,ny])
mask_sparce[:,::sparcity,::sparcity] = True
fig, axs = plt.subplots(n_depths//2,2,figsize = (20,14));
dz = z[...,-1]
cmap = matplotlib.cm.get_cmap('inferno')
normalize = matplotlib.colors.Normalize(vmin=np.nanmin(dz), vmax=np.nanmax(dz))
colors = [cmap(normalize(dz[i])) for i in range(dz.shape[0])]
fig.subplots_adjust(bottom=0.1,top=0.94);
fig.suptitle("$z$ Flowmap",fontsize=15)
p0 = axs[-1,0].get_position().get_points().flatten()
p1 = axs[-1,1].get_position().get_points().flatten()
ax_cbar = fig.add_axes([p0[0], 0.05, p1[2]-p0[0], 0.02])
cbar = matplotlib.colorbar.ColorbarBase(ax_cbar, cmap=cmap, norm=normalize, orientation='horizontal')
for d,ax in zip(np.arange(n_depths),axs.flatten()):
m = Basemap(projection='cyl',llcrnrlat=minlat-scale*latspan,urcrnrlat=maxlat+scale*latspan,\
llcrnrlon=minlon-scale*lonspan,urcrnrlon=maxlon+scale*lonspan,resolution='i',ax=ax)
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
x,y = m(lon[d*N:(d+1)*N,:],lat[d*N:(d+1)*N,:])
dz_d = dz[d*N:(d+1)*N]
x,y = x[:,0],y[:,0]
ax.set_title("Depth = {:.2f}m".format(z[d*N,0]))
X,Y,dZ_d = [var.reshape([n,n])for var in [x,y,dz_d]]
c = m.contourf(X.T,Y.T,dZ_d.T,100,vmin=np.nanmin(dz), vmax=np.nanmax(dz),cmap=cmap);
fig.savefig("flowmap.png")
mask_sparce = mask_sparce.ravel()
net_mask = np.logical_and(mask_sparce,mask_land)
#m = Basemap(projection='cyl',resolution='c')
fig, axs = plt.subplots(n_depths//2,2,figsize = (20,13));
for d,ax in zip(np.arange(n_depths),axs.flatten()):
m = Basemap(projection='cyl',llcrnrlat=minlat-scale*latspan,urcrnrlat=maxlat+scale*latspan,\
llcrnrlon=minlon-scale*lonspan,urcrnrlon=maxlon+scale*lonspan,resolution='i',ax=ax)
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
m.drawmapboundary(fill_color='aqua')
#m.shadedrelief()
x,y = m(lon[d*N:(d+1)*N],lat[d*N:(d+1)*N])
mask_local = net_mask[d*N:(d+1)*N]
# draw parallels and meridians.
#m.drawparallels(np.arange(-90.,91.,30.))
#m.drawmeridians(np.arange(-180.,181.,60.))
for particle in range(N):
if mask_local[particle]:
m.plot(x[particle,:],y[particle,:])
m.plot(x[particle,0],y[particle,0],'k.')
ax.set_title("Initial Depth = {:.2f}m".format(z[d*N,0]))
#plt.title("Equidistant Cylindrical Projection")
def plot_around_center(center_lat,center_lon,axs,d=1,lon_rad=0.2,lat_rad=0.2):
lon_gyre, lat_gyre, z_gyre = lon[d*N:(d+1)*N],lat[d*N:(d+1)*N],z[d*N:(d+1)*N]
mask_gyre = ((lon_gyre[:,0]-center_lon)/lon_rad)**2 + ((lat_gyre[:,0]-center_lat)/lat_rad)**2 < 1
lon_gyre, lat_gyre, z_gyre = lon_gyre[mask_gyre,:], lat_gyre[mask_gyre,:], z_gyre[mask_gyre,:]
mask_vert = (abs(lon_gyre[:,0] - center_lon) == np.min(abs(lon_gyre[:,0] - center_lon)))
mask_horz = (abs(lat_gyre[:,0] - center_lat) == np.min(abs(lat_gyre[:,0] - center_lat)))
#m = Basemap(projection='cyl',resolution='c')
ax = axs[0]
m = Basemap(projection='cyl',llcrnrlat=minlat-scale*latspan,urcrnrlat=maxlat+scale*latspan,\
llcrnrlon=minlon-scale*lonspan,urcrnrlon=maxlon+scale*lonspan,resolution='i',ax=ax)
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
m.drawmapboundary(fill_color='aqua')
#m.shadedrelief()
x,y = m(center_lon,center_lat)
m.plot(x, y, 'rx')
reflats = np.linspace(minlat,maxlat,5)
reflons = np.linspace(minlon,maxlon,5)
m.drawparallels(reflats,labels=np.ones_like(reflats,dtype=bool))
m.drawmeridians(reflons,labels=np.ones_like(reflons,dtype=bool))
x,y = m(lon_gyre, lat_gyre)
#indices = np.random.randint(0,x.shape[0],[10])
indices = np.arange(x.shape[0])
m.plot(x[:,0],y[:,0],'r.')
m.plot(x[mask_vert,0],y[mask_vert,0],'b.')
m.plot(x[mask_horz,0],y[mask_horz,0],'b.')
ax.set_title("Initial Depth = {:.2f}m".format(z[d*N,0]),y=1.1);
axs[3].set_title("Lat vs $z$")
axs[3].plot(lat_gyre[mask_vert,:].T,z_gyre[mask_vert,:].T,alpha=0.7)
axs[3].set_prop_cycle(None)
for i in range(len(lat_gyre[mask_vert,0])):
axs[3].plot((lat_gyre[mask_vert,0].T)[i],(z_gyre[mask_vert,0].T)[i],'o')
axs[3].set_xlabel("Latitude")
axs[2].set_title("Lon vs $z$")
axs[2].plot(lon_gyre[mask_horz,:].T,z_gyre[mask_horz,:].T,alpha=0.7)
axs[2].set_prop_cycle(None)
for i in range(len(lat_gyre[mask_horz,0])):
axs[2].plot((lon_gyre[mask_horz,0].T)[i],(z_gyre[mask_horz,0].T)[i],'o')
axs[2].set_xlabel("Longitude")
axs[1].set_title("Mean $z$ vs $t$")
axs[1].plot(z_gyre[indices,:].T,alpha=0.2)
axs[1].plot(np.mean(z_gyre[indices,:].T,axis=1),'k')
axs[1].set_xlabel("Time")
for ax in axs[1:]:
ax.set_ylim(ax.get_ylim()[::-1])
ax.set_ylabel("Depth")
return axs
center_lat = minlat + 0.33*latspan
center_lon = minlon + (1/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=0);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (2/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=0);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (3/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=0);
plt.tight_layout()
plt.show()
d=1
center_lat = minlat + 0.33*latspan
center_lon = minlon + (1/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (2/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (3/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
d=2
center_lat = minlat + 0.33*latspan
center_lon = minlon + (1/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (2/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (3/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
d=3
center_lat = minlat + 0.33*latspan
center_lon = minlon + (1/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (2/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
center_lat = minlat + 0.33*latspan
center_lon = minlon + (3/5)*lonspan
fig, axs = plt.subplots(4,1,figsize=(10,20));
axs = plot_around_center(center_lat,center_lon,axs,d=d);
plt.tight_layout()
plt.show()
cmap = matplotlib.cm.get_cmap('inferno')
normalize = matplotlib.colors.Normalize(vmin=np.nanmin(z), vmax=np.nanmax(z))
colors = [[cmap(normalize(z[i,j]))for j in range(z.shape[1])] for i in range(z.shape[0])]
# fig, ax = plt.subplots(figsize=(10,10))
# ax.scatter(x, y, color=colors)
# # Optionally add a colorbar
# cax, _ = matplotlib.colorbar.make_axes(ax)
# cbar = matplotlib.colorbar.ColorbarBase(cax, cmap=cmap, norm=normalize)
fig, axs = plt.subplots(n_depths//2,2,figsize = (20,14));
fig.subplots_adjust(bottom=0.1,top=0.94);
p0 = axs[-1,0].get_position().get_points().flatten()
p1 = axs[-1,1].get_position().get_points().flatten()
ax_cbar = fig.add_axes([p0[0], 0.07, p1[2]-p0[0], 0.02])
cbar = matplotlib.colorbar.ColorbarBase(ax_cbar, cmap=cmap, norm=normalize, orientation='horizontal')
bmaps = []
points = []
lines = []
for d,ax in zip(np.arange(n_depths),axs.flatten()):
bmaps.append(Basemap(projection='cyl',llcrnrlat=minlat-scale*latspan,urcrnrlat=maxlat+scale*latspan,\
llcrnrlon=minlon-scale*lonspan,urcrnrlon=maxlon+scale*lonspan,resolution='i',ax=ax))
m = bmaps[-1]
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
if d == 0:
x,y = m(lon,lat)
#m.drawmapboundary(fill_color='aqua')
for particle in range(N):
points.append(m.plot(x[d*N + particle,0],y[d*N+particle,0],'.',color=colors[d*N+particle][0])[0])
lines.append(m.plot(x[d*N + particle,0],y[d*N+particle,0],'k-',alpha=0.3)[0])
ax.set_title("Initial Depth = {:.2f}m".format(z[d*N,0]))
print("Backgrounds generated")
def init():
#point.set_data([], [])
return points+lines
# animation function. This is called sequentially
def animate(i):
for j,point in enumerate(points):
point.set_data(x[j,i], y[j,i])
point.set_color(colors[j][i])
lines[j].set_data(x[j,:i], y[j,:i])
fig.suptitle("T = {:.2f}".format(t[i,0]),fontsize=15)
return points+lines
# call the animator. blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(plt.gcf(), animate, init_func=init,
frames=298, interval=30, blit=True)
HTML(anim.to_html5_video())
cmap = matplotlib.cm.get_cmap('inferno')
normalize = matplotlib.colors.Normalize(vmin=np.nanmin(z), vmax=np.nanmax(z))
mask = np.logical_and(mask_sparce,mask_land)
#mask[:] = False
#mask[:] = True
for d in range(n_depths):
fig, ax = plt.subplots(dpi=200);
fig.subplots_adjust(bottom=0.1,top=0.94);
p0 = ax.get_position().get_points().flatten()
p1 = ax.get_position().get_points().flatten()
ax_cbar = fig.add_axes([p0[0], 0.07, p1[2]-p0[0], 0.02])
cbar = matplotlib.colorbar.ColorbarBase(ax_cbar, cmap=cmap, norm=normalize, orientation='horizontal')
m = Basemap(projection='cyl',llcrnrlat=minlat-scale*latspan,urcrnrlat=maxlat+scale*latspan,\
llcrnrlon=minlon-scale*lonspan,urcrnrlon=maxlon+scale*lonspan,resolution='h',ax=ax)
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
#m.shadedrelief()
mask_local = mask[d*N:(d+1)*N]
z_local = z[d*N:(d+1)*N,:]
z_local = z_local[mask_local,:]
x_local,y_local = m(lon[d*N:(d+1)*N,:],lat[d*N:(d+1)*N,:])
x_local,y_local = x_local[mask_local,:],y_local[mask_local,:]
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
points = []
lines = []
#m.drawmapboundary(fill_color='aqua')
for particle in range(x_local.shape[0]):
points.append(m.plot(x_local[particle,0],y_local[particle,0],'.',\
color=cmap(normalize(z_local[particle,0])))[0])
lines.append(m.plot(x_local[particle,0],y_local[particle,0],'k-',alpha=0.3)[0])
ax.set_title("Initial Depth = {:.2f}m".format(z[d*N,0]))
#print("Backgrounds generated")
def init():
#point.set_data([], [])
return points+lines
# animation function. This is called sequentially
def animate(i):
for j in range(x_local.shape[0]):
points[j].set_data(x_local[j,i], y_local[j,i])
points[j].set_color(cmap(normalize(z_local[j,i])))
lines[j].set_data(x_local[j,:i], y_local[j,:i])
timestring = simstart + datetime.timedelta(seconds=t[i,0])
fig.suptitle(timestring.strftime(format="%d %b %y, %H:%M:%S"),fontsize=15)
return points+lines
# call the animator. blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=t.shape[0], interval=30, blit=True)
display(HTML(anim.to_html5_video()))
#anim.save("d_100_res_{}.mp4".format(int(z[d*N,0])),dpi=200)
cmap = matplotlib.cm.get_cmap('viridis')
normalize = matplotlib.colors.Normalize(vmin=np.nanmin(z), vmax=np.nanmax(z))
d=0
fig, ax = plt.subplots();
fig.subplots_adjust(bottom=0.1,top=0.94);
p0 = ax.get_position().get_points().flatten()
p1 = ax.get_position().get_points().flatten()
ax_cbar = fig.add_axes([p0[0], 0.07, p1[2]-p0[0], 0.02])
cbar = matplotlib.colorbar.ColorbarBase(ax_cbar, cmap=cmap, norm=normalize, orientation='horizontal')
m = Basemap(projection='cyl',llcrnrlat=minlat-scale*latspan,urcrnrlat=maxlat+scale*latspan,\
llcrnrlon=minlon-scale*lonspan,urcrnrlon=maxlon+scale*lonspan,resolution='c',ax=ax)
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
#m.shadedrelief()
mask_local = mask[d*N:(d+1)*N]
z_local = z[d*N:(d+1)*N,:]
x_local,y_local = m(lon[d*N:(d+1)*N,:],lat[d*N:(d+1)*N,:])
#x_local,y_local = x_local[mask_local,:],y_local[mask_local,:]
m.drawcoastlines()
m.fillcontinents(color='coral',lake_color='aqua')
#m.drawmapboundary(fill_color='aqua')
for particle in range(x_local.shape[0]):
points.append(m.plot(x_local[particle,0],y_local[particle,0],'.',color=colors[particle][0])[0])
lines.append(m.plot(x_local[particle,0],y_local[particle,0],'k-',alpha=0.3)[0])
ax.set_title("Initial Depth = {:.2f}m".format(z[d*N,0]))
#print("Backgrounds generated")
def init():
for point in points:
point.set_data([], [])
for line in lines:
line.set_data([], [])
return points+lines
# animation function. This is called sequentially
def animate(i):
print(i)
for j in range(x_local.shape[0]):
if j==127 and i==100:
plt.plot(x_local[:,i], y_local[:,i])
plt.show()
points[j].set_data(x_local[j,i], y_local[j,i])
points[j].set_color(cmap(normalize(z_local[j,i])))
lines[j].set_data(x_local[j,:i], y_local[j,:i])
fig.suptitle("T = {:.2f}".format(t[i,0]),fontsize=15)
return points+lines
# call the animator. blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=20, interval=30, blit=True)
anim.save("d_{}.mp4".format(int(z_local[0,0])),dpi=200)