%matplotlib inline
from matplotlib import pylab
import matplotlib.pyplot as plt
import netCDF4 as NC
import numpy as np

# Load the data. Path name can be changed to look at different data.  
f = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/tide_fix/nov2009/all_forcing_crash/output.abort.nc','r');
#fU2 = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/nu40_crash/SalishSea_4h_20020915_20020921_grid_U.nc','r');
grid = NC.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea.nc','r')
bathy = grid.variables['Bathymetry'][:,:]


#Print some info about the variables. 
for dim in f.dimensions.values():
    print dim

f.variables.keys()

for var in f.variables.values():
    print var

u_vel = f.variables['vozocrtx']  #u currents
t=f.variables['time_counter']
ssh=f.variables['sossheig']
sal=f.variables['vosaline']
v_vel=f.variables['vomecrty'] # v currents
w_vel=f.variables['vovecrtz'] 
#bathy lat lon
b_lat = grid.variables['nav_lat']
b_lon = grid.variables['nav_lon']



#quick plot of the u field
ibad=220; jbad=184
depth=3

plt.figure(figsize=(14,8))
cmin=-2; cmax=2;
ax=[200,250,150,220]
#print time
tnow=t[0]/86400
print tnow

#plot U
U=u_vel[0,depth,:,:]
print U.min()
print U.max()
mu =U == 0
U= np.ma.array(U,mask=mu)
pylab.subplot(1,2,1)
pylab.pcolor(U,vmin=cmin,vmax=cmax)
#pylab.plot(ibad-1,jbad-1,marker='o') #marker for unstable point in the ocean.output file
pylab.colorbar()
#pylab.axis(ax)

#plot V
pylab.subplot(1,2,2)
V=v_vel[0,depth,:,:]
mu =V == 0
V= np.ma.array(V,mask=mu)
pylab.pcolor(V,vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.plot(ibad-1,jbad-1,marker='o') #marker for unstable point in the ocean.output file
#pylab.axis(ax)

print u_vel[0,3,jbad-1,ibad-1]
print v_vel[0,3,jbad-1,ibad-1]

#Vertical velocity
cmin=-0.005; cmax=0.005
depth=3
W=w_vel[0,depth,:,:]
print W.min()
print W.max()
mu =W == 0
W= np.ma.array(W,mask=mu)
pylab.pcolor(W,vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.axis(ax)

print w_vel[0,3,jbad-1,ibad-1]

# load the bathymetry to look for shallow depths in this region. 
plt.figure(figsize=(14,8))
cmin=0; cmax=5;
mu =bathy == 0
bathy= np.ma.array(bathy,mask=mu)
pylab.pcolor(bathy,vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.plot(ibad-1,jbad-1,marker='o') #marker for unstable point in the ocean.output file
pylab.axis(ax)




# look at sea surface height
plt.figure(figsize=(14,8))
mu =ssh == 0
ssh= np.ma.array(ssh,mask=mu)
pylab.pcolor(ssh[0,:,:],vmin=-4,vmax=4)
pylab.colorbar()
pylab.plot(ibad-1,jbad-1,marker='o') #marker for unstable point in the ocean.output file
#pylab.axis(ax)
print ssh.min()
print ssh.max()
print bathy[jbad-1,ibad-1]
print ssh[0,jbad-1,ibad-1]

# look at salinity
plt.figure(figsize=(14,8))
mu =sal == 0
sal= np.ma.array(sal,mask=mu)
pylab.pcolor(sal[0,depth,:,:],vmin=0,vmax=32)
pylab.colorbar()
pylab.plot(ibad-1,jbad-1,marker='o') #marker for unstable point in the ocean.output file
#pylab.axis(ax)
print sal[0,3,jbad-1,ibad-1]