Context Navigation

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

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

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

Download in other formats: