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

Last change on this file since 792 was 792, checked in by acolaitis, 12 years ago

PYTHON GRAPHICS

  • Added option --finddevil which finds the steepest pressure drop in the PSFC field of the file, and add a cross in it's position.

This new function is intended to provide lon and lat of a dust devil to planetoplot. This could be used to drive a slice plot at the corresponding location.

  • Removed dumpbdy for LES files (no relaxation zone)
  • Added --mark mode to tiled plot
  • Added routines in mymath to convert back and forth between greyscale Image objects and 2D arrays
File size: 70.2 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    print time
68    for i in range(len(time)-1):
69       if (time[i] > time[i+1]): a=i
70    if a >= 0 and a < (len(time)-1)/2.:
71       print "Sorry, time axis is not regular."
72       print "Contourf needs regular axis... recasting"
73       for i in range(a+1):
74          time[i]=time[i]-24.
75    if a >= 0 and a >= (len(time)-1)/2.:
76       print "Sorry, time axis is not regular."
77       print "Contourf needs regular axis... recasting"
78       for i in range((len(time)-1) - a):
79          time[a+1+i]=time[a+1+i]+24.
80    return time
81
82## Author: AS, AC, JL
83def whatkindfile (nc):
84    typefile = 'gcm' # default
85    if 'controle' in nc.variables:             typefile = 'gcm'
86    elif 'phisinit' in nc.variables:           typefile = 'gcm'
87    elif 'phis' in nc.variables:               typefile = 'gcm'
88    elif 'time_counter' in nc.variables:       typefile = 'earthgcm'
89    elif hasattr(nc,'START_DATE'):             typefile = 'meso' 
90    elif 'HGT_M' in nc.variables:              typefile = 'geo'
91    elif hasattr(nc,'institution'):
92      if "European Centre" in getattr(nc,'institution'):  typefile = 'ecmwf'
93    return typefile
94
95## Author: AS
96def getfield (nc,var):
97    ## this allows to get much faster (than simply referring to nc.variables[var])
98    import numpy as np
99    dimension = len(nc.variables[var].dimensions)
100    #print "   Opening variable with", dimension, "dimensions ..."
101    if dimension == 2:    field = nc.variables[var][:,:]
102    elif dimension == 3:  field = nc.variables[var][:,:,:]
103    elif dimension == 4:  field = nc.variables[var][:,:,:,:]
104    elif dimension == 1:  field = nc.variables[var][:]
105    # if there are NaNs in the ncdf, they should be loaded as a masked array which will be
106    # recasted as a regular array later in reducefield
107    if (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
108       print "Warning: netcdf as nan values but is not loaded as a Masked Array."
109       print "recasting array type"
110       out=np.ma.masked_invalid(field)
111       out.set_fill_value([np.NaN])
112    else:
113    # missing values from zrecast or hrecast are -1e-33
114       masked=np.ma.masked_where(field < -1e30,field)
115       masked2=np.ma.masked_where(field > 1e35,field)
116       masked.set_fill_value([np.NaN]) ; masked2.set_fill_value([np.NaN])
117       mask = np.ma.getmask(masked) ; mask2 = np.ma.getmask(masked2)
118       if (True in np.array(mask)):
119          out=masked
120          print "Masked array... Missing value is NaN"
121       elif (True in np.array(mask2)):
122          out=masked2
123          print "Masked array... Missing value is NaN"
124#       else:
125#       # missing values from api are 1e36
126#          masked=np.ma.masked_where(field > 1e35,field)
127#          masked.set_fill_value([np.NaN])
128#          mask = np.ma.getmask(masked)
129#          if (True in np.array(mask)):out=masked
130#          else:out=field
131       else:
132#       # missing values from MRAMS files are 0.100E+32
133          masked=np.ma.masked_where(field > 1e30,field)
134          masked.set_fill_value([np.NaN])
135          mask = np.ma.getmask(masked)
136          if (True in np.array(mask)):out=masked
137          else:out=field
138#       else:out=field
139    return out
140
141## Author: AC
142# Compute the norm of the winds or return an hodograph
143# The corresponding variable to call is UV or uvmet (to use api)
144def windamplitude (nc,mode):
145    import numpy as np
146    varinfile = nc.variables.keys()
147    if "U" in varinfile: zu=getfield(nc,'U')
148    elif "Um" in varinfile: zu=getfield(nc,'Um')
149    else: errormess("you need slopex or U or Um in your file.")
150    if "V" in varinfile: zv=getfield(nc,'V')
151    elif "Vm" in varinfile: zv=getfield(nc,'Vm')
152    else: errormess("you need V or Vm in your file.")
153    znt,znz,zny,znx = np.array(zu).shape
154    if hasattr(nc,'WEST-EAST_PATCH_END_UNSTAG'):znx=getattr(nc, 'WEST-EAST_PATCH_END_UNSTAG')
155    zuint = np.zeros([znt,znz,zny,znx])
156    zvint = np.zeros([znt,znz,zny,znx])
157    if "U" in varinfile:
158       if hasattr(nc,'SOUTH-NORTH_PATCH_END_STAG'): zny_stag=getattr(nc, 'SOUTH-NORTH_PATCH_END_STAG')
159       if hasattr(nc,'WEST-EAST_PATCH_END_STAG'): znx_stag=getattr(nc, 'WEST-EAST_PATCH_END_STAG')
160       if zny_stag == zny: zvint=zv
161       else:
162          for yy in np.arange(zny):      zvint[:,:,yy,:] = (zv[:,:,yy,:] + zv[:,:,yy+1,:])/2.
163       if znx_stag == znx: zuint=zu
164       else:
165          for xx in np.arange(znx):      zuint[:,:,:,xx] = (zu[:,:,:,xx] + zu[:,:,:,xx+1])/2.
166    else:
167       zuint=zu
168       zvint=zv
169    if mode=='amplitude': return np.sqrt(zuint**2 + zvint**2)
170    if mode=='hodograph': return zuint,zvint
171    if mode=='hodograph_2': return None, 360.*np.arctan(zvint/zuint)/(2.*np.pi)
172
173## Author: AC
174# Compute the enrichment factor of non condensible gases
175# The corresponding variable to call is enfact
176# enrichment factor is computed as in Yuan Lian et al. 2012
177# i.e. you need to have VL2 site at LS 135 in your data
178# this only requires co2col so that you can concat.nc at low cost
179def enrichment_factor(nc,lon,lat,time):
180    import numpy as np
181    from myplot import reducefield
182    varinfile = nc.variables.keys()
183    if "co2col" in varinfile: co2col=getfield(nc,'co2col')
184    else: print "error, you need co2col var in your file"
185    if "ps" in varinfile: ps=getfield(nc,'ps')
186    else: print "error, you need ps var in your file"
187    dimension = len(nc.variables['co2col'].dimensions)
188    if dimension == 2: 
189      zny,znx = np.array(co2col).shape
190      znt=1
191    elif dimension == 3: znt,zny,znx = np.array(co2col).shape
192    mmrarcol = np.zeros([znt,zny,znx])
193    enfact = np.zeros([znt,zny,znx])
194    grav=3.72
195    mmrarcol[:,:,:] = 1. - grav*co2col[:,:,:]/ps[:,:,:]
196# Computation with reference argon mmr at VL2 Ls 135 (as in Yuan Lian et al 2012)
197    lonvl2=np.zeros([1,2])
198    latvl2=np.zeros([1,2])
199    timevl2=np.zeros([1,2])
200    lonvl2[0,0]=-180
201    lonvl2[0,1]=180
202    latvl2[:,:]=48.16
203    timevl2[:,:]=135.
204    indexlon  = getsindex(lonvl2,0,lon)
205    indexlat  = getsindex(latvl2,0,lat)
206    indextime = getsindex(timevl2,0,time)
207    mmrvl2, error = reducefield( mmrarcol, d4=indextime, d1=indexlon, d2=indexlat)
208    print "VL2 Ls 135 mmr arcol:", mmrvl2
209    enfact[:,:,:] = mmrarcol[:,:,:]/mmrvl2
210    return enfact
211
212## Author: AC
213# Compute the norm of the slope angles
214# The corresponding variable to call is SLOPEXY
215def slopeamplitude (nc):
216    import numpy as np
217    varinfile = nc.variables.keys()
218    if "slopex" in varinfile: zu=getfield(nc,'slopex')
219    elif "SLOPEX" in varinfile: zu=getfield(nc,'SLOPEX')
220    else: errormess("you need slopex or SLOPEX in your file.") 
221    if "slopey" in varinfile: zv=getfield(nc,'slopey')
222    elif "SLOPEY" in varinfile: zv=getfield(nc,'SLOPEY')
223    else: errormess("you need slopey or SLOPEY in your file.")
224    znt,zny,znx = np.array(zu).shape
225    zuint = np.zeros([znt,zny,znx])
226    zvint = np.zeros([znt,zny,znx])
227    zuint=zu
228    zvint=zv
229    return np.sqrt(zuint**2 + zvint**2)
230
231## Author: AC
232# Compute the temperature difference between surface and first level.
233# API is automatically called to get TSURF and TK.
234# The corresponding variable to call is DELTAT
235def deltat0t1 (nc):
236    import numpy as np
237    varinfile = nc.variables.keys()
238    if "tsurf" in varinfile: zu=getfield(nc,'tsurf')
239    elif "TSURF" in varinfile: zu=getfield(nc,'TSURF')
240    else: errormess("You need tsurf or TSURF in your file")
241    if "tk" in varinfile: zv=getfield(nc,'tk')
242    elif "TK" in varinfile: zv=getfield(nc,'TK')
243    else: errormess("You need tk or TK in your file. (might need to use API. try to add -i 4 -l XXX)")
244    znt,zny,znx = np.array(zu).shape
245    zuint = np.zeros([znt,zny,znx])
246    zuint=zu - zv[:,0,:,:]
247    return zuint
248
249## Author: AS + TN + AC
250def reducefield (input,d4=None,d3=None,d2=None,d1=None,yint=False,alt=None,anomaly=False,redope=None,mesharea=None,unidim=999):
251    ### we do it the reverse way to be compliant with netcdf "t z y x" or "t y x" or "y x"
252    ### it would be actually better to name d4 d3 d2 d1 as t z y x
253    ### ... note, anomaly is only computed over d1 and d2 for the moment
254    import numpy as np
255    from mymath import max,mean,min,sum,getmask
256    csmooth = 12 ## a fair amount of grid points (too high results in high computation time)
257    if redope is not None:
258       if   redope == "mint":     input = min(input,axis=0) ; d1 = None
259       elif redope == "maxt":     input = max(input,axis=0) ; d1 = None
260       elif redope == "edge_y1":  input = input[:,:,0,:]    ; d2 = None
261       elif redope == "edge_y2":  input = input[:,:,-1,:]   ; d2 = None
262       elif redope == "edge_x1":  input = input[:,:,:,0]    ; d1 = None
263       elif redope == "edge_x2":  input = input[:,:,:,-1]   ; d1 = None
264       else:                      errormess("not supported. but try lines in reducefield beforehand.")
265       #elif redope == "minz":     input = min(input,axis=1) ; d2 = None
266       #elif redope == "maxz":     input = max(input,axis=1) ; d2 = None
267       #elif redope == "miny":     input = min(input,axis=2) ; d3 = None
268       #elif redope == "maxy":     input = max(input,axis=2) ; d3 = None
269       #elif redope == "minx":     input = min(input,axis=3) ; d4 = None
270       #elif redope == "maxx":     input = max(input,axis=3) ; d4 = None
271    dimension = np.array(input).ndim
272    shape = np.array(np.array(input).shape)
273    #print 'd1,d2,d3,d4: ',d1,d2,d3,d4
274    if anomaly: print 'ANOMALY ANOMALY'
275    output = input
276    error = False
277    #### this is needed to cope the case where d4,d3,d2,d1 are single integers and not arrays
278    if d4 is not None and not isinstance(d4, np.ndarray): d4=[d4]
279    if d3 is not None and not isinstance(d3, np.ndarray): d3=[d3]
280    if d2 is not None and not isinstance(d2, np.ndarray): d2=[d2]
281    if d1 is not None and not isinstance(d1, np.ndarray): d1=[d1]
282    ### now the main part
283    if dimension == 2:
284        #### this is needed for 1d-type files (where dim=2 but axes are time-vert and not lat-lon)
285        if unidim==1: d2=d4 ; d1=d3 ; d4 = None ; d3 = None
286        if mesharea is None: mesharea=np.ones(shape)
287        if   max(d2) >= shape[0]: error = True
288        elif max(d1) >= shape[1]: error = True
289        elif d1 is not None and d2 is not None:
290          try:
291            totalarea = np.ma.masked_where(getmask(output),mesharea)
292            totalarea = mean(totalarea[d2,:],axis=0);totalarea = mean(totalarea[d1])
293          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
294          output = output*mesharea; output = mean(output[d2,:],axis=0); output = mean(output[d1])/totalarea
295        elif d1 is not None:         output = mean(input[:,d1],axis=1)
296        elif d2 is not None:
297          try:
298            totalarea = np.ma.masked_where(getmask(output),mesharea)
299            totalarea = mean(totalarea[d2,:],axis=0)
300          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
301          output = output*mesharea; output = mean(output[d2,:],axis=0)/totalarea
302    elif dimension == 3:
303        if mesharea is None: mesharea=np.ones(shape[[1,2]])
304        if   max(d4) >= shape[0]: error = True
305        elif max(d2) >= shape[1]: error = True
306        elif max(d1) >= shape[2]: error = True
307        elif d4 is not None and d2 is not None and d1 is not None:
308          output = mean(input[d4,:,:],axis=0)
309          try:
310            totalarea = np.ma.masked_where(getmask(output),mesharea)
311            totalarea = mean(totalarea[d2,:],axis=0);totalarea = mean(totalarea[d1])
312          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
313          output = output*mesharea; output = mean(output[d2,:],axis=0); output = mean(output[d1])/totalarea
314        elif d4 is not None and d2 is not None:
315          output = mean(input[d4,:,:],axis=0)
316          try:
317            totalarea = np.ma.masked_where(getmask(output),mesharea)
318            totalarea = mean(totalarea[d2,:],axis=0)
319          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
320          output = output*mesharea; output = mean(output[d2,:],axis=0)/totalarea
321        elif d4 is not None and d1 is not None:    output = mean(input[d4,:,:],axis=0); output=mean(output[:,d1],axis=1)
322        elif d2 is not None and d1 is not None:
323          try:
324            totalarea = np.tile(mesharea,(output.shape[0],1,1))
325            totalarea = np.ma.masked_where(getmask(output),totalarea)
326            totalarea = mean(totalarea[:,d2,:],axis=1);totalarea = mean(totalarea[:,d1],axis=1)
327          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
328          output = output*mesharea; output = mean(output[:,d2,:],axis=1); output = mean(output[:,d1],axis=1)/totalarea
329        elif d1 is not None:   output = mean(input[:,:,d1],axis=2)
330        elif d2 is not None:   
331          try:
332            totalarea = np.tile(mesharea,(output.shape[0],1,1))
333            totalarea = np.ma.masked_where(getmask(output),totalarea)
334            totalarea = mean(totalarea[:,d2,:],axis=1)
335          except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
336          output = output*mesharea; output = mean(output[:,d2,:],axis=1)/totalarea
337        elif d4 is not None:   output = mean(input[d4,:,:],axis=0)
338    elif dimension == 4:
339        if mesharea is None: mesharea=np.ones(shape[[2,3]]) # mesharea=np.random.random_sample(shape[[2,3]])*5. + 2. # pour tester
340        if   max(d4) >= shape[0]: error = True
341        elif max(d3) >= shape[1]: error = True
342        elif max(d2) >= shape[2]: error = True
343        elif max(d1) >= shape[3]: error = True
344        elif d4 is not None and d3 is not None and d2 is not None and d1 is not None:
345             output = mean(input[d4,:,:,:],axis=0)
346             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
347             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.)
348             try:
349               totalarea = np.ma.masked_where(np.isnan(output),mesharea)
350               totalarea = mean(totalarea[d2,:],axis=0); totalarea = mean(totalarea[d1])
351             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
352             output = output*mesharea; output = mean(output[d2,:],axis=0); output = mean(output[d1])/totalarea
353        elif d4 is not None and d3 is not None and d2 is not None: 
354             output = mean(input[d4,:,:,:],axis=0)
355             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
356             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.)
357             try:
358               totalarea = np.ma.masked_where(np.isnan(output),mesharea)
359               totalarea = mean(totalarea[d2,:],axis=0)
360             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
361             output = output*mesharea; output = mean(output[d2,:],axis=0)/totalarea
362        elif d4 is not None and d3 is not None and d1 is not None: 
363             output = mean(input[d4,:,:,:],axis=0)
364             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
365             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
366             output = mean(output[:,d1],axis=1)
367        elif d4 is not None and d2 is not None and d1 is not None: 
368             output = mean(input[d4,:,:,:],axis=0)
369             if anomaly:
370                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
371             try:
372               totalarea = np.tile(mesharea,(output.shape[0],1,1))
373               totalarea = np.ma.masked_where(getmask(output),mesharea)
374               totalarea = mean(totalarea[:,d2,:],axis=1); totalarea = mean(totalarea[:,d1],axis=1)
375             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
376             output = output*mesharea; output = mean(output[:,d2,:],axis=1); output = mean(output[:,d1],axis=1)/totalarea
377             #noperturb = smooth1d(output,window_len=7)
378             #lenlen = len(output) ; output = output[1:lenlen-7] ; yeye = noperturb[4:lenlen-4]
379             #plot(output) ; plot(yeye) ; show() ; plot(output-yeye) ; show()
380        elif d3 is not None and d2 is not None and d1 is not None:
381             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
382             if anomaly:
383                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
384             try:
385               totalarea = np.tile(mesharea,(output.shape[0],1,1))
386               totalarea = np.ma.masked_where(getmask(output),totalarea)
387               totalarea = mean(totalarea[:,d2,:],axis=1); totalarea = mean(totalarea[:,d1],axis=1)
388             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
389             output = output*mesharea; output = mean(output[:,d2,:],axis=1); output = mean(output[:,d1],axis=1)/totalarea
390        elif d4 is not None and d3 is not None: 
391             output = mean(input[d4,:,:,:],axis=0) 
392             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
393             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
394        elif d4 is not None and d2 is not None: 
395             output = mean(input[d4,:,:,:],axis=0)
396             try:
397               totalarea = np.tile(mesharea,(output.shape[0],1,1))
398               totalarea = np.ma.masked_where(getmask(output),mesharea)
399               totalarea = mean(totalarea[:,d2,:],axis=1)
400             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
401             output = output*mesharea; output = mean(output[:,d2,:],axis=1)/totalarea
402        elif d4 is not None and d1 is not None: 
403             output = mean(input[d4,:,:,:],axis=0) 
404             output = mean(output[:,:,d1],axis=2)
405        elif d3 is not None and d2 is not None:
406             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
407             try:
408               totalarea = np.tile(mesharea,(output.shape[0],1,1))
409               totalarea = np.ma.masked_where(getmask(output),mesharea)
410               totalarea = mean(totalarea[:,d2,:],axis=1)
411             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
412             output = output*mesharea; output = mean(output[:,d2,:],axis=1)/totalarea
413        elif d3 is not None and d1 is not None: 
414             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
415             output = mean(output[:,:,d1],axis=2)
416        elif d2 is not None and d1 is not None:
417             try:
418               totalarea = np.tile(mesharea,(output.shape[0],output.shape[1],1,1))
419               totalarea = np.ma.masked_where(getmask(output),totalarea)
420               totalarea = mean(totalarea[:,:,d2,:],axis=2); totalarea = mean(totalarea[:,:,d1],axis=1)
421             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
422             output = output*mesharea; output = mean(output[:,:,d2,:],axis=2); output = mean(output[:,:,d1],axis=2)/totalarea
423        elif d1 is not None:        output = mean(input[:,:,:,d1],axis=3)
424        elif d2 is not None:
425             try:
426               totalarea = np.tile(mesharea,(output.shape[0],output.shape[1],1,output.shape[3]))
427               totalarea = np.ma.masked_where(getmask(output),totalarea)
428               totalarea = mean(totalarea[:,:,d2,:],axis=2)
429             except: print "(problem with areas. I skip this)" ; mesharea = 1. ; totalarea = 1.
430             output = output*mesharea; output = mean(output[:,:,d2,:],axis=2)/totalarea
431        elif d3 is not None:        output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
432        elif d4 is not None:        output = mean(input[d4,:,:,:],axis=0)
433    dimension2 = np.array(output).ndim
434    shape2 = np.array(output).shape
435    print 'REDUCEFIELD dim,shape: ',dimension,shape,' >>> ',dimension2,shape2
436    return output, error
437
438## Author: AC + AS
439def reduce_zaxis (input,ax=None,yint=False,vert=None,indice=None):
440    from mymath import max,mean
441    from scipy import integrate
442    import numpy as np
443    if yint and vert is not None and indice is not None:
444       if type(input).__name__=='MaskedArray':
445          input.set_fill_value([np.NaN])
446          output = integrate.trapz(input.filled(),x=vert[indice],axis=ax)
447       else:
448          output = integrate.trapz(input,x=vert[indice],axis=ax)
449    else:
450       output = mean(input,axis=ax)
451    return output
452
453## Author: AS + TN
454def definesubplot ( numplot, fig ):
455    from matplotlib.pyplot import rcParams
456    rcParams['font.size'] = 12. ## default (important for multiple calls)
457    if numplot <= 0:
458        subv = 99999
459        subh = 99999
460    elif numplot == 1: 
461        subv = 1
462        subh = 1
463    elif numplot == 2:
464        subv = 1 #2
465        subh = 2 #1
466        fig.subplots_adjust(wspace = 0.35)
467        rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
468    elif numplot == 3:
469        subv = 3
470        subh = 1
471        fig.subplots_adjust(hspace = 0.5)
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    #print meanlat, meanlon
886    if   char == "cyl":     m = Basemap(rsphere=radius,projection='cyl',\
887                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
888    elif char == "moll":    m = Basemap(rsphere=radius,projection='moll',lon_0=meanlon)
889    elif char == "ortho":   m = Basemap(rsphere=radius,projection='ortho',lon_0=ortholon,lat_0=ortholat)
890    elif char == "lcc":     m = Basemap(rsphere=radius,projection='lcc',lat_1=meanlat,lat_0=meanlat,lon_0=meanlon,\
891                              width=zewidth,height=zeheight)
892                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
893    elif char == "npstere": m = Basemap(rsphere=radius,projection='npstere', boundinglat=blat, lon_0=0.)
894    elif char == "spstere": m = Basemap(rsphere=radius,projection='spstere', boundinglat=blat, lon_0=180.)
895    elif char == "nplaea":  m = Basemap(rsphere=radius,projection='nplaea', boundinglat=wlat[0], lon_0=meanlon)
896    elif char == "laea":    m = Basemap(rsphere=radius,projection='laea',lon_0=meanlon,lat_0=meanlat,lat_ts=meanlat,\
897                              width=zewidth,height=zeheight)
898                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
899    elif char == "nsper":   m = Basemap(rsphere=radius,projection='nsper',lon_0=meanlon,lat_0=meanlat,satellite_height=h*1000.)
900    elif char == "merc":    m = Basemap(rsphere=radius,projection='merc',lat_ts=0.,\
901                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
902    fontsizemer = int(mpl.rcParams['font.size']*3./4.)
903    if char in ["cyl","lcc","merc","nsper","laea"]:   step = findstep(wlon)
904    else:                                             step = 10.
905    steplon = step*2.
906    zecolor ='grey'
907    zelinewidth = 1
908    zelatmax = 80
909    if meanlat > 75.: zelatmax = 90. ; step = step/2.
910    # to show gcm grid:
911    #zecolor = 'r'
912    #zelinewidth = 1
913    #step = 180./48.
914    #steplon = 360./64.
915    #zelatmax = 90. - step/3
916    if char not in ["moll"]:
917        if wlon[1]-wlon[0] < 2.:  ## LOCAL MODE
918            m.drawmeridians(np.r_[-1.:1.:0.05], labels=[0,0,0,1], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, fmt='%5.2f')
919            m.drawparallels(np.r_[-1.:1.:0.05], labels=[1,0,0,0], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, fmt='%5.2f')
920        else:  ## GLOBAL OR REGIONAL MODE
921            m.drawmeridians(np.r_[-180.:180.:steplon], labels=[0,0,0,1], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
922            m.drawparallels(np.r_[-90.:90.:step], labels=[1,0,0,0], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
923    if back: 
924      if back not in ["coast","sea"]:   m.warpimage(marsmap(back),scale=0.75)
925      elif back == "coast":             m.drawcoastlines()
926      elif back == "sea":               m.drawlsmask(land_color='white',ocean_color='aqua')
927            #if not back:
928            #    if not var:                                        back = "mola"    ## if no var:         draw mola
929            #    elif typefile in ['mesoapi','meso','geo'] \
930            #       and proj not in ['merc','lcc','nsper','laea']:  back = "molabw"  ## if var but meso:   draw molabw
931            #    else:                                              pass             ## else:              draw None
932    return m
933
934## Author: AS
935#### test temporaire
936def putpoints (map,plot):
937    #### from http://www.scipy.org/Cookbook/Matplotlib/Maps
938    # lat/lon coordinates of five cities.
939    lats = [18.4]
940    lons = [-134.0]
941    points=['Olympus Mons']
942    # compute the native map projection coordinates for cities.
943    x,y = map(lons,lats)
944    # plot filled circles at the locations of the cities.
945    map.plot(x,y,'bo')
946    # plot the names of those five cities.
947    wherept = 0 #1000 #50000
948    for name,xpt,ypt in zip(points,x,y):
949       plot.text(xpt+wherept,ypt+wherept,name)
950    ## le nom ne s'affiche pas...
951    return
952
953## Author: AS
954def calculate_bounds(field,vmin=None,vmax=None):
955    import numpy as np
956    from mymath import max,min,mean
957    ind = np.where(field < 9e+35)
958    fieldcalc = field[ ind ] # la syntaxe compacte ne marche si field est un tuple
959    ###
960    dev = np.std(fieldcalc)*3.0
961    ###
962    if vmin is None:  zevmin = mean(fieldcalc) - dev
963    else:             zevmin = vmin
964    ###
965    if vmax is None:  zevmax = mean(fieldcalc) + dev
966    else:             zevmax = vmax
967    if vmin == vmax:
968                      zevmin = mean(fieldcalc) - dev  ### for continuity
969                      zevmax = mean(fieldcalc) + dev  ### for continuity           
970    ###
971    if zevmin < 0. and min(fieldcalc) > 0.: zevmin = 0.
972    print "BOUNDS field ", min(fieldcalc), max(fieldcalc), " //// adopted", zevmin, zevmax
973    return zevmin, zevmax
974
975## Author: AS
976def bounds(what_I_plot,zevmin,zevmax):
977    from mymath import max,min,mean
978    ### might be convenient to add the missing value in arguments
979    #what_I_plot[ what_I_plot < zevmin ] = zevmin#*(1. + 1.e-7)
980    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1. - 1.e-7) ] = zevmin*(1. - 1.e-7)
981    else:          what_I_plot[ what_I_plot < zevmin*(1. + 1.e-7) ] = zevmin*(1. + 1.e-7)
982    #print "NEW MIN ", min(what_I_plot)
983    what_I_plot[ what_I_plot > 9e+35  ] = -9e+35
984    what_I_plot[ what_I_plot > zevmax ] = zevmax*(1. - 1.e-7)
985    #print "NEW MAX ", max(what_I_plot)
986    return what_I_plot
987
988## Author: AS
989def nolow(what_I_plot):
990    from mymath import max,min
991    lim = 0.15*0.5*(abs(max(what_I_plot))+abs(min(what_I_plot)))
992    print "NO PLOT BELOW VALUE ", lim
993    what_I_plot [ abs(what_I_plot) < lim ] = 1.e40 
994    return what_I_plot
995
996
997## Author : AC
998def hole_bounds(what_I_plot,zevmin,zevmax):
999    import numpy as np
1000    zi=0
1001    for i in what_I_plot:
1002        zj=0
1003        for j in i:
1004            if ((j < zevmin) or (j > zevmax)):what_I_plot[zi,zj]=np.NaN
1005            zj=zj+1
1006        zi=zi+1
1007
1008    return what_I_plot
1009
1010## Author: AS
1011def zoomset (wlon,wlat,zoom):
1012    dlon = abs(wlon[1]-wlon[0])/2.
1013    dlat = abs(wlat[1]-wlat[0])/2.
1014    [wlon,wlat] = [ [wlon[0]+zoom*dlon/100.,wlon[1]-zoom*dlon/100.],\
1015                    [wlat[0]+zoom*dlat/100.,wlat[1]-zoom*dlat/100.] ]
1016    print "ZOOM %",zoom,wlon,wlat
1017    return wlon,wlat
1018
1019## Author: AS
1020def fmtvar (whichvar="def"):
1021    fmtvar    =     { \
1022             "MIXED":        "%.0f",\
1023             "UPDRAFT":      "%.0f",\
1024             "DOWNDRAFT":    "%.0f",\
1025             "TK":           "%.0f",\
1026             "T":            "%.0f",\
1027             #"ZMAX_TH":      "%.0f",\
1028             #"WSTAR":        "%.0f",\
1029             # Variables from TES ncdf format
1030             "T_NADIR_DAY":  "%.0f",\
1031             "T_NADIR_NIT":  "%.0f",\
1032             # Variables from tes.py ncdf format
1033             "TEMP_DAY":     "%.0f",\
1034             "TEMP_NIGHT":   "%.0f",\
1035             # Variables from MCS and mcs.py ncdf format
1036             "DTEMP":        "%.0f",\
1037             "NTEMP":        "%.0f",\
1038             "DNUMBINTEMP":  "%.0f",\
1039             "NNUMBINTEMP":  "%.0f",\
1040             # other stuff
1041             "TPOT":         "%.0f",\
1042             "TSURF":        "%.0f",\
1043             "U_OUT1":       "%.0f",\
1044             "T_OUT1":       "%.0f",\
1045             "def":          "%.1e",\
1046             "PTOT":         "%.0f",\
1047             "PSFC":         "%.1f",\
1048             "HGT":          "%.1e",\
1049             "USTM":         "%.2f",\
1050             "HFX":          "%.0f",\
1051             "ICETOT":       "%.1e",\
1052             "TAU_ICE":      "%.2f",\
1053             "TAUICE":       "%.2f",\
1054             "VMR_ICE":      "%.1e",\
1055             "MTOT":         "%.1f",\
1056             "ANOMALY":      "%.1f",\
1057             "W":            "%.2f",\
1058             "WMAX_TH":      "%.1f",\
1059             "WSTAR":        "%.1f",\
1060             "QSURFICE":     "%.0f",\
1061             "UM":           "%.0f",\
1062             "WIND":         "%.0f",\
1063             "UVMET":         "%.0f",\
1064             "UV":         "%.0f",\
1065             "ALBBARE":      "%.2f",\
1066             "TAU":          "%.1f",\
1067             "CO2":          "%.2f",\
1068             "ENFACT":       "%.1f",\
1069             "QDUST":        "%.6f",\
1070             #### T.N.
1071             "TEMP":         "%.0f",\
1072             "VMR_H2OICE":   "%.0f",\
1073             "VMR_H2OVAP":   "%.0f",\
1074             "TAUTES":       "%.2f",\
1075             "TAUTESAP":     "%.2f",\
1076
1077                    }
1078    if "TSURF" in whichvar: whichvar = "TSURF"
1079    if whichvar not in fmtvar:
1080        whichvar = "def"
1081    return fmtvar[whichvar]
1082
1083## Author: AS
1084####################################################################################################################
1085### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
1086def defcolorb (whichone="def"):
1087    whichcolorb =    { \
1088             "def":          "spectral",\
1089             "HGT":          "spectral",\
1090             "HGT_M":        "spectral",\
1091             "TK":           "gist_heat",\
1092             "TPOT":         "Paired",\
1093             "TSURF":        "RdBu_r",\
1094             "QH2O":         "PuBu",\
1095             "USTM":         "YlOrRd",\
1096             "WIND":         "YlOrRd",\
1097             #"T_nadir_nit":  "RdBu_r",\
1098             #"T_nadir_day":  "RdBu_r",\
1099             "HFX":          "RdYlBu",\
1100             "ICETOT":       "YlGnBu_r",\
1101             #"MTOT":         "PuBu",\
1102             "CCNQ":         "YlOrBr",\
1103             "CCNN":         "YlOrBr",\
1104             "TEMP":         "Jet",\
1105             "TAU_ICE":      "Blues",\
1106             "TAUICE":       "Blues",\
1107             "VMR_ICE":      "Blues",\
1108             "W":            "jet",\
1109             "WMAX_TH":      "spectral",\
1110             "ANOMALY":      "RdBu_r",\
1111             "QSURFICE":     "hot_r",\
1112             "ALBBARE":      "spectral",\
1113             "TAU":          "YlOrBr_r",\
1114             "CO2":          "YlOrBr_r",\
1115             "MIXED":        "GnBu",\
1116             #### T.N.
1117             "MTOT":         "spectral",\
1118             "H2O_ICE_S":    "RdBu",\
1119             "VMR_H2OICE":   "PuBu",\
1120             "VMR_H2OVAP":   "PuBu",\
1121             "WATERCAPTAG":  "Blues",\
1122                     }
1123#W --> spectral ou jet
1124#spectral BrBG RdBu_r
1125    #print "predefined colorbars"
1126    if "TSURF" in whichone: whichone = "TSURF"
1127    if whichone not in whichcolorb:
1128        whichone = "def"
1129    return whichcolorb[whichone]
1130
1131## Author: AS
1132def definecolorvec (whichone="def"):
1133        whichcolor =    { \
1134                "def":          "black",\
1135                "vis":          "yellow",\
1136                "vishires":     "green",\
1137                "molabw":       "yellow",\
1138                "mola":         "black",\
1139                "gist_heat":    "white",\
1140                "hot":          "tk",\
1141                "gist_rainbow": "black",\
1142                "spectral":     "black",\
1143                "gray":         "red",\
1144                "PuBu":         "black",\
1145                "titan":        "red",\
1146                        }
1147        if whichone not in whichcolor:
1148                whichone = "def"
1149        return whichcolor[whichone]
1150
1151## Author: AS
1152def marsmap (whichone="vishires"):
1153        from os import uname
1154        mymachine = uname()[1]
1155        ### not sure about speed-up with this method... looks the same
1156        if "lmd.jussieu.fr" in mymachine:   domain = "/u/aslmd/WWW/maps/"
1157        elif "aymeric-laptop" in mymachine: domain = "/home/aymeric/Dropbox/Public/"
1158        else:                               domain = "http://www.lmd.jussieu.fr/~aslmd/maps/"
1159        whichlink =     { \
1160                #"vis":         "http://maps.jpl.nasa.gov/pix/mar0kuu2.jpg",\
1161                #"vishires":    "http://www.lmd.jussieu.fr/~aslmd/maps/MarsMap_2500x1250.jpg",\
1162                #"geolocal":    "http://dl.dropbox.com/u/11078310/geolocal.jpg",\
1163                #"mola":        "http://www.lns.cornell.edu/~seb/celestia/mars-mola-2k.jpg",\
1164                #"molabw":      "http://dl.dropbox.com/u/11078310/MarsElevation_2500x1250.jpg",\
1165                "thermalday":   domain+"thermalday.jpg",\
1166                "thermalnight": domain+"thermalnight.jpg",\
1167                "tesalbedo":    domain+"tesalbedo.jpg",\
1168                "vis":         domain+"mar0kuu2.jpg",\
1169                "vishires":    domain+"MarsMap_2500x1250.jpg",\
1170                "geolocal":    domain+"geolocal.jpg",\
1171                "mola":        domain+"mars-mola-2k.jpg",\
1172                "molabw":      domain+"MarsElevation_2500x1250.jpg",\
1173                "clouds":      "http://www.johnstonsarchive.net/spaceart/marswcloudmap.jpg",\
1174                "jupiter":     "http://www.mmedia.is/~bjj/data/jupiter_css/jupiter_css.jpg",\
1175                "jupiter_voy": "http://www.mmedia.is/~bjj/data/jupiter/jupiter_vgr2.jpg",\
1176                #"bw":          domain+"EarthElevation_2500x1250.jpg",\
1177                "bw":          "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg",\
1178                "contrast":    "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg",\
1179                "nice":        "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg",\
1180                "blue":        "http://eoimages.gsfc.nasa.gov/ve/2430/land_ocean_ice_2048.jpg",\
1181                "blueclouds":  "http://eoimages.gsfc.nasa.gov/ve/2431/land_ocean_ice_cloud_2048.jpg",\
1182                "justclouds":  "http://eoimages.gsfc.nasa.gov/ve/2432/cloud_combined_2048.jpg",\
1183                "pluto":       "http://www.boulder.swri.edu/~buie/pluto/pluto_all.png",\
1184                "triton":      "http://laps.noaa.gov/albers/sos/neptune/triton/triton_rgb_cyl_www.jpg",\
1185                "titan":       "http://laps.noaa.gov/albers/sos/saturn/titan/titan_rgb_cyl_www.jpg",\
1186                #"titan":       "http://laps.noaa.gov/albers/sos/celestia/titan_50.jpg",\
1187                "titanuni":    "http://maps.jpl.nasa.gov/pix/sat6fss1.jpg",\
1188                "venus":       "http://laps.noaa.gov/albers/sos/venus/venus4/venus4_rgb_cyl_www.jpg",\
1189                "cosmic":      "http://laps.noaa.gov/albers/sos/universe/wmap/wmap_rgb_cyl_www.jpg",\
1190                        }
1191        ### see http://www.mmedia.is/~bjj/planetary_maps.html
1192        if whichone not in whichlink: 
1193                print "marsmap: choice not defined... you'll get the default one... "
1194                whichone = "vishires" 
1195        return whichlink[whichone]
1196
1197#def earthmap (whichone):
1198#       if   whichone == "contrast":    whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg"
1199#       elif whichone == "bw":          whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg"
1200#       elif whichone == "nice":        whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg"
1201#       return whichlink
1202
1203## Author: AS
1204def latinterv (area="Whole"):
1205    list =    { \
1206        "Europe":                [[ 20., 80.],[- 50.,  50.]],\
1207        "Central_America":       [[-10., 40.],[ 230., 300.]],\
1208        "Africa":                [[-20., 50.],[- 50.,  50.]],\
1209        "Whole":                 [[-90., 90.],[-180., 180.]],\
1210        "Southern_Hemisphere":   [[-90., 60.],[-180., 180.]],\
1211        "Northern_Hemisphere":   [[-60., 90.],[-180., 180.]],\
1212        "Tharsis":               [[-30., 60.],[-170.,- 10.]],\
1213        "Whole_No_High":         [[-60., 60.],[-180., 180.]],\
1214        "Chryse":                [[-60., 60.],[- 60.,  60.]],\
1215        "North_Pole":            [[ 50., 90.],[-180., 180.]],\
1216        "Close_North_Pole":      [[ 75., 90.],[-180., 180.]],\
1217        "Far_South_Pole":        [[-90.,-40.],[-180., 180.]],\
1218        "South_Pole":            [[-90.,-50.],[-180., 180.]],\
1219        "Close_South_Pole":      [[-90.,-75.],[-180., 180.]],\
1220        "Sirenum_Crater_large":  [[-46.,-34.],[-166.,-151.]],\
1221        "Sirenum_Crater_small":  [[-36.,-26.],[-168.,-156.]],\
1222        "Rupes":                 [[ 72., 90.],[-120.,- 20.]],\
1223        "Xanadu":                [[-40., 20.],[  40., 120.]],\
1224        "Hyperboreae":           [[ 80., 87.],[- 70.,- 10.]],\
1225              }
1226    if area not in list:   area = "Whole"
1227    [olat,olon] = list[area]
1228    return olon,olat
1229
1230## Author: TN
1231def separatenames (name):
1232  from numpy import concatenate
1233  # look for comas in the input name to separate different names (files, variables,etc ..)
1234  if name is None:
1235     names = None
1236  else:
1237    names = []
1238    stop = 0
1239    currentname = name
1240    while stop == 0:
1241      indexvir = currentname.find(',')
1242      if indexvir == -1:
1243        stop = 1
1244        name1 = currentname
1245      else:
1246        name1 = currentname[0:indexvir]
1247      names = concatenate((names,[name1]))
1248      currentname = currentname[indexvir+1:len(currentname)]
1249  return names
1250
1251
1252## Author: TN
1253def readslices(saxis):
1254  from numpy import empty
1255  if saxis == None:
1256     zesaxis = None
1257  else:
1258     zesaxis = empty((len(saxis),2))
1259     for i in range(len(saxis)):
1260        a = separatenames(saxis[i])
1261        if len(a) == 1:
1262           zesaxis[i,:] = float(a[0])
1263        else:
1264           zesaxis[i,0] = float(a[0])
1265           zesaxis[i,1] = float(a[1])
1266           
1267  return zesaxis
1268
1269## Author: TN
1270def readdata(data,datatype,coord1,coord2):
1271## Read sparse data
1272  from numpy import empty
1273  if datatype == 'txt':
1274     if len(data[coord1].shape) == 1:
1275         return data[coord1][:]
1276     elif len(data[coord1].shape) == 2:
1277         return data[coord1][:,int(coord2)-1]
1278     else:
1279         errormess('error in readdata')
1280  elif datatype == 'sav':
1281     return data[coord1][coord2]
1282  else:
1283     errormess(datatype+' type is not supported!')
1284
1285
1286## Author: AS
1287def bidimfind(lon2d,lat2d,vlon,vlat,file=None):
1288   import numpy as np
1289   import matplotlib.pyplot as mpl
1290   if vlat is None:    array = (lon2d - vlon)**2
1291   elif vlon is None:  array = (lat2d - vlat)**2
1292   else:               array = (lon2d - vlon)**2 + (lat2d - vlat)**2
1293   idy,idx = np.unravel_index( np.argmin(array), lon2d.shape )
1294   if vlon is not None:
1295      if (np.abs(lon2d[idy,idx]-vlon)) > 5: errormess("longitude not found ",printvar=lon2d)
1296   if vlat is not None:
1297      if (np.abs(lat2d[idy,idx]-vlat)) > 5: errormess("latitude not found ",printvar=lat2d)
1298   if file is not None:
1299      print idx,idy,lon2d[idy,idx],vlon
1300      print idx,idy,lat2d[idy,idx],vlat
1301      var = file.variables["HGT"][:,:,:]
1302      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)
1303      mpl.show()
1304   return idy,idx
1305
1306## Author: TN
1307def getsindex(saxis,index,axis):
1308# input  : all the desired slices and the good index
1309# output : all indexes to be taken into account for reducing field
1310  import numpy as np
1311  if ( np.array(axis).ndim == 2):
1312      axis = axis[:,0]
1313  if saxis is None:
1314      zeindex = None
1315  else:
1316      aaa = int(np.argmin(abs(saxis[index,0] - axis)))
1317      bbb = int(np.argmin(abs(saxis[index,1] - axis)))
1318      [imin,imax] = np.sort(np.array([aaa,bbb]))
1319      zeindex = np.array(range(imax-imin+1))+imin
1320      # because -180 and 180 are the same point in longitude,
1321      # we get rid of one for averaging purposes.
1322      if axis[imin] == -180 and axis[imax] == 180:
1323         zeindex = zeindex[0:len(zeindex)-1]
1324         print "INFO: whole longitude averaging asked, so last point is not taken into account."
1325  return zeindex
1326     
1327## Author: TN
1328def define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,indextime,what_I_plot,dim0,vertmode,redope):
1329# Purpose of define_axis is to find x and y axis scales in a smart way
1330# x axis priority: 1/time 2/lon 3/lat 4/vertical
1331# To be improved !!!...
1332   from numpy import array,swapaxes
1333   x = None
1334   y = None
1335   count = 0
1336   what_I_plot = array(what_I_plot)
1337   shape = what_I_plot.shape
1338   if indextime is None and len(time) > 1:
1339      print "AXIS is time"
1340      x = time
1341      count = count+1
1342   if indexlon is None and len(lon) > 1 and redope not in ['edge_x1','edge_x2']:
1343      print "AXIS is lon"
1344      if count == 0: x = lon
1345      else: y = lon
1346      count = count+1
1347   if indexlat is None and len(lat) > 1 and redope not in ['edge_y1','edge_y2']:
1348      print "AXIS is lat"
1349      if count == 0: x = lat
1350      else: y = lat
1351      count = count+1
1352   if indexvert is None and ((dim0 == 4) or (y is None)):
1353      print "AXIS is vert"
1354      if vertmode == 0: # vertical axis is as is (GCM grid)
1355         if count == 0: x=range(len(vert))
1356         else: y=range(len(vert))
1357         count = count+1
1358      else: # vertical axis is in kms
1359         if count == 0: x = vert
1360         else: y = vert
1361         count = count+1
1362   x = array(x)
1363   y = array(y)
1364   print "CHECK SHAPE: what_I_plot, x, y", what_I_plot.shape, x.shape, y.shape
1365   if len(shape) == 1:
1366       if shape[0] != len(x):           print "WARNING: shape[0] != len(x). Correcting." ; what_I_plot = what_I_plot[0:len(x)]
1367       if len(y.shape) > 0:             y = ()
1368   elif len(shape) == 2:
1369       if shape[1] == len(y) and shape[0] == len(x) and shape[0] != shape[1]:
1370           print "INFO: swapaxes: ",what_I_plot.shape,shape ; what_I_plot = swapaxes(what_I_plot,0,1)
1371       else:
1372           if shape[0] != len(y):       print "WARNING: shape[0] != len(y). Correcting." ; what_I_plot = what_I_plot[0:len(y),:]
1373           elif shape[1] != len(x):     print "WARNING: shape[1] != len(x). Correcting." ; what_I_plot = what_I_plot[:,0:len(x)]
1374   elif len(shape) == 3:
1375       if vertmode < 0: print "not supported. must check array dimensions at some point. not difficult to implement though."
1376   return what_I_plot,x,y
1377
1378# Author: TN + AS + AC
1379def determineplot(slon, slat, svert, stime, redope):
1380    nlon = 1 # number of longitudinal slices -- 1 is None
1381    nlat = 1
1382    nvert = 1
1383    ntime = 1
1384    nslices = 1
1385    if slon is not None:
1386        length=len(slon[:,0])
1387        nslices = nslices*length
1388        nlon = len(slon)
1389    if slat is not None:
1390        length=len(slat[:,0])
1391        nslices = nslices*length
1392        nlat = len(slat)
1393    if svert is not None:
1394        length=len(svert[:,0])
1395        nslices = nslices*length
1396        nvert = len(svert)
1397    if stime is not None:
1398        length=len(stime[:,0])
1399        nslices = nslices*length
1400        ntime = len(stime)
1401    #else:
1402    #    nslices = 2 
1403    mapmode = 0
1404    if slon is None and slat is None and redope not in ['edge_x1','edge_x2','edge_y1','edge_y2']:
1405       mapmode = 1 # in this case we plot a map, with the given projection
1406    return nlon, nlat, nvert, ntime, mapmode, nslices
1407
1408## Author : AS
1409def maplatlon( lon,lat,field,\
1410               proj="cyl",colorb="jet",ndiv=10,zeback="molabw",trans=0.6,title="",\
1411               vecx=None,vecy=None,stride=2 ):
1412    ### an easy way to map a field over lat/lon grid
1413    import numpy as np
1414    import matplotlib.pyplot as mpl
1415    from matplotlib.cm import get_cmap
1416    ## get lon and lat in 2D version. get lat/lon intervals
1417    numdim = len(np.array(lon).shape)
1418    if numdim == 2:     [lon2d,lat2d] = [lon,lat]
1419    elif numdim == 1:   [lon2d,lat2d] = np.meshgrid(lon,lat)
1420    else:               errormess("lon and lat arrays must be 1D or 2D")
1421    [wlon,wlat] = latinterv()
1422    ## define projection and background. define x and y given the projection
1423    m = define_proj(proj,wlon,wlat,back=zeback,blat=None,blon=None)
1424    x, y = m(lon2d, lat2d)
1425    ## define field. bound field.
1426    what_I_plot = np.transpose(field)
1427    zevmin, zevmax = calculate_bounds(what_I_plot)  ## vmin=min(what_I_plot_frame), vmax=max(what_I_plot_frame))
1428    what_I_plot = bounds(what_I_plot,zevmin,zevmax)
1429    ## define contour field levels. define color palette
1430    ticks = ndiv + 1
1431    zelevels = np.linspace(zevmin,zevmax,ticks)
1432    palette = get_cmap(name=colorb)
1433    ## contour field
1434    m.contourf( x, y, what_I_plot, zelevels, cmap = palette, alpha = trans )
1435    ## draw colorbar
1436    if proj in ['moll','cyl']:        zeorientation="horizontal" ; zepad = 0.07
1437    else:                             zeorientation="vertical" ; zepad = 0.03
1438    #daformat = fmtvar(fvar.upper())
1439    daformat = "%.0f"
1440    zecb = mpl.colorbar( fraction=0.05,pad=zepad,format=daformat,orientation=zeorientation,\
1441                 ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional' ) 
1442    ## give a title
1443    if zeorientation == "horizontal": zecb.ax.set_xlabel(title)
1444    else:                             ptitle(title)
1445    ## draw vector
1446    if vecx is not None and vecy is not None:
1447       [vecx_frame,vecy_frame] = m.rotate_vector( np.transpose(vecx), np.transpose(vecy), lon2d, lat2d ) ## for metwinds
1448       vectorfield(vecx_frame, vecy_frame, x, y, stride=stride, csmooth=2,\
1449                                             scale=30., factor=500., color=definecolorvec(colorb), key=True)
1450    ## scale regle la reference du vecteur. factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
1451    return
1452## Author : AC
1453## Handles calls to specific computations (e.g. wind norm, enrichment factor...)
1454def select_getfield(zvarname=None,znc=None,ztypefile=None,mode=None,ztsat=None,ylon=None,ylat=None,yalt=None,ytime=None,analysis=None):
1455      from mymath import get_tsat
1456 
1457      ## Specific variables are described here:
1458      # for the mesoscale:
1459      specificname_meso = ['UV','uv','uvmet','slopexy','SLOPEXY','deltat','DELTAT','hodograph','tk','hodograph_2']
1460      # for the gcm:
1461      specificname_gcm = ['enfact']
1462
1463      ## Check for variable in file:
1464      if mode == 'check':     
1465          varname = zvarname
1466          varinfile=znc.variables.keys()
1467          logical_novarname = zvarname not in znc.variables
1468          logical_nospecificname_meso = not ((ztypefile in ['meso']) and (zvarname in specificname_meso))
1469          logical_nospecificname_gcm = not ((ztypefile in ['gcm']) and (zvarname in specificname_gcm))
1470          if ( logical_novarname and logical_nospecificname_meso and logical_nospecificname_gcm ):
1471              if len(varinfile) == 1:   varname = varinfile[0]
1472              else:                     varname = False
1473          ## Return the variable name:
1474          return varname
1475
1476      ## Get the corresponding variable:
1477      if mode == 'getvar':
1478          plot_x = None ; plot_y = None ;
1479          ### ----------- 1. saturation temperature
1480          if zvarname in ["temp","t","T_nadir_nit","T_nadir_day","temp_day","temp_night"] and ztsat:
1481              tt=getfield(znc,zvarname) ; print "computing Tsat-T, I ASSUME Z-AXIS IS PRESSURE"
1482              if type(tt).__name__=='MaskedArray':  tt.set_fill_value([np.NaN]) ; tinput=tt.filled()
1483              else:                                 tinput=tt
1484              all_var=get_tsat(yalt,tinput,zlon=ylon,zlat=ylat,zalt=yalt,ztime=ytime)
1485          ### ----------- 2. wind amplitude
1486          elif ((zvarname in ['UV','uv','uvmet']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1487              all_var=windamplitude(znc,'amplitude')
1488          elif ((zvarname in ['hodograph','hodograph_2']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1489              plot_x, plot_y = windamplitude(znc,zvarname)
1490              if plot_x is not None: all_var=plot_x # dummy
1491              else: all_var=plot_y ; plot_x = None ; plot_y = None # Hodograph type 2 is not 'xy' mode
1492          elif ((zvarname in ['slopexy','SLOPEXY']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1493              all_var=slopeamplitude(znc)
1494          ### ------------ 3. Near surface instability
1495          elif ((zvarname in ['DELTAT','deltat']) and (ztypefile in ['meso']) and (zvarname not in znc.variables)):
1496              all_var=deltat0t1(znc)
1497          ### ------------ 4. Enrichment factor
1498          elif ((ztypefile in ['gcm']) and (zvarname in ['enfact'])):
1499              all_var=enrichment_factor(znc,ylon,ylat,ytime)
1500          ### ------------ 5. teta -> temp
1501          elif ((ztypefile in ['meso']) and (zvarname in ['tk']) and ('tk' not in znc.variables.keys())):
1502              all_var=teta_to_tk(znc)
1503          else:
1504          ### -----------  999. Normal case
1505              all_var = getfield(znc,zvarname)
1506          if analysis is not None:
1507             if analysis in ['histo','density','histodensity']: plot_y=all_var ; plot_x = plot_y
1508             elif analysis == 'fft': plot_y, plot_x = spectrum(all_var,ytime,yalt,ylat,ylon) ; all_var = plot_y
1509          return all_var, plot_x, plot_y
1510
1511# Author : A.C
1512# FFT is computed before reducefield voluntarily, because we dont want to compute
1513# ffts on averaged fields (which would kill all waves). Instead, we take the fft everywhere
1514# (which is not efficient but it is still ok) and then, make the average (if the user wants to)
1515def spectrum(var,time,vert,lat,lon):
1516    import numpy as np
1517    fft=np.fft.fft(var,axis=1)
1518    N=len(vert)
1519    step=(vert[1]-vert[0])*1000.
1520    print "step is: ",step
1521    fftfreq=np.fft.fftfreq(N,d=step)
1522    fftfreq=np.fft.fftshift(fftfreq) # spatial FFT => this is the wavenumber
1523    fft=np.fft.fftshift(fft)
1524    fftfreq = 1./fftfreq # => wavelength (div by 0 expected, don't panic)
1525    fft=np.abs(fft) # => amplitude spectrum
1526#    fft=np.abs(fft)**2 # => power spectrum
1527    return fft,fftfreq
1528
1529# Author : A.C.
1530# Computes temperature from potential temperature for mesoscale files, without the need to use API, i.e. using natural vertical grid
1531def teta_to_tk(nc):
1532    import numpy as np
1533    varinfile = nc.variables.keys() 
1534    p0=610.
1535    t0=220.
1536    r_cp=1./3.89419
1537    if "T" in varinfile: zteta=getfield(nc,'T')
1538    else: errormess("you need T in your file.")
1539    if "PTOT" in varinfile: zptot=getfield(nc,'PTOT')
1540    else: errormess("you need PTOT in your file.")
1541    zt=(zteta+220.)*(zptot/p0)**(r_cp)
1542    return zt
1543
1544# Author : A.C.
1545# Find the lon and lat index of the dust devil with the largest pressure gradient
1546# Steps :
1547# 1/ convert the chosen PSFC frame to an image of the PSFC anomaly with respect to the mean
1548# 2/ apply the Sobel operator
1549# (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.)
1550# 3/ find the maximum of the resulting field
1551# 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
1552# 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)
1553# 6/ in this slab, find the point at which the surface pressure is minimum
1554def find_devil(nc,indextime):
1555    import numpy as np
1556    from scipy import ndimage
1557    from mymath import array2image,image2array
1558
1559    varinfile = nc.variables.keys()
1560    if "PSFC" not in varinfile: errormess("You need PSFC in your file to find dust devils")
1561    else: psfc_full=getfield(nc,'PSFC')
1562    psfc,error=reducefield( psfc_full, d4=indextime)
1563    psfcim=array2image(1000.*(psfc-psfc.mean()))
1564    sx = ndimage.sobel(psfcim, axis=0, mode='constant') ; sy = ndimage.sobel(psfcim, axis=1, mode='constant')
1565    sob = np.hypot(sx, sy)
1566    zemax=np.max(sob)
1567    goodvalues = sob[sob >= zemax/2]
1568    ix = np.in1d(sob.ravel(), goodvalues).reshape(sob.shape)
1569    idxs,idys=np.where(ix)
1570    maxvalue = sob[sob == zemax]
1571    ixmax = np.in1d(sob.ravel(), maxvalue[0]).reshape(sob.shape)
1572    idxmax,idymax=np.where(ixmax)
1573    valok=[]
1574    for i in np.arange(len(idxs)):
1575        a=np.sqrt((idxmax-idxs[i])**2 + (idymax-idys[i])**2)
1576        if 0 < a <= 5.*np.sqrt(2.): valok.append(goodvalues[i])
1577    ix = np.in1d(sob.ravel(), valok).reshape(sob.shape)
1578    idxs,idys=np.where(ix)
1579    hyperslab=psfc[np.min(idxs):np.max(idxs),np.min(idys):np.max(idys)]
1580    idxsub,idysub=np.where(hyperslab==hyperslab.min())
1581    idx=idxsub[0]+np.min(idxs) ; idy=idysub[0]+np.min(idys)
1582    return np.int(idx),np.int(idy)
Note: See TracBrowser for help on using the repository browser.