{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# NCL_polyg_2.py\nConcepts illustrated:\n  - Drawing a Lambert Conformal U.S. map color-coded by climate divisions\n  - Color-coding climate divisions based on precipitation values\n  - Drawing the climate divisions of the U.S.\n  - Zooming in on a particular area on a Lambert Conformal map\n  - Drawing a border around filled polygons\n  - Masking the ocean in a map plot\n  - Masking land in a map plot\n  - Increasing the font size of text\n  - Adding text to a plot\n  - Drawing a custom labelbar on a map\n  - Creating a red-yellow-blue color map\n\nSee following URLs to see the reproduced NCL plot & script:\n    - Original NCL script: https://www.ncl.ucar.edu/Applications/Scripts/polyg_2.ncl\n    - Original NCL plot: https://www.ncl.ucar.edu/Applications/Images/polyg_2_lg.png\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import packages:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import xarray as xr\nimport cartopy.crs as ccrs\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as colors\nimport matplotlib.cm as cm\nfrom geocat.viz import util as gvutil\n\nimport geocat.datafiles as gdf\nimport cartopy.io.shapereader as shpreader\nimport shapely.geometry as sgeom\nfrom mpl_toolkits.axes_grid1.inset_locator import inset_axes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read in data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Open climate division datafile and add to xarray\nds = xr.open_dataset(gdf.get(\"netcdf_files/climdiv_prcp_1899-1999.nc\"),\n                     decode_times=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initialize color map and bounds for each color\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "colormap = colors.ListedColormap([\n    'mediumpurple', 'mediumblue', 'royalblue', 'cornflowerblue', 'lightblue',\n    'lightseagreen', 'yellowgreen', 'green', 'wheat', 'tan', 'gold', 'orange',\n    'red', 'firebrick'\n])\n\n# Values represent average number of inches of rain\ncolorbounds = [0, 5, 10, 15, 20, 25, 30, 35, 40, 50, 60, 70, 80, 90, 100]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define helper function to determine which color to fill the divisions based on precipitation data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def findDivColor(colorbounds, pdata):\n\n    for x in range(len(colorbounds)):\n\n        if pdata >= colorbounds[len(colorbounds) - 1]:\n            return colormap.colors[x - 1]\n        if pdata >= colorbounds[x]:\n            continue\n        else:\n            # Index is 'x-1' because colorbounds is one item longer than colormap\n            return colormap.colors[x - 1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create plot\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Create plot figure\nfig = plt.figure(figsize=(8, 8))\n\n# Add axes for lambert conformal map\n# Set dimensions of axes with [X0, Y0, width, height] argument. Each value is a fraction of total figure size.\nax = plt.axes([.05, -.05, .9, 1],\n              projection=ccrs.LambertConformal(),\n              frameon=False)\n\n# Set latitude and longitude extent of map\nax.set_extent([-119, -74, 18, 50], ccrs.Geodetic())\n\n# Set shape name of map (which depicts the United States)\nshapename = 'admin_1_states_provinces_lakes_shp'\nstates_shp = shpreader.natural_earth(resolution='110m',\n                                     category='cultural',\n                                     name=shapename)\n\n# Set title and title fontsize of plot using gvutil function instead of matplotlib function call\ngvutil.set_titles_and_labels(\n    ax,\n    maintitle=\n    \"Average Annual Precipiation \\n Computed for the period 1899-1999 \\n NCDC climate division data \\n\",\n    maintitlefontsize=18)\n\n# Add outlines of each state within the United States\nfor state in shpreader.Reader(states_shp).geometries():\n\n    facecolor = 'white'\n    edgecolor = 'black'\n\n    ax.add_geometries([state],\n                      ccrs.PlateCarree(),\n                      facecolor=facecolor,\n                      edgecolor=edgecolor)\n\n# For each variable (climate division) in data set, create outline on map and fill with random color\nfor varname, da in ds.data_vars.items():\n\n    # This condition is included because first item in xarray only has one attribute, 'current date'\n    if hasattr(da, 'state_name'):\n        # Get number of years of data by dividing number of months recorded (length of array) by 12 (12 months per year)\n        numYears = len(da.values) / 12\n\n        # Get precipitation data for each climate division:\n        # Rather than looping through the whole array to find the sum of each 12 values (a year's worth of data),\n        # adding each sum to an array, and then finding the average of the values in the array, as seen in the NCL\n        # script, the one-line python method involves summing the dataset values and then dividing it by numYears (calculated in line 99)\n        precipitationdata = sum(da.values) / numYears\n\n        # Get borders of each climate division\n        lat = da.lat\n        lon = da.lon\n\n        # Get color of climate division\n        color = findDivColor(colorbounds, precipitationdata)\n\n        # Use \"shapely geometry\" module to create division outlines from lat/lon coordinates\n        track = sgeom.LineString(zip(lon, lat))\n\n        # Add division outlines to map\n        im = ax.add_geometries([track],\n                               ccrs.PlateCarree(),\n                               facecolor=color,\n                               edgecolor='black',\n                               linewidths=.5)\n\n# Create and plot colorbar\n\n# Map colors to bounds\nnorm = colors.BoundaryNorm(colorbounds, colormap.N)\n\n# Add inset axes (axes within pre-existing axes) to hold colorbar\naxins1 = inset_axes(ax, width=\"75%\", height=\"3%\", loc='lower center')\n\n# Add colorbar to plot\ncb = fig.colorbar(cm.ScalarMappable(cmap=colormap, norm=norm),\n                  cax=axins1,\n                  boundaries=colorbounds,\n                  ticks=colorbounds,\n                  spacing='uniform',\n                  orientation='horizontal',\n                  label='inches')\n\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}