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

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

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