Context Navigation

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

Last change on this file since 428 was 428, checked in by acolaitis, 13 years ago
PYTHON. Very simple modifications to allow for slices with LES files. Performance is not optimized and turbulence computations are not implemented. More to come in time.
File size: 40.8 KB

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

Note: See TracBrowser for help on using the repository browser.

Download in other formats: