{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://xarray.pydata.org/en/stable/_static/dataset-diagram-logo.png\" align=\"center\" width=\"50%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Xarray 3: Datasets and Plotting\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "1. The Xarray `Dataset`\n",
    "1. Open and select a variable from the ERA-5 dataset\n",
    "1. Convert geopotential to height\n",
    "1. Create a contour plot\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "| Concepts | Importance | Notes |\n",
    "| --- | --- | --- |\n",
    "| Xarray Lessons 1-2| Necessary | |\n",
    "\n",
    "* **Time to learn**: 15 minutes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "from metpy.units import units\n",
    "import metpy.calc as mpcalc\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Work with an Xarray `Dataset`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An Xarray [Dataset](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.html) consists of one or more `DataArray`s. Let's open the same NetCDF file as in the first Xarray notebook, but as a `Dataset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('/spare11/atm533/data/2012103000_z500_era5.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what this Dataset looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's represented somewhat similarly to a DataArray, but now we have support for multiple data variables. In this case,\n",
    "1. There is just one data variable: z (Geopotential)\n",
    "2. The coordinate variables are longitude, latitude, and time\n",
    "3. The data variable has three dimensions (time x latitude x longitude). \n",
    "4. Although time is a dimension, there is only one time value on that coordinate.\n",
    "5. The Dataset itself has attributes. The attributes of the coordinate and data variables are separately accessible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Select a data variable from the Dataset.</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = ds['z'] # Similar to selecting a Series in Pandas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can programmatically retrieve various attributes of the data variable that we browsed interactively above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z.dims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that Z is a 3-dimensional DataArray, with time as the first dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z.units"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert geopotential to height"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Print out the representation of our data variable. We can see it's a `DataArray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (Z) # Text-based, as opposed to HTML-prettified output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The units are already part of the data variable. We can use `MetPy`'s `units` and `calc` libraries to  [convert geopotential to height](https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.geopotential_to_height.html). We'll assign a new object to the output of the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HGHT = mpcalc.geopotential_to_height(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Print out its representation. Note that the array values have changed, and there is a unit attached to the `DataArray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (HGHT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\"><b>Note: </b>Why are we using `print`  to show what the HGHT object looks like? If we didn't, the array's values would be rendered in an HTML-friendly, aka *pretty-printed* manner ... while it looks nice, it can take frightfully long to render if the array is large.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can make a map with contours of 500 hPa height!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a contour plot of gridded data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use Cartopy and Matplotlib to plot our 500 hPa height contours. Matplotlib has two contour methods:\n",
    "1. Line contours: ax.contour\n",
    "1. Filled contours: ax.contourf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We will draw contour lines in meters, with a contour interval of 60 m."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### As we've done before, let's first define some variables relevant to Cartopy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cLon = -80\n",
    "cLat = 35\n",
    "lonW = -100\n",
    "lonE = -60\n",
    "latS = 20\n",
    "latN = 50\n",
    "proj_map = ccrs.LambertConformal(central_longitude=cLon, central_latitude=cLat)\n",
    "proj_data = ccrs.PlateCarree() # Our data is lat-lon; thus its native projection is Plate Carree.\n",
    "res = '50m'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now define the range of our contour values and a contour interval. 60 m is standard for 500 hPa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minVal = 4680\n",
    "maxVal = 6060\n",
    "cint = 60\n",
    "cintervals = np.arange(minVal, maxVal, cint)\n",
    "cintervals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matplotlib's contour methods require three arrays to be passed to them ... x- and y- arrays (longitude and latitude in our case), and a 2-d array (corresponding to x- and y-) of our data variable. So we need to extract the latitude and longitude coordinate variables from our DataArray. We'll also extract the third coordinate value, time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lats = ds.latitude \n",
    "lons = ds.longitude\n",
    "times = ds.time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're set to make our map. After creating our figure, setting the bounds of our map domain, and adding cartographic features, we will plot the contours. We'll assign the output of the contour method to an object, which will then be used to label the contour lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(18,12))\n",
    "ax = plt.subplot(1,1,1,projection=proj_map)\n",
    "ax.set_extent ([lonW,lonE,latS,latN])\n",
    "ax.add_feature(cfeature.COASTLINE.with_scale(res))\n",
    "ax.add_feature(cfeature.STATES.with_scale(res))\n",
    "CL = ax.contour(lons,lats,HGHT,cintervals,transform=proj_data,linewidths=1.25,colors='green')\n",
    "ax.clabel(CL, inline_spacing=0.2, fontsize=11, fmt='%.0f');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-danger\"><h3>Didn't work! The end of the traceback reads:</h3>\n",
    "    <code>TypeError: Input z must be 2D, not 3D</code>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "<h4>In other words, our geopotential height array is not two dimensional in x and y ... recall that it also has *time* as its first dimension! Even though there is only one time value in this dimension, we need to eliminate this dimension from Z.</h4>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll learn other ways to abstract out dimensions other than x and y (remember, there could be additional dimensions in our Datasets or DataArrays, such as vertical levels or ensemble members) in later notebooks, but for now, we can simply select the first (and only) element from the time dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HGHT = HGHT[0,:,:] # Select the first element of the leading (time) dimension, and all lons and lats."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HGHT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now try again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(18,12))\n",
    "ax = plt.subplot(1,1,1,projection=proj_map)\n",
    "ax.set_extent ([lonW,lonE,latS,latN])\n",
    "ax.add_feature(cfeature.COASTLINE.with_scale(res))\n",
    "ax.add_feature(cfeature.STATES.with_scale(res))\n",
    "CL = ax.contour(lons,lats,HGHT,cintervals,transform=proj_data,linewidths=1.25,colors='green')\n",
    "ax.clabel(CL, inline_spacing=0.2, fontsize=11, fmt='%.0f');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### There we have it! There are definitely some things we can improve about this plot (such as the lack of an informative title, and the fact that our contour labels seem not to be appearing on every contour line), but we'll get to that soon!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Summary\n",
    "* An Xarray `Dataset` consists of one or more `DataArrays`.\n",
    "* MetPy can perform unit-aware calculations and conversions on Xarray objects.\n",
    "* Since Xarray `DataArray`s often have more than x- and y- dimensions, care must be taken to reduce dimensionality for the purposes of Matplotlib functions such as `contour` and `contourf`.\n",
    "\n",
    "### What's Next?\n",
    "In the next notebook, we'll work with multiple `Dataset`s.\n",
    "\n",
    "## Resources and References\n",
    "1. [The Xarray Dataset](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.html)\n",
    "1. [Units in MetPy](https://unidata.github.io/MetPy/latest/tutorials/unit_tutorial.html#sphx-glr-tutorials-unit-tutorial-py)\n",
    "1. [MetPy Calculations](https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.html#module-metpy.calc)\n",
    "1. [Matplotlib Contour Line Plots](https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.contour.html)\n",
    "1. [Matplotlib Contour Fill Plots](https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.contourf.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 August 2022 Environment",
   "language": "python",
   "name": "aug22"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}