-
Notifications
You must be signed in to change notification settings - Fork 100
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
adding a mostly blank file which will be used as main driver file
- Loading branch information
drewmresnick
committed
Feb 10, 2021
1 parent
7d29489
commit ca56b80
Showing
1 changed file
with
4 additions
and
205 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,206 +1,5 @@ | ||
| #Load libraries | ||
| import xarray as xr | ||
| import numpy as np | ||
| import pandas as pd | ||
| #import eccodes | ||
| from sklearn.cluster import KMeans | ||
| from sklearn.decomposition import PCA | ||
| import cartopy.crs as ccrs | ||
| from cartopy import feature | ||
| import matplotlib.pyplot as plt | ||
| from matplotlib.colors import LinearSegmentedColormap | ||
| # key PyWR functions are imported here | ||
| from PyWR import * | ||
| #driver file | ||
| import os | ||
| import glob | ||
|
|
||
|
|
||
| #The data used in this diagnostic was downloaded from the IRI: | ||
| #https://iridl.ldeo.columbia.edu/?Set-Language=en | ||
| #If data of a similar nature is desired, refer to the prep_data.py file for instructions on how | ||
| #to download anomaly data from the http server | ||
|
|
||
| #need to be changed to where the files are located in framework | ||
| reanalysis = xr.open_dataset('WUS/data/hgt_NNRP_rean.nc', decode_cf = True, decode_times = True).stack(time=['T'], grid=['Y', 'X']) | ||
| rainfall = xr.open_dataset('WUS/data/rainfall_cpc.nc', decode_cf = True, decode_times = True).stack(time=['T'], grid=['Y', 'X']) | ||
| t2m = xr.open_dataset('WUS/data/t2m_cpc.nc', decode_cf = True, decode_times = True).stack(time=['T'], grid=['Y', 'X']) | ||
| uwnd = xr.open_dataset('WUS/data/u_NNRP_rean.nc', decode_cf = True, decode_times = True).stack(time=['T'], grid=['Y', 'X']) | ||
| uwnd = xr.open_dataset('WUS/data/v_NNRP_rean.nc', decode_cf = True, decode_times = True).stack(time=['T'], grid=['Y', 'X']) | ||
|
|
||
| #get rid of dummy pressure coordinate | ||
| reanalysis=reanalysis.isel(P=0) | ||
| uwnd=uwnd.isel(P=0) | ||
| vwnd=vwnd.isel(P=0) | ||
|
|
||
|
|
||
| #viewing the data | ||
| print(reanalysis) | ||
| print(rainfall) | ||
| print(t2m) | ||
| print(vwnd) | ||
| print(uwnd) | ||
|
|
||
|
|
||
| #DIMENSION REDUCTION: choose a percentage of variance explained that we will require | ||
| n_eof = get_number_eof(X=reanalysis['adif'].values, var_to_explain=0.9, plot=True) | ||
|
|
||
| #project the data onto the leading EOFs to get the principal component time series | ||
| #We will retain the PCA model for use later. The reanalysis_pc variable is now indexed [time, EOF] | ||
| pca_model = PCA(n_components=n_eof).fit(reanalysis['adif'].values) | ||
| reanalysis_pc = pca_model.transform(reanalysis['adif'].values) | ||
|
|
||
|
|
||
| #REANALYSIS WEATHER TYPING: perform the clustering. We will manually specify the number of clusters we want to create and the number of simulations we want to run | ||
| ncluster = 6 # use 6 WTs | ||
| n_sim = 50 # typically 25-50 -- try 25 for quick preliminary computation only | ||
|
|
||
| centroids, wtypes = loop_kmeans(X=reanalysis_pc, n_cluster=ncluster, n_sim=n_sim) | ||
| class_idx, best_part = get_classifiability_index(centroids) | ||
| print('The classifiability index is {}'.format(class_idx)) | ||
|
|
||
|
|
||
| #Now that we have identified a suitable partition, we can use it to keep only the corresponding centroid and set of weather type labels. Use centroids to define KMeans object | ||
| best_fit = KMeans(n_clusters=ncluster, init=centroids[best_part, :, :], n_init=1, max_iter=1).fit(reanalysis_pc) | ||
|
|
||
| # start with reanalysis | ||
| reanalysis_composite = reanalysis.copy() | ||
| model_clust = best_fit.fit_predict(reanalysis_pc) # get centroids | ||
| weather_types = xr.DataArray( | ||
| model_clust, | ||
| coords = {'time': reanalysis_composite['time']}, | ||
| dims='time' | ||
| ) | ||
| reanalysis_composite['WT'] = weather_types | ||
| reanalysis_composite = reanalysis_composite.groupby('WT').mean(dim='time').unstack('grid')['adif'] | ||
| reanalysis_composite['M'] = 0 | ||
|
|
||
| wt_anomalies = [] # initialize empty list | ||
| wt_anomalies.append(reanalysis_composite) | ||
|
|
||
| wt_anomalies = xr.concat(wt_anomalies, dim='M') # join together | ||
| wt_anomalies['WT'] = wt_anomalies['WT'] + 1 # start from 1 | ||
|
|
||
|
|
||
| #FIGURE: prepare a figure with rainfall and temperature composites | ||
| #Hashed out options for adding wind arrows, and additional plot labels | ||
|
|
||
| X, Y = np.meshgrid(reanalysis['adif'].X, reanalysis['adif'].Y) | ||
| map_proj = ccrs.PlateCarree() #ccrs.Orthographic(-110, 10) | ||
| data_proj = ccrs.PlateCarree() | ||
| wt_unique = np.unique(wt_anomalies['WT']) | ||
| figsize = (14, 8) | ||
|
|
||
| #WT proportions | ||
| wt=weather_types.to_dataframe(name='WT') | ||
| wt=wt+1 | ||
| #wt.to_netcdf('data/t2m_cpc.nc', format="NETCDF4") | ||
| wt_counts = wt.groupby('WT').size().div(wt['WT'].size) | ||
| wt_counts | ||
|
|
||
| xmin,xmax = reanalysis['X'].min(), reanalysis['X'].max() | ||
| ymin,ymax = reanalysis['Y'].min(), reanalysis['Y'].max() | ||
|
|
||
| # Set up the Figure | ||
| plt.rcParams.update({'font.size': 12}) | ||
| fig, axes = plt.subplots( | ||
| nrows=3, ncols=len(wt_unique), subplot_kw={'projection': map_proj}, | ||
| figsize=figsize, sharex=True, sharey=True | ||
| ) | ||
|
|
||
| # Loop through | ||
| for i,w in enumerate(wt_unique): | ||
| def selector(ds): | ||
| times = wt.loc[wt['WT'] == w].index | ||
| ds = ds.sel(time = np.in1d(ds.unstack('time')['T'], times)) | ||
| ds = ds.mean(dim = 'time') | ||
| return(ds) | ||
|
|
||
| # Top row: geopotential height anomalies | ||
| ax = axes[0, i] | ||
| ax.set_title('WT {}: {:.1%} of days'.format(w, wt_counts.values[i])) | ||
| C0 = selector(reanalysis['adif']).unstack('grid').plot.contourf( | ||
| transform = data_proj, | ||
| ax=ax, | ||
| cmap='PuOr', | ||
| extend="both", | ||
| levels=np.linspace(-2e2, 2e2, 21), | ||
| add_colorbar=False, | ||
| add_labels=False | ||
| ) | ||
| ax.coastlines() | ||
| ax.add_feature(feature.BORDERS) | ||
| #ax.set_extent([-95, -65, -12, 12]) | ||
|
|
||
| # # add wind arrows | ||
| # U = selector(uwnd).adif.values | ||
| # V = selector(vwnd).adif.values | ||
| # magnitude = np.sqrt(U**2 + V**2) | ||
| # strongest = magnitude > np.percentile(magnitude, 50) | ||
| # Q = ax.quiver( | ||
| # X[strongest], Y[strongest], U[strongest], V[strongest], | ||
| # transform=data_proj, | ||
| # width=0.001, scale=0.8,units='xy' | ||
| # ) | ||
|
|
||
| # Middle row: rainfall anomalies | ||
| ax = axes[1, i] | ||
| C1 = selector(rainfall['adif']).unstack('grid').plot.contourf( | ||
| transform = data_proj, | ||
| ax=ax, | ||
| cmap = 'BrBG', | ||
| extend="both", | ||
| levels=np.linspace(-2, 2, 13), | ||
| add_colorbar=False, | ||
| add_labels=False | ||
| ) | ||
| ax.coastlines() | ||
| ax.add_feature(feature.BORDERS) | ||
| #ax.set_extent([-95, -75, -9, 5]) | ||
|
|
||
| # Bottom row: tepmperature anomalies | ||
| ax = axes[2, i] | ||
| C2 = selector(t2m['asum']).unstack('grid').plot.contourf( | ||
| transform = data_proj, | ||
| ax=ax, | ||
| cmap = 'RdBu_r', | ||
| extend="both", | ||
| levels=np.linspace(-2, 2, 13), | ||
| add_colorbar=False, | ||
| add_labels=False | ||
| ) | ||
| ax.coastlines() | ||
| ax.add_feature(feature.BORDERS) | ||
| #ax.set_extent([-95, -70, -9, 5]) | ||
| ax.tick_params(colors='b') | ||
|
|
||
| # # Add Colorbar | ||
| plt.tight_layout() | ||
| fig.subplots_adjust(right=0.94) | ||
| cax0 = fig.add_axes([0.97, 0.65, 0.0075, 0.3]) | ||
| cax1 = fig.add_axes([0.97, 0.33, 0.0075, 0.3]) | ||
| cax2 = fig.add_axes([0.97, 0.01, 0.0075, 0.3]) | ||
| cbar0 = fig.colorbar(C0, cax = cax0) | ||
| cbar0.formatter.set_powerlimits((4, 4)) | ||
| cbar0.update_ticks() | ||
| cbar0.set_label(r'$zg_{500}$ anomaly [$m^2$/$s^2$]', rotation=270) | ||
| cbar0.ax.get_yaxis().labelpad = 20 | ||
| cbar1 = fig.colorbar(C1, cax=cax1) | ||
| cbar1.set_label('Precip. anomaly [mm/d]', rotation=270) | ||
| cbar1.ax.get_yaxis().labelpad = 20 | ||
| cbar2 = fig.colorbar(C2, cax=cax2) | ||
| cbar2.set_label('T2m anomaly [$^o$C]', rotation=270) | ||
| cbar2.ax.get_yaxis().labelpad = 20 | ||
|
|
||
| # Format these axes | ||
|
|
||
|
|
||
| #Add plot labels | ||
| # letters = string.ascii_lowercase | ||
| # for i, ax in enumerate(axes.flat): | ||
| # label = '({})'.format(letters[i]) | ||
| # t = ax.text(0.05, 0.9, label, fontsize=11, transform=ax.transAxes) | ||
| # t.set_bbox(dict(facecolor='white', edgecolor='gray')) | ||
|
|
||
| # Add a quiver key | ||
| #k = plt.quiverkey(Q, 0.9, 0.7, 1, '1 m/s', labelpos='E', coordinates='figure') | ||
|
|
||
| fig.savefig('figs/wt_composite.pdf', bbox_inches='tight') #this needs to be changed to appropriate folder in framework | ||
| plt.show() | ||
| os.environ["pr_file"] = os.environ |