03_GriddedDiagnostics_Divergence_Isotachs_ERA5

import xarray as xr
import pandas as pd
import numpy as np
from datetime import datetime as dt
from metpy.units import units
import metpy.calc as mpcalc
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt

# 1. Set the bounds of your map.
# Set the four values in the code cell below. Test the values.
latN = 50
latS = 20
lonW = -100
lonE = -60

if ( latN > latS ) and ( lonE > lonW):
    cLat = (latN + latS)/2
    cLon = (lonW + lonE)/2
else:
    e = ("Northern (eastern) latitude (longitude) must be greater than southern (western). Go back and re-adjust!")
    raise ValueError(f"Error in lat/lon specification: {e}")

# Recall that in ERA5, longitudes run between 0 and 360, not -180 and 180
if (lonW < 0 ):
    lonW = lonW + 360
if (lonE < 0 ):
    lonE = lonE + 360
    
# 2. Select the map projection for your figure.
proj_map = ccrs.LambertConformal(central_latitude=cLat, central_longitude=cLon)

# 3. Select the projection on which the dataset is based
proj_data = ccrs.PlateCarree()

# 4. Select the resolution of the Cartopy cartographic features, and, if desired, MetPy's county shapefiles.

# Uncomment / comment as desired
#res = '10m' # Most detailed, best for small regions (e.g. NYS)
res = '50m' # Medium detail, best for medium-sized regions (e.g. CONUS)
#res = '110m' # Least detailed, best for large/global maps

#res_county = '5m' # Higher-res MetPy US County shapefiles
res_county = '20m' # Lower-res MetPy US County shapefiles

# Apply the latitude/longitude ranges and extend the data region if desired

expand_lon = 1
expand_lat = 1
latRange = np.arange(latS - expand_lat,latN + expand_lat,.25) # expand the data range a bit beyond the plot range
lonRange = np.arange((lonW - expand_lon),(lonE + expand_lon),.25) # Need to match longitude values to those of the coordinate variable

# Specify desired pressure levels: in this case, a list of one or more
plevel_list = [300]

startYear = 1979
startMonth = 2
startDay = 19
startHour = 12
startMinute = 0
startDateTime = dt(startYear,startMonth,startDay, startHour, startMinute)

endYear = 1979
endMonth = 2
endDay = 19
endHour = 18
endMinute = 0
endDateTime = dt(endYear,endMonth,endDay, endHour, endMinute)

delta_time = endDateTime - startDateTime
time_range_max = 7*86400

if (delta_time.total_seconds() > time_range_max):
    raise RuntimeError("Your time range must not exceed 7 days. Go back and try again.")

# Create a list of date and times based on what we specified for the initial and final times, 
#  using Pandas’ date_range function
dateRange = pd.date_range(startDateTime, endDateTime,freq="6h")

print ("Summary of subsetted dimensions: \n\n")
print (f'Longitude Range: {lonRange} \n')
print (f'Latitude Range: {latRange} \n')
print (f'Pressure levels: {plevel_list} \n')
print(f'Date Range: {dateRange}')

endDate = dt(2023,1,10)
if (endDateTime <= endDate): # Use WeatherBench archive

    cloud_source = True
    ds = xr.open_dataset(
     'gs://weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721.zarr', 
     chunks={'time': 48},
     consolidated=True,
     engine='zarr'
    )
    
    # Attach units to the pressure coordinate
    ds.coords['level'].attrs['units'] = 'hPa'
    
    # Rename the variable names in this dataset so they use their corresponding short_name attributes
    # Construct a dictionary whose keys are the original variable names, and whose values are their 
    # short names
    
    rename_data_vars = {}
    seen_short_names = set()
    
    for var_name in ds.data_vars:
        short = ds[var_name].attrs.get("short_name")
        # Only rename if:
        # 1. short_name exists
        # 2. We haven't already used it
        if short and short not in seen_short_names:
            rename_data_vars[var_name] = short
            seen_short_names.add(short)
    # Apply renaming
    ds = ds.rename(rename_data_vars)

else: # Use local archive
    import glob, os
    cloud_source = False
    input_directory = '/free/ktyle/era5'
    files = glob.glob(os.path.join(input_directory,'*_era5.nc'))
    ds = xr.open_mfdataset(files,coords='minimal',compat='override')
    # Rename two of the coordinate variables so they match what is in the WeatherBench archive
    ds = ds.rename({'valid_time': 'time', 'pressure_level': 'level'})

# Perform the initial subset
ds = ds.sel(time = dateRange, longitude = lonRange, latitude = latRange, level = plevel_list)

ds

Z = ds['z']
U = ds['u']
V = ds['v']

U

p_level = 300
Zp = Z.sel(level=p_level).compute()
Up = U.sel(level=p_level).compute()
Vp = V.sel(level=p_level).compute()

# Calculate windspeed and convert to knots.
WSPDp = mpcalc.wind_speed(Up, Vp)
WSPDkts = WSPDp.metpy.convert_units('kts')

wspdInc = 25
wspdContours = np.arange (100, 200, wspdInc)

div = mpcalc.divergence(Up, Vp)

div.sel(latitude=40, longitude=285) # remember, ERA5 represents W Hemisphere longitudes from 180 to 360

sigma = 7.0 # this depends on how noisy your data is, adjust as necessary

divSmooth = mpcalc.smooth_gaussian(div, sigma)

scale = 1e6
divSmooth = divSmooth * scale
minDiv = divSmooth.min().values
maxDiv = divSmooth.max().values
print (minDiv, maxDiv)

divInc = 20
negDivContours = np.arange (-160, 0, divInc)
posDivContours = np.arange (10, 160, divInc)

Zdam = mpcalc.geopotential_to_height(Zp).metpy.convert_units('dam')

if (p_level >= 700):
    incZ = 3
elif (p_level >= 400):
    incZ = 6
elif (p_level >= 150):
    incZ = 9
else: 
    incZ = 12

zContours = np.arange(0,3000,incZ)

Ukts = Up.metpy.convert_units('kts')
Vkts = Vp.metpy.convert_units('kts')

lats = Up.latitude
lons = Up.longitude

constrainLat, constrainLon = (0.7, 6.5)

for time in dateRange:
    print(f'Processing {time}')
    
    # Create the time strings, for the figure title as well as for the file name.
    timeStr = dt.strftime(time,format="%Y-%m-%d %H%M UTC")
    timeStrFile = dt.strftime(time,format="%Y%m%d%H")
    
    tl1 = f'ERA5, Valid at: {timeStr}'
    tl2 = f'{p_level} hPa heights, isotachs and wind barbs and divergence x 10**6'

    title_line = f'{tl1}\n{tl2}\n'
    
    fig = plt.figure(figsize=(24,18)) # Increase size to adjust for the constrained lats/lons
    ax = fig.add_subplot(1,1,1,projection=proj_map)
    ax.set_extent ([lonW+constrainLon,lonE-constrainLon,latS+constrainLat,latN-constrainLat])
    ax.add_feature(cfeature.COASTLINE.with_scale(res))
    ax.add_feature(cfeature.STATES.with_scale(res),edgecolor='brown')
    
    # Need to use Xarray's sel method here to specify the current time for any DataArray you are plotting.
    # 1. Contour fills of isotachs
    CF = ax.contourf(lons,lats,WSPDkts.sel(time=time),levels=wspdContours,transform=proj_data,cmap='gist_gray_r',extend='max')
    cbar = plt.colorbar(CF,shrink=0.5)
    cbar.ax.tick_params(labelsize=16)
    cbar.ax.set_ylabel("Windspeed (kts)",fontsize=16)

    # 2. Contour lines of divergence (first negative values, then positive)
    cDivNeg = ax.contour(lons, lats, divSmooth.sel(time=time), levels=negDivContours, colors='blue',linestyles='dotted',linewidths=3, transform=proj_data)
    ax.clabel(cDivNeg, inline=1, fontsize=12,fmt='%.0f')
    cDivPos = ax.contour(lons, lats, divSmooth.sel(time=time), levels=posDivContours, colors='red',linewidths=2, transform=proj_data)
    ax.clabel(cDivPos, inline=1, fontsize=12,fmt='%.0f')
    
    # 3. Contour lines of geopotential heights
    cZ = ax.contour(lons, lats, Zdam.sel(time=time), levels=zContours, colors='black', transform=proj_data)
    ax.clabel(cZ, inline=1, fontsize=12, fmt='%.0f')
    
    # 4. wind barbs
    # 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.
    # Also need to sample (skip) a selected # of points to keep the plot readable.
    # Remember to use Xarray's sel method here as well to specify the current time.
    skip = 6
    ax.barbs(lons[::skip],lats[::skip],Ukts.sel(time=time)[::skip,::skip].values, Vkts.sel(time=time)[::skip,::skip].values, color='purple',transform=proj_data)

    title = ax.set_title(title_line,fontsize=16)
    
    # Generate a string for the file name and save the graphic to your current directory.
    fileName = f'{timeStrFile}_ERA5_{p_level}_WndDivHght'
    fig.savefig(fileName)