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

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

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