def composite_profile(self,va,plot_time,plot_latlon,wrfouts,outpath, dom=1,mean=1,std=1,xlim=0,ylim=0,fig=0,ax=0, locname=0,ml=-2): """ Loop over wrfout files. Get profile of variable Plot all members Optional standard deviation Optional mean If ax, save image to that axis ml : member level. negative number that corresponds to the folder in absolute string for naming purposes. """ # Set up figure if isinstance(fig,M.figure.Figure): self.fig = fig self.ax = ax else: self.fig, self.ax = plt.subplots() # Plot settings if xlim: xmin, xmax, xint = xlim if ylim: if ylim[1] > ylim[0]: # top to bottom P_top, P_bot, dp = [y*100.0 for y in ylim] else: P_bot, P_top, dp = [y*100.0 for y in ylim] else: P_bot = 100000.0 P_top = 20000.0 dp = 10000.0 # plevs = N.arange(P_bot,P_top,dp) # Get wrfout prototype for information W = WRFOut(wrfouts[0]) lat, lon = plot_latlon datestr = utils.string_from_time('output',plot_time) t_idx = W.get_time_idx(plot_time,) y, x, exact_lat, exact_lon = utils.getXY(W.lats1D,W.lons1D,lat,lon) slices = {'t': t_idx, 'la': y, 'lo': x} #var_slices = {'t': t_idx, 'lv':0, 'la':y, 'lo':x} # Initialise array # Number of levels and ensembles: nPlevs = W.z_dim data = W.get(va,slices) nvarlevs = data.shape[1] nens = len(wrfouts) # 2D: (profile,member) profile_arr = N.zeros((nvarlevs,nens)) composite_P = N.zeros((nPlevs,nens)) # Set up legend labels = [] colourlist = utils.generate_colours(M,nens) # M.rcParams['axes.color_cycle'] = colourlist # Collect profiles for n,wrfout in enumerate(wrfouts): W = WRFOut(wrfout) # Get pressure levels composite_P[:,n] = W.get('pressure',slices)[0,:,0,0] #elev = self.W.get('HGT',H_slices) #pdb.set_trace() # Grab variable profile_arr[:,n] = W.get(va,slices)[0,:,0,0] # Plot variable on graph self.ax.plot(profile_arr[:,n],composite_P[:,n],color=colourlist[n]) member = wrfout.split('/')[ml] labels.append(member) # if locname=='KOAX': pdb.set_trace() # Compute mean, std etc if mean: profile_mean = N.mean(profile_arr,axis=1) profile_mean_P = N.mean(composite_P,axis=1) self.ax.plot(profile_mean,profile_mean_P,color='black') if std: # Assume mean P across ensemble is correct level # to plot standard deviation profile_std = N.std(profile_arr,axis=1) std_upper = profile_mean + profile_std std_lower = profile_mean - profile_std self.ax.plot(std_upper,profile_mean_P,'k--') self.ax.plot(std_lower,profile_mean_P,'k--') if not locname: fname = '_'.join(('profile_comp',va,datestr,'{0:03d}'.format(x),'{0:03d}'.format(y))) + '.png' else: fname = '_'.join(('profile_comp',va,datestr,locname)) + '.png' # Set semi-log graph self.ax.set_yscale('log') # Plot limits, ticks yticks = N.arange(P_bot,P_top+dp*100.0,-100*100.0) # yticks = N.arange(P_bot,P_top+dp,-dp*100) # self.ax.set_yticks(yticks) ylabels = ["%4u" %(p/100.0) for p in yticks] self.ax.set_yticks(yticks) self.ax.set_yticklabels(ylabels) # import pdb; pdb.set_trace() # self.ax.yaxis.tick_right() #plt.axis([-20,50,105000.0,20000.0]) #plt.xlabel(r'Temperature ($^{\circ}$C) at 1000 hPa') #plt.xticks(xticks,['' if tick%10!=0 else str(tick) for tick in xticks]) self.ax.set_ylabel('Pressure (hPa)') #yticks = N.arange(self.P_bot,P_t-1,-10**4) #plt.yticks(yticks,yticks/100) if xlim: self.ax.set_xlim([xmin,xmax]) xticks = N.arange(xmin,xmax+xint,xint) self.ax.set_xticks(xticks) # Flip y axis # Limits are already set self.ax.set_ylim([P_bot,P_top]) # plt.tight_layout(self.fig) #plt.autoscale(enable=1,axis='x') #ax = plt.gca() #ax.relim() #ax.autoscale_view() #plt.draw() self.ax.legend(labels,loc=2,fontsize=6) self.save(outpath,fname) plt.close(self.fig)
# Get all permutations of files for perm in itertools.combinations(files,2): DATA[perm] = {} f1, f2 = perm W1 = WRFOut(f1) W2 = WRFOut(f2) # Make sure times are the same in both files if not W1.wrf_times == W2.wrf_times: print("Times are not identical between input files.") raise Exception # Find indices of each time t_idx = [] for t in ts: t_idx.append(W1.get_time_idx(t)) DATA[perm]['times'] = ts DATA[perm]['values'] = PLOTS[ptype](nc0,nc1,t_idx, energy,lower,upper) if d_return and not d_save: return DATA elif d_save and not d_return: self.save(DATA,d_save,d_fname) elif d_return and d_save: self.save(DATA,d_save,d_fname) return DATA def DE_xyz(self,nc0,nc1,t_idx,energy,*args): """
def compute_diff_energy(ptype, energy, files, times, upper=None, lower=None, d_save=True, d_return=True, d_fname='diff_energy_data'): """ This method computes difference kinetic energy (DKE) or different total energy (DTE, including temp) between WRFout files for a given depth of the atmosphere, at given time intervals :param ptype: 'sum_z' or 'sum_xyz'. 'sum_z' integrates vertically between lower and upper hPa and creates a time series. 'sum_xyz' integrates over the 3D space (again between the upper and lower bounds) and creates 2D arrays. :param energy: 'kinetic' or 'total' :param upper: upper limit of vertical integration :param lower: lower limit of vertical integration :param files: abs paths to all wrfout files :param times: times for computations - tuple format :param d_save: save dictionary to folder (path to folder) :param d_return: return dictionary (True or False) :param d_fname: custom filename :returns: N.ndarray -- time series or list of 2D arrays """ if d_save and not isinstance(d_save, basestring): d_save = os.environ['HOME'] # First, save or output? Can't be neither! if not d_save and not d_return: print("Pick save or output, otherwise it's a waste of computer" "power") raise Exception print("Saving pickle file to {0}".format(d_save)) # Look up the method to use depending on type of plot PLOTS = {'sum_z': self.DE_z, 'sum_xyz': self.DE_xyz} print('Get sequence of time') # Creates sequence of times ts = utils.get_sequence(times) # Dictionary of data DATA = {} print('Get permutations') # Get all permutations of files nperm = len(list(itertools.combinations(files, 2))) print('Start loop') # pdb.set_trace() for n, perm in enumerate(itertools.combinations(files, 2)): print("No. {0} from {1} permutations".format(n, nperm)) perm_start = time.time() DATA[str(n)] = {} f1, f2 = perm W1 = WRFOut(f1) W2 = WRFOut(f2) print('WRFOuts loaded.') #pdb.set_trace() # Make sure times are the same in both files if not N.all(N.array(W1.wrf_times) == N.array(W2.wrf_times)): print("Times are not identical between input files.") raise Exception else: print( "Passed check for identical timestamps between " "NetCDF files") # Find indices of each time print('Finding time indices') t_idx = [] for t in ts: t_idx.append(W1.get_time_idx(t)) print("Calculating values now...") DATA[str(n)]['times'] = ts DATA[str(n)]['values'] = [] for t in t_idx: DATA[str(n)]['values'].append(PLOTS[ptype](W1.nc, W2.nc, t, energy, lower, upper)) DATA[str(n)]['file1'] = f1 DATA[str(n)]['file2'] = f2 print "Calculation #{0} took {1:2.2f} seconds.".format( n, time.time() - perm_start) if d_return and not d_save: return DATA elif d_save and not d_return: #self.save_data(DATA,d_save,d_fname) self.save_data(DATA, d_save, d_fname) #self.json_data(DATA,d_save,d_fname) return elif d_return and d_save: #self.save_data(DATA,d_save,d_fname) self.save_data(DATA, d_save, d_fname) #self.json_data(DATA,d_save,d_fname) return DATA
def compute_diff_energy(ptype,energy,files,times,upper=None,lower=None, d_save=True,d_return=True,d_fname='diff_energy_data'): """ This method computes difference kinetic energy (DKE) or different total energy (DTE, including temp) between WRFout files for a given depth of the atmosphere, at given time intervals :param ptype: '2D' or '3D'. '2D' integrates vertically between lower and upper hPa and creates a time series. '3D' integrates over the 3D space (again between the upper and lower bounds) and creates 2D arrays. :param energy: 'DKE' or 'DTE' :param upper: upper limit of vertical integration :param lower: lower limit of vertical integration :param files: abs paths to all wrfout files :param times: times for computations - tuple format :param d_save: save dictionary to folder (path to folder) :param d_return: return dictionary (True or False) :param d_fname: custom filename :returns: N.ndarray -- time series or list of 2D arrays """ if d_save and not isinstance(d_save,basestring): d_save = os.environ['HOME'] # First, save or output? Can't be neither! if not d_save and not d_return: print("Pick save or output, otherwise it's a waste of computer" "power") raise Exception print("Saving pickle file to {0}".format(d_save)) # Look up the method to use depending on type of plot PLOTS = {'1D':DE_z, '3D':DE_xyz} print('Get sequence of time') # Creates sequence of times ts = utils.get_sequence(times) # Dictionary of data DATA = {} print('Get permutations') # Get all permutations of files nperm = len(list(itertools.combinations(files,2))) print('Start loop') # pdb.set_trace() for n, perm in enumerate(itertools.combinations(files,2)): print("No. {0} from {1} permutations".format(n,nperm)) perm_start = time.time() DATA[str(n)] = {} f1, f2 = perm W1 = WRFOut(f1) W2 = WRFOut(f2) print('WRFOuts loaded.') #pdb.set_trace() # Make sure times are the same in both files if not N.all(N.array(W1.wrf_times) == N.array(W2.wrf_times)): print("Times are not identical between input files.") raise Exception else: print("Passed check for identical timestamps between " "NetCDF files") # Find indices of each time print('Finding time indices') t_idx = [] for t in ts: t_idx.append(W1.get_time_idx(t)) print("Calculating values now...") DATA[str(n)]['times'] = ts DATA[str(n)]['values'] = [] for t in t_idx: DATA[str(n)]['values'].append(PLOTS[ptype](W1.nc,W2.nc,t, energy,lower,upper)) DATA[str(n)]['file1'] = f1 DATA[str(n)]['file2'] = f2 # import pdb; pdb.set_trace() print "Calculation #{0} took {1:2.1f} seconds.".format(n,time.time()-perm_start) if d_return and not d_save: return DATA elif d_save and not d_return: #self.save_data(DATA,d_save,d_fname) utils.save_data(DATA,d_save,d_fname) #self.json_data(DATA,d_save,d_fname) return elif d_return and d_save: #self.save_data(DATA,d_save,d_fname) utils.save_data(DATA,d_save,d_fname) #self.json_data(DATA,d_save,d_fname) return DATA