{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# NCL_conOncon_5.py\n",
    "This script illustrates the following concepts:\n",
    "   - Overlaying individual contour lines on a polar stereographic map\n",
    "   - Drawing a spaghetti contour plot\n",
    "   - Increasing the thickness of contour lines\n",
    "   - Explicitly setting contour levels\n",
    "   - Changing the color of a contour line\n",
    "\n",
    "See following URLs to see the reproduced NCL plot & script:\n",
    "   - [Original NCL script](https://www.ncl.ucar.edu/Applications/Scripts/conOncon_5.ncl)\n",
    "   - [Original NCL plot](https://www.ncl.ucar.edu/Applications/Images/conOncon_5_lg.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import packages:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import cartopy.feature as cfeature\n",
    "import cartopy.crs as ccrs\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as mticker\n",
    "\n",
    "import geocat.datafiles as gdf\n",
    "import geocat.viz as gv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read in data:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Open a netCDF data file using xarray default engine and\n",
    "# load the data into xarrays\n",
    "ds = xr.open_dataset(gdf.get(\"netcdf_files/HGT500_MON_1958-1997.nc\"),\n",
    "                     decode_times=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Generate a figure\n",
    "fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "# Create an axis with a polar stereographic projection\n",
    "ax = plt.axes(projection=ccrs.NorthPolarStereo())\n",
    "\n",
    "# Add land feature to map\n",
    "ax.add_feature(cfeature.LAND, facecolor='lightgray')\n",
    "\n",
    "# Set map boundary to include latitudes between 0 and 40 and longitudes\n",
    "# between -180 and 180 only\n",
    "gv.set_map_boundary(ax, [-180, 180], [0, 40], south_pad=1)\n",
    "\n",
    "# Set draw_labels to False so that you can manually manipulate it later\n",
    "gl = ax.gridlines(ccrs.PlateCarree(),\n",
    "                  draw_labels=False,\n",
    "                  linestyle=\"--\",\n",
    "                  linewidth=1,\n",
    "                  color='darkgray',\n",
    "                  zorder=2)\n",
    "\n",
    "# Manipulate latitude and longitude gridline numbers and spacing\n",
    "gl.ylocator = mticker.FixedLocator(np.arange(0, 90, 15))\n",
    "gl.xlocator = mticker.FixedLocator(np.arange(-180, 180, 30))\n",
    "\n",
    "# Manipulate longitude labels (0, 30 E, 60 E, ..., 30 W, etc.)\n",
    "ticks = np.arange(0, 210, 30)\n",
    "etick = ['0'] + [\n",
    "    r'%dE' % tick for tick in ticks if (tick != 0) & (tick != 180)\n",
    "] + ['180']\n",
    "wtick = [r'%dW' % tick for tick in ticks if (tick != 0) & (tick != 180)]\n",
    "labels = etick + wtick\n",
    "xticks = np.arange(0, 360, 30)\n",
    "yticks = np.full_like(xticks, -5)  # Latitude where the labels will be drawn\n",
    "for xtick, ytick, label in zip(xticks, yticks, labels):\n",
    "    if label == '180':\n",
    "        ax.text(xtick,\n",
    "                ytick,\n",
    "                label,\n",
    "                fontsize=14,\n",
    "                horizontalalignment='center',\n",
    "                verticalalignment='top',\n",
    "                transform=ccrs.Geodetic())\n",
    "    elif label == '0':\n",
    "        ax.text(xtick,\n",
    "                ytick,\n",
    "                label,\n",
    "                fontsize=14,\n",
    "                horizontalalignment='center',\n",
    "                verticalalignment='bottom',\n",
    "                transform=ccrs.Geodetic())\n",
    "    else:\n",
    "        ax.text(xtick,\n",
    "                ytick,\n",
    "                label,\n",
    "                fontsize=14,\n",
    "                horizontalalignment='center',\n",
    "                verticalalignment='center',\n",
    "                transform=ccrs.Geodetic())\n",
    "\n",
    "# Get slice of data at the 0th timestep - plot this contour line separately\n",
    "# because it will be thicker than the other contour lines\n",
    "p = ds.HGT.isel(time=0)\n",
    "\n",
    "# Use geocat-viz utility function to handle the no-shown-data\n",
    "# artifact of 0 and 360-degree longitudes\n",
    "slon = gv.xr_add_cyclic_longitudes(p, \"lon\")\n",
    "\n",
    "# Plot contour data at pressure level 5500 at the first timestep\n",
    "p = slon.plot.contour(ax=ax,\n",
    "                      transform=ccrs.PlateCarree(),\n",
    "                      linewidths=1.5,\n",
    "                      levels=[5500],\n",
    "                      colors='black',\n",
    "                      add_labels=False)\n",
    "\n",
    "# Create a color list for each of the next 18 contours\n",
    "colorlist = [\n",
    "    \"crimson\", \"green\", \"blue\", \"yellow\", \"cyan\", \"hotpink\", \"crimson\",\n",
    "    \"skyblue\", \"navy\", \"lightyellow\", \"mediumorchid\", \"orange\", \"slateblue\",\n",
    "    \"palegreen\", \"magenta\", \"springgreen\", \"pink\", \"forestgreen\", \"violet\"\n",
    "]\n",
    "\n",
    "# Iterate through 18 different timesteps\n",
    "for x in range(18):\n",
    "\n",
    "    # Get a slice of data at the 12*x+1 timestep\n",
    "    p = ds.HGT.isel(time=12 * x + 1)\n",
    "\n",
    "    # Use geocat-viz utility function to handle the no-shown-data artifact\n",
    "    # of 0 and 360-degree longitudes\n",
    "    slon = gv.xr_add_cyclic_longitudes(p, \"lon\")\n",
    "\n",
    "    # Plot contour data at pressure level 5500 for the 12*x+1 timestep\n",
    "    p = slon.plot.contour(ax=ax,\n",
    "                          transform=ccrs.PlateCarree(),\n",
    "                          linewidths=0.5,\n",
    "                          levels=[5500],\n",
    "                          colors=colorlist[x],\n",
    "                          add_labels=False)\n",
    "\n",
    "# Use geocat.viz.util convenience function to add titles\n",
    "gv.set_titles_and_labels(ax,\n",
    "                         maintitle=r\"$\\bf{Spaghetti}$\" + \" \" + r\"$\\bf{Plot}$\",\n",
    "                         lefttitle=slon.long_name,\n",
    "                         righttitle=slon.units)\n",
    "\n",
    "# Make tight layout\n",
    "plt.tight_layout()\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}