source: trunk/UTIL/PYTHON/myplot.py @ 771

Last change on this file since 771 was 771, checked in by acolaitis, 12 years ago

Sorry, typo in previous commit

File size: 67.5 KB
Line 
1## Author: AS
2def errormess(text,printvar=None):
3    print text
4    if printvar is not None: print printvar
5    exit()
6    return
7
8## Author: AS
9def adjust_length (tab, zelen):
10    from numpy import ones
11    if tab is None:
12        outtab = ones(zelen) * -999999
13    else:
14        if zelen != len(tab):
15            print "not enough or too much values... setting same values all variables"
16            outtab = ones(zelen) * tab[0]
17        else:
18            outtab = tab
19    return outtab
20
21## Author: AS
22def getname(var=False,var2=False,winds=False,anomaly=False):
23    if var and winds:     basename = var + '_UV'
24    elif var:             basename = var
25    elif winds:           basename = 'UV'
26    else:                 errormess("please set at least winds or var",printvar=nc.variables)
27    if anomaly:           basename = 'd' + basename
28    if var2:              basename = basename + '_' + var2
29    return basename
30
31## Author: AS + AC
32def localtime(time,lon,nc): # lon is the mean longitude of the plot, not of the domain. central lon of domain is taken from cen_lon
33    import numpy as np
34    dt_hour=1.
35    start=0.
36    if lon is not None:
37       if lon[0,1]!=lon[0,0]: mean_lon_plot = 0.5*(lon[0,1]-lon[0,0])
38       else: mean_lon_plot=lon[0,0]
39    elif hasattr(nc, 'CEN_LON'): mean_lon_plot=getattr(nc, 'CEN_LON')
40    else: mean_lon_plot=0.
41    if hasattr(nc,'TITLE'):
42       title=getattr(nc, 'TITLE')
43       if hasattr(nc,'DT') and 'MRAMS' in title: dt_hour=getattr(nc, 'DT')/60. #LMD MMM is 1 output/hour (and not 1 output/timestep)
44                        # MRAMS is 1 output/timestep, unless stride is added in ic.py
45       if hasattr(nc,'START_DATE'):
46          start_date=getattr(nc, 'START_DATE')
47          start_hour=np.float(start_date[11:13])
48          start_minute=np.float(start_date[14:16])/60.
49          start=start_hour+start_minute # start is the local time of simu at longitude 0
50          ltst = start + time*dt_hour + mean_lon_plot / 15.
51       else: ltst = time*dt_hour + mean_lon_plot / 15.
52    else: ltst = time*dt_hour + mean_lon_plot / 15.
53    ltst = int (ltst * 10) / 10.
54    ltst = ltst % 24
55    return ltst
56
57## Author: AC
58def check_localtime(time):
59    a=-1
60    print time
61    for i in range(len(time)-1):
62       if (time[i] > time[i+1]): a=i
63    if a >= 0 and a < (len(time)-1)/2.:
64       print "Sorry, time axis is not regular."
65       print "Contourf needs regular axis... recasting"
66       for i in range(a+1):
67          time[i]=time[i]-24.
68    if a >= 0 and a >= (len(time)-1)/2.:
69       print "Sorry, time axis is not regular."
70       print "Contourf needs regular axis... recasting"
71       for i in range((len(time)-1) - a):
72          time[a+1+i]=time[a+1+i]+24.
73    return time
74
75## Author: AS, AC, JL
76def whatkindfile (nc):
77    typefile = 'gcm' # default
78    if 'controle' in nc.variables:             typefile = 'gcm'
79    elif 'phisinit' in nc.variables:           typefile = 'gcm'
80    elif 'phis' in nc.variables:               typefile = 'gcm'
81    elif 'time_counter' in nc.variables:       typefile = 'earthgcm'
82    elif hasattr(nc,'START_DATE'):             typefile = 'meso' 
83    elif 'HGT_M' in nc.variables:              typefile = 'geo'
84    elif hasattr(nc,'institution'):
85      if "European Centre" in getattr(nc,'institution'):  typefile = 'ecmwf'
86    return typefile
87
88## Author: AS
89def getfield (nc,var):
90    ## this allows to get much faster (than simply referring to nc.variables[var])
91    import numpy as np
92    dimension = len(nc.variables[var].dimensions)
93    #print "   Opening variable with", dimension, "dimensions ..."
94    if dimension == 2:    field = nc.variables[var][:,:]
95    elif dimension == 3:  field = nc.variables[var][:,:,:]
96    elif dimension == 4:  field = nc.variables[var][:,:,:,:]
97    elif dimension == 1:  field = nc.variables[var][:]
98    # if there are NaNs in the ncdf, they should be loaded as a masked array which will be
99    # recasted as a regular array later in reducefield
100    if (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
101       print "Warning: netcdf as nan values but is not loaded as a Masked Array."
102       print "recasting array type"
103       out=np.ma.masked_invalid(field)
104       out.set_fill_value([np.NaN])
105    else:
106    # missing values from zrecast or hrecast are -1e-33
107       masked=np.ma.masked_where(field < -1e30,field)
108       masked2=np.ma.masked_where(field > 1e35,field)
109       masked.set_fill_value([np.NaN]) ; masked2.set_fill_value([np.NaN])
110       mask = np.ma.getmask(masked) ; mask2 = np.ma.getmask(masked2)
111       if (True in np.array(mask)):
112          out=masked
113          print "Masked array... Missing value is NaN"
114       elif (True in np.array(mask2)):
115          out=masked2
116          print "Masked array... Missing value is NaN"
117#       else:
118#       # missing values from api are 1e36
119#          masked=np.ma.masked_where(field > 1e35,field)
120#          masked.set_fill_value([np.NaN])
121#          mask = np.ma.getmask(masked)
122#          if (True in np.array(mask)):out=masked
123#          else:out=field
124       else:
125#       # missing values from MRAMS files are 0.100E+32
126          masked=np.ma.masked_where(field > 1e30,field)
127          masked.set_fill_value([np.NaN])
128          mask = np.ma.getmask(masked)
129          if (True in np.array(mask)):out=masked
130          else:out=field
131#       else:out=field
132    return out
133
134## Author: AC
135# Compute the norm of the winds or return an hodograph
136# The corresponding variable to call is UV or uvmet (to use api)
137def windamplitude (nc,mode):
138    import numpy as np
139    varinfile = nc.variables.keys()
140    if "U" in varinfile: zu=getfield(nc,'U')
141    elif "Um" in varinfile: zu=getfield(nc,'Um')
142    else: errormess("you need slopex or U or Um in your file.")
143    if "V" in varinfile: zv=getfield(nc,'V')
144    elif "Vm" in varinfile: zv=getfield(nc,'Vm')
145    else: errormess("you need V or Vm in your file.")
146    znt,znz,zny,znx = np.array(zu).shape
147    if hasattr(nc,'WEST-EAST_PATCH_END_UNSTAG'):znx=getattr(nc, 'WEST-EAST_PATCH_END_UNSTAG')
148    zuint = np.zeros([znt,znz,zny,znx])
149    zvint = np.zeros([znt,znz,zny,znx])
150    if "U" in varinfile:
151       if hasattr(nc,'SOUTH-NORTH_PATCH_END_STAG'): zny_stag=getattr(nc, 'SOUTH-NORTH_PATCH_END_STAG')
152       if hasattr(nc,'WEST-EAST_PATCH_END_STAG'): znx_stag=getattr(nc, 'WEST-EAST_PATCH_END_STAG')
153       if zny_stag == zny: zvint=zv
154       else:
155          for yy in np.arange(zny):      zvint[:,:,yy,:] = (zv[:,:,yy,:] + zv[:,:,yy+1,:])/2.
156       if znx_stag == znx: zuint=zu
157       else:
158          for xx in np.arange(znx):      zuint[:,:,:,xx] = (zu[:,:,:,xx] + zu[:,:,:,xx+1])/2.
159    else:
160       zuint=zu
161       zvint=zv
162    if mode=='amplitude': return np.sqrt(zuint**2 + zvint**2)
163    if mode=='hodograph': return zuint,zvint
164    if mode=='hodograph_2': return None, 360.*np.arctan(zvint/zuint)/(2.*np.pi)
165
166## Author: AC
167# Compute the enrichment factor of non condensible gases
168# The corresponding variable to call is enfact
169# enrichment factor is computed as in Yuan Lian et al. 2012
170# i.e. you need to have VL2 site at LS 135 in your data
171# this only requires co2col so that you can concat.nc at low cost
172def enrichment_factor(nc,lon,lat,time):
173    import numpy as np
174    from myplot import reducefield
175    varinfile = nc.variables.keys()
176    if "co2col" in varinfile: co2col=getfield(nc,'co2col')
177    else: print "error, you need co2col var in your file"
178    if "ps" in varinfile: ps=getfield(nc,'ps')
179    else: print "error, you need ps var in your file"
180    dimension = len(nc.variables['co2col'].dimensions)
181    if dimension == 2: 
182      zny,znx = np.array(co2col).shape
183      znt=1
184    elif dimension == 3: znt,zny,znx = np.array(co2col).shape
185    mmrarcol = np.zeros([znt,zny,znx])
186    enfact = np.zeros([znt,zny,znx])
187    grav=3.72
188    mmrarcol[:,:,:] = 1. - grav*co2col[:,:,:]/ps[:,:,:]
189# Computation with reference argon mmr at VL2 Ls 135 (as in Yuan Lian et al 2012)
190    lonvl2=np.zeros([1,2])
191    latvl2=np.zeros([1,2])
192    timevl2=np.zeros([1,2])
193    lonvl2[0,0]=-180
194    lonvl2[0,1]=180
195    latvl2[:,:]=48.16
196    timevl2[:,:]=135.
197    indexlon  = getsindex(lonvl2,0,lon)
198    indexlat  = getsindex(latvl2,0,lat)
199    indextime = getsindex(timevl2,0,time)
200    mmrvl2, error = reducefield( mmrarcol, d4=indextime, d1=indexlon, d2=indexlat)
201    print "VL2 Ls 135 mmr arcol:", mmrvl2
202    enfact[:,:,:] = mmrarcol[:,:,:]/mmrvl2
203    return enfact
204
205## Author: AC
206# Compute the norm of the slope angles
207# The corresponding variable to call is SLOPEXY
208def slopeamplitude (nc):
209    import numpy as np
210    varinfile = nc.variables.keys()
211    if "slopex" in varinfile: zu=getfield(nc,'slopex')
212    elif "SLOPEX" in varinfile: zu=getfield(nc,'SLOPEX')
213    else: errormess("you need slopex or SLOPEX in your file.") 
214    if "slopey" in varinfile: zv=getfield(nc,'slopey')
215    elif "SLOPEY" in varinfile: zv=getfield(nc,'SLOPEY')
216    else: errormess("you need slopey or SLOPEY in your file.")
217    znt,zny,znx = np.array(zu).shape
218    zuint = np.zeros([znt,zny,znx])
219    zvint = np.zeros([znt,zny,znx])
220    zuint=zu
221    zvint=zv
222    return np.sqrt(zuint**2 + zvint**2)
223
224## Author: AC
225# Compute the temperature difference between surface and first level.
226# API is automatically called to get TSURF and TK.
227# The corresponding variable to call is DELTAT
228def deltat0t1 (nc):
229    import numpy as np
230    varinfile = nc.variables.keys()
231    if "tsurf" in varinfile: zu=getfield(nc,'tsurf')
232    elif "TSURF" in varinfile: zu=getfield(nc,'TSURF')
233    else: errormess("You need tsurf or TSURF in your file")
234    if "tk" in varinfile: zv=getfield(nc,'tk')
235    elif "TK" in varinfile: zv=getfield(nc,'TK')
236    else: errormess("You need tk or TK in your file. (might need to use API. try to add -i 4 -l XXX)")
237    znt,zny,znx = np.array(zu).shape
238    zuint = np.zeros([znt,zny,znx])
239    zuint=zu - zv[:,0,:,:]
240    return zuint
241
242## Author: AS + TN + AC
243def reducefield (input,d4=None,d3=None,d2=None,d1=None,yint=False,alt=None,anomaly=False,redope=None,mesharea=None,unidim=999):
244    ### we do it the reverse way to be compliant with netcdf "t z y x" or "t y x" or "y x"
245    ### it would be actually better to name d4 d3 d2 d1 as t z y x
246    ### ... note, anomaly is only computed over d1 and d2 for the moment
247    import numpy as np
248    from mymath import max,mean,min,sum,getmask
249    csmooth = 12 ## a fair amount of grid points (too high results in high computation time)
250    if redope is not None:
251       if   redope == "mint":     input = min(input,axis=0) ; d1 = None
252       elif redope == "maxt":     input = max(input,axis=0) ; d1 = None
253       elif redope == "edge_y1":  input = input[:,:,0,:]    ; d2 = None
254       elif redope == "edge_y2":  input = input[:,:,-1,:]   ; d2 = None
255       elif redope == "edge_x1":  input = input[:,:,:,0]    ; d1 = None
256       elif redope == "edge_x2":  input = input[:,:,:,-1]   ; d1 = None
257       else:                      errormess("not supported. but try lines in reducefield beforehand.")
258       #elif redope == "minz":     input = min(input,axis=1) ; d2 = None
259       #elif redope == "maxz":     input = max(input,axis=1) ; d2 = None
260       #elif redope == "miny":     input = min(input,axis=2) ; d3 = None
261       #elif redope == "maxy":     input = max(input,axis=2) ; d3 = None
262       #elif redope == "minx":     input = min(input,axis=3) ; d4 = None
263       #elif redope == "maxx":     input = max(input,axis=3) ; d4 = None
264    dimension = np.array(input).ndim
265    shape = np.array(np.array(input).shape)
266    #print 'd1,d2,d3,d4: ',d1,d2,d3,d4
267    if anomaly: print 'ANOMALY ANOMALY'
268    output = input
269    error = False
270    #### this is needed to cope the case where d4,d3,d2,d1 are single integers and not arrays
271    if d4 is not None and not isinstance(d4, np.ndarray): d4=[d4]
272    if d3 is not None and not isinstance(d3, np.ndarray): d3=[d3]
273    if d2 is not None and not isinstance(d2, np.ndarray): d2=[d2]
274    if d1 is not None and not isinstance(d1, np.ndarray): d1=[d1]
275    ### now the main part
276    if dimension == 2:
277        #### this is needed for 1d-type files (where dim=2 but axes are time-vert and not lat-lon)
278        if unidim==1: d2=d4 ; d1=d3 ; d4 = None ; d3 = None
279        if mesharea is None: mesharea=np.ones(shape)
280        if   max(d2) >= shape[0]: error = True
281        elif max(d1) >= shape[1]: error = True
282        elif d1 is not None and d2 is not None:
283          try:
284            totalarea = np.ma.masked_where(getmask(output),mesharea)
285            totalarea = mean(totalarea[d2,:],axis=0);totalarea = mean(totalarea[d1])
286          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
287          output = output*mesharea; output = mean(output[d2,:],axis=0); output = mean(output[d1])/totalarea
288        elif d1 is not None:         output = mean(input[:,d1],axis=1)
289        elif d2 is not None:
290          try:
291            totalarea = np.ma.masked_where(getmask(output),mesharea)
292            totalarea = mean(totalarea[d2,:],axis=0)
293          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
294          output = output*mesharea; output = mean(output[d2,:],axis=0)/totalarea
295    elif dimension == 3:
296        if mesharea is None: mesharea=np.ones(shape[[1,2]])
297        if   max(d4) >= shape[0]: error = True
298        elif max(d2) >= shape[1]: error = True
299        elif max(d1) >= shape[2]: error = True
300        elif d4 is not None and d2 is not None and d1 is not None:
301          output = mean(input[d4,:,:],axis=0)
302          try:
303            totalarea = np.ma.masked_where(getmask(output),mesharea)
304            totalarea = mean(totalarea[d2,:],axis=0);totalarea = mean(totalarea[d1])
305          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
306          output = output*mesharea; output = mean(output[d2,:],axis=0); output = mean(output[d1])/totalarea
307        elif d4 is not None and d2 is not None:
308          output = mean(input[d4,:,:],axis=0)
309          try:
310            totalarea = np.ma.masked_where(getmask(output),mesharea)
311            totalarea = mean(totalarea[d2,:],axis=0)
312          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
313          output = output*mesharea; output = mean(output[d2,:],axis=0)/totalarea
314        elif d4 is not None and d1 is not None:    output = mean(input[d4,:,:],axis=0); output=mean(output[:,d1],axis=1)
315        elif d2 is not None and d1 is not None:
316          try:
317            totalarea = np.tile(mesharea,(output.shape[0],1,1))
318            totalarea = np.ma.masked_where(getmask(output),totalarea)
319            totalarea = mean(totalarea[:,d2,:],axis=1);totalarea = mean(totalarea[:,d1],axis=1)
320          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
321          output = output*mesharea; output = mean(output[:,d2,:],axis=1); output = mean(output[:,d1],axis=1)/totalarea
322        elif d1 is not None:   output = mean(input[:,:,d1],axis=2)
323        elif d2 is not None:   
324          try:
325            totalarea = np.tile(mesharea,(output.shape[0],1,1))
326            totalarea = np.ma.masked_where(getmask(output),totalarea)
327            totalarea = mean(totalarea[:,d2,:],axis=1)
328          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
329          output = output*mesharea; output = mean(output[:,d2,:],axis=1)/totalarea
330        elif d4 is not None:   output = mean(input[d4,:,:],axis=0)
331    elif dimension == 4:
332        if mesharea is None: mesharea=np.ones(shape[[2,3]]) # mesharea=np.random.random_sample(shape[[2,3]])*5. + 2. # pour tester
333        if   max(d4) >= shape[0]: error = True
334        elif max(d3) >= shape[1]: error = True
335        elif max(d2) >= shape[2]: error = True
336        elif max(d1) >= shape[3]: error = True
337        elif d4 is not None and d3 is not None and d2 is not None and d1 is not None:
338             output = mean(input[d4,:,:,:],axis=0)
339             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
340             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.)
341             try:
342               totalarea = np.ma.masked_where(np.isnan(output),mesharea)
343               totalarea = mean(totalarea[d2,:],axis=0); totalarea = mean(totalarea[d1])
344             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
345             output = output*mesharea; output = mean(output[d2,:],axis=0); output = mean(output[d1])/totalarea
346        elif d4 is not None and d3 is not None and d2 is not None: 
347             output = mean(input[d4,:,:,:],axis=0)
348             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
349             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.)
350             try:
351               totalarea = np.ma.masked_where(np.isnan(output),mesharea)
352               totalarea = mean(totalarea[d2,:],axis=0)
353             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
354             output = output*mesharea; output = mean(output[d2,:],axis=0)/totalarea
355        elif d4 is not None and d3 is not None and d1 is not None: 
356             output = mean(input[d4,:,:,:],axis=0)
357             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
358             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
359             output = mean(output[:,d1],axis=1)
360        elif d4 is not None and d2 is not None and d1 is not None: 
361             output = mean(input[d4,:,:,:],axis=0)
362             if anomaly:
363                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
364             try:
365               totalarea = np.tile(mesharea,(output.shape[0],1,1))
366               totalarea = np.ma.masked_where(getmask(output),mesharea)
367               totalarea = mean(totalarea[:,d2,:],axis=1); totalarea = mean(totalarea[:,d1],axis=1)
368             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
369             output = output*mesharea; output = mean(output[:,d2,:],axis=1); output = mean(output[:,d1],axis=1)/totalarea
370             #noperturb = smooth1d(output,window_len=7)
371             #lenlen = len(output) ; output = output[1:lenlen-7] ; yeye = noperturb[4:lenlen-4]
372             #plot(output) ; plot(yeye) ; show() ; plot(output-yeye) ; show()
373        elif d3 is not None and d2 is not None and d1 is not None:
374             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
375             if anomaly:
376                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
377             try:
378               totalarea = np.tile(mesharea,(output.shape[0],1,1))
379               totalarea = np.ma.masked_where(getmask(output),totalarea)
380               totalarea = mean(totalarea[:,d2,:],axis=1); totalarea = mean(totalarea[:,d1],axis=1)
381             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
382             output = output*mesharea; output = mean(output[:,d2,:],axis=1); output = mean(output[:,d1],axis=1)/totalarea
383        elif d4 is not None and d3 is not None: 
384             output = mean(input[d4,:,:,:],axis=0) 
385             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
386             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
387        elif d4 is not None and d2 is not None: 
388             output = mean(input[d4,:,:,:],axis=0)
389             try:
390               totalarea = np.tile(mesharea,(output.shape[0],1,1))
391               totalarea = np.ma.masked_where(getmask(output),mesharea)
392               totalarea = mean(totalarea[:,d2,:],axis=1)
393             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
394             output = output*mesharea; output = mean(output[:,d2,:],axis=1)/totalarea
395        elif d4 is not None and d1 is not None: 
396             output = mean(input[d4,:,:,:],axis=0) 
397             output = mean(output[:,:,d1],axis=2)
398        elif d3 is not None and d2 is not None:
399             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
400             try:
401               totalarea = np.tile(mesharea,(output.shape[0],1,1))
402               totalarea = np.ma.masked_where(getmask(output),mesharea)
403               totalarea = mean(totalarea[:,d2,:],axis=1)
404             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
405             output = output*mesharea; output = mean(output[:,d2,:],axis=1)/totalarea
406        elif d3 is not None and d1 is not None: 
407             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
408             output = mean(output[:,:,d1],axis=2)
409        elif d2 is not None and d1 is not None:
410             try:
411               totalarea = np.tile(mesharea,(output.shape[0],output.shape[1],1,1))
412               totalarea = np.ma.masked_where(getmask(output),totalarea)
413               totalarea = mean(totalarea[:,:,d2,:],axis=2); totalarea = mean(totalarea[:,:,d1],axis=1)
414             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
415             output = output*mesharea; output = mean(output[:,:,d2,:],axis=2); output = mean(output[:,:,d1],axis=2)/totalarea
416        elif d1 is not None:        output = mean(input[:,:,:,d1],axis=3)
417        elif d2 is not None:
418             try:
419               totalarea = np.tile(mesharea,(output.shape[0],output.shape[1],1,output.shape[3]))
420               totalarea = np.ma.masked_where(getmask(output),totalarea)
421               totalarea = mean(totalarea[:,:,d2,:],axis=2)
422             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
423             output = output*mesharea; output = mean(output[:,:,d2,:],axis=2)/totalarea
424        elif d3 is not None:        output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
425        elif d4 is not None:        output = mean(input[d4,:,:,:],axis=0)
426    dimension2 = np.array(output).ndim
427    shape2 = np.array(output).shape
428    print 'REDUCEFIELD dim,shape: ',dimension,shape,' >>> ',dimension2,shape2
429    return output, error
430
431## Author: AC + AS
432def reduce_zaxis (input,ax=None,yint=False,vert=None,indice=None):
433    from mymath import max,mean
434    from scipy import integrate
435    import numpy as np
436    if yint and vert is not None and indice is not None:
437       if type(input).__name__=='MaskedArray':
438          input.set_fill_value([np.NaN])
439          output = integrate.trapz(input.filled(),x=vert[indice],axis=ax)
440       else:
441          output = integrate.trapz(input,x=vert[indice],axis=ax)
442    else:
443       output = mean(input,axis=ax)
444    return output
445
446## Author: AS + TN
447def definesubplot ( numplot, fig ):
448    from matplotlib.pyplot import rcParams
449    rcParams['font.size'] = 12. ## default (important for multiple calls)
450    if numplot <= 0:
451        subv = 99999
452        subh = 99999
453    elif numplot == 1: 
454        subv = 1
455        subh = 1
456    elif numplot == 2:
457        subv = 1 #2
458        subh = 2 #1
459        fig.subplots_adjust(wspace = 0.35)
460        rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
461    elif numplot == 3:
462        subv = 3
463        subh = 1
464        fig.subplots_adjust(hspace = 0.5)
465        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
466    elif numplot == 4:
467        subv = 2
468        subh = 2
469        fig.subplots_adjust(wspace = 0.4, hspace = 0.3)
470        rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
471    elif numplot <= 6:
472        subv = 2
473        subh = 3
474        #fig.subplots_adjust(wspace = 0.4, hspace = 0.0)
475        fig.subplots_adjust(wspace = 0.5, hspace = 0.3)
476        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
477    elif numplot <= 8:
478        subv = 2
479        subh = 4
480        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
481        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
482    elif numplot <= 9:
483        subv = 3
484        subh = 3
485        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
486        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
487    elif numplot <= 12:
488        subv = 3
489        subh = 4
490        fig.subplots_adjust(wspace = 0, hspace = 0.1)
491        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
492    elif numplot <= 16:
493        subv = 4
494        subh = 4
495        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
496        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
497    else:
498        print "number of plot supported: 1 to 16"
499        exit()
500    return subv,subh
501
502## Author: AS
503def getstralt(nc,nvert):
504    varinfile = nc.variables.keys()
505    if 'vert' not in varinfile:
506        stralt = "_lvl" + str(nvert)
507    else:
508        zelevel = int(nc.variables['vert'][nvert])
509        if abs(zelevel) < 10000.:   strheight=str(zelevel)+"m"
510        else:                       strheight=str(int(zelevel/1000.))+"km"
511        if 'altitude'       in nc.dimensions:   stralt = "_"+strheight+"-AMR"
512        elif 'altitude_abg' in nc.dimensions:   stralt = "_"+strheight+"-ALS"
513        elif 'bottom_top'   in nc.dimensions:   stralt = "_"+strheight
514        elif 'pressure'     in nc.dimensions:   stralt = "_"+str(zelevel)+"Pa"
515        else:                                   stralt = ""
516    return stralt
517
518## Author: AS
519def getlschar ( namefile, getaxis=False ):
520    from netCDF4 import Dataset
521    from timestuff import sol2ls
522    from numpy import array
523    from string import rstrip
524    import os as daos
525    namefiletest = rstrip( rstrip( rstrip( namefile, chars="_z"), chars="_zabg"), chars="_p")
526    testexist = daos.path.isfile(namefiletest)
527    zetime = None
528    if testexist: 
529      namefile = namefiletest
530      #### we assume that wrfout is next to wrfout_z and wrfout_zabg
531      nc  = Dataset(namefile)
532      zetime = None
533      days_in_month = [61, 66, 66, 65, 60, 54, 50, 46, 47, 47, 51, 56]
534      plus_in_month = [ 0, 61,127,193,258,318,372,422,468,515,562,613]
535      if 'Times' in nc.variables:
536        zetime = nc.variables['Times'][0]
537        shape = array(nc.variables['Times']).shape
538        if shape[0] < 2: zetime = None
539    if zetime is not None \
540       and 'vert' not in nc.variables:
541        ##### strangely enough this does not work for api or ncrcat results!
542        zesol = plus_in_month[ int(zetime[5]+zetime[6])-1 ] + int(zetime[8]+zetime[9]) - 1 ##les sols GCM commencent a 0
543        dals = int( 10. * sol2ls ( zesol ) ) / 10.
544        ###
545        zetime2 = nc.variables['Times'][1]
546        one  = int(zetime[11]+zetime[12]) + int(zetime[14]+zetime[15])/37.
547        next = int(zetime2[11]+zetime2[12]) + int(zetime2[14]+zetime2[15])/37. 
548        zehour    = one
549        zehourin  = abs ( next - one )
550        if not getaxis:
551            lschar = "_Ls"+str(dals)
552        else:
553            zelen = len(nc.variables['Times'][:])
554            yeye = range(zelen) ; lsaxis = range(zelen) ; solaxis = range(zelen) ; ltaxis = range(zelen)
555            for iii in yeye:
556               zetime = nc.variables['Times'][iii] 
557               ltaxis[iii] = int(zetime[11]+zetime[12]) + int(zetime[14]+zetime[15])/37.
558               solaxis[iii] = ltaxis[iii] / 24. + plus_in_month[ int(zetime[5]+zetime[6])-1 ] + int(zetime[8]+zetime[9]) - 1 ##les sols GCM commencent a 0
559               lsaxis[iii] = sol2ls ( solaxis[iii] )
560               if ltaxis[iii] < ltaxis[iii-1]: ltaxis[iii] = ltaxis[iii] + 24.
561               #print ltaxis[iii], solaxis[iii], lsaxis[iii], getattr( nc, 'JULDAY' )
562            lschar = lsaxis ; zehour = solaxis ; zehourin = ltaxis
563    else:
564        lschar=""
565        zehour = 0
566        zehourin = 1 
567    return lschar, zehour, zehourin
568
569## Author: AS
570def getprefix (nc):
571    prefix = 'LMD_MMM_'
572    prefix = prefix + 'd'+str(getattr(nc,'GRID_ID'))+'_'
573    prefix = prefix + str(int(getattr(nc,'DX')/1000.))+'km_'
574    return prefix
575
576## Author: AS
577def getproj (nc):
578    typefile = whatkindfile(nc)
579    if typefile in ['meso','geo']:
580        ### (il faudrait passer CEN_LON dans la projection ?)
581        map_proj = getattr(nc, 'MAP_PROJ')
582        cen_lat  = getattr(nc, 'CEN_LAT')
583        if map_proj == 2:
584            if cen_lat > 10.:   
585                proj="npstere"
586                #print "NP stereographic polar domain"
587            else:           
588                proj="spstere"
589                #print "SP stereographic polar domain"
590        elif map_proj == 1: 
591            #print "lambert projection domain"
592            proj="lcc"
593        elif map_proj == 3: 
594            #print "mercator projection"
595            proj="merc"
596        else:
597            proj="merc"
598    elif typefile in ['gcm']:        proj="cyl"    ## pb avec les autres (de trace derriere la sphere ?)
599    else:                            proj="ortho"
600    return proj   
601
602## Author: AS
603def ptitle (name):
604    from matplotlib.pyplot import title
605    title(name)
606    print name
607
608## Author: AS
609def polarinterv (lon2d,lat2d):
610    import numpy as np
611    wlon = [np.min(lon2d),np.max(lon2d)]
612    ind = np.array(lat2d).shape[0] / 2  ## to get a good boundlat and to get the pole
613    wlat = [np.min(lat2d[ind,:]),np.max(lat2d[ind,:])]
614    return [wlon,wlat]
615
616## Author: AS
617def simplinterv (lon2d,lat2d):
618    import numpy as np
619    return [[np.min(lon2d),np.max(lon2d)],[np.min(lat2d),np.max(lat2d)]]
620
621## Author: AS
622def wrfinterv (lon2d,lat2d):
623    nx = len(lon2d[0,:])-1
624    ny = len(lon2d[:,0])-1
625    lon1 = lon2d[0,0] 
626    lon2 = lon2d[nx,ny] 
627    lat1 = lat2d[0,0] 
628    lat2 = lat2d[nx,ny] 
629    if abs(0.5*(lat1+lat2)) > 60.:  wider = 0.5 * (abs(lon1)+abs(lon2)) * 0.1
630    else:                           wider = 0.
631    if lon1 < lon2:  wlon = [lon1, lon2 + wider] 
632    else:            wlon = [lon2, lon1 + wider]
633    if lat1 < lat2:  wlat = [lat1, lat2]
634    else:            wlat = [lat2, lat1]
635    return [wlon,wlat]
636
637## Author: AS
638def makeplotres (filename,res=None,pad_inches_value=0.25,folder='',disp=True,ext='png',erase=False):
639    import  matplotlib.pyplot as plt
640    from os import system
641    addstr = ""
642    if res is not None:
643        res = int(res)
644        addstr = "_"+str(res)
645    name = filename+addstr+"."+ext
646    if folder != '':      name = folder+'/'+name
647    plt.savefig(name,dpi=res,bbox_inches='tight',pad_inches=pad_inches_value)
648    if disp:              display(name)
649    if ext in ['eps','ps','svg']:   system("tar czvf "+name+".tar.gz "+name+" ; rm -f "+name)
650    if erase:   system("mv "+name+" to_be_erased")             
651    return
652
653## Author: AS + AC
654def dumpbdy (field,n,stag=None,condition=False,onlyx=False,onlyy=False):
655    nx = len(field[0,:])-1
656    ny = len(field[:,0])-1
657    if condition:
658      if stag == 'U': nx = nx-1
659      if stag == 'V': ny = ny-1
660      if stag == 'W': nx = nx+1 #special les case when we dump stag on W
661    if onlyx:    result = field[:,n:nx-n]
662    elif onlyy:  result = field[n:ny-n,:]
663    else:        result = field[n:ny-n,n:nx-n]
664    return result
665
666## Author: AS + AC
667def getcoorddef ( nc ):   
668    import numpy as np
669    ## getcoord2d for predefined types
670    typefile = whatkindfile(nc)
671    if typefile in ['meso']:
672        if '9999' not in getattr(nc,'START_DATE') :   
673            ## regular mesoscale
674            [lon2d,lat2d] = getcoord2d(nc) 
675        else:                     
676            ## idealized mesoscale                     
677            nx=getattr(nc,'WEST-EAST_GRID_DIMENSION')
678            ny=getattr(nc,'SOUTH-NORTH_GRID_DIMENSION')
679            dlat=getattr(nc,'DX')
680            ## this is dirty because Martian-specific
681            # ... but this just intended to get "lat-lon" like info
682            falselon = np.arange(-dlat*(nx-1)/2.,dlat*(nx-1)/2.,dlat)/60000.
683            falselat = np.arange(-dlat*(ny-1)/2.,dlat*(ny-1)/2.,dlat)/60000.
684            [lon2d,lat2d] = np.meshgrid(falselon,falselat) ## dummy coordinates
685            print "WARNING: domain plot artificially centered on lat,lon 0,0"
686    elif typefile in ['gcm','earthgcm','ecmwf']: 
687        #### n est ce pas nc.variables ?
688        if "longitude" in nc.dimensions:  dalon = "longitude"
689        elif "lon" in nc.dimensions:      dalon = "lon"
690        else:                             dalon = "nothing"
691        if "latitude" in nc.dimensions:   dalat = "latitude"
692        elif "lat" in nc.dimensions:      dalat = "lat"
693        else:                             dalat = "nothing"
694        [lon2d,lat2d] = getcoord2d(nc,nlat=dalat,nlon=dalon,is1d=True)
695    elif typefile in ['geo']:
696        [lon2d,lat2d] = getcoord2d(nc,nlat='XLAT_M',nlon='XLONG_M')
697    return lon2d,lat2d   
698
699## Author: AS
700def getcoord2d (nc,nlat='XLAT',nlon='XLONG',is1d=False):
701    import numpy as np
702    if nlon == "nothing" or nlat == "nothing":
703        print "NO LAT LON FIELDS. I AM TRYING MY BEST. I ASSUME GLOBAL FIELD."
704        lon = np.linspace(-180.,180.,getdimfromvar(nc)[-1])
705        lat = np.linspace(-90.,90.,getdimfromvar(nc)[-2])
706        [lon2d,lat2d] = np.meshgrid(lon,lat)
707    else:
708        if is1d:
709            lat = nc.variables[nlat][:]
710            lon = nc.variables[nlon][:]
711            [lon2d,lat2d] = np.meshgrid(lon,lat)
712        else:
713            lat = nc.variables[nlat][0,:,:]
714            lon = nc.variables[nlon][0,:,:]
715            [lon2d,lat2d] = [lon,lat]
716    return lon2d,lat2d
717
718## Author: AS
719def getdimfromvar (nc):
720    varinfile = nc.variables.keys()
721    dim = nc.variables[varinfile[-1]].shape ## usually the last variable is 4D or 3D
722    return dim
723
724## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
725def smooth1d(x,window_len=11,window='hanning'):
726    import numpy
727    """smooth the data using a window with requested size.
728    This method is based on the convolution of a scaled window with the signal.
729    The signal is prepared by introducing reflected copies of the signal
730    (with the window size) in both ends so that transient parts are minimized
731    in the begining and end part of the output signal.
732    input:
733        x: the input signal
734        window_len: the dimension of the smoothing window; should be an odd integer
735        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
736            flat window will produce a moving average smoothing.
737    output:
738        the smoothed signal
739    example:
740    t=linspace(-2,2,0.1)
741    x=sin(t)+randn(len(t))*0.1
742    y=smooth(x)
743    see also:
744    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
745    scipy.signal.lfilter
746    TODO: the window parameter could be the window itself if an array instead of a string   
747    """
748    x = numpy.array(x)
749    if x.ndim != 1:
750        raise ValueError, "smooth only accepts 1 dimension arrays."
751    if x.size < window_len:
752        raise ValueError, "Input vector needs to be bigger than window size."
753    if window_len<3:
754        return x
755    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
756        raise ValueError, "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"
757    s=numpy.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
758    #print(len(s))
759    if window == 'flat': #moving average
760        w=numpy.ones(window_len,'d')
761    else:
762        w=eval('numpy.'+window+'(window_len)')
763    y=numpy.convolve(w/w.sum(),s,mode='valid')
764    return y
765
766## Author: AS
767def smooth (field, coeff):
768        ## actually blur_image could work with different coeff on x and y
769        if coeff > 1:   result = blur_image(field,int(coeff))
770        else:           result = field
771        return result
772
773## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
774def gauss_kern(size, sizey=None):
775        import numpy as np
776        # Returns a normalized 2D gauss kernel array for convolutions
777        size = int(size)
778        if not sizey:
779                sizey = size
780        else:
781                sizey = int(sizey)
782        x, y = np.mgrid[-size:size+1, -sizey:sizey+1]
783        g = np.exp(-(x**2/float(size)+y**2/float(sizey)))
784        return g / g.sum()
785
786## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
787def blur_image(im, n, ny=None) :
788        from scipy.signal import convolve
789        # blurs the image by convolving with a gaussian kernel of typical size n.
790        # The optional keyword argument ny allows for a different size in the y direction.
791        g = gauss_kern(n, sizey=ny)
792        improc = convolve(im, g, mode='same')
793        return improc
794
795## Author: AS
796def getwinddef (nc):   
797    ###
798    varinfile = nc.variables.keys()
799    if 'Um' in varinfile:   [uchar,vchar] = ['Um','Vm'] #; print "this is API meso file"
800    elif 'U' in varinfile:  [uchar,vchar] = ['U','V']   #; print "this is RAW meso file"
801    elif 'u' in varinfile:  [uchar,vchar] = ['u','v']   #; print "this is GCM file"
802    elif 'vitu' in varinfile:  [uchar,vchar] = ['vitu','vitv']   #; print "this is GCM v5 file"
803     ### you can add choices here !
804    else:                   [uchar,vchar] = ['not found','not found']
805    ###
806    if uchar in ['U']:         metwind = False ## geometrical (wrt grid)
807    else:                      metwind = True  ## meteorological (zon/mer)
808    if metwind is False:       print "Not using meteorological winds. You trust numerical grid as being (x,y)"
809    ###
810    return uchar,vchar,metwind
811
812## Author: AS
813def vectorfield (u, v, x, y, stride=3, scale=15., factor=250., color='black', csmooth=1, key=True):
814    ## scale regle la reference du vecteur
815    ## factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
816    import  matplotlib.pyplot               as plt
817    import  numpy                           as np
818    #posx = np.min(x) - np.std(x) / 10.
819    #posy = np.min(y) - np.std(y) / 10.
820    posx = np.min(x) 
821    posy = np.min(y) - 4.*np.std(y) / 10.
822    u = smooth(u,csmooth)
823    v = smooth(v,csmooth)
824    widthvec = 0.003 #0.005 #0.003
825    q = plt.quiver( x[::stride,::stride],\
826                    y[::stride,::stride],\
827                    u[::stride,::stride],\
828                    v[::stride,::stride],\
829                    angles='xy',color=color,pivot='middle',\
830                    scale=factor,width=widthvec )
831    if color in ['white','yellow']:     kcolor='black'
832    else:                               kcolor=color
833    if key: p = plt.quiverkey(q,posx,posy,scale,\
834                   str(int(scale)),coordinates='data',color=kcolor,labelpos='S',labelsep = 0.03)
835    return 
836
837## Author: AS
838def display (name):
839    from os import system
840    system("display "+name+" > /dev/null 2> /dev/null &")
841    return name
842
843## Author: AS
844def findstep (wlon):
845    steplon = int((wlon[1]-wlon[0])/4.)  #3
846    step = 120.
847    while step > steplon and step > 15. :       step = step / 2.
848    if step <= 15.:
849        while step > steplon and step > 5.  :   step = step - 5.
850    if step <= 5.:
851        while step > steplon and step > 1.  :   step = step - 1.
852    if step <= 1.:
853        step = 1. 
854    return step
855
856## Author: AS
857def define_proj (char,wlon,wlat,back=None,blat=None,blon=None):
858    from    mpl_toolkits.basemap            import Basemap
859    import  numpy                           as np
860    import  matplotlib                      as mpl
861    from mymath import max
862    meanlon = 0.5*(wlon[0]+wlon[1])
863    meanlat = 0.5*(wlat[0]+wlat[1])
864    zewidth = np.abs(wlon[0]-wlon[1])*60000.*np.cos(3.14*meanlat/180.)
865    zeheight = np.abs(wlat[0]-wlat[1])*60000.
866    if blat is None:
867        ortholat=meanlat
868        if   wlat[0] >= 80.:   blat =  -40. 
869        elif wlat[1] <= -80.:  blat = -40.
870        elif wlat[1] >= 0.:    blat = wlat[0] 
871        elif wlat[0] <= 0.:    blat = wlat[1]
872    else:  ortholat=blat
873    if blon is None:  ortholon=meanlon
874    else:             ortholon=blon
875    h = 50.  ## en km
876    radius = 3397200.
877    #print meanlat, meanlon
878    if   char == "cyl":     m = Basemap(rsphere=radius,projection='cyl',\
879                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
880    elif char == "moll":    m = Basemap(rsphere=radius,projection='moll',lon_0=meanlon)
881    elif char == "ortho":   m = Basemap(rsphere=radius,projection='ortho',lon_0=ortholon,lat_0=ortholat)
882    elif char == "lcc":     m = Basemap(rsphere=radius,projection='lcc',lat_1=meanlat,lat_0=meanlat,lon_0=meanlon,\
883                              width=zewidth,height=zeheight)
884                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
885    elif char == "npstere": m = Basemap(rsphere=radius,projection='npstere', boundinglat=blat, lon_0=0.)
886    elif char == "spstere": m = Basemap(rsphere=radius,projection='spstere', boundinglat=blat, lon_0=180.)
887    elif char == "nplaea":  m = Basemap(rsphere=radius,projection='nplaea', boundinglat=wlat[0], lon_0=meanlon)
888    elif char == "laea":    m = Basemap(rsphere=radius,projection='laea',lon_0=meanlon,lat_0=meanlat,lat_ts=meanlat,\
889                              width=zewidth,height=zeheight)
890                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
891    elif char == "nsper":   m = Basemap(rsphere=radius,projection='nsper',lon_0=meanlon,lat_0=meanlat,satellite_height=h*1000.)
892    elif char == "merc":    m = Basemap(rsphere=radius,projection='merc',lat_ts=0.,\
893                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
894    fontsizemer = int(mpl.rcParams['font.size']*3./4.)
895    if char in ["cyl","lcc","merc","nsper","laea"]:   step = findstep(wlon)
896    else:                                             step = 10.
897    steplon = step*2.
898    zecolor ='grey'
899    zelinewidth = 1
900    zelatmax = 80
901    if meanlat > 75.: zelatmax = 90. ; step = step/2.
902    # to show gcm grid:
903    #zecolor = 'r'
904    #zelinewidth = 1
905    #step = 180./48.
906    #steplon = 360./64.
907    #zelatmax = 90. - step/3
908    if char not in ["moll"]:
909        if wlon[1]-wlon[0] < 2.:  ## LOCAL MODE
910            m.drawmeridians(np.r_[-1.:1.:0.05], labels=[0,0,0,1], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, fmt='%5.2f')
911            m.drawparallels(np.r_[-1.:1.:0.05], labels=[1,0,0,0], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, fmt='%5.2f')
912        else:  ## GLOBAL OR REGIONAL MODE
913            m.drawmeridians(np.r_[-180.:180.:steplon], labels=[0,0,0,1], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
914            m.drawparallels(np.r_[-90.:90.:step], labels=[1,0,0,0], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
915    if back: m.warpimage(marsmap(back),scale=0.75)
916            #if not back:
917            #    if not var:                                        back = "mola"    ## if no var:         draw mola
918            #    elif typefile in ['mesoapi','meso','geo'] \
919            #       and proj not in ['merc','lcc','nsper','laea']:  back = "molabw"  ## if var but meso:   draw molabw
920            #    else:                                              pass             ## else:              draw None
921    return m
922
923## Author: AS
924#### test temporaire
925def putpoints (map,plot):
926    #### from http://www.scipy.org/Cookbook/Matplotlib/Maps
927    # lat/lon coordinates of five cities.
928    lats = [18.4]
929    lons = [-134.0]
930    points=['Olympus Mons']
931    # compute the native map projection coordinates for cities.
932    x,y = map(lons,lats)
933    # plot filled circles at the locations of the cities.
934    map.plot(x,y,'bo')
935    # plot the names of those five cities.
936    wherept = 0 #1000 #50000
937    for name,xpt,ypt in zip(points,x,y):
938       plot.text(xpt+wherept,ypt+wherept,name)
939    ## le nom ne s'affiche pas...
940    return
941
942## Author: AS
943def calculate_bounds(field,vmin=None,vmax=None):
944    import numpy as np
945    from mymath import max,min,mean
946    ind = np.where(field < 9e+35)
947    fieldcalc = field[ ind ] # la syntaxe compacte ne marche si field est un tuple
948    ###
949    dev = np.std(fieldcalc)*3.0
950    ###
951    if vmin is None:  zevmin = mean(fieldcalc) - dev
952    else:             zevmin = vmin
953    ###
954    if vmax is None:  zevmax = mean(fieldcalc) + dev
955    else:             zevmax = vmax
956    if vmin == vmax:
957                      zevmin = mean(fieldcalc) - dev  ### for continuity
958                      zevmax = mean(fieldcalc) + dev  ### for continuity           
959    ###
960    if zevmin < 0. and min(fieldcalc) > 0.: zevmin = 0.
961    print "BOUNDS field ", min(fieldcalc), max(fieldcalc), " //// adopted", zevmin, zevmax
962    return zevmin, zevmax
963
964## Author: AS
965def bounds(what_I_plot,zevmin,zevmax):
966    from mymath import max,min,mean
967    ### might be convenient to add the missing value in arguments
968    #what_I_plot[ what_I_plot < zevmin ] = zevmin#*(1. + 1.e-7)
969    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1. - 1.e-7) ] = zevmin*(1. - 1.e-7)
970    else:          what_I_plot[ what_I_plot < zevmin*(1. + 1.e-7) ] = zevmin*(1. + 1.e-7)
971    #print "NEW MIN ", min(what_I_plot)
972    what_I_plot[ what_I_plot > 9e+35  ] = -9e+35
973    what_I_plot[ what_I_plot > zevmax ] = zevmax*(1. - 1.e-7)
974    #print "NEW MAX ", max(what_I_plot)
975    return what_I_plot
976
977## Author: AS
978def nolow(what_I_plot):
979    from mymath import max,min
980    lim = 0.15*0.5*(abs(max(what_I_plot))+abs(min(what_I_plot)))
981    print "NO PLOT BELOW VALUE ", lim
982    what_I_plot [ abs(what_I_plot) < lim ] = 1.e40 
983    return what_I_plot
984
985
986## Author : AC
987def hole_bounds(what_I_plot,zevmin,zevmax):
988    import numpy as np
989    zi=0
990    for i in what_I_plot:
991        zj=0
992        for j in i:
993            if ((j < zevmin) or (j > zevmax)):what_I_plot[zi,zj]=np.NaN
994            zj=zj+1
995        zi=zi+1
996
997    return what_I_plot
998
999## Author: AS
1000def zoomset (wlon,wlat,zoom):
1001    dlon = abs(wlon[1]-wlon[0])/2.
1002    dlat = abs(wlat[1]-wlat[0])/2.
1003    [wlon,wlat] = [ [wlon[0]+zoom*dlon/100.,wlon[1]-zoom*dlon/100.],\
1004                    [wlat[0]+zoom*dlat/100.,wlat[1]-zoom*dlat/100.] ]
1005    print "ZOOM %",zoom,wlon,wlat
1006    return wlon,wlat
1007
1008## Author: AS
1009def fmtvar (whichvar="def"):
1010    fmtvar    =     { \
1011             "MIXED":        "%.0f",\
1012             "UPDRAFT":      "%.0f",\
1013             "DOWNDRAFT":    "%.0f",\
1014             "TK":           "%.0f",\
1015             "T":            "%.0f",\
1016             #"ZMAX_TH":      "%.0f",\
1017             #"WSTAR":        "%.0f",\
1018             # Variables from TES ncdf format
1019             "T_NADIR_DAY":  "%.0f",\
1020             "T_NADIR_NIT":  "%.0f",\
1021             # Variables from tes.py ncdf format
1022             "TEMP_DAY":     "%.0f",\
1023             "TEMP_NIGHT":   "%.0f",\
1024             # Variables from MCS and mcs.py ncdf format
1025             "DTEMP":        "%.0f",\
1026             "NTEMP":        "%.0f",\
1027             "DNUMBINTEMP":  "%.0f",\
1028             "NNUMBINTEMP":  "%.0f",\
1029             # other stuff
1030             "TPOT":         "%.0f",\
1031             "TSURF":        "%.0f",\
1032             "U_OUT1":       "%.0f",\
1033             "T_OUT1":       "%.0f",\
1034             "def":          "%.1e",\
1035             "PTOT":         "%.0f",\
1036             "PSFC":         "%.1f",\
1037             "HGT":          "%.1e",\
1038             "USTM":         "%.2f",\
1039             "HFX":          "%.0f",\
1040             "ICETOT":       "%.1e",\
1041             "TAU_ICE":      "%.2f",\
1042             "TAUICE":       "%.2f",\
1043             "VMR_ICE":      "%.1e",\
1044             "MTOT":         "%.1f",\
1045             "ANOMALY":      "%.1f",\
1046             "W":            "%.2f",\
1047             "WMAX_TH":      "%.1f",\
1048             "WSTAR":        "%.1f",\
1049             "QSURFICE":     "%.0f",\
1050             "UM":           "%.0f",\
1051             "WIND":         "%.0f",\
1052             "UVMET":         "%.0f",\
1053             "UV":         "%.0f",\
1054             "ALBBARE":      "%.2f",\
1055             "TAU":          "%.1f",\
1056             "CO2":          "%.2f",\
1057             "ENFACT":       "%.1f",\
1058             "QDUST":        "%.6f",\
1059             #### T.N.
1060             "TEMP":         "%.0f",\
1061             "VMR_H2OICE":   "%.0f",\
1062             "VMR_H2OVAP":   "%.0f",\
1063             "TAUTES":       "%.2f",\
1064             "TAUTESAP":     "%.2f",\
1065
1066                    }
1067    if "TSURF" in whichvar: whichvar = "TSURF"
1068    if whichvar not in fmtvar:
1069        whichvar = "def"
1070    return fmtvar[whichvar]
1071
1072## Author: AS
1073####################################################################################################################
1074### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
1075def defcolorb (whichone="def"):
1076    whichcolorb =    { \
1077             "def":          "spectral",\
1078             "HGT":          "spectral",\
1079             "HGT_M":        "spectral",\
1080             "TK":           "gist_heat",\
1081             "TPOT":         "Paired",\
1082             "TSURF":        "RdBu_r",\
1083             "QH2O":         "PuBu",\
1084             "USTM":         "YlOrRd",\
1085             "WIND":         "YlOrRd",\
1086             #"T_nadir_nit":  "RdBu_r",\
1087             #"T_nadir_day":  "RdBu_r",\
1088             "HFX":          "RdYlBu",\
1089             "ICETOT":       "YlGnBu_r",\
1090             #"MTOT":         "PuBu",\
1091             "CCNQ":         "YlOrBr",\
1092             "CCNN":         "YlOrBr",\
1093             "TEMP":         "Jet",\
1094             "TAU_ICE":      "Blues",\
1095             "TAUICE":       "Blues",\
1096             "VMR_ICE":      "Blues",\
1097             "W":            "jet",\
1098             "WMAX_TH":      "spectral",\
1099             "ANOMALY":      "RdBu_r",\
1100             "QSURFICE":     "hot_r",\
1101             "ALBBARE":      "spectral",\
1102             "TAU":          "YlOrBr_r",\
1103             "CO2":          "YlOrBr_r",\
1104             "MIXED":        "GnBu",\
1105             #### T.N.
1106             "MTOT":         "spectral",\
1107             "H2O_ICE_S":    "RdBu",\
1108             "VMR_H2OICE":   "PuBu",\
1109             "VMR_H2OVAP":   "PuBu",\
1110             "WATERCAPTAG":  "Blues",\
1111                     }
1112#W --> spectral ou jet
1113#spectral BrBG RdBu_r
1114    #print "predefined colorbars"
1115    if "TSURF" in whichone: whichone = "TSURF"
1116    if whichone not in whichcolorb:
1117        whichone = "def"
1118    return whichcolorb[whichone]
1119
1120## Author: AS
1121def definecolorvec (whichone="def"):
1122        whichcolor =    { \
1123                "def":          "black",\
1124                "vis":          "yellow",\
1125                "vishires":     "yellow",\
1126                "molabw":       "yellow",\
1127                "mola":         "black",\
1128                "gist_heat":    "white",\
1129                "hot":          "tk",\
1130                "gist_rainbow": "black",\
1131                "spectral":     "black",\
1132                "gray":         "red",\
1133                "PuBu":         "black",\
1134                "titan":        "red",\
1135                        }
1136        if whichone not in whichcolor:
1137                whichone = "def"
1138        return whichcolor[whichone]
1139
1140## Author: AS
1141def marsmap (whichone="vishires"):
1142        from os import uname
1143        mymachine = uname()[1]
1144        ### not sure about speed-up with this method... looks the same
1145        if "lmd.jussieu.fr" in mymachine:   domain = "/u/aslmd/WWW/maps/"
1146        elif "aymeric-laptop" in mymachine: domain = "/home/aymeric/Dropbox/Public/"
1147        else:                               domain = "http://www.lmd.jussieu.fr/~aslmd/maps/"
1148        whichlink =     { \
1149                #"vis":         "http://maps.jpl.nasa.gov/pix/mar0kuu2.jpg",\
1150                #"vishires":    "http://www.lmd.jussieu.fr/~aslmd/maps/MarsMap_2500x1250.jpg",\
1151                #"geolocal":    "http://dl.dropbox.com/u/11078310/geolocal.jpg",\
1152                #"mola":        "http://www.lns.cornell.edu/~seb/celestia/mars-mola-2k.jpg",\
1153                #"molabw":      "http://dl.dropbox.com/u/11078310/MarsElevation_2500x1250.jpg",\
1154                "thermalday":   domain+"thermalday.jpg",\
1155                "thermalnight": domain+"thermalnight.jpg",\
1156                "tesalbedo":    domain+"tesalbedo.jpg",\
1157                "vis":         domain+"mar0kuu2.jpg",\
1158                "vishires":    domain+"MarsMap_2500x1250.jpg",\
1159                "geolocal":    domain+"geolocal.jpg",\
1160                "mola":        domain+"mars-mola-2k.jpg",\
1161                "molabw":      domain+"MarsElevation_2500x1250.jpg",\
1162                "clouds":      "http://www.johnstonsarchive.net/spaceart/marswcloudmap.jpg",\
1163                "jupiter":     "http://www.mmedia.is/~bjj/data/jupiter_css/jupiter_css.jpg",\
1164                "jupiter_voy": "http://www.mmedia.is/~bjj/data/jupiter/jupiter_vgr2.jpg",\
1165                #"bw":          domain+"EarthElevation_2500x1250.jpg",\
1166                "bw":          "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg",\
1167                "contrast":    "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg",\
1168                "nice":        "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg",\
1169                "blue":        "http://eoimages.gsfc.nasa.gov/ve/2430/land_ocean_ice_2048.jpg",\
1170                "blueclouds":  "http://eoimages.gsfc.nasa.gov/ve/2431/land_ocean_ice_cloud_2048.jpg",\
1171                "justclouds":  "http://eoimages.gsfc.nasa.gov/ve/2432/cloud_combined_2048.jpg",\
1172                "pluto":       "http://www.boulder.swri.edu/~buie/pluto/pluto_all.png",\
1173                "triton":      "http://laps.noaa.gov/albers/sos/neptune/triton/triton_rgb_cyl_www.jpg",\
1174                "titan":       "http://laps.noaa.gov/albers/sos/saturn/titan/titan_rgb_cyl_www.jpg",\
1175                #"titan":       "http://laps.noaa.gov/albers/sos/celestia/titan_50.jpg",\
1176                "titanuni":    "http://maps.jpl.nasa.gov/pix/sat6fss1.jpg",\
1177                "venus":       "http://laps.noaa.gov/albers/sos/venus/venus4/venus4_rgb_cyl_www.jpg",\
1178                "cosmic":      "http://laps.noaa.gov/albers/sos/universe/wmap/wmap_rgb_cyl_www.jpg",\
1179                        }
1180        ### see http://www.mmedia.is/~bjj/planetary_maps.html
1181        if whichone not in whichlink: 
1182                print "marsmap: choice not defined... you'll get the default one... "
1183                whichone = "vishires" 
1184        return whichlink[whichone]
1185
1186#def earthmap (whichone):
1187#       if   whichone == "contrast":    whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg"
1188#       elif whichone == "bw":          whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg"
1189#       elif whichone == "nice":        whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg"
1190#       return whichlink
1191
1192## Author: AS
1193def latinterv (area="Whole"):
1194    list =    { \
1195        "Europe":                [[ 20., 80.],[- 50.,  50.]],\
1196        "Central_America":       [[-10., 40.],[ 230., 300.]],\
1197        "Africa":                [[-20., 50.],[- 50.,  50.]],\
1198        "Whole":                 [[-90., 90.],[-180., 180.]],\
1199        "Southern_Hemisphere":   [[-90., 60.],[-180., 180.]],\
1200        "Northern_Hemisphere":   [[-60., 90.],[-180., 180.]],\
1201        "Tharsis":               [[-30., 60.],[-170.,- 10.]],\
1202        "Whole_No_High":         [[-60., 60.],[-180., 180.]],\
1203        "Chryse":                [[-60., 60.],[- 60.,  60.]],\
1204        "North_Pole":            [[ 50., 90.],[-180., 180.]],\
1205        "Close_North_Pole":      [[ 75., 90.],[-180., 180.]],\
1206        "Far_South_Pole":        [[-90.,-40.],[-180., 180.]],\
1207        "South_Pole":            [[-90.,-50.],[-180., 180.]],\
1208        "Close_South_Pole":      [[-90.,-75.],[-180., 180.]],\
1209        "Sirenum_Crater_large":  [[-46.,-34.],[-166.,-151.]],\
1210        "Sirenum_Crater_small":  [[-36.,-26.],[-168.,-156.]],\
1211        "Rupes":                 [[ 72., 90.],[-120.,- 20.]],\
1212        "Xanadu":                [[-40., 20.],[  40., 120.]],\
1213              }
1214    if area not in list:   area = "Whole"
1215    [olat,olon] = list[area]
1216    return olon,olat
1217
1218## Author: TN
1219def separatenames (name):
1220  from numpy import concatenate
1221  # look for comas in the input name to separate different names (files, variables,etc ..)
1222  if name is None:
1223     names = None
1224  else:
1225    names = []
1226    stop = 0
1227    currentname = name
1228    while stop == 0:
1229      indexvir = currentname.find(',')
1230      if indexvir == -1:
1231        stop = 1
1232        name1 = currentname
1233      else:
1234        name1 = currentname[0:indexvir]
1235      names = concatenate((names,[name1]))
1236      currentname = currentname[indexvir+1:len(currentname)]
1237  return names
1238
1239
1240## Author: TN
1241def readslices(saxis):
1242  from numpy import empty
1243  if saxis == None:
1244     zesaxis = None
1245  else:
1246     zesaxis = empty((len(saxis),2))
1247     for i in range(len(saxis)):
1248        a = separatenames(saxis[i])
1249        if len(a) == 1:
1250           zesaxis[i,:] = float(a[0])
1251        else:
1252           zesaxis[i,0] = float(a[0])
1253           zesaxis[i,1] = float(a[1])
1254           
1255  return zesaxis
1256
1257## Author: TN
1258def readdata(data,datatype,coord1,coord2):
1259## Read sparse data
1260  from numpy import empty
1261  if datatype == 'txt':
1262     if len(data[coord1].shape) == 1:
1263         return data[coord1][:]
1264     elif len(data[coord1].shape) == 2:
1265         return data[coord1][:,int(coord2)-1]
1266     else:
1267         errormess('error in readdata')
1268  elif datatype == 'sav':
1269     return data[coord1][coord2]
1270  else:
1271     errormess(datatype+' type is not supported!')
1272
1273
1274## Author: AS
1275def bidimfind(lon2d,lat2d,vlon,vlat,file=None):
1276   import numpy as np
1277   import matplotlib.pyplot as mpl
1278   if vlat is None:    array = (lon2d - vlon)**2
1279   elif vlon is None:  array = (lat2d - vlat)**2
1280   else:               array = (lon2d - vlon)**2 + (lat2d - vlat)**2
1281   idy,idx = np.unravel_index( np.argmin(array), lon2d.shape )
1282   if vlon is not None:
1283      if (np.abs(lon2d[idy,idx]-vlon)) > 5: errormess("longitude not found ",printvar=lon2d)
1284   if vlat is not None:
1285      if (np.abs(lat2d[idy,idx]-vlat)) > 5: errormess("latitude not found ",printvar=lat2d)
1286   if file is not None:
1287      print idx,idy,lon2d[idy,idx],vlon
1288      print idx,idy,lat2d[idy,idx],vlat
1289      var = file.variables["HGT"][:,:,:]
1290      mpl.contourf(var[0,:,:],30,cmap = mpl.get_cmap(name="Greys_r") ) ; mpl.axis('off') ; mpl.plot(idx,idy,'mx',mew=4.0,ms=20.0)
1291      mpl.show()
1292   return idy,idx
1293
1294## Author: TN
1295def getsindex(saxis,index,axis):
1296# input  : all the desired slices and the good index
1297# output : all indexes to be taken into account for reducing field
1298  import numpy as np
1299  if ( np.array(axis).ndim == 2):
1300      axis = axis[:,0]
1301  if saxis is None:
1302      zeindex = None
1303  else:
1304      aaa = int(np.argmin(abs(saxis[index,0] - axis)))
1305      bbb = int(np.argmin(abs(saxis[index,1] - axis)))
1306      [imin,imax] = np.sort(np.array([aaa,bbb]))
1307      zeindex = np.array(range(imax-imin+1))+imin
1308      # because -180 and 180 are the same point in longitude,
1309      # we get rid of one for averaging purposes.
1310      if axis[imin] == -180 and axis[imax] == 180:
1311         zeindex = zeindex[0:len(zeindex)-1]
1312         print "INFO: whole longitude averaging asked, so last point is not taken into account."
1313  return zeindex
1314     
1315## Author: TN
1316def define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,indextime,what_I_plot,dim0,vertmode,redope):
1317# Purpose of define_axis is to find x and y axis scales in a smart way
1318# x axis priority: 1/time 2/lon 3/lat 4/vertical
1319# To be improved !!!...
1320   from numpy import array,swapaxes
1321   x = None
1322   y = None
1323   count = 0
1324   what_I_plot = array(what_I_plot)
1325   shape = what_I_plot.shape
1326   if indextime is None and len(time) > 1:
1327      print "AXIS is time"
1328      x = time
1329      count = count+1
1330   if indexlon is None and len(lon) > 1 and redope not in ['edge_x1','edge_x2']:
1331      print "AXIS is lon"
1332      if count == 0: x = lon
1333      else: y = lon
1334      count = count+1
1335   if indexlat is None and len(lat) > 1 and redope not in ['edge_y1','edge_y2']:
1336      print "AXIS is lat"
1337      if count == 0: x = lat
1338      else: y = lat
1339      count = count+1
1340   if indexvert is None and ((dim0 == 4) or (y is None)):
1341      print "AXIS is vert"
1342      if vertmode == 0: # vertical axis is as is (GCM grid)
1343         if count == 0: x=range(len(vert))
1344         else: y=range(len(vert))
1345         count = count+1
1346      else: # vertical axis is in kms
1347         if count == 0: x = vert
1348         else: y = vert
1349         count = count+1
1350   x = array(x)
1351   y = array(y)
1352   print "CHECK SHAPE: what_I_plot, x, y", what_I_plot.shape, x.shape, y.shape
1353   if len(shape) == 1:
1354       if shape[0] != len(x):           print "WARNING: shape[0] != len(x). Correcting." ; what_I_plot = what_I_plot[0:len(x)]
1355       if len(y.shape) > 0:             y = ()
1356   elif len(shape) == 2:
1357       if shape[1] == len(y) and shape[0] == len(x) and shape[0] != shape[1]:
1358           print "INFO: swapaxes: ",what_I_plot.shape,shape ; what_I_plot = swapaxes(what_I_plot,0,1)
1359       else:
1360           if shape[0] != len(y):       print "WARNING: shape[0] != len(y). Correcting." ; what_I_plot = what_I_plot[0:len(y),:]
1361           elif shape[1] != len(x):     print "WARNING: shape[1] != len(x). Correcting." ; what_I_plot = what_I_plot[:,0:len(x)]
1362   elif len(shape) == 3:
1363       if vertmode < 0: print "not supported. must check array dimensions at some point. not difficult to implement though."
1364   return what_I_plot,x,y
1365
1366# Author: TN + AS + AC
1367def determineplot(slon, slat, svert, stime, redope):
1368    nlon = 1 # number of longitudinal slices -- 1 is None
1369    nlat = 1
1370    nvert = 1
1371    ntime = 1
1372    nslices = 1
1373    if slon is not None:
1374        length=len(slon[:,0])
1375        nslices = nslices*length
1376        nlon = len(slon)
1377    if slat is not None:
1378        length=len(slat[:,0])
1379        nslices = nslices*length
1380        nlat = len(slat)
1381    if svert is not None:
1382        length=len(svert[:,0])
1383        nslices = nslices*length
1384        nvert = len(svert)
1385    if stime is not None:
1386        length=len(stime[:,0])
1387        nslices = nslices*length
1388        ntime = len(stime)
1389    #else:
1390    #    nslices = 2 
1391    mapmode = 0
1392    if slon is None and slat is None and redope not in ['edge_x1','edge_x2','edge_y1','edge_y2']:
1393       mapmode = 1 # in this case we plot a map, with the given projection
1394    return nlon, nlat, nvert, ntime, mapmode, nslices
1395
1396## Author : AS
1397def maplatlon( lon,lat,field,\
1398               proj="cyl",colorb="jet",ndiv=10,zeback="molabw",trans=0.6,title="",\
1399               vecx=None,vecy=None,stride=2 ):
1400    ### an easy way to map a field over lat/lon grid
1401    import numpy as np
1402    import matplotlib.pyplot as mpl
1403    from matplotlib.cm import get_cmap
1404    ## get lon and lat in 2D version. get lat/lon intervals
1405    numdim = len(np.array(lon).shape)
1406    if numdim == 2:     [lon2d,lat2d] = [lon,lat]
1407    elif numdim == 1:   [lon2d,lat2d] = np.meshgrid(lon,lat)
1408    else:               errormess("lon and lat arrays must be 1D or 2D")
1409    [wlon,wlat] = latinterv()
1410    ## define projection and background. define x and y given the projection
1411    m = define_proj(proj,wlon,wlat,back=zeback,blat=None,blon=None)
1412    x, y = m(lon2d, lat2d)
1413    ## define field. bound field.
1414    what_I_plot = np.transpose(field)
1415    zevmin, zevmax = calculate_bounds(what_I_plot)  ## vmin=min(what_I_plot_frame), vmax=max(what_I_plot_frame))
1416    what_I_plot = bounds(what_I_plot,zevmin,zevmax)
1417    ## define contour field levels. define color palette
1418    ticks = ndiv + 1
1419    zelevels = np.linspace(zevmin,zevmax,ticks)
1420    palette = get_cmap(name=colorb)
1421    ## contour field
1422    m.contourf( x, y, what_I_plot, zelevels, cmap = palette, alpha = trans )
1423    ## draw colorbar
1424    if proj in ['moll','cyl']:        zeorientation="horizontal" ; zepad = 0.07
1425    else:                             zeorientation="vertical" ; zepad = 0.03
1426    #daformat = fmtvar(fvar.upper())
1427    daformat = "%.0f"
1428    zecb = mpl.colorbar( fraction=0.05,pad=zepad,format=daformat,orientation=zeorientation,\
1429                 ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional' ) 
1430    ## give a title
1431    if zeorientation == "horizontal": zecb.ax.set_xlabel(title)
1432    else:                             ptitle(title)
1433    ## draw vector
1434    if vecx is not None and vecy is not None:
1435       [vecx_frame,vecy_frame] = m.rotate_vector( np.transpose(vecx), np.transpose(vecy), lon2d, lat2d ) ## for metwinds
1436       vectorfield(vecx_frame, vecy_frame, x, y, stride=stride, csmooth=2,\
1437                                             scale=30., factor=500., color=definecolorvec(colorb), key=True)
1438    ## scale regle la reference du vecteur. factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
1439    return
1440## Author : AC
1441## Handles calls to specific computations (e.g. wind norm, enrichment factor...)
1442def select_getfield(zvarname=None,znc=None,ztypefile=None,mode=None,ztsat=None,ylon=None,ylat=None,yalt=None,ytime=None,analysis=None):
1443      from mymath import get_tsat
1444 
1445      ## Specific variables are described here:
1446      # for the mesoscale:
1447      specificname_meso = ['UV','uv','uvmet','slopexy','SLOPEXY','deltat','DELTAT','hodograph','tk','hodograph_2']
1448      # for the gcm:
1449      specificname_gcm = ['enfact']
1450
1451      ## Check for variable in file:
1452      if mode == 'check':     
1453          varname = zvarname
1454          varinfile=znc.variables.keys()
1455          logical_novarname = zvarname not in znc.variables
1456          logical_nospecificname_meso = not ((ztypefile in ['meso']) and (zvarname in specificname_meso))
1457          logical_nospecificname_gcm = not ((ztypefile in ['gcm']) and (zvarname in specificname_gcm))
1458          if ( logical_novarname and logical_nospecificname_meso and logical_nospecificname_gcm ):
1459              if len(varinfile) == 1:   varname = varinfile[0]
1460              else:                     varname = False
1461          ## Return the variable name:
1462          return varname
1463
1464      ## Get the corresponding variable:
1465      if mode == 'getvar':
1466          plot_x = None ; plot_y = None ;
1467          ### ----------- 1. saturation temperature
1468          if zvarname in ["temp","t","T_nadir_nit","T_nadir_day","temp_day","temp_night"] and ztsat:
1469              tt=getfield(znc,zvarname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
1470              if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
1471              else:                                 tinput=tt
1472              all_var=get_tsat(yalt,tinput,zlon=ylon,zlat=ylat,zalt=yalt,ztime=ytime)
1473          ### ----------- 2. wind amplitude
1474          elif ((zvarname in ['UV','uv','uvmet']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1475              all_var=windamplitude(znc,'amplitude')
1476          elif ((zvarname in ['hodograph','hodograph_2']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1477              plot_x, plot_y = windamplitude(znc,zvarname)
1478              if plot_x is not None: all_var=plot_x # dummy
1479              else: all_var=plot_y ; plot_x = None ; plot_y = None # Hodograph type 2 is not 'xy' mode
1480          elif ((zvarname in ['slopexy','SLOPEXY']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1481              all_var=slopeamplitude(znc)
1482          ### ------------ 3. Near surface instability
1483          elif ((zvarname in ['DELTAT','deltat']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1484              all_var=deltat0t1(znc)
1485          ### ------------ 4. Enrichment factor
1486          elif ((ztypefile in ['gcm']) and (zvarname in ['enfact'])):
1487              all_var=enrichment_factor(znc,ylon,ylat,ytime)
1488          ### ------------ 5. teta -> temp
1489          elif ((ztypefile in ['meso']) and (zvarname in ['tk']) and ('tk' not in znc.variables.keys())):
1490              all_var=teta_to_tk(znc)
1491          else:
1492          ### -----------  999. Normal case
1493              all_var = getfield(znc,zvarname)
1494          if analysis is not None:
1495             if analysis in ['histo','density','histodensity']: plot_y=all_var ; plot_x = plot_y
1496             elif analysis == 'fft': plot_y, plot_x = spectrum(all_var,ytime,yalt,ylat,ylon) ; all_var = plot_y
1497          return all_var, plot_x, plot_y
1498
1499# Author : A.C
1500# FFT is computed before reducefield voluntarily, because we dont want to compute
1501# ffts on averaged fields (which would kill all waves). Instead, we take the fft everywhere
1502# (which is not efficient but it is still ok) and then, make the average (if the user wants to)
1503def spectrum(var,time,vert,lat,lon):
1504    import numpy as np
1505    fft=np.fft.fft(var,axis=1)
1506    N=len(vert)
1507    step=(vert[1]-vert[0])*1000.
1508    print "step is: ",step
1509    fftfreq=np.fft.fftfreq(N,d=step)
1510    fftfreq=np.fft.fftshift(fftfreq) # spatial FFT => this is the wavenumber
1511    fft=np.fft.fftshift(fft)
1512    fftfreq = 1./fftfreq # => wavelength (div by 0 expected, don't panic)
1513    fft=np.abs(fft) # => amplitude spectrum
1514#    fft=np.abs(fft)**2 # => power spectrum
1515    return fft,fftfreq
1516
1517# Author : A.C.
1518# Computes temperature from potential temperature for mesoscale files, without the need to use API, i.e. using natural vertical grid
1519def teta_to_tk(nc):
1520    import numpy as np
1521    varinfile = nc.variables.keys() 
1522    p0=610.
1523    t0=220.
1524    r_cp=1./3.89419
1525    if "T" in varinfile: zteta=getfield(nc,'T')
1526    else: errormess("you need T in your file.")
1527    if "PTOT" in varinfile: zptot=getfield(nc,'PTOT')
1528    else: errormess("you need PTOT in your file.")
1529    zt=(zteta+220.)*(zptot/p0)**(r_cp)
1530    return zt
Note: See TracBrowser for help on using the repository browser.