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

Last change on this file since 647 was 647, checked in by tnavarro, 13 years ago

Corrected bug in reducefield for masked arrays with grid area. Possibility to see stream function for a lat/vert slice with --stream option. Output of both data and axis with save txt option in 1D. Small bug corrected: title is the file name for multiple plots with multiple plots. Cleanup in myplot.py

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