def plot_domains(wrfouts, labels, latlons, outpath, colour=0): """ wrfouts : list of wrfout file paths latlons : dictionary of Nlim,Elim,Slim,Wlim for plot """ fig, ax = plt.subplots(1) # Create basemap first of all #basemap_res = getattr(self.C,'basemap_res',self.D.basemap_res) basemap_res = 'h' m = Basemap(projection='merc', llcrnrlon=latlons['Wlim'], llcrnrlat=latlons['Slim'], urcrnrlon=latlons['Elim'], urcrnrlat=latlons['Nlim'], lat_0=latlons['lat0'], lon_0=latlons['lon0'], resolution=basemap_res, area_thresh=500, ax=ax) m.drawcoastlines() m.drawstates() m.drawcountries() if not isinstance(colour, collections.Sequence): colours = [ 'k', ] * len(wrfouts) else: colours = colour # Get corners of each domain for gridlabel, fpath, colour in zip(labels, wrfouts, colours): W = WRFOut(fpath) print("Plotting domain {0} for {1}".format(gridlabel, fpath)) #Nlim, Elim, Slim, Wlim = W.get_limits() x, y = m(W.lons, W.lats) xl = len(x[0, :]) midpt = len(y[0, :]) / 2 ax.annotate(gridlabel, color=colour, fontsize=10, xy=(x[0, -(0.12 * xl)], y[0, midpt]), bbox=dict(fc='white'), alpha=1, va='center', ha='left') m.plot(x[0, :], y[0, :], colour, lw=2) ax.plot(x[:, 0], y[:, 0], colour, lw=2) ax.plot(x[len(y) - 1, :], y[len(y) - 1, :], colour, lw=2) ax.plot(x[:, len(x) - 1], y[:, len(x) - 1], colour, lw=2) # fpath = os.path.join(self.C.output_root,'domains.png') fname = 'domains.png' fpath = os.path.join(outpath, fname) fig.savefig(fpath) print("Saved to " + fpath)
def get_wrfout(self, wrf_sd=0, wrf_nc=0, dom=0, path_only=0): """Returns the WRFOut instance, given arguments: Optional inputs: wrf_sd : subdirectory for wrf file wrf_nc : filename for wrf file dom : domain for wrf file path_only : only return absolute path """ # Check configuration to see if wrfout files should be # sought inside subdirectories. descend = getattr(self.C, 'wrf_folders_descend', 1) if wrf_sd and wrf_nc: wrfpath = os.path.join(self.C.wrfout_root, wrf_sd, wrf_nc) elif wrf_sd: wrfdir = os.path.join(self.C.wrfout_root, wrf_sd) # print wrfdir wrfpath = utils.wrfout_files_in(wrfdir, dom=dom, unambiguous=1, descend=descend) else: wrfdir = os.path.join(self.C.wrfout_root) wrfpath = utils.wrfout_files_in(wrfdir, dom=dom, unambiguous=1, descend=descend) if path_only: return wrfpath else: return WRFOut(wrfpath)
def spaghetti(self, t, lv, va, contour, wrfouts, outpath, da=0, dom=0): """ wrfouts : list of wrfout files Only change dom if there are multiple domains. """ m, x, y = self.basemap_setup() time_idx = self.W.get_time_idx(t) colours = utils.generate_colours(M, len(wrfouts)) # import pdb; pdb.set_trace() if lv == 2000: lv_idx = None else: print("Only support surface right now") raise Exception lat_sl, lon_sl = self.get_limited_domain(da) slices = {'t': time_idx, 'lv': lv_idx, 'la': lat_sl, 'lo': lon_sl} # self.ax.set_color_cycle(colours) ctlist = [] for n, wrfout in enumerate(wrfouts): self.W = WRFOut(wrfout) data = self.W.get(va, slices)[0, ...] # m.contour(x,y,data,levels=[contour,]) ct = m.contour(x, y, data, colors=[ colours[n], ], levels=[ contour, ], label=wrfout.split('/')[-2]) print("Plotting contour level {0} for {1} from file \n {2}".format( contour, va, wrfout)) # ctlist.append(ct) # self.ax.legend() # labels = [w.split('/')[-2] for w in wrfouts] # print labels # self.fig.legend(handles=ctlist) # plt.legend(handles=ctlist,labels=labels) #labels,ncol=3, loc=3, # bbox_to_anchor=[0.5,1.5]) datestr = utils.string_from_time('output', t, tupleformat=0) lv_na = utils.get_level_naming(va, lv) naming = ['spaghetti', va, lv_na, datestr] if dom: naming.append(dom) fname = self.create_fname(*naming) self.save(outpath, fname)
def colocate_WRF_map(self,wrfdir): searchfor = os.path.join(wrfdir,'wrfout*') # pdb.set_trace() files = glob.glob(searchfor) W = WRFOut(files[0]) limits = {} limits['Nlim'] = W.lats.max() limits['Slim'] = W.lats.min() limits['Wlim'] = W.lons.min() limits['Elim'] = W.lons.max() return limits
def std(ncfiles,va,tidx,lvidx): """ Find standard deviation in along axis of ensemble members. Returns matrix x-y for plotting """ for n, nc in enumerate(ncfiles): W = WRFOut(nc) slices = {'lv':lvidx, 't':tidx} va_array = W.get(va,slices) dims = va_array.shape if n==0: all_members = N.zeros([len(ncfiles),1,1,dims[-2],dims[-1]]) all_members[n,0,0,:,:] = va_array[...] std = N.std(all_members,axis=0).reshape([dims[-2],dims[-1]]) # pdb.set_trace() return std
def plot_skewT(self, plot_time, plot_latlon, dom=1, save_output=0, composite=0): wrfouts = self.wrfout_files_in(self.C.wrfout_root) for wrfout in wrfouts: if not composite: W = WRFOut(wrfout) ST = SkewT(self.C, W) ST.plot_skewT(plot_time, plot_latlon, dom, save_output) nice_time = utils.string_from_time('title', plot_time) print("Plotted Skew-T for time {0} at {1}".format( nice_time, plot_latlon)) else: #ST = SkewT(self.C) pass
def plot_variable2D(self,va,pt,en,lv,p2p,na=0,da=0): """Plot a longitude--latitude cross-section (bird's-eye-view). Use Basemap to create geographical data ======== REQUIRED ======== va = variable(s) pt = plot time(s) nc = ensemble member(s) lv = level(s) p2p = path to plots ======== OPTIONAL ======== da = smaller domain area(s), needs dictionary || DEFAULT = 0 na = naming scheme for plot files || DEFAULT = get what you're given """ va = self.get_sequence(va) pt = self.get_sequence(pt,SoS=1) en = self.get_sequence(en) lv = self.get_sequence(lv) da = self.get_sequence(da) perms = self.make_iterator(va,pt,en,lv,da) # Find some way of looping over wrfout files first, avoiding need # to create new W instances # print("Beginning plotting of {0} figures.".format(len(list(perms)))) #pdb.set_trace() for x in perms: va,pt,en,lv,da = x W = WRFOut(en) # wrfout file class using path F = BirdsEye(self.C,W,p2p) # 2D figure class F.plot2D(va,pt,en,lv,da,na) # Plot/save figure pt_s = utils.string_from_time('title',pt) print("Plotting from file {0}: \n variable = {1}" " time = {2}, level = {3}, area = {4}.".format(en,va,pt_s,lv,da))
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)
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
raise Exception # Look up the method to use depending on type of plot PLOTS = {'sum_z':self.DE_z, 'sum_xyz':self.DE_xyz} # Creates sequence of times ts = self.get_sequence(times) # Dictionary of data DATA = {} # 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)
def plot_diff_energy(self, ptype, energy, time, folder, fname, p2p, plotname, V, no_title=0, ax=0): """ folder : directory holding computed data fname : naming scheme of required files p2p : root directory for plots V : constant values to contour at """ sw = 0 DATA = self.load_data(folder, fname, format='pickle') if isinstance(time, collections.Sequence): time = calendar.timegm(time) #for n,t in enumerate(times): for pn, perm in enumerate(DATA): f1 = DATA[perm]['file1'] f2 = DATA[perm]['file2'] if sw == 0: # Get times and info about nc files # First time to save power W1 = WRFOut(f1) permtimes = DATA[perm]['times'] sw = 1 # Find array for required time x = N.where(N.array(permtimes) == time)[0][0] data = DATA[perm]['values'][x][0] if not pn: stack = data else: stack = N.dstack((data, stack)) stack_average = N.average(stack, axis=2) if ax: kwargs1 = {'ax': ax} kwargs2 = {'save': 0} #birdseye plot with basemap of DKE/DTE F = BirdsEye(self.C, W1, **kwargs1) # 2D figure class #F.plot2D(va,t,en,lv,da,na) # Plot/save figure tstr = utils.string_from_time('output', time) fname_t = ''.join((plotname, '_{0}'.format(tstr))) # fpath = os.path.join(p2p,fname_t) fig_obj = F.plot_data(stack_average, 'contourf', p2p, fname_t, time, V, no_title=no_title, **kwargs2) if ax: return fig_obj