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

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

UTIL PYTHON. A bug fix for previous commit. Ensure that what is coming from getfieldred has not changed dimension.

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