diff --git a/scripts/plotting_tools.py b/scripts/plotting_tools.py new file mode 100644 index 0000000..9c52c23 --- /dev/null +++ b/scripts/plotting_tools.py @@ -0,0 +1,320 @@ +import os +from os.path import join +from glob import glob + +import pandas as pd +import xarray as xr +import numpy as np + +from matplotlib import pyplot as plt +import cartopy.crs as ccrs +import cartopy.feature as cfeature + +import logging + +ptypes = ['snow', 'rain', 'frzr', 'icep'] + +def get_tle_files(base_path, valid_time, n_members=1, min_lead_time=1, init_step=1): + ''' + Get ptype files at a given valid time. + + Arguments: + base_path (str): path to directory + valid_time (str): YYYY-mm-dd HHMM + n_members (int): for gathering an ensemble of times, default is no ensemble + min_lead_time (int): ? + init_step (int): not implemented + ''' + date = pd.to_datetime(valid_time) + file_names = [] + if n_members == 1: + dt = date - pd.Timedelta(f"1h") + date_str = f"MILES_ptype_hrrr_{dt.strftime('%Y-%m-%d_%H%M')}" + return [join(base_path, dt.strftime("%Y%m%d"), dt.strftime("%H%M"), f"{date_str}_f01.nc")] + for time_delta in range(n_members, 0, -1): + dt = date - pd.Timedelta(f"{time_delta}h") + date_str = f"MILES_ptype_hrrr_{dt.strftime('%Y-%m-%d_%H%M')}" + path = join(base_path, dt.strftime("%Y%m%d"), dt.strftime("%H%M"), f"{date_str}_f{str(time_delta).zfill(2)}.nc") + file_names.append(path) + return file_names + + +def load_data(files, variables): + def process(ds): + ds = ds[variables] + precip_sum = ds[['crain', 'csnow', 'cicep', 'cfrzr']].to_array().sum(dim='variable') + for ptype in ptypes: + ds[f'ML_c{ptype}'] = xr.where(precip_sum >= 1, x=ds[f"ML_c{ptype}"], y=np.nan) + ds[f'ML_{ptype}'] = xr.where(precip_sum >= 1, x=ds[f"ML_{ptype}"], y=np.nan) + ds[f'ML_{ptype}_epi'] = xr.where(precip_sum >= 1, x=ds[f"ML_{ptype}_epi"], y=np.nan) + ds[f'ML_{ptype}_ale'] = xr.where(precip_sum >= 1, x=ds[f"ML_{ptype}_ale"], y=np.nan) + ds[ptype] = ds[f'ML_c{ptype}'].where(ds[f'ML_c{ptype}'] >= 1) * ds[f'ML_{ptype}'] # mask where precip sum >= 1 and where ML_cptype >= 1 + ds[f'c{ptype}'] = xr.where(ds[f'c{ptype}'] == 1, 1, np.nan) + return ds + def process_ensemble(ds): + ds = ds[variables] + precip_sum = ds[['crain', 'csnow', 'cicep', 'cfrzr']].to_array().sum(dim='variable') + for ptype in ptypes: + ds[f'ML_{ptype}'] = xr.where(precip_sum >= 1, x=ds[f'ML_{ptype}'], y=0) + ds[f'ML_{ptype}_epi'] = xr.where(precip_sum >= 1, x=ds[f"ML_{ptype}_epi"], y=0) + ds[f'ML_{ptype}_ale'] = xr.where(precip_sum >= 1, x=ds[f"ML_{ptype}_ale"], y=0) + concat = xr.concat([ds[f'ML_{ptype}'] for ptype in ptypes], dim='ptype') + concat_tle = concat.mean("time") + max_idx = concat_tle.argmax(dim='ptype') + for i, ptype in enumerate(ptypes): + ds[f'ML_c{ptype}'] = xr.where(max_idx == i, 1, np.nan) # set categorical values + ds[ptype] = ds[f'ML_c{ptype}'] * ds[f'ML_{ptype}'] # set categorical multiplied by prob + ds = ds.mean("time").apply(lambda x: x.where(x != 0, np.nan)) # sum and set 0 to nan in one line + return ds + + if len(files) == 1: + ds = xr.open_dataset(files[0]).squeeze() + return process(ds) + else: + ds = xr.open_mfdataset(files, parallel=True) + ds = process_ensemble(ds) + return ds + + +def plot_hrrr_ptype(ds): + '''plot hrrr precip categorizations''' + projection = ccrs.LambertConformal(central_longitude=-97.5, standard_parallels=(38.5, 38.5)) + fig, ax = plt.subplots(figsize=(7, 7), subplot_kw={'projection': projection}) + ax.set_extent([-108, -91, 37, 47.5], crs=ccrs.PlateCarree()) + cmaps = ["Greens", "Blues", "cool", "Reds"] + ptyped = ['rain', 'snow', 'icep', 'frzr'] # ordered in terms of importance, most important plotted last + lat = ds['latitude'] + lon = ds['longitude'] + + for i, ptype in enumerate(ptyped): + h = ax.pcolormesh(lon, lat, ds[f'c{ptype}'] * 0.9, transform=ccrs.PlateCarree(), cmap=cmaps[i], vmin=0, vmax=1) + + ax.add_feature(cfeature.STATES, linewidth=0.5) + return ax + + +def plot_ptype(ds, ptype=None): + '''plots probabilities of each precipitation type based on the type with maximum probability''' + projection = ccrs.LambertConformal(central_longitude=-97.5, standard_parallels=(38.5, 38.5)) + #projection = ccrs.PlateCarree() + fig, ax = plt.subplots(figsize=(7, 7), subplot_kw={'projection': projection}) + ax.set_extent([-108, -91, 37, 47.5], crs=ccrs.PlateCarree()) + cmaps = ["Blues", "Greens", "cool", "Reds"] + ptyped = ['snow', 'rain', 'icep', 'frzr'] + lat = ds['latitude'] + lon = ds['longitude'] + + for i, ptype in enumerate(ptyped): + h = ax.pcolormesh(lon, lat, ds[ptype], transform=ccrs.PlateCarree(), cmap=cmaps[i], vmin=0, vmax=1) + + ax.add_feature(cfeature.STATES, linewidth=0.5) + return ax + + +def plot_probability(ds, ptype): + '''plots probabilities where a certain precipitation type occurs''' + projection = ccrs.LambertConformal(central_longitude=-97.5, standard_parallels=(38.5, 38.5)) + extent = [-108, -91, 37, 47.5] + #projection=ccrs.PlateCarree() + fig, ax = plt.subplots(figsize=(7, 7), subplot_kw={'projection': projection}) + ax.set_extent([-108, -91, 37, 47.5], crs=ccrs.PlateCarree()) + ax.add_feature(cfeature.STATES) + + lat = ds['latitude'] + lon = ds['longitude'] + prob_data = ds[f'ML_{ptype}'] + pcm = ax.pcolormesh(lon, lat, prob_data, transform=ccrs.PlateCarree(), cmap='GnBu') + + cbar = plt.colorbar(pcm, ax=ax, pad=0.025, fraction=0.042) + cbar.set_label(f'Probability of {ptype}') + return ax + + +def plot_uncertainty(ds, ptype): + '''plot epistemic and aleatoric uncertainty where a certain precipitation type occurs''' + projection = ccrs.LambertConformal(central_longitude=-97.5, standard_parallels=(38.5, 38.5)) + extent = [-108, -91, 37, 47.5] + fig, axs = plt.subplots(1, 2, figsize=(15, 7), subplot_kw={'projection': projection}) + + for ax in axs: + ax.set_extent(extent, crs=ccrs.PlateCarree()) + ax.add_feature(cfeature.STATES, linewidth=0.4) + + lat = ds['latitude'] + lon = ds['longitude'] + + ale_data = ds[f'ML_{ptype}_ale'] + epi_data = ds[f'ML_{ptype}_epi'] + + pcm_ale = axs[0].pcolormesh(lon, lat, ale_data, transform=ccrs.PlateCarree(), vmin=0) + plt.colorbar(pcm_ale, ax=axs[0], fraction=0.042, pad=0.025,) + + pcm_epi = axs[1].pcolormesh(lon, lat, epi_data, transform=ccrs.PlateCarree(), vmin=0) + plt.colorbar(pcm_epi, ax=axs[1], fraction=0.042, pad=0.025) + + plt.tight_layout() + return axs + + +def plot_winds(ds, base_plot, ptype=None): + '''plots u and v wind quivers on top of another plot''' + ax = base_plot(ds, ptype) + u = ds['u10'].values + v = ds['v10'].values + lat = ds['latitude'].values + lon = ds['longitude'].values + if isinstance(ax, np.ndarray): # case of multiple subplots (possibly add loop later) + ax[0].quiver(lon[::20, ::20], lat[::20, ::20], u[::20, ::20], v[::20, ::20], width=0.0015, transform=ccrs.PlateCarree()) + ax[1].quiver(lon[::20, ::20], lat[::20, ::20], u[::20, ::20], v[::20, ::20], width=0.0015, transform=ccrs.PlateCarree()) + else: + ax.quiver(lon[::20, ::20], lat[::20, ::20], u[::20, ::20], v[::20, ::20], width=0.0015, transform=ccrs.PlateCarree()) + return ax + + +def plot_temp(ds, base_plot, ptype=None): + '''plots temp contours on top of another plot''' + ax = base_plot(ds, ptype) + lat = ds['latitude'] + lon = ds['longitude'] + temp_data = ds['t2m'] - 273.15 # Convert from Kelvin to Celsius + if isinstance(ax, np.ndarray): + for a in ax: + contour(a, lon, lat, temp_data) + else: + contour(ax, lon, lat, temp_data) + return ax + + +def plot_dpt(ds, base_plot, ptype=None): + '''plots dewpoint contours on top of another plot''' + ax = base_plot(ds, ptype) + lat = ds['latitude'] + lon = ds['longitude'] + dpt_data = ds['d2m'] - 273.15 # Convert from Kelvin to Celsius + + if isinstance(ax, np.ndarray): + for a in ax: + contour(a, lon, lat, dpt_data) + else: + contour(ax, lon, lat, dpt_data) + return ax + + +def contour(ax, lon, lat, data): + contours = ax.contour(lon, lat, data, transform=ccrs.PlateCarree(), levels=10) + #zero_cel = ax.contour(lon, lat, dpt_data, transform=ccrs.PlateCarree(), levels=[0], colors='black', linewidths=2) + ax.clabel(contours, fontsize=10, colors='black') + #cbar = plt.colorbar(contours, orientation="horizontal", fraction=0.15, pad=0.025) + #cbar.set_label('Dew Pt Temperature (°C)') + + +def plot_sp(ds, base_plot, ptype=None): + '''plot surface pressure contours on top of base plot''' + ax = base_plot(ds, ptype) + lat = ds['latitude'] + lon = ds['longitude'] + sp_data = ds['sp'] + + if isinstance(ax, np.ndarray): + for a in ax: + contour(a, lon, lat, sp_data) + else: + contour(ax, lon, lat, sp_data) + return ax + + +def save(title, outname, ax=None): + print(f'saving {outname}') + if isinstance(ax, np.ndarray): + ax[0].set_title(title[0]) + ax[1].set_title(title[1]) + else: + plt.title(title) + plt.savefig(outname, bbox_inches='tight', dpi=300) + plt.close() + + +if __name__ == '__main__': + logger = logging.getLogger(__name__) + logging.basicConfig(level=logging.INFO) + + # Example usage: + base_path = '/glade/derecho/scratch/cbecker/ptype_real_time/winter_2023_2024/hrrr' + valid_time = '2023-12-26 0100' + time = valid_time[:-5] + n_members = 1 + output_dir = './plots/' + + # check if path exists + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + variables = [ + 'u10', 'v10', 'ML_rain', 'ML_crain', 'ML_snow', 'ML_csnow', 'ML_frzr', 'ML_cfrzr', 'ML_icep', 'ML_cicep', + 'crain', 'csnow', 'cfrzr', 'cicep', 'ML_rain_ale', 'ML_rain_epi', 'ML_snow_ale', 'ML_snow_epi', 'ML_frzr_ale', + 'ML_frzr_epi', 'ML_icep_ale', 'ML_icep_epi', 'd2m', 't2m', 'sp' + ] + title = f'' + outname = f'' + + if n_members > 1: + title += f'Time Lagged Ensemble' + outname += f'time_lagged_' + + files = get_tle_files(base_path, valid_time, n_members) + ds = load_data(files, variables) + logger.info(f'{len(files)} files loaded') + + title_ptype = f'{title} {valid_time} Precip\nSnow = Blue, Rain = Green, Sleet = Purple, Freezing Rain = Red' + out_ptype = f'{output_dir}{outname}ptype_{time}' + + plot_hrrr_ptype(ds) + save(title_ptype, f'{out_ptype}_hrrr.png') + + plot_ptype(ds) + save(title_ptype, f'{out_ptype}.png') + + plot_winds(ds, plot_ptype) + save(title_ptype, f'{out_ptype}_quivers.png') + + plot_temp(ds, plot_ptype) + save(title_ptype, f'{out_ptype}_temp.png') + + plot_dpt(ds, plot_ptype) + save(title_ptype, f'{out_ptype}_dpt.png') + + plot_sp(ds, plot_ptype) + save(f'{title_ptype}', f'{out_ptype}_sp.png') + + for ptype in ptypes: + logger.info(f'plotting prob {ptype}') + title_prob = f'{title} probability {ptype} {valid_time}' + out_prob = f'{output_dir}{outname}prob_{time}_{ptype}' + plot_probability(ds, ptype) + save(title_prob, f'{out_prob}.png') + + plot_winds(ds, plot_probability, ptype=ptype) + save(f'{title_prob}', f'{out_prob}_quivers.png') + + plot_temp(ds, plot_probability, ptype=ptype) + save(f'{title_prob}', f'{out_prob}_temp.png') + + plot_dpt(ds, plot_probability, ptype=ptype) + save(f'{title_prob}', f'{out_prob}_dpt.png') + + logger.info(f'plotting uncertainty {ptype}') + title_uncert = [f'{title} Aleatoric Uncertainty {ptype} {valid_time}', f'{title} Epistemic Uncertainty {ptype} {valid_time}'] + out_uncert = f'{output_dir}{outname}uncert_{time}_{ptype}' + ax = plot_uncertainty(ds, ptype) + save(title_uncert, f'{out_uncert}.png', ax) + + ax = plot_winds(ds, plot_uncertainty, ptype=ptype) + save(title_uncert, f'{out_uncert}_quivers.png', ax) + + ax = plot_temp(ds, plot_uncertainty, ptype=ptype) + save(title_uncert, f'{out_uncert}_temp.png', ax) + + ax = plot_dpt(ds, plot_uncertainty, ptype=ptype) + save(title_uncert, f'{out_uncert}_dpt.png', ax) + + \ No newline at end of file