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

Last change on this file since 748 was 724, checked in by aslmd, 13 years ago

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