ERA5 Data and Graphics Preparation for the Science-on-a-Sphere

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

warnings.simplefilter("ignore")

# You define these
caseDir = 'Bliz2026'
prodDir = '850_T_Z_Wind'

# These remain as is
graphicsDir = caseDir + '/' + prodDir + '/2048/'
labelsDir = caseDir + '/' + prodDir + '/labels/'
mediaDir = caseDir + '/' + prodDir + '/media/'
playlistDir = caseDir + '/' + prodDir + '/playlist/'

! mkdir -p {graphicsDir} {labelsDir} {mediaDir} {playlistDir}

# Remove all PNGs in the graphics directory - COMMENT OUT if you don't want to do this!
! rm -f  {graphicsDir}*.png

# Set the bounds of the map ... for the SoS, we will cover the entire globe.
## Do not change these values ##
latN = 90
latS = -90
lonW = 0
lonE = 360

latRange = np.arange(latS, latN + 0.1,.5) # ensure +90 is included
lonRange = np.arange(lonW, lonE, .5) 
## 

# Select the map projection for your figure.
proj_map = ccrs.PlateCarree() # For SoS maps

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

# Select the resolution of the Cartopy cartographic features: coarsest Natural Earth, and no counties

res = '110m' # Least detailed, best for large/global maps

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

startYear = 2026
startMonth = 2
startDay = 20
startHour = 0
startMinute = 0
startDateTime = dt(startYear,startMonth,startDay, startHour, startMinute)

endYear = 2026
endMonth = 2
endDay = 24
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")

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 temporal and vertical level subset
ds = ds.sel(time = dateRange, longitude = lonRange, latitude = latRange, level = plevel_list)

ds

def reorder_longitudes(da):
    """
    This function reorders the longitudes in a global grid from 0 --> 360 to -180 --> 180
    """
    da = da.assign_coords(longitude=(((da.longitude + 180 ) % 360 ) - 180))
    return da.sortby(da['longitude'])

# Select only those variables that you need for your product. Comment out or remove those you don't need.
## 4-D variables
T = reorder_longitudes(ds['t'])
U = reorder_longitudes(ds['u'])
V = reorder_longitudes(ds['v'])
Z = reorder_longitudes(ds['z'])

# 3-D variables
#SLP = reorder_longitudes(ds['msl'])

p_level = 850
# 4-d variables; as above, comment out or delete those you don't need
%time Tp = T.sel(level=p_level).compute()
%time Up = U.sel(level=p_level).compute()
%time Vp = V.sel(level=p_level).compute()
%time Zp = Z.sel(level=p_level).compute()

# 3-d variables; similarly, comment out or delete those you don't need
#%time SLPp = SLP.compute()

lons, lats, times= Tp.longitude, Tp.latitude, Tp.time

lons

Zdam = mpcalc.geopotential_to_height(Zp).metpy.convert_units('dam')
Ukts = Up.metpy.convert_units('kts')
Vkts = Vp.metpy.convert_units('kts')
Tc = Tp.metpy.convert_units('degC')

Tc

Tc.min().values,Tc.max().values

tContours = np.arange(-45,39,3)

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)

T0 = Tc.isel(time=0)
U0 = Ukts.isel(time=0)
V0 = Vkts.isel(time=0)
Z0 = Zdam.isel(time=0)

Z0

coast_border_col = 'dimgrey'
state_col = 'lightgrey'
wind_col = 'whitesmoke'

dpi = 100
fig = plt.figure(figsize=(2048/dpi, 1024/dpi))
ax = fig.add_subplot(1,1,1,projection=ccrs.PlateCarree(central_longitude=180), frameon=False)

ax.set_global()
ax.add_feature(cfeature.COASTLINE.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.BORDERS.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.STATES.with_scale(res), edgecolor=state_col)

# Temperature (T) contour fills

# Note we don't need the transform argument since the map/data projections are the same, but we'll leave it in
CF = ax.contourf(lons,lats,T0,levels=tContours,cmap=plt.get_cmap('coolwarm'), extend='both', transform=proj_data) 

# Height (Z) contour lines
CL = ax.contour(lons,lats,Z0,zContours,linewidths=1.25,colors='yellow', transform=proj_data)
ax.clabel(CL, inline_spacing=0.2, fontsize=8, fmt='%.0f')
fig.tight_layout(pad=.01)

# 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.
skip = 8
# Default length of wind barbs is 7. Shorten them a bit.
ax.barbs(lons[::skip],lats[::skip],U0[::skip,::skip].values, V0[::skip,::skip].values, color=wind_col, length = 5, transform=proj_data);

dpi = 100
fig = plt.figure(figsize=(2048/dpi, 1024/dpi))
ax = fig.add_subplot(1,1,1,projection=ccrs.Orthographic(central_longitude=180), frameon=False)

ax.set_global()
ax.add_feature(cfeature.COASTLINE.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.BORDERS.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.STATES.with_scale(res), edgecolor=state_col)

# Temperature (T) contour fills

# Note we don't need the transform argument since the map/data projections are the same, but we'll leave it in
CF = ax.contourf(lons,lats,T0,levels=tContours,cmap=plt.get_cmap('coolwarm'), extend='both', transform=proj_data) 

# Height (Z) contour lines
CL = ax.contour(lons,lats,Z0,zContours,linewidths=1.25,colors='yellow', transform=proj_data)
ax.clabel(CL, inline_spacing=0.2, fontsize=8, fmt='%.0f')
fig.tight_layout(pad=.01)

# 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.
skip = 8
# Default length of wind barbs is 7. Shorten them a bit.
ax.barbs(lons[::skip],lats[::skip],U0[::skip,::skip].values, V0[::skip,::skip].values, color=wind_col, length = 5, transform=proj_data);

# add cyclic points to data and longitudes
# latitudes are unchanged in 1-dimension
Z_cyc, lons_cyc = cutil.add_cyclic_point(Z0, lons)
T_cyc, lons_cyc = cutil.add_cyclic_point(T0, lons)
U_cyc, lons_cyc = cutil.add_cyclic_point(U0, lons)
V_cyc, lons_cyc = cutil.add_cyclic_point(V0, lons)

Z_cyc

lons_cyc

dpi = 100
fig = plt.figure(figsize=(2048/dpi, 1024/dpi))
ax = fig.add_subplot(1,1,1, projection=ccrs.Orthographic(central_longitude=180), frameon=False)

ax.set_global()
ax.add_feature(cfeature.COASTLINE.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.BORDERS.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.STATES.with_scale(res), edgecolor=state_col)

# Temperature (T) contour fills

# Note we don't need the transform argument since the map/data projections are the same, but we'll leave it in
CF = ax.contourf(lons_cyc,lats,T_cyc,levels=tContours,cmap=plt.get_cmap('coolwarm'), extend='both', transform=proj_data) 

# Height (Z) contour lines
CL = ax.contour(lons_cyc,lats,Z_cyc,zContours,linewidths=1.25,colors='yellow', transform=proj_data)
ax.clabel(CL, inline_spacing=0.2, fontsize=8, fmt='%.0f')
fig.tight_layout(pad=.01)

# 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.
skip = 8
# Default length of wind barbs is 7. Shorten them a bit.
ax.barbs(lons_cyc[::skip],lats[::skip],U_cyc[::skip,::skip], V_cyc[::skip,::skip], color=wind_col, length = 5, transform=proj_data);

dpi = 100
fig = plt.figure(figsize=(2048/dpi, 1024/dpi))
ax = fig.add_subplot(1,1,1, projection=ccrs.PlateCarree(), frameon=False)

ax.set_global()
ax.add_feature(cfeature.COASTLINE.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.BORDERS.with_scale(res), edgecolor=coast_border_col)
ax.add_feature(cfeature.STATES.with_scale(res), edgecolor=state_col, linewidth=0.75)

# Temperature (T) contour fills

# Note we don't need the transform argument since the map/data projections are the same, but we'll leave it in
CF = ax.contourf(lons_cyc,lats,T_cyc,levels=tContours,cmap=plt.get_cmap('coolwarm'), extend='both', transform=proj_data) 

# Height (Z) contour lines
CL = ax.contour(lons_cyc,lats,Z_cyc,zContours,linewidths=1.25,colors='yellow', transform=proj_data)
ax.clabel(CL, inline_spacing=0.2, fontsize=8, fmt='%.0f')
fig.tight_layout(pad=.01)

# 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.
skip = 8
# Default length of wind barbs is 7. Shorten them a bit.
ax.barbs(lons_cyc[::skip],lats[::skip],U_cyc[::skip,::skip], V_cyc[::skip,::skip], color=wind_col, length = 5, transform=proj_data)
# Save this figure to your current directory
fig.savefig('test_ERA5_SOS.png')

res = '110m'
dpi = 100
for size in (128, 800):
    if (size == 128):
        sizeStr = 'small'
    else:
        sizeStr ='big'
        
    fig = plt.figure(figsize=(size/dpi, size/dpi))
    ax = plt.subplot(1,1,1,projection=ccrs.Orthographic(central_longitude=-90), frameon=False)
    ax.set_global()
    ax.add_feature(cfeature.COASTLINE.with_scale(res))
    tl1 = caseDir
    tl2 = prodDir
    ax.set_title(f"{tl1}\n{tl2}", color='purple', fontsize=8)

    # Contour fills

    CF = ax.contourf(lons_cyc,lats,T_cyc, levels=tContours,cmap=plt.get_cmap('coolwarm'), extend='both', transform=proj_data) 

    fig.tight_layout(pad=.01)
    fig.savefig(f'{mediaDir}thumbnail_{sizeStr}.jpg')

# draw a new figure and replot the colorbar there
fig,ax = plt.subplots(figsize=(14,2), dpi=125)
# set tick and ticklabel color
tick_color='black'
label_color='orange'
cbar = fig.colorbar(CF, ax=ax, orientation='horizontal')
cbar.ax.xaxis.set_tick_params(color=tick_color, labelcolor=label_color, labelsize=8)
# Remove the Axes object ... essentially the contours and cartographic features ... from the figure.
ax.remove()
# All that remains is the colorbar ... save it to disk. Make the background transparent.
fig.savefig(f'{labelsDir}colorbar.png',transparent=True)

labelFname = f'{labelsDir}labels.txt'

def make_sos_map (timeIdx):
    
    res = '110m'
    dpi = 100
    fig = plt.figure(figsize=(2048/dpi, 1024/dpi))
    ax = plt.subplot(1,1,1,projection=ccrs.PlateCarree(), frameon=False)
    ax.set_global()
    ax.add_feature(cfeature.COASTLINE.with_scale(res), edgecolor=coast_border_col)
    ax.add_feature(cfeature.BORDERS.with_scale(res), edgecolor=coast_border_col)
    ax.add_feature(cfeature.STATES.with_scale(res), edgecolor=state_col, linewidth=0.75 )

#   add cyclic points to data and longitudes
#   latitudes are unchanged in 1-dimension
    Z_cyc, lons_cyc= cutil.add_cyclic_point(Zdam.isel(time=timeIdx), lons)
    T_cyc, lons_cyc = cutil.add_cyclic_point(Tc.isel(time=timeIdx), lons)
    U_cyc, lons_cyc = cutil.add_cyclic_point(Ukts.isel(time=timeIdx), lons)
    V_cyc, lons_cyc = cutil.add_cyclic_point(Vkts.isel(time=timeIdx), lons)

#   Temperature (T) contour fills

#   Note we don't need the transform argument since the map/data projections are the same, but we'll leave it in
    CF = ax.contourf(lons_cyc,lats,T_cyc,levels=tContours,cmap=plt.get_cmap('coolwarm'), extend='both', transform=proj_data) 

#   Height (Z) contour lines
    CL = ax.contour(lons_cyc,lats,Z_cyc,zContours,linewidths=1.25,colors='yellow', transform=proj_data)
    ax.clabel(CL, inline_spacing=0.2, fontsize=8, fmt='%.0f')
    fig.tight_layout(pad=.01)

#   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.
    skip = 8
#   Default length of wind barbs is 7. Shorten them a bit.
    ax.barbs(lons_cyc[::skip],lats[::skip],U_cyc[::skip,::skip], V_cyc[::skip,::skip], color=wind_col, length = 5, transform=proj_data)
    frameNum = f'{timeIdx}'.zfill(2)
    figName = f'{graphicsDir}ERA5_{caseDir}_{prodDir}_{frameNum}.png'
    fig.savefig(figName)
    
    # Reduce the size of the PNG image via the Linux pngquant utility. The -f option overwites the resulting file if it already exists.
    # The output file will end in "-fs8.png"
    ! pngquant -f {figName}
    
    # Remove the original PNG
    ! rm -f {figName}
    
    # Do not show the graphic in the notebook
    plt.close()

nTimes = len(dateRange)
nTimes

labelFileObject = open(labelFname, 'w')
nFrames = 2
nFrames = nTimes

for timeIdx in range(0, nFrames, 1):
    
    make_sos_map(timeIdx)
    
    # Construct the title string and write it to the file
    valTime = pd.to_datetime(times[timeIdx].values).strftime("%m/%d/%Y %H UTC")
    tl1 = f'ERA5 {p_level} hPa Z/T/Wind {valTime} \n' # \n is the newline character
    labelFileObject.write(tl1)
    print(tl1)
    
# Close the text file
labelFileObject.close()
print("Finished with SoS graphic loop")

playlistFname = f'{playlistDir}playlist.sos'

# Set the author (creator) name
creaName = "Rovin Lazyle" # Put your name here!

plFileObject = open(playlistFname, 'w')
subCat = "ERA5"
cRet = "\n" # New line character code
lBr = "{"
rBr = "}"


nameStr = f'ERA5 {caseDir}{prodDir}'
descriptionStr = f'{lBr}{caseDir}{prodDir}{rBr}'
creaName = "Rovin Lazyle" # Put your name here!

plFileObject.write(f"name = {nameStr}{cRet}")
plFileObject.write(f"description = {descriptionStr}{cRet}")

plFileObject.write(f"pip = ../labels/colorbar.png{cRet}")
plFileObject.write(f"pipheight = 10{cRet}")
plFileObject.write(f"pipvertical = -35{cRet}")

plFileObject.write(f"label = ../labels/labels.txt{cRet}")
plFileObject.write(f"layer = Grids{cRet}")
plFileObject.write(f"layerdata = ../2048{cRet}")

plFileObject.write(f"firstdwell = 2000{cRet}")
plFileObject.write(f"lastdwell = 3000{cRet}")
plFileObject.write(f"fps = 8{cRet}")

plFileObject.write(f"zrotationenable = 1{cRet}")
plFileObject.write(f"zfps = 30{cRet}")

plFileObject.write(f"subcategory = {subCat}{cRet}")
plFileObject.write(f"Creator = {creaName}{cRet}")
                   
plFileObject.close()

! cat {playlistFname}