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

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

PYTHON

M 393 mymath.py
----------------- Cosmetic change

M 393 make_netcdf.py
----------------- Cosmetic change

M 393 planetoplot.py
----------------- corrected bug with varname and --tsat

M 393 myplot.py
----------------- added possibility for the script to read ncdf files with NaN value and treat them correctly without messing with the mathematical operations.
  • Property svn:executable set to *
File size: 22.1 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           numplot=None,\
18           var=None,\
19           colorb="def",\
20           winds=True,\
21           addchar=None,\
22           interv=[0,1],\
23           vmin=None,\
24           vmax=None,\
25           tile=False,\
26           zoom=None,\
27           display=True,\
28           itstep=None,\
29           hole=False,\
30           save="gui",\
31           anomaly=False,\
32           var2=None,\
33           ndiv=10,\
34           first=1,\
35           mult=1.,\
36           zetitle="fill",\
37           slon=None,\
38           slat=None,\
39           svert=None,\
40           stime=None,\
41           outputname=None,\
42           resolution=200,\
43           ope=None,\
44           fileref=None,\
45           minop=0.,\
46           maxop=0.,\
47           titleref="fill",\
48           invert_y=False,\
49           xaxis=[None,None],\
50           yaxis=[None,None],\
51           ylog=False,\
52           yintegral=False,\
53           blat=None,\
54           tsat=False):
55
56
57    ####################################################################################################################
58    ### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
59
60    #################################
61    ### Load librairies and functions
62    from netCDF4 import Dataset
63    from myplot import getcoord2d,define_proj,makeplotres,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy,\
64                       fmtvar,definecolorvec,defcolorb,getprefix,putpoints,calculate_bounds,errormess,definesubplot,\
65                       zoomset,getcoorddef,getwinddef,whatkindfile,reducefield,bounds,getstralt,getfield,smooth,nolow,\
66                       getname,localtime,polarinterv,getsindex,define_axis,determineplot,readslices
67    from mymath import deg,max,min,mean,get_tsat
68    import matplotlib as mpl
69    from matplotlib.pyplot import contour,contourf, subplot, figure, rcParams, savefig, colorbar, pcolor, show, plot, clabel, title
70    from matplotlib.cm import get_cmap
71    import numpy as np
72    from numpy.core.defchararray import find
73
74    ################################
75    ### Preliminary stuff
76    ################################
77    print "********************************************"
78    print "********** WELCOME TO PLANETOPLOT **********"
79    print "********************************************"
80    if not isinstance(namefiles, np.ndarray): namefiles = [namefiles]
81    if not isinstance(var, np.ndarray):       var = [var]
82    if ope is not None and len(var) > 1:      errormess("you are using an operation... please set only one var !")
83
84    ################################
85    ### Which plot needs to be done?
86    ################################
87    nlon, nlat, nvert, ntime, mapmode, nslices = determineplot(slon, slat, svert, stime)
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    if numplot is None:  numplot = zelen*nslices
93    print "********** FILES, SLICES, VARS, TOTAL PLOTS: ", len(namefiles), nslices, len(var), numplot
94    print "********** MAPMODE: ", mapmode
95    if fileref is not None:       all_var  = [[]]*(zelen+2) ;  all_varname = [[]]*(zelen+2)
96    else:                         all_var  = [[]]*zelen     ;  all_var2  = [[]]*zelen ;  all_title = [[]]*zelen ;  all_varname = [[]]*zelen
97 
98    #################################################################################################
99    ### Loop over the files + vars initially separated by commas to be plotted on the same figure ###
100    #################################################################################################
101    k = 0
102    firstfile = True
103    for nnn in range(len(namefiles)):
104     for vvv in range(len(var)): 
105
106      print "********** LOOP..... THIS IS SUBPLOT NUMBER.....",k
107
108      ######################
109      ### Load NETCDF object
110      namefile = namefiles[nnn] ; print "********** THE NAMEFILE IS....", namefile
111      nc  = Dataset(namefile)
112
113      ##################################
114      ### Initial checks and definitions
115      ### ... TYPEFILE
116      typefile = whatkindfile(nc)                                 
117      if firstfile:                 typefile0 = typefile
118      elif typefile != typefile0:   errormess("Not the same kind of files !", [typefile0, typefile])
119      ### ... VAR
120      varname=var[vvv]
121      print "********** THE VAR IS....",varname, var2
122      if varname not in nc.variables: varname = False
123      ### ... WINDS
124      if winds:                                                   
125         [uchar,vchar,metwind] = getwinddef(nc)             
126         if uchar == 'not found': winds = False
127      if not varname and not winds: errormess("please set at least winds or var",printvar=nc.variables)
128      ### ... COORDINATES, could be moved below
129      [lon2d,lat2d] = getcoorddef(nc)                       
130      ### ... PROJECTION
131      if proj == None:   proj = getproj(nc)                 
132
133##########################################################
134      if typefile == "gcm":
135          lat = nc.variables["latitude"][:] 
136          lon = nc.variables["longitude"][:]
137          if "Time" in nc.variables:
138             time = nc.variables["Time"][:]
139          elif "time" in nc.variables:
140             time = nc.variables["time"][:]
141          else:
142             print "no time axis found."
143             exit()
144          vert = nc.variables["altitude"][:]
145      elif typefile in ['meso','mesoapi']:
146          if varname in ['PHTOT','W']:    vertdim='BOTTOM-TOP_PATCH_END_STAG'
147          else:                           vertdim='BOTTOM-TOP_PATCH_END_UNSTAG'
148          if varname in ['V']:  latdim='SOUTH-NORTH_PATCH_END_STAG'
149          else:                 latdim='SOUTH-NORTH_PATCH_END_UNSTAG'
150          if varname in ['U']:  londim='WEST-EAST_PATCH_END_STAG'
151          else:                 londim='WEST-EAST_PATCH_END_UNSTAG'
152          lon = np.arange(0,getattr(nc,londim),1)
153          lat = np.arange(0,getattr(nc,latdim),1)
154          if vertmode is None:  vertmode=0
155          if vertmode == 0:     vert = np.arange(0,getattr(nc,vertdim),1)
156          else:                 vert = nc.variables["vert"][:]
157          time = np.arange(0,len(nc.variables["Times"]),1)
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         ### ... NAME FOR TITLE and FILE (have to do something for multiple vars...)
176         basename = getname(var=varname,winds=winds,anomaly=anomaly)
177         basename = basename + getstralt(nc,level)  ## can be moved elsewhere for a more generic routine
178
179      ##### SPECIFIC
180      if varname in ["temp","t","T_nadir_nit","T_nadir_day"] and tsat:
181          tt=getfield(nc,varname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
182          if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
183          else:                                 tinput=tt
184          all_var[k]=get_tsat(vert,tinput,zlon=lon,zlat=lat,zalt=vert,ztime=time)
185          all_varname[k] = varname
186      else:
187      ##### GENERAL STUFF HERE
188          all_var[k] = getfield(nc,varname)
189          all_varname[k] = varname
190      if var2: all_var2[k] = getfield(nc,var2)
191     
192      print "********** all_var[k].shape", all_var[k].shape
193      k += 1
194      firstfile = False
195      #### End of for namefile in namefiles
196
197    ##################################
198    ### Operation on files
199    ### ... for the moment only OK when 1 var only
200    if ope is not None:
201             print ope
202             if ope in ["-","+"]:
203                if fileref is not None:   all_var[k] = getfield(Dataset(fileref),all_varname[k-1]) ; all_varname[k] = all_varname[k-1]
204                else:                     errormess("fileref is missing!") 
205                if ope == "-":     all_var[k+1]= all_var[k-1] - all_var[k]
206                elif ope == "+":   all_var[k+1]= all_var[k-1] + all_var[k]
207                all_varname[k+1] = all_varname[k]
208                numplot = numplot+2
209             elif ope in ["cat"]:
210                tab = all_var[0];k = 0
211                while k != len(namefiles)-1: 
212                    tab = np.append(tab,all_var[k],axis=0)
213                    k += 1
214                time = np.arange(0,len(tab),1) ### AS: time reference is simple, should be better
215                all_var[0] = np.array(tab);numplot = 1
216             elif ope is not None:
217                errormess(ope+" : non-implemented operation. Available op. are + and -.")
218
219    ##################################
220    ### Open a figure and set subplots
221    fig = figure()
222    subv,subh = definesubplot( numplot, fig ) 
223 
224    #################################
225    ### Time loop for plotting device
226    found_lct = False;nplot = 1;itime = first;error = False
227    if itstep is None and numplot > 0: itstep = int(24./numplot)
228    elif numplot <= 0:                 itstep = 1 
229   
230    print "********************************************"
231    while error is False:
232       print "********** nplot", nplot, "itime",itime,"error",error
233       
234       ### Which local time ?
235       ltst = localtime ( interv[0]+itime*interv[1], 0.5*(wlon[0]+wlon[1]) )
236
237       ### General plot settings
238       #print itime, int(ltst), numplot, nplot
239       if numplot >= 1: 
240           if nplot > numplot: break
241           if numplot > 1:     
242               if typefile not in ['geo']:  subplot(subv,subh,nplot)
243           found_lct = True
244       ### If only one local time is requested (numplot < 0)
245       elif numplot <= 0: 
246           if int(ltst) + numplot != 0: 
247                 itime += 1 
248                 if found_lct is True: break     ## because it means LT was found at previous iteration
249                 else:                 continue  ## continue to iterate to find the correct LT
250           else: 
251                 found_lct = True
252
253       ### Map projection                   
254       if mapmode == 1:
255           m = define_proj(proj,wlon,wlat,back=back,blat=blat)
256           x, y = m(lon2d, lat2d)
257
258####################################################################
259       ## get all indexes to be taken into account for this subplot and then reduce field
260       ## 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
261       indexlon  = getsindex(slon,(nplot-1)%nlon,lon)
262       indexlat  = getsindex(slat,((nplot-1)//nlon)%nlat,lat)
263       indexvert = getsindex(svert,((nplot-1)//(nlon*nlat))%nvert,vert) 
264       if mapmode == 1 and stime is None: 
265           indextime = itime
266           if len(var) != 1 or len(namefiles) != 1: indextime = first
267       else:
268           indextime = getsindex(stime,((nplot-1)//(nlon*nlat*nvert))%ntime,time) 
269
270       if fileref is not None:
271           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(namefiles)+2)  ## OK only 1 var, see test in the beginning
272       else:
273           yeah = len(namefiles)*len(var)
274           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%yeah
275       #print nlon, nlat, nvert, ntime 
276       #print slon, slat, svert, stime       
277       #print index_f
278       #print "index lon,lat,vert,time", max(indexlon), max(indexlat), max(indexvert), max(indextime)
279       #if varname != all_varname[index_f]:  indextime = first
280####################################################################
281
282       ticks = ndiv + 1
283
284       #### Contour plot
285       if var2:
286           what_I_plot, error = reducefield(all_var2[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert, yint=yintegral, alt=vert)
287           #what_I_plot = what_I_plot*mult
288           if not error:
289              if mapmode == 1:
290                  if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
291              zevmin, zevmax = calculate_bounds(what_I_plot)
292              zelevels = np.linspace(zevmin,zevmax,ticks) #20)
293              if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
294              if mapmode == 0:
295                  x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
296                        indextime,what_I_plot.shape, len(all_var2[index_f].shape),vertmode)
297              ### If we plot a 2-D field
298              if len(what_I_plot.shape) is 2:
299                  cs = contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
300                  #if typefile in ['gcm']: clabel(cs,fmt = '%1.2e')
301              ### If we plot a 1-D field
302              elif len(what_I_plot.shape) is 1:
303                  plot(what_I_plot,x) 
304           else:
305              errormess("There is an error in reducing field !")
306
307       #### Shaded plot
308       varname = all_varname[index_f]
309       if varname:
310           what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert , yint=yintegral, alt=vert)
311           what_I_plot = what_I_plot*mult
312           if not error: 
313               fvar = varname
314               ###
315               if anomaly:
316                   what_I_plot = 100. * ((what_I_plot / smooth(what_I_plot,12)) - 1.)
317                   fvar = 'anomaly'
318               #if mult != 1:     
319               #    fvar = str(mult) + "*" + var
320               ###
321               if mapmode == 1:
322                   if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
323               elif mapmode == 0:
324                   what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
325                         indextime,what_I_plot, len(all_var[index_f].shape),vertmode)
326               if (fileref is not None) and (index_f is numplot-1):
327                  zevmin, zevmax = calculate_bounds(what_I_plot,vmin=minop,vmax=maxop)
328               else:
329                  zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
330               if colorb in ["def","nobar"]:   palette = get_cmap(name=defcolorb(fvar.upper()))
331               else:                           palette = get_cmap(name=colorb)
332               ##### 2D field
333               if len(what_I_plot.shape) is 2:
334                 if not tile:
335                     if not hole: what_I_plot = bounds(what_I_plot,zevmin,zevmax)
336                     #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
337                     zelevels = np.linspace(zevmin,zevmax,num=ticks)
338                     #contourf(what_I_plot, zelevels, cmap = palette )
339
340                     if mapmode == 1:
341                       m.contourf( x, y, what_I_plot, zelevels, cmap = palette)
342                     elif mapmode == 0:
343                       contourf( x, y, what_I_plot, zelevels, cmap = palette)
344
345                     zxmin, zxmax = xaxis
346                     zymin, zymax = yaxis
347                     if zxmin is not None: mpl.pyplot.xlim(xmin=zxmin)
348                     if zxmax is not None: mpl.pyplot.xlim(xmax=zxmax)
349                     if zymin is not None: mpl.pyplot.ylim(ymin=zymin)
350                     if zymax is not None: mpl.pyplot.ylim(ymax=zymax)
351                     if invert_y:     lima,limb = mpl.pyplot.ylim() ; mpl.pyplot.ylim(limb,lima)
352                     if ylog:         mpl.pyplot.semilogy()
353
354                 else:
355                     if hole:  what_I_plot = nolow(what_I_plot)
356                     pcolor( x, y, what_I_plot, cmap = palette, vmin=zevmin, vmax=zevmax )
357                 if colorb != 'nobar' and varname != 'HGT' :       
358                     if (fileref is not None) and (index_f is numplot-1):
359                        colorbar(fraction=0.05,pad=0.03,format="%.2f",\
360                                           ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,20])),\
361                                           extend='neither',spacing='proportional')
362                     else:
363                        colorbar(fraction=0.05,pad=0.03,format=fmtvar(fvar.upper()),\
364                                           ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,20])),\
365                                           extend='neither',spacing='proportional')
366
367                                           # both min max neither
368               ##### 1D field
369               elif len(what_I_plot.shape) is 1:
370                 plot(x,what_I_plot)
371               #### Other cases: (maybe plot 3-D field one day or movie ??)
372               else:
373                 print "WARNING!!! ",len(what_I_plot.shape),"-D PLOT NOT SUPPORTED !!!"
374                 print "field dimensions: ", what_I_plot.shape
375                 exit()
376               
377           else:
378               errormess("There is an error in reducing field !")
379
380       ### Vector plot --- a adapter
381       if winds:
382           vecx, error = reducefield( getfield(nc,uchar), d4=indextime, d3=indexvert , yint=yintegral , alt=vert)
383           vecy, error = reducefield( getfield(nc,vchar), d4=indextime, d3=indexvert , yint=yintegral , alt=vert)
384           #what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert )
385           if not error:
386               if typefile in ['mesoapi','meso']:   
387                   [vecx,vecy] = [dumpbdy(vecx,6,stag=uchar), dumpbdy(vecy,6,stag=vchar)]
388                   key = True
389               elif typefile in ['gcm']:           
390                   key = False
391               if metwind:  [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
392               if varname == False:       colorvec = definecolorvec(back)
393               else:                  colorvec = definecolorvec(colorb)
394               vectorfield(vecx, vecy,\
395                          x, y, stride=stride, csmooth=2,\
396                          #scale=15., factor=300., color=colorvec, key=key)
397                          scale=20., factor=250., color=colorvec, key=key)
398                                            #200.         ## or csmooth=stride
399   
400       ### Next subplot
401       basename = getname(var=varname,winds=winds,anomaly=anomaly)
402       if typefile in ['mesoapi','meso']:
403            plottitle = basename
404            if addchar:  plottitle = plottitle + addchar + "_LT"+str(ltst)
405            else:        plottitle = plottitle + "_LT"+str(ltst)
406       else:
407            if fileref is not None:
408                if index_f is numplot-1:
409                   plottitle = basename+' '+"fig(1) "+ope+" fig(2)"
410                elif index_f is numplot-2:
411                   plottitle = basename+' '+fileref
412                else:
413                   plottitle = basename+' '+namefiles[0]#index_f]
414            else:
415                plottitle = basename+' '+namefiles[0]#index_f]
416       if mult != 1:           plottitle = str(mult) + "*" + plottitle
417       if zetitle != "fill":   
418          plottitle = zetitle
419          if titleref is "fill":
420             titleref=zetitle
421          if fileref is not None:
422             if index_f is numplot-2:
423                plottitle = titleref
424             if index_f is numplot-1:
425                plottitle = "fig(1) "+ope+" fig(2)"
426#       if indexlon is not None:
427#         plottitle = plottitle + " lon: " + str(min(lon[indexlon])) +" "+ str(max(lon[indexlon]))
428#       if indexlat is not None:
429#         plottitle = plottitle + " lat: " + str(min(lat[indexlat])) +" "+ str(max(lat[indexlat]))
430#       if indexvert is not None:
431#         plottitle = plottitle + " vert: " + str(min(vert[indexvert])) +" "+ str(max(vert[indexvert]))
432#       if indextime is not None:
433#         plottitle = plottitle + " time: " + str(min(time[indextime])) +" "+ str(max(time[indextime]))
434       title( plottitle )
435       itime += itstep
436       if nplot >= numplot:
437          error = True
438       nplot += 1
439
440 
441
442
443
444
445
446     
447    ##########################################################################
448    ### Save the figure in a file in the data folder or an user-defined folder
449    if outputname is None:
450       if typefile in ['meso','mesoapi']:   prefix = getprefix(nc)
451       elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
452       else:                                prefix = ''
453    ###
454       zeplot = prefix + basename
455       if addchar:         zeplot = zeplot + addchar
456       if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
457    ###
458       if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
459       else:               zeplot = target + "/" + zeplot 
460    ###
461    else:
462       zeplot=outputname
463
464    if found_lct:     
465        pad_inches_value = 0.35
466        print "********** SAVE ", save
467        if save == 'png': 
468            if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
469            makeplotres(zeplot,res=resolution,pad_inches_value=pad_inches_value,disp=False)
470        elif save in ['eps','svg','pdf']:
471            makeplotres(zeplot,         pad_inches_value=pad_inches_value,disp=False,ext=save)
472        elif save == 'gui':
473            show()
474        else: 
475            print "INFO: save mode not supported. using gui instead."
476            show()
477    else:   print "!!! Local time not found"
478
479    ###############
480    ### Now the end
481    return zeplot
Note: See TracBrowser for help on using the repository browser.