Ejemplo n.º 1
1
Archivo: sal.py Proyecto: NicWayand/WEM
    def __init__(self,Wctrl_fpath,Wmod_fpath,vrbl,utc,lv=False,
                    accum_hr=False,radar_datadir=False,thresh=False,
                    footprint=500):
        self.utc = utc
        self.C = {}
        self.M = {}
        self.thresh = thresh
        self.footprint = footprint

        self.M['WRFOut'] = WRFOut(Wmod_fpath)
        self.dx = self.M['WRFOut'].dx
        self.compute_d(self.M['WRFOut'])

        if Wctrl_fpath is False and (vrbl=='REFL_comp' or vrbl=='cref'):
            use_radar_obs = True
            self.C['data'] = self.get_radar_verif(utc,radar_datadir)
        else:
            use_radar_obs = False
            Wctrl = WRFOut(Wctrl_fpath)

        # Get 2D grids for ctrl and model
        if vrbl == 'accum_precip':
            if not accum_hr:
                raise Exception("Need to set accumulation hours.")
            self.C['data'] = Wctrl.compute_accum_rain(utc,accum_hr)[0,0,:,:]
            self.M['data'] = self.M['WRFOut'].compute_accum_rain(utc,accum_hr)[0,0,:,:]
        else:
            self.M['data'] = self.M['WRFOut'].get(vrbl,level=lv,utc=utc)[0,0,:,:]
            if not use_radar_obs:
                self.C['data'] = Wctrl.get(vrbl,level=lv,utc=utc)[0,0,:,:]

        # Set negative values to 0 
        # if vrbl == 'REFL_comp':
        self.C['data'][self.C['data']<0] = 0
        self.M['data'][self.M['data']<0] = 0
        self.vrbl = vrbl
       
        """
        fig, ax = plt.subplots(1)
        ax.pcolor(self.C['data'])
        fig.savefig('/home/jrlawson/public_html/bowecho/SALtests/obs_pcolor.png')
        plt.close(fig)
        fig, ax = plt.subplots(1)
        ax.pcolor(self.M['data'])
        fig.savefig('/home/jrlawson/public_html/bowecho/SALtests/mod_pcolor.png')
        plt.close(fig)
        import pdb; pdb.set_trace()
        """

        self.identify_objects()
        self.compute_amplitude()
        self.compute_location()
        self.compute_structure()

        print("S = {0}    A = {1}     L = {2}".format(self.S,self.A,self.L))
Ejemplo n.º 2
0
    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)
Ejemplo n.º 3
0
    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)
Ejemplo n.º 4
0
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)
Ejemplo n.º 5
0
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
Ejemplo n.º 6
0
def std(ncfiles,vrbl,utc=False,level=False,other=False,axis=0):
    """
    Find standard deviation in along axis of ensemble
    members. Returns matrix x-y for plotting
    """
    for n, nc in enumerate(ncfiles):
        # print("Ensemble member {0} loaded.".format(n))
        W = WRFOut(nc)
        vrbl_array = W.get(vrbl,utc=utc,level=level,other=other)
        if vrbl=='cref':
            vrbl_array[vrbl_array<0] = 0

        if n==0:
            dims = [len(ncfiles),] + list(vrbl_array.shape)
            all_members = N.zeros(dims)
        all_members[n,...] = vrbl_array[...]
        # import pdb; pdb.set_trace()

    std = N.std(all_members,axis=axis)
    return std
Ejemplo n.º 7
0
    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
Ejemplo n.º 8
0
def std_ttest(ncfiles1,ncfiles2,vrbl,utc=False,level=False,other=False):
    """
    Find standard deviation in along axis of ensemble
    members. Returns matrix x-y for plotting. Returns sig test
    """
    std = []
    std_ave = []

    for ncfiles in (ncfiles1, ncfiles2):
        for n, nc in enumerate(ncfiles):
            # print("Ensemble member {0} loaded.".format(n))
            W = WRFOut(nc)
            vrbl_array = W.get(vrbl,utc=utc,level=level,other=other)
            if vrbl=='cref':
                vrbl_array[vrbl_array<0] = 0

            if n==0:
                dims = [len(ncfiles),] + list(vrbl_array.shape)
                all_members = N.zeros(dims)
            all_members[n,...] = vrbl_array[...]
            # sample[n,...] = vrbl_array[
            # import pdb; pdb.set_trace()

        sample=all_members
        std.append(N.std(sample,axis=0)[0,0,:,:])
        std_ave.append(N.std(all_members, axis=None))
        
    # N.random.shuffle(std[0])
    # N.random.shuffle(std[1])

    choice0 = N.random.choice(std[0].flatten(),size=1000)
    choice1 = N.random.choice(std[1].flatten(),size=1000)
    # import pdb; pdb.set_trace()

    # tstat, pvalue = scipy.stats.ttest_rel(std[0],std[1],axis=None)
    tstat, pvalue = scipy.stats.ttest_rel(choice0,choice1,axis=None)

    return std_ave[0], std_ave[1], tstat, pvalue
Ejemplo n.º 9
0
    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)
Ejemplo n.º 10
0
 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
Ejemplo n.º 11
0
    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))
Ejemplo n.º 12
0
class BirdsEye(Figure):
    def __init__(self,config,wrfout,ax=0,fig=0):
        super(BirdsEye,self).__init__(config,wrfout,ax=ax,fig=fig)

    def get_contouring(self,vrbl='user',lv='user',**kwargs):
        """
        Returns colourmap and contouring levels
        
        Options keyword arguments:
        V   :   manually override contour levels
        """
        
        data = self.data.reshape((self.la_n,self.lo_n))
        
        # List of args and dictionary of kwargs
        plotargs = [self.x,self.y,data]
        plotkwargs = kwargs

        # cmap = getattr(kwargs,'cmap',plt.cm.jet)


        if vrbl=='user':
            pass
            
        else:
            S = Scales(vrbl,lv)
            if S.cm:
                plotkwargs['cmap'] = S.cm
            if isinstance(S.clvs,N.ndarray):
                plotkwargs['levels'] = S.clvs
            
            # if self.mplcommand == 'contour':
                # multiplier = S.get_multiplier(vrbl,lv)
        if 'clvs' in kwargs:
            if isinstance(kwargs['clvs'],N.ndarray):
                plotkwargs['levels'] = kwargs['clvs']
                kwargs.pop('clvs')
        # pdb.set_trace()
        return plotargs, plotkwargs
            
    def plot_data(self,data,mplcommand,p2p,fname,pt,no_title=1,save=1,**kwargs):
        """
        Generic method that plots any matrix of data on a map

        Inputs:
        data        :   lat/lon matrix of data
        vrbl        :   variable type for contouring convention
        m           :   basemap instance
        mplcommand  :   contour or contourf etc
        p2p         :   path to plots
        fname       :   filename for plot
        V           :   scale for contours
        no_title    :   switch to turn off title
        save        :   whether to save to file
        """
        # INITIALISE
        # self.fig = plt.figure()
        # self.fig = self.figsize(8,8,self.fig)     # Create a default figure size if not set by user
        # self.fig.set_size_inches(5,5)
        self.bmap,self.x,self.y = self.basemap_setup()#ax=self.ax)
        self.mplcommand = mplcommand
        self.data = data

        self.la_n = self.data.shape[-2]
        self.lo_n = self.data.shape[-1]
        
        # if plottype == 'contourf':
            # f1 = self.bmap.contourf(*plotargs,**plotkwargs)
        # elif plottype == 'contour':
            # plotkwargs['colors'] = 'k'
            # f1 = self.bmap.contour(*plotargs,**plotkwargs)
            # scaling_func = M.ticker.FuncFormatter(lambda x, pos:'{0:d}'.format(int(x*multiplier)))
            # plt.clabel(f1, inline=1, fmt=scaling_func, fontsize=9, colors='k')

        plotargs, plotkwargs = self.get_contouring(**kwargs)
        # pdb.set_trace()
        if self.mplcommand == 'contour':
            f1 = self.bmap.contour(*plotargs,**plotkwargs)
        elif self.mplcommand == 'contourf':
            f1 = self.bmap.contourf(*plotargs,**plotkwargs)
        elif self.mplcommand == 'pcolor':
            f1 = self.bmap.pcolor(*plotargs,**plotkwargs)
        elif self.mplcommand == 'pcolormesh':
            f1 = self.bmap.pcolormesh(*plotargs,**plotkwargs)
        elif self.mplcommand == 'scatter':
            f1 = self.bmap.scatter(*plotargs,**plotkwargs)
        else:
            print("Specify plot type.")
            raise Exception

        # LABELS, TITLES etc
        """
        Change these to hasattr!
        """
        #if self.C.plot_titles:
        if not no_title:
            title = utils.string_from_time('title',pt,tupleformat=0)
            plt.title(title)
        plot_colorbar = 1
        if plot_colorbar:
            self.fig.colorbar(f1,orientation='horizontal')
        # plt.show(self.fig)
        # div0 = make_axes_locatable(self.ax)
        # cax0 = div0.append_axes("bottom", size="20%", pad=0.05)
        # cb0 = self.fig.colorbar(f1, cax=cax0)


        # SAVE FIGURE
        datestr = utils.string_from_time('output',pt,tupleformat=0)
        # self.fname = self.create_fname(fpath) # No da variable here
        if save:
            self.save(p2p,fname)
        
        plt.close(self.fig)
        return f1

        # print("Plot saved to {0}.".format(os.path.join(p2p,fname)))

    #def plot2D(self,va,**kwargs):
    def plot2D(self,vrbl,t,lv,dom,outpath,bounding=0,smooth=1,
                plottype='contourf',save=1,return_data=0):
        """
        Inputs:

        vrbl        :   variable string
        t           :   date/time in (YYYY,MM,DD,HH,MM,SS) or datenum format
                        If tuple of two dates, it's start time and
                        end time, e.g. for finding max/average.
        lv          :   level
        dom         :   domain
        outpath     :   absolute path to output
        bounding    :   list of four floats (Nlim, Elim, Slim, Wlim):
            Nlim    :   northern limit
            Elim    :   eastern limit
            Slim    :   southern limit
            Wlim    :   western limit
        smooth      :   smoothing. 1 is off. integer greater than one is
                        the degree of smoothing, to be specified.
        save        :   whether to save to file
        """
        # INITIALISE
        self.fig.set_size_inches(8,8)
        self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth)
        self.mplcommand = plottype
        
        # Make sure smooth=0 is corrected to 1
        # They are both essentially 'off'.
        if smooth==0:
            smooth = 1

        # Get indices for time, level, lats, lons
        
        if isinstance(t,collections.Sequence) and len(t)!=6:
            # List of two dates, start and end
            # pdb.set_trace()

            it_idx = self.W.get_time_idx(t[0])
            ft_idx = self.W.get_time_idx(t[1])
            assert ft_idx > it_idx
            tidx = slice(it_idx,ft_idx,None)
            title = "range"
            datestr = "range"
        else:
            tidx = self.W.get_time_idx(t)
            title = utils.string_from_time('title',t)
            datestr = utils.string_from_time('output',t)

        
        # Until pressure coordinates are fixed TODO
        lvidx = 0
        latidx, lonidx = self.get_limited_domain(bounding,smooth=smooth)
        
        # if vc == 'surface':
        #     lv_idx = 0
        # elif lv == 'all':
        #     lv_idx = 'all'
        # else:
        #     print("Need to sort other levels")
        #     raise Exception

        # FETCH DATA
        ncidx = {'t': tidx, 'lv': lvidx, 'la': latidx, 'lo': lonidx}
        self.data = self.W.get(vrbl,ncidx)#,**vardict)

        self.la_n = self.data.shape[-2]
        self.lo_n = self.data.shape[-1]

        # COLORBAR, CONTOURING
        plotargs, plotkwargs = self.get_contouring(vrbl,lv)

        # S = Scales(vrbl,lv)

        # multiplier = S.get_multiplier(vrbl,lv)

        # if S.cm:
            # plotargs = (self.x,self.y,data.reshape((la_n,lo_n)),S.clvs)
            # cmap = S.cm
        # elif isinstance(S.clvs,N.ndarray):
            # if plottype == 'contourf':
                # plotargs = (self.x,self.y,data.reshape((la_n,lo_n)),S.clvs)
                # cmap = plt.cm.jet
            # else:
                # plotargs = (self.x,self.y,data.reshape((la_n,lo_n)),S.clvs)
        # else:
            # plotargs = (self.x,self.y,data.reshape((la_n,lo_n)))
            # cmap = plt.cm.jet
        # pdb.set_trace()

        if self.mplcommand == 'contourf':
            # f1 = self.bmap.contourf(*plotargs,cmap=cmap)
            f1 = self.bmap.contourf(*plotargs,**plotkwargs)
        elif self.mplcommand == 'contour':
            plotkwargs['colors'] = 'k'
            f1 = self.bmap.contour(*plotargs,**kwargs)
            scaling_func = M.ticker.FuncFormatter(lambda x, pos:'{0:d}'.format(int(x*multiplier)))
            plt.clabel(f1, inline=1, fmt=scaling_func, fontsize=9, colors='k')

        # LABELS, TITLES etc
        if self.C.plot_titles:
            plt.title(title)
        if self.mplcommand == 'contourf' and self.C.colorbar:
            plt.colorbar(f1,orientation='horizontal')
        
        # SAVE FIGURE
        # pdb.set_trace()
        lv_na = utils.get_level_naming(vrbl,lv)
        naming = [vrbl,lv_na,datestr]
        if dom:
            naming.append(dom)
        self.fname = self.create_fname(*naming)
        if save:
            self.save(outpath,self.fname)
        plt.close()
        if isinstance(self.data,N.ndarray):
            return self.data.reshape((self.la_n,self.lo_n))


    def plot_streamlines(self,lv,pt,outpath,da=0):
        m,x,y = self.basemap_setup()

        time_idx = self.W.get_time_idx(pt)

        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}

        if lv == 2000:
            u = self.W.get('U10',slices)[0,:,:]
            v = self.W.get('V10',slices)[0,:,:]
        else:
            u = self.W.get('U',slices)[0,0,:,:]
            v = self.W.get('V',slices)[0,0,:,:]
        # pdb.set_trace()
        
        #div = N.sum(N.dstack((N.gradient(u)[0],N.gradient(v)[1])),axis=2)*10**4
        #vort = (N.gradient(v)[0] - N.gradient(u)[1])*10**4
        #pdb.set_trace()
        lv_na = utils.get_level_naming('wind',lv=2000)

        m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],u,v,
                        density=2.5,linewidth=0.75,color='k')
        #div_Cs = N.arange(-30,31,1)
        #divp = m.contourf(x,y,vort,alpha=0.6)
        #divp = m.contour(x,y,vort)

        #plt.colorbar(divp,orientation='horizontal')
        if self.C.plot_titles:
            title = utils.string_from_time('title',pt)
            m.title(title)
        datestr = utils.string_from_time('output',pt)
        na = ('streamlines',lv_na,datestr)
        fname = self.create_fname(*na)
        self.save(outpath,fname)

    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)
Ejemplo n.º 13
0
class BirdsEye(Figure):
    def __init__(self, wrfout, ax=0, fig=0):
        super(BirdsEye, self).__init__(wrfout, ax=ax, fig=fig)

    def get_plot_arguments(self, cmap=False, clvs=False):
        """
        Returns colourmap and contouring levels

        Options keyword arguments:
        clvs    :   manually override contour levels
        """
        data = self.data.reshape((self.la_n, self.lo_n))

        # List of args and dictionary of kwargs
        plotargs = [self.x, self.y, data]
        plotkwargs = {}

        # if self.mplcommand == 'contour':
        # multiplier = S.get_multiplier(vrbl,lv)
        if clvs is not False:
            plotkwargs['levels'] = clvs

        if cmap is not False:
            # cmap = eval('M.cm.{0}'.format(cmap))
            plotkwargs['cmap'] = cmap
        # import pdb; pdb.set_trace()
        return plotargs, plotkwargs

    # Old plot_data
    def plot2D(self,
               data,
               fname,
               outdir,
               plottype='contourf',
               save=1,
               smooth=1,
               lats=False,
               lons=False,
               clvs=False,
               cmap=False,
               title=False,
               colorbar=True):
        """
        Generic method that plots any matrix of data on a map

        Inputs:
        data        :   2D matrix of data
        outdir      :   path to plots
        outf        :   filename for output (with or without .png)

        Optional:
        plottype    :   matplotlib function for plotting
        smooth      :   Gaussian smooth by this many grid spaces
        clvs        :   scale for contours
        title       :   title on plot
        save        :   whether to save to file
        """
        # INITIALISE
        self.data = data
        self.bmap, self.x, self.y = self.basemap_setup(
            smooth=smooth,
            lats=lats,
            lons=lons,
        )  #ax=self.ax)
        self.la_n = self.data.shape[-2]
        self.lo_n = self.data.shape[-1]

        plotargs, plotkwargs = self.get_plot_arguments(clvs=clvs, cmap=cmap)

        # import pdb; pdb.set_trace()
        if plottype == 'contour':
            f1 = self.bmap.contour(*plotargs, **plotkwargs)
        elif plottype == 'contourf':
            f1 = self.bmap.contourf(*plotargs, **plotkwargs)
        elif plottype == 'pcolor':
            f1 = self.bmap.pcolor(*plotargs, **plotkwargs)
        elif plottype == 'pcolormesh':
            f1 = self.bmap.pcolormesh(*plotargs, **plotkwargs)
        elif plottype == 'scatter':
            f1 = self.bmap.scatter(*plotargs, **plotkwargs)
        else:
            print("Specify correct plot type.")
            raise Exception

        if isinstance(title, basestring):
            plt.title(title)
        if colorbar:
            self.fig.colorbar(f1, orientation='vertical')
        if save:
            self.save(outdir, fname)

        plt.close(self.fig)

    def plot_streamlines(self,
                         U,
                         V,
                         outdir,
                         fname,
                         lats=False,
                         lons=False,
                         smooth=1,
                         title=False,
                         lw_speed=False):
        """
        Plot streamlines.

        U       :   U-component of wind (nx x ny)
        V       :   V-component of wind (same dimensions)

        lw_speed    :   linewidth is proportional to wind speed
        """
        m, x, y = self.basemap_setup()

        if lw_speed:
            wind = N.sqrt(U**2 + V**2)
            lw = 5 * wind / wind.max()
        else:
            lw = 1

        if smooth > 1:
            U = stats.gauss_smooth(U, smooth)
            V = stats.gauss_smooth(V, smooth)

        m.streamplot(x[self.W.x_dim / 2, :],
                     y[:, self.W.y_dim / 2],
                     U,
                     V,
                     density=1.8,
                     linewidth=lw,
                     color='k',
                     arrowsize=3)

        if isinstance(title, basestring):
            self.ax.set_title(title)

        self.save(outdir, fname)

    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)
Ejemplo n.º 14
0
    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
Ejemplo n.º 15
0
class BirdsEye(Figure):
    def __init__(self,nc=False,ax=0,fig=0):
        super(BirdsEye,self).__init__(nc=nc,ax=ax,fig=fig)

    def get_plot_arguments(self,cmap=False,clvs=False,color=False):
        """
        Returns colourmap and contouring levels

        Options keyword arguments:
        clvs    :   manually override contour levels
        """
        # import pdb; pdb.set_trace()
        # data = self.data.reshape((self.la_n,self.lo_n))
        data = self.data

        # List of args and dictionary of kwargs
        plotargs = [self.x,self.y,data]
        plotkwargs = {}

        # if self.mplcommand == 'contour':
            # multiplier = S.get_multiplier(vrbl,lv)
        if clvs is not False:
            plotkwargs['levels'] = clvs

        if cmap is not False:
            # cmap = eval('M.cm.{0}'.format(cmap))
            plotkwargs['cmap'] = cmap
        return plotargs, plotkwargs

    def axes_of_dilatation(self,xdata,ydata,fname,outdir,
                    lats=False,lons=False,smooth=False, locations=False,
                    x=False,y=False,m=False,
                    Nlim=False,Elim=False,Slim=False,Wlim=False):

        if x is False and y is False and m is False:
            if not Nlim:
                self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
                                                    lons=lons,)#ax=self.ax)
            else:
                self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
                                                    lons=lons,Nlim=Nlim,Elim=Elim,
                                                    Slim=Slim,Wlim=Wlim)

        else:
            self.bmap = m
            self.x = x
            self.y = y

        # self.la_n = self.data.shape[-2]
        # self.lo_n = self.data.shape[-1]
        # import pdb; pdb.set_trace()
        # b = 10
        # self.bmap.quiver(self.x[::b],self.y[::b],xdata[::b,::b],ydata[::b,::b],headwidth=0, units='xy',scale=10)
        self.bmap.streamplot(self.x,self.y,xdata,ydata)
        # m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],U,V,
                        # density=1.8,linewidth=lw,color='k',arrowsize=3)
    
        self.save(outdir,fname)
        plt.close(self.fig)

    # Old plot_data
    def plot2D(self,data,fname,outdir,plottype='contourf',
                    save=True,smooth=1,lats=False,lons=False,
                    clvs=False,cmap=False,title=False,cb=True,
                    locations=False,m=False,x=False,y=False,
                    Nlim=False,Elim=False,Slim=False,Wlim=False,
                    color='k',inline=False,lw=False,extend=False,
                    cblabel=False):

        """
        Generic method that plots any matrix of data on a map

        Inputs:
        data        :   2D matrix of data
        outdir      :   path to plots
        outf        :   filename for output (with or without .png)

        Optional:
        plottype    :   matplotlib function for plotting
        smooth      :   Gaussian smooth by this many grid spaces
        clvs        :   scale for contours
        title       :   title on plot
        save        :   whether to save to file

        :param locations:       Locations to plot on the basemap.
                                Format: locations = {'label':(lat,lon),etc}
        :type locations:        dict
        """
        # INITIALISE
        self.data = data
        if x is False and y is False and m is False:
            if not Nlim:
                self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
                                                    lons=lons,)#ax=self.ax)
            else:
                self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
                                                    lons=lons,Nlim=Nlim,Elim=Elim,
                                                    Slim=Slim,Wlim=Wlim)

        else:
            self.bmap = m
            self.x = x
            self.y = y

        # self.la_n = self.data.shape[-2]
        # self.lo_n = self.data.shape[-1]

        plotargs, plotkwargs = self.get_plot_arguments(clvs=clvs,cmap=cmap,color=color)

        # import pdb; pdb.set_trace()
        if plottype == 'contour':
            plotkwargs['colors'] = color
            plotkwargs['inline'] = inline
            if lw:
                plotkwargs['lw'] = lw
            f1 = self.bmap.contour(*plotargs,**plotkwargs)
            if inline:
                plt.clabel(f1,inline=True,fmt='%d',color='black',fontsize=9)
        elif plottype == 'contourf':
            if isinstance(extend,str):
                plotkwargs['extend'] = extend
            f1 = self.bmap.contourf(*plotargs,**plotkwargs)
        elif plottype == 'pcolor':
            f1 = self.bmap.pcolor(*plotargs,**plotkwargs)
        elif plottype == 'pcolormesh':
            f1 = self.bmap.pcolormesh(*plotargs,**plotkwargs)
        elif plottype == 'scatter':
            f1 = self.bmap.scatter(*plotargs,**plotkwargs)
        # elif plottype == 'quiver':
            # f1 = self.bmap.quiver(*plotargs,**plotkwargs) 
        else:
            print("Specify correct plot type.")
            raise Exception

        if isinstance(locations,dict):
            for k,v in locations.iteritems():
                if isinstance(v,tuple) and len(v) == 2:
                    xpt, ypt = self.bmap(v[1],v[0])
                    # bbox_style = {'boxstyle':'square','fc':'white','alpha':0.5}
                    self.bmap.plot(xpt,ypt,'ko',markersize=3,zorder=100)
                    self.ax.text(xpt,ypt,k,ha='left',fontsize=7)
                    # self.ax.text(xpt-15000,ypt,k,bbox=bbox_style,ha='left',fontsize=7)
                else:
                    print("Not a valid location argument.")
                    raise Exception

        if isinstance(title,basestring):
            plt.title(title)
        if cb != False:
            if cb==True:
                cb1 = plt.colorbar(f1,orientation='vertical',ax=self.ax)
            elif cb=='horizontal':
                cb1 = plt.colorbar(f1,orientation='horizontal',ax=self.ax)
            elif cb == 'only':
                save = False
                self.fig,self.ax = plt.subplots(figsize=(4,0.8))
                cb1 = plt.colorbar(f1,cax=self.ax,orientation='horizontal')
                if isinstance(cblabel,str):
                    cb1.set_label(cblabel)
                self.save(outdir,fname+'_cb')
            else:
                cb1 = plt.colorbar(f1,orientation='vertical',cax=cb)
        if cb and isinstance(cblabel,str):
            cb1.set_label(cblabel)
        if save:
            self.save(outdir,fname)
            plt.close(self.fig)
        else:
            return f1

    def plot_streamlines(self,U,V,outdir,fname,lats=False,lons=False,smooth=1,
                            title=False,lw_speed=False,density=1.8):
        """
        Plot streamlines.

        U       :   U-component of wind (nx x ny)
        V       :   V-component of wind (same dimensions)

        lw_speed    :   linewidth is proportional to wind speed
        """
        m,x,y = self.basemap_setup(lats=lats,lons=lons)

        if lw_speed:
            wind = N.sqrt(U**2 + V**2)
            lw = 5*wind/wind.max()
        else:
            lw = 1

        if smooth>1:
            U = stats.gauss_smooth(U,smooth)
            V = stats.gauss_smooth(V,smooth)

        # m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],U,V,
                        # density=density,linewidth=lw,color='k',arrowsize=3)
        m.streamplot(x,y,U,V,
                        density=density,linewidth=lw,color='k',arrowsize=3)

        if isinstance(title,basestring):
            self.ax.set_title(title)

        self.save(outdir,fname)

    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 make_subplot_label(ax,label):
        ax.text(0.05,0.85,label,transform=ax.transAxes,
            bbox={'facecolor':'white'},fontsize=15,zorder=1000)
        return
Ejemplo n.º 16
0
    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)
Ejemplo n.º 17
0
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
Ejemplo n.º 18
0
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
Ejemplo n.º 19
0
            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)
Ejemplo n.º 20
0
class BirdsEye(Figure):
    def __init__(self,wrfout,ax=0,fig=0):
        super(BirdsEye,self).__init__(wrfout,ax=ax,fig=fig)

    def get_plot_arguments(self,cmap=False,clvs=False):
        """
        Returns colourmap and contouring levels

        Options keyword arguments:
        clvs    :   manually override contour levels
        """
        data = self.data.reshape((self.la_n,self.lo_n))

        # List of args and dictionary of kwargs
        plotargs = [self.x,self.y,data]
        plotkwargs = {}

        # if self.mplcommand == 'contour':
            # multiplier = S.get_multiplier(vrbl,lv)
        if clvs is not False:
            plotkwargs['levels'] = clvs

        if cmap is not False:
            # cmap = eval('M.cm.{0}'.format(cmap))
            plotkwargs['cmap'] = cmap
        # import pdb; pdb.set_trace()
        return plotargs, plotkwargs

    # Old plot_data
    def plot2D(self,data,fname,outdir,plottype='contourf',
                    save=1,smooth=1,lats=False,lons=False,
                    clvs=False,cmap=False,title=False,colorbar=True):

        """
        Generic method that plots any matrix of data on a map

        Inputs:
        data        :   2D matrix of data
        outdir      :   path to plots
        outf        :   filename for output (with or without .png)

        Optional:
        plottype    :   matplotlib function for plotting
        smooth      :   Gaussian smooth by this many grid spaces
        clvs        :   scale for contours
        title       :   title on plot
        save        :   whether to save to file
        """
        # INITIALISE
        self.data = data
        self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
                                                    lons=lons,)#ax=self.ax)
        self.la_n = self.data.shape[-2]
        self.lo_n = self.data.shape[-1]

        plotargs, plotkwargs = self.get_plot_arguments(clvs=clvs,cmap=cmap)

        # import pdb; pdb.set_trace()
        if plottype == 'contour':
            f1 = self.bmap.contour(*plotargs,**plotkwargs)
        elif plottype == 'contourf':
            f1 = self.bmap.contourf(*plotargs,**plotkwargs)
        elif plottype == 'pcolor':
            f1 = self.bmap.pcolor(*plotargs,**plotkwargs)
        elif plottype == 'pcolormesh':
            f1 = self.bmap.pcolormesh(*plotargs,**plotkwargs)
        elif plottype == 'scatter':
            f1 = self.bmap.scatter(*plotargs,**plotkwargs)
        else:
            print("Specify correct plot type.")
            raise Exception

        if isinstance(title,basestring):
            plt.title(title)
        if colorbar:
            self.fig.colorbar(f1,orientation='vertical')
        if save:
            self.save(outdir,fname)

        plt.close(self.fig)

    def plot_streamlines(self,U,V,outdir,fname,lats=False,lons=False,smooth=1,
                            title=False,lw_speed=False):
        """
        Plot streamlines.

        U       :   U-component of wind (nx x ny)
        V       :   V-component of wind (same dimensions)

        lw_speed    :   linewidth is proportional to wind speed
        """
        m,x,y = self.basemap_setup()

        if lw_speed:
            wind = N.sqrt(U**2 + V**2)
            lw = 5*wind/wind.max()
        else:
            lw = 1

        if smooth>1:
            U = stats.gauss_smooth(U,smooth)
            V = stats.gauss_smooth(V,smooth)

        m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],U,V,
                        density=1.8,linewidth=lw,color='k',arrowsize=3)

        if isinstance(title,basestring):
            self.ax.set_title(title)

        self.save(outdir,fname)

    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)