#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from salishsea_tools import viz_tools, places, visualisations
import netCDF4 as nc # unless you prefer xarray
import datetime as dt
import os
import glob
import cmocean

get_ipython().run_line_magic('matplotlib', 'inline')

A note on indexing: python starts indexing with 0. Also, when you "slice" an array, you use the first index you want to return and the first index you want to leave out. Example:
# In[2]:


x=np.arange(0,10) # create a numpy array containing the numbers 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
print('x:',x)
y=x[4:8]  # put the 4-indexed (4) through 7-indexed (7) elements in y
print('y:',y)
print('last element:', x[-1])
print('all of x:',x[:])


# In[3]:


places.PLACES['S3']


# ### Load a file from the 201812 hindcast
A note on model results. We currently have two hindcasts available:
- 201812:  /results/SalishSea/nowcast-green.201812/
- 201905:  /results2/SalishSea/nowcast-green.201905/

These directories contain subdirectories for each day of the run, with results files inside those directories.
The file naming convention is as follows: \
SalishSea_timeResolution_startdate_enddate_filetype.nc \
1h indicates hourly average output (time dimension length 24)
1d indicates daily average output (time dimesnion length 1)
ptrc_T files contain biological variables like nitrate, phytoplankton, and zooplankton concentrations
grid_T files contain temperature (votemper) and salinity (vosaline)
# In[4]:


f=nc.Dataset('/results/SalishSea/nowcast-green.201812/01apr15/SalishSea_1h_20150401_20150401_ptrc_T.nc')


# In[5]:


f2=nc.Dataset('/results/SalishSea/nowcast-green.201812/25feb15/SalishSea_1h_20150225_20150225_ptrc_T.nc')


# In[6]:


print(f.variables.keys())


# In[7]:


fe3t=nc.Dataset('/results/SalishSea/nowcast-green.201812/01apr15/SalishSea_1h_20150401_20150401_carp_T.nc')


# In[8]:


fe3t_2=nc.Dataset('/results/SalishSea/nowcast-green.201812/25feb15/SalishSea_1h_20150225_20150225_carp_T.nc')


# In[9]:


print(fe3t.variables.keys())


# In[10]:


# return times as datetime objects:
torig=dt.datetime.strptime(f.variables['time_centered'].time_origin,'%Y-%m-%d %H:%M:%S')
print(torig)
times=np.array([torig + dt.timedelta(seconds=ii) for ii in f.variables['time_centered'][:]])


# In[11]:


times[0:3]


# In[12]:


# load model mesh
with nc.Dataset('/data/eolson/results/MEOPAR/NEMO-forcing-new/grid/mesh_mask201702.nc') as fm:
    print(fm.variables.keys())
    tmask=fm.variables['tmask'][:,:,:,:]
    navlon=fm.variables['nav_lon'][:,:]
    navlat=fm.variables['nav_lat'][:,:]


# ### Depth Profile

# In[13]:


np.shape(tmask)


# In[16]:


tmask[0,:,ij,ii]==0


# In[19]:


np.ma.masked_where(tmask[0,:,ij,ii]==0,f.variables['deptht'][:])


# In[15]:


fig,ax=plt.subplots(1,2,figsize=(6,6))
fig.subplots_adjust(wspace=.5) # space the axes out more
il=12 # hour
# use location 'S3':
ij,ii=places.PLACES['S3']['NEMO grid ji']
ax[0].plot(np.ma.masked_where(tmask[0,:,ij,ii]==0,f.variables['microzooplankton'][il,:,ij,ii]),f.variables['deptht'][:],'b-x',label='microzo')
ax[0].plot(np.ma.masked_where(tmask[0,:,ij,ii]==0,f.variables['diatoms'][il,:,ij,ii]),f.variables['deptht'][:],'g-x',label='diatoms')
ax[0].set_ylim(300,0)
ax[0].legend()
ax[0].set_xlabel('Concentration ($\mu$M)')
ax[0].set_ylabel('Depth (m)')
ax[1].plot(np.ma.masked_where(tmask[0,:,ij,ii]==0,f.variables['nitrate'][il,:,ij,ii]),f.variables['deptht'][:],'-x',color='orange',label='NO$_3$')
ax[1].plot(np.ma.masked_where(tmask[0,:,ij,ii]==0,f.variables['silicon'][il,:,ij,ii]),f.variables['deptht'][:],'-x',color='r',label='dissolved Si')
ax[1].set_ylim(300,0)
ax[1].set_ylabel('Depth (m)')
ax[1].legend()
ax[1].set_xlabel('Concentration ($\mu$M)')


# In[16]:


# mask example
amask=np.array((1,1,1,0,0))
anarray=np.array((6,5,4,3,2))
print(amask)
print(anarray)
print(np.ma.masked_where(amask,anarray))


# In[18]:


print(np.ma.masked_where(anarray>4,anarray))


# In[ ]:


tmask[0,:,ij,ii]


# In[4]:


import numpy as np


# In[5]:


x=np.array([0,1,2,3])
.5*(x[0:-1]+x[1:])


# ### Surface, Integrated plots - Aerial view

# In[14]:


# with pcolormesh: no smoothing
cmap0=cmocean.cm.thermal
cmap0.set_bad('tan')
cmap1=cmocean.cm.haline
cmap1.set_bad('k')
il=5
fig,ax=plt.subplots(1,2,figsize=(12,5))
m0=ax[0].pcolormesh(np.ma.masked_where(tmask[0,0,:,:]==0,f.variables['microzooplankton'][il,0,:,:]),cmap=cmap0)
viz_tools.set_aspect(ax[0],coords='grid')
ax[0].set_title('Surface Microzooplanton ($\mu$M) \n in Grid Coordinates')
fig.colorbar(m0,ax=ax[0])
# vertical sum of microzo in mmol/m3 * vertical grid thickness in m:
intuz=np.sum(f.variables['microzooplankton'][il,:,:,:]*fe3t.variables['e3t'][il,:,:,:]*tmask[0,:,:,:],0)
avguz=intuz/np.sum(fe3t.variables['e3t'][il,:,:,:]*tmask[0,:,:,:],0)
m1=ax[1].pcolormesh(navlon,navlat,np.ma.masked_where(tmask[0,0,:,:]==0,intuz),cmap=cmap1,shading='nearest')
viz_tools.set_aspect(ax[0],coords='map')
ax[1].set_title('Depth-Integrated Microzooplanton (mmol/m$^2$) \n in Geographic Coordinates');
fig.colorbar(m1,ax=ax[1])


# In[22]:


# with pcolormesh: no smoothing
cmap0=cmocean.cm.thermal
cmap0.set_bad('tan')
cmap1=cmocean.cm.haline
cmap1.set_bad('k')
il=5
fig,ax=plt.subplots(1,2,figsize=(12,5))
m0=ax[0].pcolormesh(np.ma.masked_where(tmask[0,0,:,:]==0,f2.variables['microzooplankton'][il,0,:,:]),cmap=cmap0)
viz_tools.set_aspect(ax[0],coords='grid')
ax[0].set_title('February 25 Surface Microzooplanton ($\mu$M) \n in Grid Coordinates')
fig.colorbar(m0,ax=ax[0])
# vertical sum of microzo in mmol/m3 * vertical grid thickness in m:
intuz=np.sum(f2.variables['microzooplankton'][il,:,:,:]*fe3t_2.variables['e3t'][il,:,:,:]*tmask[0,:,:,:],0)
avguz=intuz/np.sum(fe3t_2.variables['e3t'][il,:,:,:]*tmask[0,:,:,:],0)
m1=ax[1].pcolormesh(navlon,navlat,np.ma.masked_where(tmask[0,0,:,:]==0,intuz),cmap=cmap1,shading='nearest')
viz_tools.set_aspect(ax[0],coords='map')
ax[1].set_title('February 25 Depth-Integrated Microzooplanton (mmol/m$^2$) \n in Geographic Coordinates');
fig.colorbar(m1,ax=ax[1])


# In[24]:


cmap0=cmocean.cm.balance
cmap0.set_bad('tan')
fig,ax=plt.subplots(2,3,figsize=(12,12))
ax=ax.flatten()
tlist=[3, 5, 6, 7, 9,14]
for ii,il in enumerate(tlist):
    m0=ax[ii].pcolormesh(np.ma.masked_where(tmask[0,0,:,:]==0,
                    f.variables['microzooplankton'][il,0,:,:]-f.variables['microzooplankton'][0,0,:,:]),cmap=cmap0,
                        vmin=-.1,vmax=.1)
    viz_tools.set_aspect(ax[ii],coords='grid')
    ax[ii].set_title('Difference in \n Surface Microzooplanton ($\mu$M) \n in Grid Coordinates, hour='+str(il))
    fig.colorbar(m0,ax=ax[ii])


# In[25]:


# with pcolormesh: no smoothing
cmap0=cmocean.cm.thermal
cmap1=cmocean.cm.haline
il=5
fig,ax=plt.subplots(1,1,figsize=(8,8))
m0=ax.pcolormesh(np.ma.masked_where(tmask[0,0,:,:]==0,f.variables['microzooplankton'][il,0,:,:]),cmap=cmap0)
viz_tools.set_aspect(ax,coords='grid')
ax.set_title('Surface Microzooplanton ($\mu$M) \n in Grid Coordinates')
ax.set_ylim(500,550)
ax.set_xlim(180,230)
fig.colorbar(m0)


# In[26]:


# with contourf: smoothing, but creates a smaller file
cmap0=cmocean.cm.algae
cmap1=cmocean.cm.amp
cmap1.set_bad('k') # does nothing here
il=5
fig,ax=plt.subplots(1,2,figsize=(12,5))
m0=ax[0].contourf(np.ma.masked_where(tmask[0,0,:,:]==0,f.variables['microzooplankton'][il,0,:,:]),levels=50,cmap=cmap0)
viz_tools.set_aspect(ax[0],coords='grid')
ax[0].set_title('Surface Microzooplanton ($\mu$M) \n in Grid Coordinates')
fig.colorbar(m0,ax=ax[0])
# vertical sum of microzo in mmol/m3 * vertical grid thickness in m:
intuz=np.sum(f.variables['microzooplankton'][il,:,:,:]*fe3t.variables['e3t'][il,:,:,:],0)
m1=ax[1].contourf(navlon,navlat,np.ma.masked_where(tmask[0,0,:,:]==0,intuz),cmap=cmap1,levels=50)
viz_tools.set_aspect(ax[0],coords='map')
ax[1].set_title('Depth-Integrated Microzooplanton (mmol/m$^2$) \n in Geographic Coordinates');
ax[1].set_facecolor('gray')
fig.colorbar(m1,ax=ax[1])


# In[27]:


# with contourf: smoothing, but creates a smaller file
il=5
fig,ax=plt.subplots(1,1,figsize=(8,8))
m0=ax.contourf(np.ma.masked_where(tmask[0,0,:,:]==0,f.variables['microzooplankton'][il,0,:,:]),levels=50,cmap=cmap0)
viz_tools.set_aspect(ax,coords='grid')
ax.set_title('Surface Microzooplanton ($\mu$M) \n in Grid Coordinates')
ax.set_ylim(500,550)
ax.set_xlim(180,230)
fig.colorbar(m0)


# ### Thalweg plot
#   method using contour_thalweg from visualisations.py in tools repo

# In[28]:


#open bathy file and meshmask
fbathy=nc.Dataset('/data/eolson/results/MEOPAR/NEMO-forcing-new/grid/bathymetry_201702.nc')
fmesh=nc.Dataset('/data/eolson/results/MEOPAR/NEMO-forcing-new/grid/mesh_mask201702.nc')


# In[29]:


fig,ax=plt.subplots(1,1,figsize=(12,3))
cb=visualisations.contour_thalweg(ax,f.variables['microzooplankton'][il,...],fbathy,fmesh,clevels=50,cmap=cmocean.cm.amp)


# In[30]:


fbathy.close()
fmesh.close()


# In[31]:


f.close()
fe3t.close()


# In[32]:


f2.close()


# In[ ]: