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

Last change on this file since 425 was 425, checked in by aslmd, 13 years ago

GRAPHICS: cosmetic changes. corrected a bug in treating local times with two input meso files.

  • Property svn:executable set to *
File size: 23.3 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']:
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          if vertmode is None:  vertmode=0
153          if vertmode == 0:     vert = np.arange(0,getattr(nc,vertdim),1)
154          else:                 vert = nc.variables["vert"][:]
155          time = np.arange(0,len(nc.variables["Times"]),1)
156       #if firstfile:
157       #   lat0 = lat
158       #elif len(lat0) != len(lat):
159       #   errormess("Not the same latitude lengths !", [len(lat0), len(lat)])
160       #elif sum((lat == lat0) == False) != 0:
161       #   errormess("Not the same latitudes !", [lat,lat0])
162       ## Faire d'autre checks sur les compatibilites entre fichiers!!
163##########################################################
164
165      if firstfile:
166         ##########################
167         ### Define plot boundaries
168         ### todo: possible areas in latinterv in argument (ex: "Far_South_Pole")
169         if proj in ["npstere","spstere"]: [wlon,wlat] = polarinterv(lon2d,lat2d)
170         elif proj in ["lcc","laea"]:      [wlon,wlat] = wrfinterv(lon2d,lat2d)
171         else:                             [wlon,wlat] = simplinterv(lon2d,lat2d)
172         if zoom:                          [wlon,wlat] = zoomset(wlon,wlat,zoom) 
173
174      all_varname[k] = varname
175      all_namefile[k] = namefile
176      all_time[k] = time
177      if var2: all_var2[k] = getfield(nc,var2)
178      ##### SPECIFIC
179      if varname in ["temp","t","T_nadir_nit","T_nadir_day","temp_day","temp_night"] and tsat:
180          tt=getfield(nc,varname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
181          if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
182          else:                                 tinput=tt
183          all_var[k]=get_tsat(vert,tinput,zlon=lon,zlat=lat,zalt=vert,ztime=time)
184      else:
185      ##### GENERAL STUFF HERE
186          all_var[k] = getfield(nc,varname)
187      print "********** all_var[k].shape", all_var[k].shape
188      k += 1
189      firstfile = False
190      #### End of for namefile in namefiles
191
192    ##################################
193    ### Operation on files
194    if ope is not None:
195        print ope
196        if "var" not in ope:
197             if len(var) > 1: errormess("for this operation... please set only one var !")
198             if ope in ["-","+"]:
199                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]
200                else:                     errormess("fileref is missing!") 
201                if ope == "-":     all_var[k+1]= all_var[k-1] - all_var[k]
202                elif ope == "+":   all_var[k+1]= all_var[k-1] + all_var[k]
203                all_varname[k+1] = all_varname[k] ; all_time[k+1] = all_time[k] ; all_namefile[k+1] = all_namefile[k] ; numplot = numplot+2
204             elif ope in ["cat"]:
205                tab = all_var[0];k = 1
206                while k != len(namefiles)-1: 
207                    tab = np.append(tab,all_var[k],axis=0) ; k += 1
208                all_time[0] = np.arange(0,len(tab),1) ### AS: time reference is too simplistic, should be better
209                all_var[0] = np.array(tab) ; numplot = 1
210             else: errormess(ope+" : non-implemented operation. Check pp.py --help")
211        else:
212             if len(namefiles) > 1: errormess("for this operation... please set only one file !") 
213             if len(var) > 2:       errormess("not sure this works for more than 2 vars... please check.")
214             if   ope in ["div_var"]: all_var[k] = all_var[k-2] / all_var[k-1] ; insert = '_div_'
215             elif ope in ["mul_var"]: all_var[k] = all_var[k-2] * all_var[k-1] ; insert = '_mul_'
216             elif ope in ["add_var"]: all_var[k] = all_var[k-2] + all_var[k-1] ; insert = '_add_'
217             elif ope in ["sub_var"]: all_var[k] = all_var[k-2] - all_var[k-1] ; insert = '_sub_'
218             else:                    errormess(ope+" : non-implemented operation. Check pp.py --help")
219             numplot = numplot + 1 ; all_time[k] = all_time[k-1] ; all_namefile[k] = all_namefile[k-1]
220             all_varname[k] = all_varname[k-2] + insert + all_varname[k-1] 
221
222    ##################################
223    ### Open a figure and set subplots
224    fig = figure()
225    subv,subh = definesubplot( numplot, fig ) 
226 
227    #################################
228    ### Time loop for plotting device
229    nplot = 1;error = False
230    print "********************************************"
231    while error is False:
232       print "********** nplot", nplot, "error",error
233     
234       ### General plot settings
235       if nplot > numplot: break
236       if numplot > 1:  subplot(subv,subh,nplot)
237
238       ### Map projection                   
239       if mapmode == 1:
240           m = define_proj(proj,wlon,wlat,back=back,blat=blat)
241           x, y = m(lon2d, lat2d)
242
243       ####################################################################
244       ## get all indexes to be taken into account for this subplot and then reduce field
245       ## 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
246       indexlon  = getsindex(slon,(nplot-1)%nlon,lon)
247       indexlat  = getsindex(slat,((nplot-1)//nlon)%nlat,lat)
248       indexvert = getsindex(svert,((nplot-1)//(nlon*nlat))%nvert,vert) 
249
250       if ope is not None:
251           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
252           elif "var" in ope:           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(var)+1)        ## OK only 1 file, see test in the beginning
253       else:                            yeah = len(namefiles)*len(var) ; index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%yeah
254       time = all_time[index_f]
255       if stime is not None:
256           if stime[0][0] < 0:
257               if typefile in ['mesoapi','meso']:
258                   for i in range(len(time)):  time[i] = localtime ( interv[0]+time[i]*interv[1], 0.5*(wlon[0]+wlon[1]) )
259                   print "OK... WORKING WITH LOCAL TIMES"
260               else: errormess("local times not supported for GCM files. not too hard to modify the code though.")
261       if mapmode == 1 and stime is None:   indextime = 1
262       else:                                indextime = getsindex(stime,((nplot-1)//(nlon*nlat*nvert))%ntime,time)
263       if typefile in ['mesoapi','meso'] and indextime is not None:  ltst = localtime ( interv[0]+indextime*interv[1], 0.5*(wlon[0]+wlon[1]) ) 
264       print "********** index lon, lat, vert, time ",indexlon,indexlat,indexvert,indextime
265       ####################################################################
266
267       ticks = ndiv + 1
268
269       #### Contour plot
270       if var2:
271           what_I_plot, error = reducefield(all_var2[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert, yint=yintegral, alt=vert)
272           #what_I_plot = what_I_plot*mult
273           if not error:
274              if mapmode == 1:
275                  if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
276              zevmin, zevmax = calculate_bounds(what_I_plot)
277              zelevels = np.linspace(zevmin,zevmax,ticks) #20)
278              if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
279              if mapmode == 0:
280                  x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
281                        indextime,what_I_plot.shape, len(all_var2[index_f].shape),vertmode)
282              ### If we plot a 2-D field
283              if len(what_I_plot.shape) is 2:
284                  cs = contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
285                  #if typefile in ['gcm']: clabel(cs,fmt = '%1.2e')
286              ### If we plot a 1-D field
287              elif len(what_I_plot.shape) is 1:
288                  plot(what_I_plot,x) 
289           else:
290              errormess("There is an error in reducing field !")
291
292       #### Shaded plot
293       varname = all_varname[index_f]
294       if varname:
295           what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert , yint=yintegral, alt=vert, anomaly=anomaly)
296           what_I_plot = what_I_plot*mult
297           if not error: 
298               fvar = varname
299               if anomaly: fvar = 'anomaly'
300               ##### MAPMODE-SPECIFIC SETTINGS #####
301               if mapmode == 1:
302                   if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
303               elif mapmode == 0:
304                   what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
305                         indextime,what_I_plot, len(all_var[index_f].shape),vertmode)
306                   zxmin, zxmax = xaxis ; zymin, zymax = yaxis
307                   if zxmin is not None: mpl.pyplot.xlim(xmin=zxmin)
308                   if zxmax is not None: mpl.pyplot.xlim(xmax=zxmax)
309                   if zymin is not None: mpl.pyplot.ylim(ymin=zymin) 
310                   if zymax is not None: mpl.pyplot.ylim(ymax=zymax)
311                   if invert_y:     lima,limb = mpl.pyplot.ylim() ; mpl.pyplot.ylim(limb,lima)
312                   if ylog:         mpl.pyplot.semilogy()
313               if (fileref is not None) and (index_f is numplot-1):    zevmin, zevmax = calculate_bounds(what_I_plot,vmin=minop,vmax=maxop)
314               else:                                                   zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
315               if colorb in ["def","nobar"]:   palette = get_cmap(name=defcolorb(fvar.upper()))
316               elif (fileref is not None) and (index_f is numplot-1): palette = get_cmap(name="RdBu_r")
317               else:                           palette = get_cmap(name=colorb)
318               ##### 2D field
319               if len(what_I_plot.shape) is 2:
320                 if hole: what_I_plot = hole_bounds(what_I_plot,zevmin,zevmax)
321                 else: what_I_plot = bounds(what_I_plot,zevmin,zevmax)
322                 if flagnolow: what_I_plot = nolow(what_I_plot)
323                 if not tile:
324                     #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
325                     zelevels = np.linspace(zevmin,zevmax,num=ticks)
326                     if mapmode == 1:       m.contourf( x, y, what_I_plot, zelevels, cmap = palette)
327                     elif mapmode == 0:     contourf( x, y, what_I_plot, zelevels, cmap = palette)
328                 else:
329                     if mapmode == 1:       m.pcolor( x, y, what_I_plot, cmap = palette, vmin=zevmin, vmax=zevmax )
330                     elif mapmode == 0:     pcolor( x, y, what_I_plot, cmap = palette, vmin=zevmin, vmax=zevmax )
331                 if colorb != 'nobar' and varname != 'HGT' :       
332                     if (fileref is not None) and (index_f is numplot-1):
333                        colorbar(fraction=0.05,pad=0.03,format="%.3f",\
334                                           ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,20])),\
335                                           extend='neither',spacing='proportional')
336                     else:
337                        colorbar(fraction=0.05,pad=0.03,format=fmtvar(fvar.upper()),\
338                                           ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,20])),\
339                                           extend='neither',spacing='proportional')
340                                           # both min max neither
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 !!!"
349                 print "field dimensions: ", what_I_plot.shape
350                 exit()
351               
352           else:
353               errormess("There is an error in reducing field !")
354
355       ### Vector plot --- a adapter
356       if winds:
357           vecx, error = reducefield( getfield(nc,uchar), d4=indextime, d3=indexvert , yint=yintegral , alt=vert)
358           vecy, error = reducefield( getfield(nc,vchar), d4=indextime, d3=indexvert , yint=yintegral , alt=vert)
359           #what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert )
360           if not error:
361               if typefile in ['mesoapi','meso']:   
362                   [vecx,vecy] = [dumpbdy(vecx,6,stag=uchar), dumpbdy(vecy,6,stag=vchar)]
363                   key = True
364               elif typefile in ['gcm']:           
365                   key = False
366               if metwind:  [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
367               if varname == False:       colorvec = definecolorvec(back)
368               else:                  colorvec = definecolorvec(colorb)
369               vectorfield(vecx, vecy,\
370                          x, y, stride=stride, csmooth=2,\
371                          #scale=15., factor=300., color=colorvec, key=key)
372                          scale=20., factor=250., color=colorvec, key=key)
373                                            #200.         ## or csmooth=stride
374   
375       ### Next subplot
376       basename = getname(var=varname,winds=winds,anomaly=anomaly)
377       basename = basename + getstralt(nc,level) 
378       if typefile in ['mesoapi','meso']:
379            if slon is not None: basename = basename + "_lon_" + str(int(lon2d[indices[1],indices[0]]))
380            if slat is not None: basename = basename + "_lat_" + str(int(lat2d[indices[1],indices[0]]))
381            plottitle = basename
382            if addchar: 
383                [addchar,gogol,gogol2] = getlschar ( all_namefile[index_f] )
384                plottitle = plottitle + addchar + "_LT"+str(ltst)
385            else:        plottitle = plottitle + "_LT"+str(ltst)
386       else:
387            if fileref is not None:
388                if index_f is numplot-1:     plottitle = basename+' '+"fig(1) "+ope+" fig(2)"
389                elif index_f is numplot-2:   plottitle = basename+' '+fileref
390                else:                        plottitle = basename+' '+namefiles[0]#index_f]
391            else:                            plottitle = basename+' '+namefiles[0]#index_f]
392       if mult != 1:                         plottitle = str(mult) + "*" + plottitle
393       if zetitle != "fill":                 
394          plottitle = zetitle
395          if titleref is "fill":             titleref=zetitle
396          if fileref is not None:
397             if index_f is numplot-2:        plottitle = titleref
398             if index_f is numplot-1:        plottitle = "fig(1) "+ope+" fig(2)"
399#       if indexlon is not None:      plottitle = plottitle + " lon: " + str(min(lon[indexlon])) +" "+ str(max(lon[indexlon]))
400#       if indexlat is not None:      plottitle = plottitle + " lat: " + str(min(lat[indexlat])) +" "+ str(max(lat[indexlat]))
401#       if indexvert is not None:     plottitle = plottitle + " vert: " + str(min(vert[indexvert])) +" "+ str(max(vert[indexvert]))
402#       if indextime is not None:     plottitle = plottitle + " time: " + str(min(time[indextime])) +" "+ str(max(time[indextime]))
403       title( plottitle )
404       if nplot >= numplot: error = True
405       nplot += 1
406
407 
408
409
410
411
412
413     
414    ##########################################################################
415    ### Save the figure in a file in the data folder or an user-defined folder
416    if outputname is None:
417       if typefile in ['meso','mesoapi']:   prefix = getprefix(nc)
418       elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
419       else:                                prefix = ''
420    ###
421       zeplot = prefix + basename
422       if addchar:         zeplot = zeplot + addchar
423       if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
424    ###
425       if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
426       else:               zeplot = target + "/" + zeplot 
427    ###
428    else:
429       zeplot=outputname
430
431    pad_inches_value = 0.35
432    print "********** SAVE ", save
433    if save == 'png': 
434        if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
435        makeplotres(zeplot,res=resolution,pad_inches_value=pad_inches_value,disp=False)
436    elif save in ['eps','svg','pdf']:     makeplotres(zeplot,pad_inches_value=pad_inches_value,disp=False,ext=save)
437    elif save == 'gui':                   show()
438    elif save == 'txt':                   print "Saved results in txt file." 
439    else: 
440         print "INFO: save mode not supported. using gui instead."
441         show()
442
443    ###############
444    ### Now the end
445    return zeplot
Note: See TracBrowser for help on using the repository browser.