{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 04_GriddedDiagnostics_Frontogenesis_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 frontogenesis.\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 ('/cfsr/data/%s/g.%s.0p5.anl.nc' % (startYear, startYear))\n",
    "dsT = xr.open_dataset ('/cfsr/data/%s/t.%s.0p5.anl.nc' % (startYear, startYear))\n",
    "dsU = xr.open_dataset ('/cfsr/data/%s/u.%s.0p5.anl.nc' % (startYear, startYear))\n",
    "dsV = xr.open_dataset ('/cfsr/data/%s/v.%s.0p5.anl.nc' % (startYear, startYear))\n",
    "dsW = xr.open_dataset ('/cfsr/data/%s/w.%s.0p5.anl.nc' % (startYear, startYear))\n",
    "dsQ = xr.open_dataset ('/cfsr/data/%s/q.%s.0p5.anl.nc' % (startYear, startYear))\n",
    "dsSLP = xr.open_dataset ('/cfsr/data/%s/pmsl.%s.0p5.anl.nc' % (startYear, startYear))"
   ]
  },
  {
   "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": "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": [
    "# <span style=\"color:purple\">3) Calculate and visualize 2-dimensional frontogenesis.</span>"
   ]
  },
  {
   "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": "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": "markdown",
   "metadata": {},
   "source": [
    "#### Unit conversions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "UKts = U.metpy.convert_units('kts')\n",
    "VKts = V.metpy.convert_units('kts')"
   ]
  },
  {
   "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": [
    "#### Let's explore the `frontogenesis` diagnostic: https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.frontogenesis.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### We first need to compute [potential temperature](https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.potential_temperature.html). First, let's attach units to our chosen pressure level, so that MetPy can calculate potential temperature accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Attach units to pLevel for use in MetPy with new variable, P:\n",
    "P = pLevel*units['hPa']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Now calculate potential temperature, and verify that it looks reasonable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Thta = mpcalc.potential_temperature(P, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Thta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Now calculate frontogenesis, and examine the output `DataArray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frnt = mpcalc.frontogenesis(Thta, U, V)\n",
    "frnt"
   ]
  },
  {
   "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": [
    "frnt.plot(figsize=(15,10),cmap='coolwarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the white pixel near 44 N, 83 W. This likely is a `NaN` somehow resulting from the frontogenesis calculation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The units are in degrees K (or C) per meter per second. Traditionally, we plot frontogenesis on more of a meso- or synoptic scale ... degrees K/C per 100 km per 3 hours. Let's simply multiply by the conversion factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A conversion factor to get frontogensis units of K per 100 km per 3 h\n",
    "convert_to_per_100km_3h = 1000*100*3600*3\n",
    "frnt = frnt * convert_to_per_100km_3h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the range and scale of values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frnt.min().values, frnt.max().values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A scale of 1 (i.e., 1x10^0, AKA 1e0) looks appropriate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = 1e0"
   ]
  },
  {
   "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": [
    "frntInc = 3\n",
    "negFrntContours = np.arange (-18, 0, frntInc)\n",
    "posFrntContours = np.arange (3, 21, frntInc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now, let's plot frontogenesis on the map. "
   ]
  },
  {
   "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 frontogenesis ($^\\circ$C / 100 km / 3 hr)\"\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",
    "    # 1a. Contour lines of warm (positive temperature) advection.\n",
    "    # Don't forget to multiply by the scaling factor!\n",
    "    cPosTAdv = ax.contour(lons, lats, frnt.sel(time=time)*scale, levels=posFrntContours, 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, frnt.sel(time=time)*scale, levels=negFrntContours, 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 = timeStrFile + '_CFSR_' + levStr + '_Frnt.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": [
    "Use [MetPy's implementation of a Gaussian smoother](https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.smooth_gaussian.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = 10.0 # this depends on how noisy your data is, adjust as necessary\n",
    "frntSmth = mpcalc.smooth_gaussian(frnt, sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frntSmth.plot(figsize=(15,10),cmap='coolwarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the smoothed field."
   ]
  },
  {
   "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 frontogenesis ($^\\circ$C / 100 km / 3 hr)\"\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",
    "    # 1a. Contour lines of warm (positive temperature) advection.\n",
    "    # Don't forget to multiply by the scaling factor!\n",
    "    cPosTAdv = ax.contour(lons, lats, frntSmth.sel(time=time)*scale, levels=posFrntContours, 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, frntSmth.sel(time=time)*scale, levels=negFrntContours, 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 = timeStrFile + '_CFSR_' + levStr + '_Frnt.png'\n",
    "    fig.savefig(fileName)\n"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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
}