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

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

use of python function instead of fortran code for GCM ls axtime. No need to compile timestuff anymore

  • Property svn:executable set to *
File size: 47.6 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 times 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])
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 ope in ["cat"] and nnn > 0:    count = time[-1] + 1  ## so that a cat is possible with simple subscripts
257              else:                             count = 0
258              if "Times" in nc.dimensions:   time = count + np.arange(0,len(nc.dimensions["Times"]),1)
259              elif "Time" in nc.dimensions:  time = count + np.arange(0,len(nc.dimensions["Time"]),1)
260              else:                          time = count + np.arange(0,1,1)
261          if axtime in ["lt"]:
262              ftime = np.zeros(len(time))
263              for i in range(len(time)): 
264                 ftime[i] = localtime ( interv[0]+time[i]*interv[1], 0.5*(wlon[0]+wlon[1]) )
265              time=ftime
266              time=check_localtime(time)
267              print "LOCAL TIMES.... ", time
268
269          ###
270          if typefile in ['geo']:   vert = [0.] ; stime = readslices(str(0))
271          else:
272              if vertmode is None:  vertmode=0
273              if vertmode == 0:     vert = np.arange(0,getattr(nc,vertdim),1)
274              elif vertmode == -1:  vert = nc.variables["PHTOT"][0,:,0,0]/3.72 ; vert = np.array(vert[0:len(vert)-1]) #; print vert
275              elif vertmode == 1 or vertmode == 2:  vert = nc.variables["vert"][:]        ## pressure in Pa
276              else:                                 vert = nc.variables["vert"][:]/1000.  ## altitude in km
277       #if firstfile:
278       #   lat0 = lat
279       #elif len(lat0) != len(lat):
280       #   errormess("Not the same latitude lengths !", [len(lat0), len(lat)])
281       #elif sum((lat == lat0) == False) != 0:
282       #   errormess("Not the same latitudes !", [lat,lat0])
283       ## Faire d'autre checks sur les compatibilites entre fichiers!!
284##########################################################
285##########################################################
286##########################################################
287
288      all_varname[k] = varname
289      all_namefile[k] = namefile
290      all_time[k] = time
291      if var2: 
292         all_var2[k] = getfield(nc,var2)
293         if ((var2 in ['slopexy','SLOPEXY']) and (typefile in ['meso']) and (var2 not in nc.variables)): all_var2[k] = slopeamplitude(nc)
294      if winds: all_windu[k] = getfield(nc,uchar) ; all_windv[k] = getfield(nc,vchar)
295      #####################
296      ##### GETFIELDS #####
297      #####################
298      ##### SPECIFIC CASES
299      ##### 1. saturation temperature
300      if varname in ["temp","t","T_nadir_nit","T_nadir_day","temp_day","temp_night"] and tsat:
301          tt=getfield(nc,varname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
302          if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
303          else:                                 tinput=tt
304          all_var[k]=get_tsat(vert,tinput,zlon=lon,zlat=lat,zalt=vert,ztime=time)
305      ##### 2. wind amplitude
306      elif ((varname in ['UV','uv','uvmet']) and (typefile in ['meso']) and (varname not in nc.variables)):
307          all_var[k]=windamplitude(nc)
308      elif ((varname in ['slopexy','SLOPEXY']) and (typefile in ['meso']) and (varname not in nc.variables)):
309          all_var[k]=slopeamplitude(nc)
310      elif ((varname in ['DELTAT','deltat']) and (typefile in ['meso']) and (varname not in nc.variables)):
311          all_var[k]=deltat0t1(nc)
312      else:
313      ##### GENERAL STUFF HERE
314          all_var[k] = getfield(nc,varname)
315      print "**** GOT SUBDATA:",k," NAMEFILE:",namefile," VAR:",varname, var2 ; k += 1 ; firstfile = False
316      #### End of for namefile in namefiles
317
318    ##################################
319    ### Operation on files
320    if ope is not None:
321        print "********** OPERATION: ",ope
322        if "var" not in ope:
323             if len(var) > 1: errormess("for this operation... please set only one var !")
324             if ope in ["-","+","-%"]:
325                if fileref is not None:   
326                   ncref = Dataset(fileref)
327                   if ((all_varname[k-1] in ['UV','uv','uvmet']) and (typefile in ['meso']) and (all_varname[k-1] not in ncref.variables)):
328                      all_var[k] = windamplitude(ncref)
329                   elif ((all_varname[k-1] in ['slopexy','SLOPEXY']) and (typefile in ['meso']) and (all_varname[k-1] not in ncref.variables)):
330                      all_var[k] = slopeamplitude(ncref)
331                   elif ((all_varname[k-1] in ['DELTAT','deltat']) and (typefile in ['meso']) and (all_varname[k-1] not in ncref.variables)):
332                      all_var[k]=deltat0t1(ncref)
333                   else:
334                      all_var[k] = getfield(ncref,all_varname[k-1])
335                   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]
336                   if winds: all_windu[k] = getfield(ncref,uchar) ; all_windv[k] = getfield(ncref,vchar)
337                else: errormess("fileref is missing!") 
338                if ope == "-":     all_var[k+1]= all_var[k-1] - all_var[k]
339                elif ope == "+":   all_var[k+1]= all_var[k-1] + all_var[k]
340                elif ope == "-%":
341                    masked = np.ma.masked_where(all_var[k] == 0,all_var[k])
342                    masked.set_fill_value([np.NaN])
343                    all_var[k+1]= 100.*(all_var[k-1] - masked)/masked
344                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
345                if winds: all_windu[k+1] = all_windu[k-1]-all_windu[k] ; all_windv[k+1] = all_windv[k-1] - all_windv[k]
346             elif ope in ["cat"]:
347                tabtime = all_time[0];tab = all_var[0];k = 1
348                if var2: tab2 = all_var2[0]
349                while k != len(namefiles) and len(all_time[k]) != 0:
350                    if var2: tab2 = np.append(tab2,all_var2[k],axis=0) 
351                    tabtime = np.append(tabtime,all_time[k]) ; tab = np.append(tab,all_var[k],axis=0) ; k += 1
352                all_time[0] = np.array(tabtime) ; all_var[0] = np.array(tab) ; numplot = 1
353                if var2: all_var2[0] = np.array(tab2)
354             else: errormess(ope+" : non-implemented operation. Check pp.py --help")
355        else:
356             if len(namefiles) > 1: errormess("for this operation... please set only one file !") 
357             if len(var) > 2:       errormess("not sure this works for more than 2 vars... please check.")
358             if   ope in ["div_var"]: all_var[k] = all_var[k-2] / all_var[k-1] ; insert = '_div_'
359             elif ope in ["mul_var"]: all_var[k] = all_var[k-2] * all_var[k-1] ; insert = '_mul_'
360             elif ope in ["add_var"]: all_var[k] = all_var[k-2] + all_var[k-1] ; insert = '_add_'
361             elif ope in ["sub_var"]: all_var[k] = all_var[k-2] - all_var[k-1] ; insert = '_sub_'
362             else:                    errormess(ope+" : non-implemented operation. Check pp.py --help")
363             numplot = numplot + 1 ; all_time[k] = all_time[k-1] ; all_namefile[k] = all_namefile[k-1]
364             all_varname[k] = all_varname[k-2] + insert + all_varname[k-1] 
365    ##################################
366    ### Open a figure and set subplots
367    fig = figure()
368    subv,subh = definesubplot( numplot, fig ) 
369    if ope in ['-','-%']: subv,subh = 2,2
370 
371    #################################
372    ### Time loop for plotting device
373    nplot = 1 ; error = False 
374    if lstyle is not None: linecycler = cycle(lstyle)
375    else: linecycler = cycle(["-","--","-.",":"])
376    print "********************************************"
377    while error is False:
378     
379       print "********** NPLOT", nplot
380       if nplot > numplot: break
381
382       ####################################################################
383       ## get all indexes to be taken into account for this subplot and then reduce field
384       ## 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
385       indexlon  = getsindex(sslon,(nplot-1)%nlon,lon)
386       indexlat  = getsindex(sslat,((nplot-1)//nlon)%nlat,lat)
387       indexvert = getsindex(svert,((nplot-1)//(nlon*nlat))%nvert,vert) 
388       if ope is not None:
389           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
390           elif "var" in ope:           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(var)+1)        ## OK only 1 file, see test in the beginning
391           elif "cat" in ope:           index_f = 0
392       else:                            yeah = len(namefiles)*len(var) ; index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%yeah
393       time = all_time[index_f]
394       if mrate is not None:                 indextime = None 
395       else:                                 indextime = getsindex(stime,((nplot-1)//(nlon*nlat*nvert))%ntime,time)
396       ltst = None 
397       if typefile in ['meso'] and indextime is not None:  ltst = localtime ( interv[0]+indextime*interv[1], 0.5*(wlon[0]+wlon[1]) ) 
398       print "********** INDEX LON:",indexlon," LAT:",indexlat," VERT:",indexvert," TIME:",indextime
399       ##var = nc.variables["phisinit"][:,:]
400       ##contourf(np.transpose(var),30,cmap = get_cmap(name="Greys_r") ) ; axis('off') ; plot(indexlat,indexlon,'mx',mew=4.0,ms=20.0)
401       ##show()
402       ##exit()
403       #truc = True
404       #truc = False
405       #if truc: indexvert = None
406       ####################################################################
407       ########## REDUCE FIELDS
408       ####################################################################
409       error = False
410       varname = all_varname[index_f]
411       if varname:   ### what is shaded.
412           what_I_plot, error = reducefield( all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat, d3=indexvert, \
413                                             yint=yintegral, alt=vert, anomaly=anomaly, redope=redope, mesharea=area )
414           if mult != 2718.:  what_I_plot = what_I_plot*mult
415           else:              what_I_plot = np.log10(what_I_plot) ; print "log plot"
416                     
417       if var2:      ### what is contoured.
418           what_I_plot_contour, error = reducefield( all_var2[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert, \
419                                             yint=yintegral, alt=vert )
420       if winds:     ### what is plot as vectors.
421           vecx, error = reducefield( all_windu[index_f], d4=indextime, d3=indexvert, yint=yintegral, alt=vert)
422           vecy, error = reducefield( all_windv[index_f], d4=indextime, d3=indexvert, yint=yintegral, alt=vert)
423           if varname in [uchar,vchar]: what_I_plot = np.sqrt( np.square(vecx) + np.square(vecy) ) ; varname = "wind"
424
425       #####################################################################
426       #if truc:
427       #   nx = what_I_plot.shape[2] ; ny = what_I_plot.shape[1] ; nz = what_I_plot.shape[0]
428       #   for k in range(nz): print k,' over ',nz ; what_I_plot[k,:,:] = what_I_plot[k,:,:] / smooth(what_I_plot[k,:,:],12)
429       #   for iii in range(nx):
430       #    for jjj in range(ny):
431       #     deviation = what_I_plot[:,jjj,iii] ; mx = max(deviation) ; mn = min(deviation)
432       #     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
433       #     else:                                                   what_I_plot[0,jjj,iii]     = 0.
434       #     if np.abs(what_I_plot[0,jjj,iii]) > 1.5:
435       #         print iii,jjj,what_I_plot[0,jjj,iii],int(abs(1.-mx)*100.),int(abs(1.-mn)*100.)
436       #         plot(rel)
437       #   show()
438       #   anomaly = True ### pour avoir les bons reglages plots
439       #   what_I_plot = what_I_plot[0,:,:] 
440       #####################################################################
441
442       ####################################################################
443       ### General plot settings
444       changesubplot = (numplot > 1) and (len(what_I_plot.shape) != 1)  ## default for 1D plots: superimposed. to be reworked for better flexibility.
445       if changesubplot: subplot(subv,subh,nplot) #; subplots_adjust(wspace=0,hspace=0)
446       ####################################################################
447       if error:
448               errormess("There is an error in reducing field !")
449       else:
450               ticks = ndiv + 1 
451               fvar = varname
452               if anomaly: fvar = 'anomaly'
453               ###
454               if mapmode == 0:    ### could this be moved inside imov loop ?
455                   itime=indextime
456                   if len(what_I_plot.shape) == 3: itime=[0]
457                   m = None ; x = None ; y = None
458                   latyeah = lat ; lonyeah = lon
459                   if typefile in ['meso']:
460                       # now that the section is determined we can set the real lat
461                       # ... or for now, a temptative one.
462                       if slon is not None: latyeah = lat2d[:,0]
463                       if slat is not None: lonyeah = lon2d[0,:]
464                   what_I_plot, x, y = define_axis(lonyeah,latyeah,vert,time,indexlon,indexlat,indexvert,\
465                         itime,what_I_plot, len(all_var[index_f].shape),vertmode)
466               ###
467               if (fileref is not None) and (index_f == numplot-1):    zevmin, zevmax = calculate_bounds(what_I_plot,vmin=minop,vmax=maxop)
468               else:                                                   zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
469               if (fileref is not None) and (index_f == numplot-1):    colorb = "RdBu_r"
470               if colorb in ["def","nobar","onebar"]:                  palette = get_cmap(name=defcolorb(fvar.upper()))
471               else:                                                   palette = get_cmap(name=colorb)
472               #palette = cm.GMT_split
473               ##### 1. ELIMINATE 0D or >3D CASES
474               if len(what_I_plot.shape) == 0:   
475                 print "VALUE VALUE VALUE VALUE ::: ",what_I_plot
476                 save = 'donothing'
477               elif len(what_I_plot.shape) >= 4:
478                 print "WARNING!!! ",len(what_I_plot.shape),"-D PLOT NOT SUPPORTED !!! dimensions: ",what_I_plot.shape
479                 errormess("Are you sure you did not forget to prescribe a dimension ?")
480               ##### 2. HANDLE simple 1D/2D field and movies of 1D/2D fields
481               else:
482                 if mrate is not None: iend=len(time)-1
483                 else:                 iend=0
484                 imov = 0 
485                 if len(what_I_plot.shape) == 3:
486                    if var2:               which = "contour" ## have to start with contours rather than shading
487                    else:                  which = "regular"
488                    if mrate is None:      errormess("3D field. Use --rate RATE for movie or specify --time TIME. Exit.")
489                 elif len(what_I_plot.shape) == 2:
490                    if var2:               which = "contour" ## have to start with contours rather than shading
491                    else:                  which = "regular"
492                    if mrate is not None:  which = "unidim"
493                 elif len(what_I_plot.shape) == 1:
494                    which = "unidim"
495                    if what_I_plot.shape[-1] == 1:      print "VALUE VALUE VALUE VALUE ::: ", what_I_plot[0] ; save = 'donothing'
496                 ##### IMOV LOOP #### IMOV LOOP
497                 while imov <= iend:
498                    print "-> frame ",imov+1, which
499                    if which == "regular":   
500                        if mrate is None:                                   what_I_plot_frame = what_I_plot
501                        else:                                               what_I_plot_frame = what_I_plot[imov,:,:]
502                        if winds:
503                            if mrate is None:                                   vecx_frame = vecx ; vecy_frame = vecy
504                            else:                                               vecx_frame = vecx[imov,:,:] ; vecy_frame = vecy[imov,:,:]
505                    elif which == "contour": 
506                        if mrate is None or what_I_plot_contour.ndim < 3:   what_I_plot_frame = what_I_plot_contour
507                        else:                                               what_I_plot_frame = what_I_plot_contour[imov,:,:]
508                    elif which == "unidim":
509                        if mrate is None:                                   what_I_plot_frame = what_I_plot
510                        else:                                               what_I_plot_frame = what_I_plot[:,imov]  ## because swapaxes
511                    #if mrate is not None:     
512                    if mapmode == 1: 
513                            m = define_proj(proj,wlon,wlat,back=back,blat=blat,blon=blon)  ## this is dirty, defined above but out of imov loop
514                            x, y = m(lon2d, lat2d)                                         ## this is dirty, defined above but out of imov loop
515                    if typefile in ['meso'] and mapmode == 1:   what_I_plot_frame = dumpbdy(what_I_plot_frame,6,condition=True)
516#                   if typefile in ['mesoideal']:    what_I_plot_frame = dumpbdy(what_I_plot_frame,0,stag='W',condition=dumped_vert_stag)
517
518                    if which == "unidim":
519                        if lbls is not None: lbl=lbls[nplot-1] 
520                        else:
521                           lbl = ""
522                           if indexlat is not None:  lbl = lbl + " ix" + str(indexlat[0])
523                           if indexlon is not None:  lbl = lbl + " iy" + str(indexlon[0])
524                           if indexvert is not None: lbl = lbl + " iz" + str(indexvert[0])
525                           if indextime is not None: lbl = lbl + " it" + str(indextime[0])
526                           if lbl == "": lbl = all_namefile[index_f]
527
528                        if mrate is not None: x = y  ## because swapaxes...
529                        #what_I_plot_frame = np.diff(what_I_plot_frame, n=1) ; x = x[1:]
530                     
531                        zeline = next(linecycler) ## "-" for simple lines
532                        if tile:      zemarker = 'x'
533                        else:         zemarker = None 
534                        this_is_a_regular_plot = (indexvert is not None) or (indextime is None) or (indexlat is None) or (indexlon is None)
535                        if this_is_a_regular_plot:   plot(x,what_I_plot_frame,zeline,label=lbl,marker=zemarker)  ## vertical profile
536                        else:                        plot(what_I_plot_frame,x,zeline,label=lbl,marker=zemarker)  ## regular plot
537                        if nplot > 1: legend(loc='best')
538                        if indextime is None and axtime is not None and xlab is None:    xlabel(axtime.upper()) ## define the right label
539                        if save == 'txt':  writeascii(np.transpose([x,np.transpose(what_I_plot)]),'profile'+str(nplot*1000+imov)+'.txt')
540
541                    elif which == "regular":
542                   
543                        # plot stream lines if there is a stream file and a vert/lat slice. Might not work with movies ??
544                        if streamflag and sslat is None and svert is None:
545                             streamfile = all_namefile[index_f].replace('.nc','_stream.nc')
546                             teststream = os.path.exists(streamfile)
547                             if teststream:
548                                print 'INFO: Using stream file',streamfile, 'for stream lines'
549                                ncstream = Dataset(streamfile)
550                                psi = getfield(ncstream,'psi')
551                                psi = psi[0,:,:,0] # time and longitude are dummy dimensions
552                                if psi.shape[1] != len(x) or psi.shape[0] != len(y):
553                                    errormess('stream file does not fit! Dimensions: '+str(psi.shape)+' '+str(x.shape)+' '+str(y.shape))
554                                zelevels = np.arange(-1.e10,1.e10,1.e9)
555                                zemin = np.min(abs(zelevels))
556                                zemax = np.max(abs(zelevels))
557                                zewidth  =  (abs(zelevels)-zemin)*(5.- 0.5)/(zemax - zemin) + 0.5 # linewidth ranges from 5 to 0.5
558                                cs = contour( x,y,psi, zelevels, colors='k', linewidths = zewidth)
559                                clabel(cs, inline=True, fontsize = 4.*rcParams['font.size']/5., fmt="%1.1e")
560                             else:
561                                print 'WARNING: STREAM FILE',streamfile, 'DOES NOT EXIST !'
562                             
563                        if hole:         what_I_plot_frame = hole_bounds(what_I_plot_frame,zevmin,zevmax)
564                        else:            what_I_plot_frame = bounds(what_I_plot_frame,zevmin,zevmax)
565                        if flagnolow:    what_I_plot_frame = nolow(what_I_plot_frame)
566                        if not tile:
567                            #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
568                            zelevels = np.linspace(zevmin,zevmax,num=ticks)
569                            #what_I_plot_frame = smooth(what_I_plot_frame,100)
570                            if mapmode == 1:       m.contourf( x, y, what_I_plot_frame, zelevels, cmap = palette, alpha=trans)
571                            elif mapmode == 0:     contourf( x, y, what_I_plot_frame, zelevels, cmap = palette, alpha=trans)
572                            if cross is not None and mapmode == 1:
573                               idx,idy=m(cross[0][0],cross[0][1])
574                               mpl.pyplot.plot([idx],[idy],'+k',mew=0.5,ms=18)
575                        else:
576                            if mapmode == 1:       m.pcolor( x, y, what_I_plot_frame, cmap = palette, vmin=zevmin, vmax=zevmax, alpha=trans)
577                            elif mapmode == 0:     pcolor( x, y, what_I_plot_frame, cmap = palette, vmin=zevmin, vmax=zevmax, alpha=trans)
578                       
579
580                        if colorb not in ['nobar','onebar']:       
581                            if (fileref is not None) and (index_f == numplot-1):   daformat = "%.3f" 
582                            elif mult != 1:                                        daformat = "%.1f"
583                            else:                                                  daformat = fmtvar(fvar.upper())
584                            #if proj in ['moll','cyl']:  zeorientation="horizontal" ; zepad = 0.05
585                            if proj in ['moll']:        zeorientation="horizontal" ; zepad = 0.05
586                            else:                       zeorientation="vertical" ; zepad = 0.03
587                            zecb = colorbar( fraction=0.05,pad=zepad,format=daformat,orientation=zeorientation,\
588                                      ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional' ) 
589                            if zeorientation == "horizontal" and zetitle[0] != "fill": zecb.ax.set_xlabel(zetitle[index_f]) ; zetitle[index_f]=""
590                        if winds:
591                            if typefile in ['meso']:
592                                [vecx_frame,vecy_frame] = [dumpbdy(vecx_frame,6,stag=uchar,condition=True), dumpbdy(vecy_frame,6,stag=vchar,condition=True)]
593                                key = True
594                                if fvar in ['UV','uv','uvmet']: key = False
595                            elif typefile in ['gcm']:
596                                key = False
597                            if metwind and mapmode == 1:   [vecx_frame,vecy_frame] = m.rotate_vector(vecx_frame, vecy_frame, lon2d, lat2d)
598                            if var:       colorvec = definecolorvec(colorb) 
599                            else:         colorvec = definecolorvec(back) 
600                            vectorfield(vecx_frame, vecy_frame, x, y, stride=stride, csmooth=2,\
601                                             #scale=15., factor=300., color=colorvec, key=key)
602                                             scale=20., factor=250., color=colorvec, key=key)
603                                                              #200.         ## or csmooth=stride
604                        if ope == '-' and nplot == numplot: # this should work as long as ope is '-' guarantees 3 plots for 4 panels without contour
605                            subplot(subv,subh,nplot+1)
606                            rcParams["legend.fontsize"] = 'xx-large'
607                            if indexlat is None:
608                                latmin = -50.; latmax = 50. # latitude range for histogram of difference
609                                zeindexlat = (lat<latmax)*(lat>latmin)
610                                # this follows the define_axis logic in myplot.py:
611                                if indextime is None or indexlon is None: what_I_plot_frame = what_I_plot_frame[zeindexlat,:]
612                                else: what_I_plot_frame = what_I_plot_frame[:,zeindexlat]
613                            toplot = np.ravel(what_I_plot_frame[np.isnan(what_I_plot_frame)==False])
614                            zebins = np.linspace(zevmin,zevmax,num=30)
615                            hist(toplot,bins=zebins,histtype='step',linewidth=2,color='k',normed=True)
616                            zestd = np.std(toplot);zemean = np.mean(toplot)
617                            zebins = np.linspace(zevmin,zevmax,num=300)
618                            zegauss = (1./(zestd * np.sqrt(2 * np.pi)) * np.exp( - (zebins - zemean)**2 / (2 * zestd**2) ) )
619                            plot(zebins, zegauss, linewidth=1, color='r',label="mean: "+str(zemean)[0:5]+"\nstd: "+str(zestd)[0:5])
620                            legend(loc=0,frameon=False)
621                            title("Histogram fig(1) "+ope+" fig(2)")
622                            subplot(subv,subh,nplot) # go back to last plot for title of contour difference
623                           
624                    elif which == "contour":
625                        zevminc, zevmaxc = calculate_bounds(what_I_plot_frame, vmin=min(what_I_plot_frame), vmax=max(what_I_plot_frame))
626                        zelevels = np.linspace(zevminc,zevmaxc,ticks/2) #20)
627                        if var2 == 'HGT':     zelevels = np.arange(-10000.,30000.,1000.)
628                        elif var2 == 'tpot':  zelevels = np.arange(270,370,5)
629                        elif var2 == 'tk':    zelevels = np.arange(150,250,5)
630                        if mapmode == 0:   
631                            what_I_plot_frame, x, y = define_axis( lonyeah,latyeah,vert,time,indexlon,indexlat,indexvert,\
632                                                              itime,what_I_plot_frame, len(all_var2[index_f].shape),vertmode )
633                            ## is this needed? only if len(all_var2[index_f].shape) != len(all_var[index_f].shape)
634                            cs = contour( x,y,what_I_plot_frame, zelevels, colors='k', linewidths = 0.33)#, alpha=0.5, linestyles='solid')
635                            clabel(cs, inline=1, fontsize = 4.*rcParams['font.size']/5., fmt=fmtvar(var2.upper()))
636                        elif mapmode == 1:  cs = m.contour( x,y,what_I_plot_frame, zelevels, colors='k', linewidths = 0.33)#, alpha=0.5, linestyles='solid')
637                    if which in ["regular","unidim"]:
638
639                        if nplot > 1 and which == "unidim":
640                           pass  ## because we superimpose nplot instances
641                        else:
642                           # Axis directives for movie frames [including the first one).
643                           zxmin, zxmax = xaxis ; zymin, zymax = yaxis
644                           if zxmin is not None: mpl.pyplot.xlim(xmin=zxmin)
645                           if zxmax is not None: mpl.pyplot.xlim(xmax=zxmax)
646                           if zymin is not None: mpl.pyplot.ylim(ymin=zymin)
647                           if zymax is not None: mpl.pyplot.ylim(ymax=zymax)
648                           if ylog:      mpl.pyplot.semilogy()
649                           if invert_y:  ax = mpl.pyplot.gca() ; ax.set_ylim(ax.get_ylim()[::-1])
650                           if xlab is not None: xlabel(xlab)
651                           if ylab is not None: ylabel(ylab)
652
653                        if mrate is not None:
654                           ### THIS IS A MENCODER MOVIE
655                           if mrate > 0:
656                             figframe=mpl.pyplot.gcf()
657                             if mquality:   figframe.set_dpi(600.)
658                             else:          figframe.set_dpi(200.)
659                             mframe=fig2img(figframe)
660                             if imov == 0:
661                                moviename='movie' ;W,H = figframe.canvas.get_width_height()
662                                video = VideoSink((H,W), moviename, rate=mrate, byteorder="rgba")
663                             video.run(mframe) ; close()
664                             if imov == iend: video.close()                           
665                           ### THIS IS A WEBPAGE MOVIE
666                           else:
667                             nameframe = "image"+str(1000+imov)
668                             makeplotres(nameframe,res=100.,disp=False) ; close()
669                             if imov == 0: myfile = open("zepics", 'w')
670                             myfile.write("modImages["+str(imov)+"] = '"+nameframe+"_100.png';"+ '\n')
671                             if imov == iend:
672                                 myfile.write("first_image = 0;"+ '\n')
673                                 myfile.write("last_image = "+str(iend)+";"+ '\n')
674                                 myfile.close()
675                        if var2 and which == "regular":  which = "contour"
676                        imov = imov+1
677                    elif which == "contour":
678                        which = "regular"
679
680       ### Next subplot
681       zevarname = varname
682       if redope is not None: zevarname = zevarname + "_" + redope
683       basename = getname(var=zevarname,var2=var2,winds=winds,anomaly=anomaly)
684       if len(what_I_plot.shape) > 3:
685           basename = basename + getstralt(nc,level) 
686       if mrate is not None: basename = "movie_" + basename
687       if typefile in ['meso']:
688            if sslon is not None: basename = basename + "_lon_" + str(int(round(lonp)))
689            if sslat is not None: basename = basename + "_lat_" + str(int(round(latp)))
690            plottitle = basename
691            ### dans le nouveau systeme time=ls,sol,lt cette ligne pourrait ne servir a rien (ou deplacer au dessus)
692            if addchar and indextime is not None:   [addchar,gogol,gogol2] = getlschar ( all_namefile[index_f] )  ;  plottitle = plottitle + addchar
693            ### en fait redope is None doit etre remplace par : n'est ni maxt ni mint
694            if redope is None and ltst is not None and ( (mapmode == 0) or (proj in ["lcc","laea","merc","nsper"]) ):  plottitle = plottitle + "_LT" + str(ltst)
695       else:
696            if fileref is not None:
697                if index_f is numplot-1:     plottitle = basename+' '+"fig(1) "+ope+" fig(2)"
698                elif index_f is numplot-2:   plottitle = basename+' '+fileref
699                else:                        plottitle = basename+' '+all_namefile[index_f]
700            else:                            plottitle = basename+' '+all_namefile[index_f]
701       if mult != 1:                         plottitle = '{:.0e}'.format(mult) + "*" + plottitle
702       if zetitle[0] != "fill":                 
703          if index_f < len(zetitle): plottitle = zetitle[index_f]
704          else: plottitle = zetitle[0]
705          if titleref is "fill":             titleref=zetitle[index_f]
706          if fileref is not None:
707             if index_f is numplot-2:        plottitle = titleref
708             if index_f is numplot-1:        plottitle = "fig(1) "+ope+" fig(2)"
709#       if indexlon is not None:      plottitle = plottitle + " lon: " + str(min(lon[indexlon])) +" "+ str(max(lon[indexlon]))
710#       if indexlat is not None:      plottitle = plottitle + " lat: " + str(min(lat[indexlat])) +" "+ str(max(lat[indexlat]))
711#       if indexvert is not None:     plottitle = plottitle + " vert: " + str(min(vert[indexvert])) +" "+ str(max(vert[indexvert]))
712#       if indextime is not None:     plottitle = plottitle + " time: " + str(min(time[indextime])) +" "+ str(max(time[indextime]))
713       if colorb != "onebar": title( plottitle )
714       if nplot >= numplot: error = True
715       nplot += 1
716
717    if colorb == "onebar":
718        cax = axes([0.1, 0.2, 0.8, 0.03]) # a ameliorer
719        zecb = colorbar(cax=cax, orientation="horizontal", format=fmtvar(fvar.upper()),\
720                 ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional')
721        if zetitle[0] != "fill": zecb.ax.set_xlabel(zetitle[index_f]) ; zetitle[index_f]=""
722
723
724    ##########################################################################
725    ### Save the figure in a file in the data folder or an user-defined folder
726    if outputname is None:
727       if typefile in ['meso']:   prefix = getprefix(nc)
728       elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
729       else:                                prefix = ''
730    ###
731       zeplot = prefix + basename
732       if addchar:         zeplot = zeplot + addchar
733       if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
734    ###
735       if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
736       else:               zeplot = target + "/" + zeplot 
737    ###
738    else:
739       zeplot=outputname
740
741    if mrate is None:
742        pad_inches_value = 0.35
743        print "********** SAVE ", save
744        if save == 'png': 
745            if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
746            makeplotres(zeplot,res=resolution,pad_inches_value=pad_inches_value,disp=False)
747        elif save in ['eps','svg','pdf']:     makeplotres(zeplot,pad_inches_value=pad_inches_value,disp=False,ext=save)
748        elif save == 'gui':                   show()
749        elif save == 'donothing':             pass
750        elif save == 'txt':                   print "Saved results in txt file." 
751        else: 
752            print "INFO: save mode not supported. using gui instead."
753            show()
754
755    ###################################
756    #### Getting more out of this video -- PROBLEMS WITH CREATED VIDEOS
757    #
758    #if mrate is not None:
759    #    print "Re-encoding movie.. first pass"
760    #    video.first_pass(filename=moviename,quality=mquality,rate=mrate)
761    #    print "Re-encoding movie.. second pass"
762    #    video.second_pass(filename=moviename,quality=mquality,rate=mrate)   
763
764    ###############
765    ### Now the end
766    return zeplot
Note: See TracBrowser for help on using the repository browser.