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

Last change on this file since 817 was 817, checked in by aslmd, 12 years ago

UTIL PYTHON : improved projections. improved local time ticks. small bug fix with vertical coordinates. added a keyword nocolorb.

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