Geostrophic and Ageostrophic Wind

from datetime import datetime, timedelta

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import metpy.calc as mpcalc
import numpy as np
from scipy import ndimage
from siphon.catalog import TDSCatalog
from xarray.backends import NetCDF4DataStore
import xarray as xr

dt = datetime(2016, 8, 22, 18)
forecast_hour = 3
h = timedelta(hours=forecast_hour)

# Assemble our URL to the THREDDS Data Server catalog,
# and access our desired dataset within via NCSS
base_url = 'https://www.ncei.noaa.gov/thredds/catalog/model-gfs-g4-anl-files-old/'
cat = TDSCatalog(f'{base_url}{dt:%Y%m}/{dt:%Y%m%d}/catalog.xml')
ncss = cat.datasets[f'gfsanl_4_{dt:%Y%m%d}_{dt:%H}'
                    f'00_00{forecast_hour}.grb2'].subset()

# Create NCSS query for our desired time, region, and data variables
query = ncss.query()

query.lonlat_box(north=50, south=30, east=-80, west=-115)
query.time(dt + h)
query.variables('Geopotential_height_isobaric',
                'u-component_of_wind_isobaric',
                'v-component_of_wind_isobaric')
query.vertical_level(100000)

data = ncss.get_data(query)

ds = xr.open_dataset(NetCDF4DataStore(data)).metpy.parse_cf()

# Pull out variables you want to use
height = ds.Geopotential_height_isobaric.squeeze()
u_wind = ds['u-component_of_wind_isobaric'].squeeze().metpy.quantify()
v_wind = ds['v-component_of_wind_isobaric'].squeeze().metpy.quantify()

vtime = height.time.values.squeeze().astype('datetime64[ms]').astype('O')

lat = ds.latitude
lon = ds.longitude

# Combine 1D latitude and longitudes into a 2D grid of locations
lon_2d, lat_2d = np.meshgrid(lon, lat)

# Smooth height data
height = mpcalc.smooth_n_point(height, 9, 3)

# Compute the geostrophic wind
geo_wind_u, geo_wind_v = mpcalc.geostrophic_wind(height)

# Calculate ageostrophic wind components
ageo_wind_u = u_wind - geo_wind_u
ageo_wind_v = v_wind - geo_wind_v

# Create new figure
fig = plt.figure(figsize=(15, 10), facecolor='black')

# Add the map and set the extent
ax = plt.axes(projection=ccrs.PlateCarree())
ax.set_extent([-105., -93., 35., 43.])
ax.patch.set_fill(False)

# Add state boundaries to plot
ax.add_feature(cfeature.STATES, edgecolor='white', linewidth=2)

# Contour the heights every 10 m
contours = np.arange(10, 200, 10)
# Because we have a very local graphics area, the contours have joints
# to smooth those out we can use `ndimage.zoom`
zoom_500 = mpcalc.zoom_xarray(height, 5)
c = ax.contour(zoom_500.longitude, zoom_500.latitude, zoom_500, levels=contours,
               colors='red', linewidths=4)
ax.clabel(c, fontsize=12, inline=1, inline_spacing=3, fmt='%i')

# Set up parameters for quiver plot. The slices below are used to
# subset the data (here taking every 4th point in x and y). The
# quiver_kwargs are parameters to control the appearance of the
# quiver so that they stay consistent between the calls.
quiver_slices = (slice(None, None, 2), slice(None, None, 2))
quiver_kwargs = {'headlength': 4, 'headwidth': 3, 'angles': 'uv',
                 'scale_units': 'xy', 'scale': 20}

# Plot the wind vectors
wind = ax.quiver(lon_2d[quiver_slices], lat_2d[quiver_slices],
                 u_wind[quiver_slices], v_wind[quiver_slices],
                 color='blue', **quiver_kwargs)
geo = ax.quiver(lon_2d[quiver_slices], lat_2d[quiver_slices],
                geo_wind_u[quiver_slices], geo_wind_v[quiver_slices],
                color='darkorchid', **quiver_kwargs)
ageo = ax.quiver(lon_2d[quiver_slices], lat_2d[quiver_slices],
                 ageo_wind_u[quiver_slices], ageo_wind_v[quiver_slices],
                 color='lime', **quiver_kwargs)

# Add a title to the plot
plt.title('1000mb Geopotential Heights(m), Wind(blue), '
          'Geostrophic Wind(purple), and \n Ageostrophic Wind(green) '
          f'for {vtime:%d %B %Y %H:%MZ}', color='white', size=14)