{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# NCL_overlay_6.py\nThis script illustrates the following concepts:\n   - Overlaying filled contours, streamlines, and vectors over the same map\n   - Adding various map elements to a figure\n   - Using inset_axes() to create additional axes for color bars\n   - Creating custom label formats for colorbars\n   - Creating a quiverkey\n   - Assigning a colormap to contour and quiver plots\n   - Add arrows to streamlines\n   - Using zorder to specify the order in which elements will be drawn\n   \nSee following URLs to see the reproduced NCL plot & script:\n    - Original NCL script: https://www.ncl.ucar.edu/Applications/Scripts/overlay_6.ncl\n    - Original NCL plots: https://www.ncl.ucar.edu/Applications/Images/overlay_6_lg.png\n\nDifferences between NCL example and this one:\n    In the NCL version of this plot the vectors for the winds are nearly\n    uniform in length. Given the reference vector in that figure, the wind\n    speeds appear to be near 20 units. A histogram reveals that this is not a\n    true representation of the data as the magnitudes of the majority of wind\n    vectors are between 3 and 6 units with only a handful being greater than 13\n    and only one near 20. Because of this, we have chosen not to manipulate the\n    vector glyphs to appear more uniform as this would poorly represent the\n    data and be misleading. The lengths of the vectors in this examples are\n    proportional to the wind magnitudes where in the NCL examples they are not,\n    which is why the reference vector in this Python example is longer than the\n    NCL example and why the lengths vary more between the minimum and maximum\n    wind speeds.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import packages:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport xarray as xr\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as mpatches\nfrom mpl_toolkits.axes_grid1.inset_locator import inset_axes\nimport cartopy.crs as ccrs\nimport cartopy.feature as cfeature\nimport matplotlib.colors as mcolors\n\nimport geocat.datafiles as gdf\nfrom geocat.viz import cmaps as gvcmaps"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read in data:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Open a netCDF data file using xarray default engine and load the data into xarrays\nuf = xr.open_dataset(gdf.get(\"netcdf_files/Ustorm.cdf\"))\nvf = xr.open_dataset(gdf.get(\"netcdf_files/Vstorm.cdf\"))\npf = xr.open_dataset(gdf.get(\"netcdf_files/Pstorm.cdf\"))\ntf = xr.open_dataset(gdf.get(\"netcdf_files/Tstorm.cdf\"))\nu500f = xr.open_dataset(gdf.get(\"netcdf_files/U500storm.cdf\"))\nv500f = xr.open_dataset(gdf.get(\"netcdf_files/V500storm.cdf\"))\n\np = pf.p.isel(timestep=0).drop('timestep')\nt = tf.t.isel(timestep=0).drop('timestep')\nu = uf.u.isel(timestep=0).drop('timestep')\nv = vf.v.isel(timestep=0).drop('timestep')\nu500 = u500f.u.isel(timestep=0).drop('timestep')\nv500 = v500f.v.isel(timestep=0).drop('timestep')\ntime = vf.timestep\n\n# Convert Pa to hPa\np = p / 100\n# Convert K to F\nt = (t - 273.15) * 9 / 5 + 32"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create plot:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(8, 7))\nproj = ccrs.LambertAzimuthalEqualArea(central_longitude=-100,\n                                      central_latitude=40)\n\n# Set axis projection\nax = plt.axes([0, 0.2, 0.8, 0.7], projection=proj)\n# Create inset axes for color bars\ncax1 = inset_axes(ax,\n                  width='5%',\n                  height='100%',\n                  loc='lower right',\n                  bbox_to_anchor=(0.125, 0, 1, 1),\n                  bbox_transform=ax.transAxes,\n                  borderpad=0)\n\ncax2 = inset_axes(ax,\n                  width='100%',\n                  height='7%',\n                  loc='lower left',\n                  bbox_to_anchor=(0, -0.15, 1, 1),\n                  bbox_transform=ax.transAxes,\n                  borderpad=0)\n\n# Set extent to include roughly the United States\nax.set_extent((-128, -58, 18, 65), crs=ccrs.PlateCarree())\n\n#\n# Add map features\n#\n# Using the zorder keyword, we can specify the order of the layering. Lower\n# numbers are plotted before higher ones. For example, the coastlines have\n# zorder=2 while the filled contours have zorder=1. This will draw the\n# coastlines on top of the filled contours.\ntransparent = (0, 0, 0, 0)  # RGBA value for a transparent color for lakes\nax.add_feature(cfeature.OCEAN, color='lightskyblue', zorder=0)\nax.add_feature(cfeature.LAND, color='silver', zorder=0)\nax.add_feature(cfeature.LAKES,\n               linewidth=0.5,\n               edgecolor=transparent,\n               facecolor='white',\n               zorder=0)\nax.add_feature(cfeature.LAKES,\n               linewidth=0.5,\n               edgecolor='black',\n               facecolor=transparent,\n               zorder=2)\nax.add_feature(cfeature.COASTLINE, linewidth=0.5, zorder=2)\n\n#\n# Plot pressure level contour\n#\np_cmap = gvcmaps.StepSeq25\npressure = p.plot.contourf(ax=ax,\n                           transform=ccrs.PlateCarree(),\n                           cmap=p_cmap,\n                           levels=np.arange(975, 1050, 5),\n                           add_colorbar=False,\n                           add_labels=False,\n                           zorder=1)\nplt.colorbar(pressure, cax=cax1, ticks=np.arange(980, 1045, 5))\n\n# Format color bar label\ncax1.yaxis.set_label_text(label='\\n'.join('Sea Level Pressure'),\n                          fontsize=14,\n                          rotation=0)\ncax1.yaxis.set_label_coords(-0.5, 0.9)\ncax1.tick_params(size=0)\n\n#\n# Overlay streamlines\n#\nwith np.errstate(\n        invalid='ignore'\n):  # Indeed not needed, just to get rid of warnings about numpy's NaN comparisons\n    streams = ax.streamplot(u500.lon,\n                            u500.lat,\n                            u500.data,\n                            v500.data,\n                            transform=ccrs.PlateCarree(),\n                            color='black',\n                            arrowstyle='-',\n                            linewidth=0.5,\n                            density=2,\n                            zorder=5)\n\n# Divide streamlines into segments\nseg = streams.lines.get_segments()\n# Determine how many arrows on each streamline, the placement, and angles of the arrows\nperiod = 7\narrow_x = np.array([seg[i][0, 0] for i in range(0, len(seg), period)])\narrow_y = np.array([seg[i][0, 1] for i in range(0, len(seg), period)])\narrow_dx = np.array(\n    [seg[i][1, 0] - seg[i][0, 0] for i in range(0, len(seg), period)])\narrow_dy = np.array(\n    [seg[i][1, 1] - seg[i][0, 1] for i in range(0, len(seg), period)])\n# Add arrows to streamlines\nq = ax.quiver(arrow_x,\n              arrow_y,\n              arrow_dx,\n              arrow_dy,\n              color='black',\n              angles='xy',\n              scale=1,\n              units='y',\n              minshaft=3,\n              headwidth=4,\n              headlength=2,\n              headaxislength=2,\n              zorder=5)\n\n#\n# Overlay wind vectors\n#\n# First thin the data so the vector grid is less cluttered\nlon_size = u['lon'].size\nlat_size = u['lat'].size\n\nx = u['lon'].data[0:lon_size:2]\ny = u['lat'].data[0:lat_size:2]\nu = u.data[0:lat_size:2, 0:lon_size:2]\nv = v.data[0:lat_size:2, 0:lon_size:2]\nt = t.data[0:lat_size:2, 0:lon_size:2]\n\n# Import and modify color map for vectors\nwind_cmap = gvcmaps.amwg_blueyellowred\nbounds = np.arange(-30, 120, 10)  # Sets where boundarys on color map will be\nnorm = mcolors.BoundaryNorm(bounds, wind_cmap.N)  # Assigns colors to values\n# Draw wind vectors\nwith np.errstate(\n        invalid='ignore'\n):  # Indeed not needed, just to get rid of warnings about numpy's NaN comparisons\n    Q = ax.quiver(x,\n                  y,\n                  u,\n                  v,\n                  t,\n                  transform=ccrs.PlateCarree(),\n                  norm=norm,\n                  headwidth=5,\n                  cmap=wind_cmap,\n                  zorder=4)\n   \n# Add color bar\nplt.colorbar(Q,\n             cax=cax2,\n             ticks=np.arange(-20, 110, 10),\n             orientation='horizontal')\n\n# Format color bar label\ncax2.xaxis.set_label_text(label='Surface Temperature', fontsize=14)\ncax2.xaxis.set_label_position('top')\ncax2.tick_params(size=0)\n\n# Add quiverkey and white patch behind it\nax.add_patch(\n    mpatches.Rectangle(xy=[0.85, 0],\n                       width=0.15,\n                       height=0.0925,\n                       facecolor='white',\n                       transform=ax.transAxes,\n                       zorder=3))\nax.quiverkey(Q, 0.925, 0.025, 20, label='20', zorder=4)\n\n# Add title\nax.set_title('January 1996 Snow Storm\\n1996 01 05 00:00 + 0',\n             fontweight='bold',\n             fontsize=16)\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
}