{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# NCL_vector_4.py\nPlot U & V vectors globally, colored according to temperature\n\nThis script illustrates the following concepts:\n  - Coloring vectors based on temperature data\n  - Changing the scale of the vectors on the plot\n\nSee following URLs to see the reproduced NCL plot & script:\n    - Original NCL script: https://www.ncl.ucar.edu/Applications/Scripts/vector_4.ncl\n    - Original NCL plot: https://www.ncl.ucar.edu/Applications/Images/vector_4_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\nfrom matplotlib import pyplot as plt\nimport cartopy\nimport cartopy.crs as ccrs\n\nimport geocat.datafiles as gdf\nfrom geocat.viz import cmaps as gvcmaps\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 a netCDF data file using xarray default engine and load the data into xarrays\nfile_in = xr.open_dataset(gdf.get(\"netcdf_files/83.nc\"))\n\n# Extract slices of lon and lat for first timestamp and 13th lev\nds = file_in.isel(time=0, lev=12, lon=slice(0, -1, 5), lat=slice(2, -1, 3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Because there is no equivalent to ``CurlyVector`` in ``geocat.viz`` yet,\n# this plot does not look as identical as the NCL version.\n\n# Generate figure (set its size (width, height) in inches)\nfig, ax = plt.subplots(figsize=(10, 7.25))\n\n# Generate axes using Cartopy projection\nax = plt.axes(projection=ccrs.PlateCarree())\n\n# Import an NCL colormap and truncate it for a range and color levels\nplt.cm.register_cmap(\n    'BlAqGrYeOrReVi200',\n    gvutil.truncate_colormap(gvcmaps.BlAqGrYeOrReVi200,\n                             minval=0.03,\n                             maxval=0.95,\n                             n=16))\ncmap = plt.cm.get_cmap('BlAqGrYeOrReVi200', 16)\n\n# Draw vector plot\n# (there is no matplotlib equivalent to \"CurlyVector\" yet)\nQ = plt.quiver(ds['lon'],\n               ds['lat'],\n               ds['U'].data,\n               ds['V'].data,\n               ds['T'].data,\n               cmap=cmap,\n               zorder=1,\n               pivot=\"middle\",\n               width=0.001)\nplt.clim(228, 292)\n\n# Draw legend for vector plot\nax.add_patch(\n    plt.Rectangle((150, -140),\n                  30,\n                  30,\n                  facecolor='white',\n                  edgecolor='black',\n                  clip_on=False))\nqk = ax.quiverkey(Q,\n                  0.93,\n                  0.06,\n                  10,\n                  r'10 $m/s$',\n                  labelpos='N',\n                  coordinates='figure',\n                  color='black')\n\n# Use geocat.viz.util convenience function to add minor and major tick lines\ngvutil.add_major_minor_ticks(ax, labelsize=12)\n\n# Use geocat.viz.util convenience function to make plots look like NCL plots by using latitude, longitude tick labels\ngvutil.add_lat_lon_ticklabels(ax)\n\n# Set major and minor ticks\nplt.xticks(range(-180, 181, 30))\nplt.yticks(range(-90, 91, 30))\n\n# Use geocat.viz.util convenience function to add titles to left and right of the plot axis.\ngvutil.set_titles_and_labels(ax,\n                             maintitle=\"Vectors colored by a scalar map\",\n                             lefttitle=\"Temperature\",\n                             righttitle=\"$^{\\circ}$K\")\n\ncax = plt.axes((0.225, 0.075, 0.55, 0.025))\ncbar = fig.colorbar(Q,\n                    ax=ax,\n                    cax=cax,\n                    orientation='horizontal',\n                    ticks=range(232, 289, 8),\n                    drawedges=True)\n\n# Turn on continent shading\nax.add_feature(cartopy.feature.LAND,\n               edgecolor='lightgray',\n               facecolor='lightgray',\n               zorder=0)\n\n# Generate plot!\nplt.tight_layout()\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
}