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

Last change on this file since 381 was 380, checked in by aslmd, 14 years ago

PYTHON: added an operation cat (especially aimed at time series with multiple files, working well with wrfout files. updated meso.py with last gcm.py changes. added a tip in help to use regular expressions in option -f.

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