source: trunk/UTIL/PYTHON/planetoplot_v0/winds.py @ 1062

Last change on this file since 1062 was 943, checked in by aslmd, 12 years ago

clean and organized UTIL/PYTHON folder

  • Property svn:executable set to *
File size: 18.1 KB
RevLine 
[180]1#!/usr/bin/env python
2
[202]3### A. Spiga -- LMD -- 30/06/2011 to 10/07/2011
[184]4### Thanks to A. Colaitis for the parser trick
[180]5
[184]6
7####################################
8####################################
9### The main program to plot vectors
[182]10def winds (namefile,\
11           nvert,\
[184]12           proj=None,\
[182]13           back=None,\
14           target=None,
[184]15           stride=3,\
[225]16           numplot=2,\
[187]17           var=None,\
[241]18           colorb="def",\
[195]19           winds=True,\
[199]20           addchar=None,\
21           interv=[0,1],\
22           vmin=None,\
[201]23           vmax=None,\
[204]24           tile=False,\
[233]25           zoom=None,\
26           display=True,\
[238]27           itstep=None,\
[240]28           hole=False,\
29           save="gui",\
[241]30           anomaly=False,\
[244]31           var2=None,\
[254]32           ndiv=10,\
[287]33           first=1,\
34           mult=1.,\
35           zetitle="fill"):
[184]36
[192]37    ####################################################################################################################
38    ### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
39
[184]40    #################################
41    ### Load librairies and functions
[180]42    from netCDF4 import Dataset
[240]43    from myplot import getcoord2d,define_proj,makeplotres,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy,\
[233]44                       fmtvar,definecolorvec,defcolorb,getprefix,putpoints,calculate_bounds,errormess,definesubplot,\
[251]45                       zoomset,getcoorddef,getwinddef,whatkindfile,reducefield,bounds,getstralt,getfield,smooth,nolow,\
46                       getname,localtime,polarinterv
[199]47    from mymath import deg,max,min,mean
[240]48    from matplotlib.pyplot import contour,contourf, subplot, figure, rcParams, savefig, colorbar, pcolor, show
[192]49    from matplotlib.cm import get_cmap
[182]50    import numpy as np
[202]51    from numpy.core.defchararray import find
[184]52
53    ######################
54    ### Load NETCDF object
[240]55    nc  = Dataset(namefile) 
[233]56   
57    ##################################
58    ### Initial checks and definitions
59    typefile = whatkindfile(nc)                                  ## TYPEFILE
60    if var not in nc.variables: var = False                      ## VAR
[238]61    if winds:                                                    ## WINDS
62        [uchar,vchar,metwind] = getwinddef(nc)             
63        if uchar == 'not found': winds = False
[251]64    if not var and not winds: errormess("please set at least winds or var",printvar=nc.variables)
[233]65    [lon2d,lat2d] = getcoorddef(nc)                              ## COORDINATES, could be moved below
66    if proj == None:   proj = getproj(nc)                        ## PROJECTION
[184]67
[233]68    ##########################
69    ### Define plot boundaries
[251]70    ### todo: possible areas in latinterv in argument (ex: "Far_South_Pole")
71    if proj in ["npstere","spstere"]: [wlon,wlat] = polarinterv(lon2d,lat2d)
[207]72    elif proj in ["lcc","laea"]:      [wlon,wlat] = wrfinterv(lon2d,lat2d)
73    else:                             [wlon,wlat] = simplinterv(lon2d,lat2d)
[233]74    if zoom:                          [wlon,wlat] = zoomset(wlon,wlat,zoom) 
[184]75
[186]76    #########################################
77    ### Name for title and graphics save file
[251]78    basename = getname(var=var,winds=winds,anomaly=anomaly)
[233]79    basename = basename + getstralt(nc,nvert)  ## can be moved elsewhere for a more generic routine
[186]80
[184]81    ##################################
82    ### Open a figure and set subplots
[180]83    fig = figure()
[345]84    subv,subh = definesubplot( numplot, fig ) 
[233]85 
[184]86    #################################
87    ### Time loop for plotting device
[186]88    found_lct = False
[287]89    nplot = 1
[254]90    itime = first
[233]91    error = False
[241]92    if itstep is None and numplot > 0: itstep = int(24./numplot)
[243]93    elif numplot <= 0:                 itstep = 1
[233]94    while error is False: 
[184]95
[197]96       ### Which local time ?
[251]97       ltst = localtime ( interv[0]+itime*interv[1], 0.5*(wlon[0]+wlon[1]) )
[197]98
[184]99       ### General plot settings
[233]100       #print itime, int(ltst), numplot, nplot
101       if numplot >= 1: 
102           if nplot > numplot: break
103           if numplot > 1:     
[345]104               if typefile not in ['geo']:  subplot(subv,subh,nplot)
[186]105           found_lct = True
[233]106       ### If only one local time is requested (numplot < 0)
[186]107       elif numplot <= 0: 
[241]108           if int(ltst) + numplot != 0: 
[233]109                 itime += 1 
110                 if found_lct is True: break     ## because it means LT was found at previous iteration
111                 else:                 continue  ## continue to iterate to find the correct LT
112           else: 
113                 found_lct = True
[184]114
115       ### Map projection
[180]116       m = define_proj(proj,wlon,wlat,back=back)
117       x, y = m(lon2d, lat2d)
[184]118
119       #### Contour plot
[241]120       if var2:
121           what_I_contour, error = reducefield( getfield(nc,var2), d4=itime, d3=nvert )
122           if not error:
123               if typefile in ['mesoapi','meso']:    what_I_contour = dumpbdy(what_I_contour,6)
124               zevmin, zevmax = calculate_bounds(what_I_contour)
125               zelevels = np.linspace(zevmin,zevmax,num=20)
126               if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
127               contour( x, y, what_I_contour, zelevels, colors='k', linewidths = 0.33 ) #colors='w' )# , alpha=0.5)
[287]128           else:
129               errormess("There is an error in reducing field !")
[241]130
131       #### Shaded plot
[184]132       if var:
[233]133           what_I_plot, error = reducefield( getfield(nc,var), d4=itime, d3=nvert )
[287]134           what_I_plot = what_I_plot*mult
[233]135           if not error: 
[240]136               fvar = var
137               ###
138               if anomaly:
[244]139                   what_I_plot = 100. * ((what_I_plot / smooth(what_I_plot,12)) - 1.)
[240]140                   fvar = 'anomaly'
[287]141               #if mult != 1:     
142               #    fvar = str(mult) + "*" + var
[240]143               ###
144               if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
[233]145               zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
[241]146               if colorb in ["def","nobar"]:   palette = get_cmap(name=defcolorb(fvar))
147               else:                           palette = get_cmap(name=colorb)
[233]148               if not tile:
[238]149                   if not hole: what_I_plot = bounds(what_I_plot,zevmin,zevmax)
[287]150                   #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
151                   zelevels = np.linspace(zevmin,zevmax)
[233]152                   contourf( x, y, what_I_plot, zelevels, cmap = palette )
[241]153               else:
154                   if hole:  what_I_plot = nolow(what_I_plot) 
155                   pcolor( x, y, what_I_plot, cmap = palette, \
156                           vmin=zevmin, vmax=zevmax )
[243]157               if colorb != 'nobar' and var != 'HGT':             
[251]158                         colorbar(fraction=0.05,pad=0.03,format=fmtvar(fvar),\
[233]159                                           ticks=np.linspace(zevmin,zevmax,ndiv+1),\
[240]160                                           extend='neither',spacing='proportional')
161                                           # both min max neither
[287]162           else:
163               errormess("There is an error in reducing field !")
164 
[184]165       ### Vector plot
[191]166       if winds:
[233]167           vecx, error = reducefield( getfield(nc,uchar), d4=itime, d3=nvert )
168           vecy, error = reducefield( getfield(nc,vchar), d4=itime, d3=nvert )
169           if not error:
170               if typefile in ['mesoapi','meso']:   
[240]171                   [vecx,vecy] = [dumpbdy(vecx,6,stag=uchar), dumpbdy(vecy,6,stag=vchar)]
[233]172                   key = True
173               elif typefile in ['gcm']:           
174                   key = False
175               if metwind:  [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
176               if var == False:       colorvec = definecolorvec(back)
177               else:                  colorvec = definecolorvec(colorb)
178               vectorfield(vecx, vecy,\
179                          x, y, stride=stride, csmooth=2,\
[247]180                          #scale=15., factor=300., color=colorvec, key=key)
181                          scale=20., factor=250., color=colorvec, key=key)
[233]182                                            #200.         ## or csmooth=stride
183               
[184]184       ### Next subplot
[202]185       plottitle = basename
[225]186       if typefile in ['mesoapi','meso']:
[233]187            if addchar:  plottitle = plottitle + addchar + "_LT"+str(ltst)
188            else:        plottitle = plottitle + "_LT"+str(ltst)
[287]189       if mult != 1:           plottitle = str(mult) + "*" + plottitle
190       if zetitle != "fill":   plottitle = zetitle
[195]191       ptitle( plottitle )
[233]192       itime += itstep
193       nplot += 1
[184]194
195    ##########################################################################
196    ### Save the figure in a file in the data folder or an user-defined folder
[202]197    if typefile in ['meso','mesoapi']:   prefix = getprefix(nc)   
[317]198    elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
[202]199    else:                                prefix = ''
200    ###
201    zeplot = prefix + basename
202    if addchar:         zeplot = zeplot + addchar
203    if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
204    ###
205    if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
206    else:               zeplot = target + "/" + zeplot 
207    ###
[240]208    if found_lct:     
209        pad_inches_value = 0.35
210        if save == 'png': 
[295]211            if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
[240]212            makeplotres(zeplot,res=200.,pad_inches_value=pad_inches_value,disp=False)
213        elif save in ['eps','svg','pdf']:
214            makeplotres(zeplot,         pad_inches_value=pad_inches_value,disp=False,ext=save)
215        elif save == 'gui':
216            show()
217        else: 
218            print "save mode not supported. using gui instead."
219            show()
[287]220    else:   print "Local time not found"
[202]221
222    ###############
223    ### Now the end
[193]224    return zeplot
[180]225
[233]226##############################
227### A specific stuff for below
228def adjust_length (tab, zelen):
229    from numpy import ones
230    if tab is None:
231        outtab = ones(zelen) * -999999
232    else:
233        if zelen != len(tab):
234            print "not enough or too much values... setting same values all variables"
235            outtab = ones(zelen) * tab[0]
236        else:
237            outtab = tab
238    return outtab
239
[184]240###########################################################################################
241###########################################################################################
242### What is below relate to running the file as a command line executable (very convenient)
[180]243if __name__ == "__main__":
244    import sys
[184]245    from optparse import OptionParser    ### to be replaced by argparse
[186]246    from api_wrapper import api_onelevel
[195]247    from netCDF4 import Dataset
248    from myplot import getlschar
[240]249    from os import system
[193]250
251    #############################
252    ### Get options and variables
[182]253    parser = OptionParser()
[244]254    parser.add_option('-f', '--file',   action='append',dest='namefile', type="string",  default=None,  help='[NEEDED] name of WRF file (append)')
255    parser.add_option('-l', '--level',  action='store',dest='nvert',     type="float",   default=0,     help='level (def=0)(-i 2: p,mbar)(-i 3,4: z,km)')
256    parser.add_option('-p', '--proj',   action='store',dest='proj',      type="string",  default=None,  help='projection')
257    parser.add_option('-b', '--back',   action='store',dest='back',      type="string",  default=None,  help='background image (def: None)')
258    parser.add_option('-t', '--target', action='store',dest='target',    type="string",  default=None,  help='destination folder')
259    parser.add_option('-s', '--stride', action='store',dest='stride',    type="int",     default=3,     help='stride vectors (def=3)')
260    parser.add_option('-v', '--var',    action='append',dest='var',      type="string",  default=None,  help='variable color-shaded (append)')
261    parser.add_option('-n', '--num',    action='store',dest='numplot',   type="int",     default=2,     help='plot number (def=2)(<0: plot LT -*numplot*)')
262    parser.add_option('-i', '--interp', action='store',dest='interp',    type="int",     default=None,  help='interpolation (2: p, 3: z-amr, 4:z-als)')
263    parser.add_option('-c', '--color',  action='store',dest='colorb',    type="string",  default="def", help='change colormap (nobar: no colorbar)')
264    parser.add_option('-x', '--no-vect',action='store_false',dest='winds',               default=True,  help='no wind vectors')
265    parser.add_option('-m', '--min',    action='append',dest='vmin',     type="float",   default=None,  help='bounding minimum value (append)')   
266    parser.add_option('-M', '--max',    action='append',dest='vmax',     type="float",   default=None,  help='bounding maximum value (append)') 
267    parser.add_option('-T', '--tiled',  action='store_true',dest='tile',                 default=False, help='draw a tiled plot (no blank zone)')
268    parser.add_option('-z', '--zoom',   action='store',dest='zoom',      type="float",   default=None,  help='zoom factor in %')
269    parser.add_option('-N', '--no-api', action='store_true',dest='nocall',               default=False, help='do not recreate api file')
270    parser.add_option('-d', '--display',action='store_false',dest='display',             default=True,  help='do not pop up created images')
[254]271    parser.add_option('-e', '--itstep', action='store',dest='itstep',    type="int",     default=None,  help='stride time (def=4)')
[244]272    parser.add_option('-H', '--hole',   action='store_true',dest='hole',                 default=False, help='holes above max and below min')
273    parser.add_option('-S', '--save',   action='store',dest='save',      type="string",  default="gui", help='save mode (png,eps,svg,pdf or gui)(def=gui)')
274    parser.add_option('-a', '--anomaly',action='store_true',dest='anomaly',              default=False, help='compute and plot relative anomaly in %')
[287]275    parser.add_option('-w', '--with',   action='store',dest='var2',      type="string",  default=None,  help='variable contoured')
276    parser.add_option('--div',          action='store',dest='ndiv',      type="int",     default=10,    help='number of divisions in colorbar (def: 10)')
277    parser.add_option('-F', '--first',  action='store',dest='first',     type="int",     default=1,     help='first subscript to plot (def: 1)')
278    parser.add_option('--mult',         action='store',dest='mult',      type="float",   default=1.,    help='a multiplicative factor to plotted field')
279    parser.add_option('--title',        action='store',dest='zetitle',   type="string",  default="fill",help='customize the whole title')
[201]280    #parser.add_option('-V', action='store', dest='comb',        type="float",   default=None,  help='a defined combination of variables')
[182]281    (opt,args) = parser.parse_args()
282    if opt.namefile is None: 
283        print "I want to eat one file at least ! Use winds.py -f name_of_my_file. Or type winds.py -h"
284        exit()
[241]285    if opt.var is None and opt.anomaly is True:
286        print "Cannot ask to compute anomaly if no variable is set"
287        exit()   
[182]288    print "Options:", opt
[233]289
290    listvar = '' 
[225]291    if opt.var is None:
[233]292        zerange = [-999999]
[225]293    else:
294        zelen = len(opt.var)
[233]295        zerange = range(zelen)
296        #if zelen == 1: listvar = opt.var[0] + ','
297        #else         :
298        for jjj in zerange: listvar += opt.var[jjj] + ','
[225]299        listvar = listvar[0:len(listvar)-1]
[233]300        vmintab = adjust_length (opt.vmin, zelen) 
301        vmaxtab = adjust_length (opt.vmax, zelen)
[193]302
[225]303    for i in range(len(opt.namefile)):
[184]304
[225]305        zefile = opt.namefile[i]
306        print zefile   
307        zelevel = opt.nvert   
308        stralt = None
309        [lschar,zehour,zehourin] = getlschar ( zefile )  ## getlschar from wrfout (or simply return "" if another file)
310   
311        #####################################################
312        ### Call Fortran routines for vertical interpolations
313        if opt.interp is not None:
314            if opt.nvert is 0 and opt.interp is 4:  zelevel = 0.010
315            ### winds or no winds
316            if opt.winds            :  zefields = 'uvmet'
317            else                    :  zefields = ''
318            ### var or no var
[233]319            #if opt.var is None      :  pass
320            if zefields == ''       :  zefields = listvar
[225]321            else                    :  zefields = zefields + "," + listvar
[241]322            if opt.var2 is not None : zefields = zefields + "," + opt.var2 
[225]323            print zefields
324            zefile = api_onelevel (  path_to_input   = '', \
325                                     input_name      = zefile, \
326                                     fields          = zefields, \
327                                     interp_method   = opt.interp, \
328                                     onelevel        = zelevel, \
329                                     nocall          = opt.nocall )
330            print zefile
331            zelevel = 0 ## so that zelevel could play again the role of nvert
[193]332
[233]333        if opt.var is None: zerange = [-999999]
334        else:               zerange = range(zelen) 
335        for jjj in zerange:
336            if jjj == -999999: 
337                argvar  = None
338                argvmin = None
339                argvmax = None
340            else:
341                argvar = opt.var[jjj]
342                if vmintab[jjj] != -999999:  argvmin = vmintab[jjj]
343                else:                        argvmin = None
344                if vmaxtab[jjj] != -999999:  argvmax = vmaxtab[jjj] 
345                else:                        argvmax = None
[225]346            #############
347            ### Main call
348            name = winds (zefile,int(zelevel),\
[233]349                proj=opt.proj,back=opt.back,target=opt.target,stride=opt.stride,var=argvar,\
350                numplot=opt.numplot,colorb=opt.colorb,winds=opt.winds,\
351                addchar=lschar,interv=[zehour,zehourin],vmin=argvmin,vmax=argvmax,\
352                tile=opt.tile,zoom=opt.zoom,display=opt.display,\
[240]353                itstep=opt.itstep,hole=opt.hole,save=opt.save,\
[287]354                anomaly=opt.anomaly,var2=opt.var2,ndiv=opt.ndiv,first=opt.first,\
355                mult=opt.mult,zetitle=opt.zetitle)
[225]356            print 'Done: '+name
[240]357            system("rm -f to_be_erased")
358 
[225]359        #########################################################
360        ### Generate a .sh file with the used command saved in it
361        command = ""
362        for arg in sys.argv: command = command + arg + ' '
363        f = open(name+'.sh', 'w')
364        f.write(command)
Note: See TracBrowser for help on using the repository browser.