source: trunk/UTIL/PYTHON/planetoplot_v1/myplot.py @ 1062

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

clean and organized UTIL/PYTHON folder

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