{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# NCL_polyg_19.py\nThis script illustrates the following concepts:\n   - Adding lines and polygons to a map\n   - Adding a map to another map as an annotation\n   - Coloring shapefile outlines based on an array of values\n   - Drawing a custom colorbar on a map\n   - Using functions for cleaner code\n   - Overlaying a shape from one shapefile over another\nSee following URLs to see the reproduced NCL plot & script:\n    - Original NCL script: https://www.ncl.ucar.edu/Applications/Scripts/polyg_19.ncl\n    - Original NCL plot: https://www.ncl.ucar.edu/Applications/Images/polyg_19_lg.png\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import packages:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nfrom matplotlib.collections import PatchCollection\nfrom matplotlib.patches import Polygon\nimport matplotlib.colors as colors\nfrom mpl_toolkits.axes_grid1.inset_locator import inset_axes\nimport matplotlib.cm as cm\nimport shapefile as shp\nimport numpy as np\n\nimport geocat.datafiles as gdf\nfrom geocat.viz import util as gvutil"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read in data:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Open all shapefiles and associated .dbf, .shp, and .prj files so sphinx can run and generate documents\nfile1 = open(gdf.get(\"shape_files/gadm36_USA_1.dbf\"), 'r')\nfile2 = open(gdf.get(\"shape_files/gadm36_USA_1.shp\"), 'r')\nfile3 = open(gdf.get(\"shape_files/gadm36_USA_1.shx\"), 'r')\nfile4 = open(gdf.get(\"shape_files/gadm36_USA_1.prj\"), 'r')\n\nfile5 = open(gdf.get(\"shape_files/gadm36_USA_2.dbf\"), 'r')\nfile6 = open(gdf.get(\"shape_files/gadm36_USA_2.shp\"), 'r')\nfile7 = open(gdf.get(\"shape_files/gadm36_USA_2.shx\"), 'r')\nfile8 = open(gdf.get(\"shape_files/gadm36_USA_2.prj\"), 'r')\n\nfile9 = open(gdf.get(\"shape_files/gadm36_PRI_0.dbf\"), 'r')\nfile10 = open(gdf.get(\"shape_files/gadm36_PRI_0.shp\"), 'r')\nfile11 = open(gdf.get(\"shape_files/gadm36_PRI_0.shx\"), 'r')\nfile12 = open(gdf.get(\"shape_files/gadm36_PRI_0.prj\"), 'r')\n\n# Open the text file with the population data\nstate_population_file = open(gdf.get(\"ascii_files/us_state_population.txt\"),\n                             'r')\n# Open shapefiles\nus = shp.Reader(gdf.get(\"shape_files/gadm36_USA_1.dbf\"))\nusdetailed = shp.Reader(gdf.get(\"shape_files/gadm36_USA_2.dbf\"))\npr = shp.Reader(gdf.get(\"shape_files/gadm36_PRI_0.dbf\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set colormap data and colormap bounds:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "colormap = colors.ListedColormap([\n    'lightpink', 'wheat', 'palegreen', 'powderblue', 'thistle', 'lightcoral',\n    'peru', 'dodgerblue', 'slateblue', 'firebrick', 'sienna', 'olivedrab',\n    'steelblue', 'navy'\n])\n\ncolorbounds = [0, 1, 2.5, 3, 4, 5, 6, 7, 8, 9, 10, 12, 25, 38, 40]\n\nnorm = colors.BoundaryNorm(colorbounds, colormap.N)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define helper function to get the populations of each state\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def getStatePopulations(state_population_file):\n\n    population_dict = {}\n    Lines = state_population_file.read().splitlines()\n    for line in Lines:\n        nameandpop = line.split(\" \")\n        if nameandpop[-1].isnumeric():\n            name = nameandpop[0]\n            pop = (int)(nameandpop[-1]) / 1000000\n            population_dict[name] = pop\n    return population_dict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define helper function to get the color of each state based on its population\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": [
        "Define helper function to remove ticks from axes\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def removeTicks(axis):\n\n    axis.get_xaxis().set_visible(False)\n    axis.get_yaxis().set_visible(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define helper function to plot and color each state\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def plotRegion(region, axis, xlim, puertoRico, waterBody):\n\n    # Create empty lists for filled polygons or\" patches\" and \"water_patches\"\n    patches = []\n    water_patches = []\n    \n    # Plot each shape within a region (ex. mainland Alaska and all of it's surrounding Alaskan islands)\n    for i in range(len(region.shape.parts)):\n\n        i_start = region.shape.parts[i]\n        if i == len(region.shape.parts) - 1:\n            i_end = len(region.shape.points)\n        else:\n            i_end = region.shape.parts[i + 1]\n\n        # Create new empty lists to hold lat and lon coordinates\n        x = []\n        y = []\n\n        # Get every coordinate within every shape (as long as it is within the x coordinate limits)\n        for i in region.shape.points[i_start:i_end]:\n            if xlim[0] is not None and i[0] < xlim[0]:\n                continue\n            if xlim[1] is not None and i[0] > xlim[1]:\n                continue\n            else:\n                x.append(i[0])\n                y.append(i[1])\n\n        # Plot outline of each region\n        axis.plot(x, y, color='black', linewidth=0.1, zorder=1)\n\n        # Fill each state with color:\n        # Determine the region to be plotted (Puerto Rico or United States)\n        if waterBody is False:\n            if puertoRico is False:\n                abbrevname = shape.record[-1].split(\".\")\n                abbrevstate = abbrevname[1]\n            else:\n                abbrevstate = 'PR'\n\n        # If the region being plotted is a state with a population\n        if waterBody is False:\n            pop = population_dict[abbrevstate]\n            color = findDivColor(colorbounds, pop)\n            # Set characteristics and measurements of each filled polygon \"patch\"\n            patches.append(\n                Polygon(np.vstack((x, y)).T, True, color=color, linewidth=0.1))\n\n        # If the region being plotted is a body of water with no population\n        else:\n            # Set characteristics and measurements of each filled polygon \"patch\"\n            water_patches.append(\n                Polygon(np.vstack((x, y)).T, True, color='white', linewidth=.7))\n\n    pc = PatchCollection(patches,\n                         match_original=True,\n                         edgecolor='k',\n                         linewidths=0.1,\n                         zorder=2)\n    # Plot filled region on axis\n    axis.add_collection(pc)\n\n    wpc = PatchCollection(water_patches,\n                         match_original=True,\n                         edgecolor='white',\n                         linewidth=.8,\n                         zorder=3)\n    # Plot filled region on axis\n    axis.add_collection(wpc)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Create figure\nfig = plt.figure(figsize=(8, 8))\nspec = gridspec.GridSpec(ncols=1,\n                         nrows=2,\n                         hspace=0.05,\n                         wspace=0.1,\n                         figure=fig,\n                         height_ratios=[2, 1])\n\n# Create upper axis\nax1 = fig.add_subplot(spec[0, 0], frameon=False)\nremoveTicks(ax1)\n\n# Create lower axis\nax2 = fig.add_subplot(spec[1, 0], frameon=False)\nremoveTicks(ax2)\n\n# Create three inset axes on lower axis for Alaska, Hawaii, and Puerto Rico respectively\naxin1 = ax2.inset_axes([0.0, 0.7, 0.30, 0.80], frameon=False)\nremoveTicks(axin1)\naxin2 = ax2.inset_axes([0.40, 0.7, 0.20, 0.40], frameon=False)\nremoveTicks(axin2)\naxin3 = ax2.inset_axes([0.70, 0.7, 0.30, 0.30], frameon=False)\nremoveTicks(axin3)\n\n# Get population of each state\npopulation_dict = getStatePopulations(state_population_file)\n\n# Plot every shape in the US shapefile\nfor shape in us.shapeRecords():\n\n    if shape.record[3] == 'Alaska':\n        plotRegion(shape, axin1, [None, 100], puertoRico=False, waterBody=False)\n    elif shape.record[3] == 'Hawaii':\n        plotRegion(shape,\n                   axin2, [-161, None],\n                   puertoRico=False,\n                   waterBody=False)\n    else:\n        plotRegion(shape, ax1, [None, None], puertoRico=False, waterBody=False)\n\n# Plot every shape in the puerto rico shapefile\nfor shape in pr.shapeRecords():\n    plotRegion(shape, axin3, [None, None], puertoRico=True, waterBody=False)\n\n# Plot every body of water shape in the detailed US shapefile\nfor shape in usdetailed.shapeRecords():\n    if shape.record[9] == 'Water body':\n        plotRegion(shape, ax1, [None, None], puertoRico=False, waterBody=True)\n\n# Set title using helper function from geocat-viz\ntitle = r\"$\\bf{Population}$\" + \" \" + r\"$\\bf{in}$\" + \" \" + r\"$\\bf{Millions}$\" + \" \" + r\"$\\bf{(2014)}$\"\ngvutil.set_titles_and_labels(ax1, maintitle=title, maintitlefontsize=18)\n\n# Create fourth inset axis for colorbar\naxin4 = inset_axes(ax2, width=\"115%\", height=\"12%\", loc='center')\n\n# Create colorbar\ncb = fig.colorbar(cm.ScalarMappable(cmap=colormap, norm=norm),\n                  cax=axin4,\n                  boundaries=colorbounds,\n                  ticks=colorbounds[1:-1],\n                  spacing='uniform',\n                  orientation='horizontal')\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
}