source: trunk/MESOSCALE_DEV/PLOT/PYTHON/scripts/winds.py @ 245

Last change on this file since 245 was 244, checked in by aslmd, 14 years ago

MESOSCALE: graphics. added full words options to winds.py, added ndiv keyword.

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