diff --git a/utilities/plotting/forcing_animation.py b/utilities/plotting/forcing_animation.py new file mode 100755 index 0000000000..a6bc4cffd2 --- /dev/null +++ b/utilities/plotting/forcing_animation.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python +""" +@author Nels J. Frazier +@date 01-02-2024 +@version 0.0.0 + +Utility to generate animations in gif format in both gridded form +and catchment aggregated form. Use forcing_animation --help for +descriptions of options and required inputs. + +""" +import geopandas as gpd +import matplotlib.pyplot as plt +#import PIL +#PIL.Image.MAX_IMAGE_PIXELS = 188745354+10000 +from matplotlib.animation import FuncAnimation, PillowWriter, FFMpegWriter +import seaborn as sns +import xarray as xr +from functools import partial +import cartopy.crs as ccrs +import cartopy.feature as cfeature +#import pickle +import argparse +from pyogrio import read_dataframe + +def plot_catchment_frame(ax, forecast, forcing, catchments, i): + t = forcing[i] + #each dataset is a slice of time + ts = t.time.dt.strftime("%Y-%m-%d %H:%M:%S").item() + ax.set_title(forecast+"\n"+ts+"\nPrecipitation", fontsize=20) + #don't plot new, update old!!! + ax.collections[0].set_array(t.values) + print("Plotting time ", i) + return ax + +from matplotlib.collections import QuadMesh +def plot_grid_frame(ax, forecast, forcing, cmap, i): + #Get time at int index i + t = forcing['RAINRATE'][i:i+1] + ts = t.time.dt.strftime("%Y-%m-%d %H:%M:%S").item() + ax.axes.set_title(forecast+"\n"+ts+"\nPrecipitation", fontsize=20) + #update + ax.set_array(t[0].values) + print("Plotting time ", i) + return ax + +cycles = ["Short", "Medium", "Long"] +sources = ['HRRR-RAP', 'CFS', 'GFS', 'Other'] + +parser = argparse.ArgumentParser() +#TODO add cariable arg +#TODO add cmap arg +parser.add_argument('-p', '--preview', help="Plot the first frame and view only", required=False, action="store_true") +parser.add_argument('-s', '--switch', help="switch the orientation of the colorbar from horizontal to vertical", required=False, action="store_true") +parser.add_argument('-d', '--dpi', required=False, default=100, help="Resolution to save gif", type=int) +parser.add_argument("source", choices=sources) +parser.add_argument("cycle", choices=cycles) +commands = parser.add_subparsers(dest="command") + +h_commands = commands.add_parser("hydrofabric") +h_commands.add_argument('hydrofabric', metavar="hydrofabric geopackage", type=argparse.FileType('r')) +h_commands.add_argument('data', metavar="netcdf precip", type=argparse.FileType('r')) + +g_commands = commands.add_parser("grid") +g_commands.add_argument('data', metavar="gridded netcdf precip", type=argparse.FileType('r')) + +args = parser.parse_args() + +forecast = args.cycle+" Range "+args.source +crs = '' +epsg = '' +if(args.command == 'hydrofabric'): + epsg = '5070' + crs = ccrs.epsg(epsg) + crs = ccrs.PlateCarree() +elif(args.command == 'grid'): + epsg = '4326' + # doesn't work: crs= ccrs.epsg('4326') + crs = ccrs.PlateCarree() #NJF why this one??? + +forcing = xr.open_dataset(args.data.name, decode_cf=False) +forcing.time.attrs['units'] = forcing.time.attrs['reference_time'] +forcing = xr.decode_cf(forcing) + +###Setup general figure +#cmap = sns.color_palette('flare', as_cmap=True) +#Single hughe blue, really light on the 0 side... +#cmap = sns.color_palette('Blues', as_cmap=True) +#Little better constrast, still high desat on the low end +cmap = sns.color_palette("light:b", as_cmap=True) +#Reasonably good "blue ish" cube helix +#cmap = sns.color_palette("ch:start=.2,rot=-.2", as_cmap=True) + +bkgrnd = sns.color_palette('flare', as_cmap=True) +#cmap.set_under(bkgrnd(50), alpha=0.05) +cmap.set_under("bisque", alpha=0.1) +e_color = 'none' +vmin = forcing['RAINRATE'].min().item() +vmin = 0.0009 +vmax = forcing['RAINRATE'].max().item() + +fig = plt.figure(figsize=(19.20, 10.80), dpi=100) +ax = plt.axes(projection=crs) +#plt.suptitle("Short Range") +if(args.switch): + cax = fig.add_axes([0.9, 0.1, 0.03, 0.8]) #vertical + orient = "vertical" +else: + cax = fig.add_axes([0.05, 0.05, 0.9, 0.03]) #horizontal + orient = "horizontal" + norm = plt.Normalize(vmin=vmin, vmax=vmax) +sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm) +# fake up the array of the scalar mappable. Urgh... +sm._A = [] + +cbar = fig.colorbar(sm, cax=cax, orientation=orient) +cbar.set_label(forcing['RAINRATE'].attrs['units']) +#ax = fig.gca() +ax.set_axis_off() +ax.title.set_size(20) + +interval = 100 +fps = 20 +if(args.cycle == "Short"): + interval = 5 + fps = 10 + +if(args.command == "hydrofabric"): + fs = forcing["catchment_ids"].to_series() + print( fs[ fs.duplicated() ].values) + # WAY faster reading with pyogrio and use_arrow (though just pyogrio is a significant imporvement!) + catchments = gpd.read_file(args.hydrofabric.name, layer="divides", include_fields=['divide_id'], engine='pyogrio', use_arrow=True) + # union = gpd.GeoSeries( catchments.unary_union, crs=catchments.crs) + union = gpd.read_file("domain.gpkg", engine='pyogrio') + + print("Projecting domain to epsg:4326") + union = union.to_crs(epsg='4326') + catchments = catchments.to_crs(epsg='4326') + + forcing = forcing.set_index({"catchment ids":"catchment_ids"}).to_array().squeeze().drop_vars('variable') + #Get the catchment geopandas index, strip to int for xarray index + cat_ids = catchments['divide_id'].apply(lambda x: int(x.split("-")[1])) + #This is one way to re-org the data along the geopandas index + #forcing = forcing.loc[:, cat_ids.values] + #But this should be much faster! + #FIXME issues with duplicate ids in forcing files... + forcing = forcing.drop_duplicates("catchment ids") + forcing = forcing.reindex({"catchment ids":cat_ids.values}) + # This was a good idea -- cache the axes, but won't work for general usage + # and also just doesn't work to deserialize the GeoAxes from cartopy... + # if( Path(".catchment_axes_caches").exists() ): + # with open('.catchment_axes_caches', 'rb') as cache: + # ax = pickle.load(cache) + # else: + # ax = catchments.plot(ax=ax, legend=False, edgecolor=e_color, facecolor='white') + # with open('.catchment_axes_caches', 'wb') as cache: + # pickle.dump(ax, cache) + + # epsg 5070 projected bounds: -7013714.59 967179.63 -5299735.56 2630462.94 + # This cuts off florida and texas!!! For now, just reproject to wgs84 and plot + # on the PlateCarree projection + # might consider trying + print("Plotting base...") + ax.add_geometries( union.geometry, facecolor="none", edgecolor="black", crs=crs) + #ax.set_extent( catchments.total_bounds, crs=crs ) + box = catchments.total_bounds + ax.set_xlim((box[0], box[2])) + ax.set_ylim((box[1], box[3])) + ax = catchments.plot(ax=ax, legend=False, edgecolor=e_color, facecolor='white', vmin=vmin, vmax=vmax, norm=norm) + #FIXME try the union before projecting, then re-project the resulting geometry + # Turns out, this fixes the topology error, but is slow, for now hacking in Mike's outline + #ax.add_geometries( [catchments.unary_union], crs=crs, facecolor=None, edgecolor="black") + ax.collections[0].set_cmap(cmap) + ax.collections[0].set_clim(vmin, vmax) + + #ax.add_feature(cfeature.STATES.with_scale('50m')) + if(args.preview): + plot_catchment_frame(ax, forecast, forcing, cat_ids, 6) + else: + ani = FuncAnimation(fig, partial(plot_catchment_frame, ax, forecast, forcing, cat_ids), frames=forcing.coords['time'].size, interval=interval) +elif(args.command == "grid"): + ax.add_feature(cfeature.STATES.with_scale('110m')) + ax = forcing['RAINRATE'].isel(time=0).plot(ax=ax, cmap=cmap, add_colorbar=False, vmin=vmin, vmax=vmax, norm=norm) + # print(ax) + #plt.show() + # os._exit(1) + #odr = cartopy.feature.NaturalEarthFeature(category='cultural', + #name='admin_0_boundary_lines_land', scale='50m', facecolor='none', alpha=0.7) + #ax.add_feature(odr) + if(args.preview): + plot_grid_frame( ax, forecast, forcing, cmap, 0) + else: + ani = FuncAnimation(fig, partial(plot_grid_frame, ax, forecast, forcing, cmap), frames=forcing.coords['time'].size, interval=interval) + + + +# catchments = gpd.read_file("gauge_01073000.gpkg", layer="divides", include_fields=['divide_id']) +# forcing = xr.open_dataset("data/Long_Range_Hydrofabric_Precip.nc", decode_cf=False) + +# forecast="Short Range" #TODO add product (CFS, GFS ect...) +# forecast="Long Range" #TODO add product (CFS, GFS ect...) + + +print("NUM FRAMES") +print(forcing.coords['time'].size) +#plot_frame(ax, 224) + +#plt.show() + +#ani.save("forcings_"+forecast.replace(" ","_")+".gif", dpi=1200, writer=PillowWriter(fps=20)) +#plt.tight_layout() +#plt.subplots_adjust(top=0.85) +if(args.preview): + plt.show() +else: + #ani.save("forcings_"+forecast.replace(" ","_")+".mp4", writer=FFMpegWriter(fps=fps) ) + ani.save(args.command+"_forcings_"+forecast.replace(" ","_")+".gif", writer="ffmpeg", dpi=args.dpi) +plt.close() +#plt.savefig("test.png", dpi=1200, bbox_inches='tight', +# transparent=True, +# pad_inches=0) \ No newline at end of file