Context Navigation

source: trunk/MESOSCALE/PLOT/PYTHON/scripts/winds.py @ 201

Last change on this file since 201 was 201, checked in by aslmd, 14 years ago
MESOSCALE: python. a few bug fixes and additions.
Property svn:executable set to ``*
File size: 15.2 KB

Rev	Line
[180]	1	#!/usr/bin/env python
	2
[191]	3	### A. Spiga -- LMD -- 30/06/2011 to 04/07/2011
[184]	4	### Thanks to A. Colaitis for the parser trick
[180]	5
[184]	6
	7	####################################
	8	####################################
	9	### The main program to plot vectors
[182]	10	def winds (namefile,\
	11	nvert,\
[184]	12	proj=None,\
[182]	13	back=None,\
	14	target=None,
[184]	15	stride=3,\
	16	numplot=4,\
[187]	17	var=None,\
[201]	18	colorb=True,\
[195]	19	winds=True,\
[199]	20	addchar=None,\
	21	interv=[0,1],\
	22	vmin=None,\
[201]	23	vmax=None,\
	24	tile=False):
[184]	25
[192]	26	####################################################################################################################
	27	### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
	28
[184]	29	#################################
	30	### Load librairies and functions
[180]	31	from netCDF4 import Dataset
[201]	32	from myplot import getcoord2d,define_proj,makeplotpng,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy,fmtvar
[199]	33	from mymath import deg,max,min,mean
	34	from matplotlib.pyplot import contourf, subplot, figure, rcParams, savefig, colorbar, pcolor
[192]	35	from matplotlib.cm import get_cmap
[182]	36	import numpy as np
[184]	37
	38	######################
	39	### Load NETCDF object
[180]	40	nc = Dataset(namefile)
[184]	41
	42	###################################
	43	### Recognize predefined file types
	44	if 'controle' in nc.variables: typefile = 'gcm'
[191]	45	elif 'vert' in nc.variables: typefile = 'mesoapi'
[184]	46	elif 'U' in nc.variables: typefile = 'meso'
	47	else:
	48	print "typefile not supported."
[191]	49	print nc.variables
[182]	50	exit()
[184]	51
	52	##############################################################
	53	### Try to guess the projection from wrfout if not set by user
	54	if typefile in ['mesoapi','meso']:
	55	if proj == None: proj = getproj(nc)
	56	### (il faudrait passer CEN_LON dans la projection ?)
	57	elif typefile in ['gcm']:
	58	if proj == None: proj = "cyl"
	59	## pb avec les autres (de trace derriere la sphere ?)
	60
[188]	61	############################################
	62	#### Choose underlying topography by default
	63	if not back:
	64	if not var: back = "mola" ## if no var: draw mola
	65	elif typefile in ['mesoapi','meso'] \
	66	and proj not in ['merc','lcc']: back = "molabw" ## if var but meso: draw molabw
	67	else: pass ## else: draw None
[184]	68
	69	####################################################
	70	### Get geographical coordinates and plot boundaries
	71	if typefile in ['mesoapi','meso']:
	72	[lon2d,lat2d] = getcoord2d(nc)
	73	lon2d = dumpbdy(lon2d)
	74	lat2d = dumpbdy(lat2d)
	75	elif typefile in ['gcm']:
	76	[lon2d,lat2d] = getcoord2d(nc,nlat="latitude",nlon="longitude",is1d=True)
	77	if proj == "npstere": [wlon,wlat] = latinterv("North_Pole")
	78	elif proj == "lcc": [wlon,wlat] = wrfinterv(lon2d,lat2d)
	79	else: [wlon,wlat] = simplinterv(lon2d,lat2d)
	80
[197]	81	###################
	82	#### Get local time
	83	#if typefile in ['mesoapi','meso']:
	84	# ltst = interv[0]
	85	# #ltst = int(getattr(nc, 'GMT') + 0.5*(wlon[0]+wlon[1])/15.)
	86	#elif typefile in ['gcm']:
	87	# ltst = 0
[184]	88
	89	##############################################################################
	90	### Get winds and know if those are meteorological winds (ie. zon, mer) or not
[191]	91	if winds:
	92	if typefile is 'mesoapi':
	93	[u,v] = getwinds(nc)
	94	metwind = True ## meteorological (zon/mer)
	95	elif typefile is 'gcm':
	96	[u,v] = getwinds(nc,charu='u',charv='v')
	97	metwind = True ## meteorological (zon/mer)
	98	elif typefile is 'meso':
	99	[u,v] = getwinds(nc,charu='U',charv='V')
	100	metwind = False ## geometrical (wrt grid)
	101	print "Beware ! Not using meteorological winds. You trust numerical grid as being (x,y)."
[184]	102
[191]	103	#####################################################
[184]	104	### Load the chosen variables, whether it is 2D or 3D
	105	if var:
	106	if var not in nc.variables:
	107	print "not found in file:",var
	108	exit()
	109	else:
	110	dimension = np.array(nc.variables[var]).ndim
	111	if dimension == 2: field = nc.variables[var][:,:]
	112	elif dimension == 3: field = nc.variables[var][:,:,:]
	113	elif dimension == 4: field = nc.variables[var][:,nvert,:,:]
[201]	114	dev = np.std(field)*3.0
[199]	115	if vmin is None: zevmin = mean(field) - dev
	116	else: zevmin = vmin
	117	if vmax is None: zevmax = mean(field) + dev
	118	else: zevmax = vmax
	119	print "bounds ", zevmin, zevmax
[201]	120
[191]	121	###########################
	122	### Get length of time axis
	123	if winds: nt = len(u[:,0,0,0])
	124	elif var:
	125	if dimension == 2: nt = 1
	126	else : nt = len(field[:,0,0])
	127
[186]	128	#########################################
	129	### Name for title and graphics save file
[191]	130	if winds: basename = "WINDS_"
	131	else: basename = ""
[186]	132	if var:
[191]	133	basename = basename + var
[188]	134	if typefile is 'meso': stralt = "_lvl" + str(nvert)
	135	elif typefile is 'mesoapi':
	136	zelevel = int(nc.variables['vert'][nvert])
	137	if 'altitude_abg' in nc.dimensions: stralt = "_"+str(zelevel)+"m-ALS"
	138	elif 'bottom_top' in nc.dimensions: stralt = "_"+str(zelevel)+"m-AMR"
	139	elif 'pressure' in nc.dimensions: stralt = "_"+str(zelevel)+"Pa"
	140	else: stralt = ""
	141	basename = basename + stralt
[186]	142
[184]	143	##################################
	144	### Open a figure and set subplots
[180]	145	fig = figure()
[186]	146	if numplot > 0:
	147	if numplot == 4:
	148	sub = 221
[201]	149	#fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
	150	fig.subplots_adjust(wspace = 0.3, hspace = 0.3)
[186]	151	rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
	152	elif numplot == 2:
	153	sub = 121
	154	fig.subplots_adjust(wspace = 0.3)
	155	rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
	156	elif numplot == 3:
	157	sub = 131
	158	fig.subplots_adjust(wspace = 0.3)
	159	rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
	160	elif numplot == 6:
	161	sub = 231
	162	fig.subplots_adjust(wspace = 0.4, hspace = 0.0)
	163	rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
	164	elif numplot == 8:
	165	sub = 331 #241
	166	fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
	167	rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
	168	elif numplot == 9:
	169	sub = 331
	170	fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
	171	rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
	172	elif numplot == 1:
[189]	173	sub = 99999
[186]	174	else:
	175	print "supported: 1,2,3,4,6,8"
	176	exit()
	177	### Prepare time loop
	178	if nt <= numplot or numplot == 1:
	179	tabrange = [0]
	180	numplot = 1
	181	else:
	182	tabrange = range(0,nt,int(nt/numplot)) #nt-1
	183	tabrange = tabrange[0:numplot]
	184	else:
	185	tabrange = range(0,nt,1)
	186	sub = 99999
	187	print tabrange
[184]	188
	189	#################################
	190	### Time loop for plotting device
[186]	191	found_lct = False
[184]	192	for i in tabrange:
	193
[197]	194	### Which local time ?
	195	ltst = ( interv[0] + 0.5(wlon[0]+wlon[1])/15.) + iinterv[1]
	196	ltst = int (ltst * 10) / 10.
	197	ltst = ltst % 24
	198
[184]	199	### General plot settings
[186]	200	if numplot > 1:
	201	subplot(sub)
	202	found_lct = True
	203	elif numplot == 1:
	204	found_lct = True
[188]	205	### If only one local time is requested (numplot < 0)
[186]	206	elif numplot <= 0:
[201]	207	if int(ltst) + numplot != 0: continue
	208	else: found_lct = True
[184]	209
	210	### Map projection
[180]	211	m = define_proj(proj,wlon,wlat,back=back)
	212	x, y = m(lon2d, lat2d)
[184]	213
	214	#### Contour plot
	215	if var:
	216	if typefile in ['mesoapi','meso']: what_I_plot = dumpbdy(field[i,:,:])
	217	elif typefile in ['gcm']:
	218	if dimension == 2: what_I_plot = field[:,:]
	219	elif dimension == 3: what_I_plot = field[i,:,:]
[201]	220	if not tile:
	221	zelevels = np.linspace(zevmin,zevmax)
	222	contourf( x, y, what_I_plot, cmap = get_cmap(name=colorb), levels = zelevels )
	223	else:
	224	pcolor( x, y, what_I_plot, cmap = get_cmap(name=colorb), vmin=zevmin, vmax=zevmax )
	225	if colorb: colorbar(fraction=0.05,pad=0.1,format=fmtvar(var))
[184]	226
	227	### Vector plot
[191]	228	if winds:
	229	if typefile in ['mesoapi','meso']:
	230	[vecx,vecy] = [dumpbdy(u[i,nvert,:,:]), dumpbdy(v[i,nvert,:,:])]
	231	key = True
	232	elif typefile in ['gcm']:
	233	[vecx,vecy] = [ u[i,nvert,:,:] , v[i,nvert,:,:] ]
	234	key = False
	235	if metwind: [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
	236	if var == None and back == "vishires": colorvec = 'w'
	237	else: colorvec = 'k'
	238	vectorfield(vecx, vecy,\
[186]	239	x, y, stride=stride, csmooth=stride,\
[188]	240	scale=15., factor=300., color=colorvec, key=key)
	241	#200. ## or csmooth=2
[184]	242
	243	### Next subplot
[197]	244	plottitle = basename + "_LT"+str(ltst)
[195]	245	if addchar: plottitle = plottitle + addchar
	246	ptitle( plottitle )
[180]	247	sub += 1
[184]	248
	249	##########################################################################
	250	### Save the figure in a file in the data folder or an user-defined folder
[186]	251	if not target: zeplot = namefile+"_"+basename
	252	else: zeplot = target+"/"+basename
	253	if numplot <= 0: zeplot = zeplot + "_LT"+str(abs(numplot))
	254	if found_lct: makeplotpng(zeplot,pad_inches_value=0.35)
	255	else: print "Local time not found"
[193]	256	return zeplot
[180]	257
[184]	258
	259	####################################################
	260	####################################################
	261	### A simple program to get wind vectors' components
	262	def getwinds (nc,charu='Um',charv='Vm'):
[182]	263	import numpy as np
[180]	264	u = nc.variables[charu]
	265	v = nc.variables[charv]
[182]	266	if charu == 'U': u = u[:, :, :, 0:len(u[0,0,0,:])-1]
	267	if charv == 'V': v = v[:, :, 0:len(v[0,0,:,0])-1, :]
[184]	268	### ou alors prendre les coordonnees speciales
[180]	269	return u,v
	270
[184]	271
	272
	273	###########################################################################################
	274	###########################################################################################
	275	### What is below relate to running the file as a command line executable (very convenient)
[180]	276	if __name__ == "__main__":
	277	import sys
[184]	278	from optparse import OptionParser ### to be replaced by argparse
[186]	279	from api_wrapper import api_onelevel
[195]	280	from netCDF4 import Dataset
	281	from myplot import getlschar
[193]	282
	283	#############################
	284	### Get options and variables
[182]	285	parser = OptionParser()
[188]	286	parser.add_option('-f', action='store', dest='namefile', type="string", default=None, help='[NEEDED] name of WRF file')
	287	parser.add_option('-l', action='store', dest='nvert', type="float", default=0, help='vertical level (def=0)(-i 2: p,mbar)(-i 3,4: z,km)')
[184]	288	parser.add_option('-p', action='store', dest='proj', type="string", default=None, help='projection')
[182]	289	parser.add_option('-b', action='store', dest='back', type="string", default=None, help='background')
	290	parser.add_option('-t', action='store', dest='target', type="string", default=None, help='destination folder')
[186]	291	parser.add_option('-s', action='store', dest='stride', type="int", default=3, help='stride vectors (def=3)')
[184]	292	parser.add_option('-v', action='store', dest='var', type="string", default=None, help='variable contoured')
[187]	293	parser.add_option('-n', action='store', dest='numplot', type="int", default=4, help='number of plots (def=1)(<0: 1 plot of LT -numplot)')
[188]	294	parser.add_option('-i', action='store', dest='interp', type="int", default=None, help='interpolation method (2: press, 3: z-amr, 4:z-als)')
[201]	295	parser.add_option('-c', action='store', dest='colorb', type="string", default=True, help='change colormap (and draw a colorbar)')
[191]	296	parser.add_option('-x', action='store_false', dest='winds', default=True, help='flag: no wind vectors')
[199]	297	parser.add_option('-m', action='store', dest='vmin', type="float", default=None, help='bounding minimum value for color plot')
	298	parser.add_option('-M', action='store', dest='vmax', type="float", default=None, help='bounding maximum value for color plot')
[201]	299	parser.add_option('-T', action='store_true', dest='tile', default=False, help='draw a tiled plot')
	300	#parser.add_option('-V', action='store', dest='comb', type="float", default=None, help='a defined combination of variables')
[182]	301	(opt,args) = parser.parse_args()
	302	if opt.namefile is None:
	303	print "I want to eat one file at least ! Use winds.py -f name_of_my_file. Or type winds.py -h"
	304	exit()
	305	print "Options:", opt
[186]	306	zefile = opt.namefile
	307	zelevel = opt.nvert
[188]	308	stralt = None
[197]	309	[lschar,zehour,zehourin] = getlschar ( zefile ) ## getlschar from wrfout (or simply return "" if another file)
[193]	310
	311	#####################################################
	312	### Call Fortran routines for vertical interpolations
[186]	313	if opt.interp is not None:
[187]	314	if opt.nvert is 0 and opt.interp is 4: zelevel = 0.010
[191]	315	### winds or no winds
	316	if opt.winds : zefields = 'uvmet'
	317	else : zefields = ''
	318	### var or no var
	319	if opt.var is None : pass
	320	elif zefields == '' : zefields = opt.var
	321	else : zefields = zefields + "," + opt.var
	322	print zefields
[186]	323	zefile = api_onelevel ( path_to_input = '', \
	324	input_name = zefile, \
	325	path_to_output = opt.target, \
	326	fields = zefields, \
	327	interp_method = opt.interp, \
	328	onelevel = zelevel )
[188]	329	zelevel = 0 ## so that zelevel could play again the role of nvert
[184]	330
[193]	331	#############
	332	### Main call
	333	name = winds (zefile,int(zelevel),\
[195]	334	proj=opt.proj,back=opt.back,target=opt.target,stride=opt.stride,var=opt.var,numplot=opt.numplot,colorb=opt.colorb,winds=opt.winds,\
[201]	335	addchar=lschar,interv=[zehour,zehourin],vmin=opt.vmin,vmax=opt.vmax,tile=opt.tile)
[193]	336
	337	#########################################################
	338	### Generate a .sh file with the used command saved in it
	339	command = ""
	340	for arg in sys.argv: command = command + arg + ' '
	341	f = open(name+'.sh', 'w')
	342	f.write(command)

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

Download in other formats: