Context Navigation

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

Last change on this file since 195 was 195, checked in by aslmd, 14 years ago
MESOSCALE: graphics interface. corrected a bug in displaying automatic Ls
Property svn:executable set to ``*
File size: 14.1 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,\
[191]	18	colorb=None,\
[195]	19	winds=True,\
	20	addchar=None):
[184]	21
[192]	22	####################################################################################################################
	23	### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png
	24
[184]	25	#################################
	26	### Load librairies and functions
[180]	27	from netCDF4 import Dataset
[184]	28	from myplot import getcoord2d,define_proj,makeplotpng,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy
[194]	29	from mymath import deg
[187]	30	from matplotlib.pyplot import contourf, subplot, figure, rcParams, savefig, colorbar
[192]	31	from matplotlib.cm import get_cmap
[182]	32	import numpy as np
[184]	33
	34	######################
	35	### Load NETCDF object
[180]	36	nc = Dataset(namefile)
[184]	37
	38	###################################
	39	### Recognize predefined file types
	40	if 'controle' in nc.variables: typefile = 'gcm'
[191]	41	elif 'vert' in nc.variables: typefile = 'mesoapi'
[184]	42	elif 'U' in nc.variables: typefile = 'meso'
	43	else:
	44	print "typefile not supported."
[191]	45	print nc.variables
[182]	46	exit()
[184]	47
	48	##############################################################
	49	### Try to guess the projection from wrfout if not set by user
	50	if typefile in ['mesoapi','meso']:
	51	if proj == None: proj = getproj(nc)
	52	### (il faudrait passer CEN_LON dans la projection ?)
	53	elif typefile in ['gcm']:
	54	if proj == None: proj = "cyl"
	55	## pb avec les autres (de trace derriere la sphere ?)
	56
[188]	57	############################################
	58	#### Choose underlying topography by default
	59	if not back:
	60	if not var: back = "mola" ## if no var: draw mola
	61	elif typefile in ['mesoapi','meso'] \
	62	and proj not in ['merc','lcc']: back = "molabw" ## if var but meso: draw molabw
	63	else: pass ## else: draw None
[184]	64
	65	####################################################
	66	### Get geographical coordinates and plot boundaries
	67	if typefile in ['mesoapi','meso']:
	68	[lon2d,lat2d] = getcoord2d(nc)
	69	lon2d = dumpbdy(lon2d)
	70	lat2d = dumpbdy(lat2d)
	71	elif typefile in ['gcm']:
	72	[lon2d,lat2d] = getcoord2d(nc,nlat="latitude",nlon="longitude",is1d=True)
	73	if proj == "npstere": [wlon,wlat] = latinterv("North_Pole")
	74	elif proj == "lcc": [wlon,wlat] = wrfinterv(lon2d,lat2d)
	75	else: [wlon,wlat] = simplinterv(lon2d,lat2d)
	76
[195]	77	##################
	78	### Get local time
[194]	79	if typefile in ['mesoapi','meso']:
	80	ltst = int(getattr(nc, 'GMT') + 0.5*(wlon[0]+wlon[1])/15.)
[195]	81	#zetime = nc.variables['Times'][0]
	82	#zetimestart = getattr(nc, 'START_DATE')
	83	#zeday = int(zetime[8]+zetime[9]) - int(zetimestart[8]+zetimestart[9])
	84	#if zeday < 0: lschar="" ## might have crossed a month... fix soon
	85	#else: lschar="_Ls"+str( int( sol2ls ( getattr( nc, 'JULDAY' ) + zeday ) ) )
[194]	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,:,:]
	114
[191]	115	###########################
	116	### Get length of time axis
	117	if winds: nt = len(u[:,0,0,0])
	118	elif var:
	119	if dimension == 2: nt = 1
	120	else : nt = len(field[:,0,0])
	121
[186]	122	#########################################
	123	### Name for title and graphics save file
[191]	124	if winds: basename = "WINDS_"
	125	else: basename = ""
[186]	126	if var:
[191]	127	basename = basename + var
[188]	128	if typefile is 'meso': stralt = "_lvl" + str(nvert)
	129	elif typefile is 'mesoapi':
	130	zelevel = int(nc.variables['vert'][nvert])
	131	if 'altitude_abg' in nc.dimensions: stralt = "_"+str(zelevel)+"m-ALS"
	132	elif 'bottom_top' in nc.dimensions: stralt = "_"+str(zelevel)+"m-AMR"
	133	elif 'pressure' in nc.dimensions: stralt = "_"+str(zelevel)+"Pa"
	134	else: stralt = ""
	135	basename = basename + stralt
[186]	136
[184]	137	##################################
	138	### Open a figure and set subplots
[180]	139	fig = figure()
[186]	140	if numplot > 0:
	141	if numplot == 4:
	142	sub = 221
	143	fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
	144	rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
	145	elif numplot == 2:
	146	sub = 121
	147	fig.subplots_adjust(wspace = 0.3)
	148	rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
	149	elif numplot == 3:
	150	sub = 131
	151	fig.subplots_adjust(wspace = 0.3)
	152	rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
	153	elif numplot == 6:
	154	sub = 231
	155	fig.subplots_adjust(wspace = 0.4, hspace = 0.0)
	156	rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
	157	elif numplot == 8:
	158	sub = 331 #241
	159	fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
	160	rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
	161	elif numplot == 9:
	162	sub = 331
	163	fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
	164	rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
	165	elif numplot == 1:
[189]	166	sub = 99999
[186]	167	else:
	168	print "supported: 1,2,3,4,6,8"
	169	exit()
	170	### Prepare time loop
	171	if nt <= numplot or numplot == 1:
	172	tabrange = [0]
	173	numplot = 1
	174	else:
	175	tabrange = range(0,nt,int(nt/numplot)) #nt-1
	176	tabrange = tabrange[0:numplot]
	177	else:
	178	tabrange = range(0,nt,1)
	179	sub = 99999
	180	print tabrange
[184]	181
	182	#################################
	183	### Time loop for plotting device
[186]	184	found_lct = False
[184]	185	for i in tabrange:
	186
	187	### General plot settings
[186]	188	if numplot > 1:
	189	subplot(sub)
	190	found_lct = True
	191	elif numplot == 1:
	192	found_lct = True
[188]	193	### If only one local time is requested (numplot < 0)
[186]	194	elif numplot <= 0:
	195	if (ltst+i)%24 + numplot != 0: continue
	196	else: found_lct = True
[184]	197
	198	### Map projection
[180]	199	m = define_proj(proj,wlon,wlat,back=back)
	200	x, y = m(lon2d, lat2d)
[184]	201
	202	#### Contour plot
	203	if var:
	204	if typefile in ['mesoapi','meso']: what_I_plot = dumpbdy(field[i,:,:])
	205	elif typefile in ['gcm']:
	206	if dimension == 2: what_I_plot = field[:,:]
	207	elif dimension == 3: what_I_plot = field[i,:,:]
[192]	208	contourf( x, y, what_I_plot, 30, cmap = get_cmap(name=colorb) )
[187]	209	if colorb: colorbar(fraction=0.05,pad=0.1)
[184]	210
	211	### Vector plot
[191]	212	if winds:
	213	if typefile in ['mesoapi','meso']:
	214	[vecx,vecy] = [dumpbdy(u[i,nvert,:,:]), dumpbdy(v[i,nvert,:,:])]
	215	key = True
	216	elif typefile in ['gcm']:
	217	[vecx,vecy] = [ u[i,nvert,:,:] , v[i,nvert,:,:] ]
	218	key = False
	219	if metwind: [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
	220	if var == None and back == "vishires": colorvec = 'w'
	221	else: colorvec = 'k'
	222	vectorfield(vecx, vecy,\
[186]	223	x, y, stride=stride, csmooth=stride,\
[188]	224	scale=15., factor=300., color=colorvec, key=key)
	225	#200. ## or csmooth=2
[184]	226
	227	### Next subplot
[195]	228	plottitle = basename + "_LT"+str((ltst+i)%24)
	229	if addchar: plottitle = plottitle + addchar
	230	ptitle( plottitle )
[180]	231	sub += 1
[184]	232
	233	##########################################################################
	234	### Save the figure in a file in the data folder or an user-defined folder
[186]	235	if not target: zeplot = namefile+"_"+basename
	236	else: zeplot = target+"/"+basename
	237	if numplot <= 0: zeplot = zeplot + "_LT"+str(abs(numplot))
	238	if found_lct: makeplotpng(zeplot,pad_inches_value=0.35)
	239	else: print "Local time not found"
[193]	240	return zeplot
[180]	241
[184]	242
	243	####################################################
	244	####################################################
	245	### A simple program to get wind vectors' components
	246	def getwinds (nc,charu='Um',charv='Vm'):
[182]	247	import numpy as np
[180]	248	u = nc.variables[charu]
	249	v = nc.variables[charv]
[182]	250	if charu == 'U': u = u[:, :, :, 0:len(u[0,0,0,:])-1]
	251	if charv == 'V': v = v[:, :, 0:len(v[0,0,:,0])-1, :]
[184]	252	### ou alors prendre les coordonnees speciales
[180]	253	return u,v
	254
[184]	255
	256
	257	###########################################################################################
	258	###########################################################################################
	259	### What is below relate to running the file as a command line executable (very convenient)
[180]	260	if __name__ == "__main__":
	261	import sys
[184]	262	from optparse import OptionParser ### to be replaced by argparse
[186]	263	from api_wrapper import api_onelevel
[195]	264	from netCDF4 import Dataset
	265	from myplot import getlschar
[193]	266
	267	#############################
	268	### Get options and variables
[182]	269	parser = OptionParser()
[188]	270	parser.add_option('-f', action='store', dest='namefile', type="string", default=None, help='[NEEDED] name of WRF file')
	271	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]	272	parser.add_option('-p', action='store', dest='proj', type="string", default=None, help='projection')
[182]	273	parser.add_option('-b', action='store', dest='back', type="string", default=None, help='background')
	274	parser.add_option('-t', action='store', dest='target', type="string", default=None, help='destination folder')
[186]	275	parser.add_option('-s', action='store', dest='stride', type="int", default=3, help='stride vectors (def=3)')
[184]	276	parser.add_option('-v', action='store', dest='var', type="string", default=None, help='variable contoured')
[187]	277	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]	278	parser.add_option('-i', action='store', dest='interp', type="int", default=None, help='interpolation method (2: press, 3: z-amr, 4:z-als)')
[192]	279	parser.add_option('-c', action='store', dest='colorb', type="string", default=None, help='change colormap (and draw a colorbar)')
[191]	280	parser.add_option('-x', action='store_false', dest='winds', default=True, help='flag: no wind vectors')
[182]	281	(opt,args) = parser.parse_args()
	282	if opt.namefile is None:
	283	print "I want to eat one file at least ! Use winds.py -f name_of_my_file. Or type winds.py -h"
	284	exit()
	285	print "Options:", opt
[186]	286	zefile = opt.namefile
	287	zelevel = opt.nvert
[188]	288	stralt = None
[195]	289	lschar = getlschar ( zefile ) ## getlschar from wrfout (or simply return "" if another file)
[193]	290
	291	#####################################################
	292	### Call Fortran routines for vertical interpolations
[186]	293	if opt.interp is not None:
[187]	294	if opt.nvert is 0 and opt.interp is 4: zelevel = 0.010
[191]	295	### winds or no winds
	296	if opt.winds : zefields = 'uvmet'
	297	else : zefields = ''
	298	### var or no var
	299	if opt.var is None : pass
	300	elif zefields == '' : zefields = opt.var
	301	else : zefields = zefields + "," + opt.var
	302	print zefields
[186]	303	zefile = api_onelevel ( path_to_input = '', \
	304	input_name = zefile, \
	305	path_to_output = opt.target, \
	306	fields = zefields, \
	307	interp_method = opt.interp, \
	308	onelevel = zelevel )
[188]	309	zelevel = 0 ## so that zelevel could play again the role of nvert
[184]	310
[193]	311	#############
	312	### Main call
	313	name = winds (zefile,int(zelevel),\
[195]	314	proj=opt.proj,back=opt.back,target=opt.target,stride=opt.stride,var=opt.var,numplot=opt.numplot,colorb=opt.colorb,winds=opt.winds,\
	315	addchar=lschar)
[193]	316
	317	#########################################################
	318	### Generate a .sh file with the used command saved in it
	319	command = ""
	320	for arg in sys.argv: command = command + arg + ' '
	321	f = open(name+'.sh', 'w')
	322	f.write(command)

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

Download in other formats: