source: trunk/UTIL/PYTHON/planetoplot.py @ 427

Last change on this file since 427 was 427, checked in by acolaitis, 13 years ago

PYTHON. Corrected bugs related to contour plots with overlines, and mymath.mean() with no Axis specified. Now working on LES format for pp.py... :)

  • Property svn:executable set to *
File size: 23.2 KB
Line 
1#######################
2##### PLANETOPLOT #####
3#######################
4
5### A. Spiga     -- LMD -- 06~09/2011 -- General building and mapping capabilities
6### T. Navarro   -- LMD -- 10~11/2011 -- Improved use for GCM and added sections + 1Dplot capabilities
7### A. Colaitis  -- LMD --    11/2011 -- Mostly minor improvements and inter-plot operation capabilities + zrecast interpolation for gcm
8### A. Spiga     -- LMD --    11/2011 -- Extended multivar subplot capabilities + cosmetic changes + general cleaning and tests
9
10def planetoplot (namefiles,\
11           level=0,\
12           vertmode=0,\
13           proj=None,\
14           back=None,\
15           target=None,
16           stride=3,\
17           var=None,\
18           colorb="def",\
19           winds=False,\
20           addchar=None,\
21           interv=[0,1],\
22           vmin=None,\
23           vmax=None,\
24           tile=False,\
25           zoom=None,\
26           display=True,\
27           hole=False,\
28           save="gui",\
29           anomaly=False,\
30           var2=None,\
31           ndiv=10,\
32           mult=1.,\
33           zetitle="fill",\
34           slon=None,\
35           slat=None,\
36           svert=None,\
37           stime=None,\
38           outputname=None,\
39           resolution=200,\
40           ope=None,\
41           fileref=None,\
42           minop=0.,\
43           maxop=0.,\
44           titleref="fill",\
45           invert_y=False,\
46           xaxis=[None,None],\
47           yaxis=[None,None],\
48           ylog=False,\
49           yintegral=False,\
50           blat=None,\
51           tsat=False,\
52           flagnolow=False):
53
54
55    ####################################################################################################################
56    ### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
57
58    #################################
59    ### Load librairies and functions
60    from netCDF4 import Dataset
61    from myplot import getcoord2d,define_proj,makeplotres,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy,\
62                       fmtvar,definecolorvec,defcolorb,getprefix,putpoints,calculate_bounds,errormess,definesubplot,\
63                       zoomset,getcoorddef,getwinddef,whatkindfile,reducefield,bounds,getstralt,getfield,smooth,nolow,\
64                       getname,localtime,polarinterv,getsindex,define_axis,determineplot,readslices,bidimfind,getlschar,hole_bounds
65    from mymath import deg,max,min,mean,get_tsat,writeascii
66    import matplotlib as mpl
67    from matplotlib.pyplot import contour,contourf, subplot, figure, rcParams, savefig, colorbar, pcolor, show, plot, clabel, title
68    from matplotlib.cm import get_cmap
69    import numpy as np
70    from numpy.core.defchararray import find
71
72    ################################
73    ### Preliminary stuff
74    ################################
75    print "********************************************"
76    print "********** WELCOME TO PLANETOPLOT **********"
77    print "********************************************"
78    if not isinstance(namefiles, np.ndarray): namefiles = [namefiles]
79    if not isinstance(var, np.ndarray):       var = [var]
80
81    ################################
82    ### Which plot needs to be done?
83    ################################
84    nlon, nlat, nvert, ntime, mapmode, nslices = determineplot(slon, slat, svert, stime)
85    vlon = None ; vlat = None
86    if slon is not None: vlon = slon[0][0]
87    if slat is not None: vlat = slat[0][0]
88    if mapmode == 0:       winds=False
89    elif mapmode == 1:     
90        if svert is None:  svert = readslices(str(level)) ; nvert=1
91    zelen = len(namefiles)*len(var)
92    numplot = zelen*nslices
93    print "********** FILES, SLICES, VARS, TOTAL PLOTS: ", len(namefiles), nslices, len(var), numplot
94    print "********** MAPMODE: ", mapmode
95    if ope is not None:
96        if fileref is not None:       zelen = zelen + 2
97        elif "var" in ope:            zelen = zelen + 1
98    all_var  = [[]]*zelen ; all_var2  = [[]]*zelen ; all_title = [[]]*zelen ; all_varname = [[]]*zelen ; all_namefile = [[]]*zelen ; all_time = [[]]*zelen
99 
100    #################################################################################################
101    ### Loop over the files + vars initially separated by commas to be plotted on the same figure ###
102    #################################################################################################
103    k = 0 ; firstfile = True
104    for nnn in range(len(namefiles)):
105     for vvv in range(len(var)): 
106
107      print "********** LOOP..... THIS IS SUBPLOT NUMBER.....",k
108
109      ######################
110      ### Load NETCDF object
111      namefile = namefiles[nnn] ; print "********** THE NAMEFILE IS....", namefile
112      nc  = Dataset(namefile)
113
114      ##################################
115      ### Initial checks and definitions
116      ### ... TYPEFILE
117      typefile = whatkindfile(nc)                                 
118      if firstfile:                 typefile0 = typefile
119      elif typefile != typefile0:   errormess("Not the same kind of files !", [typefile0, typefile])
120      ### ... VAR
121      varname=var[vvv]
122      print "********** THE VAR IS....",varname, var2
123      if varname not in nc.variables: varname = False
124      ### ... WINDS
125      if winds:                                                   
126         [uchar,vchar,metwind] = getwinddef(nc)             
127         if uchar == 'not found': winds = False
128      if not varname and not winds: errormess("please set at least winds or var",printvar=nc.variables)
129      ### ... COORDINATES, could be moved below
130      [lon2d,lat2d] = getcoorddef(nc)
131      ### ... PROJECTION
132      if proj == None:   proj = getproj(nc)                 
133
134##########################################################
135      if typefile == "gcm":
136          lat = nc.variables["latitude"][:] ; lon = nc.variables["longitude"][:]
137          if "Time" in nc.variables:      time = nc.variables["Time"][:]
138          elif "time" in nc.variables:    time = nc.variables["time"][:]
139          else:                           errormess("no time axis found.")
140          vert = nc.variables["altitude"][:]
141      elif typefile in ['meso','mesoapi','geo']:
142          if vlon is not None or vlat is not None:   indices = bidimfind(lon2d,lat2d,vlon,vlat) ; print '********** INDICES: ', indices
143          if slon is not None: slon[0][0] = indices[0] ; slon[0][1] = indices[0]
144          if slat is not None: slat[0][0] = indices[1] ; slat[0][1] = indices[1]
145          if varname in ['PHTOT','W']:    vertdim='BOTTOM-TOP_PATCH_END_STAG'
146          else:                           vertdim='BOTTOM-TOP_PATCH_END_UNSTAG'
147          if varname in ['V']:  latdim='SOUTH-NORTH_PATCH_END_STAG'
148          else:                 latdim='SOUTH-NORTH_PATCH_END_UNSTAG'
149          if varname in ['U']:  londim='WEST-EAST_PATCH_END_STAG'
150          else:                 londim='WEST-EAST_PATCH_END_UNSTAG'
151          lon = np.arange(0,getattr(nc,londim),1) ; lat = np.arange(0,getattr(nc,latdim),1)
152          time = np.arange(0,len(nc.variables["Times"]),1)
153          if typefile in ['geo']:   vert = [0.] ; stime = readslices(str(0))
154          else:
155              if vertmode is None:  vertmode=0
156              if vertmode == 0:     vert = np.arange(0,getattr(nc,vertdim),1)
157              else:                 vert = nc.variables["vert"][:]
158       #if firstfile:
159       #   lat0 = lat
160       #elif len(lat0) != len(lat):
161       #   errormess("Not the same latitude lengths !", [len(lat0), len(lat)])
162       #elif sum((lat == lat0) == False) != 0:
163       #   errormess("Not the same latitudes !", [lat,lat0])
164       ## Faire d'autre checks sur les compatibilites entre fichiers!!
165##########################################################
166
167      if firstfile:
168         ##########################
169         ### Define plot boundaries
170         ### todo: possible areas in latinterv in argument (ex: "Far_South_Pole")
171         if proj in ["npstere","spstere"]: [wlon,wlat] = polarinterv(lon2d,lat2d)
172         elif proj in ["lcc","laea"]:      [wlon,wlat] = wrfinterv(lon2d,lat2d)
173         else:                             [wlon,wlat] = simplinterv(lon2d,lat2d)
174         if zoom:                          [wlon,wlat] = zoomset(wlon,wlat,zoom) 
175
176      all_varname[k] = varname
177      all_namefile[k] = namefile
178      all_time[k] = time
179      if var2: all_var2[k] = getfield(nc,var2)
180      ##### SPECIFIC
181      if varname in ["temp","t","T_nadir_nit","T_nadir_day","temp_day","temp_night"] and tsat:
182          tt=getfield(nc,varname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
183          if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
184          else:                                 tinput=tt
185          all_var[k]=get_tsat(vert,tinput,zlon=lon,zlat=lat,zalt=vert,ztime=time)
186      else:
187      ##### GENERAL STUFF HERE
188          all_var[k] = getfield(nc,varname)
189      print "********** all_var[k].shape", all_var[k].shape
190      k += 1
191      firstfile = False
192      #### End of for namefile in namefiles
193
194    ##################################
195    ### Operation on files
196    if ope is not None:
197        print ope
198        if "var" not in ope:
199             if len(var) > 1: errormess("for this operation... please set only one var !")
200             if ope in ["-","+"]:
201                if fileref is not None:   all_var[k] = getfield(Dataset(fileref),all_varname[k-1]) ; all_varname[k] = all_varname[k-1] ; all_time[k] = all_time[k-1] ; all_namefile[k] = all_namefile[k-1]
202                else:                     errormess("fileref is missing!") 
203                if ope == "-":     all_var[k+1]= all_var[k-1] - all_var[k]
204                elif ope == "+":   all_var[k+1]= all_var[k-1] + all_var[k]
205                all_varname[k+1] = all_varname[k] ; all_time[k+1] = all_time[k] ; all_namefile[k+1] = all_namefile[k] ; numplot = numplot+2
206             elif ope in ["cat"]:
207                tab = all_var[0];k = 1
208                while k != len(namefiles)-1: 
209                    tab = np.append(tab,all_var[k],axis=0) ; k += 1
210                all_time[0] = np.arange(0,len(tab),1) ### AS: time reference is too simplistic, should be better
211                all_var[0] = np.array(tab) ; numplot = 1
212             else: errormess(ope+" : non-implemented operation. Check pp.py --help")
213        else:
214             if len(namefiles) > 1: errormess("for this operation... please set only one file !") 
215             if len(var) > 2:       errormess("not sure this works for more than 2 vars... please check.")
216             if   ope in ["div_var"]: all_var[k] = all_var[k-2] / all_var[k-1] ; insert = '_div_'
217             elif ope in ["mul_var"]: all_var[k] = all_var[k-2] * all_var[k-1] ; insert = '_mul_'
218             elif ope in ["add_var"]: all_var[k] = all_var[k-2] + all_var[k-1] ; insert = '_add_'
219             elif ope in ["sub_var"]: all_var[k] = all_var[k-2] - all_var[k-1] ; insert = '_sub_'
220             else:                    errormess(ope+" : non-implemented operation. Check pp.py --help")
221             numplot = numplot + 1 ; all_time[k] = all_time[k-1] ; all_namefile[k] = all_namefile[k-1]
222             all_varname[k] = all_varname[k-2] + insert + all_varname[k-1] 
223
224    ##################################
225    ### Open a figure and set subplots
226    fig = figure()
227    subv,subh = definesubplot( numplot, fig ) 
228 
229    #################################
230    ### Time loop for plotting device
231    nplot = 1;error = False
232    print "********************************************"
233    while error is False:
234       print "********** NPLOT", nplot
235     
236       ### General plot settings
237       if nplot > numplot: break
238       if numplot > 1:  subplot(subv,subh,nplot)
239
240       ### Map projection                   
241       if mapmode == 1:
242           m = define_proj(proj,wlon,wlat,back=back,blat=blat)
243           x, y = m(lon2d, lat2d)
244
245       ####################################################################
246       ## get all indexes to be taken into account for this subplot and then reduce field
247       ## We plot 1) all lon slices 2) all lat slices 3) all vert slices 4) all time slices and then go to the next slice
248       indexlon  = getsindex(slon,(nplot-1)%nlon,lon)
249       indexlat  = getsindex(slat,((nplot-1)//nlon)%nlat,lat)
250       indexvert = getsindex(svert,((nplot-1)//(nlon*nlat))%nvert,vert) 
251       if ope is not None:
252           if fileref is not None:      index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(namefiles)+2)  ## OK only 1 var,  see test in the beginning
253           elif "var" in ope:           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(var)+1)        ## OK only 1 file, see test in the beginning
254       else:                            yeah = len(namefiles)*len(var) ; index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%yeah
255       time = all_time[index_f]
256       if stime is not None:
257           if stime[0][0] < 0:
258               if typefile in ['mesoapi','meso']:
259                   for i in range(len(time)):  time[i] = localtime ( interv[0]+time[i]*interv[1], 0.5*(wlon[0]+wlon[1]) )
260                   print "OK... WORKING WITH LOCAL TIMES"
261               else: errormess("local times not supported for GCM files. not too hard to modify the code though.")
262       if mapmode == 1 and stime is None:   indextime = 1
263       else:                                indextime = getsindex(stime,((nplot-1)//(nlon*nlat*nvert))%ntime,time)
264       ltst = None 
265       if typefile in ['mesoapi','meso'] and indextime is not None:  ltst = localtime ( interv[0]+indextime*interv[1], 0.5*(wlon[0]+wlon[1]) ) 
266       print "********** index lon, lat, vert, time ",indexlon,indexlat,indexvert,indextime
267       ####################################################################
268
269       ticks = ndiv + 1
270
271       #### Contour plot
272       if var2:
273           what_I_plot, error = reducefield(all_var2[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert, yint=yintegral, alt=vert)
274           #what_I_plot = what_I_plot*mult
275           if not error:
276              if mapmode == 1:
277                  if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
278              zevmin, zevmax = calculate_bounds(what_I_plot)
279              zelevels = np.linspace(zevmin,zevmax,ticks) #20)
280              if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
281              if mapmode == 0:
282                  what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
283                        indextime,what_I_plot, len(all_var2[index_f].shape),vertmode)
284              ### If we plot a 2-D field
285              if len(what_I_plot.shape) is 2:
286                  cs = contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
287#                  clabel(cs,#zelevels[::2],
288#                     inline=3,
289#                     fmt='%1.1e',
290#                     fontsize=7)
291              ### If we plot a 1-D field
292              elif len(what_I_plot.shape) is 1:
293                  plot(what_I_plot,x) 
294           else:
295              errormess("There is an error in reducing field !")
296
297       #### Shaded plot
298       varname = all_varname[index_f]
299       if varname:
300           what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert , yint=yintegral, alt=vert, anomaly=anomaly)
301           what_I_plot = what_I_plot*mult
302           if not error: 
303               fvar = varname
304               if anomaly: fvar = 'anomaly'
305               ##### MAPMODE-SPECIFIC SETTINGS #####
306               if mapmode == 1:
307                   if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
308               elif mapmode == 0:
309                   what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
310                         indextime,what_I_plot, len(all_var[index_f].shape),vertmode)
311                   zxmin, zxmax = xaxis ; zymin, zymax = yaxis
312                   if zxmin is not None: mpl.pyplot.xlim(xmin=zxmin)
313                   if zxmax is not None: mpl.pyplot.xlim(xmax=zxmax)
314                   if zymin is not None: mpl.pyplot.ylim(ymin=zymin) 
315                   if zymax is not None: mpl.pyplot.ylim(ymax=zymax)
316                   if invert_y:     lima,limb = mpl.pyplot.ylim() ; mpl.pyplot.ylim(limb,lima)
317                   if ylog:         mpl.pyplot.semilogy()
318               if (fileref is not None) and (index_f is numplot-1):    zevmin, zevmax = calculate_bounds(what_I_plot,vmin=minop,vmax=maxop)
319               else:                                                   zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
320               if colorb in ["def","nobar"]:   palette = get_cmap(name=defcolorb(fvar.upper()))
321               elif (fileref is not None) and (index_f is numplot-1): palette = get_cmap(name="RdBu_r")
322               else:                           palette = get_cmap(name=colorb)
323               ##### 2D field
324               if len(what_I_plot.shape) is 2:
325                 if hole: what_I_plot = hole_bounds(what_I_plot,zevmin,zevmax)
326                 else: what_I_plot = bounds(what_I_plot,zevmin,zevmax)
327                 if flagnolow: what_I_plot = nolow(what_I_plot)
328                 if not tile:
329                     #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
330                     zelevels = np.linspace(zevmin,zevmax,num=ticks)
331                     if mapmode == 1:       m.contourf( x, y, what_I_plot, zelevels, cmap = palette)
332                     elif mapmode == 0:     contourf( x, y, what_I_plot, zelevels, cmap = palette)
333                 else:
334                     if mapmode == 1:       m.pcolor( x, y, what_I_plot, cmap = palette, vmin=zevmin, vmax=zevmax )
335                     elif mapmode == 0:     pcolor( x, y, what_I_plot, cmap = palette, vmin=zevmin, vmax=zevmax )
336                 if colorb != 'nobar':       
337                     if (fileref is not None) and (index_f is numplot-1):   daformat = "%.3f"
338                     else:                                                  daformat = fmtvar(fvar.upper())
339                     colorbar( fraction=0.05,pad=0.03,format=daformat,\
340                               ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,20])),extend='neither',spacing='proportional' ) 
341               ##### 1D field
342               elif len(what_I_plot.shape) is 1:
343                 plot(x,what_I_plot)
344                 if save == 'txt':  writeascii(np.transpose(what_I_plot),'profile'+str(nplot)+'.txt')
345
346               #### Other cases: (maybe plot 3-D field one day or movie ??)
347               else:
348                 print "WARNING!!! ",len(what_I_plot.shape),"-D PLOT NOT SUPPORTED !!! dimensions: ",what_I_plot.shape
349                 errormess("Are you sure you did not forget to prescribe a dimension ?")
350           else:
351               errormess("There is an error in reducing field !")
352
353       ### Vector plot --- a adapter
354       if winds:
355           vecx, error = reducefield( getfield(nc,uchar), d4=indextime, d3=indexvert , yint=yintegral , alt=vert)
356           vecy, error = reducefield( getfield(nc,vchar), d4=indextime, d3=indexvert , yint=yintegral , alt=vert)
357           #what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert )
358           if not error:
359               if typefile in ['mesoapi','meso']:   
360                   [vecx,vecy] = [dumpbdy(vecx,6,stag=uchar), dumpbdy(vecy,6,stag=vchar)]
361                   key = True
362               elif typefile in ['gcm']:           
363                   key = False
364               if metwind:  [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
365               if varname == False:       colorvec = definecolorvec(back)
366               else:                  colorvec = definecolorvec(colorb)
367               vectorfield(vecx, vecy,\
368                          x, y, stride=stride, csmooth=2,\
369                          #scale=15., factor=300., color=colorvec, key=key)
370                          scale=20., factor=250., color=colorvec, key=key)
371                                            #200.         ## or csmooth=stride
372   
373       ### Next subplot
374       basename = getname(var=varname,winds=winds,anomaly=anomaly)
375       basename = basename + getstralt(nc,level) 
376       if typefile in ['mesoapi','meso']:
377            if slon is not None: basename = basename + "_lon_" + str(int(lon2d[indices[1],indices[0]]))
378            if slat is not None: basename = basename + "_lat_" + str(int(lat2d[indices[1],indices[0]]))
379            plottitle = basename
380            if addchar: 
381                [addchar,gogol,gogol2] = getlschar ( all_namefile[index_f] )
382                plottitle = plottitle + addchar + "_LT"
383            else:        plottitle = plottitle + "_LT"
384            if ltst is not None: plottitle = plottitle + str(ltst)
385       else:
386            if fileref is not None:
387                if index_f is numplot-1:     plottitle = basename+' '+"fig(1) "+ope+" fig(2)"
388                elif index_f is numplot-2:   plottitle = basename+' '+fileref
389                else:                        plottitle = basename+' '+namefiles[0]#index_f]
390            else:                            plottitle = basename+' '+namefiles[0]#index_f]
391       if mult != 1:                         plottitle = str(mult) + "*" + plottitle
392       if zetitle != "fill":                 
393          plottitle = zetitle
394          if titleref is "fill":             titleref=zetitle
395          if fileref is not None:
396             if index_f is numplot-2:        plottitle = titleref
397             if index_f is numplot-1:        plottitle = "fig(1) "+ope+" fig(2)"
398#       if indexlon is not None:      plottitle = plottitle + " lon: " + str(min(lon[indexlon])) +" "+ str(max(lon[indexlon]))
399#       if indexlat is not None:      plottitle = plottitle + " lat: " + str(min(lat[indexlat])) +" "+ str(max(lat[indexlat]))
400#       if indexvert is not None:     plottitle = plottitle + " vert: " + str(min(vert[indexvert])) +" "+ str(max(vert[indexvert]))
401#       if indextime is not None:     plottitle = plottitle + " time: " + str(min(time[indextime])) +" "+ str(max(time[indextime]))
402       title( plottitle )
403       if nplot >= numplot: error = True
404       nplot += 1
405
406 
407
408
409
410
411
412     
413    ##########################################################################
414    ### Save the figure in a file in the data folder or an user-defined folder
415    if outputname is None:
416       if typefile in ['meso','mesoapi']:   prefix = getprefix(nc)
417       elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
418       else:                                prefix = ''
419    ###
420       zeplot = prefix + basename
421       if addchar:         zeplot = zeplot + addchar
422       if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
423    ###
424       if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
425       else:               zeplot = target + "/" + zeplot 
426    ###
427    else:
428       zeplot=outputname
429
430    pad_inches_value = 0.35
431    print "********** SAVE ", save
432    if save == 'png': 
433        if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
434        makeplotres(zeplot,res=resolution,pad_inches_value=pad_inches_value,disp=False)
435    elif save in ['eps','svg','pdf']:     makeplotres(zeplot,pad_inches_value=pad_inches_value,disp=False,ext=save)
436    elif save == 'gui':                   show()
437    elif save == 'txt':                   print "Saved results in txt file." 
438    else: 
439         print "INFO: save mode not supported. using gui instead."
440         show()
441
442    ###############
443    ### Now the end
444    return zeplot
Note: See TracBrowser for help on using the repository browser.