{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 03_Xarray: Overlays of variables"
   ]
  },
  {
   "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 ERA5 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 an Xarray `Dataset`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('/spare11/atm350/data/199303_era5.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Examine the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "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",
    "\n",
    "lonW = -100\n",
    "lonE = -60\n",
    "latS = 25\n",
    "latN = 55\n",
    "cLat, cLon = (latS + latN)/2, (lonW + lonE)/2\n",
    "\n",
    "# Recall that in ERA5, longitudes run between 0 and 360, not -180 and 180\n",
    "if (lonW < 0 ):\n",
    "    lonW = lonW + 360\n",
    "if (lonE < 0 ):\n",
    "    lonE = lonE + 360\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['geopotential'].sel(time=dateTime,level=plevel,latitude=latRange,longitude=lonRange)\n",
    "slp = ds['mean_sea_level_pressure'].sel(time=dateTime,latitude=latRange,longitude=lonRange) # of course, no isobaric surface for SLP"
   ]
  },
  {
   "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.latitude\n",
    "lons = z.longitude"
   ]
  },
  {
   "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 to geopotential height in 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 = mpcalc.geopotential_to_height(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": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a well-labeled multi-parameter contour plot of gridded ERA5 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\"ERA5 {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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 Jan. 2025 Environment",
   "language": "python",
   "name": "jan25"
  },
  "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}