CR.ProductRegrid.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "####################################################################################################\n",
    "# CR.ProductRegrid.ipynb\n",
    "# Karl Lapo July/2015\n",
    "####################################################################################################\n",
    "# Regrids all data to MTCLIM grid, looks @ differences between products\n",
    "####################################################################################################\n",
    "\n",
    "# must insert this statement to render the plots within the notebook\n",
    "# this is specific to the ipython notebook\n",
    "%matplotlib inline\n",
    "\n",
    "## Import statements\n",
    "# netcdf/numpy/xray\n",
    "from netCDF4 import num2date, date2num\n",
    "import numpy as np\n",
    "from datetime import datetime, timedelta\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import xray\n",
    "import scipy\n",
    "\n",
    "# OS interaction\n",
    "import sys\n",
    "import os\n",
    "\n",
    "# import subplots function for plotting\n",
    "import seaborn as sns\n",
    "sns.set_style(\"whitegrid\")\n",
    "import matplotlib\n",
    "from matplotlib.pyplot import subplots\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "import kgraph\n",
    "\n",
    "#### Directory Lists\n",
    "# Unix\n",
    "if 'linux' in sys.platform:\n",
    "    dir_pre = '/home/lapok/gdrive/'\n",
    "# Mac\n",
    "elif 'darwin' in sys.platform:\n",
    "    dir_pre = '/Users/karllapo/gdrive/'\n",
    "\n",
    "dir_data = dir_pre+'SnowHydrology/proj/CloudClimatology/data'\n",
    "dir_print = dir_pre+'SnowHydrology/proj/CloudClimatology/Graphics'\n",
    "\n",
    "# List of sub-directory names for each data set\n",
    "dir_NLDAS = '/NLDAS'\n",
    "dir_SYN = '/CERES_SYN'\n",
    "dir_grobs = '/GroundObs'\n",
    "dir_VIC = '/VIC_MTCLIM'\n",
    "dir_WRF = '/WRF'\n",
    "\n",
    "# Directory for basemap pickle files\n",
    "dir_bmap = dir_pre+'SnowHydrology/proj/CloudClimatology/data/basemap'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "####################################################################################################\n",
    "# Functions\n",
    "####################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "##### Basemap\n",
    "def build_basemap(lon,lat,dir_bmap,bmap_name='basemap.pickle',rewrite=False):\n",
    "    # Lat/Lon handling - map extent\n",
    "    bmap_dict = {}\n",
    "    bmap_dict['lat_i'] = np.min(lat)\n",
    "    bmap_dict['lon_i'] = np.min(lon)\n",
    "    bmap_dict['lat_j'] = np.max(lat)\n",
    "    bmap_dict['lon_j'] = np.max(lon)\n",
    "    \n",
    "    bmap_dict['lat_mid'] = lat[np.round(lat.size/2)]\n",
    "    bmap_dict['lon_mid'] = lon[np.round(lon.size/2)]\n",
    "    \n",
    "    bmap_dict['lat_labels'] = np.arange(np.round(bmap_dict['lat_i']), np.round(bmap_dict['lat_j']), 2)\n",
    "    bmap_dict['lon_labels'] = np.arange(np.round(bmap_dict['lon_i']), np.round(bmap_dict['lon_j']), 2)\n",
    "    \n",
    "    os.chdir(dir_bmap)\n",
    "    # Force rewriting basemap pickle file\n",
    "    if rewrite:\n",
    "        bmap = Basemap(llcrnrlon=bmap_dict['lon_i'],llcrnrlat=bmap_dict['lat_i'],\\\n",
    "                        urcrnrlon=bmap_dict['lon_j'],urcrnrlat=bmap_dict['lat_j'],\\\n",
    "                        rsphere=(6378137.00,6356752.3142),resolution='l',area_thresh=1000.,projection='lcc',\\\n",
    "                        lat_1=bmap_dict['lat_mid'],lon_0=bmap_dict['lon_mid'])\n",
    "        pickle.dump(bmap,open(bmap_name,'wb'),-1)\n",
    "    \n",
    "    else:\n",
    "        try:\n",
    "            bmap = pickle.load(open(bmap_name,'rb'))\n",
    "        except IOError as e:\n",
    "            bmap = Basemap(llcrnrlon=bmap_dict['lon_i'],llcrnrlat=bmap_dict['lat_i'],\\\n",
    "                            urcrnrlon=bmap_dict['lon_j'],urcrnrlat=bmap_dict['lat_j'],\\\n",
    "                            rsphere=(6378137.00,6356752.3142),resolution='l',area_thresh=1000.,projection='lcc',\\\n",
    "                            lat_1=bmap_dict['lat_mid'],lon_0=bmap_dict['lon_mid'])\n",
    "            pickle.dump(bmap,open(bmap_name,'wb'),-1)\n",
    "    \n",
    "    return bmap,bmap_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "##### Re-Gridding Function\n",
    "# Converts from_xray_ds spatial resolution to to_xray_ds using a nearest neighbor approach.\n",
    "# Assumes arrays are ['time','lat','lon']\n",
    "def xray_regrid(from_xray_ds,to_xray_ds,field):\n",
    "    ## Coordinates to regrid to\n",
    "    lat_rg = to_xray_ds.latitude.values\n",
    "    nlat_rg = lat_rg.size\n",
    "    lon_rg = to_xray_ds.longitude.values\n",
    "    nlon_rg = lon_rg.size\n",
    "    lonm_rg,latm_rg = np.meshgrid(lon_rg,lat_rg)\n",
    "    \n",
    "    ## Coordinates of original dataset\n",
    "    lat_g = from_xray_ds.latitude.values\n",
    "    nlat_g = lat_g.size\n",
    "    lon_g = from_xray_ds.longitude.values\n",
    "    nlon_g = lon_g.size\n",
    "    lonm_g,latm_g = np.meshgrid(lon_g,lat_g)\n",
    "    latm_g = np.reshape(latm_g[:],(nlat_g*nlon_g,1))\n",
    "    lonm_g = np.reshape(lonm_g[:],(nlat_g*nlon_g,1))\n",
    "    \n",
    "    ## Combine coordinates\n",
    "    grid_points = np.append(lonm_g[:],latm_g[:],axis=1)\n",
    "    regrid_points = (lonm_rg,latm_rg)\n",
    "    \n",
    "    ## Interpolate in space dimension, along time dimension\n",
    "    dates = from_xray_ds.time\n",
    "    regrid = np.empty((dates.size,nlat_rg,nlon_rg))\n",
    "    for ind,d in enumerate(dates):\n",
    "        ## Use loc or sel here instead of values w/ indexing\n",
    "        grid = np.reshape(np.squeeze(from_xray_ds[field].values[ind,:,:]),(nlat_g*nlon_g))\n",
    "        regrid[ind,:,:] = scipy.interpolate.griddata(grid_points,grid,regrid_points,method='nearest') \n",
    "        \n",
    "    ## Output to xray structure\n",
    "    regrid_xray_ds = xray.Dataset({field: (['time','latitude','longitude'], regrid)},\\\n",
    "                                  coords={'time': dates,'latitude': lat_rg,'longitude':lon_rg})\n",
    "    return(regrid_xray_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "##### Discrete colorbar -- from Joe Hamman (https://github.com/jhamman/tonic/blob/master/tonic/plot_utils.py#L66-L94)\n",
    "def cmap_discretize(cmap, n_colors=10):\n",
    "    \"\"\"Return discretized colormap.\n",
    "    Parameters\n",
    "    ----------\n",
    "    cmap : str or colormap object\n",
    "        Colormap to discretize.\n",
    "    n_colors : int\n",
    "        Number of discrete colors to divide `cmap` into.\n",
    "    Returns\n",
    "    ----------\n",
    "    disc_cmap : LinearSegmentedColormap\n",
    "        Discretized colormap.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        cmap = cm.get_cmap(cmap)\n",
    "    except:\n",
    "        cmap = cm.get_cmap(eval(cmap))\n",
    "    colors_i = np.concatenate((np.linspace(0, 1., n_colors), (0., 0., 0., 0.)))\n",
    "    colors_rgba = cmap(colors_i)\n",
    "    indices = np.linspace(0, 1., n_colors + 1)\n",
    "    cdict = {}\n",
    "    for ki, key in enumerate(('red', 'green', 'blue')):\n",
    "        cdict[key] = [(indices[i], colors_rgba[i - 1, ki], colors_rgba[i, ki])\n",
    "                      for i in range(n_colors + 1)]\n",
    "\n",
    "    return matplotlib.colors.LinearSegmentedColormap(cmap.name + \"_%d\" % n_colors,\n",
    "                                              cdict, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#### Function for combining xray data variables into a single array with new labeled dimension\n",
    "def combinevars(ds_in,dat_vars,new_dim_name='new_dim',combinevarname='new_var'):\n",
    "#     ds_out = xray.Dataset()\n",
    "    ds_out = xray.concat([ds_in[dv] for dv in dat_vars],dim='new_dim')\n",
    "    ds_out = ds_out.rename({'new_dim': new_dim_name})\n",
    "    ds_out.coords[new_dim_name] = dat_vars\n",
    "    ds_out = xray.Dataset({combinevarname:ds_out})\n",
    "#     ds_out.name = combinevarname\n",
    "    \n",
    "    return ds_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#### Violinplots function using xray structures\n",
    "# Formatting variables\n",
    "def vp_xray(plot_ds,plot_var,plot_interval,plot_period,\\\n",
    "            bin_centers,grouping,grouping_var,grouping_labels):\n",
    "    \n",
    "    ## Labels (seasonal or monthly)\n",
    "    month_label = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec']\n",
    "    season_label = ['DJF','MAM','JJA','SON']\n",
    "\n",
    "    # Violinplot body color\n",
    "    colsche = kgraph.SRON(plot_var.size)\n",
    "\n",
    "    # Bin positions\n",
    "    pos_buffer = 1\n",
    "    num_pos = plot_var.size\n",
    "    num_group = grouping.size\n",
    "    bin_pos = [np.linspace((num_pos+pos_buffer)*n+1,(num_pos+pos_buffer)*(n+1)-pos_buffer,\\\n",
    "                num=num_pos) for n in np.arange(num_group)]\n",
    "    fig = {}\n",
    "\n",
    "    ## Create structure for plotting\n",
    "    # Note: the dimenion and variable names are meaningless, but relevant to the ProductRegrid.ipynb perspective\n",
    "    dat_for_plot = combinevars(plot_ds,plot_var,new_dim_name='product',combinevarname='diff')\n",
    "\n",
    "    ## Loop through years provided\n",
    "    # Create iterable object\n",
    "    if plot_period.size < 2:\n",
    "        plot_period = np.linspace(plot_period,plot_period+1,num=1)\n",
    "\n",
    "    for n,YYYY in enumerate(plot_period):\n",
    "        YYYY = int(YYYY)\n",
    "        if plot_interval == 'monthly':\n",
    "            f, axes = plt.subplots(3, 4, figsize=(12, 9), sharex=True, sharey=True)\n",
    "        elif plot_interval == 'seasonal':\n",
    "            f, axes = plt.subplots(1, 4, figsize=(12, 3), sharex=True, sharey=True)\n",
    "        fig[YYYY] = f\n",
    "            \n",
    "        for m,ax in enumerate(axes.flat):\n",
    "            ax.plot(bin_pos,np.zeros_like(bin_pos),linewidth=.75,color='.75')\n",
    "\n",
    "            ## Product in interval\n",
    "            if plot_interval == 'monthly':\n",
    "                d = pd.date_range(start=(str(YYYY)+'-'+str(m+1)+'-1'),freq='M',periods=1)\n",
    "            elif plot_interval == 'seasonal':\n",
    "                if m == 0:\n",
    "                    d = pd.date_range(start=(str(YYYY-1)+'-'+str(12)+'-1'),freq='M',periods=3)\n",
    "                else:\n",
    "                    d = pd.date_range(start=(str(YYYY)+'-'+str((3*m))+'-1'),freq='M',periods=3)\n",
    "\n",
    "            # Remove figures outside date range of input variables\n",
    "            if d[-1] > dat_for_plot.time.values[-1]:\n",
    "                plt.close(f)\n",
    "                fig[YYYY]\n",
    "                del fig[YYYY]\n",
    "                return fig\n",
    "            # Locate time step in Dataset\n",
    "            temp = dat_for_plot.loc[{'time':d}]\n",
    "            temp = temp.mean('time')\n",
    "\n",
    "            ## Violin plots\n",
    "            for n,gr in enumerate(grouping):\n",
    "                if not gr in temp[grouping_var].values:\n",
    "                    continue\n",
    "                group_for_plot = temp['diff'].loc[{grouping_var:gr}].values\n",
    "                result = build_violinplot(group_for_plot.T,ax,bin_pos[n],color=colsche)\n",
    "\n",
    "            # Violinplot labels\n",
    "            handle = [matplotlib.patches.Patch(color=colsche[n], label=plot_var[n]) \\\n",
    "                      for n,p in enumerate(plot_var)]\n",
    "\n",
    "            # Labels and legend\n",
    "            if m == 0:\n",
    "                ax.legend(handles=handle,loc='best',ncol=num_pos)\n",
    "            ax.set_xticks(np.array(bin_pos)[:,0])\n",
    "            ax.set_ylim([-80,80])\n",
    "            if np.mod(((m)+4),4) == 0:\n",
    "                ax.set_ylabel('Bias (Wm$^{-2}$)')\n",
    "            if m >= 8 and plot_interval == 'monthly':        \n",
    "                ax.set_xticklabels(grouping_labels, rotation=60)\n",
    "                ax.set_xlabel('elevation bin (m)')\n",
    "            elif plot_interval == 'seasonal':\n",
    "                ax.set_xticklabels(grouping_labels, rotation=60)\n",
    "                ax.set_xlabel('elevation bin (m)')\n",
    "            if plot_interval == 'monthly':\n",
    "                ax.set_title(month_label[m]) \n",
    "            elif plot_interval == 'seasonal':\n",
    "                ax.set_title(season_label[m])\n",
    "        f.tight_layout()\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#### Hack the shit out of the violinplot because it has some really silly limitations\n",
    "# Callable method for doing the kernel estimate\n",
    "def violin_local(ax,vpstats, positions=None, widths=1, extrema=False):\n",
    "    \"\"\"Drawing function for violin plots.\n",
    "\n",
    "    Call signature::\n",
    "\n",
    "      violin(vpstats, positions=None, vert=True, widths=0.5,\n",
    "             showmeans=False, showextrema=True, showmedians=False):\n",
    "\n",
    "    Draw a violin plot for each column of `vpstats`. Each filled area\n",
    "    extends to represent the entire data range, with optional lines at the\n",
    "    mean, the median, the minimum, and the maximum.\n",
    "    \"\"\"\n",
    "    # Collection\n",
    "    artists={}\n",
    "    \n",
    "    # Statistical quantities to be plotted on the violins\n",
    "    means = []\n",
    "    mins = []\n",
    "    maxes = []\n",
    "    medians = []\n",
    "\n",
    "    # Collections to be returned\n",
    "    artists = {}\n",
    "    if np.shape(np.shape(vpstats))[0] > 1:\n",
    "        N = len(vpstats)\n",
    "    else:\n",
    "        N = 1\n",
    "    datashape_message = (\"List of violinplot statistics and `{0}` \"\n",
    "                         \"values must have the same length\")\n",
    "\n",
    "    # Validate positions\n",
    "    if positions is None:\n",
    "        positions = range(1, N + 1)\n",
    "    elif len(np.atleast_1d(positions)) != N:\n",
    "        raise ValueError(datashape_message.format(\"positions\"))\n",
    "\n",
    "    # Validate widths\n",
    "    if np.isscalar(widths):\n",
    "        widths = [widths] * N\n",
    "    elif len(np.atleast_1d(widths)) != N:\n",
    "        raise ValueError(datashape_message.format(\"widths\"))\n",
    "\n",
    "    # Calculate ranges for statistics lines\n",
    "    pmins = -0.25 * np.array(widths) + positions\n",
    "    pmaxes = 0.25 * np.array(widths) + positions\n",
    "\n",
    "    # Check whether we are rendering vertically or horizontally\n",
    "    fill = ax.fill_betweenx\n",
    "    perp_lines = ax.hlines\n",
    "    par_lines = ax.vlines\n",
    "\n",
    "    # Render violins\n",
    "    bodies = []\n",
    "    for stats, pos, width in zip(vpstats, positions, widths):\n",
    "        # The 0.5 factor reflects the fact that we plot from v-p to\n",
    "        # v+p\n",
    "        vals = np.array(stats['vals'])\n",
    "        vals = 0.5 * width * vals / vals.max()\n",
    "        bodies += [fill(stats['coords'],\n",
    "                        -vals + pos,\n",
    "                        vals + pos,\n",
    "                        facecolor='y',\n",
    "                        alpha=0.3)]\n",
    "        means.append(stats['mean'])\n",
    "        mins.append(stats['min'])\n",
    "        maxes.append(stats['max'])\n",
    "        medians.append(stats['median'])\n",
    "    artists['bodies'] = bodies\n",
    "    \n",
    "    # Render means\n",
    "    artists['cmeans'] = perp_lines(means, pmins, pmaxes, colors='r')\n",
    "    \n",
    "    # Render extremagi\n",
    "    if extrema:\n",
    "        artists['cmaxes'] = perp_lines(maxes, pmins, pmaxes, colors='r')\n",
    "        artists['cmins'] = perp_lines(mins, pmins, pmaxes, colors='r')\n",
    "        artists['cbars'] = par_lines(positions, mins, maxes, colors='r')\n",
    "    \n",
    "    return artists\n",
    "\n",
    "def _kde_method(X, coords):\n",
    "    kde = matplotlib.mlab.GaussianKDE(X, None)\n",
    "    return kde.evaluate(coords)\n",
    "\n",
    "# Actually build the violinplot statistics and removes nans\n",
    "def build_violinplot(X,ax,pos=None,color=None,dim=1):            \n",
    "    import kgraph\n",
    "    # Dimension to operate along\n",
    "    if dim < 0 or dim > np.shape(X.shape):\n",
    "        raise ValueError('dim must be in range of dimensions of X')\n",
    "    # Number of arrays provided for violinplots\n",
    "    nx = X.shape[dim]\n",
    "    \n",
    "    if color is None:\n",
    "        color = kgraph.SRON(nx)\n",
    "    elif not len(np.atleast_1d(color)) == nx:\n",
    "        raise ValueError('color variable does not match data (length mismatch)')\n",
    "    \n",
    "    # Validate positions\n",
    "    if pos is None:\n",
    "        pos = range(1, N + 1)\n",
    "    elif not len(np.atleast_1d(pos)) == nx:\n",
    "        raise ValueError('pos variable does not match data (length mismatch)')\n",
    "    \n",
    "    ## violinplots\n",
    "    # pop the desired axis to the front\n",
    "    X = np.rollaxis(X,dim)\n",
    "    for n in np.arange(0,nx):\n",
    "        # Remove NaNs\n",
    "        xplot = X[n]\n",
    "        xplot = xplot[~np.isnan(xplot)]\n",
    "    \n",
    "        if xplot.any():\n",
    "            # Dictionary of violin statistics\n",
    "            vpstats = matplotlib.cbook.violin_stats(xplot,_kde_method,points=100)\n",
    "            position = [pos[n]]*1\n",
    "            result = violin_local(ax,vpstats,positions=position)\n",
    "            \n",
    "            ## Format\n",
    "            # Violin body and label\n",
    "            for pc in result['bodies']:\n",
    "                pc.set_facecolor(color[n])\n",
    "                pc.set_edgecolor(color[n])\n",
    "                pc.set_alpha(None)\n",
    "            # Violin lines\n",
    "            violinplotlines = ['cbars','cmins','cmeans','cmaxes']\n",
    "            for vpl in violinplotlines:\n",
    "                if vpl in result:\n",
    "                    result[vpl].set_color('k')\n",
    "    \n",
    "## todo:\n",
    "# [x] handle position argument of none\n",
    "# [x] handle X as a matrix (loop within function instead of outside of it)\n",
    "# [x] argument handling which dimension to group the violinplots along\n",
    "# [x] integrate into xray violin function\n",
    "# [ ] submit issue for matplotlib to use a smarter violin_stat and violin (handle single violin, handle nans in stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "####################################\n",
    "## Read previously processed data ##\n",
    "####################################\n",
    "# ///// See CC.CA.StatisticsMaps.Master for details on creation of xray data \n",
    "\n",
    "###########\n",
    "## NLDAS ##\n",
    "os.chdir(dir_data+dir_NLDAS)\n",
    "nldas = xray.open_dataset('CA.NLDAS.irrad.monthly.nc')\n",
    "nldas = nldas.rename({'DLWRF_110_SFC':'LWdwn','DSWRF_110_SFC':'SWdwn'})\n",
    "\n",
    "#########\n",
    "## SYN ##\n",
    "os.chdir(dir_data+dir_SYN)\n",
    "syn = xray.open_dataset('CA.syn.irrad.monthly.nc')\n",
    "############ No flipping -- SYN reformatted on 08/20/15 for CR.Composite. OTHER SCRIPTS WILL NEED THIS FIX\n",
    "\n",
    "############\n",
    "## MTCLIM ##\n",
    "os.chdir(dir_data+dir_VIC)\n",
    "mtclim = xray.open_dataset('CA.MTCLIM.irrad.monthly.nc')\n",
    "\n",
    "#########################\n",
    "## Ground Observations ##\n",
    "os.chdir(dir_data+dir_grobs)\n",
    "grobs = xray.open_dataset('CA.grobs.irrad.monthly.nc')\n",
    "grobs.SWdwn.values[grobs.SWdwn.values == 0] = np.nan\n",
    "grobs = grobs.rename({'lon':'longitude','lat':'latitude'})\n",
    "grobs.longitude.values = -grobs.longitude.values\n",
    "\n",
    "#########\n",
    "## WRF ##\n",
    "os.chdir(dir_data+dir_WRF)\n",
    "wrf = xray.open_dataset('CA.WRF.irrad.monthly.nc')\n",
    "wrf = wrf.rename({'lon':'longitude','lat':'latitude'})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "### Re-grid to mtclim spatial scale\n",
    "# Q: do all coordinates represent the same thing? (e.g., lat-lon = lower left corner? \n",
    "# Should I convert to center of grid labeling scheme?\n",
    "regrid_ds = {}\n",
    "regrid_ds['nldas'] = xray_regrid(nldas,wrf,'SWdwn')\n",
    "regrid_ds['syn'] = xray_regrid(syn,wrf,'SWdwn')\n",
    "regrid_ds['mtclim'] = xray_regrid(mtclim,wrf,'SWdwn')\n",
    "regrid_ds['wrf'] = wrf\n",
    "\n",
    "# Describe product names for iteration\n",
    "pr_names = ['mtclim','nldas','syn','wrf']\n",
    "num_products = np.size(pr_names)\n",
    "\n",
    "### Aligned xray Dataset structure\n",
    "align_ds = xray.Dataset()\n",
    "for pr in pr_names:\n",
    "    align_ds[pr] = regrid_ds[pr].SWdwn\n",
    "    \n",
    "os.chdir(dir_data)\n",
    "align_ds.to_netcdf('SWdwn.AllProducts.Align.nc')\n",
    "\n",
    "print(align_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "##########################################\n",
    "###### Plot differences from mtclim ######\n",
    "## plotting variables, coordinates, basemap\n",
    "# colors\n",
    "SWmin_delta = -50\n",
    "SWmax_delta = 50\n",
    "cmap = cmap_discretize(cm.gnuplot2,15)\n",
    "cmap_delta = cmap_discretize(cm.RdBu_r,11)\n",
    "\n",
    "# Build basemap\n",
    "lat = regrid_ds['mtclim'].latitude.values\n",
    "lon = regrid_ds['mtclim'].longitude.values\n",
    "\n",
    "bmp,bmd = build_basemap(lon,lat,dir_bmap,'CA.Domain.bmp.pickle',rewrite=True)\n",
    "lat_labels = bmd['lat_labels']\n",
    "lon_labels = bmd['lon_labels']\n",
    "\n",
    "## Loop through dates\n",
    "for d in pd.date_range(start='2010-9-01',end='2010-10-01',freq='M'):\n",
    "    print(d)\n",
    "    # Create figure properties\n",
    "    fig = plt.figure(figsize=(12,6))\n",
    "#     gs = matplotlib.gridspec.GridSpec(2,5,width_ratios=[16,16,16,16,1])\n",
    "    gs = matplotlib.gridspec.GridSpec(2,num_products+1,width_ratios=np.append(np.ones(num_products)*16,np.array(1.)))\n",
    "\n",
    "    ## Color range\n",
    "    SWmax = 0\n",
    "    SWmin = 500\n",
    "    for pr in pr_names:\n",
    "        if regrid_ds[pr].SWdwn.loc[d:d].any() \\\n",
    "            and not np.isnan(np.nanmax(np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values))).any():\n",
    "            SWmax = max(np.nanmax(np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values)),SWmax)\n",
    "            SWmax = np.round(SWmax/10)*10\n",
    "        \n",
    "        if regrid_ds[pr].SWdwn.loc[d:d].any() \\\n",
    "            and np.nanmin(np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values)) > 0 \\\n",
    "            and not np.isnan(np.nanmin(np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values))).any():\n",
    "            SWmin = min(np.nanmin(np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values)),SWmin)\n",
    "            SWmin = np.round(SWmin/10)*10\n",
    "    dSW = 10\n",
    "    \n",
    "    for ind,pr in enumerate(pr_names):\n",
    "        ## Monthly value for each product\n",
    "        ax = plt.subplot(gs[0,ind])\n",
    "        # Lat/Lon handling - product coords\n",
    "        lon_rad,lat_rad = np.meshgrid(regrid_ds[pr].longitude.values,regrid_ds[pr].latitude.values)\n",
    "        if regrid_ds[pr].SWdwn.loc[d:d].any():\n",
    "            SW_for_plot = np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values)\n",
    "            SW_for_plot = np.ma.masked_where(np.isnan(SW_for_plot),SW_for_plot)\n",
    "            im_avg = bmp.pcolormesh(lon_rad,lat_rad,SW_for_plot,\\\n",
    "                        cmap=cmap,vmin=SWmin,vmax=SWmax,shading='flat',latlon=True)\n",
    "        ax.set_title((pr))\n",
    "\n",
    "        ## Format\n",
    "        if ind == 0:\n",
    "            bmp.drawparallels(lat_labels,labels=[1,0,0,0])\n",
    "            bmp.drawmeridians(lon_labels,labels=[0,0,0,1]) \n",
    "        else:\n",
    "            bmp.drawparallels(lat_labels)\n",
    "            bmp.drawmeridians(lon_labels) \n",
    "        \n",
    "        # political boundaries.\n",
    "        bmp.drawstates()\n",
    "        bmp.drawcoastlines()\n",
    "        bmp.drawcounties()\n",
    "    \n",
    "        ## Difference from mtclim values\n",
    "        if not pr == 'mtclim':\n",
    "            ax = plt.subplot(gs[1,ind])\n",
    "            # Lat/Lon handling - product coords\n",
    "            lon_rad,lat_rad = np.meshgrid(regrid_ds[pr].longitude.values,regrid_ds[pr].latitude.values)\n",
    "            if regrid_ds[pr].SWdwn.loc[d:d].any():\n",
    "                SWdiff_for_plot = np.squeeze(np.squeeze(regrid_ds[pr].SWdwn.loc[d:d].values-\\\n",
    "                                          regrid_ds['mtclim'].SWdwn.loc[d:d].values))\n",
    "                SWdiff_for_plot = np.ma.masked_where(np.isnan(SWdiff_for_plot),SWdiff_for_plot)\n",
    "                im_dif = bmp.pcolormesh(lon_rad,lat_rad,SWdiff_for_plot,\\\n",
    "                        cmap=cmap_delta,vmin=SWmin_delta,vmax=SWmax_delta,shading='flat',latlon=True)\n",
    "            ax.set_title((pr+\" - mtclim\"))\n",
    "            \n",
    "            ## Format\n",
    "            if ind == 1:\n",
    "                bmp.drawparallels(lat_labels,labels=[1,0,0,0])\n",
    "            else:\n",
    "                bmp.drawparallels(lat_labels)\n",
    "            bmp.drawmeridians(lon_labels,labels=[0,0,0,1]) \n",
    "            \n",
    "            # political boundaries.\n",
    "            bmp.drawstates()\n",
    "            bmp.drawcoastlines()\n",
    "            bmp.drawcounties()\n",
    "            \n",
    "    ## Final formatting\n",
    "    plt.tight_layout\n",
    "    \n",
    "    # Colorbar - monthly values\n",
    "    caxi=plt.subplot(gs[0,-1])\n",
    "    cbar = plt.colorbar(im_avg, cax=caxi, orientation = \"vertical\",\\\n",
    "                        ticks=np.arange(SWmin,SWmax+dSW,dSW),spacing='proportional')\n",
    "    cbar.ax.set_ylabel(('Irradiance (Wm$^{-2}$)'))\n",
    "    \n",
    "    # Colorbar - difference\n",
    "    caxi=plt.subplot(gs[1,-1])\n",
    "    cbar = plt.colorbar(im_dif, cax=caxi, orientation = \"vertical\",spacing='proportional')\n",
    "    cbar.ax.set_ylabel(('Difference (Wm$^{-2}$)'))\n",
    "    \n",
    "    # Output\n",
    "    os.chdir(dir_print+'/MonthlyProductDiff')\n",
    "    outdate = pd.to_datetime(d) \n",
    "    outdate = outdate.strftime('%Y_%m')\n",
    "    fname = 'ProdDiff.'+outdate+'.png'\n",
    "    fig.savefig(fname)\n",
    "#     plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "################################################\n",
    "###### Summary of differences from mtclim ######\n",
    "pr_for_plot = ['nldas','syn','wrf']\n",
    "colsche =   np.array([[68.,119.,170.],\\\n",
    "            [221.,204.,119.],\\\n",
    "            [204.,102.,119.]])\n",
    "colsche = np.divide(colsche,256.)\n",
    "month_label = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec']\n",
    "\n",
    "## Loop through dates\n",
    "for YYYY in np.arange(2009,2010):\n",
    "    f, axes = plt.subplots(3, 4, figsize=(15, 9))#, sharex=True, sharey=True)\n",
    "    for ax,m in zip(axes.flat, np.arange(1,13)):            \n",
    "        # Get values for year\n",
    "        if m == 12:\n",
    "            d = pd.date_range(start=(str(YYYY)+'-'+str(m)+'-1'),end=(str(YYYY+1)+'-'+str(1)+'-1'),freq='M')\n",
    "        else:\n",
    "            d = pd.date_range(start=(str(YYYY)+'-'+str(m)+'-1'),end=(str(YYYY)+'-'+str(m+1)+'-1'),freq='M')\n",
    "        ## Difference from mtclim values\n",
    "        for n,pr in enumerate(pr_for_plot):        \n",
    "            if regrid_ds[pr].SWdwn.loc[d[0]:d[0]].any():\n",
    "                SWdiff_for_plot = np.squeeze(np.squeeze(regrid_ds[pr].SWdwn.loc[d[0]:d[0]].values-\\\n",
    "                                          regrid_ds['mtclim'].SWdwn.loc[d[0]:d[0]].values))\n",
    "                ax.plot(np.squeeze(wrf.elev[0,:,:]),SWdiff_for_plot,'.',color=colsche[n],label=pr)\n",
    "            ax.set_title(pr+\" - mtclim\")\n",
    "            ax.plot(np.array([0,5000]),np.array([0,0]),'k--')\n",
    "        # Legend\n",
    "        if m == 1:\n",
    "            pr_names_for_legend = [pr for pr in pr_for_plot if not pr == 'grobs']\n",
    "            handle = [matplotlib.patches.Patch(color=colsche[n], label=pr_names_for_legend[n]) \\\n",
    "                      for n,p in enumerate(pr_names_for_legend)]\n",
    "            ax.legend(handles=handle,ncol=3,labelspacing=.25,columnspacing=.3,markerscale=.1)\n",
    "        # Labels\n",
    "        ax.set_title(month_label[m-1])\n",
    "        if np.mod((m+3),4) == 0:\n",
    "            ax.set_ylabel('diff. from mtclim (Wm$^{-2}$)')\n",
    "        if m > 8:        \n",
    "            ax.set_xlabel('Elevation (m)')\n",
    "        else:\n",
    "            ax.tick_params(axis='x',which='both',labelbottom='off')\n",
    "        ax.set_ylim(-50,100)\n",
    "        ax.set_xlim(0,3500)\n",
    "            \n",
    "    ## Final formatting\n",
    "    plt.tight_layout\n",
    "    \n",
    "    # Output\n",
    "    os.chdir(dir_print)\n",
    "    fname = 'ProdDiff_elevation.'+str(YYYY)+'.png'\n",
    "    f.savefig(fname)\n",
    "    if YYYY < 2006:\n",
    "        plt.close(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "######################################################\n",
    "## Monthly diff shortwave w/ elevation - histograms ##\n",
    "######################################################\n",
    "#### Organize data for plotting\n",
    "## Difference structure\n",
    "diff_ds = xray.Dataset()\n",
    "for pr in pr_names:\n",
    "    if pr == 'mtclim':\n",
    "        continue\n",
    "    # Difference\n",
    "    diff_ds[pr] = align_ds[pr]-align_ds['mtclim']\n",
    "    \n",
    "## Digitize elevation\n",
    "bins = np.append(0,np.arange(500,3000,500))\n",
    "bin_centers = bins[1:]\n",
    "bin_labels = bin_centers\n",
    "\n",
    "nlat = diff_ds.latitude.size\n",
    "nlon = diff_ds.longitude.size\n",
    "ntime = diff_ds.time.size\n",
    "diff_ds['elev'] = (('latitude','longitude'),wrf.elev.values[0,:,:])\n",
    "\n",
    "elev = diff_ds.elev.values.flatten()\n",
    "digitized = np.digitize(elev, bins)\n",
    "bin_centers = np.append(bin_centers,np.nan)\n",
    "\n",
    "## Reshape to 2D array of time x space\n",
    "time = diff_ds.time\n",
    "plot_ds = xray.Dataset()\n",
    "for pr in pr_names:\n",
    "    if pr == 'mtclim':\n",
    "        continue\n",
    "    plot_ds[pr] = (('time','digi_elev'),diff_ds[pr].values.reshape((ntime,nlat*nlon)))    \n",
    "\n",
    "plot_ds.coords['time'] = ('time',time)\n",
    "plot_ds.coords['digi_elev'] = ('digi_elev',bin_centers[digitized-1])\n",
    "\n",
    "#### Function required variables\n",
    "bin_centers = bin_centers[:-1]\n",
    "grouping = bin_centers\n",
    "grouping_var = 'digi_elev'\n",
    "grouping_labels = bin_labels\n",
    "plot_var = np.array(['nldas','syn','wrf'])\n",
    "print(plot_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#### All years -- monthly\n",
    "plot_interval = 'monthly'\n",
    "plot_period = np.arange(2003,2013)\n",
    "fig = vp_xray(plot_ds,plot_var,plot_interval,plot_period,\\\n",
    "            bin_centers,\\\n",
    "            grouping,grouping_var,grouping_labels)\n",
    "# Output\n",
    "os.chdir(dir_print)\n",
    "for f in fig:\n",
    "    fname = 'ProdDiffElev.Violin.Monthly.'+str(f)+'.png'\n",
    "    fig[f].savefig(fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#### All years -- seasonal\n",
    "plot_interval = 'seasonal'\n",
    "fig = vp_xray(plot_ds,plot_var,plot_interval,plot_period,\\\n",
    "            bin_centers,\\\n",
    "            grouping,grouping_var,grouping_labels)\n",
    "# Output\n",
    "os.chdir(dir_print)\n",
    "for f in fig:\n",
    "    fname = 'ProdDiffElev.Violin.Seasonal.'+str(f)+'.png'\n",
    "    fig[f].savefig(fname)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}