{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://xarray.pydata.org/en/stable/_static/dataset-diagram-logo.png\" align=\"center\" width=\"30%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 03_Xarray: Multiple Files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "1. Work with multiple Xarray `Dataset`s\n",
    "2. Subset the Dataset along its dimensions\n",
    "3. Perform unit conversions\n",
    "4. Create a well-labeled multi-parameter contour plot of gridded CFSR reanalysis data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from datetime import datetime as dt\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 two Xarray `Dataset`s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "year = 1993\n",
    "ds_slp = xr.open_dataset(f'/cfsr/data/{year}/pmsl.1993.0p5.anl.nc')\n",
    "ds_g = xr.open_dataset(f'/cfsr/data/{year}/g.1993.0p5.anl.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look at the sizes of these datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (f'Size of SLP dataset: {ds_slp.nbytes / 1e9} GB')\n",
    "print (f'Size of geopotential height dataset: {ds_g.nbytes / 1e9} GB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting Datasets are on the order of 50 GB in size altogether, but we aren't noticing any latency thus far when accessing this Dataset. Xarray employs what's called [lazy loading](http://xarray.pydata.org/en/stable/io.html). Key point:\n",
    "<blockquote>\n",
    "    \"Data is always loaded lazily from netCDF files. You can manipulate, slice and subset Dataset and DataArray objects, and no array values are loaded into memory until you try to perform some sort of actual computation.\"</blockquote>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subset the Datasets along their dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We noticed in the previous notebook that our contour labels were not appearing with every contour line. This is because we passed the entire horizontal extent (all latitudes and longitudes) to the `ax.contour` method. Since our intent is to plot only over a regional subset, we will use the `sel` method on the latitude and longitude dimensions as well as time and isobaric surface.\n",
    "\n",
    "We'll retrieve two data variables, geopotential and sea-level pressure, from the Dataset.\n",
    "\n",
    "We'll also use Datetime and string methods to more dynamically assign various dimensional specifications, as well as aid in figure-labeling later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Areal extent\n",
    "lonW = -100\n",
    "lonE = -60\n",
    "latS = 25\n",
    "latN = 55\n",
    "cLat, cLon = (latS + latN)/2, (lonW + lonE)/2\n",
    "\n",
    "expand = 1\n",
    "latRange = np.arange(latS - expand,latN + expand,.5) # expand the data range a bit beyond the plot range\n",
    "lonRange = np.arange((lonW - expand),(lonE + expand),.5) # Need to match longitude values to those of the coordinate variable\n",
    "\n",
    "# Vertical level specificaton\n",
    "plevel = 500\n",
    "levelStr = str(plevel)\n",
    "\n",
    "# Date/Time specification\n",
    "Year = 1993\n",
    "Month = 3\n",
    "Day = 14\n",
    "Hour = 12\n",
    "Minute = 0\n",
    "dateTime = dt(Year,Month,Day, Hour, Minute)\n",
    "timeStr = dateTime.strftime(\"%Y-%m-%d %H%M UTC\")\n",
    "\n",
    "# Data variable selection\n",
    "\n",
    "Z = ds_g['g'].sel(time=dateTime,lev=plevel,lat=latRange,lon=lonRange)\n",
    "SLP = ds_slp['pmsl'].sel(time=dateTime,lat=latRange,lon=lonRange) # of course, no isobaric surface for SLP\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "levelStr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timeStr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at some of the attributes"
   ]
  },
  {
   "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": [
    "As a result of selecting just a single time and isobaric surface, both of those dimensions have been abstracted out of the DataArray."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z.units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SLP.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SLP.dims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SLP.units"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define our subsetted coordinate arrays of lat and lon. Pull them from either of the two DataArrays. We'll need to pass these into the contouring functions later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lats = Z.lat\n",
    "lons = Z.lon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lons"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform unit conversions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert geopotential meters to decameters, and Pascals to hectoPascals."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We take the DataArrays and apply [MetPy's unit conversion method](https://unidata.github.io/MetPy/latest/tutorials/xarray_tutorial.html#units)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SLP = SLP.metpy.convert_units('hPa')\n",
    "Z = Z.metpy.convert_units('dam')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">These conversions involve actual calculations, so in that cell, we are not *lazy-loading*. But that is the first instance where actual bytes of anything other than metadata get transfered ... <i>and we've already subset our original 48 GB and 2 GB Datasets into something much smaller</i>.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SLP.nbytes / 1e6, Z.nbytes / 1e6 # size in MB for the subsetted DataArrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a well-labeled multi-parameter contour plot of gridded CFSR reanalysis data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will make contour fills of 500 hPa height in decameters, with a contour interval of 6 dam, and contour lines of SLP in hPa, contour interval = 4."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we've done before, let's first define some variables relevant to Cartopy. Recall that we already defined the areal extent up above when we did the data subsetting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 = 474\n",
    "maxVal = 606\n",
    "cint = 6\n",
    "Zcintervals = np.arange(minVal, maxVal, cint)\n",
    "Zcintervals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minVal = 900\n",
    "maxVal = 1080\n",
    "cint = 4\n",
    "SLPcintervals = np.arange(minVal, maxVal, cint)\n",
    "SLPcintervals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the map, with filled contours of 500 hPa geopotential heights, and contour lines of SLP."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a meaningful title string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tl1 = f\"CFSR {levelStr} hPa heights (dam, filled contours) and SLP (lines, hPa)\"\n",
    "tl2 = f\"Valid at: {timeStr}\"\n",
    "title_line = (tl1 + '\\n' + tl2 + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_map = ccrs.LambertConformal(central_longitude=cLon, central_latitude=cLat)\n",
    "proj_data = ccrs.PlateCarree()\n",
    "res = '50m'\n",
    "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",
    "CF = ax.contourf(lons,lats,Z, levels=Zcintervals,transform=proj_data,cmap=plt.get_cmap('coolwarm'))\n",
    "cbar = plt.colorbar(CF,shrink=0.5)\n",
    "cbar.ax.tick_params(labelsize=16)\n",
    "cbar.ax.set_ylabel(\"Height (dam)\",fontsize=16)\n",
    "\n",
    "CL = ax.contour(lons,lats,SLP,SLPcintervals,transform=proj_data,linewidths=1.25,colors='green')\n",
    "ax.clabel(CL, inline_spacing=0.2, fontsize=11, fmt='%.0f')\n",
    "title = plt.title(title_line,fontsize=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're just missing the outer longitudes at higher latitudes. We could do one of two things to resolve this:\n",
    "\n",
    "1. Re-subset our original datset by extending the longitudinal range\n",
    "1. Slightly constrain the map plotting region"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try the latter here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constrainLon = 7 # trial and error!\n",
    "\n",
    "proj_map = ccrs.LambertConformal(central_longitude=cLon, central_latitude=cLat)\n",
    "proj_data = ccrs.PlateCarree()\n",
    "res = '50m'\n",
    "fig = plt.figure(figsize=(18,12))\n",
    "ax = plt.subplot(1,1,1,projection=proj_map)\n",
    "ax.set_extent ([lonW+constrainLon,lonE-constrainLon,latS,latN])\n",
    "ax.add_feature(cfeature.COASTLINE.with_scale(res))\n",
    "ax.add_feature(cfeature.STATES.with_scale(res))\n",
    "CF = ax.contourf(lons,lats,Z,levels=Zcintervals,transform=proj_data,cmap=plt.get_cmap('coolwarm'))\n",
    "cbar = plt.colorbar(CF,shrink=0.5)\n",
    "cbar.ax.tick_params(labelsize=16)\n",
    "cbar.ax.set_ylabel(\"Height (dam)\",fontsize=16)\n",
    "\n",
    "CL = ax.contour(lons,lats,SLP,SLPcintervals,transform=proj_data,linewidths=1.25,colors='green')\n",
    "ax.clabel(CL, inline_spacing=0.2, fontsize=11, fmt='%.0f')\n",
    "title = plt.title(title_line,fontsize=16)"
   ]
  }
 ],
 "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
}