updating to accomodate plotting issues, in hopes that others dont have to bash...

updating to accomodate plotting issues, in hopes that others dont have to bash their heads against a wall

plot_hrrr_nc-gridded.py

-################################################################################
-# Written by Jared Rennie, NCICS
-# Derived this iPython Notebook, from Kevin Tyle
+########################################
+# Written by Jared Rennie
+# Derived from Kevin Tyle
 # https://github.com/ktyle/python_pangeo_ams2021/blob/main/notebooks/03_Pangeo_HRRR.ipynb
-################################################################################
+########################################
 
 # Import Packages
 import sys, time, datetime, os
@@ -26,17 +26,18 @@ import cartopy.io.shapereader as shpreader
 
 import rioxarray
 from shapely.geometry import mapping
+from metpy.plots import USCOUNTIES
+
+import warnings
+warnings.filterwarnings("ignore")
 
 # Declare directories
 main_directory='.'
+plot_directory=main_directory+'/results_plots'
 
 # Define Plotting Info/Bounds
 dpi=300
 plt.style.use('dark_background')
-minLat = 33.75    
-maxLat = 36.75   
-minLon = -85 
-maxLon = -75
 
 ####################
 # BEGIN PROGRAM
@@ -45,7 +46,7 @@ start=time.time()
 
 # Read in Arguments 
 if len(sys.argv) < 6:
-    sys.exit("USAGE: <YYYY> <MM> <DD> <HH> <ELEMENT>\nExample: ipython plot_hrrr_nc-gridded.py 2021 08 25 12 T/Td")
+    sys.exit("USAGE: <YYYY> <MM> <DD> <HH> <ELEMENT>\nExample: ipython plot_hrrr_nc-gridded.py 2021 08 11 12 T/Td")
 
 # Get Date Info
 year= "%04i" % int(sys.argv[1])
@@ -111,20 +112,24 @@ lat1 = 38.5
 slat = 38.5
 projData= ccrs.LambertConformal(central_longitude=lon1,central_latitude=lat1,standard_parallels=[slat])
 
-# Get Data
+# Get Temperature Data
 airTemp = ds[var].metpy.convert_units('degF')
 
 # Clip By Shapefile
 airTemp=airTemp.rio.set_spatial_dims(x_dim="projection_x_coordinate", y_dim="projection_y_coordinate", inplace=True)
 airTemp=airTemp.rio.write_crs(projData.proj4_params, inplace=True)
-
-# One thing I saw was the reproject option in rioxarray, but I get a "float16 key error when trying this"
-#airTemp=airTemp.rio.reproject(projection)
-
 airTempC = airTemp.rio.clip(geo_shapefile.geometry.apply(mapping), geo_shapefile.crs, drop=False)
 
 # PLOT
 print("\nPLOT ANALYSIS")
+dpi=300
+plt.style.use('dark_background')
+
+# Set Bounds
+minLat = 33.75    
+maxLat = 36.75   
+minLon = -85 
+maxLon = -75
 
 # Get Plotting Coordinates
 x = airTempC.projection_x_coordinate
@@ -140,24 +145,20 @@ cm.set_array(airTemp)
 cm.set_clim(vmin, vmax)
 
 # Set Up Figure
-# When trying to set projection, ccrs.Mercator() does not work
-outProjection=projData
-#outProjection=ccrs.Mercator()
+outProjection=ccrs.Mercator()
 fig= plt.figure(num=1, figsize=(8,4), dpi=dpi, facecolor='w', edgecolor='k')
 ax = fig.add_axes([0, 0, 1, 1], projection=outProjection)
 ax.set_facecolor('black')
 ax.set_extent([minLon, maxLon, minLat, maxLat], crs=ccrs.PlateCarree())  
 
 # Add Boundaries
-#ax.add_feature(cfeature.COASTLINE,linewidth=0.5,edgecolor='white',facecolor='black',zorder=5)
-#ax.add_feature(cfeature.OCEAN,linewidth=0.5,edgecolor='white',facecolor='black',zorder=5)
-#ax.add_feature(cfeature.STATES,linewidth=0.5,edgecolor='white',facecolor='black',zorder=5)
+#ax.add_feature(cfeature.COASTLINE,linewidth=1.5,edgecolor='white',facecolor='black',zorder=6)
+#ax.add_feature(cfeature.OCEAN,linewidth=1.5,edgecolor='white',facecolor='black',zorder=6)
+#ax.add_feature(cfeature.STATES,linewidth=1.5,edgecolor='white',facecolor=None,zorder=90)
+ax.add_feature(USCOUNTIES.with_scale('500k'),linewidth=0.25,edgecolor='black',facecolor='none',zorder=9)
 
 # Plot Data
-# The transform options do not work
-ax.pcolormesh(x, y, airTempC,cmap=color_map,norm=norm,vmin=vmin,vmax=vmax,zorder=8)
-#ax.pcolormesh(x, y, airTempC,transform=ccrs.PlateCarree(),cmap=color_map,norm=norm,vmin=vmin,vmax=vmax,zorder=8)
-#ax.pcolormesh(x, y, airTempC,transform=ccrs.Mercator(),cmap=color_map,norm=norm,vmin=vmin,vmax=vmax,zorder=8)
+ax.pcolormesh(x, y, airTempC,transform=projData,cmap=color_map,norm=norm,vmin=vmin,vmax=vmax,zorder=8)
 
 # Add Colormap
 cax = fig.add_axes([0.1, -0.035, 0.8, 0.03])
@@ -166,14 +167,103 @@ cbar.ax.tick_params(labelsize=15)
 cbar.set_label(unit,size=15)
 
 # Add Titles
-plt.suptitle(outElement+' Analysis For '+init_name,fontsize=17,y=1.10)
+plt.suptitle(outElement+' Analysis For '+init_name,fontsize=17,y=1.05)
 plt.annotate('Source: HRRR, init= '+init_name+'\nMade By Jared Rennie @jjrennie',xy=(1.045, -3.91), xycoords='axes fraction', fontsize=5,backgroundcolor='black',color='white',horizontalalignment='right', verticalalignment='bottom')
 
 # Save Figure
-plt.savefig(main_directory+"/hrrr-nc-grd_"+var.lower()+"_init-"+init_time+"-anls-00.png",bbox_inches='tight') 
+plt.savefig(plot_directory+"/hrrr-nc-grd_"+var.lower()+"_init-"+init_time+"-anls-00.png",bbox_inches='tight') 
 plt.clf()
 plt.close()
 
+####################
+# PART TWO: Forecast
+####################
+# Access HRRR from AWS ... projection dimensions are in url2
+# Note That 'fcst' stands for Forecast 
+fs = s3fs.S3FileSystem(anon=True)
+url1 = 's3://hrrrzarr/sfc/' + date + '/' + date + '_' + hour + 'z_fcst.zarr/' + level + '/' + var + '/' + level
+url2 = 's3://hrrrzarr/sfc/' + date + '/' + date + '_' + hour + 'z_fcst.zarr/' + level + '/' + var
+
+file1 = s3fs.S3Map(url1, s3=fs)
+file2 = s3fs.S3Map(url2, s3=fs)
+
+# Open Datasets
+print("\nOPEN FORECAST DATASETS")
+print ("\t",url1)
+print ("\t",url2)
+ds = xr.open_mfdataset([file1,file2], engine='zarr')
+
+# Plot Each Forecast Hour
+print("\nPLOT FORECASTS")
+
+init_name=year+month+day+': '+hour+'UTC'
+init_time=year+month+day+hour
+
+for time_counter in range(0,len(ds.time.values)):
+    # Organize Forecast Time
+    fcstTime=ds.time.values[time_counter]
+    fcstYear= "%04i" % int(fcstTime.astype('datetime64[h]').astype(str)[0:4])
+    fcstMonth="%02i" % int(fcstTime.astype('datetime64[h]').astype(str)[5:7])
+    fcstDay="%02i" % int(fcstTime.astype('datetime64[h]').astype(str)[8:10])
+    fcstHour="%02i" % int(fcstTime.astype('datetime64[h]').astype(str)[11:13])
+    
+    forecast_name=fcstYear+fcstMonth+fcstDay+': '+fcstHour+'UTC'
+    forecast_time=fcstYear+fcstMonth+fcstDay+fcstHour
+    forecast_hour= "%02i" % (time_counter+1)
+    print("\tT= ",forecast_hour,": ",forecast_name)
+
+    # Get Temperature Data
+    airTemp = ds[var][time_counter].metpy.convert_units('degF')
+
+    # Clip By Shapefile
+    airTemp=airTemp.rio.set_spatial_dims(x_dim="projection_x_coordinate", y_dim="projection_y_coordinate", inplace=True)
+    airTemp=airTemp.rio.write_crs(projData.proj4_params, inplace=True)
+    airTempC = airTemp.rio.clip(geo_shapefile.geometry.apply(mapping), geo_shapefile.crs, drop=False)
+
+    # Get Plotting Coordinates
+    x = airTempC.projection_x_coordinate
+    y = airTempC.projection_y_coordinate
+
+    # Set Up ColorBar
+    vmin=np.min(bounds)
+    vmax=np.max(bounds)
+    extend='both'
+    norm = mcolors.BoundaryNorm(boundaries=bounds, ncolors=256)
+    cm = plt.cm.ScalarMappable(cmap=color_map)
+    cm.set_array(airTemp)
+    cm.set_clim(vmin, vmax)
+
+    # Set Up Figure
+    outProjection=ccrs.Mercator()
+    fig= plt.figure(num=1, figsize=(8,4), dpi=dpi, facecolor='w', edgecolor='k')
+    ax = fig.add_axes([0, 0, 1, 1], projection=outProjection)
+    ax.set_facecolor('black')
+    ax.set_extent([minLon, maxLon, minLat, maxLat], crs=ccrs.PlateCarree())  
+
+    # Add Boundaries
+    #ax.add_feature(cfeature.COASTLINE,linewidth=1.5,edgecolor='white',facecolor='black',zorder=6)
+    #ax.add_feature(cfeature.OCEAN,linewidth=1.5,edgecolor='white',facecolor='black',zorder=6)
+    #ax.add_feature(cfeature.STATES,linewidth=1.5,edgecolor='white',facecolor=None,zorder=90)
+    ax.add_feature(USCOUNTIES.with_scale('500k'),linewidth=0.25,edgecolor='black',facecolor='none',zorder=9)
+
+    # Plot Data
+    ax.pcolormesh(x, y, airTempC,transform=projData,cmap=color_map,norm=norm,vmin=vmin,vmax=vmax,zorder=8)
+
+    # Add Colormap
+    cax = fig.add_axes([0.1, -0.035, 0.8, 0.03])
+    cbar=plt.colorbar(cm, cax=cax,boundaries=bounds,orientation='horizontal',extend=extend,spacing='uniform')
+    cbar.ax.tick_params(labelsize=15)
+    cbar.set_label(unit,size=15)
+
+    # Add Titles
+    plt.suptitle(outElement+' Forecast For: '+forecast_name,fontsize=17,y=1.05)
+    plt.annotate('Source: HRRR, init= '+init_name+' | fcst= '+forecast_name+'\nMade By Jared Rennie @jjrennie',xy=(1.045, -3.91), xycoords='axes fraction', fontsize=5,backgroundcolor='black',color='white',horizontalalignment='right', verticalalignment='bottom')
+
+    # Save Figure
+    plt.savefig(plot_directory+"/hrrr-nc-grd_"+var.lower()+"_init-"+init_time+"-fcst-"+forecast_hour+".png",bbox_inches='tight') 
+    plt.clf()
+    plt.close()
+
 ####################
 # DONE
 ####################