{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 01_GriddedDiagnostics_TempAdvection_CFSR: Compute derived quantities using MetPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## In this notebook, we'll cover the following:\n",
    "1. Select a date and access various CFSR Datasets\n",
    "2. Subset the desired Datasets along their dimensions\n",
    "3. Calculate and visualize diagnostic quantities.\n",
    "4. Smooth the diagnostic field."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <span style=\"color:purple\">0) Preliminaries </span>"
   ]
  },
  {
   "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": [
    "# <span style=\"color:purple\">1) Specify a starting and ending date/time, and access several CFSR Datasets</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "startYear = 2007\n",
    "startMonth = 2\n",
    "startDay = 14\n",
    "startHour = 12\n",
    "startMinute = 0\n",
    "startDateTime = dt(startYear,startMonth,startDay, startHour, startMinute)\n",
    "\n",
    "endYear = 2007\n",
    "endMonth = 2\n",
    "endDay = 14\n",
    "endHour = 12\n",
    "endMinute = 0\n",
    "endDateTime = dt(endYear,endMonth,endDay, endHour, endMinute)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Xarray `Dataset` objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsZ = xr.open_dataset (f'/cfsr/data/{startYear}/g.{startYear}.0p5.anl.nc')\n",
    "dsT = xr.open_dataset (f'/cfsr/data/{startYear}/t.{startYear}.0p5.anl.nc')\n",
    "dsU = xr.open_dataset (f'/cfsr/data/{startYear}/u.{startYear}.0p5.anl.nc')\n",
    "dsV = xr.open_dataset (f'/cfsr/data/{startYear}/v.{startYear}.0p5.anl.nc')\n",
    "dsW = xr.open_dataset (f'/cfsr/data/{startYear}/w.{startYear}.0p5.anl.nc')\n",
    "dsQ = xr.open_dataset (f'/cfsr/data/{startYear}/q.{startYear}.0p5.anl.nc')\n",
    "dsSLP = xr.open_dataset (f'/cfsr/data/{startYear}/pmsl.{startYear}.0p5.anl.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <span style=\"color:purple\">2) Specify a date/time range, and subset the desired `Dataset`s along their dimensions.</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a list of date and times based on what we specified for the initial and final times, using Pandas' date_range function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dateList = pd.date_range(startDateTime, endDateTime,freq=\"6H\")\n",
    "dateList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Areal extent\n",
    "lonW = -90\n",
    "lonE = -60\n",
    "latS = 35\n",
    "latN = 50\n",
    "cLat, cLon = (latS + latN)/2, (lonW + lonE)/2\n",
    "latRange = np.arange(latS,latN+.5,.5) # expand the data range a bit beyond the plot range\n",
    "lonRange = np.arange(lonW,lonE+.5,.5) # Need to match longitude values to those of the coordinate variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specify the pressure level. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vertical level specificaton\n",
    "pLevel = 850\n",
    "levStr = f'{pLevel}'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will display temperature and wind (and ultimately, temperature advection), so pick the relevant variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now create objects for our desired DataArrays based on the coordinates we have subsetted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data variable selection\n",
    "T = dsT['t'].sel(time=dateList,lev=pLevel,lat=latRange,lon=lonRange)\n",
    "U = dsU['u'].sel(time=dateList,lev=pLevel,lat=latRange,lon=lonRange)\n",
    "V = dsV['v'].sel(time=dateList,lev=pLevel,lat=latRange,lon=lonRange)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In order to calculate many meteorologically-relevant quantities that depend on distances between gridpoints, we need horizontal distance between gridpoints in meters. MetPy can infer this from datasets such as our CFSR that are lat-lon based, but we need to explicitly assign a coordinate reference system first."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <span style=\"color:purple\">3) Calculate and visualize diagnostic quantities.</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crsCFSR = {'grid_mapping_name': 'latitude_longitude', 'earth_radius': 6371229.0}\n",
    "T = T.metpy.assign_crs(crsCFSR)\n",
    "U = U.metpy.assign_crs(crsCFSR)\n",
    "V = V.metpy.assign_crs(crsCFSR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T.metpy.assign_crs(crsCFSR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define our subsetted coordinate arrays of lat and lon. Pull them from any of the DataArrays. We'll need to pass these into the contouring functions later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lats = T.lat\n",
    "lons = T.lon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First, let's just plot contour lines of temperature and wind barbs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Unit conversions\n",
    "While we will simply redefine our temperature object so it is in Celsius, let's create two new objects for U and V in knots, since we will want to preserve the original units (m/s) when it comes time to calculate temperature advection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = T.metpy.convert_units('degC')\n",
    "UKts = U.metpy.convert_units('kts')\n",
    "VKts = V.metpy.convert_units('kts')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an array of values for the temperature contours: since we are using contour lines, simply define a range large enough to encompass the expected values, and set a contour interval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T.min().values, T.max().values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cint = 2\n",
    "minT, maxT = (-30,22)\n",
    "TContours = np.arange(minT, maxT, cint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Make the map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constrainLat, constrainLon = (0.5, 4.0)\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": [
    "Although there is just a single time in our time range, we'll still employ a loop here in case we wanted to include multiple times later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for time in dateList:\n",
    "    print(\"Processing\", time)\n",
    "    \n",
    "    # Create the time strings, for the figure title as well as for the file name.\n",
    "    timeStr = dt.strftime(time,format=\"%Y-%m-%d %H%M UTC\")\n",
    "    timeStrFile = dt.strftime(time,format=\"%Y%m%d%H\")\n",
    "    \n",
    "    tl1 = str('CFSR, Valid at: '+ timeStr)\n",
    "    tl2 = levStr + \" hPa temperature (°C) and winds (kts)\"\n",
    "    \n",
    "    title_line = (tl1 + '\\n' + tl2 + '\\n')\n",
    "    \n",
    "    fig = plt.figure(figsize=(21,15)) # Increase size to adjust for the constrained lats/lons\n",
    "    ax = plt.subplot(1,1,1,projection=proj_map)\n",
    "    ax.set_extent ([lonW+constrainLon,lonE-constrainLon,latS+constrainLat,latN-constrainLat])\n",
    "    ax.add_feature(cfeature.COASTLINE.with_scale(res))\n",
    "    ax.add_feature(cfeature.STATES.with_scale(res),edgecolor='brown')\n",
    "    \n",
    "    # Need to use Xarray's sel method here to specify the current time for any DataArray you are plotting.\n",
    "    \n",
    "    # 1. Contour lines of temperature\n",
    "    cT = ax.contour(lons, lats, T.sel(time=time), levels=TContours, colors='orange', linewidths=3, transform=proj_data)\n",
    "    ax.clabel(cT, inline=1, fontsize=12, fmt='%.0f')\n",
    "    \n",
    "    # 4. wind barbs\n",
    "    # Plotting wind barbs uses the ax.barbs method. Here, you can't pass in the DataArray directly; you can only pass in the array's values.\n",
    "    # Also need to sample (skip) a selected # of points to keep the plot readable.\n",
    "    # Remember to use Xarray's sel method here as well to specify the current time.\n",
    "    skip = 2\n",
    "    ax.barbs(lons[::skip],lats[::skip],UKts.sel(time=time)[::skip,::skip].values, VKts.sel(time=time)[::skip,::skip].values, color='purple',zorder=2,transform=proj_data)\n",
    "\n",
    "    title = plt.title(title_line,fontsize=16)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Look at the map and you should easily be able to identify the areas of prominent warm and cold advection. Instead of qualitatively assessing a diagnostic quantity such as temperature advection, how about we quantify it? For that, we will use MetPy's <i>diagnostic library</i>, which we have imported as `mpcalc`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [MetPy Diagnostics](https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Let's explore the `advection` diagnostic: https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.advection.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### In order to calculate horizontal tempareture advection, we need a <i>scalar</i> quantity (temperature, the \"thing\" being advected), and a <i>vector</i> field (wind, the \"thing\" doing the advecting). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For MetPy's advection function, we pass in the scalar, followed by the vector ... in this case, T followed by U and V."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate temperature advection by the horizontal wind:\n",
    "tAdv = mpcalc.advection(T, U, V)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the output `DataArray`, containing the calculated values (and units) of temperature advection. Note how these values scale ... on the order of 10**-5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tAdv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get a quick visualization, using Xarray's built-in interface to Matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tAdv.plot(figsize=(15,10),cmap='coolwarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We have three strategies to make the resulting map look nicer:\n",
    "1. Scale up the values by an appropriate power of ten\n",
    "2. Focus on the more extreme values of temperature advection: thus, do not contour values that are close to zero.\n",
    "3. Smooth the values of temperature advection ... especially important in datasets with a high degree of horizontal resolution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take care of the first two strategies, and then assess the need for smoothing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scale these values up by 1e5 (or 1 * 10**5, or 100,000) and find the min/max values. Use these to inform the setting of the contour fill intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = 1e5\n",
    "minTAdv = (tAdv*scale).min().values\n",
    "maxTAdv = (tAdv*scale).max().values\n",
    "print (minTAdv, maxTAdv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Usually, we wish to avoid plotting the \"zero\" contour line for diagnostic quantities such as divergence, advection, and frontogenesis. Thus, create two lists of values ... one for negative and one for positive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "advInc = 50\n",
    "negAdvContours = np.arange (-300, 0, advInc)\n",
    "posAdvContours = np.arange (50, 350, advInc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "negAdvContours"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now, let's plot temperature advection on the map. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for time in dateList:\n",
    "    print(f'Processing {time}')\n",
    "    \n",
    "    # Create the time strings, for the figure title as well as for the file name.\n",
    "    timeStr = dt.strftime(time,format=\"%Y-%m-%d %H%M UTC\")\n",
    "    timeStrFile = dt.strftime(time,format=\"%Y%m%d%H\")\n",
    "    \n",
    "    tl1 = f'CFSR, Valid at: {timeStr}'\n",
    "    tl2 = f'{levStr} hPa temperature advection (°C x 10**5 / s)'\n",
    "    \n",
    "    title_line = f'{tl1}\\n{tl2}\\n'\n",
    "    \n",
    "    fig = plt.figure(figsize=(21,15)) # Increase size to adjust for the constrained lats/lons\n",
    "    ax = fig.add_subplot(1,1,1,projection=proj_map)\n",
    "    ax.set_extent ([lonW+constrainLon,lonE-constrainLon,latS+constrainLat,latN-constrainLat])\n",
    "    ax.add_feature(cfeature.COASTLINE.with_scale(res))\n",
    "    ax.add_feature(cfeature.STATES.with_scale(res),edgecolor='brown')\n",
    "    \n",
    "    # Need to use Xarray's sel method here to specify the current time for any DataArray you are plotting.\n",
    "    \n",
    "    # 1a. Contour lines of warm (positive temperature) advection.\n",
    "    # Don't forget to multiply by the scaling factor!\n",
    "    cPosTAdv = ax.contour(lons, lats, tAdv.sel(time=time)*scale, levels=posAdvContours, colors='red', linewidths=3, transform=proj_data)\n",
    "    ax.clabel(cPosTAdv, inline=1, fontsize=12, fmt='%.0f')\n",
    "    \n",
    "    # 1b. Contour lines of cold (negative temperature) advection\n",
    "    cNegTAdv = ax.contour(lons, lats, tAdv.sel(time=time)*scale, levels=negAdvContours, colors='blue', linewidths=3, transform=proj_data)\n",
    "    ax.clabel(cNegTAdv, inline=1, fontsize=12, fmt='%.0f')\n",
    "    \n",
    "    title = plt.title(title_line,fontsize=16)\n",
    "    #Generate a string for the file name and save the graphic to your current directory.\n",
    "    fileName = f'{timeStrFile}_CFSR_{levStr}_TAdv.png'\n",
    "    fig.savefig(fileName)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <span style=\"color:purple\">4. Smooth the diagnostic field.</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is not a terribly \"noisy\" field. But let's demonstrate the technique of employing a smoothing function. We will use a [Gaussian filter](https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.gaussian_filter.html),  from the [SciPy library](https://scipy.org/scipylib/), which, along with NumPy and Pandas, is one of the core packages in Python's scientific software ecosystem.\n",
    "\n",
    "MetPy now includes the Gaussian smoother as one of its diagnostics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = 8.0 # this depends on how noisy your data is, adjust as necessary\n",
    "\n",
    "tAdvSmooth = mpcalc.smooth_gaussian(tAdv, sigma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we did before with the unsmoothed values, let's look at the resulting `DataArray` and its extrema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tAdvSmooth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = 1e5\n",
    "minTAdv = (tAdvSmooth * scale).min().values\n",
    "maxTAdv = (tAdvSmooth * scale).max().values\n",
    "print (minTAdv, maxTAdv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not surprisingly, the magnitude of the extrema decreased."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the smoothed advection field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for time in dateList:\n",
    "    print(f'Processing {time}')\n",
    "    \n",
    "    # Create the time strings, for the figure title as well as for the file name.\n",
    "    timeStr = dt.strftime(time,format=\"%Y-%m-%d %H%M UTC\")\n",
    "    timeStrFile = dt.strftime(time,format=\"%Y%m%d%H\")\n",
    "    \n",
    "    tl1 = f'CFSR, Valid at: {timeStr}'\n",
    "    tl2 = f'{levStr} hPa temperature advection ($^\\circ$C x 10**5 / s), smoothed'\n",
    "\n",
    "    title_line = f'{tl1}\\n{tl2}\\n'\n",
    "    \n",
    "    fig = plt.figure(figsize=(21,15)) # Increase size to adjust for the constrained lats/lons\n",
    "    ax = fig.add_subplot(1,1,1,projection=proj_map)\n",
    "    ax.set_extent ([lonW+constrainLon,lonE-constrainLon,latS+constrainLat,latN-constrainLat])\n",
    "    ax.add_feature(cfeature.COASTLINE.with_scale(res))\n",
    "    ax.add_feature(cfeature.STATES.with_scale(res),edgecolor='brown')\n",
    "    \n",
    "    # Need to use Xarray's sel method here to specify the current time for any DataArray you are plotting.\n",
    "    \n",
    "    # 1a. Contour lines of warm (positive temperature) advection.\n",
    "    # Don't forget to multiply by the scaling factor!\n",
    "    cPosTAdv = ax.contour(lons, lats, tAdvSmooth.sel(time=time)*scale, levels=posAdvContours, colors='red', linewidths=3, transform=proj_data)\n",
    "    ax.clabel(cPosTAdv, inline=1, fontsize=12, fmt='%.0f')\n",
    "    \n",
    "    # 1b. Contour lines of cold (negative temperature) advection\n",
    "    cNegTAdv = ax.contour(lons, lats, tAdvSmooth.sel(time=time)*scale, levels=negAdvContours, colors='blue', linewidths=3, transform=proj_data)\n",
    "    ax.clabel(cNegTAdv, inline=1, fontsize=12, fmt='%.0f')\n",
    "    \n",
    "    title = plt.title(title_line,fontsize=16)\n",
    "    #Generate a string for the file name and save the graphic to your current directory.\n",
    "    fileName = f'{timeStrFile}_CFSR_{levStr}_TAdvSmth.png'\n",
    "    fig.savefig(fileName)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 Jan. 2024 Environment",
   "language": "python",
   "name": "jan24"
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}