In [58]:
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

Load Data

In [45]:
data = io.loadmat('./trajectory_100.mat')
In [46]:
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,:]

Set limits

In [47]:
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

Mask out particles on land

In [48]:
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)
/home/mdoshi/miniconda3/envs/calypsoenv/lib/python3.7/site-packages/ipykernel/__main__.py:2: RuntimeWarning: invalid value encountered in greater
  from ipykernel import kernelapp as app

Select only a sparce subset

In [49]:
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

Generate "flow map"

In [67]:
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);
In [42]:
fig.savefig("flowmap.png")
In [50]:
mask_sparce = mask_sparce.ravel()

Plot some trajectories

In [51]:
net_mask = np.logical_and(mask_sparce,mask_land)

Plot

In [52]:
#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")

Plot Vertical motion of gyre area

Setup

In [33]:
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

Select all points around "center"

Surface

In [35]:
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 = 75m

In [37]:
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 = 138m

In [38]:
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 = 200m

In [40]:
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()

Videos

All depths together and sparce

In [41]:
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)
---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
<ipython-input-41-76bf1c258ba8> in <module>
      1 cmap = matplotlib.cm.get_cmap('inferno')
      2 normalize = matplotlib.colors.Normalize(vmin=np.nanmin(z), vmax=np.nanmax(z))
----> 3 colors = [[cmap(normalize(z[i,j]))for j in range(z.shape[1])] for i in range(z.shape[0])]
      4 
      5 # fig, ax = plt.subplots(figsize=(10,10))

<ipython-input-41-76bf1c258ba8> in <listcomp>(.0)
      1 cmap = matplotlib.cm.get_cmap('inferno')
      2 normalize = matplotlib.colors.Normalize(vmin=np.nanmin(z), vmax=np.nanmax(z))
----> 3 colors = [[cmap(normalize(z[i,j]))for j in range(z.shape[1])] for i in range(z.shape[0])]
      4 
      5 # fig, ax = plt.subplots(figsize=(10,10))

<ipython-input-41-76bf1c258ba8> in <listcomp>(.0)
      1 cmap = matplotlib.cm.get_cmap('inferno')
      2 normalize = matplotlib.colors.Normalize(vmin=np.nanmin(z), vmax=np.nanmax(z))
----> 3 colors = [[cmap(normalize(z[i,j]))for j in range(z.shape[1])] for i in range(z.shape[0])]
      4 
      5 # fig, ax = plt.subplots(figsize=(10,10))

~/miniconda3/envs/calypsoenv/lib/python3.7/site-packages/matplotlib/colors.py in __call__(self, X, alpha, bytes)
    491 
    492         if xa.dtype.kind == "f":
--> 493             xa *= self.N
    494             # Negative values are out of range, but astype(int) would truncate
    495             # them towards zero.

KeyboardInterrupt: 
In [132]:
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())
Backgrounds generated
Out[132]:

One at a time

In [76]:
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)

ALL POINTS (never run)

In [48]:
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)
0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19