Context Navigation

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

Last change on this file since 205 was 204, checked in by aslmd, 13 years ago
MESOSCALE: python: a small late commit to add zooming capabilities.
Property svn:executable set to ``*
File size: 16.1 KB

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

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

Download in other formats: