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

Last change on this file since 696 was 687, checked in by aslmd, 13 years ago

UTIL PYTHON 1. make planetoplot work for API-generated files directly 2. if there is a problem with areas the program keeps on moving and print a warning while setting mesharea and totalarea to 1 so that it is harmless

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