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

Last change on this file since 418 was 418, checked in by acolaitis, 13 years ago

PYTHON. Moved the nolow mode to an option (--nolow) and re-arranged the way -H worked in planetoplot, by introducing a new function (hole_bounds). Bounds is called in normal cases and hole_bounds in -H cases. Minor corrections in mcs.py for missing values comming from zrecast.
File size: 40.2 KB
Line 
1## Author: AS
2def errormess(text,printvar=None):
3    print text
4    if printvar is not None: print printvar
5    exit()
6    return
7
8## Author: AS
9def adjust_length (tab, zelen):
10    from numpy import ones
11    if tab is None:
12        outtab = ones(zelen) * -999999
13    else:
14        if zelen != len(tab):
15            print "not enough or too much values... setting same values all variables"
16            outtab = ones(zelen) * tab[0]
17        else:
18            outtab = tab
19    return outtab
20
21## Author: AS
22def getname(var=False,winds=False,anomaly=False):
23    if var and winds:     basename = var + '_UV'
24    elif var:             basename = var
25    elif winds:           basename = 'UV'
26    else:                 errormess("please set at least winds or var",printvar=nc.variables)
27    if anomaly:           basename = 'd' + basename
28    return basename
29
30## Author: AS
31def localtime(utc,lon):
32    ltst = utc + lon / 15.
33    ltst = int (ltst * 10) / 10.
34    ltst = ltst % 24
35    return ltst
36
37## Author: AS
38def whatkindfile (nc):
39    if 'controle' in nc.variables:   typefile = 'gcm'
40    elif 'phisinit' in nc.variables: typefile = 'gcm' 
41    elif 'vert' in nc.variables:     typefile = 'mesoapi'
42    elif 'U' in nc.variables:        typefile = 'meso'
43    elif 'HGT_M' in nc.variables:    typefile = 'geo'
44    #else:                            errormess("whatkindfile: typefile not supported.")
45    else:                            typefile = 'gcm' # for lslin-ed files from gcm
46    return typefile
47
48## Author: AS
49def getfield (nc,var):
50    ## this allows to get much faster (than simply referring to nc.variables[var])
51    import numpy as np
52    dimension = len(nc.variables[var].dimensions)
53    #print "   Opening variable with", dimension, "dimensions ..."
54    if dimension == 2:    field = nc.variables[var][:,:]
55    elif dimension == 3:  field = nc.variables[var][:,:,:]
56    elif dimension == 4:  field = nc.variables[var][:,:,:,:]
57    # if there are NaNs in the ncdf, they should be loaded as a masked array which will be
58    # recasted as a regular array later in reducefield
59    if (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
60       print "Warning: netcdf as nan values but is not loaded as a Masked Array."
61       print "recasting array type"
62       out=np.ma.masked_invalid(field)
63       out.set_fill_value([np.NaN])
64    else:
65       out=field
66    return out
67
68## Author: AS + TN + AC
69def reducefield (input,d4=None,d3=None,d2=None,d1=None,yint=False,alt=None,anomaly=False):
70    ### we do it the reverse way to be compliant with netcdf "t z y x" or "t y x" or "y x"
71    ### it would be actually better to name d4 d3 d2 d1 as t z y x
72    ### ... note, anomaly is only computed over d1 and d2 for the moment
73    import numpy as np
74    from mymath import max,mean
75    csmooth = 12 ## a fair amount of grid points
76    dimension = np.array(input).ndim
77    shape = np.array(input).shape
78    #print 'd1,d2,d3,d4: ',d1,d2,d3,d4
79    print 'IN  REDUCEFIELD dim,shape: ',dimension,shape
80    if anomaly: print 'ANOMALY ANOMALY'
81    output = input
82    error = False
83    #### this is needed to cope the case where d4,d3,d2,d1 are single integers and not arrays
84    if d4 is not None and not isinstance(d4, np.ndarray): d4=[d4]
85    if d3 is not None and not isinstance(d3, np.ndarray): d3=[d3]
86    if d2 is not None and not isinstance(d2, np.ndarray): d2=[d2]
87    if d1 is not None and not isinstance(d1, np.ndarray): d1=[d1]
88    ### now the main part
89    if dimension == 2:
90        if   d2 >= shape[0]: error = True
91        elif d1 >= shape[1]: error = True
92        elif d1 is not None and d2 is not None:  output = mean(input[d2,:],axis=0); output = mean(output[d1],axis=0)
93        elif d1 is not None:         output = mean(input[:,d1],axis=1)
94        elif d2 is not None:         output = mean(input[d2,:],axis=0)
95    elif dimension == 3:
96        if   max(d4) >= shape[0]: error = True
97        elif max(d2) >= shape[1]: error = True
98        elif max(d1) >= shape[2]: error = True
99        elif d4 is not None and d2 is not None and d1 is not None: 
100            output = mean(input[d4,:,:],axis=0); output = mean(output[d2,:],axis=0); output = mean(output[d1],axis=0)
101        elif d4 is not None and d2 is not None:    output = mean(input[d4,:,:],axis=0); output=mean(output[d2,:],axis=0)
102        elif d4 is not None and d1 is not None:    output = mean(input[d4,:,:],axis=0); output=mean(output[:,d1],axis=1)
103        elif d2 is not None and d1 is not None:    output = mean(input[:,d2,:],axis=1); output=mean(output[:,d1],axis=1)
104        elif d1 is not None:                       output = mean(input[:,:,d1],axis=2)
105        elif d2 is not None:                       output = mean(input[:,d2,:],axis=1)
106        elif d4 is not None:                       output = mean(input[d4,:,:],axis=0)
107    elif dimension == 4:
108        if   max(d4) >= shape[0]: error = True
109        elif max(d3) >= shape[1]: error = True
110        elif max(d2) >= shape[2]: error = True
111        elif max(d1) >= shape[3]: error = True
112        elif d4 is not None and d3 is not None and d2 is not None and d1 is not None:
113             output = mean(input[d4,:,:,:],axis=0)
114             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
115             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
116             output = mean(output[d2,:],axis=0)
117             output = mean(output[d1],axis=0)
118        elif d4 is not None and d3 is not None and d2 is not None: 
119             output = mean(input[d4,:,:,:],axis=0)
120             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
121             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
122             output = mean(output[d2,:],axis=0)
123        elif d4 is not None and d3 is not None and d1 is not None: 
124             output = mean(input[d4,:,:,:],axis=0)
125             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
126             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
127             output = mean(output[:,d1],axis=1)
128        elif d4 is not None and d2 is not None and d1 is not None: 
129             output = mean(input[d4,:,:,:],axis=0)
130             if anomaly:
131                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
132             output = mean(output[:,d2,:],axis=1)
133             output = mean(output[:,d1],axis=1)
134             #noperturb = smooth1d(output,window_len=7)
135             #lenlen = len(output) ; output = output[1:lenlen-7] ; yeye = noperturb[4:lenlen-4]
136             #plot(output) ; plot(yeye) ; show() ; plot(output-yeye) ; show()
137        elif d3 is not None and d2 is not None and d1 is not None: 
138             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3) 
139             if anomaly:
140                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
141             output = mean(output[:,d2,:],axis=1)
142             output = mean(output[:,d1],axis=1) 
143        elif d4 is not None and d3 is not None: 
144             output = mean(input[d4,:,:,:],axis=0) 
145             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
146             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
147        elif d4 is not None and d2 is not None: 
148             output = mean(input[d4,:,:,:],axis=0) 
149             output = mean(output[:,d2,:],axis=1)
150        elif d4 is not None and d1 is not None: 
151             output = mean(input[d4,:,:,:],axis=0) 
152             output = mean(output[:,:,d1],axis=2)
153        elif d3 is not None and d2 is not None: 
154             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
155             output = mean(output[:,d2,:],axis=1)
156        elif d3 is not None and d1 is not None: 
157             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
158             output = mean(output[:,:,d1],axis=0)
159        elif d2 is not None and d1 is not None: 
160             output = mean(input[:,:,d2,:],axis=2)
161             output = mean(output[:,:,d1],axis=2)
162        elif d1 is not None:        output = mean(input[:,:,:,d1],axis=3)
163        elif d2 is not None:        output = mean(input[:,:,d2,:],axis=2)
164        elif d3 is not None:        output = reduce_zaxis(input[:,d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
165        elif d4 is not None:        output = mean(input[d4,:,:,:],axis=0)
166    dimension = np.array(output).ndim
167    shape = np.array(output).shape
168    print 'OUT REDUCEFIELD dim,shape: ',dimension,shape
169    return output, error
170
171## Author: AC + AS
172def reduce_zaxis (input,ax=None,yint=False,vert=None,indice=None):
173    from mymath import max,mean
174    from scipy import integrate
175    if yint and vert is not None and indice is not None:
176       if type(input).__name__=='MaskedArray':
177          input.set_fill_value([np.NaN])
178          output = integrate.trapz(input.filled(),x=vert[indice],axis=ax)
179       else:
180          output = integrate.trapz(input,x=vert[indice],axis=ax)
181    else:
182       output = mean(input,axis=ax)
183    return output
184
185## Author: AS + TN
186def definesubplot ( numplot, fig ):
187    from matplotlib.pyplot import rcParams
188    rcParams['font.size'] = 12. ## default (important for multiple calls)
189    if numplot <= 0:
190        subv = 99999
191        subh = 99999
192    elif numplot == 1: 
193        subv = 99999
194        subh = 99999
195    elif numplot == 2:
196        subv = 1
197        subh = 2
198        fig.subplots_adjust(wspace = 0.35)
199        rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
200    elif numplot == 3:
201        subv = 2
202        subh = 2
203        fig.subplots_adjust(wspace = 0.5)
204        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
205    elif   numplot == 4:
206        subv = 2
207        subh = 2
208        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
209        rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
210    elif numplot <= 6:
211        subv = 2
212        subh = 3
213        fig.subplots_adjust(wspace = 0.4, hspace = 0.0)
214        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
215    elif numplot <= 8:
216        subv = 2
217        subh = 4
218        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
219        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
220    elif numplot <= 9:
221        subv = 3
222        subh = 3
223        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
224        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
225    elif numplot <= 12:
226        subv = 3
227        subh = 4
228        fig.subplots_adjust(wspace = 0.1, hspace = 0.1)
229        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
230    elif numplot <= 16:
231        subv = 4
232        subh = 4
233        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
234        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
235    else:
236        print "number of plot supported: 1 to 16"
237        exit()
238    return subv,subh
239
240## Author: AS
241def getstralt(nc,nvert):
242    typefile = whatkindfile(nc)
243    if typefile is 'meso':                     
244        stralt = "_lvl" + str(nvert)
245    elif typefile is 'mesoapi':
246        zelevel = int(nc.variables['vert'][nvert])
247        if abs(zelevel) < 10000.:   strheight=str(zelevel)+"m"
248        else:                       strheight=str(int(zelevel/1000.))+"km"
249        if 'altitude'       in nc.dimensions:   stralt = "_"+strheight+"-AMR"
250        elif 'altitude_abg' in nc.dimensions:   stralt = "_"+strheight+"-ALS"
251        elif 'bottom_top'   in nc.dimensions:   stralt = "_"+strheight
252        elif 'pressure'     in nc.dimensions:   stralt = "_"+str(zelevel)+"Pa"
253        else:                                   stralt = ""
254    else:
255        stralt = ""
256    return stralt
257
258## Author: AS
259def getlschar ( namefile ):
260    from netCDF4 import Dataset
261    from timestuff import sol2ls
262    from numpy import array
263    from string import rstrip
264    namefile = rstrip( rstrip( rstrip( namefile, chars="_z"), chars="_zabg"), chars="_p")
265    #### we assume that wrfout is next to wrfout_z and wrfout_zabg
266    nc  = Dataset(namefile)
267    zetime = None
268    if 'Times' in nc.variables:
269        zetime = nc.variables['Times'][0]
270        shape = array(nc.variables['Times']).shape
271        if shape[0] < 2: zetime = None
272    if zetime is not None \
273       and 'vert' not in nc.variables:
274        #### strangely enough this does not work for api or ncrcat results!
275        zetimestart = getattr(nc, 'START_DATE')
276        zeday = int(zetime[8]+zetime[9]) - int(zetimestart[8]+zetimestart[9])
277        if zeday < 0:    lschar=""  ## might have crossed a month... fix soon
278        else:            lschar="_Ls"+str( int( 10. * sol2ls ( getattr( nc, 'JULDAY' ) + zeday ) ) / 10. )
279        ###
280        zetime2 = nc.variables['Times'][1]
281        one  = int(zetime[11]+zetime[12]) + int(zetime[14]+zetime[15])/37.
282        next = int(zetime2[11]+zetime2[12]) + int(zetime2[14]+zetime2[15])/37. 
283        zehour    = one
284        zehourin  = abs ( next - one )
285    else:
286        lschar=""
287        zehour = 0
288        zehourin = 1 
289    return lschar, zehour, zehourin
290
291## Author: AS
292def getprefix (nc):
293    prefix = 'LMD_MMM_'
294    prefix = prefix + 'd'+str(getattr(nc,'GRID_ID'))+'_'
295    prefix = prefix + str(int(getattr(nc,'DX')/1000.))+'km_'
296    return prefix
297
298## Author: AS
299def getproj (nc):
300    typefile = whatkindfile(nc)
301    if typefile in ['mesoapi','meso','geo']:
302        ### (il faudrait passer CEN_LON dans la projection ?)
303        map_proj = getattr(nc, 'MAP_PROJ')
304        cen_lat  = getattr(nc, 'CEN_LAT')
305        if map_proj == 2:
306            if cen_lat > 10.:   
307                proj="npstere"
308                #print "NP stereographic polar domain"
309            else:           
310                proj="spstere"
311                #print "SP stereographic polar domain"
312        elif map_proj == 1: 
313            #print "lambert projection domain"
314            proj="lcc"
315        elif map_proj == 3: 
316            #print "mercator projection"
317            proj="merc"
318        else:
319            proj="merc"
320    elif typefile in ['gcm']:  proj="cyl"    ## pb avec les autres (de trace derriere la sphere ?)
321    else:                      proj="ortho"
322    return proj   
323
324## Author: AS
325def ptitle (name):
326    from matplotlib.pyplot import title
327    title(name)
328    print name
329
330## Author: AS
331def polarinterv (lon2d,lat2d):
332    import numpy as np
333    wlon = [np.min(lon2d),np.max(lon2d)]
334    ind = np.array(lat2d).shape[0] / 2  ## to get a good boundlat and to get the pole
335    wlat = [np.min(lat2d[ind,:]),np.max(lat2d[ind,:])]
336    return [wlon,wlat]
337
338## Author: AS
339def simplinterv (lon2d,lat2d):
340    import numpy as np
341    return [[np.min(lon2d),np.max(lon2d)],[np.min(lat2d),np.max(lat2d)]]
342
343## Author: AS
344def wrfinterv (lon2d,lat2d):
345    nx = len(lon2d[0,:])-1
346    ny = len(lon2d[:,0])-1
347    lon1 = lon2d[0,0] 
348    lon2 = lon2d[nx,ny] 
349    lat1 = lat2d[0,0] 
350    lat2 = lat2d[nx,ny] 
351    if abs(0.5*(lat1+lat2)) > 60.:  wider = 0.5 * (abs(lon1)+abs(lon2)) * 0.1
352    else:                           wider = 0.
353    if lon1 < lon2:  wlon = [lon1, lon2 + wider] 
354    else:            wlon = [lon2, lon1 + wider]
355    if lat1 < lat2:  wlat = [lat1, lat2]
356    else:            wlat = [lat2, lat1]
357    return [wlon,wlat]
358
359## Author: AS
360def makeplotres (filename,res=None,pad_inches_value=0.25,folder='',disp=True,ext='png',erase=False):
361    import  matplotlib.pyplot as plt
362    from os import system
363    addstr = ""
364    if res is not None:
365        res = int(res)
366        addstr = "_"+str(res)
367    name = filename+addstr+"."+ext
368    if folder != '':      name = folder+'/'+name
369    plt.savefig(name,dpi=res,bbox_inches='tight',pad_inches=pad_inches_value)
370    if disp:              display(name)
371    if ext in ['eps','ps','svg']:   system("tar czvf "+name+".tar.gz "+name+" ; rm -f "+name)
372    if erase:   system("mv "+name+" to_be_erased")             
373    return
374
375## Author: AS
376def dumpbdy (field,n,stag=None):
377    nx = len(field[0,:])-1
378    ny = len(field[:,0])-1
379    if stag == 'U': nx = nx-1
380    if stag == 'V': ny = ny-1
381    return field[n:ny-n,n:nx-n]
382
383## Author: AS
384def getcoorddef ( nc ):   
385    import numpy as np
386    ## getcoord2d for predefined types
387    typefile = whatkindfile(nc)
388    if typefile in ['mesoapi','meso']:
389        [lon2d,lat2d] = getcoord2d(nc)
390        lon2d = dumpbdy(lon2d,6)
391        lat2d = dumpbdy(lat2d,6)
392    elif typefile in ['gcm']: 
393        [lon2d,lat2d] = getcoord2d(nc,nlat="latitude",nlon="longitude",is1d=True)
394    elif typefile in ['geo']:
395        [lon2d,lat2d] = getcoord2d(nc,nlat='XLAT_M',nlon='XLONG_M')
396    return lon2d,lat2d   
397
398## Author: AS
399def getcoord2d (nc,nlat='XLAT',nlon='XLONG',is1d=False):
400    import numpy as np
401    if is1d:
402        lat = nc.variables[nlat][:]
403        lon = nc.variables[nlon][:]
404        [lon2d,lat2d] = np.meshgrid(lon,lat)
405    else:
406        lat = nc.variables[nlat][0,:,:]
407        lon = nc.variables[nlon][0,:,:]
408        [lon2d,lat2d] = [lon,lat]
409    return lon2d,lat2d
410
411## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
412def smooth1d(x,window_len=11,window='hanning'):
413    import numpy
414    """smooth the data using a window with requested size.
415    This method is based on the convolution of a scaled window with the signal.
416    The signal is prepared by introducing reflected copies of the signal
417    (with the window size) in both ends so that transient parts are minimized
418    in the begining and end part of the output signal.
419    input:
420        x: the input signal
421        window_len: the dimension of the smoothing window; should be an odd integer
422        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
423            flat window will produce a moving average smoothing.
424    output:
425        the smoothed signal
426    example:
427    t=linspace(-2,2,0.1)
428    x=sin(t)+randn(len(t))*0.1
429    y=smooth(x)
430    see also:
431    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
432    scipy.signal.lfilter
433    TODO: the window parameter could be the window itself if an array instead of a string   
434    """
435    x = numpy.array(x)
436    if x.ndim != 1:
437        raise ValueError, "smooth only accepts 1 dimension arrays."
438    if x.size < window_len:
439        raise ValueError, "Input vector needs to be bigger than window size."
440    if window_len<3:
441        return x
442    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
443        raise ValueError, "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"
444    s=numpy.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
445    #print(len(s))
446    if window == 'flat': #moving average
447        w=numpy.ones(window_len,'d')
448    else:
449        w=eval('numpy.'+window+'(window_len)')
450    y=numpy.convolve(w/w.sum(),s,mode='valid')
451    return y
452
453## Author: AS
454def smooth (field, coeff):
455        ## actually blur_image could work with different coeff on x and y
456        if coeff > 1:   result = blur_image(field,int(coeff))
457        else:           result = field
458        return result
459
460## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
461def gauss_kern(size, sizey=None):
462        import numpy as np
463        # Returns a normalized 2D gauss kernel array for convolutions
464        size = int(size)
465        if not sizey:
466                sizey = size
467        else:
468                sizey = int(sizey)
469        x, y = np.mgrid[-size:size+1, -sizey:sizey+1]
470        g = np.exp(-(x**2/float(size)+y**2/float(sizey)))
471        return g / g.sum()
472
473## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
474def blur_image(im, n, ny=None) :
475        from scipy.signal import convolve
476        # blurs the image by convolving with a gaussian kernel of typical size n.
477        # The optional keyword argument ny allows for a different size in the y direction.
478        g = gauss_kern(n, sizey=ny)
479        improc = convolve(im, g, mode='same')
480        return improc
481
482## Author: AS
483def getwinddef (nc):   
484    ## getwinds for predefined types
485    typefile = whatkindfile(nc)
486    ###
487    if typefile is 'mesoapi':    [uchar,vchar] = ['Um','Vm']
488    elif typefile is 'gcm':      [uchar,vchar] = ['u','v']
489    elif typefile is 'meso':     [uchar,vchar] = ['U','V']
490    else:                        [uchar,vchar] = ['not found','not found']
491    ###
492    if typefile in ['meso']:     metwind = False ## geometrical (wrt grid)
493    else:                        metwind = True  ## meteorological (zon/mer)
494    if metwind is False:         print "Not using meteorological winds. You trust numerical grid as being (x,y)"
495    ###
496    return uchar,vchar,metwind
497
498## Author: AS
499def vectorfield (u, v, x, y, stride=3, scale=15., factor=250., color='black', csmooth=1, key=True):
500    ## scale regle la reference du vecteur
501    ## factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
502    import  matplotlib.pyplot               as plt
503    import  numpy                           as np
504    posx = np.min(x) - np.std(x) / 10.
505    posy = np.min(y) - np.std(y) / 10.
506    u = smooth(u,csmooth)
507    v = smooth(v,csmooth)
508    widthvec = 0.003 #0.005 #0.003
509    q = plt.quiver( x[::stride,::stride],\
510                    y[::stride,::stride],\
511                    u[::stride,::stride],\
512                    v[::stride,::stride],\
513                    angles='xy',color=color,pivot='middle',\
514                    scale=factor,width=widthvec )
515    if color in ['white','yellow']:     kcolor='black'
516    else:                               kcolor=color
517    if key: p = plt.quiverkey(q,posx,posy,scale,\
518                   str(int(scale)),coordinates='data',color=kcolor,labelpos='S',labelsep = 0.03)
519    return 
520
521## Author: AS
522def display (name):
523    from os import system
524    system("display "+name+" > /dev/null 2> /dev/null &")
525    return name
526
527## Author: AS
528def findstep (wlon):
529    steplon = int((wlon[1]-wlon[0])/4.)  #3
530    step = 120.
531    while step > steplon and step > 15. :       step = step / 2.
532    if step <= 15.:
533        while step > steplon and step > 5.  :   step = step - 5.
534    if step <= 5.:
535        while step > steplon and step > 1.  :   step = step - 1.
536    if step <= 1.:
537        step = 1. 
538    return step
539
540## Author: AS
541def define_proj (char,wlon,wlat,back=None,blat=None):
542    from    mpl_toolkits.basemap            import Basemap
543    import  numpy                           as np
544    import  matplotlib                      as mpl
545    from mymath import max
546    meanlon = 0.5*(wlon[0]+wlon[1])
547    meanlat = 0.5*(wlat[0]+wlat[1])
548    if blat is None:
549        ortholat=meanlat
550        if   wlat[0] >= 80.:   blat =  40. 
551        elif wlat[1] <= -80.:  blat = -40.
552        elif wlat[1] >= 0.:    blat = wlat[0] 
553        elif wlat[0] <= 0.:    blat = wlat[1]
554    else:  ortholat=blat
555    #print "blat ", blat
556    h = 50.  ## en km
557    radius = 3397200.
558    if   char == "cyl":     m = Basemap(rsphere=radius,projection='cyl',\
559                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
560    elif char == "moll":    m = Basemap(rsphere=radius,projection='moll',lon_0=meanlon)
561    elif char == "ortho":   m = Basemap(rsphere=radius,projection='ortho',lon_0=meanlon,lat_0=ortholat)
562    elif char == "lcc":     m = Basemap(rsphere=radius,projection='lcc',lat_1=meanlat,lat_0=meanlat,lon_0=meanlon,\
563                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
564    elif char == "npstere": m = Basemap(rsphere=radius,projection='npstere', boundinglat=blat, lon_0=0.)
565    elif char == "spstere": m = Basemap(rsphere=radius,projection='spstere', boundinglat=blat, lon_0=180.)
566    elif char == "nplaea":  m = Basemap(rsphere=radius,projection='nplaea', boundinglat=wlat[0], lon_0=meanlon)
567    elif char == "laea":    m = Basemap(rsphere=radius,projection='laea',lon_0=meanlon,lat_0=meanlat,lat_ts=meanlat,\
568                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
569    elif char == "nsper":   m = Basemap(rsphere=radius,projection='nsper',lon_0=meanlon,lat_0=meanlat,satellite_height=h*1000.)
570    elif char == "merc":    m = Basemap(rsphere=radius,projection='merc',lat_ts=0.,\
571                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
572    fontsizemer = int(mpl.rcParams['font.size']*3./4.)
573    if char in ["cyl","lcc","merc","nsper","laea"]:   step = findstep(wlon)
574    else:                                             step = 10.
575    steplon = step*2.
576    #if back in ["geolocal"]:                         
577    #    step = np.min([5.,step])
578    #    steplon = step
579    m.drawmeridians(np.r_[-180.:180.:steplon], labels=[0,0,0,1], color='grey', fontsize=fontsizemer)
580    m.drawparallels(np.r_[-90.:90.:step], labels=[1,0,0,0], color='grey', fontsize=fontsizemer)
581    if back: m.warpimage(marsmap(back),scale=0.75)
582            #if not back:
583            #    if not var:                                        back = "mola"    ## if no var:         draw mola
584            #    elif typefile in ['mesoapi','meso','geo'] \
585            #       and proj not in ['merc','lcc','nsper','laea']:  back = "molabw"  ## if var but meso:   draw molabw
586            #    else:                                              pass             ## else:              draw None
587    return m
588
589## Author: AS
590#### test temporaire
591def putpoints (map,plot):
592    #### from http://www.scipy.org/Cookbook/Matplotlib/Maps
593    # lat/lon coordinates of five cities.
594    lats = [18.4]
595    lons = [-134.0]
596    points=['Olympus Mons']
597    # compute the native map projection coordinates for cities.
598    x,y = map(lons,lats)
599    # plot filled circles at the locations of the cities.
600    map.plot(x,y,'bo')
601    # plot the names of those five cities.
602    wherept = 0 #1000 #50000
603    for name,xpt,ypt in zip(points,x,y):
604       plot.text(xpt+wherept,ypt+wherept,name)
605    ## le nom ne s'affiche pas...
606    return
607
608## Author: AS
609def calculate_bounds(field,vmin=None,vmax=None):
610    import numpy as np
611    from mymath import max,min,mean
612    ind = np.where(field < 9e+35)
613    fieldcalc = field[ ind ] # la syntaxe compacte ne marche si field est un tuple
614    ###
615    dev = np.std(fieldcalc)*3.0
616    ###
617    if vmin is None:
618        zevmin = mean(fieldcalc) - dev
619    else:             zevmin = vmin
620    ###
621    if vmax is None:  zevmax = mean(fieldcalc) + dev
622    else:             zevmax = vmax
623    if vmin == vmax:
624                      zevmin = mean(fieldcalc) - dev  ### for continuity
625                      zevmax = mean(fieldcalc) + dev  ### for continuity           
626    ###
627    if zevmin < 0. and min(fieldcalc) > 0.: zevmin = 0.
628    print "BOUNDS field ", min(fieldcalc), max(fieldcalc)
629    print "BOUNDS adopted ", zevmin, zevmax
630    return zevmin, zevmax
631
632## Author: AS
633def bounds(what_I_plot,zevmin,zevmax):
634    from mymath import max,min,mean
635    ### might be convenient to add the missing value in arguments
636    #what_I_plot[ what_I_plot < zevmin ] = zevmin#*(1. + 1.e-7)
637    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1. - 1.e-7) ] = zevmin*(1. - 1.e-7)
638    else:          what_I_plot[ what_I_plot < zevmin*(1. + 1.e-7) ] = zevmin*(1. + 1.e-7)
639    print "NEW MIN ", min(what_I_plot)
640    what_I_plot[ what_I_plot > 9e+35  ] = -9e+35
641    what_I_plot[ what_I_plot > zevmax ] = zevmax
642    print "NEW MAX ", max(what_I_plot)
643    return what_I_plot
644
645## Author: AS
646def nolow(what_I_plot):
647    from mymath import max,min
648    lim = 0.15*0.5*(abs(max(what_I_plot))+abs(min(what_I_plot)))
649    print "NO PLOT BELOW VALUE ", lim
650    what_I_plot [ abs(what_I_plot) < lim ] = 1.e40 
651    return what_I_plot
652
653
654## Author : AC
655def hole_bounds(what_I_plot,zevmin,zevmax):
656    import numpy as np
657    zi=0
658    for i in what_I_plot:
659        zj=0
660        for j in i:
661            if ((j < zevmin) or (j > zevmax)):what_I_plot[zi,zj]=np.NaN
662            zj=zj+1
663        zi=zi+1
664
665    return what_I_plot
666
667## Author: AS
668def zoomset (wlon,wlat,zoom):
669    dlon = abs(wlon[1]-wlon[0])/2.
670    dlat = abs(wlat[1]-wlat[0])/2.
671    [wlon,wlat] = [ [wlon[0]+zoom*dlon/100.,wlon[1]-zoom*dlon/100.],\
672                    [wlat[0]+zoom*dlat/100.,wlat[1]-zoom*dlat/100.] ]
673    print "ZOOM %",zoom,wlon,wlat
674    return wlon,wlat
675
676## Author: AS
677def fmtvar (whichvar="def"):
678    fmtvar    =     { \
679             "TK":           "%.0f",\
680# Variables from TES ncdf format
681             "T_NADIR_DAY":  "%.0f",\
682             "T_NADIR_NIT":  "%.0f",\
683# Variables from tes.py ncdf format
684             "TEMP_DAY":     "%.0f",\
685             "TEMP_NIGHT":   "%.0f",\
686# Variables from MCS and mcs.py ncdf format
687             "DTEMP":     "%.0f",\
688             "NTEMP":   "%.0f",\
689             "TPOT":         "%.0f",\
690             "TSURF":        "%.0f",\
691             "def":          "%.1e",\
692             "PTOT":         "%.0f",\
693             "HGT":          "%.1e",\
694             "USTM":         "%.2f",\
695             "HFX":          "%.0f",\
696             "ICETOT":       "%.1e",\
697             "TAU_ICE":      "%.2f",\
698             "VMR_ICE":      "%.1e",\
699             "MTOT":         "%.1f",\
700             "ANOMALY":      "%.1f",\
701             "W":            "%.1f",\
702             "WMAX_TH":      "%.1f",\
703             "QSURFICE":     "%.0f",\
704             "UM":           "%.0f",\
705             "ALBBARE":      "%.2f",\
706             "TAU":          "%.1f",\
707             "CO2":          "%.2f",\
708             #### T.N.
709             "TEMP":         "%.0f",\
710             "VMR_H2OICE":   "%.0f",\
711             "VMR_H2OVAP":   "%.0f",\
712             "TAUTES":       "%.2f",\
713             "TAUTESAP":     "%.2f",\
714
715                    }
716    if whichvar not in fmtvar:
717        whichvar = "def"
718    return fmtvar[whichvar]
719
720## Author: AS
721####################################################################################################################
722### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
723def defcolorb (whichone="def"):
724    whichcolorb =    { \
725             "def":          "spectral",\
726             "HGT":          "spectral",\
727             "TK":           "gist_heat",\
728             "TSURF":        "RdBu_r",\
729             "QH2O":         "PuBu",\
730             "USTM":         "YlOrRd",\
731             #"T_nadir_nit":  "RdBu_r",\
732             #"T_nadir_day":  "RdBu_r",\
733             "HFX":          "RdYlBu",\
734             "ICETOT":       "YlGnBu_r",\
735             #"MTOT":         "PuBu",\
736             "CCNQ":         "YlOrBr",\
737             "CCNN":         "YlOrBr",\
738             "TEMP":         "Jet",\
739             "TAU_ICE":      "Blues",\
740             "VMR_ICE":      "Blues",\
741             "W":            "jet",\
742             "WMAX_TH":      "spectral",\
743             "ANOMALY":      "RdBu_r",\
744             "QSURFICE":     "hot_r",\
745             "ALBBARE":      "spectral",\
746             "TAU":          "YlOrBr_r",\
747             "CO2":          "YlOrBr_r",\
748             #### T.N.
749             "MTOT":         "Jet",\
750             "H2O_ICE_S":    "RdBu",\
751             "VMR_H2OICE":   "PuBu",\
752             "VMR_H2OVAP":   "PuBu",\
753                     }
754#W --> spectral ou jet
755#spectral BrBG RdBu_r
756    #print "predefined colorbars"
757    if whichone not in whichcolorb:
758        whichone = "def"
759    return whichcolorb[whichone]
760
761## Author: AS
762def definecolorvec (whichone="def"):
763        whichcolor =    { \
764                "def":          "black",\
765                "vis":          "yellow",\
766                "vishires":     "yellow",\
767                "molabw":       "yellow",\
768                "mola":         "black",\
769                "gist_heat":    "white",\
770                "hot":          "tk",\
771                "gist_rainbow": "black",\
772                "spectral":     "black",\
773                "gray":         "red",\
774                "PuBu":         "black",\
775                        }
776        if whichone not in whichcolor:
777                whichone = "def"
778        return whichcolor[whichone]
779
780## Author: AS
781def marsmap (whichone="vishires"):
782        from os import uname
783        mymachine = uname()[1]
784        ### not sure about speed-up with this method... looks the same
785        if "lmd.jussieu.fr" in mymachine: domain = "/u/aslmd/WWW/maps/"
786        else:                             domain = "http://www.lmd.jussieu.fr/~aslmd/maps/"
787        whichlink =     { \
788                #"vis":         "http://maps.jpl.nasa.gov/pix/mar0kuu2.jpg",\
789                #"vishires":    "http://www.lmd.jussieu.fr/~aslmd/maps/MarsMap_2500x1250.jpg",\
790                #"geolocal":    "http://dl.dropbox.com/u/11078310/geolocal.jpg",\
791                #"mola":        "http://www.lns.cornell.edu/~seb/celestia/mars-mola-2k.jpg",\
792                #"molabw":      "http://dl.dropbox.com/u/11078310/MarsElevation_2500x1250.jpg",\
793                "vis":         domain+"mar0kuu2.jpg",\
794                "vishires":    domain+"MarsMap_2500x1250.jpg",\
795                "geolocal":    domain+"geolocal.jpg",\
796                "mola":        domain+"mars-mola-2k.jpg",\
797                "molabw":      domain+"MarsElevation_2500x1250.jpg",\
798                "clouds":      "http://www.johnstonsarchive.net/spaceart/marswcloudmap.jpg",\
799                "jupiter":     "http://www.mmedia.is/~bjj/data/jupiter_css/jupiter_css.jpg",\
800                "jupiter_voy": "http://www.mmedia.is/~bjj/data/jupiter/jupiter_vgr2.jpg",\
801                "bw":          "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg",\
802                "contrast":    "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg",\
803                "nice":        "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg",\
804                "blue":        "http://eoimages.gsfc.nasa.gov/ve/2430/land_ocean_ice_2048.jpg",\
805                "blueclouds":  "http://eoimages.gsfc.nasa.gov/ve/2431/land_ocean_ice_cloud_2048.jpg",\
806                "justclouds":  "http://eoimages.gsfc.nasa.gov/ve/2432/cloud_combined_2048.jpg",\
807                        }
808        ### see http://www.mmedia.is/~bjj/planetary_maps.html
809        if whichone not in whichlink: 
810                print "marsmap: choice not defined... you'll get the default one... "
811                whichone = "vishires" 
812        return whichlink[whichone]
813
814#def earthmap (whichone):
815#       if   whichone == "contrast":    whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg"
816#       elif whichone == "bw":          whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg"
817#       elif whichone == "nice":        whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg"
818#       return whichlink
819
820## Author: AS
821def latinterv (area="Whole"):
822    list =    { \
823        "Europe":                [[ 20., 80.],[- 50.,  50.]],\
824        "Central_America":       [[-10., 40.],[ 230., 300.]],\
825        "Africa":                [[-20., 50.],[- 50.,  50.]],\
826        "Whole":                 [[-90., 90.],[-180., 180.]],\
827        "Southern_Hemisphere":   [[-90., 60.],[-180., 180.]],\
828        "Northern_Hemisphere":   [[-60., 90.],[-180., 180.]],\
829        "Tharsis":               [[-30., 60.],[-170.,- 10.]],\
830        "Whole_No_High":         [[-60., 60.],[-180., 180.]],\
831        "Chryse":                [[-60., 60.],[- 60.,  60.]],\
832        "North_Pole":            [[ 50., 90.],[-180., 180.]],\
833        "Close_North_Pole":      [[ 75., 90.],[-180., 180.]],\
834        "Far_South_Pole":        [[-90.,-40.],[-180., 180.]],\
835        "South_Pole":            [[-90.,-50.],[-180., 180.]],\
836        "Close_South_Pole":      [[-90.,-75.],[-180., 180.]],\
837              }
838    if area not in list:   area = "Whole"
839    [olat,olon] = list[area]
840    return olon,olat
841
842## Author: TN
843def separatenames (name):
844  from numpy import concatenate
845  # look for comas in the input name to separate different names (files, variables,etc ..)
846  if name is None:
847     names = None
848  else:
849    names = []
850    stop = 0
851    currentname = name
852    while stop == 0:
853      indexvir = currentname.find(',')
854      if indexvir == -1:
855        stop = 1
856        name1 = currentname
857      else:
858        name1 = currentname[0:indexvir]
859      names = concatenate((names,[name1]))
860      currentname = currentname[indexvir+1:len(currentname)]
861  return names
862
863## Author: TN [old]
864def getopmatrix (kind,n):
865  import numpy as np
866  # compute matrix of operations between files
867  # the matrix is 'number of files'-square
868  # 1: difference (row minus column), 2: add
869  # not 0 in diag : just plot
870  if n == 1:
871    opm = np.eye(1)
872  elif kind == 'basic':
873    opm = np.eye(n)
874  elif kind == 'difference1': # show differences with 1st file
875    opm = np.zeros((n,n))
876    opm[0,:] = 1
877    opm[0,0] = 0
878  elif kind == 'difference2': # show differences with 1st file AND show 1st file
879    opm = np.zeros((n,n))
880    opm[0,:] = 1
881  else:
882    opm = np.eye(n)
883  return opm
884 
885## Author: TN [old]
886def checkcoherence (nfiles,nlat,nlon,ntime):
887  if (nfiles > 1):
888     if (nlat > 1):
889        errormess("what you asked is not possible !")
890  return 1
891
892## Author: TN
893def readslices(saxis):
894  from numpy import empty
895  if saxis == None:
896     zesaxis = None
897  else:
898     zesaxis = empty((len(saxis),2))
899     for i in range(len(saxis)):
900        a = separatenames(saxis[i])
901        if len(a) == 1:
902           zesaxis[i,:] = float(a[0])
903        else:
904           zesaxis[i,0] = float(a[0])
905           zesaxis[i,1] = float(a[1])
906           
907  return zesaxis
908
909## Author: AS
910def bidimfind(lon2d,lat2d,vlon,vlat):
911   import numpy as np
912   if vlat is None:    array = (lon2d - vlon)**2
913   elif vlon is None:  array = (lat2d - vlat)**2
914   else:               array = (lon2d - vlon)**2 + (lat2d - vlat)**2
915   idy,idx = np.unravel_index( np.argmin(array), lon2d.shape )
916   if vlon is not None:
917       #print lon2d[idy,idx],vlon
918       if (np.abs(lon2d[idy,idx]-vlon)) > 5: errormess("longitude not found ",printvar=lon2d)
919   if vlat is not None:
920       #print lat2d[idy,idx],vlat
921       if (np.abs(lat2d[idy,idx]-vlat)) > 5: errormess("latitude not found ",printvar=lat2d)
922   return idx,idy
923
924## Author: TN
925def getsindex(saxis,index,axis):
926# input  : all the desired slices and the good index
927# output : all indexes to be taken into account for reducing field
928  import numpy as np
929  if ( np.array(axis).ndim == 2):
930      axis = axis[:,0]
931  if saxis is None:
932      zeindex = None
933  else:
934      aaa = int(np.argmin(abs(saxis[index,0] - axis)))
935      bbb = int(np.argmin(abs(saxis[index,1] - axis)))
936      [imin,imax] = np.sort(np.array([aaa,bbb]))
937      zeindex = np.array(range(imax-imin+1))+imin
938      # because -180 and 180 are the same point in longitude,
939      # we get rid of one for averaging purposes.
940      if axis[imin] == -180 and axis[imax] == 180:
941         zeindex = zeindex[0:len(zeindex)-1]
942         print "INFO: whole longitude averaging asked, so last point is not taken into account."
943  return zeindex
944     
945## Author: TN
946def define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,indextime,what_I_plot,dim0,vertmode):
947# Purpose of define_axis is to find x and y axis scales in a smart way
948# x axis priority: 1/time 2/lon 3/lat 4/vertical
949# To be improved !!!...
950   from numpy import array,swapaxes
951   x = None
952   y = None
953   count = 0
954   what_I_plot = array(what_I_plot)
955   shape = what_I_plot.shape
956   if indextime is None:
957      print "AXIS is time"
958      x = time
959      count = count+1
960   if indexlon is None:
961      print "AXIS is lon"
962      if count == 0: x = lon
963      else: y = lon
964      count = count+1
965   if indexlat is None:
966      print "AXIS is lat"
967      if count == 0: x = lat
968      else: y = lat
969      count = count+1
970   if indexvert is None and dim0 is 4:
971      print "AXIS is vert"
972      if vertmode == 0: # vertical axis is as is (GCM grid)
973         if count == 0: x=range(len(vert))
974         else: y=range(len(vert))
975         count = count+1
976      else: # vertical axis is in kms
977         if count == 0: x = vert
978         else: y = vert
979         count = count+1
980   x = array(x)
981   y = array(y)
982   print "CHECK: what_I_plot.shape", what_I_plot.shape
983   print "CHECK: x.shape, y.shape", x.shape, y.shape
984   if len(shape) == 1:
985      if shape[0] != len(x):
986         print "WARNING HERE !!!"
987         x = y
988   elif len(shape) == 2:
989      if shape[1] == len(y) and shape[0] == len(x) and shape[0] != shape[1]:
990         what_I_plot = swapaxes(what_I_plot,0,1)
991         print "INFO: swapaxes", what_I_plot.shape, shape
992   return what_I_plot,x,y
993
994# Author: TN + AS
995def determineplot(slon, slat, svert, stime):
996    nlon = 1 # number of longitudinal slices -- 1 is None
997    nlat = 1
998    nvert = 1
999    ntime = 1
1000    nslices = 1
1001    if slon is not None:
1002        nslices = nslices*len(slon)
1003        nlon = len(slon)
1004    if slat is not None:
1005        nslices = nslices*len(slat)
1006        nlat = len(slat)
1007    if svert is not None:
1008        nslices = nslices*len(svert)
1009        nvert = len(svert)
1010    if stime is not None:
1011        nslices = nslices*len(stime)
1012        ntime = len(stime)
1013    #else:
1014    #    nslices = 2 
1015
1016    mapmode = 0
1017    if slon is None and slat is None:
1018       mapmode = 1 # in this case we plot a map, with the given projection
1019
1020    return nlon, nlat, nvert, ntime, mapmode, nslices
Note: See TracBrowser for help on using the repository browser.