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

Last change on this file since 426 was 426, checked in by aslmd, 13 years ago

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