% matplotlib inline

from pylab import *
import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap # plot on map projections

from netCDF4 import Dataset as netcdf # netcdf4-python module

"""
Explanation: Process U-Wind: Mean and Std
In this notebook, we will do a little complicated operations
* read 4D u-Wind data
* calculate mean and stadard deviation along the axis of time
* visualize based on the library of basemap
Data
wind data can be downlaed from https://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis2.html
This u-wind is a 4D data, includding [months|levels|lat|lon]. The presure levels in hPa. 
Moreover, the wind data with scaling and offset. when using them, have to restore them to oringal values.
1. Load basic libs
End of explanation
"""


ncset = netcdf(r'data/uwnd3.mon.mean.nc')
print(ncset)

"""
Explanation: 2. Read u-wind data and pick variables
2.1 Use print to check variable information.
Actually, you can also use numdump infile.nc -h to check the same inforamtion
End of explanation
"""


ncset.set_auto_mask(False)

lon = ncset['lon'][:]  
lat = ncset['lat'][:] 
lev = ncset['level'][:]    
u   = ncset['uwnd'][504:624,:]   # for the period 1990-1999.

print(u.shape)

print(lev)

"""
Explanation: 2.2 Read data
Have to set_auto_mask(False) to automatically scaling and offseting, or may cause problem.
End of explanation
"""


plot(u[:,1,0, 0])

"""
Explanation: 2.3  Have a quick shot on first grid
End of explanation
"""


u_10y = np.mean(u, axis=0) # calculate mean for all years and months
u_10y.shape

"""
Explanation: 3. Calculate Mean and STD in time
3.1 Mean
End of explanation
"""


u_10y_std=np.std(u, axis=0)
u_10y.shape

"""
Explanation: 3.2  STD
End of explanation
"""


[lons, lats] = meshgrid(lon,lat)

m = Basemap(projection='robin', lon_0=0)
m.drawcoastlines()

# draw parallels and meridians.
m.drawparallels(np.arange(-90.,120.,30.))
m.drawmeridians(np.arange(0.,360.,60.))
m.drawmapboundary(fill_color='aqua')

minu = floor(np.min(u_10y[0]))
maxu = ceil(np.max(u_10y[0]))

h = m.pcolormesh(lons, lats, u_10y[0], shading='flat',latlon=True, cmap='jet', vmin=minu, vmax=maxu)
m.colorbar(h, location='bottom', pad="15%", label='[$^oC$]')


plt.title('U1000 Mean between 1990-1999 [m/s]')

[lons, lats] = meshgrid(lon,lat)

m = Basemap(projection='robin', lon_0=0)
m.drawcoastlines()

# draw parallels and meridians.
m.drawparallels(np.arange(-90.,120.,30.))
m.drawmeridians(np.arange(0.,360.,60.))
m.drawmapboundary(fill_color='aqua')

minu = floor(np.min(u_10y_std[0]))
maxu = ceil(np.max(u_10y_std[0]))

h = m.pcolormesh(lons, lats, u_10y_std[0], shading='flat',latlon=True, cmap='jet', vmin=minu, vmax=maxu)
m.colorbar(h, location='bottom', pad="15%", label='[$^oC$]')

plt.title('U1000 STD between 1990-1999 [m/s]')

"""
Explanation: 3.3 Visualize Mean and STD at 1000hPa (the first level)
End of explanation
"""


