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

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

Correction for different convention of naming the time axis in different ncdf files...

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