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

Last change on this file since 647 was 647, checked in by tnavarro, 13 years ago

Corrected bug in reducefield for masked arrays with grid area. Possibility to see stream function for a lat/vert slice with --stream option. Output of both data and axis with save txt option in 1D. Small bug corrected: title is the file name for multiple plots with multiple plots. Cleanup in myplot.py

  • Property svn:executable set to *
File size: 47.7 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~12/2011 -- Extended multivar subplot capabilities + cosmetic changes + general cleaning and tests
9### A. Colaitis  -- LMD --    12/2011 -- Added movie capability [mencoder must be installed]
10### A. Spiga     -- LMD --    12/2011 -- Added HTML animated page capability + general tests of consistency [winds, etc...] + consistent generic movie loop
11### J. Leconte   -- LMD --    02/2012 -- Added area weighted averaging. Compatibility with terrestrial gcm.
12### A. Spiga     -- LMD --    03/2012 -- Cleaning and improved comments
13def planetoplot (namefiles,\
14           level=0,\
15           vertmode=0,\
16           proj=None,\
17           back=None,\
18           target=None,
19           stride=3,\
20           var=None,\
21           colorb="def",\
22           winds=False,\
23           addchar=None,\
24           interv=[0,1],\
25           vmin=None,\
26           vmax=None,\
27           tile=False,\
28           zoom=None,\
29           display=True,\
30           hole=False,\
31           save="gui",\
32           anomaly=False,\
33           var2=None,\
34           ndiv=10,\
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           blon=None,\
55           tsat=False,\
56           flagnolow=False,\
57           mrate=None,\
58           mquality=False,\
59           trans=1,\
60           zarea=None,\
61           axtime=None,\
62           redope=None,\
63           seevar=False,\
64           xlab=None,\
65           ylab=None,\
66           lbls=None,\
67           lstyle=None,\
68           cross=None,\
69           streamflag=False):
70
71    ####################################################################################################################
72    ### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
73
74    #################################
75    ### Load librairies and functions
76    from netCDF4 import Dataset
77    from myplot import getcoord2d,define_proj,makeplotres,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy,\
78                       fmtvar,definecolorvec,defcolorb,getprefix,putpoints,calculate_bounds,errormess,definesubplot,\
79                       zoomset,getcoorddef,getwinddef,whatkindfile,reducefield,bounds,getstralt,getfield,smooth,nolow,\
80                       getname,localtime,check_localtime,polarinterv,getsindex,define_axis,determineplot,readslices,bidimfind,getlschar,hole_bounds,\
81                       windamplitude,slopeamplitude,deltat0t1
82    from mymath import deg,max,min,mean,get_tsat,writeascii,fig2data,fig2img
83    import matplotlib as mpl
84    from matplotlib.pyplot import contour,contourf,hist, text,subplot, figure, rcParams, savefig, colorbar, pcolor, show, plot, clabel, title, close, legend, xlabel, axis, ylabel, subplots_adjust, axes, clabel
85    from matplotlib.cm import get_cmap
86    #from mpl_toolkits.basemap import cm
87    import numpy as np
88    from numpy.core.defchararray import find
89    from videosink import VideoSink
90    from timestuff import sol2ls
91    import subprocess
92    #from singlet import singlet
93    from itertools import cycle
94    import os
95
96#########################
97### PRELIMINARY STUFF ###
98#########################
99### we say hello
100    print "********************************************"
101    print "********** WELCOME TO PLANETOPLOT **********"
102    print "********************************************"
103### we ensure namefiles and var are arrays
104    if not isinstance(namefiles, np.ndarray): namefiles = [namefiles]
105    if not isinstance(var, np.ndarray):       var = [var]
106### we initialize a few variables
107    initime=-1 ; sslon = None ; sslat = None
108    if slon is not None: sslon = np.zeros([1,2])
109    if slat is not None: sslat = np.zeros([1,2])
110    k = 0 ; firstfile = True ; count = 0
111### we avoid specific cases not yet implemented
112    if mrate is not None and len(var) > 1: errormess("multivar not allowed in movies. should be fixed soon!")
113### we prepare number of plot fields [zelen] and plot instances [numplot] according to user choices
114### --> we support multifile and multivar plots : files + vars separated by commas are on the same figure
115    nlon, nlat, nvert, ntime, mapmode, nslices = determineplot(slon, slat, svert, stime)
116    zelen = len(namefiles)*len(var)
117    numplot = zelen*nslices
118    print "********** FILES, SLICES, VARS, TOTAL PLOTS: ", len(namefiles), nslices, len(var), numplot
119    print "********** MAPMODE: ", mapmode
120### we correct number of plot fields for possible operation (substract, etc...)
121    if ope is not None:
122        if fileref is not None:       zelen = zelen + 2
123        elif "var" in ope:            zelen = zelen + 1
124### we define the arrays for plot fields -- which will be placed within multiplot loops
125    all_var  = [[]]*zelen ; all_var2  = [[]]*zelen
126    all_title = [[]]*zelen ; all_varname = [[]]*zelen ; all_namefile = [[]]*zelen ; all_time = [[]]*zelen
127    all_windu = [[]]*zelen ; all_windv = [[]]*zelen
128 
129#############################
130### LOOP OVER PLOT FIELDS ###
131#############################
132    for nnn in range(len(namefiles)):
133     for vvv in range(len(var)): 
134
135    ### we load each NETCDF objects in namefiles
136      namefile = namefiles[nnn] 
137      nc  = Dataset(namefile)
138    ### we explore the variables in the file
139      varinfile = nc.variables.keys()
140      if seevar: print varinfile ; exit()
141    ### we define the type of file we have (gcm, meso, etc...)
142      typefile = whatkindfile(nc)
143      if firstfile:                 typefile0 = typefile
144      elif typefile != typefile0:   errormess("Not the same kind of files !", [typefile0, typefile])
145    ### we care for input file being 1D simulations
146      if "longitude" in nc.dimensions and "latitude" in nc.dimensions: is1d = len(nc.variables["longitude"][:])*len(nc.variables["latitude"][:])
147      elif "lon" in nc.dimensions and "lat" in nc.dimensions: is1d = len(nc.variables["lon"][:])*len(nc.variables["lat"][:])
148      if typefile in ['gcm','earthgcm'] and is1d == 1:       mapmode=0 ; winds=False
149    ### we create default vert and time prescriptions if not here in case mapping mode is on (lat/lon maps)
150      if redope is None and mapmode == 1:
151          if svert is None:  svert = readslices(str(level)) ; nvert=1
152          if stime is None and mrate is None:
153             stime = readslices(str(0)) ; ntime=1 ## this is a default choice
154             print "WELL... nothing about time axis. I took default: first time reference stored in file."
155    ### we get the names of variables to be read. in case only one available, we choose this one.
156      varname=var[vvv] 
157      if ((varname not in nc.variables) and not ((typefile in ['meso']) and (varname in ['UV','uv','uvmet','slopexy','SLOPEXY','deltat','DELTAT']))): 
158          if len(varinfile) == 1:   varname = varinfile[0] 
159          else:                     varname = False
160    ### we get the names of wind variables to be read (if any)
161      if winds:                                                   
162         [uchar,vchar,metwind] = getwinddef(nc)             
163         if uchar == 'not found': winds = False
164    ### we tell the user that either no var or no wind is not acceptable
165      if not varname and not winds: errormess("please set at least winds or var",printvar=nc.variables)
166    ### we get the coordinates lat/lon to be used
167      [lon2d,lat2d] = getcoorddef(nc)
168    ### we get an adapted map projection if none is provided by the user
169      if proj == None:   proj = getproj(nc)                 
170    ### we define plot boundaries given projection or user choices
171      if firstfile:
172         if proj in ["npstere","spstere"]: [wlon,wlat] = polarinterv(lon2d,lat2d)
173         elif proj in ["lcc","laea"]:      [wlon,wlat] = wrfinterv(lon2d,lat2d)
174         else:                             [wlon,wlat] = simplinterv(lon2d,lat2d)
175         if zoom:                          [wlon,wlat] = zoomset(wlon,wlat,zoom)
176         elif zarea is not None:           [wlon,wlat] = latinterv(area=zarea)
177
178##########################################################
179############ LOAD 4D DIMENSIONS : x, y, z, t #############
180##########################################################
181      if typefile in ["gcm","earthgcm","ecmwf"]:
182          if slon is not None: sslon = slon 
183          if slat is not None: sslat = slat 
184          ### LAT
185          if "latitude" in nc.variables:  lat = nc.variables["latitude"][:]
186          elif "lat" in nc.variables:     lat = nc.variables["lat"][:]
187          ### LON
188          if "longitude" in nc.variables: lon = nc.variables["longitude"][:]
189          elif "lon" in nc.variables:     lon = nc.variables["lon"][:]
190          ### ALT
191          if "altitude" in nc.variables:  vert = nc.variables["altitude"][:]
192          elif "Alt" in nc.variables:     vert = nc.variables["Alt"][:]
193          elif "lev" in nc.variables:     vert = nc.variables["lev"][:]
194          else:                           vert = [0.]
195          ### AIRE (to weight means with the area)
196          if "aire" in nc.variables:      area = nc.variables["aire"][:,:] 
197          else:                           area = None
198          ### TIME
199          if "Time" in nc.variables:            time = nc.variables["Time"][:]
200          elif "time" in nc.variables:          time = nc.variables["time"][:]
201          elif "time_counter" in nc.variables:  time = nc.variables["time_counter"][:]/86400. #### time counter cinverstion from s-> days
202          else:                                 time = [0.] #errormess("no time axis found.")
203          if axtime in ["ls"]:
204              print "converting to Ls ..."
205              for iii in range(len(time)):
206                time[iii] = sol2ls(time[iii]) # to use with timestuff
207                if iii > 0:
208                  while abs(time[iii]-time[iii-1]) > 300:
209                    time[iii] = time[iii]+360
210          # for 1D plots (no need for longitude computation):
211          if axtime in ["lt"]:
212              if initime == -1: initime=input("Please type initial local time:")
213              time = (initime+time*24)%24
214              print "LOCAL TIMES.... ", time
215      elif typefile in ['meso','geo']:
216          area = None ## not active for the moment
217          ###### STUFF TO GET THE CORRECT LAT/LON FROM MESOSCALE FILES WITH 2D LAT/LON ARRAYS
218          ###### principle: calculate correct indices then repopulate slon and slat
219          if slon is not None or slat is not None:
220              if firstfile and save == 'png' and typefile == 'meso' and "HGT" in varinfile and display:   iwantawhereplot = nc     #show a topo map with a cross on the chosen point
221              else:                                                                           iwantawhereplot = None   #do not show anything, just select indices
222              numlon = 1 ; numlat = 1 
223              if slon is not None:   numlon = slon.shape[1]   
224              if slat is not None:   numlat = slat.shape[1]
225              indices = np.ones([numlon,numlat,2]) ; vlon = None ; vlat = None
226              for iii in range(numlon): 
227               for jjj in range(numlat):
228                 if slon is not None:  vlon = slon[0][iii]  ### note: slon[:][0] does not work
229                 if slat is not None:  vlat = slat[0][jjj]  ### note: slon[:][0] does not work
230                 indices[iii,jjj,:] = bidimfind(lon2d,lat2d,vlon,vlat,file=iwantawhereplot) 
231                 lonp,latp = ( lon2d[indices[iii,jjj,0],indices[iii,jjj,1]] , lat2d[indices[iii,jjj,0],indices[iii,jjj,1]] )
232              ### possible bug here if several --lat
233              for iii in range(numlon):
234               for jjj in range(numlat):
235                 if slon is not None: sslon[0][iii] = indices[iii,0,1] #...this is idx
236                 if slat is not None: sslat[0][jjj] = indices[0,jjj,0] #...this is idy
237              lonp,latp = ( lon2d[indices[0,0,0],indices[0,0,1]] , lat2d[indices[0,0,0],indices[0,0,1]] )
238          ######
239          if typefile in ['meso'] and mapmode == 1: lon2d = dumpbdy(lon2d,6) ; lat2d = dumpbdy(lat2d,6)  ### important to do that now and not before
240          ######
241          if varname in ['PHTOT','W']:    vertdim='BOTTOM-TOP_PATCH_END_STAG'
242          else:                           vertdim='BOTTOM-TOP_PATCH_END_UNSTAG'
243          if (var2 is not None and var2 not in ['PHTOT','W']): dumped_vert_stag=True ; vertdim='BOTTOM-TOP_PATCH_END_UNSTAG'
244          else:                                                dumped_vert_stag=False
245          if varname in ['V']:  latdim='SOUTH-NORTH_PATCH_END_STAG'
246          else:                 latdim='SOUTH-NORTH_PATCH_END_UNSTAG'
247          if varname in ['U']:  londim='WEST-EAST_PATCH_END_STAG'
248          else:                 londim='WEST-EAST_PATCH_END_UNSTAG'
249          lon = np.arange(0,getattr(nc,londim),1) ; lat = np.arange(0,getattr(nc,latdim),1)
250          ###
251          if axtime in ["ls","sol"]:
252              lstab, soltab, lttab = getlschar ( namefile, getaxis = True )
253              if axtime == "ls":      time = lstab
254              elif axtime == "sol":   time = soltab
255          else:
256              if "Times" in nc.dimensions:   time = count + np.arange(0,len(nc.dimensions["Times"]),1)
257              elif "Time" in nc.dimensions:  time = count + np.arange(0,len(nc.dimensions["Time"]),1)
258              else:                          time = count + np.arange(0,1,1)
259              if ope in ["cat"]:
260                 if nnn > 0:  count = time[-1] + 1  ## so that a cat is possible with simple subscripts
261                 else:        count = 0
262              else: count = 0
263          if axtime in ["lt"]:
264              ftime = np.zeros(len(time))
265              for i in range(len(time)): 
266                 ftime[i] = localtime ( interv[0]+time[i]*interv[1], 0.5*(wlon[0]+wlon[1]) )
267              time=ftime
268              time=check_localtime(time)
269              print "LOCAL TIMES.... ", time
270
271          ###
272          if typefile in ['geo']:   vert = [0.] ; stime = readslices(str(0))
273          else:
274              if vertmode is None:  vertmode=0
275              if vertmode == 0:     vert = np.arange(0,getattr(nc,vertdim),1)
276              elif vertmode == -1:  vert = nc.variables["PHTOT"][0,:,0,0]/3.72 ; vert = np.array(vert[0:len(vert)-1]) #; print vert
277              elif vertmode == 1 or vertmode == 2:  vert = nc.variables["vert"][:]        ## pressure in Pa
278              else:                                 vert = nc.variables["vert"][:]/1000.  ## altitude in km
279       #if firstfile:
280       #   lat0 = lat
281       #elif len(lat0) != len(lat):
282       #   errormess("Not the same latitude lengths !", [len(lat0), len(lat)])
283       #elif sum((lat == lat0) == False) != 0:
284       #   errormess("Not the same latitudes !", [lat,lat0])
285       ## Faire d'autre checks sur les compatibilites entre fichiers!!
286##########################################################
287##########################################################
288##########################################################
289
290      all_varname[k] = varname
291      all_namefile[k] = namefile
292      all_time[k] = time
293      if var2: 
294         all_var2[k] = getfield(nc,var2)
295         if ((var2 in ['slopexy','SLOPEXY']) and (typefile in ['meso']) and (var2 not in nc.variables)): all_var2[k] = slopeamplitude(nc)
296      if winds: all_windu[k] = getfield(nc,uchar) ; all_windv[k] = getfield(nc,vchar)
297      #####################
298      ##### GETFIELDS #####
299      #####################
300      ##### SPECIFIC CASES
301      ##### 1. saturation temperature
302      if varname in ["temp","t","T_nadir_nit","T_nadir_day","temp_day","temp_night"] and tsat:
303          tt=getfield(nc,varname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
304          if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
305          else:                                 tinput=tt
306          all_var[k]=get_tsat(vert,tinput,zlon=lon,zlat=lat,zalt=vert,ztime=time)
307      ##### 2. wind amplitude
308      elif ((varname in ['UV','uv','uvmet']) and (typefile in ['meso']) and (varname not in nc.variables)):
309          all_var[k]=windamplitude(nc)
310      elif ((varname in ['slopexy','SLOPEXY']) and (typefile in ['meso']) and (varname not in nc.variables)):
311          all_var[k]=slopeamplitude(nc)
312      elif ((varname in ['DELTAT','deltat']) and (typefile in ['meso']) and (varname not in nc.variables)):
313          all_var[k]=deltat0t1(nc)
314      else:
315      ##### GENERAL STUFF HERE
316          all_var[k] = getfield(nc,varname)
317      print "**** GOT SUBDATA:",k," NAMEFILE:",namefile," VAR:",varname, var2 ; k += 1 ; firstfile = False
318      #### End of for namefile in namefiles
319
320    ##################################
321    ### Operation on files
322    if ope is not None:
323        print "********** OPERATION: ",ope
324        if "var" not in ope:
325             if len(var) > 1: errormess("for this operation... please set only one var !")
326             if ope in ["-","+","-%"]:
327                if fileref is not None:   
328                   ncref = Dataset(fileref)
329                   if ((all_varname[k-1] in ['UV','uv','uvmet']) and (typefile in ['meso']) and (all_varname[k-1] not in ncref.variables)):
330                      all_var[k] = windamplitude(ncref)
331                   elif ((all_varname[k-1] in ['slopexy','SLOPEXY']) and (typefile in ['meso']) and (all_varname[k-1] not in ncref.variables)):
332                      all_var[k] = slopeamplitude(ncref)
333                   elif ((all_varname[k-1] in ['DELTAT','deltat']) and (typefile in ['meso']) and (all_varname[k-1] not in ncref.variables)):
334                      all_var[k]=deltat0t1(ncref)
335                   else:
336                      all_var[k] = getfield(ncref,all_varname[k-1])
337                   all_varname[k] = all_varname[k-1] ; all_time[k] = all_time[k-1] ; all_namefile[k] = all_namefile[k-1] ; all_var2[k] = all_var2[k-1]
338                   if winds: all_windu[k] = getfield(ncref,uchar) ; all_windv[k] = getfield(ncref,vchar)
339                else: errormess("fileref is missing!") 
340                if ope == "-":     all_var[k+1]= all_var[k-1] - all_var[k]
341                elif ope == "+":   all_var[k+1]= all_var[k-1] + all_var[k]
342                elif ope == "-%":
343                    masked = np.ma.masked_where(all_var[k] == 0,all_var[k])
344                    masked.set_fill_value([np.NaN])
345                    all_var[k+1]= 100.*(all_var[k-1] - masked)/masked
346                all_varname[k+1] = all_varname[k] ; all_time[k+1] = all_time[k] ; all_namefile[k+1] = all_namefile[k] ; all_var2[k+1] = all_var2[k] ; numplot = numplot+2
347                if winds: all_windu[k+1] = all_windu[k-1]-all_windu[k] ; all_windv[k+1] = all_windv[k-1] - all_windv[k]
348             elif ope in ["cat"]:
349                tabtime = all_time[0];tab = all_var[0];k = 1
350                if var2: tab2 = all_var2[0]
351                while k != len(namefiles) and len(all_time[k]) != 0:
352                    if var2: tab2 = np.append(tab2,all_var2[k],axis=0) 
353                    tabtime = np.append(tabtime,all_time[k]) ; tab = np.append(tab,all_var[k],axis=0) ; k += 1
354                all_time[0] = np.array(tabtime) ; all_var[0] = np.array(tab) ; numplot = 1
355                if var2: all_var2[0] = np.array(tab2)
356             else: errormess(ope+" : non-implemented operation. Check pp.py --help")
357        else:
358             if len(namefiles) > 1: errormess("for this operation... please set only one file !") 
359             if len(var) > 2:       errormess("not sure this works for more than 2 vars... please check.")
360             if   ope in ["div_var"]: all_var[k] = all_var[k-2] / all_var[k-1] ; insert = '_div_'
361             elif ope in ["mul_var"]: all_var[k] = all_var[k-2] * all_var[k-1] ; insert = '_mul_'
362             elif ope in ["add_var"]: all_var[k] = all_var[k-2] + all_var[k-1] ; insert = '_add_'
363             elif ope in ["sub_var"]: all_var[k] = all_var[k-2] - all_var[k-1] ; insert = '_sub_'
364             else:                    errormess(ope+" : non-implemented operation. Check pp.py --help")
365             numplot = numplot + 1 ; all_time[k] = all_time[k-1] ; all_namefile[k] = all_namefile[k-1]
366             all_varname[k] = all_varname[k-2] + insert + all_varname[k-1] 
367    ##################################
368    ### Open a figure and set subplots
369    fig = figure()
370    subv,subh = definesubplot( numplot, fig ) 
371    if ope in ['-','-%']: subv,subh = 2,2
372 
373    #################################
374    ### Time loop for plotting device
375    nplot = 1 ; error = False 
376    if lstyle is not None: linecycler = cycle(lstyle)
377    else: linecycler = cycle(["-","--","-.",":"])
378    print "********************************************"
379    while error is False:
380     
381       print "********** NPLOT", nplot
382       if nplot > numplot: break
383
384       ####################################################################
385       ## get all indexes to be taken into account for this subplot and then reduce field
386       ## 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
387       indexlon  = getsindex(sslon,(nplot-1)%nlon,lon)
388       indexlat  = getsindex(sslat,((nplot-1)//nlon)%nlat,lat)
389       indexvert = getsindex(svert,((nplot-1)//(nlon*nlat))%nvert,vert) 
390       if ope is not None:
391           if fileref is not None:      index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(namefiles)+2)  ## OK only 1 var,  see test in the beginning
392           elif "var" in ope:           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(var)+1)        ## OK only 1 file, see test in the beginning
393           elif "cat" in ope:           index_f = 0
394       else:                            yeah = len(namefiles)*len(var) ; index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%yeah
395       time = all_time[index_f]
396       if mrate is not None:                 indextime = None 
397       else:                                 indextime = getsindex(stime,((nplot-1)//(nlon*nlat*nvert))%ntime,time)
398       ltst = None 
399       if typefile in ['meso'] and indextime is not None:  ltst = localtime ( interv[0]+indextime*interv[1], 0.5*(wlon[0]+wlon[1]) ) 
400       print "********** INDEX LON:",indexlon," LAT:",indexlat," VERT:",indexvert," TIME:",indextime
401       ##var = nc.variables["phisinit"][:,:]
402       ##contourf(np.transpose(var),30,cmap = get_cmap(name="Greys_r") ) ; axis('off') ; plot(indexlat,indexlon,'mx',mew=4.0,ms=20.0)
403       ##show()
404       ##exit()
405       #truc = True
406       #truc = False
407       #if truc: indexvert = None
408       ####################################################################
409       ########## REDUCE FIELDS
410       ####################################################################
411       error = False
412       varname = all_varname[index_f]
413       if varname:   ### what is shaded.
414           what_I_plot, error = reducefield( all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat, d3=indexvert, \
415                                             yint=yintegral, alt=vert, anomaly=anomaly, redope=redope, mesharea=area )
416           if mult != 2718.:  what_I_plot = what_I_plot*mult
417           else:              what_I_plot = np.log10(what_I_plot) ; print "log plot"
418                     
419       if var2:      ### what is contoured.
420           what_I_plot_contour, error = reducefield( all_var2[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert, \
421                                             yint=yintegral, alt=vert )
422       if winds:     ### what is plot as vectors.
423           vecx, error = reducefield( all_windu[index_f], d4=indextime, d3=indexvert, yint=yintegral, alt=vert)
424           vecy, error = reducefield( all_windv[index_f], d4=indextime, d3=indexvert, yint=yintegral, alt=vert)
425           if varname in [uchar,vchar]: what_I_plot = np.sqrt( np.square(vecx) + np.square(vecy) ) ; varname = "wind"
426
427       #####################################################################
428       #if truc:
429       #   nx = what_I_plot.shape[2] ; ny = what_I_plot.shape[1] ; nz = what_I_plot.shape[0]
430       #   for k in range(nz): print k,' over ',nz ; what_I_plot[k,:,:] = what_I_plot[k,:,:] / smooth(what_I_plot[k,:,:],12)
431       #   for iii in range(nx):
432       #    for jjj in range(ny):
433       #     deviation = what_I_plot[:,jjj,iii] ; mx = max(deviation) ; mn = min(deviation)
434       #     if iii > 6 and iii < nx-6 and jjj > 6 and jjj < ny-6:   what_I_plot[0,jjj,iii],rel = singlet(deviation,vert/1000.)  ### z must be in km
435       #     else:                                                   what_I_plot[0,jjj,iii]     = 0.
436       #     if np.abs(what_I_plot[0,jjj,iii]) > 1.5:
437       #         print iii,jjj,what_I_plot[0,jjj,iii],int(abs(1.-mx)*100.),int(abs(1.-mn)*100.)
438       #         plot(rel)
439       #   show()
440       #   anomaly = True ### pour avoir les bons reglages plots
441       #   what_I_plot = what_I_plot[0,:,:] 
442       #####################################################################
443
444       ####################################################################
445       ### General plot settings
446       changesubplot = (numplot > 1) and (len(what_I_plot.shape) != 1)  ## default for 1D plots: superimposed. to be reworked for better flexibility.
447       if changesubplot: subplot(subv,subh,nplot) #; subplots_adjust(wspace=0,hspace=0)
448       ####################################################################
449       if error:
450               errormess("There is an error in reducing field !")
451       else:
452               ticks = ndiv + 1 
453               fvar = varname
454               if anomaly: fvar = 'anomaly'
455               ###
456               if mapmode == 0:    ### could this be moved inside imov loop ?
457                   itime=indextime
458                   if len(what_I_plot.shape) == 3: itime=[0]
459                   m = None ; x = None ; y = None
460                   latyeah = lat ; lonyeah = lon
461                   if typefile in ['meso']:
462                       # now that the section is determined we can set the real lat
463                       # ... or for now, a temptative one.
464                       if slon is not None: latyeah = lat2d[:,0]
465                       if slat is not None: lonyeah = lon2d[0,:]
466                   what_I_plot, x, y = define_axis(lonyeah,latyeah,vert,time,indexlon,indexlat,indexvert,\
467                         itime,what_I_plot, len(all_var[index_f].shape),vertmode)
468               ###
469               if (fileref is not None) and (index_f == numplot-1):    zevmin, zevmax = calculate_bounds(what_I_plot,vmin=minop,vmax=maxop)
470               else:                                                   zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
471               if (fileref is not None) and (index_f == numplot-1):    colorb = "RdBu_r"
472               if colorb in ["def","nobar","onebar"]:                  palette = get_cmap(name=defcolorb(fvar.upper()))
473               else:                                                   palette = get_cmap(name=colorb)
474               #palette = cm.GMT_split
475               ##### 1. ELIMINATE 0D or >3D CASES
476               if len(what_I_plot.shape) == 0:   
477                 print "VALUE VALUE VALUE VALUE ::: ",what_I_plot
478                 save = 'donothing'
479               elif len(what_I_plot.shape) >= 4:
480                 print "WARNING!!! ",len(what_I_plot.shape),"-D PLOT NOT SUPPORTED !!! dimensions: ",what_I_plot.shape
481                 errormess("Are you sure you did not forget to prescribe a dimension ?")
482               ##### 2. HANDLE simple 1D/2D field and movies of 1D/2D fields
483               else:
484                 if mrate is not None: iend=len(time)-1
485                 else:                 iend=0
486                 imov = 0 
487                 if len(what_I_plot.shape) == 3:
488                    if var2:               which = "contour" ## have to start with contours rather than shading
489                    else:                  which = "regular"
490                    if mrate is None:      errormess("3D field. Use --rate RATE for movie or specify --time TIME. Exit.")
491                 elif len(what_I_plot.shape) == 2:
492                    if var2:               which = "contour" ## have to start with contours rather than shading
493                    else:                  which = "regular"
494                    if mrate is not None:  which = "unidim"
495                 elif len(what_I_plot.shape) == 1:
496                    which = "unidim"
497                    if what_I_plot.shape[-1] == 1:      print "VALUE VALUE VALUE VALUE ::: ", what_I_plot[0] ; save = 'donothing'
498                 ##### IMOV LOOP #### IMOV LOOP
499                 while imov <= iend:
500                    print "-> frame ",imov+1, which
501                    if which == "regular":   
502                        if mrate is None:                                   what_I_plot_frame = what_I_plot
503                        else:                                               what_I_plot_frame = what_I_plot[imov,:,:]
504                        if winds:
505                            if mrate is None:                                   vecx_frame = vecx ; vecy_frame = vecy
506                            else:                                               vecx_frame = vecx[imov,:,:] ; vecy_frame = vecy[imov,:,:]
507                    elif which == "contour": 
508                        if mrate is None or what_I_plot_contour.ndim < 3:   what_I_plot_frame = what_I_plot_contour
509                        else:                                               what_I_plot_frame = what_I_plot_contour[imov,:,:]
510                    elif which == "unidim":
511                        if mrate is None:                                   what_I_plot_frame = what_I_plot
512                        else:                                               what_I_plot_frame = what_I_plot[:,imov]  ## because swapaxes
513                    #if mrate is not None:     
514                    if mapmode == 1: 
515                            m = define_proj(proj,wlon,wlat,back=back,blat=blat,blon=blon)  ## this is dirty, defined above but out of imov loop
516                            x, y = m(lon2d, lat2d)                                         ## this is dirty, defined above but out of imov loop
517                    if typefile in ['meso'] and mapmode == 1:   what_I_plot_frame = dumpbdy(what_I_plot_frame,6,condition=True)
518#                   if typefile in ['mesoideal']:    what_I_plot_frame = dumpbdy(what_I_plot_frame,0,stag='W',condition=dumped_vert_stag)
519
520                    if which == "unidim":
521                        if lbls is not None: lbl=lbls[nplot-1] 
522                        else:
523                           lbl = ""
524                           if indexlat is not None:  lbl = lbl + " ix" + str(indexlat[0])
525                           if indexlon is not None:  lbl = lbl + " iy" + str(indexlon[0])
526                           if indexvert is not None: lbl = lbl + " iz" + str(indexvert[0])
527                           if indextime is not None: lbl = lbl + " it" + str(indextime[0])
528                           if lbl == "": lbl = all_namefile[index_f]
529
530                        if mrate is not None: x = y  ## because swapaxes...
531                        #what_I_plot_frame = np.diff(what_I_plot_frame, n=1) ; x = x[1:]
532                     
533                        zeline = next(linecycler) ## "-" for simple lines
534                        if tile:      zemarker = 'x'
535                        else:         zemarker = None 
536                        this_is_a_regular_plot = (indexvert is not None) or (indextime is None) or (indexlat is None) or (indexlon is None)
537                        if this_is_a_regular_plot:   plot(x,what_I_plot_frame,zeline,label=lbl,marker=zemarker)  ## vertical profile
538                        else:                        plot(what_I_plot_frame,x,zeline,label=lbl,marker=zemarker)  ## regular plot
539                        if nplot > 1: legend(loc='best')
540                        if indextime is None and axtime is not None and xlab is None:    xlabel(axtime.upper()) ## define the right label
541                        if save == 'txt':  writeascii(np.transpose([x,np.transpose(what_I_plot)]),'profile'+str(nplot*1000+imov)+'.txt')
542
543                    elif which == "regular":
544                   
545                        # plot stream lines if there is a stream file and a vert/lat slice. Might not work with movies ??
546                        if streamflag and sslat is None and svert is None:
547                             streamfile = all_namefile[index_f].replace('.nc','_stream.nc')
548                             teststream = os.path.exists(streamfile)
549                             if teststream:
550                                print 'INFO: Using stream file',streamfile, 'for stream lines'
551                                ncstream = Dataset(streamfile)
552                                psi = getfield(ncstream,'psi')
553                                psi = psi[0,:,:,0] # time and longitude are dummy dimensions
554                                if psi.shape[1] != len(x) or psi.shape[0] != len(y):
555                                    errormess('stream file does not fit! Dimensions: '+str(psi.shape)+' '+str(x.shape)+' '+str(y.shape))
556                                zelevels = np.arange(-1.e10,1.e10,1.e9)
557                                zemin = np.min(abs(zelevels))
558                                zemax = np.max(abs(zelevels))
559                                zewidth  =  (abs(zelevels)-zemin)*(5.- 0.5)/(zemax - zemin) + 0.5 # linewidth ranges from 5 to 0.5
560                                cs = contour( x,y,psi, zelevels, colors='k', linewidths = zewidth)
561                                clabel(cs, inline=True, fontsize = 4.*rcParams['font.size']/5., fmt="%1.1e")
562                             else:
563                                print 'WARNING: STREAM FILE',streamfile, 'DOES NOT EXIST !'
564                             
565                        if hole:         what_I_plot_frame = hole_bounds(what_I_plot_frame,zevmin,zevmax)
566                        else:            what_I_plot_frame = bounds(what_I_plot_frame,zevmin,zevmax)
567                        if flagnolow:    what_I_plot_frame = nolow(what_I_plot_frame)
568                        if not tile:
569                            #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
570                            zelevels = np.linspace(zevmin,zevmax,num=ticks)
571                            #what_I_plot_frame = smooth(what_I_plot_frame,100)
572                            if mapmode == 1:       m.contourf( x, y, what_I_plot_frame, zelevels, cmap = palette, alpha=trans)
573                            elif mapmode == 0:     contourf( x, y, what_I_plot_frame, zelevels, cmap = palette, alpha=trans)
574                            if cross is not None and mapmode == 1:
575                               idx,idy=m(cross[0][0],cross[0][1])
576                               mpl.pyplot.plot([idx],[idy],'+k',mew=0.5,ms=18)
577                        else:
578                            if mapmode == 1:       m.pcolor( x, y, what_I_plot_frame, cmap = palette, vmin=zevmin, vmax=zevmax, alpha=trans)
579                            elif mapmode == 0:     pcolor( x, y, what_I_plot_frame, cmap = palette, vmin=zevmin, vmax=zevmax, alpha=trans)
580                       
581
582                        if colorb not in ['nobar','onebar']:       
583                            if (fileref is not None) and (index_f == numplot-1):   daformat = "%.3f" 
584                            elif mult != 1:                                        daformat = "%.1f"
585                            else:                                                  daformat = fmtvar(fvar.upper())
586                            #if proj in ['moll','cyl']:  zeorientation="horizontal" ; zepad = 0.05
587                            if proj in ['moll']:        zeorientation="horizontal" ; zepad = 0.05
588                            else:                       zeorientation="vertical" ; zepad = 0.03
589                            zecb = colorbar( fraction=0.05,pad=zepad,format=daformat,orientation=zeorientation,\
590                                      ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional' ) 
591                            if zeorientation == "horizontal" and zetitle[0] != "fill": zecb.ax.set_xlabel(zetitle[index_f]) ; zetitle[index_f]=""
592                        if winds:
593                            if typefile in ['meso']:
594                                [vecx_frame,vecy_frame] = [dumpbdy(vecx_frame,6,stag=uchar,condition=True), dumpbdy(vecy_frame,6,stag=vchar,condition=True)]
595                                key = True
596                                if fvar in ['UV','uv','uvmet']: key = False
597                            elif typefile in ['gcm']:
598                                key = False
599                            if metwind and mapmode == 1:   [vecx_frame,vecy_frame] = m.rotate_vector(vecx_frame, vecy_frame, lon2d, lat2d)
600                            if var:       colorvec = definecolorvec(colorb) 
601                            else:         colorvec = definecolorvec(back) 
602                            vectorfield(vecx_frame, vecy_frame, x, y, stride=stride, csmooth=2,\
603                                             #scale=15., factor=300., color=colorvec, key=key)
604                                             scale=20., factor=250., color=colorvec, key=key)
605                                                              #200.         ## or csmooth=stride
606                        if ope == '-' and nplot == numplot: # this should work as long as ope is '-' guarantees 3 plots for 4 panels without contour
607                            subplot(subv,subh,nplot+1)
608                            rcParams["legend.fontsize"] = 'xx-large'
609                            if indexlat is None:
610                                latmin = -50.; latmax = 50. # latitude range for histogram of difference
611                                zeindexlat = (lat<latmax)*(lat>latmin)
612                                # this follows the define_axis logic in myplot.py:
613                                if indextime is None or indexlon is None: what_I_plot_frame = what_I_plot_frame[zeindexlat,:]
614                                else: what_I_plot_frame = what_I_plot_frame[:,zeindexlat]
615                            toplot = np.ravel(what_I_plot_frame[np.isnan(what_I_plot_frame)==False])
616                            zebins = np.linspace(zevmin,zevmax,num=30)
617                            hist(toplot,bins=zebins,histtype='step',linewidth=2,color='k',normed=True)
618                            zestd = np.std(toplot);zemean = np.mean(toplot)
619                            zebins = np.linspace(zevmin,zevmax,num=300)
620                            zegauss = (1./(zestd * np.sqrt(2 * np.pi)) * np.exp( - (zebins - zemean)**2 / (2 * zestd**2) ) )
621                            plot(zebins, zegauss, linewidth=1, color='r',label="mean: "+str(zemean)[0:5]+"\nstd: "+str(zestd)[0:5])
622                            legend(loc=0,frameon=False)
623                            title("Histogram fig(1) "+ope+" fig(2)")
624                            subplot(subv,subh,nplot) # go back to last plot for title of contour difference
625                           
626                    elif which == "contour":
627                        zevminc, zevmaxc = calculate_bounds(what_I_plot_frame, vmin=min(what_I_plot_frame), vmax=max(what_I_plot_frame))
628                        zelevels = np.linspace(zevminc,zevmaxc,ticks/2) #20)
629                        if var2 == 'HGT':     zelevels = np.arange(-10000.,30000.,1000.)
630                        elif var2 == 'tpot':  zelevels = np.arange(270,370,5)
631                        elif var2 == 'tk':    zelevels = np.arange(150,250,5)
632                        if mapmode == 0:   
633                            what_I_plot_frame, x, y = define_axis( lonyeah,latyeah,vert,time,indexlon,indexlat,indexvert,\
634                                                              itime,what_I_plot_frame, len(all_var2[index_f].shape),vertmode )
635                            ## is this needed? only if len(all_var2[index_f].shape) != len(all_var[index_f].shape)
636                            cs = contour( x,y,what_I_plot_frame, zelevels, colors='k', linewidths = 0.33)#, alpha=0.5, linestyles='solid')
637                            clabel(cs, inline=1, fontsize = 4.*rcParams['font.size']/5., fmt=fmtvar(var2.upper()))
638                        elif mapmode == 1:  cs = m.contour( x,y,what_I_plot_frame, zelevels, colors='k', linewidths = 0.33)#, alpha=0.5, linestyles='solid')
639                    if which in ["regular","unidim"]:
640
641                        if nplot > 1 and which == "unidim":
642                           pass  ## because we superimpose nplot instances
643                        else:
644                           # Axis directives for movie frames [including the first one).
645                           zxmin, zxmax = xaxis ; zymin, zymax = yaxis
646                           if zxmin is not None: mpl.pyplot.xlim(xmin=zxmin)
647                           if zxmax is not None: mpl.pyplot.xlim(xmax=zxmax)
648                           if zymin is not None: mpl.pyplot.ylim(ymin=zymin)
649                           if zymax is not None: mpl.pyplot.ylim(ymax=zymax)
650                           if ylog:      mpl.pyplot.semilogy()
651                           if invert_y:  ax = mpl.pyplot.gca() ; ax.set_ylim(ax.get_ylim()[::-1])
652                           if xlab is not None: xlabel(xlab)
653                           if ylab is not None: ylabel(ylab)
654
655                        if mrate is not None:
656                           ### THIS IS A MENCODER MOVIE
657                           if mrate > 0:
658                             figframe=mpl.pyplot.gcf()
659                             if mquality:   figframe.set_dpi(600.)
660                             else:          figframe.set_dpi(200.)
661                             mframe=fig2img(figframe)
662                             if imov == 0:
663                                moviename='movie' ;W,H = figframe.canvas.get_width_height()
664                                video = VideoSink((H,W), moviename, rate=mrate, byteorder="rgba")
665                             video.run(mframe) ; close()
666                             if imov == iend: video.close()                           
667                           ### THIS IS A WEBPAGE MOVIE
668                           else:
669                             nameframe = "image"+str(1000+imov)
670                             makeplotres(nameframe,res=100.,disp=False) ; close()
671                             if imov == 0: myfile = open("zepics", 'w')
672                             myfile.write("modImages["+str(imov)+"] = '"+nameframe+"_100.png';"+ '\n')
673                             if imov == iend:
674                                 myfile.write("first_image = 0;"+ '\n')
675                                 myfile.write("last_image = "+str(iend)+";"+ '\n')
676                                 myfile.close()
677                        if var2 and which == "regular":  which = "contour"
678                        imov = imov+1
679                    elif which == "contour":
680                        which = "regular"
681
682       ### Next subplot
683       zevarname = varname
684       if redope is not None: zevarname = zevarname + "_" + redope
685       basename = getname(var=zevarname,var2=var2,winds=winds,anomaly=anomaly)
686       if len(what_I_plot.shape) > 3:
687           basename = basename + getstralt(nc,level) 
688       if mrate is not None: basename = "movie_" + basename
689       if typefile in ['meso']:
690            if sslon is not None: basename = basename + "_lon_" + str(int(round(lonp)))
691            if sslat is not None: basename = basename + "_lat_" + str(int(round(latp)))
692            plottitle = basename
693            ### dans le nouveau systeme time=ls,sol,lt cette ligne pourrait ne servir a rien (ou deplacer au dessus)
694            if addchar and indextime is not None:   [addchar,gogol,gogol2] = getlschar ( all_namefile[index_f] )  ;  plottitle = plottitle + addchar
695            ### en fait redope is None doit etre remplace par : n'est ni maxt ni mint
696            if redope is None and ltst is not None and ( (mapmode == 0) or (proj in ["lcc","laea","merc","nsper"]) ):  plottitle = plottitle + "_LT" + str(ltst)
697       else:
698            if fileref is not None:
699                if index_f is numplot-1:     plottitle = basename+' '+"fig(1) "+ope+" fig(2)"
700                elif index_f is numplot-2:   plottitle = basename+' '+fileref
701                else:                        plottitle = basename+' '+all_namefile[index_f]
702            else:                            plottitle = basename+' '+all_namefile[index_f]
703       if mult != 1:                         plottitle = '{:.0e}'.format(mult) + "*" + plottitle
704       if zetitle[0] != "fill":                 
705          if index_f < len(zetitle): plottitle = zetitle[index_f]
706          else: plottitle = zetitle[0]
707          if titleref is "fill":             titleref=zetitle[index_f]
708          if fileref is not None:
709             if index_f is numplot-2:        plottitle = titleref
710             if index_f is numplot-1:        plottitle = "fig(1) "+ope+" fig(2)"
711#       if indexlon is not None:      plottitle = plottitle + " lon: " + str(min(lon[indexlon])) +" "+ str(max(lon[indexlon]))
712#       if indexlat is not None:      plottitle = plottitle + " lat: " + str(min(lat[indexlat])) +" "+ str(max(lat[indexlat]))
713#       if indexvert is not None:     plottitle = plottitle + " vert: " + str(min(vert[indexvert])) +" "+ str(max(vert[indexvert]))
714#       if indextime is not None:     plottitle = plottitle + " time: " + str(min(time[indextime])) +" "+ str(max(time[indextime]))
715       if colorb != "onebar": title( plottitle )
716       if nplot >= numplot: error = True
717       nplot += 1
718
719    if colorb == "onebar":
720        cax = axes([0.1, 0.2, 0.8, 0.03]) # a ameliorer
721        zecb = colorbar(cax=cax, orientation="horizontal", format=fmtvar(fvar.upper()),\
722                 ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional')
723        if zetitle[0] != "fill": zecb.ax.set_xlabel(zetitle[index_f]) ; zetitle[index_f]=""
724
725
726    ##########################################################################
727    ### Save the figure in a file in the data folder or an user-defined folder
728    if outputname is None:
729       if typefile in ['meso']:   prefix = getprefix(nc)
730       elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
731       else:                                prefix = ''
732    ###
733       zeplot = prefix + basename
734       if addchar:         zeplot = zeplot + addchar
735       if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
736    ###
737       if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
738       else:               zeplot = target + "/" + zeplot 
739    ###
740    else:
741       zeplot=outputname
742
743    if mrate is None:
744        pad_inches_value = 0.35
745        print "********** SAVE ", save
746        if save == 'png': 
747            if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
748            makeplotres(zeplot,res=resolution,pad_inches_value=pad_inches_value,disp=False)
749        elif save in ['eps','svg','pdf']:     makeplotres(zeplot,pad_inches_value=pad_inches_value,disp=False,ext=save)
750        elif save == 'gui':                   show()
751        elif save == 'donothing':             pass
752        elif save == 'txt':                   print "Saved results in txt file." 
753        else: 
754            print "INFO: save mode not supported. using gui instead."
755            show()
756
757    ###################################
758    #### Getting more out of this video -- PROBLEMS WITH CREATED VIDEOS
759    #
760    #if mrate is not None:
761    #    print "Re-encoding movie.. first pass"
762    #    video.first_pass(filename=moviename,quality=mquality,rate=mrate)
763    #    print "Re-encoding movie.. second pass"
764    #    video.second_pass(filename=moviename,quality=mquality,rate=mrate)   
765
766    ###############
767    ### Now the end
768    return zeplot
Note: See TracBrowser for help on using the repository browser.