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

# # Plot variable on original FESOM grid

# In many cases it is fine to first interpolate FESOM data to the regular grid and then plot them. Usually it will allow you to make plots faster. However sometimes you would like to work with original FESOM triangular mesh. 

# If you don't need interactivity on the plots, comment:
# 
#     #%matplotlib notebook

# In[4]:


import sys
sys.path.append("../")

import pyfesom as pf
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
get_ipython().run_line_magic('matplotlib', 'notebook')
from matplotlib import cm
from netCDF4 import Dataset


# Load mesh

# In[5]:


meshpath  ='../../../../FESOM/mesh/'
mesh = pf.load_mesh(meshpath, get3d=True, usepickle=True)


# Load netCDF file that contains your data:

# In[6]:


fl = Dataset('../../../../FESOM/data/fesom.2007.oce.nc')


# There are several variables that we can plot:

# In[7]:


fl.variables.keys()


# For every timestep there is a vector of values, each coresponding to the node of the model mesh.

# In[8]:


fl.variables['temp'].shape


# In order to get values at some depth one should use *get_data* function from *pyfesom*. The input variables for this function are:
# 
#     data : array
#         complete 3d data for one timestep
#     mesh : fesom_mesh object
#         mesh representation
#     depth : int
#         desired depth
#         
# As an output you will get: 
# 
#     level_data : array
#         2d array (actually vector) with data from the desired level.
#     elem_no_nan : array
#         array with triangles (defined as triplets of node indexes) with
#         not NaN elements.
#         
# The actual model level will be printed out.

# Here is the result of requesting data from the level that is closest to 500 meters:

# In[9]:


level_data, elem_no_nan = pf.get_data(fl.variables['temp'][0,:], mesh, 500)


# The information from *level_data* and *elem_no_nan* is enought to make the plot:

# Define the map:

# In[16]:


map = Basemap(projection='robin',lon_0=0, resolution='c')
x, y = map(mesh.x2, mesh.y2)


# Plot the data from the last timestep of our dataset at 100m depth:

# In[11]:


get_ipython().run_cell_magic('time', '', 'level_data, elem_no_nan = pf.get_data(fl.variables[\'temp\'][-1,:],mesh,100)\n\nplt.figure(figsize=(10,7))\nmap.drawmapboundary(fill_color=\'0.9\')\nmap.drawcoastlines()\n\nlevels = np.arange(-3., 30., 1)\nplt.tricontourf(x, y, elem_no_nan[::], level_data, levels = levels, cmap=cm.Spectral_r, extend=\'both\')\ncbar = plt.colorbar(orientation=\'horizontal\', pad=0.03);\ncbar.set_label("Temperature, $^{\\circ}$C")\nplt.title(\'Temperature at 100m depth\')\nplt.tight_layout()\n')


# If your color range covers the whole range of your data everything is fine. But if you try to narrow the range, the places with values outside your range become masked out:

# In[12]:


get_ipython().run_cell_magic('time', '', 'level_data, elem_no_nan = pf.get_data(fl.variables[\'temp\'][-1,:],mesh,100)\n\nplt.figure(figsize=(10,7))\nmap.drawmapboundary(fill_color=\'0.9\')\nmap.drawcoastlines()\n\nlevels = np.arange(-3., 20., 1)\nplt.tricontourf(x, y, elem_no_nan[::], level_data, levels = levels, cmap=cm.Spectral_r, extend=\'both\')\ncbar = plt.colorbar(orientation=\'horizontal\', pad=0.03);\ncbar.set_label("Temperature, $^{\\circ}$C")\nplt.title(\'Temperature at 100m depth\')\nplt.tight_layout()\n')


# The way around it is to chnage your data :) 

# In[13]:


get_ipython().run_cell_magic('time', '', 'level_data, elem_no_nan = pf.get_data(fl.variables[\'temp\'][-1,:],mesh,100)\n\nplt.figure(figsize=(10,7))\nmap.drawmapboundary(fill_color=\'0.9\')\nmap.drawcoastlines()\n\nlevels = np.arange(-3., 20., 1)\n\neps=(levels.max()-levels.min())/50.\nlevel_data[level_data<=levels.min()]=levels.min()+eps\nlevel_data[level_data>=levels.max()]=levels.max()-eps\nplt.tricontourf(x, y, elem_no_nan[::], level_data, levels = levels, cmap=cm.Spectral_r, extend=\'both\')\ncbar = plt.colorbar(orientation=\'horizontal\', pad=0.03);\ncbar.set_label("Temperature, $^{\\circ}$C")\nplt.title(\'Temperature at 100m depth\')\nplt.tight_layout()\n')


# The things above were easy - since it was just doing contour plots. But sometimes it is nessesary to have a look at the data on original grid. Since properties in FESOM 1 are defined on the verticies of the triangles, the values for making the colors of the triangles will be calculated as a mean of the values in the verticies. 

# In[39]:


get_ipython().run_cell_magic('time', '', 'level_data, elem_no_nan = pf.get_data(fl.variables[\'temp\'][-1,:],mesh,100)\n\nplt.figure(figsize=(10,7))\nmap.drawmapboundary(fill_color=\'0.9\')\nmap.drawcoastlines()\n\nplt.tripcolor(x, y, elem_no_nan, \\\n              level_data, \\\n              edgecolors=\'k\',\\\n              lw = 0.01,\n             cmap=cm.Spectral_r,\n             vmin = -3,\n             vmax = 30)\ncbar = plt.colorbar(orientation=\'horizontal\', pad=0.03);\ncbar.set_label("Temperature, $^{\\circ}$C")\nplt.title(\'Temperature at 100m depth\')\nplt.tight_layout()\n')


# As you can see this takes considerably longer time. If you would like to have a look at the values exactly on the points were they are defined, you can try to use technique described in the *fesom_with_datashader.ipynb* tootorial.

# In[46]:


m = Basemap(projection='mill',llcrnrlat=45,urcrnrlat=60,\
            llcrnrlon=-10,urcrnrlon=15,resolution='c')
x2, y2 = m(mesh.x2, mesh.y2)


# In[47]:


get_ipython().run_cell_magic('time', '', 'level_data, elem_no_nan = pf.get_data(fl.variables[\'temp\'][-1,:],mesh,0)\n\nplt.figure(figsize=(10,7))\nm.drawmapboundary(fill_color=\'0.9\')\n#m.drawcoastlines()\n\nplt.tripcolor(x2, y2, elem_no_nan, \\\n              level_data, \\\n              edgecolors=\'k\',\\\n              lw = 0.1,\n             cmap=cm.Spectral_r,\n             vmin = 5,\n             vmax = 15)\ncbar = plt.colorbar(orientation=\'horizontal\', pad=0.03);\ncbar.set_label("Temperature, $^{\\circ}$C")\nplt.title(\'Temperature at 0m depth\')\nplt.tight_layout()\n')


# In[ ]: