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

Last change on this file since 400 was 399, checked in by aslmd, 14 years ago

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