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

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

PYTHON. 1/ mcs.py now works on a list of 3D and 4D variables given by -v or --var (it will also put aps and bps in the diagfi). 2/ added option --intas which is a template for the -i 2 gcm interpolation. -i 2 --intas mcs will interpolate the field on the same pressure levels as mcs. Works also with --intas tes. 3/ Corrected a small bug in make_gcm_netcdf
File size: 39.9 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## Author: AS
654def zoomset (wlon,wlat,zoom):
655    dlon = abs(wlon[1]-wlon[0])/2.
656    dlat = abs(wlat[1]-wlat[0])/2.
657    [wlon,wlat] = [ [wlon[0]+zoom*dlon/100.,wlon[1]-zoom*dlon/100.],\
658                    [wlat[0]+zoom*dlat/100.,wlat[1]-zoom*dlat/100.] ]
659    print "ZOOM %",zoom,wlon,wlat
660    return wlon,wlat
661
662## Author: AS
663def fmtvar (whichvar="def"):
664    fmtvar    =     { \
665             "TK":           "%.0f",\
666# Variables from TES ncdf format
667             "T_NADIR_DAY":  "%.0f",\
668             "T_NADIR_NIT":  "%.0f",\
669# Variables from tes.py ncdf format
670             "TEMP_DAY":     "%.0f",\
671             "TEMP_NIGHT":   "%.0f",\
672# Variables from MCS and mcs.py ncdf format
673             "DTEMP":     "%.0f",\
674             "NTEMP":   "%.0f",\
675             "TPOT":         "%.0f",\
676             "TSURF":        "%.0f",\
677             "def":          "%.1e",\
678             "PTOT":         "%.0f",\
679             "HGT":          "%.1e",\
680             "USTM":         "%.2f",\
681             "HFX":          "%.0f",\
682             "ICETOT":       "%.1e",\
683             "TAU_ICE":      "%.2f",\
684             "VMR_ICE":      "%.1e",\
685             "MTOT":         "%.1f",\
686             "ANOMALY":      "%.1f",\
687             "W":            "%.1f",\
688             "WMAX_TH":      "%.1f",\
689             "QSURFICE":     "%.0f",\
690             "UM":           "%.0f",\
691             "ALBBARE":      "%.2f",\
692             "TAU":          "%.1f",\
693             "CO2":          "%.2f",\
694             #### T.N.
695             "TEMP":         "%.0f",\
696             "VMR_H2OICE":   "%.0f",\
697             "VMR_H2OVAP":   "%.0f",\
698             "TAUTES":       "%.2f",\
699             "TAUTESAP":     "%.2f",\
700
701                    }
702    if whichvar not in fmtvar:
703        whichvar = "def"
704    return fmtvar[whichvar]
705
706## Author: AS
707####################################################################################################################
708### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
709def defcolorb (whichone="def"):
710    whichcolorb =    { \
711             "def":          "spectral",\
712             "HGT":          "spectral",\
713             "TK":           "gist_heat",\
714             "TSURF":        "RdBu_r",\
715             "QH2O":         "PuBu",\
716             "USTM":         "YlOrRd",\
717             #"T_nadir_nit":  "RdBu_r",\
718             #"T_nadir_day":  "RdBu_r",\
719             "HFX":          "RdYlBu",\
720             "ICETOT":       "YlGnBu_r",\
721             #"MTOT":         "PuBu",\
722             "CCNQ":         "YlOrBr",\
723             "CCNN":         "YlOrBr",\
724             "TEMP":         "Jet",\
725             "TAU_ICE":      "Blues",\
726             "VMR_ICE":      "Blues",\
727             "W":            "jet",\
728             "WMAX_TH":      "spectral",\
729             "ANOMALY":      "RdBu_r",\
730             "QSURFICE":     "hot_r",\
731             "ALBBARE":      "spectral",\
732             "TAU":          "YlOrBr_r",\
733             "CO2":          "YlOrBr_r",\
734             #### T.N.
735             "MTOT":         "Jet",\
736             "H2O_ICE_S":    "RdBu",\
737             "VMR_H2OICE":   "PuBu",\
738             "VMR_H2OVAP":   "PuBu",\
739                     }
740#W --> spectral ou jet
741#spectral BrBG RdBu_r
742    #print "predefined colorbars"
743    if whichone not in whichcolorb:
744        whichone = "def"
745    return whichcolorb[whichone]
746
747## Author: AS
748def definecolorvec (whichone="def"):
749        whichcolor =    { \
750                "def":          "black",\
751                "vis":          "yellow",\
752                "vishires":     "yellow",\
753                "molabw":       "yellow",\
754                "mola":         "black",\
755                "gist_heat":    "white",\
756                "hot":          "tk",\
757                "gist_rainbow": "black",\
758                "spectral":     "black",\
759                "gray":         "red",\
760                "PuBu":         "black",\
761                        }
762        if whichone not in whichcolor:
763                whichone = "def"
764        return whichcolor[whichone]
765
766## Author: AS
767def marsmap (whichone="vishires"):
768        from os import uname
769        mymachine = uname()[1]
770        ### not sure about speed-up with this method... looks the same
771        if "lmd.jussieu.fr" in mymachine: domain = "/u/aslmd/WWW/maps/"
772        else:                             domain = "http://www.lmd.jussieu.fr/~aslmd/maps/"
773        whichlink =     { \
774                #"vis":         "http://maps.jpl.nasa.gov/pix/mar0kuu2.jpg",\
775                #"vishires":    "http://www.lmd.jussieu.fr/~aslmd/maps/MarsMap_2500x1250.jpg",\
776                #"geolocal":    "http://dl.dropbox.com/u/11078310/geolocal.jpg",\
777                #"mola":        "http://www.lns.cornell.edu/~seb/celestia/mars-mola-2k.jpg",\
778                #"molabw":      "http://dl.dropbox.com/u/11078310/MarsElevation_2500x1250.jpg",\
779                "vis":         domain+"mar0kuu2.jpg",\
780                "vishires":    domain+"MarsMap_2500x1250.jpg",\
781                "geolocal":    domain+"geolocal.jpg",\
782                "mola":        domain+"mars-mola-2k.jpg",\
783                "molabw":      domain+"MarsElevation_2500x1250.jpg",\
784                "clouds":      "http://www.johnstonsarchive.net/spaceart/marswcloudmap.jpg",\
785                "jupiter":     "http://www.mmedia.is/~bjj/data/jupiter_css/jupiter_css.jpg",\
786                "jupiter_voy": "http://www.mmedia.is/~bjj/data/jupiter/jupiter_vgr2.jpg",\
787                "bw":          "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg",\
788                "contrast":    "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg",\
789                "nice":        "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg",\
790                "blue":        "http://eoimages.gsfc.nasa.gov/ve/2430/land_ocean_ice_2048.jpg",\
791                "blueclouds":  "http://eoimages.gsfc.nasa.gov/ve/2431/land_ocean_ice_cloud_2048.jpg",\
792                "justclouds":  "http://eoimages.gsfc.nasa.gov/ve/2432/cloud_combined_2048.jpg",\
793                        }
794        ### see http://www.mmedia.is/~bjj/planetary_maps.html
795        if whichone not in whichlink: 
796                print "marsmap: choice not defined... you'll get the default one... "
797                whichone = "vishires" 
798        return whichlink[whichone]
799
800#def earthmap (whichone):
801#       if   whichone == "contrast":    whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg"
802#       elif whichone == "bw":          whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg"
803#       elif whichone == "nice":        whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg"
804#       return whichlink
805
806## Author: AS
807def latinterv (area="Whole"):
808    list =    { \
809        "Europe":                [[ 20., 80.],[- 50.,  50.]],\
810        "Central_America":       [[-10., 40.],[ 230., 300.]],\
811        "Africa":                [[-20., 50.],[- 50.,  50.]],\
812        "Whole":                 [[-90., 90.],[-180., 180.]],\
813        "Southern_Hemisphere":   [[-90., 60.],[-180., 180.]],\
814        "Northern_Hemisphere":   [[-60., 90.],[-180., 180.]],\
815        "Tharsis":               [[-30., 60.],[-170.,- 10.]],\
816        "Whole_No_High":         [[-60., 60.],[-180., 180.]],\
817        "Chryse":                [[-60., 60.],[- 60.,  60.]],\
818        "North_Pole":            [[ 50., 90.],[-180., 180.]],\
819        "Close_North_Pole":      [[ 75., 90.],[-180., 180.]],\
820        "Far_South_Pole":        [[-90.,-40.],[-180., 180.]],\
821        "South_Pole":            [[-90.,-50.],[-180., 180.]],\
822        "Close_South_Pole":      [[-90.,-75.],[-180., 180.]],\
823              }
824    if area not in list:   area = "Whole"
825    [olat,olon] = list[area]
826    return olon,olat
827
828## Author: TN
829def separatenames (name):
830  from numpy import concatenate
831  # look for comas in the input name to separate different names (files, variables,etc ..)
832  if name is None:
833     names = None
834  else:
835    names = []
836    stop = 0
837    currentname = name
838    while stop == 0:
839      indexvir = currentname.find(',')
840      if indexvir == -1:
841        stop = 1
842        name1 = currentname
843      else:
844        name1 = currentname[0:indexvir]
845      names = concatenate((names,[name1]))
846      currentname = currentname[indexvir+1:len(currentname)]
847  return names
848
849## Author: TN [old]
850def getopmatrix (kind,n):
851  import numpy as np
852  # compute matrix of operations between files
853  # the matrix is 'number of files'-square
854  # 1: difference (row minus column), 2: add
855  # not 0 in diag : just plot
856  if n == 1:
857    opm = np.eye(1)
858  elif kind == 'basic':
859    opm = np.eye(n)
860  elif kind == 'difference1': # show differences with 1st file
861    opm = np.zeros((n,n))
862    opm[0,:] = 1
863    opm[0,0] = 0
864  elif kind == 'difference2': # show differences with 1st file AND show 1st file
865    opm = np.zeros((n,n))
866    opm[0,:] = 1
867  else:
868    opm = np.eye(n)
869  return opm
870 
871## Author: TN [old]
872def checkcoherence (nfiles,nlat,nlon,ntime):
873  if (nfiles > 1):
874     if (nlat > 1):
875        errormess("what you asked is not possible !")
876  return 1
877
878## Author: TN
879def readslices(saxis):
880  from numpy import empty
881  if saxis == None:
882     zesaxis = None
883  else:
884     zesaxis = empty((len(saxis),2))
885     for i in range(len(saxis)):
886        a = separatenames(saxis[i])
887        if len(a) == 1:
888           zesaxis[i,:] = float(a[0])
889        else:
890           zesaxis[i,0] = float(a[0])
891           zesaxis[i,1] = float(a[1])
892           
893  return zesaxis
894
895## Author: AS
896def bidimfind(lon2d,lat2d,vlon,vlat):
897   import numpy as np
898   if vlat is None:    array = (lon2d - vlon)**2
899   elif vlon is None:  array = (lat2d - vlat)**2
900   else:               array = (lon2d - vlon)**2 + (lat2d - vlat)**2
901   idy,idx = np.unravel_index( np.argmin(array), lon2d.shape )
902   if vlon is not None:
903       #print lon2d[idy,idx],vlon
904       if (np.abs(lon2d[idy,idx]-vlon)) > 5: errormess("longitude not found ",printvar=lon2d)
905   if vlat is not None:
906       #print lat2d[idy,idx],vlat
907       if (np.abs(lat2d[idy,idx]-vlat)) > 5: errormess("latitude not found ",printvar=lat2d)
908   return idx,idy
909
910## Author: TN
911def getsindex(saxis,index,axis):
912# input  : all the desired slices and the good index
913# output : all indexes to be taken into account for reducing field
914  import numpy as np
915  if ( np.array(axis).ndim == 2):
916      axis = axis[:,0]
917  if saxis is None:
918      zeindex = None
919  else:
920      aaa = int(np.argmin(abs(saxis[index,0] - axis)))
921      bbb = int(np.argmin(abs(saxis[index,1] - axis)))
922      [imin,imax] = np.sort(np.array([aaa,bbb]))
923      zeindex = np.array(range(imax-imin+1))+imin
924      # because -180 and 180 are the same point in longitude,
925      # we get rid of one for averaging purposes.
926      if axis[imin] == -180 and axis[imax] == 180:
927         zeindex = zeindex[0:len(zeindex)-1]
928         print "INFO: whole longitude averaging asked, so last point is not taken into account."
929  return zeindex
930     
931## Author: TN
932def define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,indextime,what_I_plot,dim0,vertmode):
933# Purpose of define_axis is to find x and y axis scales in a smart way
934# x axis priority: 1/time 2/lon 3/lat 4/vertical
935# To be improved !!!...
936   from numpy import array,swapaxes
937   x = None
938   y = None
939   count = 0
940   what_I_plot = array(what_I_plot)
941   shape = what_I_plot.shape
942   if indextime is None:
943      print "AXIS is time"
944      x = time
945      count = count+1
946   if indexlon is None:
947      print "AXIS is lon"
948      if count == 0: x = lon
949      else: y = lon
950      count = count+1
951   if indexlat is None:
952      print "AXIS is lat"
953      if count == 0: x = lat
954      else: y = lat
955      count = count+1
956   if indexvert is None and dim0 is 4:
957      print "AXIS is vert"
958      if vertmode == 0: # vertical axis is as is (GCM grid)
959         if count == 0: x=range(len(vert))
960         else: y=range(len(vert))
961         count = count+1
962      else: # vertical axis is in kms
963         if count == 0: x = vert
964         else: y = vert
965         count = count+1
966   x = array(x)
967   y = array(y)
968   print "CHECK: what_I_plot.shape", what_I_plot.shape
969   print "CHECK: x.shape, y.shape", x.shape, y.shape
970   if len(shape) == 1:
971      if shape[0] != len(x):
972         print "WARNING HERE !!!"
973         x = y
974   elif len(shape) == 2:
975      if shape[1] == len(y) and shape[0] == len(x) and shape[0] != shape[1]:
976         what_I_plot = swapaxes(what_I_plot,0,1)
977         print "INFO: swapaxes", what_I_plot.shape, shape
978   return what_I_plot,x,y
979
980# Author: TN + AS
981def determineplot(slon, slat, svert, stime):
982    nlon = 1 # number of longitudinal slices -- 1 is None
983    nlat = 1
984    nvert = 1
985    ntime = 1
986    nslices = 1
987    if slon is not None:
988        nslices = nslices*len(slon)
989        nlon = len(slon)
990    if slat is not None:
991        nslices = nslices*len(slat)
992        nlat = len(slat)
993    if svert is not None:
994        nslices = nslices*len(svert)
995        nvert = len(svert)
996    if stime is not None:
997        nslices = nslices*len(stime)
998        ntime = len(stime)
999    #else:
1000    #    nslices = 2 
1001
1002    mapmode = 0
1003    if slon is None and slat is None:
1004       mapmode = 1 # in this case we plot a map, with the given projection
1005
1006    return nlon, nlat, nvert, ntime, mapmode, nslices
Note: See TracBrowser for help on using the repository browser.