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

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

Python. Bug correction for 1D plots. This bug was leading to completely wrong profiles when asking for a unidim plot from a testphys1d output file.
File size: 59.8 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 'time_counter' in nc.variables:       typefile = 'earthgcm'
61    elif hasattr(nc,'START_DATE'):             typefile = 'meso' 
62    elif 'HGT_M' in nc.variables:              typefile = 'geo'
63    elif hasattr(nc,'institution'):
64      if "European Centre" in getattr(nc,'institution'):  typefile = 'ecmwf'
65    return typefile
66
67## Author: AS
68def getfield (nc,var):
69    ## this allows to get much faster (than simply referring to nc.variables[var])
70    import numpy as np
71    dimension = len(nc.variables[var].dimensions)
72    #print "   Opening variable with", dimension, "dimensions ..."
73    if dimension == 2:    field = nc.variables[var][:,:]
74    elif dimension == 3:  field = nc.variables[var][:,:,:]
75    elif dimension == 4:  field = nc.variables[var][:,:,:,:]
76    elif dimension == 1:  field = nc.variables[var][:]
77    # if there are NaNs in the ncdf, they should be loaded as a masked array which will be
78    # recasted as a regular array later in reducefield
79    if (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
80       print "Warning: netcdf as nan values but is not loaded as a Masked Array."
81       print "recasting array type"
82       out=np.ma.masked_invalid(field)
83       out.set_fill_value([np.NaN])
84    else:
85    # missing values from zrecast or hrecast are -1e-33
86       masked=np.ma.masked_where(field < -1e30,field)
87       masked2=np.ma.masked_where(field > 1e35,field)
88       masked.set_fill_value([np.NaN]) ; masked2.set_fill_value([np.NaN])
89       mask = np.ma.getmask(masked) ; mask2 = np.ma.getmask(masked2)
90       if (True in np.array(mask)):
91          out=masked
92          print "Masked array... Missing value is NaN"
93       elif (True in np.array(mask2)):
94          out=masked2
95          print "Masked array... Missing value is NaN"
96#       else:
97#       # missing values from api are 1e36
98#          masked=np.ma.masked_where(field > 1e35,field)
99#          masked.set_fill_value([np.NaN])
100#          mask = np.ma.getmask(masked)
101#          if (True in np.array(mask)):out=masked
102#          else:out=field
103       else:out=field
104    return out
105
106## Author: AC
107# Compute the norm of the winds
108# The corresponding variable to call is UV or uvmet (to use api)
109def windamplitude (nc):
110    import numpy as np
111    varinfile = nc.variables.keys()
112    if "U" in varinfile: zu=getfield(nc,'U')
113    elif "Um" in varinfile: zu=getfield(nc,'Um')
114    if "V" in varinfile: zv=getfield(nc,'V')
115    elif "Vm" in varinfile: zv=getfield(nc,'Vm')
116    znt,znz,zny,znx = np.array(zu).shape
117    if "U" in varinfile:znx=znx-1
118    zuint = np.zeros([znt,znz,zny,znx])
119    zvint = np.zeros([znt,znz,zny,znx])
120    if "U" in varinfile:
121       for xx in np.arange(znx):      zuint[:,:,:,xx] = (zu[:,:,:,xx] + zu[:,:,:,xx+1])/2.
122       for yy in np.arange(zny):      zvint[:,:,yy,:] = (zv[:,:,yy,:] + zv[:,:,yy+1,:])/2.
123    else:
124       zuint=zu
125       zvint=zv
126    return np.sqrt(zuint**2 + zvint**2)
127
128## Author: AC
129# Compute the enrichment factor of non condensible gases
130# The corresponding variable to call is enfact
131def enrichment_factor(nc):
132    import numpy as np
133    varinfile = nc.variables.keys()
134    if "co2" in varinfile: co2=getfield(nc,'co2')
135    else: print "error, you need co2 var in your file"
136    if "ap" in varinfile: ap=getfield(nc,'ap')
137    else: print "error, you need ap var in your file"
138    if "bp" in varinfile: bp=getfield(nc,'bp')
139    else: print "error, you need bp var in your file"
140    if "ps" in varinfile: ps=getfield(nc,'ps')
141    else: print "error, you need ps var in your file"
142    dimension = len(nc.variables['co2'].dimensions)
143    if dimension == 3: 
144      znz,zny,znx = np.array(co2).shape
145      znt=1
146    elif dimension == 4: znt,znz,zny,znx = np.array(co2).shape
147    co2col = np.zeros([znt,zny,znx])
148    arcol = np.zeros([znt,zny,znx])
149    mmrarcol = np.zeros([znt,zny,znx])
150    meanar = np.zeros([znt])
151    pplev = np.zeros([znt,znz+1,zny,znx])
152    enfact = np.zeros([znt,zny,znx])
153    grav=3.72
154    for zz in np.arange(znz):
155      pplev[:,zz,:,:] = ap[zz]+bp[zz]*ps[:,:,:]
156    pplev[:,znz,:,:]=0.
157
158    for zz in np.arange(znz):
159      co2col[:,:,:] = co2col[:,:,:] + co2[:,zz,:,:]*(pplev[:,zz,:,:]-pplev[:,zz+1,:,:])/grav
160      arcol[:,:,:] = arcol[:,:,:] + (1.-co2[:,zz,:,:])*(pplev[:,zz,:,:]-pplev[:,zz+1,:,:])/grav
161
162    mmrarcol = arcol/(arcol + co2col)
163    for xx in np.arange(znx):
164      for yy in np.arange(zny):
165          meanar[:] = meanar[:] + mmrarcol[:,yy,xx]
166    meanar = meanar/(znx*zny)
167    for tt in np.arange(znt):
168      enfact[tt,:,:] =-(meanar[tt] - mmrarcol[tt,:,:])/meanar[tt]
169
170    return enfact
171
172## Author: AC
173# Compute the norm of the slope angles
174# The corresponding variable to call is SLOPEXY
175def slopeamplitude (nc):
176    import numpy as np
177    varinfile = nc.variables.keys()
178    if "slopex" in varinfile: zu=getfield(nc,'slopex')
179    elif "SLOPEX" in varinfile: zu=getfield(nc,'SLOPEX')
180    if "slopey" in varinfile: zv=getfield(nc,'slopey')
181    elif "SLOPEY" in varinfile: zv=getfield(nc,'SLOPEY')
182    znt,zny,znx = np.array(zu).shape
183    zuint = np.zeros([znt,zny,znx])
184    zvint = np.zeros([znt,zny,znx])
185    zuint=zu
186    zvint=zv
187    return np.sqrt(zuint**2 + zvint**2)
188
189## Author: AC
190# Compute the temperature difference between surface and first level.
191# API is automatically called to get TSURF and TK.
192# The corresponding variable to call is DELTAT
193def deltat0t1 (nc):
194    import numpy as np
195    varinfile = nc.variables.keys()
196    if "tsurf" in varinfile: zu=getfield(nc,'tsurf')
197    elif "TSURF" in varinfile: zu=getfield(nc,'TSURF')
198    if "tk" in varinfile: zv=getfield(nc,'tk')
199    elif "TK" in varinfile: zv=getfield(nc,'TK')
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        if "longitude" in nc.dimensions:  dalon = "longitude"
647        elif "lon" in nc.dimensions:      dalon = "lon"
648        if "latitude" in nc.dimensions:   dalat = "latitude"
649        elif "lat" in nc.dimensions:      dalat = "lat"
650        [lon2d,lat2d] = getcoord2d(nc,nlat=dalat,nlon=dalon,is1d=True)
651    elif typefile in ['geo']:
652        [lon2d,lat2d] = getcoord2d(nc,nlat='XLAT_M',nlon='XLONG_M')
653    return lon2d,lat2d   
654
655## Author: AS
656def getcoord2d (nc,nlat='XLAT',nlon='XLONG',is1d=False):
657    import numpy as np
658    if is1d:
659        lat = nc.variables[nlat][:]
660        lon = nc.variables[nlon][:]
661        [lon2d,lat2d] = np.meshgrid(lon,lat)
662    else:
663        lat = nc.variables[nlat][0,:,:]
664        lon = nc.variables[nlon][0,:,:]
665        [lon2d,lat2d] = [lon,lat]
666    return lon2d,lat2d
667
668## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
669def smooth1d(x,window_len=11,window='hanning'):
670    import numpy
671    """smooth the data using a window with requested size.
672    This method is based on the convolution of a scaled window with the signal.
673    The signal is prepared by introducing reflected copies of the signal
674    (with the window size) in both ends so that transient parts are minimized
675    in the begining and end part of the output signal.
676    input:
677        x: the input signal
678        window_len: the dimension of the smoothing window; should be an odd integer
679        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
680            flat window will produce a moving average smoothing.
681    output:
682        the smoothed signal
683    example:
684    t=linspace(-2,2,0.1)
685    x=sin(t)+randn(len(t))*0.1
686    y=smooth(x)
687    see also:
688    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
689    scipy.signal.lfilter
690    TODO: the window parameter could be the window itself if an array instead of a string   
691    """
692    x = numpy.array(x)
693    if x.ndim != 1:
694        raise ValueError, "smooth only accepts 1 dimension arrays."
695    if x.size < window_len:
696        raise ValueError, "Input vector needs to be bigger than window size."
697    if window_len<3:
698        return x
699    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
700        raise ValueError, "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"
701    s=numpy.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
702    #print(len(s))
703    if window == 'flat': #moving average
704        w=numpy.ones(window_len,'d')
705    else:
706        w=eval('numpy.'+window+'(window_len)')
707    y=numpy.convolve(w/w.sum(),s,mode='valid')
708    return y
709
710## Author: AS
711def smooth (field, coeff):
712        ## actually blur_image could work with different coeff on x and y
713        if coeff > 1:   result = blur_image(field,int(coeff))
714        else:           result = field
715        return result
716
717## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
718def gauss_kern(size, sizey=None):
719        import numpy as np
720        # Returns a normalized 2D gauss kernel array for convolutions
721        size = int(size)
722        if not sizey:
723                sizey = size
724        else:
725                sizey = int(sizey)
726        x, y = np.mgrid[-size:size+1, -sizey:sizey+1]
727        g = np.exp(-(x**2/float(size)+y**2/float(sizey)))
728        return g / g.sum()
729
730## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
731def blur_image(im, n, ny=None) :
732        from scipy.signal import convolve
733        # blurs the image by convolving with a gaussian kernel of typical size n.
734        # The optional keyword argument ny allows for a different size in the y direction.
735        g = gauss_kern(n, sizey=ny)
736        improc = convolve(im, g, mode='same')
737        return improc
738
739## Author: AS
740def getwinddef (nc):   
741    ###
742    varinfile = nc.variables.keys()
743    if 'Um' in varinfile:   [uchar,vchar] = ['Um','Vm'] #; print "this is API meso file"
744    elif 'U' in varinfile:  [uchar,vchar] = ['U','V']   #; print "this is RAW meso file"
745    elif 'u' in varinfile:  [uchar,vchar] = ['u','v']   #; print "this is GCM file"
746     ### you can add choices here !
747    else:                   [uchar,vchar] = ['not found','not found']
748    ###
749    if uchar in ['U']:         metwind = False ## geometrical (wrt grid)
750    else:                      metwind = True  ## meteorological (zon/mer)
751    if metwind is False:       print "Not using meteorological winds. You trust numerical grid as being (x,y)"
752    ###
753    return uchar,vchar,metwind
754
755## Author: AS
756def vectorfield (u, v, x, y, stride=3, scale=15., factor=250., color='black', csmooth=1, key=True):
757    ## scale regle la reference du vecteur
758    ## factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
759    import  matplotlib.pyplot               as plt
760    import  numpy                           as np
761    #posx = np.min(x) - np.std(x) / 10.
762    #posy = np.min(y) - np.std(y) / 10.
763    posx = np.min(x) 
764    posy = np.min(y) - 4.*np.std(y) / 10.
765    u = smooth(u,csmooth)
766    v = smooth(v,csmooth)
767    widthvec = 0.003 #0.005 #0.003
768    q = plt.quiver( x[::stride,::stride],\
769                    y[::stride,::stride],\
770                    u[::stride,::stride],\
771                    v[::stride,::stride],\
772                    angles='xy',color=color,pivot='middle',\
773                    scale=factor,width=widthvec )
774    if color in ['white','yellow']:     kcolor='black'
775    else:                               kcolor=color
776    if key: p = plt.quiverkey(q,posx,posy,scale,\
777                   str(int(scale)),coordinates='data',color=kcolor,labelpos='S',labelsep = 0.03)
778    return 
779
780## Author: AS
781def display (name):
782    from os import system
783    system("display "+name+" > /dev/null 2> /dev/null &")
784    return name
785
786## Author: AS
787def findstep (wlon):
788    steplon = int((wlon[1]-wlon[0])/4.)  #3
789    step = 120.
790    while step > steplon and step > 15. :       step = step / 2.
791    if step <= 15.:
792        while step > steplon and step > 5.  :   step = step - 5.
793    if step <= 5.:
794        while step > steplon and step > 1.  :   step = step - 1.
795    if step <= 1.:
796        step = 1. 
797    return step
798
799## Author: AS
800def define_proj (char,wlon,wlat,back=None,blat=None,blon=None):
801    from    mpl_toolkits.basemap            import Basemap
802    import  numpy                           as np
803    import  matplotlib                      as mpl
804    from mymath import max
805    meanlon = 0.5*(wlon[0]+wlon[1])
806    meanlat = 0.5*(wlat[0]+wlat[1])
807    zewidth = np.abs(wlon[0]-wlon[1])*60000.*np.cos(3.14*meanlat/180.)
808    zeheight = np.abs(wlat[0]-wlat[1])*60000.
809    if blat is None:
810        ortholat=meanlat
811        if   wlat[0] >= 80.:   blat =  -40. 
812        elif wlat[1] <= -80.:  blat = -40.
813        elif wlat[1] >= 0.:    blat = wlat[0] 
814        elif wlat[0] <= 0.:    blat = wlat[1]
815    else:  ortholat=blat
816    if blon is None:  ortholon=meanlon
817    else:             ortholon=blon
818    h = 50.  ## en km
819    radius = 3397200.
820    #print meanlat, meanlon
821    if   char == "cyl":     m = Basemap(rsphere=radius,projection='cyl',\
822                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
823    elif char == "moll":    m = Basemap(rsphere=radius,projection='moll',lon_0=meanlon)
824    elif char == "ortho":   m = Basemap(rsphere=radius,projection='ortho',lon_0=ortholon,lat_0=ortholat)
825    elif char == "lcc":     m = Basemap(rsphere=radius,projection='lcc',lat_1=meanlat,lat_0=meanlat,lon_0=meanlon,\
826                              width=zewidth,height=zeheight)
827                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
828    elif char == "npstere": m = Basemap(rsphere=radius,projection='npstere', boundinglat=blat, lon_0=0.)
829    elif char == "spstere": m = Basemap(rsphere=radius,projection='spstere', boundinglat=blat, lon_0=180.)
830    elif char == "nplaea":  m = Basemap(rsphere=radius,projection='nplaea', boundinglat=wlat[0], lon_0=meanlon)
831    elif char == "laea":    m = Basemap(rsphere=radius,projection='laea',lon_0=meanlon,lat_0=meanlat,lat_ts=meanlat,\
832                              width=zewidth,height=zeheight)
833                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
834    elif char == "nsper":   m = Basemap(rsphere=radius,projection='nsper',lon_0=meanlon,lat_0=meanlat,satellite_height=h*1000.)
835    elif char == "merc":    m = Basemap(rsphere=radius,projection='merc',lat_ts=0.,\
836                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
837    fontsizemer = int(mpl.rcParams['font.size']*3./4.)
838    if char in ["cyl","lcc","merc","nsper","laea"]:   step = findstep(wlon)
839    else:                                             step = 10.
840    steplon = step*2.
841    zecolor ='grey'
842    zelinewidth = 1
843    zelatmax = 80
844    if meanlat > 75.: zelatmax = 90. ; step = step/2.
845    # to show gcm grid:
846    #zecolor = 'r'
847    #zelinewidth = 1
848    #step = 180./48.
849    #steplon = 360./64.
850    #zelatmax = 90. - step/3
851    if char not in ["moll"]:
852        m.drawmeridians(np.r_[-180.:180.:steplon], labels=[0,0,0,1], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
853        m.drawparallels(np.r_[-90.:90.:step], labels=[1,0,0,0], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
854    if back: m.warpimage(marsmap(back),scale=0.75)
855            #if not back:
856            #    if not var:                                        back = "mola"    ## if no var:         draw mola
857            #    elif typefile in ['mesoapi','meso','geo'] \
858            #       and proj not in ['merc','lcc','nsper','laea']:  back = "molabw"  ## if var but meso:   draw molabw
859            #    else:                                              pass             ## else:              draw None
860    return m
861
862## Author: AS
863#### test temporaire
864def putpoints (map,plot):
865    #### from http://www.scipy.org/Cookbook/Matplotlib/Maps
866    # lat/lon coordinates of five cities.
867    lats = [18.4]
868    lons = [-134.0]
869    points=['Olympus Mons']
870    # compute the native map projection coordinates for cities.
871    x,y = map(lons,lats)
872    # plot filled circles at the locations of the cities.
873    map.plot(x,y,'bo')
874    # plot the names of those five cities.
875    wherept = 0 #1000 #50000
876    for name,xpt,ypt in zip(points,x,y):
877       plot.text(xpt+wherept,ypt+wherept,name)
878    ## le nom ne s'affiche pas...
879    return
880
881## Author: AS
882def calculate_bounds(field,vmin=None,vmax=None):
883    import numpy as np
884    from mymath import max,min,mean
885    ind = np.where(field < 9e+35)
886    fieldcalc = field[ ind ] # la syntaxe compacte ne marche si field est un tuple
887    ###
888    dev = np.std(fieldcalc)*3.0
889    ###
890    if vmin is None:  zevmin = mean(fieldcalc) - dev
891    else:             zevmin = vmin
892    ###
893    if vmax is None:  zevmax = mean(fieldcalc) + dev
894    else:             zevmax = vmax
895    if vmin == vmax:
896                      zevmin = mean(fieldcalc) - dev  ### for continuity
897                      zevmax = mean(fieldcalc) + dev  ### for continuity           
898    ###
899    if zevmin < 0. and min(fieldcalc) > 0.: zevmin = 0.
900    print "BOUNDS field ", min(fieldcalc), max(fieldcalc), " //// adopted", zevmin, zevmax
901    return zevmin, zevmax
902
903## Author: AS
904def bounds(what_I_plot,zevmin,zevmax):
905    from mymath import max,min,mean
906    ### might be convenient to add the missing value in arguments
907    #what_I_plot[ what_I_plot < zevmin ] = zevmin#*(1. + 1.e-7)
908    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1. - 1.e-7) ] = zevmin*(1. - 1.e-7)
909    else:          what_I_plot[ what_I_plot < zevmin*(1. + 1.e-7) ] = zevmin*(1. + 1.e-7)
910    #print "NEW MIN ", min(what_I_plot)
911    what_I_plot[ what_I_plot > 9e+35  ] = -9e+35
912    what_I_plot[ what_I_plot > zevmax ] = zevmax*(1. - 1.e-7)
913    #print "NEW MAX ", max(what_I_plot)
914    return what_I_plot
915
916## Author: AS
917def nolow(what_I_plot):
918    from mymath import max,min
919    lim = 0.15*0.5*(abs(max(what_I_plot))+abs(min(what_I_plot)))
920    print "NO PLOT BELOW VALUE ", lim
921    what_I_plot [ abs(what_I_plot) < lim ] = 1.e40 
922    return what_I_plot
923
924
925## Author : AC
926def hole_bounds(what_I_plot,zevmin,zevmax):
927    import numpy as np
928    zi=0
929    for i in what_I_plot:
930        zj=0
931        for j in i:
932            if ((j < zevmin) or (j > zevmax)):what_I_plot[zi,zj]=np.NaN
933            zj=zj+1
934        zi=zi+1
935
936    return what_I_plot
937
938## Author: AS
939def zoomset (wlon,wlat,zoom):
940    dlon = abs(wlon[1]-wlon[0])/2.
941    dlat = abs(wlat[1]-wlat[0])/2.
942    [wlon,wlat] = [ [wlon[0]+zoom*dlon/100.,wlon[1]-zoom*dlon/100.],\
943                    [wlat[0]+zoom*dlat/100.,wlat[1]-zoom*dlat/100.] ]
944    print "ZOOM %",zoom,wlon,wlat
945    return wlon,wlat
946
947## Author: AS
948def fmtvar (whichvar="def"):
949    fmtvar    =     { \
950             "MIXED":        "%.0f",\
951             "UPDRAFT":      "%.0f",\
952             "DOWNDRAFT":    "%.0f",\
953             "TK":           "%.0f",\
954             "T":            "%.0f",\
955             #"ZMAX_TH":      "%.0f",\
956             #"WSTAR":        "%.0f",\
957             # Variables from TES ncdf format
958             "T_NADIR_DAY":  "%.0f",\
959             "T_NADIR_NIT":  "%.0f",\
960             # Variables from tes.py ncdf format
961             "TEMP_DAY":     "%.0f",\
962             "TEMP_NIGHT":   "%.0f",\
963             # Variables from MCS and mcs.py ncdf format
964             "DTEMP":        "%.0f",\
965             "NTEMP":        "%.0f",\
966             "DNUMBINTEMP":  "%.0f",\
967             "NNUMBINTEMP":  "%.0f",\
968             # other stuff
969             "TPOT":         "%.0f",\
970             "TSURF":        "%.0f",\
971             "U_OUT1":       "%.0f",\
972             "T_OUT1":       "%.0f",\
973             "def":          "%.1e",\
974             "PTOT":         "%.0f",\
975             "HGT":          "%.1e",\
976             "USTM":         "%.2f",\
977             "HFX":          "%.0f",\
978             "ICETOT":       "%.1e",\
979             "TAU_ICE":      "%.2f",\
980             "TAUICE":       "%.2f",\
981             "VMR_ICE":      "%.1e",\
982             "MTOT":         "%.1f",\
983             "ANOMALY":      "%.1f",\
984             "W":            "%.1f",\
985             "WMAX_TH":      "%.1f",\
986             "WSTAR":        "%.1f",\
987             "QSURFICE":     "%.0f",\
988             "UM":           "%.0f",\
989             "WIND":         "%.0f",\
990             "UVMET":         "%.0f",\
991             "UV":         "%.0f",\
992             "ALBBARE":      "%.2f",\
993             "TAU":          "%.1f",\
994             "CO2":          "%.2f",\
995             "ENFACT":       "%.1f",\
996             #### T.N.
997             "TEMP":         "%.0f",\
998             "VMR_H2OICE":   "%.0f",\
999             "VMR_H2OVAP":   "%.0f",\
1000             "TAUTES":       "%.2f",\
1001             "TAUTESAP":     "%.2f",\
1002
1003                    }
1004    if "TSURF" in whichvar: whichvar = "TSURF"
1005    if whichvar not in fmtvar:
1006        whichvar = "def"
1007    return fmtvar[whichvar]
1008
1009## Author: AS
1010####################################################################################################################
1011### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
1012def defcolorb (whichone="def"):
1013    whichcolorb =    { \
1014             "def":          "spectral",\
1015             "HGT":          "spectral",\
1016             "HGT_M":        "spectral",\
1017             "TK":           "gist_heat",\
1018             "TPOT":         "Paired",\
1019             "TSURF":        "RdBu_r",\
1020             "QH2O":         "PuBu",\
1021             "USTM":         "YlOrRd",\
1022             "WIND":         "YlOrRd",\
1023             #"T_nadir_nit":  "RdBu_r",\
1024             #"T_nadir_day":  "RdBu_r",\
1025             "HFX":          "RdYlBu",\
1026             "ICETOT":       "YlGnBu_r",\
1027             #"MTOT":         "PuBu",\
1028             "CCNQ":         "YlOrBr",\
1029             "CCNN":         "YlOrBr",\
1030             "TEMP":         "Jet",\
1031             "TAU_ICE":      "Blues",\
1032             "TAUICE":       "Blues",\
1033             "VMR_ICE":      "Blues",\
1034             "W":            "jet",\
1035             "WMAX_TH":      "spectral",\
1036             "ANOMALY":      "RdBu_r",\
1037             "QSURFICE":     "hot_r",\
1038             "ALBBARE":      "spectral",\
1039             "TAU":          "YlOrBr_r",\
1040             "CO2":          "YlOrBr_r",\
1041             #### T.N.
1042             "MTOT":         "spectral",\
1043             "H2O_ICE_S":    "RdBu",\
1044             "VMR_H2OICE":   "PuBu",\
1045             "VMR_H2OVAP":   "PuBu",\
1046             "WATERCAPTAG":  "Blues",\
1047                     }
1048#W --> spectral ou jet
1049#spectral BrBG RdBu_r
1050    #print "predefined colorbars"
1051    if "TSURF" in whichone: whichone = "TSURF"
1052    if whichone not in whichcolorb:
1053        whichone = "def"
1054    return whichcolorb[whichone]
1055
1056## Author: AS
1057def definecolorvec (whichone="def"):
1058        whichcolor =    { \
1059                "def":          "black",\
1060                "vis":          "yellow",\
1061                "vishires":     "yellow",\
1062                "molabw":       "yellow",\
1063                "mola":         "black",\
1064                "gist_heat":    "white",\
1065                "hot":          "tk",\
1066                "gist_rainbow": "black",\
1067                "spectral":     "black",\
1068                "gray":         "red",\
1069                "PuBu":         "black",\
1070                        }
1071        if whichone not in whichcolor:
1072                whichone = "def"
1073        return whichcolor[whichone]
1074
1075## Author: AS
1076def marsmap (whichone="vishires"):
1077        from os import uname
1078        mymachine = uname()[1]
1079        ### not sure about speed-up with this method... looks the same
1080        if "lmd.jussieu.fr" in mymachine:   domain = "/u/aslmd/WWW/maps/"
1081        elif "aymeric-laptop" in mymachine: domain = "/home/aymeric/Dropbox/Public/"
1082        else:                               domain = "http://www.lmd.jussieu.fr/~aslmd/maps/"
1083        whichlink =     { \
1084                #"vis":         "http://maps.jpl.nasa.gov/pix/mar0kuu2.jpg",\
1085                #"vishires":    "http://www.lmd.jussieu.fr/~aslmd/maps/MarsMap_2500x1250.jpg",\
1086                #"geolocal":    "http://dl.dropbox.com/u/11078310/geolocal.jpg",\
1087                #"mola":        "http://www.lns.cornell.edu/~seb/celestia/mars-mola-2k.jpg",\
1088                #"molabw":      "http://dl.dropbox.com/u/11078310/MarsElevation_2500x1250.jpg",\
1089                "thermalday":   domain+"thermalday.jpg",\
1090                "thermalnight": domain+"thermalnight.jpg",\
1091                "tesalbedo":    domain+"tesalbedo.jpg",\
1092                "vis":         domain+"mar0kuu2.jpg",\
1093                "vishires":    domain+"MarsMap_2500x1250.jpg",\
1094                "geolocal":    domain+"geolocal.jpg",\
1095                "mola":        domain+"mars-mola-2k.jpg",\
1096                "molabw":      domain+"MarsElevation_2500x1250.jpg",\
1097                "clouds":      "http://www.johnstonsarchive.net/spaceart/marswcloudmap.jpg",\
1098                "jupiter":     "http://www.mmedia.is/~bjj/data/jupiter_css/jupiter_css.jpg",\
1099                "jupiter_voy": "http://www.mmedia.is/~bjj/data/jupiter/jupiter_vgr2.jpg",\
1100                #"bw":          domain+"EarthElevation_2500x1250.jpg",\
1101                "bw":          "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg",\
1102                "contrast":    "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg",\
1103                "nice":        "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg",\
1104                "blue":        "http://eoimages.gsfc.nasa.gov/ve/2430/land_ocean_ice_2048.jpg",\
1105                "blueclouds":  "http://eoimages.gsfc.nasa.gov/ve/2431/land_ocean_ice_cloud_2048.jpg",\
1106                "justclouds":  "http://eoimages.gsfc.nasa.gov/ve/2432/cloud_combined_2048.jpg",\
1107                        }
1108        ### see http://www.mmedia.is/~bjj/planetary_maps.html
1109        if whichone not in whichlink: 
1110                print "marsmap: choice not defined... you'll get the default one... "
1111                whichone = "vishires" 
1112        return whichlink[whichone]
1113
1114#def earthmap (whichone):
1115#       if   whichone == "contrast":    whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg"
1116#       elif whichone == "bw":          whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg"
1117#       elif whichone == "nice":        whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg"
1118#       return whichlink
1119
1120## Author: AS
1121def latinterv (area="Whole"):
1122    list =    { \
1123        "Europe":                [[ 20., 80.],[- 50.,  50.]],\
1124        "Central_America":       [[-10., 40.],[ 230., 300.]],\
1125        "Africa":                [[-20., 50.],[- 50.,  50.]],\
1126        "Whole":                 [[-90., 90.],[-180., 180.]],\
1127        "Southern_Hemisphere":   [[-90., 60.],[-180., 180.]],\
1128        "Northern_Hemisphere":   [[-60., 90.],[-180., 180.]],\
1129        "Tharsis":               [[-30., 60.],[-170.,- 10.]],\
1130        "Whole_No_High":         [[-60., 60.],[-180., 180.]],\
1131        "Chryse":                [[-60., 60.],[- 60.,  60.]],\
1132        "North_Pole":            [[ 50., 90.],[-180., 180.]],\
1133        "Close_North_Pole":      [[ 75., 90.],[-180., 180.]],\
1134        "Far_South_Pole":        [[-90.,-40.],[-180., 180.]],\
1135        "South_Pole":            [[-90.,-50.],[-180., 180.]],\
1136        "Close_South_Pole":      [[-90.,-75.],[-180., 180.]],\
1137        "Sirenum_Crater_large":  [[-46.,-34.],[-166.,-151.]],\
1138        "Sirenum_Crater_small":  [[-36.,-26.],[-168.,-156.]],\
1139        "Rupes":                 [[ 72., 90.],[-120.,- 20.]],\
1140              }
1141    if area not in list:   area = "Whole"
1142    [olat,olon] = list[area]
1143    return olon,olat
1144
1145## Author: TN
1146def separatenames (name):
1147  from numpy import concatenate
1148  # look for comas in the input name to separate different names (files, variables,etc ..)
1149  if name is None:
1150     names = None
1151  else:
1152    names = []
1153    stop = 0
1154    currentname = name
1155    while stop == 0:
1156      indexvir = currentname.find(',')
1157      if indexvir == -1:
1158        stop = 1
1159        name1 = currentname
1160      else:
1161        name1 = currentname[0:indexvir]
1162      names = concatenate((names,[name1]))
1163      currentname = currentname[indexvir+1:len(currentname)]
1164  return names
1165
1166
1167## Author: TN
1168def readslices(saxis):
1169  from numpy import empty
1170  if saxis == None:
1171     zesaxis = None
1172  else:
1173     zesaxis = empty((len(saxis),2))
1174     for i in range(len(saxis)):
1175        a = separatenames(saxis[i])
1176        if len(a) == 1:
1177           zesaxis[i,:] = float(a[0])
1178        else:
1179           zesaxis[i,0] = float(a[0])
1180           zesaxis[i,1] = float(a[1])
1181           
1182  return zesaxis
1183
1184## Author: TN
1185def readdata(data,datatype,coord1,coord2):
1186## Read sparse data
1187  from numpy import empty
1188  if datatype == 'txt':
1189     if len(data[coord1].shape) == 1:
1190         return data[coord1][:]
1191     elif len(data[coord1].shape) == 2:
1192         return data[coord1][:,int(coord2)-1]
1193     else:
1194         errormess('error in readdata')
1195  elif datatype == 'sav':
1196     return data[coord1][coord2]
1197  else:
1198     errormess(datatype+' type is not supported!')
1199
1200
1201## Author: AS
1202def bidimfind(lon2d,lat2d,vlon,vlat,file=None):
1203   import numpy as np
1204   import matplotlib.pyplot as mpl
1205   if vlat is None:    array = (lon2d - vlon)**2
1206   elif vlon is None:  array = (lat2d - vlat)**2
1207   else:               array = (lon2d - vlon)**2 + (lat2d - vlat)**2
1208   idy,idx = np.unravel_index( np.argmin(array), lon2d.shape )
1209   if vlon is not None:
1210      if (np.abs(lon2d[idy,idx]-vlon)) > 5: errormess("longitude not found ",printvar=lon2d)
1211   if vlat is not None:
1212      if (np.abs(lat2d[idy,idx]-vlat)) > 5: errormess("latitude not found ",printvar=lat2d)
1213   if file is not None:
1214      print idx,idy,lon2d[idy,idx],vlon
1215      print idx,idy,lat2d[idy,idx],vlat
1216      var = file.variables["HGT"][:,:,:]
1217      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)
1218      mpl.show()
1219   return idy,idx
1220
1221## Author: TN
1222def getsindex(saxis,index,axis):
1223# input  : all the desired slices and the good index
1224# output : all indexes to be taken into account for reducing field
1225  import numpy as np
1226  if ( np.array(axis).ndim == 2):
1227      axis = axis[:,0]
1228  if saxis is None:
1229      zeindex = None
1230  else:
1231      aaa = int(np.argmin(abs(saxis[index,0] - axis)))
1232      bbb = int(np.argmin(abs(saxis[index,1] - axis)))
1233      [imin,imax] = np.sort(np.array([aaa,bbb]))
1234      zeindex = np.array(range(imax-imin+1))+imin
1235      # because -180 and 180 are the same point in longitude,
1236      # we get rid of one for averaging purposes.
1237      if axis[imin] == -180 and axis[imax] == 180:
1238         zeindex = zeindex[0:len(zeindex)-1]
1239         print "INFO: whole longitude averaging asked, so last point is not taken into account."
1240  return zeindex
1241     
1242## Author: TN
1243def define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,indextime,what_I_plot,dim0,vertmode):
1244# Purpose of define_axis is to find x and y axis scales in a smart way
1245# x axis priority: 1/time 2/lon 3/lat 4/vertical
1246# To be improved !!!...
1247   from numpy import array,swapaxes
1248   x = None
1249   y = None
1250   count = 0
1251   what_I_plot = array(what_I_plot)
1252   shape = what_I_plot.shape
1253   if indextime is None and len(time) > 1:
1254      print "AXIS is time"
1255      x = time
1256      count = count+1
1257   if indexlon is None and len(lon) > 1:
1258      print "AXIS is lon"
1259      if count == 0: x = lon
1260      else: y = lon
1261      count = count+1
1262   if indexlat is None and len(lat) > 1:
1263      print "AXIS is lat"
1264      if count == 0: x = lat
1265      else: y = lat
1266      count = count+1
1267   if indexvert is None and ((dim0 == 4) or (y is None)):
1268      print "AXIS is vert"
1269      if vertmode == 0: # vertical axis is as is (GCM grid)
1270         if count == 0: x=range(len(vert))
1271         else: y=range(len(vert))
1272         count = count+1
1273      else: # vertical axis is in kms
1274         if count == 0: x = vert
1275         else: y = vert
1276         count = count+1
1277   x = array(x)
1278   y = array(y)
1279   print "CHECK SHAPE: what_I_plot, x, y", what_I_plot.shape, x.shape, y.shape
1280   if len(shape) == 1:
1281       if shape[0] != len(x):           print "WARNING: shape[0] != len(x). Correcting." ; what_I_plot = what_I_plot[0:len(x)]
1282       if len(y.shape) > 0:             y = ()
1283   elif len(shape) == 2:
1284       if shape[1] == len(y) and shape[0] == len(x) and shape[0] != shape[1]:
1285           print "INFO: swapaxes: ",what_I_plot.shape,shape ; what_I_plot = swapaxes(what_I_plot,0,1)
1286       else:
1287           if shape[0] != len(y):       print "WARNING: shape[0] != len(y). Correcting." ; what_I_plot = what_I_plot[0:len(y),:]
1288           elif shape[1] != len(x):     print "WARNING: shape[1] != len(x). Correcting." ; what_I_plot = what_I_plot[:,0:len(x)]
1289   elif len(shape) == 3:
1290       if vertmode < 0: print "not supported. must check array dimensions at some point. not difficult to implement though."
1291   return what_I_plot,x,y
1292
1293# Author: TN + AS
1294def determineplot(slon, slat, svert, stime):
1295    nlon = 1 # number of longitudinal slices -- 1 is None
1296    nlat = 1
1297    nvert = 1
1298    ntime = 1
1299    nslices = 1
1300    if slon is not None:
1301        nslices = nslices*len(slon)
1302        nlon = len(slon)
1303    if slat is not None:
1304        nslices = nslices*len(slat)
1305        nlat = len(slat)
1306    if svert is not None:
1307        nslices = nslices*len(svert)
1308        nvert = len(svert)
1309    if stime is not None:
1310        nslices = nslices*len(stime)
1311        ntime = len(stime)
1312    #else:
1313    #    nslices = 2 
1314    mapmode = 0
1315    if slon is None and slat is None:
1316       mapmode = 1 # in this case we plot a map, with the given projection
1317
1318    return nlon, nlat, nvert, ntime, mapmode, nslices
1319
1320## Author: AC
1321## Reduce complexity of main script by moving the contour part here. Also allow to call it from elsewhere
1322## which can be usefull
1323#
1324#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):
1325#      from matplotlib.pyplot import contour, plot, clabel
1326#      import numpy as np
1327#      #what_I_plot = what_I_plot*mult
1328#      if not error:
1329#         if mapmode == 1:
1330#             if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
1331#         zevmin, zevmax = calculate_bounds(what_I_plot)
1332#         zelevels = np.linspace(zevmin,zevmax,ticks) #20)
1333#         if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
1334#         if mapmode == 0:
1335#             #if typefile in ['mesoideal']:    what_I_plot = dumpbdy(what_I_plot,0,stag='W')
1336#             itime=indextime
1337#             if len(what_I_plot.shape) is 3: itime=[0]
1338#             what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
1339#                   itime,what_I_plot, len(ze_var2.shape),vertmode)
1340#         ### If we plot a 2-D field
1341#         if len(what_I_plot.shape) is 2:
1342#             #zelevels=[1.]
1343#             if mapmode == 0:cs = contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
1344#             elif mapmode == 1:cs = m.contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
1345#             #clabel(cs,zelevels,inline=3,fmt='%1.1f',fontsize=7)
1346#         ### If we plot a 1-D field
1347#         elif len(what_I_plot.shape) is 1:
1348#             plot(what_I_plot,x)
1349#      else:
1350#         errormess("There is an error in reducing field !")
1351#      return error
1352
1353## Author : AS
1354def maplatlon( lon,lat,field,\
1355               proj="cyl",colorb="jet",ndiv=10,zeback="molabw",trans=0.6,title="",\
1356               vecx=None,vecy=None,stride=2 ):
1357    ### an easy way to map a field over lat/lon grid
1358    import numpy as np
1359    import matplotlib.pyplot as mpl
1360    from matplotlib.cm import get_cmap
1361    ## get lon and lat in 2D version. get lat/lon intervals
1362    numdim = len(np.array(lon).shape)
1363    if numdim == 2:     [lon2d,lat2d] = [lon,lat]
1364    elif numdim == 1:   [lon2d,lat2d] = np.meshgrid(lon,lat)
1365    else:               errormess("lon and lat arrays must be 1D or 2D")
1366    [wlon,wlat] = latinterv()
1367    ## define projection and background. define x and y given the projection
1368    m = define_proj(proj,wlon,wlat,back=zeback,blat=None,blon=None)
1369    x, y = m(lon2d, lat2d)
1370    ## define field. bound field.
1371    what_I_plot = np.transpose(field)
1372    zevmin, zevmax = calculate_bounds(what_I_plot)  ## vmin=min(what_I_plot_frame), vmax=max(what_I_plot_frame))
1373    what_I_plot = bounds(what_I_plot,zevmin,zevmax)
1374    ## define contour field levels. define color palette
1375    ticks = ndiv + 1
1376    zelevels = np.linspace(zevmin,zevmax,ticks)
1377    palette = get_cmap(name=colorb)
1378    ## contour field
1379    m.contourf( x, y, what_I_plot, zelevels, cmap = palette, alpha = trans )
1380    ## draw colorbar
1381    if proj in ['moll','cyl']:        zeorientation="horizontal" ; zepad = 0.07
1382    else:                             zeorientation="vertical" ; zepad = 0.03
1383    #daformat = fmtvar(fvar.upper())
1384    daformat = "%.0f"
1385    zecb = mpl.colorbar( fraction=0.05,pad=zepad,format=daformat,orientation=zeorientation,\
1386                 ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional' ) 
1387    ## give a title
1388    if zeorientation == "horizontal": zecb.ax.set_xlabel(title)
1389    else:                             ptitle(title)
1390    ## draw vector
1391    if vecx is not None and vecy is not None:
1392       [vecx_frame,vecy_frame] = m.rotate_vector( np.transpose(vecx), np.transpose(vecy), lon2d, lat2d ) ## for metwinds
1393       vectorfield(vecx_frame, vecy_frame, x, y, stride=stride, csmooth=2,\
1394                                             scale=30., factor=500., color=definecolorvec(colorb), key=True)
1395    ## scale regle la reference du vecteur. factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
1396    return
Note: See TracBrowser for help on using the repository browser.