{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1e1ed7-f656-41e4-bc55-759247b05ee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f26c1af",
   "metadata": {},
   "source": [
    "Download NCEP GFS 1 degree data file via OpenDAP\n",
    "\n",
    "extract a number of fields at 300, 500 hPa, and surface\n",
    "\n",
    "Calculate 500 hPa vorticity field\n",
    "\n",
    "Plot the fields in 4 different figures\n",
    "\n",
    "https://unidata.github.io/python-training/gallery/500hpa_absolute_vorticity_winds/ is a helfpul example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc654702-a22c-4914-9a5c-c26d5f0e4b17",
   "metadata": {},
   "outputs": [],
   "source": [
    "###########################################\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.ndimage as ndimage\n",
    "import xarray as xr\n",
    "\n",
    "import metpy.calc as mpcalc\n",
    "from metpy.units import units\n",
    "\n",
    "from metpy.cbook import get_test_data\n",
    "from metpy.plots import add_metpy_logo\n",
    "\n",
    "###########################################\n",
    "\n",
    "# Open the example netCDF data\n",
    "#ds = xr.open_dataset(get_test_data('gfs_output.nc', False))\n",
    "\n",
    "#ds = xr.open_dataset('/home/jovyan/GetGFS/gfs.20220223t00z.1p00.f012.nc')\n",
    "#ds = xr.open_dataset('/home/jovyan/GetGFS/gfs.20220223t00z.0p25.f012.grb2.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48a4870f-0f1e-4742-9208-14b3824c6465",
   "metadata": {},
   "source": [
    "Open GFS forecasts using OpenDAP server directly. See \n",
    "https://earthdata.nasa.gov/collaborate/open-data-services-and-software/api/opendap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9d6f2f8-a015-4847-9461-9f0ab46a1c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('http://nomads.ncep.noaa.gov:80/dods/gfs_1p00/gfs20220223/gfs_1p00_00z')\n",
    "\n",
    "#ds = xr.open_dataset('http://nomads.ncep.noaa.gov:80/dods/gfs_0p25/gfs20220223/gfs_0p25_00z')\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eba48ca-e393-4cd9-9b3e-7c54ee2e99e6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62406376",
   "metadata": {},
   "outputs": [],
   "source": [
    "TIMEstring = '2022-02-23T12:00:00.000000000'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ae028e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "latmin = 0.0\n",
    "latmax = 90.0\n",
    "lonmax = 360.0-50.0\n",
    "lonmin = 360.0-140.0\n",
    "\n",
    "precip_water = ds['pwatclm'].sel(time=TIMEstring,lat=slice(latmin,latmax),lon=slice(lonmin,lonmax))\n",
    "\n",
    "Pmsl = ndimage.gaussian_filter(ds['prmslmsl'].sel(time=TIMEstring,lat=slice(latmin,latmax),lon=slice(lonmin,lonmax)), sigma=1.5, order=0)\n",
    "\n",
    "surface_temp = ds['tmp2m'].sel(time=TIMEstring,lat=slice(latmin,latmax),lon=slice(lonmin,lonmax)).values\n",
    "\n",
    "u_300 = ds['ugrdprs'].sel(time=TIMEstring, lev=300,lat=slice(latmin,latmax),lon=slice(lonmin,lonmax))\n",
    "v_300 = \n",
    "\n",
    "winds_300 = np.sqrt( np.square(u_300.values) + np.square(v_300.values) ) # wind speed\n",
    "\n",
    "u_500 = \n",
    "v_500 = \n",
    "\n",
    "winds_500 =   # wind speed\n",
    "\n",
    "lats = ds['lat'].sel(lat=slice(latmin,latmax))\n",
    "lons = ds['lon'].sel(lon=slice(lonmin,lonmax))\n",
    "\n",
    "d2r = np.pi/180.0   # degree to radian conversion factor\n",
    "radius = 6378137.0 # earth radius in meters\n",
    "\n",
    "llats=len(lats)\n",
    "llons=len(lons)\n",
    "\n",
    "print( llats, llons) \n",
    "\n",
    "# Combine 1D latitude and longitudes into a 2D grid of locations\n",
    "\n",
    "lon_2d, lat_2d = np.meshgrid( lons, lats )\n",
    "\n",
    "u_500.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61af127f",
   "metadata": {},
   "source": [
    "Calculate vertical vorticity at 500 hPa from u and v wind components, according to formula in spherical coordinates:\n",
    "\n",
    "\\begin{equation}\n",
    "    \\frac{1}{r \\mathrm{sin}\\theta} \\left(\n",
    "    \\frac{\\partial v_\\phi \\mathrm{sin}\\theta}{\\partial \\theta} - \\frac{\\partial v_\\theta}{\\partial \\phi} ) \\right)\n",
    "\\end{equation}\n",
    "\n",
    "where $\\phi$ is the azimuthal angle in spherical coordinates and can be considered the longitute (in radian).  $\\theta$ is the polar angle and is equal to ($\\frac{\\pi}{2}$ - latitude). All angles are in radian.\n",
    "\n",
    "See https://en.wikipedia.org/wiki/Spherical_coordinate_system for info on spherical coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3de4849b",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'np' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/var/folders/ln/5x3f1b2j4yj8zzvrxsmg1bw40000gn/T/ipykernel_21321/2511698485.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mvort_500\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzeros\u001b[0m\u001b[0;34m(\u001b[0m \u001b[0mwinds_500\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mj\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mllats\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mllons\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'np' is not defined"
     ]
    }
   ],
   "source": [
    "vort_500 = np.zeros( winds_500.shape )\n",
    "\n",
    "for j in range(1, llats-1 ):\n",
    "    \n",
    "    for i in range(1, llons-1 ):\n",
    "        vort_500[j,i] = ( -(u_500[j+1,i]*np.sin((90.0-lats[j+1])*d2r)\n",
    "        -u_500[j-1,i]*np.sin((90.0-lats[j-1])*d2r))/(lats[j+1]-lats[j-1])\n",
    "        +(v_500[j,i+1]-v_500[j,i-1])/((lons[i+1]-lons[i-1])*d2r))/(radius*np.sin((90.0-lats[j])*d2r))\n",
    "    \n",
    "# Set each and west boundary values by exending from the interior points\n",
    "\n",
    "    vort_500[j,0]=vort_500[j,1]\n",
    "    vort_500[j,llons-1]=vort_500[j,llons-2]\n",
    "  \n",
    "# Set north and south boundary values by exending from the interior points\n",
    "\n",
    "for i in range(0, llons ):\n",
    "    vort_500[0,i]=vort_500[1,i]\n",
    "    vort_500[llats-1,i]=vort_500[llats-2,i]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce1a957e",
   "metadata": {},
   "source": [
    "Calculate horizontal flow divergence at 300 hPa from u and v wind components, according to formula in spherical coordinates:\n",
    "\n",
    "\\begin{equation}\n",
    "    \\frac{1}{r \\mathrm{sin}\\theta} \\left(\n",
    "    \\frac{\\partial v_\\theta \\mathrm{sin}\\theta}{\\partial \\theta} + \\frac{\\partial v_\\phi}{\\partial \\phi} ) \\right)\n",
    "\\end{equation}\n",
    "\n",
    "where $\\phi$ and $\\theta$ are defined the same way as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2968ef3c-c09e-4d30-bd62-3aed90fbec0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "div_300 = np.zeros( winds_500.shape )\n",
    "\n",
    "for j in range(1, llats-1 ):\n",
    "    for i in range(1, llons-1 ):\n",
    "        div_300[j,i] = 0.0   # fill in the formula"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15066035-f889-4df4-9a90-b8d87264be92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get height fields for 300 and 500 hPa the height data\n",
    "\n",
    "heights_300 = \n",
    "heights_500 = "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac703561-1c19-42bc-9019-8c30f6e12d4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface_temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5a12dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "###########################################\n",
    "\n",
    "# Do unit conversions to what we wish to plot\n",
    "\n",
    "vort_500 = vort_500 * 1e5 # 1e5 * Vorticity 1/s.\n",
    "\n",
    "Pmsl = Pmsl * 0.01 # hPa\n",
    "\n",
    "surface_temp = surface_temp - 273.15  # K to C "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5d43899",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup  projection for map plotting\n",
    "\n",
    "###########################################\n",
    "\n",
    "crs = ccrs.LambertConformal(central_longitude=-100.0, central_latitude=45.0)\n",
    "\n",
    "###########################################\n",
    "\n",
    "# Function used to create the map subplots\n",
    "def plot_background(ax):\n",
    "    ax.set_extent([-125.0, -70.0, 20., 55.])\n",
    "    ax.gridlines(draw_labels=True, dms=True, x_inline=False, y_inline=False)\n",
    "    ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.5)\n",
    "    ax.add_feature(cfeature.STATES, linewidth=0.5)\n",
    "    ax.add_feature(cfeature.BORDERS, linewidth=0.5)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e2fda40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the figure and plot background on different axes\n",
    "\n",
    "# https://scitools.org.uk/cartopy/docs/latest/tutorials/understanding_transform.html\n",
    "\n",
    "fig, axarr = plt.subplots(nrows=1, ncols=1, figsize=(15, 7.5), constrained_layout=True,\n",
    "                          subplot_kw={'projection': crs})\n",
    "plot_background(axarr)\n",
    "\n",
    "# Upper left plot - 300-hPa winds and geopotential heights\n",
    "cf1 = axarr.contourf(lon_2d, lat_2d, winds_300, cmap='cool', transform=ccrs.PlateCarree())\n",
    "c1 = axarr.contour(lon_2d, lat_2d, heights_300, colors='black', linewidths=2,\n",
    "                       transform=ccrs.PlateCarree())\n",
    "\n",
    "axarr.clabel(c1, fontsize=10, inline=1, inline_spacing=1, fmt='%i', rightside_up=True)\n",
    "axarr.set_title('300-hPa Wind Speeds and Heights', fontsize=16)\n",
    "cb1 = fig.colorbar(cf1, ax=axarr, location='bottom', shrink=1.0, fraction=0.05, pad=0.05)\n",
    "cb1.set_label('$m s^{-1}$', size='x-large')\n",
    "\n",
    "# Set up a 2D slice to reduce the number of wind barbs plotted \n",
    "wind_slice = (slice(None, None, 2), slice(None, None, 2))\n",
    "axarr.barbs(lon_2d[wind_slice], lat_2d[wind_slice],\n",
    "         u_300[wind_slice].values, v_300[wind_slice].values,\n",
    "         pivot='middle', color='black', transform=ccrs.PlateCarree())\n",
    "\n",
    "plt.savefig('Fig1.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6050d50",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axarr = plt.subplots(nrows=1, ncols=1, figsize=(15,7.5), constrained_layout=True,\n",
    "                          subplot_kw={'projection': crs})\n",
    "\n",
    "plot_background(axarr)\n",
    "\n",
    "\n",
    "# Upper right plot - 500mb relative vorticity and geopotential heights\n",
    "#cf2 = axarr[1].contourf(lon_2d, lat_2d, vort_500, cmap='BrBG', transform=ccrs.PlateCarree(),\n",
    "#                         zorder=0, norm=plt.Normalize(-32, 32))\n",
    "cf2 = axarr.contourf(lon_2d, lat_2d, vort_500, cmap='cool', transform=ccrs.PlateCarree(),\n",
    "                         zorder=0)\n",
    "c2 = axarr.contour(lon_2d, lat_2d, heights_500, colors='k', linewidths=2,\n",
    "                       transform=ccrs.PlateCarree())\n",
    "axarr.clabel(c2, fontsize=10, inline=1, inline_spacing=1, fmt='%i', rightside_up=True)\n",
    "axarr.set_title('500-hPa Relative Vorticity and Heights', fontsize=16)\n",
    "cb2 = fig.colorbar(cf2, ax=axarr, location='bottom',  shrink=1.0, fraction=0.05, pad=0.05)\n",
    "cb2.set_label(r'$10^{-5}$ s$^{-1}$', size='x-large')\n",
    "\n",
    "# Set up a 2D slice to reduce the number of wind barbs plotted \n",
    "wind_slice = (slice(None, None, 2), slice(None, None, 2))\n",
    "axarr.quiver(lon_2d[wind_slice], lat_2d[wind_slice],\n",
    "         u_500[wind_slice].values, v_500[wind_slice].values,\n",
    "         pivot='middle', color='black', transform=ccrs.PlateCarree())\n",
    "\n",
    "plt.savefig('Fig2.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e247c8e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lower left plot - surface temperatures\n",
    "\n",
    "fig, axarr = plt.subplots(nrows=1, ncols=1, figsize=(15,7.5), constrained_layout=True,\n",
    "                          subplot_kw={'projection': crs})\n",
    "plot_background(axarr)\n",
    "\n",
    "cf3 = axarr.contourf(lon_2d, lat_2d, surface_temp, cmap='YlOrRd',\n",
    "                         transform=ccrs.PlateCarree(), zorder=0)\n",
    "c3 = axarr.contour(lon_2d, lat_2d, Pmsl, colors='k', linewidths=2,\n",
    "                       transform=ccrs.PlateCarree())\n",
    "\n",
    "axarr.set_title('Surface Temperature (C) and Pmsl (hPa)', fontsize=16)\n",
    "cb3 = fig.colorbar(cf3, ax=axarr, location='bottom', shrink=1.0, fraction=0.05,  pad=0.05)\n",
    "#cb3.set_label('\\N{DEGREE FAHRENHEIT}', size='x-large')\n",
    "cb3.set_label('\\N{DEGREE CELSIUS}', size='x-large')\n",
    "\n",
    "# Set up a 2D slice to reduce the number of wind barbs plotted \n",
    "wind_slice = (slice(None, None, 2), slice(None, None, 2))\n",
    "'''\n",
    "axarr.barbs(lon_2d[wind_slice], lat_2d[wind_slice],\n",
    "         u_500[wind_slice].values, v_500[wind_slice].values,\n",
    "         pivot='middle', color='black', transform=ccrs.PlateCarree())\n",
    "'''\n",
    "\n",
    "plt.savefig('Fig3.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "755acb7d-4e21-419a-b0d3-a84b03735ecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axarr = plt.subplots(nrows=1, ncols=1, figsize=(15,7.5), constrained_layout=True,\n",
    "                          subplot_kw={'projection': crs})\n",
    "plot_background(axarr)\n",
    "\n",
    "# Lower right plot - precipitable water entire atmosphere\n",
    "cf4 = axarr.contourf(lon_2d, lat_2d, precip_water, cmap='Greens',\n",
    "                         transform=ccrs.PlateCarree(), zorder=0)\n",
    "axarr.set_title('Precipitable Water', fontsize=16)\n",
    "cb4 = fig.colorbar(cf4, ax=axarr, location='bottom', shrink=1.0, fraction=0.05, pad=0.05)\n",
    "cb4.set_label('kg $m^{-2}$ ', size='x-large')\n",
    "\n",
    "plt.savefig('Fig4.png')\n",
    "\n",
    "# Set height padding for plots\n",
    "# fig.set_constrained_layout_pads(w_pad=0., h_pad=0.1, hspace=0., wspace=0.)\n",
    "\n",
    "# Set figure title\n",
    "\n",
    "# write the correct time in the title according to TIMEString\n",
    "\n",
    "# Create a clean datetime object for plotting based on time of Geopotential heights\n",
    "#vtime = ds.time.data[0].astype('datetime64[ms]').astype('O')\n",
    "# Plot two titles, one on right and left side\n",
    "#plt.title('500-hPa NAM Geopotential Heights (m)'\n",
    "#          ' and Wind Barbs (kt)', loc='left')\n",
    "#plt.title('Valid Time: {}'.format(vtime), loc='right')\n",
    "\n",
    "# fig.suptitle(ds['time'][0].dt.strftime('%d %B %Y %H:%MZ').values, fontsize=24)\n",
    "\n",
    "# Display the plot\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946d6e25",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd8c271",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}