plotSUMMAhru.py

"""
Created on Thu May 12 01:34:06 2016

@author: Michael
"""
%matplotlib inline
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import calendar
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import geopandas as gp
import time
import numpy as np
import platform
import os

## set working directory
if platform.uname()[0] == 'Linux':
    wdpath=r'/home/mgou/uwhydro/summaProj/summaData/columbia_full_run'
    bsnshp=r'/home/mgou/Dropbox/postdoc/summa/columbia/data/ColumbiaBasin.shp'
    hrushp=r'/home/mgou/Dropbox/postdoc/summa/columbia/data/columbia_hru_output.shp'
    outfolder=wdpath
else:
    wdpath=r'D:\@Workspace\SUMMA'
    bsnshp=r'D:\@Workspace\SUMMA\shapefiles\ColumbiaBasin.shp'
    hrushp=r'D:\@Workspace\SUMMA\outputs\columbia_hru_output.shp'
    outfolder=wdpath

os.chdir(outfolder)

#%% some function to define basemap

def add_gridlines(axis,labelsize=15):
    gl=axis.gridlines(draw_labels=True, 
                    xlocs = [-100, -110, -115, -120, -125], 
                    ylocs = [40, 42, 44, 46, 48, 50, 52, 54])
    gl.xlabels_top = False
    gl.ylabels_right = False
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER
    gl.xlabel_style = {'size': labelsize}
    gl.ylabel_style = {'size': labelsize}
    return gl

def add_map_features(ax, states_provinces=True, country_borders=True, land=True, ocean=True,lake=False):
    if states_provinces==True:
        states_provinces = cfeature.NaturalEarthFeature(
                category='cultural',
                name='admin_1_states_provinces_lines',
                scale='50m',
                facecolor='none')
        ax.add_feature(states_provinces, edgecolor='black', zorder = 2) #linewidth = 2

    if country_borders==True:
        country_borders = cfeature.NaturalEarthFeature(
                category='cultural',
                name='admin_0_boundary_lines_land',
                scale='50m',
                facecolor='none')
        ax.add_feature(country_borders, edgecolor='black', zorder = 2, linewidth = 1)

    if land==True:
        land = cfeature.NaturalEarthFeature(
            category='physical',
            name='land',
            scale='50m',
            facecolor='gray')
        ax.add_feature(land,facecolor='lightgray', zorder = 1)

    if ocean==True:
        ocean = cfeature.NaturalEarthFeature(
            category='physical',
            name='ocean',
            scale='50m',
            facecolor='blue')
        ax.add_feature(ocean,facecolor='lightblue', zorder = 1)
        
    if lake==True:
        rivers_lakes = cfeature.NaturalEarthFeature(
            category='physical',
            name='rivers_lake_centerlines',
            scale='50m',
            facecolor='none')
        ax.add_feature(rivers_lakes,facecolor='lightblue', zorder = 2)
    



#%% plot the actual HRU output by geopandas

start = time.time()
from descartes.patch import PolygonPatch


def plot_hru(ax,gdf,nhru,src_crs,to_crs,cmap,norm,var,title='',alpha=0.5):
    """ Plot HRU Polygon geometry """
    ## define the coordinate limits
    ax.set_extent([-125.01, -109.5, 41, 53], ccrs.Geodetic())    
    ## set title
    ax.set_title(title, fontsize=23)    
    ## add grid line
    gl = add_gridlines(ax,20)
    ## add base map features
    add_map_features(ax)    
    ## add HRU ploygons geopandas method      
    for i, row in gdf.iterrows():
        # transfer the projection
        if np.isnan(var[i] or var[i]>9999):
            continue
        if var[i]<0.00001:
            var[i]=-2
        color=cmap(norm(var[i]))
        geom = to_crs.project_geometry(row['geometry'], src_crs=src_crs)
        if geom.type == 'Polygon':
            ax.add_patch(PolygonPatch(gemo, facecolor=color, linewidth=0, alpha=alpha, zorder = 999))  # linewidth=0 because boundaries are drawn separately
        elif geom.type == 'MultiPolygon':
            for poly in geom.geoms:
                ax.add_patch(PolygonPatch(poly, facecolor=color, linewidth=0, alpha=alpha, zorder = 999))
        if i >= nhru-1: 
            break
        
shpname = hrushp

gdf = gp.GeoDataFrame.from_file(shpname)
gdf.describe()

# prepare color bar
swe_cmap=plt.get_cmap('Blues')
swe_cmap.set_under(color='tan')
norm = matplotlib.colors.Normalize(vmin=0.1,vmax=2150)

mons = ["Oct 2013","Nov 2013","Dec 2013","Jan 2014","Feb 2014","Mar 2014","Apr 2014","May 2014","Jun 2014","Jul 2014","Aug 2014","Sep 2014"]
rows = [0,0,0,0,1,1,1,1,2,2,2,2]
cols = [0,1,2,3,0,1,2,3,0,1,2,3]

# define projection
crs = ccrs.Mercator(central_longitude=-120, min_latitude=40, max_latitude=53, globe=None)

# plot single one
#fig, axes = plt.subplots(nrows=3, ncols=4, figsize=(20,18), subplot_kw=dict(projection=crs))

fig = plt.figure(figsize=(10,10))

swe_cmap=plt.get_cmap('Blues')
swe_cmap.set_under(color='tan')
norm = matplotlib.colors.Normalize(vmin=-1,vmax=2150)

hrucrs = ccrs.epsg(5070)    
#hruvars = [1.0,2,3,4,5,6,7,8,9,10,11,12]

maxhru = 11723

for i in [7]: #range(12):
    #ax=axes[rows[i],cols[i]]
    ax  = plt.axes(projection=crs)
    
    colname = 'swe{}'.format(i+1)
    hruvars = gdf[colname]
    plot_hru(ax, gdf, maxhru, hrucrs, crs, swe_cmap, norm, hruvars, mons[i], alpha=1.0)
    fig.savefig('test{:02d}.png'.format(i+1),dpi=300,bbox_inches='tight')
    fig.clear()

#f, axes = plt.subplots(nrows=3, ncols=4,figsize=(20,18), subplot_kw=dict(projection=crs))

    #gdf['val'] = i
    #plot = gp.plot(gdf,column='val',ax=ax,cmap=swe_cmap,linewidth=0.0)
    #
#f.subplots_adjust(left=0.1,bottom=0.05,right=0.95,top=0.95)

## Fine-tune figure; make subplots close to each other 
#f.subplots_adjust(hspace=0)

## # Fine-tune figure; hide x ticks for top plots and y ticks for right plots
#plt.setp([a.get_xticklabels() for a in axes[0, :]], visible=False)
#plt.setp([a.get_yticklabels() for a in axes[:, 1]], visible=False)

## Make the absolute color bar
#cbar_ax_abs = fig.add_axes([0.03, 0.2, 0.015, 0.6])
#cbar_ax_abs.tick_params(labelsize=14)
#cbar_abs = matplotlib.colorbar.ColorbarBase(cbar_ax_abs, cmap=swe_cmap, norm=norm).set_label(label='Snow Water Equivalent [mm]', size=16, labelpad=-85)

end = time.time()
print('Elapsed time is {} hour.'.format((end - start)/3600.))
#fig.savefig('test300.png', dpi=300,bbox_inches='tight')
#fig.savefig('test600.png', dpi=600,bbox_inches='tight')