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

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

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