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

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

UTIL PYTHON: added maplatlon function, a user-friendly way to map fields and winds. a few cosmetic changes.

File size: 55.2 KB
RevLine 
[345]1## Author: AS
[252]2def errormess(text,printvar=None):
[233]3    print text
[399]4    if printvar is not None: print printvar
[233]5    exit()
6    return
7
[345]8## Author: AS
[349]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
[468]22def getname(var=False,var2=False,winds=False,anomaly=False):
[252]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
[468]28    if var2:              basename = basename + '_' + var2
[252]29    return basename
30
[345]31## Author: AS
[252]32def localtime(utc,lon):
33    ltst = utc + lon / 15.
34    ltst = int (ltst * 10) / 10.
35    ltst = ltst % 24
36    return ltst
37
[569]38## Author: AC
39def check_localtime(time):
40    a=-1
41    for i in range(len(time)-1):
[583]42       if (time[i] > time[i+1]): a=i
43    if a >= 0 and a < (len(time)-1)/2.:
[569]44       print "Sorry, time axis is not regular."
45       print "Contourf needs regular axis... recasting"
46       for i in range(a+1):
47          time[i]=time[i]-24.
[583]48    if a >= 0 and a >= (len(time)-1)/2.:
49       print "Sorry, time axis is not regular."
50       print "Contourf needs regular axis... recasting"
51       for i in range((len(time)-1) - a):
52          time[a+1+i]=time[a+1+i]+24.
[569]53    return time
54
[525]55## Author: AS, AC, JL
[233]56def whatkindfile (nc):
[429]57    if 'controle' in nc.variables:             typefile = 'gcm'
58    elif 'phisinit' in nc.variables:           typefile = 'gcm'
[525]59    elif 'time_counter' in nc.variables:       typefile = 'earthgcm'
[548]60    elif hasattr(nc,'START_DATE'):             typefile = 'meso' 
[429]61    elif 'HGT_M' in nc.variables:              typefile = 'geo'
[558]62    elif hasattr(nc,'institution'):
63      if "European Centre" in getattr(nc,'institution'):  typefile = 'ecmwf'
[429]64    else:                                      typefile = 'gcm' # for lslin-ed files from gcm
[233]65    return typefile
66
[345]67## Author: AS
[233]68def getfield (nc,var):
69    ## this allows to get much faster (than simply referring to nc.variables[var])
[395]70    import numpy as np
[233]71    dimension = len(nc.variables[var].dimensions)
[392]72    #print "   Opening variable with", dimension, "dimensions ..."
[233]73    if dimension == 2:    field = nc.variables[var][:,:]
74    elif dimension == 3:  field = nc.variables[var][:,:,:]
75    elif dimension == 4:  field = nc.variables[var][:,:,:,:]
[494]76    elif dimension == 1:  field = nc.variables[var][:]
[395]77    # if there are NaNs in the ncdf, they should be loaded as a masked array which will be
78    # recasted as a regular array later in reducefield
79    if (np.isnan(np.sum(field)) and (type(field).__name__ not in 'MaskedArray')):
80       print "Warning: netcdf as nan values but is not loaded as a Masked Array."
81       print "recasting array type"
82       out=np.ma.masked_invalid(field)
83       out.set_fill_value([np.NaN])
84    else:
[464]85    # missing values from zrecast or hrecast are -1e-33
[469]86       masked=np.ma.masked_where(field < -1e30,field)
[578]87       masked2=np.ma.masked_where(field > 1e35,field)
88       masked.set_fill_value([np.NaN]) ; masked2.set_fill_value([np.NaN])
89       mask = np.ma.getmask(masked) ; mask2 = np.ma.getmask(masked2)
90       if (True in np.array(mask)):
91          out=masked
92          print "Masked array... Missing value is NaN"
93       elif (True in np.array(mask2)):
94          out=masked2
95          print "Masked array... Missing value is NaN"
96#       else:
97#       # missing values from api are 1e36
98#          masked=np.ma.masked_where(field > 1e35,field)
99#          masked.set_fill_value([np.NaN])
100#          mask = np.ma.getmask(masked)
101#          if (True in np.array(mask)):out=masked
102#          else:out=field
103       else:out=field
[395]104    return out
[233]105
[571]106## Author: AC
[612]107# Compute the norm of the winds
108# The corresponding variable to call is UV or uvmet (to use api)
[572]109def windamplitude (nc):
[581]110    import numpy as np
[571]111    varinfile = nc.variables.keys()
112    if "U" in varinfile: zu=getfield(nc,'U')
113    elif "Um" in varinfile: zu=getfield(nc,'Um')
114    if "V" in varinfile: zv=getfield(nc,'V')
115    elif "Vm" in varinfile: zv=getfield(nc,'Vm')
116    znt,znz,zny,znx = np.array(zu).shape
117    if "U" in varinfile:znx=znx-1
118    zuint = np.zeros([znt,znz,zny,znx])
119    zvint = np.zeros([znt,znz,zny,znx])
120    if "U" in varinfile:
121       for xx in np.arange(znx):      zuint[:,:,:,xx] = (zu[:,:,:,xx] + zu[:,:,:,xx+1])/2.
122       for yy in np.arange(zny):      zvint[:,:,yy,:] = (zv[:,:,yy,:] + zv[:,:,yy+1,:])/2.
123    else:
124       zuint=zu
125       zvint=zv
126    return np.sqrt(zuint**2 + zvint**2)
127
[612]128## Author: AC
129# Compute the norm of the slope angles
130# The corresponding variable to call is SLOPEXY
131def slopeamplitude (nc):
132    import numpy as np
133    varinfile = nc.variables.keys()
134    if "slopex" in varinfile: zu=getfield(nc,'slopex')
135    elif "SLOPEX" in varinfile: zu=getfield(nc,'SLOPEX')
136    if "slopey" in varinfile: zv=getfield(nc,'slopey')
137    elif "SLOPEY" in varinfile: zv=getfield(nc,'SLOPEY')
138    znt,zny,znx = np.array(zu).shape
139    zuint = np.zeros([znt,zny,znx])
140    zvint = np.zeros([znt,zny,znx])
141    zuint=zu
142    zvint=zv
143    return np.sqrt(zuint**2 + zvint**2)
144
145## Author: AC
146# Compute the temperature difference between surface and first level.
147# API is automatically called to get TSURF and TK.
148# The corresponding variable to call is DELTAT
149def deltat0t1 (nc):
150    import numpy as np
151    varinfile = nc.variables.keys()
152    if "tsurf" in varinfile: zu=getfield(nc,'tsurf')
153    elif "TSURF" in varinfile: zu=getfield(nc,'TSURF')
154    if "tk" in varinfile: zv=getfield(nc,'tk')
155    elif "TK" in varinfile: zv=getfield(nc,'TK')
156    znt,zny,znx = np.array(zu).shape
157    zuint = np.zeros([znt,zny,znx])
158    zuint=zu - zv[:,0,:,:]
159    return zuint
160
[382]161## Author: AS + TN + AC
[525]162def reducefield (input,d4=None,d3=None,d2=None,d1=None,yint=False,alt=None,anomaly=False,redope=None,mesharea=None):
[252]163    ### we do it the reverse way to be compliant with netcdf "t z y x" or "t y x" or "y x"
[233]164    ### it would be actually better to name d4 d3 d2 d1 as t z y x
[405]165    ### ... note, anomaly is only computed over d1 and d2 for the moment
[233]166    import numpy as np
[525]167    from mymath import max,mean,min,sum
[422]168    csmooth = 12 ## a fair amount of grid points (too high results in high computation time)
[483]169    if redope is not None:
170       if   redope == "mint":     input = min(input,axis=0) ; d1 = None
171       elif redope == "maxt":     input = max(input,axis=0) ; d1 = None
172       else:                      errormess("not supported. but try lines in reducefield beforehand.")
173       #elif redope == "minz":     input = min(input,axis=1) ; d2 = None
174       #elif redope == "maxz":     input = max(input,axis=1) ; d2 = None
175       #elif redope == "miny":     input = min(input,axis=2) ; d3 = None
176       #elif redope == "maxy":     input = max(input,axis=2) ; d3 = None
177       #elif redope == "minx":     input = min(input,axis=3) ; d4 = None
178       #elif redope == "maxx":     input = max(input,axis=3) ; d4 = None
[233]179    dimension = np.array(input).ndim
[525]180    shape = np.array(np.array(input).shape)
[349]181    #print 'd1,d2,d3,d4: ',d1,d2,d3,d4
[405]182    if anomaly: print 'ANOMALY ANOMALY'
[233]183    output = input
184    error = False
[350]185    #### this is needed to cope the case where d4,d3,d2,d1 are single integers and not arrays
[345]186    if d4 is not None and not isinstance(d4, np.ndarray): d4=[d4]
187    if d3 is not None and not isinstance(d3, np.ndarray): d3=[d3]
188    if d2 is not None and not isinstance(d2, np.ndarray): d2=[d2]
189    if d1 is not None and not isinstance(d1, np.ndarray): d1=[d1]
190    ### now the main part
[233]191    if dimension == 2:
[525]192        if mesharea is None: mesharea=np.ones(shape)
193        if   max(d2) >= shape[0]: error = True
194        elif max(d1) >= shape[1]: error = True
195        elif d1 is not None and d2 is not None:
196          totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
197          output = output*mesharea; output = sum(output[d2,:],axis=0); output = sum(output[d1],axis=0)/totalarea
[350]198        elif d1 is not None:         output = mean(input[:,d1],axis=1)
[525]199        elif d2 is not None:         totalarea=sum(mesharea[d2,:],axis=0); output = output*mesharea; output = sum(output[d2,:],axis=0)/totalarea
[233]200    elif dimension == 3:
[525]201        if mesharea is None: mesharea=np.ones(shape[[1,2]])
[345]202        if   max(d4) >= shape[0]: error = True
203        elif max(d2) >= shape[1]: error = True
204        elif max(d1) >= shape[2]: error = True
[350]205        elif d4 is not None and d2 is not None and d1 is not None: 
[525]206          totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
207          output = mean(input[d4,:,:],axis=0)
208          output = output*mesharea; output = sum(output[d2,:],axis=0); output = sum(output[d1],axis=0)/totalarea
209        elif d4 is not None and d2 is not None:
210          output = mean(input[d4,:,:],axis=0)
211          totalarea=sum(mesharea[d2,:],axis=0); output = output*mesharea; output = sum(output[d2,:],axis=0)/totalarea
[349]212        elif d4 is not None and d1 is not None:    output = mean(input[d4,:,:],axis=0); output=mean(output[:,d1],axis=1)
[525]213        elif d2 is not None and d1 is not None:
214          totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
215          output = output*mesharea; output = sum(output[:,d2,:],axis=1); output = sum(output[:,d1],axis=1)/totalarea
216        elif d1 is not None:   output = mean(input[:,:,d1],axis=2)
217        elif d2 is not None:   totalarea=sum(mesharea[d2,:],axis=0); output = output*mesharea; output = sum(output[:,d2,:],axis=1)/totalarea
218        elif d4 is not None:   output = mean(input[d4,:,:],axis=0)
[233]219    elif dimension == 4:
[525]220        if mesharea is None: mesharea=np.ones(shape[[2,3]])
[345]221        if   max(d4) >= shape[0]: error = True
222        elif max(d3) >= shape[1]: error = True
223        elif max(d2) >= shape[2]: error = True
224        elif max(d1) >= shape[3]: error = True
[382]225        elif d4 is not None and d3 is not None and d2 is not None and d1 is not None:
[392]226             output = mean(input[d4,:,:,:],axis=0)
227             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
[427]228             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.)
[525]229             totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
230             output = output*mesharea
231             output = sum(output[d2,:],axis=0)
232             output = sum(output[d1],axis=0)/totalarea
[350]233        elif d4 is not None and d3 is not None and d2 is not None: 
[392]234             output = mean(input[d4,:,:,:],axis=0)
235             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
[525]236             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.)
237             totalarea=sum(mesharea[d2,:],axis=0)
238             output = output*mesharea; output = sum(output[d2,:],axis=0)/totalarea
[350]239        elif d4 is not None and d3 is not None and d1 is not None: 
[392]240             output = mean(input[d4,:,:,:],axis=0)
241             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
[405]242             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
[392]243             output = mean(output[:,d1],axis=1)
[350]244        elif d4 is not None and d2 is not None and d1 is not None: 
[392]245             output = mean(input[d4,:,:,:],axis=0)
[405]246             if anomaly:
247                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
[525]248             totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
249             output = output*mesharea; output = sum(output[:,d2,:],axis=1); output = sum(output[:,d1],axis=1)/totalarea
[405]250             #noperturb = smooth1d(output,window_len=7)
251             #lenlen = len(output) ; output = output[1:lenlen-7] ; yeye = noperturb[4:lenlen-4]
252             #plot(output) ; plot(yeye) ; show() ; plot(output-yeye) ; show()
[350]253        elif d3 is not None and d2 is not None and d1 is not None: 
[392]254             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3) 
[405]255             if anomaly:
256                 for k in range(output.shape[0]):  output[k,:,:] = 100. * ((output[k,:,:] / smooth(output[k,:,:],csmooth)) - 1.)
[525]257             totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
258             output = output*mesharea; output = sum(output[:,d2,:],axis=1); output = sum(output[:,d1],axis=1)/totalarea
[392]259        elif d4 is not None and d3 is not None: 
260             output = mean(input[d4,:,:,:],axis=0) 
261             output = reduce_zaxis(output[d3,:,:],ax=0,yint=yint,vert=alt,indice=d3)
[405]262             if anomaly: output = 100. * ((output / smooth(output,csmooth)) - 1.) 
[392]263        elif d4 is not None and d2 is not None: 
264             output = mean(input[d4,:,:,:],axis=0) 
[525]265             totalarea=sum(mesharea[d2,:],axis=0); output = output*mesharea; output = sum(output[:,d2,:],axis=1)/totalarea
[392]266        elif d4 is not None and d1 is not None: 
267             output = mean(input[d4,:,:,:],axis=0) 
268             output = mean(output[:,:,d1],axis=2)
269        elif d3 is not None and d2 is not None: 
270             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
[525]271             totalarea=sum(mesharea[d2,:],axis=0); output = output*mesharea; output = sum(output[:,d2,:],axis=1)/totalarea
[392]272        elif d3 is not None and d1 is not None: 
273             output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
[448]274             output = mean(output[:,:,d1],axis=2)
[392]275        elif d2 is not None and d1 is not None: 
[525]276             totalarea=sum(mesharea[d2,:],axis=0); totalarea=sum(totalarea[d1],axis=0)
277             output = output*mesharea; output = sum(output[:,:,d2,:],axis=2); output = sum(output[:,:,d1],axis=2)/totalarea
[392]278        elif d1 is not None:        output = mean(input[:,:,:,d1],axis=3)
[525]279        elif d2 is not None: 
280             totalarea=sum(mesharea[d2,:],axis=0); output = output*mesharea; output = sum(output[:,:,d2,:],axis=2)/totalarea
[437]281        elif d3 is not None:        output = reduce_zaxis(input[:,d3,:,:],ax=1,yint=yint,vert=alt,indice=d3)
[392]282        elif d4 is not None:        output = mean(input[d4,:,:,:],axis=0)
[468]283    dimension2 = np.array(output).ndim
284    shape2 = np.array(output).shape
285    print 'REDUCEFIELD dim,shape: ',dimension,shape,' >>> ',dimension2,shape2
[233]286    return output, error
287
[392]288## Author: AC + AS
289def reduce_zaxis (input,ax=None,yint=False,vert=None,indice=None):
[382]290    from mymath import max,mean
291    from scipy import integrate
[637]292    import numpy as np
[392]293    if yint and vert is not None and indice is not None:
[391]294       if type(input).__name__=='MaskedArray':
295          input.set_fill_value([np.NaN])
[392]296          output = integrate.trapz(input.filled(),x=vert[indice],axis=ax)
[391]297       else:
[396]298          output = integrate.trapz(input,x=vert[indice],axis=ax)
[382]299    else:
300       output = mean(input,axis=ax)
301    return output
302
[345]303## Author: AS + TN
[233]304def definesubplot ( numplot, fig ):
305    from matplotlib.pyplot import rcParams
306    rcParams['font.size'] = 12. ## default (important for multiple calls)
[345]307    if numplot <= 0:
308        subv = 99999
309        subh = 99999
310    elif numplot == 1: 
[568]311        subv = 1
312        subh = 1
[233]313    elif numplot == 2:
[483]314        subv = 1 #2
315        subh = 2 #1
[233]316        fig.subplots_adjust(wspace = 0.35)
317        rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
318    elif numplot == 3:
[453]319        subv = 3
320        subh = 1
[613]321        fig.subplots_adjust(hspace = 0.5)
[233]322        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
[637]323    elif numplot == 4:
[345]324        subv = 2
325        subh = 2
[610]326        fig.subplots_adjust(wspace = 0.4, hspace = 0.3)
[345]327        rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
328    elif numplot <= 6:
329        subv = 2
330        subh = 3
[638]331        #fig.subplots_adjust(wspace = 0.4, hspace = 0.0)
332        fig.subplots_adjust(wspace = 0.5, hspace = 0.3)
[233]333        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
[345]334    elif numplot <= 8:
335        subv = 2
336        subh = 4
[233]337        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
338        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
[345]339    elif numplot <= 9:
340        subv = 3
341        subh = 3
[233]342        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
343        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
[345]344    elif numplot <= 12:
345        subv = 3
346        subh = 4
[453]347        fig.subplots_adjust(wspace = 0, hspace = 0.1)
[345]348        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
349    elif numplot <= 16:
350        subv = 4
351        subh = 4
352        fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
353        rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
[233]354    else:
[345]355        print "number of plot supported: 1 to 16"
[233]356        exit()
[345]357    return subv,subh
[233]358
[345]359## Author: AS
[233]360def getstralt(nc,nvert):
[548]361    varinfile = nc.variables.keys()
362    if 'vert' not in varinfile:
[233]363        stralt = "_lvl" + str(nvert)
[548]364    else:
[233]365        zelevel = int(nc.variables['vert'][nvert])
366        if abs(zelevel) < 10000.:   strheight=str(zelevel)+"m"
367        else:                       strheight=str(int(zelevel/1000.))+"km"
368        if 'altitude'       in nc.dimensions:   stralt = "_"+strheight+"-AMR"
369        elif 'altitude_abg' in nc.dimensions:   stralt = "_"+strheight+"-ALS"
370        elif 'bottom_top'   in nc.dimensions:   stralt = "_"+strheight
371        elif 'pressure'     in nc.dimensions:   stralt = "_"+str(zelevel)+"Pa"
372        else:                                   stralt = ""
373    return stralt
374
[345]375## Author: AS
[468]376def getlschar ( namefile, getaxis=False ):
[195]377    from netCDF4 import Dataset
378    from timestuff import sol2ls
[233]379    from numpy import array
[400]380    from string import rstrip
[405]381    namefile = rstrip( rstrip( rstrip( namefile, chars="_z"), chars="_zabg"), chars="_p")
[400]382    #### we assume that wrfout is next to wrfout_z and wrfout_zabg
[195]383    nc  = Dataset(namefile)
[237]384    zetime = None
[489]385    days_in_month = [61, 66, 66, 65, 60, 54, 50, 46, 47, 47, 51, 56]
386    plus_in_month = [ 0, 61,127,193,258,318,372,422,468,515,562,613]
[400]387    if 'Times' in nc.variables:
[233]388        zetime = nc.variables['Times'][0]
389        shape = array(nc.variables['Times']).shape
390        if shape[0] < 2: zetime = None
391    if zetime is not None \
[225]392       and 'vert' not in nc.variables:
[489]393        ##### strangely enough this does not work for api or ncrcat results!
394        zesol = plus_in_month[ int(zetime[5]+zetime[6])-1 ] + int(zetime[8]+zetime[9]) - 1 ##les sols GCM commencent a 0
395        dals = int( 10. * sol2ls ( zesol ) ) / 10.
[197]396        ###
397        zetime2 = nc.variables['Times'][1]
398        one  = int(zetime[11]+zetime[12]) + int(zetime[14]+zetime[15])/37.
399        next = int(zetime2[11]+zetime2[12]) + int(zetime2[14]+zetime2[15])/37. 
400        zehour    = one
401        zehourin  = abs ( next - one )
[489]402        if not getaxis:
403            lschar = "_Ls"+str(dals)
404        else:
[468]405            zelen = len(nc.variables['Times'][:])
406            yeye = range(zelen) ; lsaxis = range(zelen) ; solaxis = range(zelen) ; ltaxis = range(zelen)
407            for iii in yeye:
[489]408               zetime = nc.variables['Times'][iii] 
[468]409               ltaxis[iii] = int(zetime[11]+zetime[12]) + int(zetime[14]+zetime[15])/37.
[489]410               solaxis[iii] = ltaxis[iii] / 24. + plus_in_month[ int(zetime[5]+zetime[6])-1 ] + int(zetime[8]+zetime[9]) - 1 ##les sols GCM commencent a 0
[468]411               lsaxis[iii] = sol2ls ( solaxis[iii] )
412               if ltaxis[iii] < ltaxis[iii-1]: ltaxis[iii] = ltaxis[iii] + 24.
[489]413               #print ltaxis[iii], solaxis[iii], lsaxis[iii], getattr( nc, 'JULDAY' )
[468]414            lschar = lsaxis ; zehour = solaxis ; zehourin = ltaxis
[195]415    else:
416        lschar=""
[197]417        zehour = 0
418        zehourin = 1 
419    return lschar, zehour, zehourin
[195]420
[345]421## Author: AS
[202]422def getprefix (nc):
423    prefix = 'LMD_MMM_'
424    prefix = prefix + 'd'+str(getattr(nc,'GRID_ID'))+'_'
425    prefix = prefix + str(int(getattr(nc,'DX')/1000.))+'km_'
426    return prefix
427
[345]428## Author: AS
[184]429def getproj (nc):
[233]430    typefile = whatkindfile(nc)
[548]431    if typefile in ['meso','geo']:
[233]432        ### (il faudrait passer CEN_LON dans la projection ?)
433        map_proj = getattr(nc, 'MAP_PROJ')
434        cen_lat  = getattr(nc, 'CEN_LAT')
435        if map_proj == 2:
436            if cen_lat > 10.:   
437                proj="npstere"
[392]438                #print "NP stereographic polar domain"
[233]439            else:           
440                proj="spstere"
[392]441                #print "SP stereographic polar domain"
[233]442        elif map_proj == 1: 
[392]443            #print "lambert projection domain"
[233]444            proj="lcc"
445        elif map_proj == 3: 
[392]446            #print "mercator projection"
[233]447            proj="merc"
448        else:
449            proj="merc"
[548]450    elif typefile in ['gcm']:        proj="cyl"    ## pb avec les autres (de trace derriere la sphere ?)
451    else:                            proj="ortho"
[184]452    return proj   
453
[345]454## Author: AS
[180]455def ptitle (name):
456    from matplotlib.pyplot import title
457    title(name)
458    print name
459
[345]460## Author: AS
[252]461def polarinterv (lon2d,lat2d):
462    import numpy as np
463    wlon = [np.min(lon2d),np.max(lon2d)]
464    ind = np.array(lat2d).shape[0] / 2  ## to get a good boundlat and to get the pole
465    wlat = [np.min(lat2d[ind,:]),np.max(lat2d[ind,:])]
466    return [wlon,wlat]
467
[345]468## Author: AS
[180]469def simplinterv (lon2d,lat2d):
470    import numpy as np
471    return [[np.min(lon2d),np.max(lon2d)],[np.min(lat2d),np.max(lat2d)]]
472
[345]473## Author: AS
[184]474def wrfinterv (lon2d,lat2d):
475    nx = len(lon2d[0,:])-1
476    ny = len(lon2d[:,0])-1
[225]477    lon1 = lon2d[0,0] 
478    lon2 = lon2d[nx,ny] 
479    lat1 = lat2d[0,0] 
480    lat2 = lat2d[nx,ny] 
[233]481    if abs(0.5*(lat1+lat2)) > 60.:  wider = 0.5 * (abs(lon1)+abs(lon2)) * 0.1
482    else:                           wider = 0.
483    if lon1 < lon2:  wlon = [lon1, lon2 + wider] 
[225]484    else:            wlon = [lon2, lon1 + wider]
485    if lat1 < lat2:  wlat = [lat1, lat2]
486    else:            wlat = [lat2, lat1]
487    return [wlon,wlat]
[184]488
[345]489## Author: AS
[240]490def makeplotres (filename,res=None,pad_inches_value=0.25,folder='',disp=True,ext='png',erase=False):
[180]491    import  matplotlib.pyplot as plt
[240]492    from os import system
493    addstr = ""
494    if res is not None:
495        res = int(res)
496        addstr = "_"+str(res)
497    name = filename+addstr+"."+ext
[186]498    if folder != '':      name = folder+'/'+name
[180]499    plt.savefig(name,dpi=res,bbox_inches='tight',pad_inches=pad_inches_value)
[240]500    if disp:              display(name)
501    if ext in ['eps','ps','svg']:   system("tar czvf "+name+".tar.gz "+name+" ; rm -f "+name)
502    if erase:   system("mv "+name+" to_be_erased")             
[180]503    return
504
[430]505## Author: AS + AC
[451]506def dumpbdy (field,n,stag=None,condition=False,onlyx=False,onlyy=False):
[447]507    nx = len(field[0,:])-1
508    ny = len(field[:,0])-1
[444]509    if condition:
510      if stag == 'U': nx = nx-1
511      if stag == 'V': ny = ny-1
512      if stag == 'W': nx = nx+1 #special les case when we dump stag on W
[451]513    if onlyx:    result = field[:,n:nx-n]
514    elif onlyy:  result = field[n:ny-n,:]
515    else:        result = field[n:ny-n,n:nx-n]
516    return result
[180]517
[444]518## Author: AS + AC
[233]519def getcoorddef ( nc ):   
[317]520    import numpy as np
[233]521    ## getcoord2d for predefined types
522    typefile = whatkindfile(nc)
[548]523    if typefile in ['meso']:
524        if '9999' not in getattr(nc,'START_DATE') :   
525            ## regular mesoscale 
526            [lon2d,lat2d] = getcoord2d(nc) 
527        else:                     
528            ## idealized mesoscale                     
529            nx=getattr(nc,'WEST-EAST_GRID_DIMENSION')
530            ny=getattr(nc,'SOUTH-NORTH_GRID_DIMENSION')
531            dlat=getattr(nc,'DX')
532            ## this is dirty because Martian-specific
533            # ... but this just intended to get "lat-lon" like info
534            falselon = np.arange(-dlat*(nx-1)/2.,dlat*(nx-1)/2.,dlat)/60000.
535            falselat = np.arange(-dlat*(ny-1)/2.,dlat*(ny-1)/2.,dlat)/60000.
536            [lon2d,lat2d] = np.meshgrid(falselon,falselat) ## dummy coordinates
537            print "WARNING: domain plot artificially centered on lat,lon 0,0"
[637]538    elif typefile in ['gcm','earthgcm','ecmwf']: 
539        if "longitude" in nc.dimensions:  dalon = "longitude"
540        elif "lon" in nc.dimensions:      dalon = "lon"
541        if "latitude" in nc.dimensions:   dalat = "latitude"
542        elif "lat" in nc.dimensions:      dalat = "lat"
543        [lon2d,lat2d] = getcoord2d(nc,nlat=dalat,nlon=dalon,is1d=True)
[233]544    elif typefile in ['geo']:
545        [lon2d,lat2d] = getcoord2d(nc,nlat='XLAT_M',nlon='XLONG_M')
546    return lon2d,lat2d   
547
[345]548## Author: AS
[184]549def getcoord2d (nc,nlat='XLAT',nlon='XLONG',is1d=False):
550    import numpy as np
551    if is1d:
552        lat = nc.variables[nlat][:]
553        lon = nc.variables[nlon][:]
554        [lon2d,lat2d] = np.meshgrid(lon,lat)
555    else:
556        lat = nc.variables[nlat][0,:,:]
557        lon = nc.variables[nlon][0,:,:]
558        [lon2d,lat2d] = [lon,lat]
559    return lon2d,lat2d
560
[405]561## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
562def smooth1d(x,window_len=11,window='hanning'):
563    import numpy
564    """smooth the data using a window with requested size.
565    This method is based on the convolution of a scaled window with the signal.
566    The signal is prepared by introducing reflected copies of the signal
567    (with the window size) in both ends so that transient parts are minimized
568    in the begining and end part of the output signal.
569    input:
570        x: the input signal
571        window_len: the dimension of the smoothing window; should be an odd integer
572        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
573            flat window will produce a moving average smoothing.
574    output:
575        the smoothed signal
576    example:
577    t=linspace(-2,2,0.1)
578    x=sin(t)+randn(len(t))*0.1
579    y=smooth(x)
580    see also:
581    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
582    scipy.signal.lfilter
583    TODO: the window parameter could be the window itself if an array instead of a string   
584    """
585    x = numpy.array(x)
586    if x.ndim != 1:
587        raise ValueError, "smooth only accepts 1 dimension arrays."
588    if x.size < window_len:
589        raise ValueError, "Input vector needs to be bigger than window size."
590    if window_len<3:
591        return x
592    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
593        raise ValueError, "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"
594    s=numpy.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
595    #print(len(s))
596    if window == 'flat': #moving average
597        w=numpy.ones(window_len,'d')
598    else:
599        w=eval('numpy.'+window+'(window_len)')
600    y=numpy.convolve(w/w.sum(),s,mode='valid')
601    return y
602
[345]603## Author: AS
[180]604def smooth (field, coeff):
605        ## actually blur_image could work with different coeff on x and y
606        if coeff > 1:   result = blur_image(field,int(coeff))
607        else:           result = field
608        return result
609
[345]610## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
[180]611def gauss_kern(size, sizey=None):
612        import numpy as np
613        # Returns a normalized 2D gauss kernel array for convolutions
614        size = int(size)
615        if not sizey:
616                sizey = size
617        else:
618                sizey = int(sizey)
619        x, y = np.mgrid[-size:size+1, -sizey:sizey+1]
620        g = np.exp(-(x**2/float(size)+y**2/float(sizey)))
621        return g / g.sum()
622
[345]623## FROM COOKBOOK http://www.scipy.org/Cookbook/SignalSmooth
[180]624def blur_image(im, n, ny=None) :
625        from scipy.signal import convolve
626        # blurs the image by convolving with a gaussian kernel of typical size n.
627        # The optional keyword argument ny allows for a different size in the y direction.
628        g = gauss_kern(n, sizey=ny)
629        improc = convolve(im, g, mode='same')
630        return improc
631
[345]632## Author: AS
[233]633def getwinddef (nc):   
634    ###
[548]635    varinfile = nc.variables.keys()
636    if 'Um' in varinfile:   [uchar,vchar] = ['Um','Vm'] #; print "this is API meso file"
637    elif 'U' in varinfile:  [uchar,vchar] = ['U','V']   #; print "this is RAW meso file"
638    elif 'u' in varinfile:  [uchar,vchar] = ['u','v']   #; print "this is GCM file"
639     ### you can add choices here !
640    else:                   [uchar,vchar] = ['not found','not found']
[233]641    ###
[548]642    if uchar in ['U']:         metwind = False ## geometrical (wrt grid)
643    else:                      metwind = True  ## meteorological (zon/mer)
644    if metwind is False:       print "Not using meteorological winds. You trust numerical grid as being (x,y)"
[233]645    ###
646    return uchar,vchar,metwind
[202]647
[345]648## Author: AS
[184]649def vectorfield (u, v, x, y, stride=3, scale=15., factor=250., color='black', csmooth=1, key=True):
650    ## scale regle la reference du vecteur
651    ## factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
652    import  matplotlib.pyplot               as plt
653    import  numpy                           as np
[638]654    #posx = np.min(x) - np.std(x) / 10.
655    #posy = np.min(y) - np.std(y) / 10.
656    posx = np.min(x) 
657    posy = np.min(y) - 4.*np.std(y) / 10.
[184]658    u = smooth(u,csmooth)
659    v = smooth(v,csmooth)
[188]660    widthvec = 0.003 #0.005 #0.003
[184]661    q = plt.quiver( x[::stride,::stride],\
662                    y[::stride,::stride],\
663                    u[::stride,::stride],\
664                    v[::stride,::stride],\
[228]665                    angles='xy',color=color,pivot='middle',\
[184]666                    scale=factor,width=widthvec )
[202]667    if color in ['white','yellow']:     kcolor='black'
668    else:                               kcolor=color
[184]669    if key: p = plt.quiverkey(q,posx,posy,scale,\
[194]670                   str(int(scale)),coordinates='data',color=kcolor,labelpos='S',labelsep = 0.03)
[184]671    return 
[180]672
[345]673## Author: AS
[180]674def display (name):
[184]675    from os import system
676    system("display "+name+" > /dev/null 2> /dev/null &")
677    return name
[180]678
[345]679## Author: AS
[180]680def findstep (wlon):
[184]681    steplon = int((wlon[1]-wlon[0])/4.)  #3
682    step = 120.
683    while step > steplon and step > 15. :       step = step / 2.
684    if step <= 15.:
685        while step > steplon and step > 5.  :   step = step - 5.
686    if step <= 5.:
687        while step > steplon and step > 1.  :   step = step - 1.
688    if step <= 1.:
689        step = 1. 
[180]690    return step
691
[345]692## Author: AS
[451]693def define_proj (char,wlon,wlat,back=None,blat=None,blon=None):
[180]694    from    mpl_toolkits.basemap            import Basemap
695    import  numpy                           as np
696    import  matplotlib                      as mpl
[240]697    from mymath import max
[180]698    meanlon = 0.5*(wlon[0]+wlon[1])
699    meanlat = 0.5*(wlat[0]+wlat[1])
[637]700    zewidth = np.abs(wlon[0]-wlon[1])*60000.*np.cos(3.14*meanlat/180.)
701    zeheight = np.abs(wlat[0]-wlat[1])*60000.
[385]702    if blat is None:
[398]703        ortholat=meanlat
[453]704        if   wlat[0] >= 80.:   blat =  -40. 
[345]705        elif wlat[1] <= -80.:  blat = -40.
706        elif wlat[1] >= 0.:    blat = wlat[0] 
707        elif wlat[0] <= 0.:    blat = wlat[1]
[398]708    else:  ortholat=blat
[451]709    if blon is None:  ortholon=meanlon
710    else:             ortholon=blon
[207]711    h = 50.  ## en km
[202]712    radius = 3397200.
[637]713    #print meanlat, meanlon
[184]714    if   char == "cyl":     m = Basemap(rsphere=radius,projection='cyl',\
[180]715                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
[184]716    elif char == "moll":    m = Basemap(rsphere=radius,projection='moll',lon_0=meanlon)
[451]717    elif char == "ortho":   m = Basemap(rsphere=radius,projection='ortho',lon_0=ortholon,lat_0=ortholat)
[184]718    elif char == "lcc":     m = Basemap(rsphere=radius,projection='lcc',lat_1=meanlat,lat_0=meanlat,lon_0=meanlon,\
[637]719                              width=zewidth,height=zeheight)
720                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
[184]721    elif char == "npstere": m = Basemap(rsphere=radius,projection='npstere', boundinglat=blat, lon_0=0.)
[395]722    elif char == "spstere": m = Basemap(rsphere=radius,projection='spstere', boundinglat=blat, lon_0=180.)
[207]723    elif char == "nplaea":  m = Basemap(rsphere=radius,projection='nplaea', boundinglat=wlat[0], lon_0=meanlon)
724    elif char == "laea":    m = Basemap(rsphere=radius,projection='laea',lon_0=meanlon,lat_0=meanlat,lat_ts=meanlat,\
[637]725                              width=zewidth,height=zeheight)
726                              #llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
[184]727    elif char == "nsper":   m = Basemap(rsphere=radius,projection='nsper',lon_0=meanlon,lat_0=meanlat,satellite_height=h*1000.)
728    elif char == "merc":    m = Basemap(rsphere=radius,projection='merc',lat_ts=0.,\
729                              llcrnrlat=wlat[0],urcrnrlat=wlat[1],llcrnrlon=wlon[0],urcrnrlon=wlon[1])
730    fontsizemer = int(mpl.rcParams['font.size']*3./4.)
[207]731    if char in ["cyl","lcc","merc","nsper","laea"]:   step = findstep(wlon)
732    else:                                             step = 10.
[238]733    steplon = step*2.
[453]734    zecolor ='grey'
735    zelinewidth = 1
736    zelatmax = 80
[637]737    if meanlat > 75.: zelatmax = 90. ; step = step/2.
[453]738    # to show gcm grid:
739    #zecolor = 'r'
740    #zelinewidth = 1
741    #step = 5.625
742    #steplon = 5.625
743    #zelatmax = 89.9
[516]744    if char not in ["moll"]:
745        m.drawmeridians(np.r_[-180.:180.:steplon], labels=[0,0,0,1], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
746        m.drawparallels(np.r_[-90.:90.:step], labels=[1,0,0,0], color=zecolor, linewidth=zelinewidth, fontsize=fontsizemer, latmax=zelatmax)
[233]747    if back: m.warpimage(marsmap(back),scale=0.75)
748            #if not back:
749            #    if not var:                                        back = "mola"    ## if no var:         draw mola
750            #    elif typefile in ['mesoapi','meso','geo'] \
751            #       and proj not in ['merc','lcc','nsper','laea']:  back = "molabw"  ## if var but meso:   draw molabw
752            #    else:                                              pass             ## else:              draw None
[180]753    return m
754
[345]755## Author: AS
[232]756#### test temporaire
757def putpoints (map,plot):
758    #### from http://www.scipy.org/Cookbook/Matplotlib/Maps
759    # lat/lon coordinates of five cities.
760    lats = [18.4]
761    lons = [-134.0]
762    points=['Olympus Mons']
763    # compute the native map projection coordinates for cities.
764    x,y = map(lons,lats)
765    # plot filled circles at the locations of the cities.
766    map.plot(x,y,'bo')
767    # plot the names of those five cities.
768    wherept = 0 #1000 #50000
769    for name,xpt,ypt in zip(points,x,y):
770       plot.text(xpt+wherept,ypt+wherept,name)
771    ## le nom ne s'affiche pas...
772    return
773
[345]774## Author: AS
[233]775def calculate_bounds(field,vmin=None,vmax=None):
776    import numpy as np
777    from mymath import max,min,mean
778    ind = np.where(field < 9e+35)
779    fieldcalc = field[ ind ] # la syntaxe compacte ne marche si field est un tuple
780    ###
781    dev = np.std(fieldcalc)*3.0
782    ###
[562]783    if vmin is None:  zevmin = mean(fieldcalc) - dev
[233]784    else:             zevmin = vmin
785    ###
786    if vmax is None:  zevmax = mean(fieldcalc) + dev
787    else:             zevmax = vmax
788    if vmin == vmax:
789                      zevmin = mean(fieldcalc) - dev  ### for continuity
790                      zevmax = mean(fieldcalc) + dev  ### for continuity           
791    ###
792    if zevmin < 0. and min(fieldcalc) > 0.: zevmin = 0.
[468]793    print "BOUNDS field ", min(fieldcalc), max(fieldcalc), " //// adopted", zevmin, zevmax
[233]794    return zevmin, zevmax
[232]795
[345]796## Author: AS
[233]797def bounds(what_I_plot,zevmin,zevmax):
[247]798    from mymath import max,min,mean
[233]799    ### might be convenient to add the missing value in arguments
[310]800    #what_I_plot[ what_I_plot < zevmin ] = zevmin#*(1. + 1.e-7)
801    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1. - 1.e-7) ] = zevmin*(1. - 1.e-7)
802    else:          what_I_plot[ what_I_plot < zevmin*(1. + 1.e-7) ] = zevmin*(1. + 1.e-7)
[451]803    #print "NEW MIN ", min(what_I_plot)
[233]804    what_I_plot[ what_I_plot > 9e+35  ] = -9e+35
[587]805    what_I_plot[ what_I_plot > zevmax ] = zevmax*(1. - 1.e-7)
[451]806    #print "NEW MAX ", max(what_I_plot)
[233]807    return what_I_plot
808
[345]809## Author: AS
[241]810def nolow(what_I_plot):
811    from mymath import max,min
812    lim = 0.15*0.5*(abs(max(what_I_plot))+abs(min(what_I_plot)))
[392]813    print "NO PLOT BELOW VALUE ", lim
[241]814    what_I_plot [ abs(what_I_plot) < lim ] = 1.e40 
815    return what_I_plot
816
[418]817
818## Author : AC
819def hole_bounds(what_I_plot,zevmin,zevmax):
820    import numpy as np
821    zi=0
822    for i in what_I_plot:
823        zj=0
824        for j in i:
825            if ((j < zevmin) or (j > zevmax)):what_I_plot[zi,zj]=np.NaN
826            zj=zj+1
827        zi=zi+1
828
829    return what_I_plot
830
[345]831## Author: AS
[233]832def zoomset (wlon,wlat,zoom):
833    dlon = abs(wlon[1]-wlon[0])/2.
834    dlat = abs(wlat[1]-wlat[0])/2.
835    [wlon,wlat] = [ [wlon[0]+zoom*dlon/100.,wlon[1]-zoom*dlon/100.],\
836                    [wlat[0]+zoom*dlat/100.,wlat[1]-zoom*dlat/100.] ]
[392]837    print "ZOOM %",zoom,wlon,wlat
[233]838    return wlon,wlat
839
[345]840## Author: AS
[201]841def fmtvar (whichvar="def"):
[204]842    fmtvar    =     { \
[502]843             "MIXED":        "%.0f",\
844             "UPDRAFT":      "%.0f",\
845             "DOWNDRAFT":    "%.0f",\
[405]846             "TK":           "%.0f",\
[637]847             "T":            "%.0f",\
[516]848             #"ZMAX_TH":      "%.0f",\
849             #"WSTAR":        "%.0f",\
[425]850             # Variables from TES ncdf format
[363]851             "T_NADIR_DAY":  "%.0f",\
[376]852             "T_NADIR_NIT":  "%.0f",\
[425]853             # Variables from tes.py ncdf format
[398]854             "TEMP_DAY":     "%.0f",\
855             "TEMP_NIGHT":   "%.0f",\
[425]856             # Variables from MCS and mcs.py ncdf format
[427]857             "DTEMP":        "%.0f",\
858             "NTEMP":        "%.0f",\
859             "DNUMBINTEMP":  "%.0f",\
860             "NNUMBINTEMP":  "%.0f",\
[425]861             # other stuff
[405]862             "TPOT":         "%.0f",\
[295]863             "TSURF":        "%.0f",\
[612]864             "U_OUT1":       "%.0f",\
865             "T_OUT1":       "%.0f",\
[204]866             "def":          "%.1e",\
867             "PTOT":         "%.0f",\
868             "HGT":          "%.1e",\
869             "USTM":         "%.2f",\
[225]870             "HFX":          "%.0f",\
[232]871             "ICETOT":       "%.1e",\
[237]872             "TAU_ICE":      "%.2f",\
[451]873             "TAUICE":       "%.2f",\
[252]874             "VMR_ICE":      "%.1e",\
[345]875             "MTOT":         "%.1f",\
[405]876             "ANOMALY":      "%.1f",\
[241]877             "W":            "%.1f",\
[287]878             "WMAX_TH":      "%.1f",\
[562]879             "WSTAR":        "%.1f",\
[287]880             "QSURFICE":     "%.0f",\
[405]881             "UM":           "%.0f",\
[490]882             "WIND":         "%.0f",\
[612]883             "UVMET":         "%.0f",\
884             "UV":         "%.0f",\
[295]885             "ALBBARE":      "%.2f",\
[317]886             "TAU":          "%.1f",\
[382]887             "CO2":          "%.2f",\
[345]888             #### T.N.
889             "TEMP":         "%.0f",\
890             "VMR_H2OICE":   "%.0f",\
891             "VMR_H2OVAP":   "%.0f",\
892             "TAUTES":       "%.2f",\
893             "TAUTESAP":     "%.2f",\
894
[204]895                    }
[518]896    if "TSURF" in whichvar: whichvar = "TSURF"
[204]897    if whichvar not in fmtvar:
898        whichvar = "def"
899    return fmtvar[whichvar]
[201]900
[345]901## Author: AS
[233]902####################################################################################################################
903### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
[202]904def defcolorb (whichone="def"):
[204]905    whichcolorb =    { \
906             "def":          "spectral",\
907             "HGT":          "spectral",\
[426]908             "HGT_M":        "spectral",\
[405]909             "TK":           "gist_heat",\
[425]910             "TPOT":         "Paired",\
[295]911             "TSURF":        "RdBu_r",\
[204]912             "QH2O":         "PuBu",\
913             "USTM":         "YlOrRd",\
[490]914             "WIND":         "YlOrRd",\
[363]915             #"T_nadir_nit":  "RdBu_r",\
916             #"T_nadir_day":  "RdBu_r",\
[225]917             "HFX":          "RdYlBu",\
[310]918             "ICETOT":       "YlGnBu_r",\
[345]919             #"MTOT":         "PuBu",\
920             "CCNQ":         "YlOrBr",\
921             "CCNN":         "YlOrBr",\
922             "TEMP":         "Jet",\
[238]923             "TAU_ICE":      "Blues",\
[451]924             "TAUICE":       "Blues",\
[252]925             "VMR_ICE":      "Blues",\
[241]926             "W":            "jet",\
[287]927             "WMAX_TH":      "spectral",\
[405]928             "ANOMALY":      "RdBu_r",\
[287]929             "QSURFICE":     "hot_r",\
[295]930             "ALBBARE":      "spectral",\
[317]931             "TAU":          "YlOrBr_r",\
[382]932             "CO2":          "YlOrBr_r",\
[345]933             #### T.N.
934             "MTOT":         "Jet",\
935             "H2O_ICE_S":    "RdBu",\
936             "VMR_H2OICE":   "PuBu",\
937             "VMR_H2OVAP":   "PuBu",\
[453]938             "WATERCAPTAG":  "Blues",\
[204]939                     }
[241]940#W --> spectral ou jet
[240]941#spectral BrBG RdBu_r
[392]942    #print "predefined colorbars"
[518]943    if "TSURF" in whichone: whichone = "TSURF"
[204]944    if whichone not in whichcolorb:
945        whichone = "def"
946    return whichcolorb[whichone]
[202]947
[345]948## Author: AS
[202]949def definecolorvec (whichone="def"):
950        whichcolor =    { \
951                "def":          "black",\
952                "vis":          "yellow",\
953                "vishires":     "yellow",\
954                "molabw":       "yellow",\
955                "mola":         "black",\
956                "gist_heat":    "white",\
957                "hot":          "tk",\
958                "gist_rainbow": "black",\
959                "spectral":     "black",\
960                "gray":         "red",\
961                "PuBu":         "black",\
962                        }
963        if whichone not in whichcolor:
964                whichone = "def"
965        return whichcolor[whichone]
966
[345]967## Author: AS
[180]968def marsmap (whichone="vishires"):
[233]969        from os import uname
970        mymachine = uname()[1]
971        ### not sure about speed-up with this method... looks the same
[511]972        if "lmd.jussieu.fr" in mymachine:   domain = "/u/aslmd/WWW/maps/"
973        elif "aymeric-laptop" in mymachine: domain = "/home/aymeric/Dropbox/Public/"
974        else:                               domain = "http://www.lmd.jussieu.fr/~aslmd/maps/"
[180]975        whichlink =     { \
[233]976                #"vis":         "http://maps.jpl.nasa.gov/pix/mar0kuu2.jpg",\
977                #"vishires":    "http://www.lmd.jussieu.fr/~aslmd/maps/MarsMap_2500x1250.jpg",\
978                #"geolocal":    "http://dl.dropbox.com/u/11078310/geolocal.jpg",\
979                #"mola":        "http://www.lns.cornell.edu/~seb/celestia/mars-mola-2k.jpg",\
980                #"molabw":      "http://dl.dropbox.com/u/11078310/MarsElevation_2500x1250.jpg",\
[453]981                "thermalday":   domain+"thermalday.jpg",\
982                "thermalnight": domain+"thermalnight.jpg",\
983                "tesalbedo":    domain+"tesalbedo.jpg",\
[233]984                "vis":         domain+"mar0kuu2.jpg",\
985                "vishires":    domain+"MarsMap_2500x1250.jpg",\
986                "geolocal":    domain+"geolocal.jpg",\
987                "mola":        domain+"mars-mola-2k.jpg",\
988                "molabw":      domain+"MarsElevation_2500x1250.jpg",\
[238]989                "clouds":      "http://www.johnstonsarchive.net/spaceart/marswcloudmap.jpg",\
990                "jupiter":     "http://www.mmedia.is/~bjj/data/jupiter_css/jupiter_css.jpg",\
991                "jupiter_voy": "http://www.mmedia.is/~bjj/data/jupiter/jupiter_vgr2.jpg",\
[558]992                #"bw":          domain+"EarthElevation_2500x1250.jpg",\
[273]993                "bw":          "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg",\
994                "contrast":    "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg",\
995                "nice":        "http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg",\
996                "blue":        "http://eoimages.gsfc.nasa.gov/ve/2430/land_ocean_ice_2048.jpg",\
[296]997                "blueclouds":  "http://eoimages.gsfc.nasa.gov/ve/2431/land_ocean_ice_cloud_2048.jpg",\
998                "justclouds":  "http://eoimages.gsfc.nasa.gov/ve/2432/cloud_combined_2048.jpg",\
[180]999                        }
[238]1000        ### see http://www.mmedia.is/~bjj/planetary_maps.html
[180]1001        if whichone not in whichlink: 
1002                print "marsmap: choice not defined... you'll get the default one... "
1003                whichone = "vishires" 
1004        return whichlink[whichone]
1005
[273]1006#def earthmap (whichone):
1007#       if   whichone == "contrast":    whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthMapAtmos_2500x1250.jpg"
1008#       elif whichone == "bw":          whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/EarthElevation_2500x1250.jpg"
1009#       elif whichone == "nice":        whichlink="http://users.info.unicaen.fr/~karczma/TEACH/InfoGeo/Images/Planets/earthmap1k.jpg"
1010#       return whichlink
[180]1011
[345]1012## Author: AS
[241]1013def latinterv (area="Whole"):
1014    list =    { \
1015        "Europe":                [[ 20., 80.],[- 50.,  50.]],\
1016        "Central_America":       [[-10., 40.],[ 230., 300.]],\
1017        "Africa":                [[-20., 50.],[- 50.,  50.]],\
[273]1018        "Whole":                 [[-90., 90.],[-180., 180.]],\
1019        "Southern_Hemisphere":   [[-90., 60.],[-180., 180.]],\
1020        "Northern_Hemisphere":   [[-60., 90.],[-180., 180.]],\
[241]1021        "Tharsis":               [[-30., 60.],[-170.,- 10.]],\
1022        "Whole_No_High":         [[-60., 60.],[-180., 180.]],\
1023        "Chryse":                [[-60., 60.],[- 60.,  60.]],\
1024        "North_Pole":            [[ 50., 90.],[-180., 180.]],\
1025        "Close_North_Pole":      [[ 75., 90.],[-180., 180.]],\
1026        "Far_South_Pole":        [[-90.,-40.],[-180., 180.]],\
1027        "South_Pole":            [[-90.,-50.],[-180., 180.]],\
1028        "Close_South_Pole":      [[-90.,-75.],[-180., 180.]],\
[637]1029        "Sirenum_Crater_large":  [[-46.,-34.],[-166.,-151.]],\
1030        "Sirenum_Crater_small":  [[-36.,-26.],[-168.,-156.]],\
1031        "Rupes":                 [[ 72., 90.],[-120.,- 20.]],\
[241]1032              }
1033    if area not in list:   area = "Whole"
1034    [olat,olon] = list[area]
1035    return olon,olat
1036
[345]1037## Author: TN
1038def separatenames (name):
1039  from numpy import concatenate
1040  # look for comas in the input name to separate different names (files, variables,etc ..)
1041  if name is None:
1042     names = None
1043  else:
1044    names = []
1045    stop = 0
1046    currentname = name
1047    while stop == 0:
1048      indexvir = currentname.find(',')
1049      if indexvir == -1:
1050        stop = 1
1051        name1 = currentname
1052      else:
1053        name1 = currentname[0:indexvir]
1054      names = concatenate((names,[name1]))
1055      currentname = currentname[indexvir+1:len(currentname)]
1056  return names
1057
1058## Author: TN [old]
1059def getopmatrix (kind,n):
1060  import numpy as np
1061  # compute matrix of operations between files
1062  # the matrix is 'number of files'-square
1063  # 1: difference (row minus column), 2: add
1064  # not 0 in diag : just plot
1065  if n == 1:
1066    opm = np.eye(1)
1067  elif kind == 'basic':
1068    opm = np.eye(n)
1069  elif kind == 'difference1': # show differences with 1st file
1070    opm = np.zeros((n,n))
1071    opm[0,:] = 1
1072    opm[0,0] = 0
1073  elif kind == 'difference2': # show differences with 1st file AND show 1st file
1074    opm = np.zeros((n,n))
1075    opm[0,:] = 1
1076  else:
1077    opm = np.eye(n)
1078  return opm
1079 
1080## Author: TN [old]
1081def checkcoherence (nfiles,nlat,nlon,ntime):
1082  if (nfiles > 1):
1083     if (nlat > 1):
1084        errormess("what you asked is not possible !")
1085  return 1
1086
1087## Author: TN
1088def readslices(saxis):
1089  from numpy import empty
1090  if saxis == None:
1091     zesaxis = None
1092  else:
1093     zesaxis = empty((len(saxis),2))
1094     for i in range(len(saxis)):
1095        a = separatenames(saxis[i])
1096        if len(a) == 1:
1097           zesaxis[i,:] = float(a[0])
1098        else:
1099           zesaxis[i,0] = float(a[0])
1100           zesaxis[i,1] = float(a[1])
1101           
1102  return zesaxis
1103
[568]1104## Author: TN
1105def readdata(data,datatype,coord1,coord2):
1106## Read sparse data
1107  from numpy import empty
[572]1108  if datatype == 'txt':
[568]1109     if len(data[coord1].shape) == 1:
1110         return data[coord1][:]
1111     elif len(data[coord1].shape) == 2:
1112         return data[coord1][:,int(coord2)-1]
1113     else:
1114         errormess('error in readdata')
[572]1115  elif datatype == 'sav':
[568]1116     return data[coord1][coord2]
1117  else:
1118     errormess(datatype+' type is not supported!')
1119
1120
[399]1121## Author: AS
[475]1122def bidimfind(lon2d,lat2d,vlon,vlat,file=None):
[399]1123   import numpy as np
[475]1124   import matplotlib.pyplot as mpl
[399]1125   if vlat is None:    array = (lon2d - vlon)**2
1126   elif vlon is None:  array = (lat2d - vlat)**2
1127   else:               array = (lon2d - vlon)**2 + (lat2d - vlat)**2
1128   idy,idx = np.unravel_index( np.argmin(array), lon2d.shape )
1129   if vlon is not None:
[475]1130      if (np.abs(lon2d[idy,idx]-vlon)) > 5: errormess("longitude not found ",printvar=lon2d)
[399]1131   if vlat is not None:
[475]1132      if (np.abs(lat2d[idy,idx]-vlat)) > 5: errormess("latitude not found ",printvar=lat2d)
1133   if file is not None:
1134      print idx,idy,lon2d[idy,idx],vlon
1135      print idx,idy,lat2d[idy,idx],vlat
1136      var = file.variables["HGT"][:,:,:]
[489]1137      mpl.contourf(var[0,:,:],30,cmap = mpl.get_cmap(name="Greys_r") ) ; mpl.axis('off') ; mpl.plot(idx,idy,'mx',mew=4.0,ms=20.0)
[475]1138      mpl.show()
1139   return idy,idx
[399]1140
[345]1141## Author: TN
[399]1142def getsindex(saxis,index,axis):
[345]1143# input  : all the desired slices and the good index
1144# output : all indexes to be taken into account for reducing field
1145  import numpy as np
[349]1146  if ( np.array(axis).ndim == 2):
1147      axis = axis[:,0]
[345]1148  if saxis is None:
1149      zeindex = None
1150  else:
1151      aaa = int(np.argmin(abs(saxis[index,0] - axis)))
1152      bbb = int(np.argmin(abs(saxis[index,1] - axis)))
1153      [imin,imax] = np.sort(np.array([aaa,bbb]))
1154      zeindex = np.array(range(imax-imin+1))+imin
1155      # because -180 and 180 are the same point in longitude,
1156      # we get rid of one for averaging purposes.
1157      if axis[imin] == -180 and axis[imax] == 180:
1158         zeindex = zeindex[0:len(zeindex)-1]
[392]1159         print "INFO: whole longitude averaging asked, so last point is not taken into account."
[345]1160  return zeindex
1161     
1162## Author: TN
1163def define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,indextime,what_I_plot,dim0,vertmode):
1164# Purpose of define_axis is to find x and y axis scales in a smart way
1165# x axis priority: 1/time 2/lon 3/lat 4/vertical
1166# To be improved !!!...
1167   from numpy import array,swapaxes
1168   x = None
1169   y = None
1170   count = 0
1171   what_I_plot = array(what_I_plot)
1172   shape = what_I_plot.shape
[477]1173   if indextime is None and len(time) > 1:
[392]1174      print "AXIS is time"
[345]1175      x = time
1176      count = count+1
[477]1177   if indexlon is None and len(lon) > 1:
[392]1178      print "AXIS is lon"
[345]1179      if count == 0: x = lon
1180      else: y = lon
1181      count = count+1
[579]1182   if indexlat is None and len(lat) > 1:
[392]1183      print "AXIS is lat"
[345]1184      if count == 0: x = lat
1185      else: y = lat
1186      count = count+1
[579]1187   if indexvert is None and ((dim0 == 4) or (y is None)):
[392]1188      print "AXIS is vert"
[345]1189      if vertmode == 0: # vertical axis is as is (GCM grid)
1190         if count == 0: x=range(len(vert))
1191         else: y=range(len(vert))
1192         count = count+1
1193      else: # vertical axis is in kms
1194         if count == 0: x = vert
1195         else: y = vert
1196         count = count+1
1197   x = array(x)
1198   y = array(y)
[468]1199   print "CHECK SHAPE: what_I_plot, x, y", what_I_plot.shape, x.shape, y.shape
[345]1200   if len(shape) == 1:
[562]1201       if shape[0] != len(x):           print "WARNING: shape[0] != len(x). Correcting." ; what_I_plot = what_I_plot[0:len(x)]
[579]1202       if len(y.shape) > 0:             y = ()
[345]1203   elif len(shape) == 2:
[562]1204       if shape[1] == len(y) and shape[0] == len(x) and shape[0] != shape[1]:
1205           print "INFO: swapaxes: ",what_I_plot.shape,shape ; what_I_plot = swapaxes(what_I_plot,0,1)
1206       else:
1207           if shape[0] != len(y):       print "WARNING: shape[0] != len(y). Correcting." ; what_I_plot = what_I_plot[0:len(y),:]
1208           elif shape[1] != len(x):     print "WARNING: shape[1] != len(x). Correcting." ; what_I_plot = what_I_plot[:,0:len(x)]
1209   elif len(shape) == 3:
1210       if vertmode < 0: print "not supported. must check array dimensions at some point. not difficult to implement though."
[345]1211   return what_I_plot,x,y
[349]1212
[428]1213# Author: TN + AS
[349]1214def determineplot(slon, slat, svert, stime):
1215    nlon = 1 # number of longitudinal slices -- 1 is None
1216    nlat = 1
1217    nvert = 1
1218    ntime = 1
1219    nslices = 1
1220    if slon is not None:
1221        nslices = nslices*len(slon)
1222        nlon = len(slon)
1223    if slat is not None:
1224        nslices = nslices*len(slat)
1225        nlat = len(slat)
1226    if svert is not None:
1227        nslices = nslices*len(svert)
1228        nvert = len(svert)
1229    if stime is not None:
1230        nslices = nslices*len(stime)
1231        ntime = len(stime)
1232    #else:
1233    #    nslices = 2 
1234    mapmode = 0
1235    if slon is None and slat is None:
1236       mapmode = 1 # in this case we plot a map, with the given projection
1237
1238    return nlon, nlat, nvert, ntime, mapmode, nslices
[440]1239
[448]1240## Author: AC
1241## Reduce complexity of main script by moving the contour part here. Also allow to call it from elsewhere
1242## which can be usefull
1243#
1244#def call_contour(what_I_plot,error,x,y,m,lon,lat,vert,time,vertmode,ze_var2,indextime,indexlon,indexlat,indexvert,yintegral,mapmode,typefile,var2,ticks):
1245#      from matplotlib.pyplot import contour, plot, clabel
1246#      import numpy as np
1247#      #what_I_plot = what_I_plot*mult
1248#      if not error:
1249#         if mapmode == 1:
1250#             if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
1251#         zevmin, zevmax = calculate_bounds(what_I_plot)
1252#         zelevels = np.linspace(zevmin,zevmax,ticks) #20)
1253#         if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
1254#         if mapmode == 0:
1255#             #if typefile in ['mesoideal']:    what_I_plot = dumpbdy(what_I_plot,0,stag='W')
1256#             itime=indextime
1257#             if len(what_I_plot.shape) is 3: itime=[0]
1258#             what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
1259#                   itime,what_I_plot, len(ze_var2.shape),vertmode)
1260#         ### If we plot a 2-D field
1261#         if len(what_I_plot.shape) is 2:
1262#             #zelevels=[1.]
1263#             if mapmode == 0:cs = contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
1264#             elif mapmode == 1:cs = m.contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
1265#             #clabel(cs,zelevels,inline=3,fmt='%1.1f',fontsize=7)
1266#         ### If we plot a 1-D field
1267#         elif len(what_I_plot.shape) is 1:
1268#             plot(what_I_plot,x)
1269#      else:
1270#         errormess("There is an error in reducing field !")
1271#      return error
[440]1272
[638]1273## Author : AS
1274def maplatlon( lon,lat,field,\
1275               proj="cyl",colorb="jet",ndiv=10,zeback="molabw",trans=0.6,title="",\
1276               vecx=None,vecy=None,stride=2 ):
1277    ### an easy way to map a field over lat/lon grid
1278    import numpy as np
1279    import matplotlib.pyplot as mpl
1280    from matplotlib.cm import get_cmap
1281    ## get lon and lat in 2D version. get lat/lon intervals
1282    numdim = len(np.array(lon).shape)
1283    if numdim == 2:     [lon2d,lat2d] = [lon,lat]
1284    elif numdim == 1:   [lon2d,lat2d] = np.meshgrid(lon,lat)
1285    else:               errormess("lon and lat arrays must be 1D or 2D")
1286    [wlon,wlat] = latinterv()
1287    ## define projection and background. define x and y given the projection
1288    m = define_proj(proj,wlon,wlat,back=zeback,blat=None,blon=None)
1289    x, y = m(lon2d, lat2d)
1290    ## define field. bound field.
1291    what_I_plot = np.transpose(field)
1292    zevmin, zevmax = calculate_bounds(what_I_plot)  ## vmin=min(what_I_plot_frame), vmax=max(what_I_plot_frame))
1293    what_I_plot = bounds(what_I_plot,zevmin,zevmax)
1294    ## define contour field levels. define color palette
1295    ticks = ndiv + 1
1296    zelevels = np.linspace(zevmin,zevmax,ticks)
1297    palette = get_cmap(name=colorb)
1298    ## contour field
1299    m.contourf( x, y, what_I_plot, zelevels, cmap = palette, alpha = trans )
1300    ## draw colorbar
1301    if proj in ['moll','cyl']:        zeorientation="horizontal" ; zepad = 0.07
1302    else:                             zeorientation="vertical" ; zepad = 0.03
1303    #daformat = fmtvar(fvar.upper())
1304    daformat = "%.0f"
1305    zecb = mpl.colorbar( fraction=0.05,pad=zepad,format=daformat,orientation=zeorientation,\
1306                 ticks=np.linspace(zevmin,zevmax,num=min([ticks/2+1,21])),extend='neither',spacing='proportional' ) 
1307    ## give a title
1308    if zeorientation == "horizontal": zecb.ax.set_xlabel(title)
1309    else:                             ptitle(title)
1310    ## draw vector
1311    if vecx is not None and vecy is not None:
1312       [vecx_frame,vecy_frame] = m.rotate_vector( np.transpose(vecx), np.transpose(vecy), lon2d, lat2d ) ## for metwinds
1313       vectorfield(vecx_frame, vecy_frame, x, y, stride=stride, csmooth=2,\
1314                                             scale=30., factor=500., color=definecolorvec(colorb), key=True)
1315    ## scale regle la reference du vecteur. factor regle toutes les longueurs (dont la reference). l'AUGMENTER pour raccourcir les vecteurs.
1316    return
Note: See TracBrowser for help on using the repository browser.