Context Navigation

source: trunk/UTIL/PYTHON/planetoplot_v2/ppplot.py @ 1196

Last change on this file since 1196 was 1159, checked in by aslmd, 11 years ago
planetoplot. corrected a /0 bug on steplon steplat
File size: 37.1 KB

Rev	Line
[910]	1	###############################################
	2	## PLANETOPLOT ##
	3	## --> PPPLOT ##
	4	###############################################
	5	## Author: Aymeric Spiga. 02-03/2013 ##
	6	###############################################
	7	# python built-in librairies
	8	import time as timelib
	9	# added librairies
	10	import numpy as np
	11	import matplotlib.pyplot as mpl
	12	from matplotlib.cm import get_cmap
	13	from mpl_toolkits.basemap import Basemap
[936]	14	from matplotlib.ticker import FormatStrFormatter,MaxNLocator
[910]	15	# personal librairies
	16	import ppcompute
	17	###############################################
	18
	19	#################################
	20	# global variables and settings #
	21	#################################
	22
	23	# matplotlib settings
	24	# http://matplotlib.org/users/customizing.html
	25	# -------------------------------
	26	mpl.rcParams['font.family'] = "serif"
[947]	27	mpl.rcParams['axes.color_cycle'] = "b,r,g,k"
[928]	28	mpl.rcParams['contour.negative_linestyle'] = "dashed" # ou "solid"
[910]	29	mpl.rcParams['verbose.level'] = "silent"
	30	mpl.rcParams['lines.linewidth'] = 1.5
	31	mpl.rcParams['lines.markersize'] = 10
	32	mpl.rcParams['xtick.major.pad'] = 10
	33	mpl.rcParams['ytick.major.pad'] = 10
	34
	35	# global variables
	36	# -------------------------------
	37	# - where settings files are located
	38	# (None means planetoplot_v2 in PYTHONPATH)
	39	whereset = None
	40	# - some good default settings.
	41	# (bounds)
	42	how_many_sigma = 3.0
	43	# (contours)
	44	ccol = 'black'
	45	cline = 0.55
[1117]	46	#cline = 0.8
[910]	47	# (vectors)
	48	widthvec = 0.002
	49	reducevec = 30.
	50	# (colorbar)
	51	zeorientation="vertical"
	52	zefrac = 0.05
	53	# (save figures)
	54	pad_inches_value=0.25
[928]	55	# (negative mode)
	56	def_negative = False
[1084]	57	# (xkcd mode)
	58	def_xkcd = False
[910]	59	###############################################
	60
[928]	61	### settings for 'negative-like' mode
	62	if def_negative:
	63	mpl.rc('figure', facecolor='k', edgecolor='k')
	64	mpl.rcParams['text.color'] = 'w'
	65	mpl.rc('axes',labelcolor='w',facecolor='k',edgecolor='w')
	66	mpl.rcParams['xtick.color'] = 'w'
	67	mpl.rcParams['ytick.color'] = 'w'
	68	mpl.rcParams['grid.color'] = 'w'
	69	mpl.rc('savefig',facecolor='k',edgecolor='k')
	70
[1084]	71	### settings for xkcd mode (only with matplotlib 1.3)
	72	### ... you must have Humori-Sans Font installed
	73	if def_xkcd:
	74	mpl.xkcd()
	75
[910]	76	##########################
	77	# executed when imported #
	78	##########################
	79	###########################################
	80	# we load user-defined automatic settings #
	81	###########################################
	82	# initialize the warning variable about file not present...
	83	files_not_present = ""
	84	# ... and the whereset variable
	85	whereset = ppcompute.findset(whereset)
	86
	87	# - variable settings
	88	# -------------------------------
	89	zefile = "set_var.txt"
[933]	90	vf = {} ; vc = {} ; vl = {} ; vu = {}
[910]	91	try:
	92	f = open(whereset+zefile, 'r')
	93	for line in f:
	94	if "#" in line: pass
	95	else:
[933]	96	var, format, colorb, label, units = line.strip().split(';')
[1093]	97	ind = var.strip().upper()
[910]	98	vf[ind] = format.strip()
	99	vc[ind] = colorb.strip()
	100	vl[ind] = label.strip()
[933]	101	vu[ind] = units.strip()
[910]	102	f.close()
	103	except IOError:
	104	files_not_present = files_not_present + zefile + " "
	105
	106	## - file settings
	107	## -------------------------------
	108	#zefile = "set_file.txt"
	109	#prefix_t = {} ; suffix_t = {} ; name_t = {} ; proj_t = {} ; vecx_t = {} ; vecy_t = {}
	110	#try:
	111	# f = open(whereset+zefile, 'r')
	112	# for line in f:
	113	# if "#" in line: pass
	114	# else:
	115	# prefix, suffix, name, proj, vecx, vecy = line.strip().split(';')
	116	# ind = name.strip()
	117	# prefix_t[ind] = prefix.strip()
	118	# suffix_t[ind] = suffix.strip()
	119	# #name_t[ind] = name.strip()
	120	# proj_t[ind] = proj.strip()
	121	# vecx_t[ind] = vecx.strip()
	122	# vecy_t[ind] = vecy.strip()
	123	# f.close()
	124	#except IOError:
	125	# files_not_present = files_not_present + zefile + " "
	126
	127	# - multiplot settings
	128	# -------------------------------
	129	zefile = "set_multiplot.txt"
	130	subv_t = {} ; subh_t = {} ; wspace_t = {} ; hspace_t = {} ; font_t = {}
	131	try:
	132	f = open(whereset+zefile, 'r')
	133	for line in f:
	134	if "#" in line: pass
	135	else:
	136	num, subv, subh, wspace, hspace, font = line.strip().split(';')
	137	ind = int(num.strip())
	138	subv_t[ind] = int(subv.strip())
	139	subh_t[ind] = int(subh.strip())
	140	wspace_t[ind] = float(wspace.strip())
	141	hspace_t[ind] = float(hspace.strip())
	142	font_t[ind] = float(font.strip())
	143	f.close()
	144	except IOError:
	145	files_not_present = files_not_present + zefile + " "
	146
	147	# - background settings
	148	# -------------------------------
	149	zefile = "set_back.txt"
	150	back = {}
	151	try:
	152	f = open(whereset+zefile, 'r')
	153	for line in f:
	154	if "#" in line: pass
	155	else:
	156	name, link = line.strip().split(';')
	157	ind = name.strip()
	158	back[ind] = link.strip()
	159	f.close()
	160	except IOError:
	161	files_not_present = files_not_present + zefile + " "
	162
	163	# - area settings
	164	# -------------------------------
	165	zefile = "set_area.txt"
	166	area = {}
	167	try:
	168	f = open(whereset+zefile, 'r')
	169	for line in f:
	170	if "#" in line: pass
	171	else:
	172	name, wlat1, wlat2, wlon1, wlon2 = line.strip().split(';')
	173	area[name.strip()] = [[float(wlon1.strip()),float(wlon2.strip())],\
	174	[float(wlat1.strip()),float(wlat2.strip())]]
	175	f.close()
	176	except IOError:
	177	files_not_present = files_not_present + zefile + " "
	178	# A note to the user about missing files
	179	if files_not_present != "":
	180	print "warning: files "+files_not_present+" not in "+whereset+" ; those presets will be missing"
	181
	182	# TBD: should change vector color with colormaps
	183	#"gist_heat": "white",\
	184	#"hot": "white",\
	185	#"gray": "red",\
	186
	187	####################
	188	# useful functions #
	189	####################
	190
[1007]	191	# continuity with matplotlib
	192	def close():
	193	mpl.close()
	194
	195	# a function for predefined figure sizes
	196	def figuref(x=16,y=9):
[1029]	197	fig = mpl.figure(figsize=(x,y))
	198	return fig
[1007]	199
[977]	200	# a function to change color lines according to a color map (idea by A. Pottier)
	201	# ... two optional arguments: color maps + a number telling how much intervals are needed
	202	def rainbow(cb="jet",num=8):
	203	ax = mpl.gca()
	204	pal = mpl.cm.get_cmap(name=cb)
	205	ax._get_lines.set_color_cycle([pal(i) for i in np.linspace(0,0.9,num)])
	206
[910]	207	# a function to define subplot
	208	# ... user can change settings in set_multiplot.txt read above
	209	# -------------------------------
	210	def definesubplot(numplot, fig):
	211	try:
	212	mpl.rcParams['font.size'] = font_t[numplot]
	213	except:
	214	mpl.rcParams['font.size'] = 18
	215	try:
	216	fig.subplots_adjust(wspace = wspace_t[numplot], hspace = hspace_t[numplot])
	217	subv, subh = subv_t[numplot],subh_t[numplot]
	218	except:
[1029]	219	print "!! WARNING !! (ignore if superpose mode) no preset found from set_multiplot.txt, or this setting file was not found."
[910]	220	subv = 1 ; subh = numplot
	221	return subv,subh
	222
	223	# a function to calculate automatically bounds (or simply prescribe those)
	224	# -------------------------------
	225	def calculate_bounds(field,vmin=None,vmax=None):
	226	# prescribed cases first
	227	zevmin = vmin
	228	zevmax = vmax
	229	# computed cases
	230	if zevmin is None or zevmax is None:
	231	# select values
	232	ind = np.where(field < 9e+35)
	233	fieldcalc = field[ ind ] # field must be a numpy array
	234	# calculate stdev and mean
	235	dev = np.std(fieldcalc)*how_many_sigma
	236	damean = ppcompute.mean(fieldcalc)
	237	# fill min/max if needed
	238	if vmin is None: zevmin = damean - dev
	239	if vmax is None: zevmax = damean + dev
	240	# special case: negative values with stddev while field is positive
[933]	241	if zevmin < 0. and ppcompute.min(fieldcalc) >= 0.: zevmin = 0.
[920]	242	# check that bounds are not too tight given the field
	243	amin = ppcompute.min(field)
	244	amax = ppcompute.max(field)
[936]	245	if np.abs(amin) < 1.e-15:
	246	cmin = 0.
	247	else:
	248	cmin = 100.*np.abs((amin - zevmin)/amin)
[920]	249	cmax = 100.*np.abs((amax - zevmax)/amax)
	250	if cmin > 150. or cmax > 150.:
	251	print "!! WARNING !! Bounds are a bit too tight. Might need to reconsider those."
	252	print "!! WARNING !! --> actual",amin,amax,"adopted",zevmin,zevmax
[910]	253	return zevmin, zevmax
	254	#### treat vmin = vmax for continuity
	255	#if vmin == vmax: zevmin = damean - dev ; zevmax = damean + dev
	256
	257	# a function to solve the problem with blank bounds !
	258	# -------------------------------
	259	def bounds(what_I_plot,zevmin,zevmax,miss=9e+35):
	260	small_enough = 1.e-7
	261	if zevmin < 0: what_I_plot[ what_I_plot < zevmin(1.-small_enough) ] = zevmin(1.-small_enough)
	262	else: what_I_plot[ what_I_plot < zevmin(1.+small_enough) ] = zevmin(1.+small_enough)
	263	what_I_plot[ what_I_plot > miss ] = -miss
	264	what_I_plot[ what_I_plot > zevmax ] = zevmax*(1.-small_enough)
	265	return what_I_plot
	266
[991]	267	# a function to change labels with modulo
	268	# ---------------------------------------
	269	def labelmodulo(ax,mod):
	270	mpl.draw()
	271	strtab = []
	272	for tick in ax.get_xaxis().get_ticklabels():
[1027]	273	onetick = tick.get_text()
	274	if len(onetick) > 0:
	275	num = float(onetick)
	276	strtab.append(num % mod)
[991]	277	ax.get_xaxis().set_ticklabels(strtab)
	278	return ax
	279
[990]	280	# a function to output an ascii file
	281	# ----------------------------------
	282	def writeascii (field=None,absc=None,name=None):
[1002]	283	field = np.array(field)
	284	absc = np.array(absc)
[990]	285	if name is None:
	286	name = "prof"
	287	for ttt in timelib.gmtime():
	288	name = name + "_" + str(ttt)
	289	if field is None or absc is None:
	290	print "!! WARNING !! Not printing the file, incorrect field or absc."
	291	else:
	292	if field.ndim == 1:
	293	myfile = open(name, 'w')
	294	for ix in range(len(absc)):
	295	myfile.write("%15.5e%15.5e\n" % ( absc[ix], field[ix] ))
	296	myfile.close()
	297	else:
	298	print "!! WARNING !! Not printing the file, 2D fields not supported yet."
	299	return
	300
[910]	301	# a generic function to show (GUI) or save a plot (PNG,EPS,PDF,...)
	302	# -------------------------------
[917]	303	def save(mode="gui",filename="plot",folder="./",includedate=True,res=150,custom=False):
[933]	304	if mode != "nothing":
	305	# a few settings
	306	possiblesave = ['eps','ps','svg','png','jpg','pdf']
	307	# now the main instructions
	308	if mode == "gui":
	309	mpl.show()
	310	elif mode in possiblesave:
	311	## name of plot
	312	name = folder+'/'+filename
	313	if includedate:
	314	for ttt in timelib.gmtime():
[1131]	315	name = name + "_" + "%03d" % (ttt)
[933]	316	name = name +"."+mode
	317	## save file
	318	print "**** Saving file in "+mode+" format... Please wait."
	319	if not custom:
	320	# ... regular plots
	321	mpl.savefig(name,dpi=res,pad_inches=pad_inches_value,bbox_inches='tight')
	322	else:
	323	# ... mapping mode, adapted space for labels etc...
[1131]	324	# TBD: not impacted by pad_inches_value. why?
[933]	325	mpl.savefig(name,dpi=res)
	326	else:
	327	print "!! ERROR !! File format not supported. Supported: ",possiblesave
[910]	328
[1050]	329	# a necessary addition for people used to matplotlib
	330	def show():
	331	mpl.show()
	332
[910]	333	##################################
	334	# a generic class to make a plot #
	335	##################################
	336	class plot():
	337
	338	# print out a help string when help is invoked on the object
	339	# -------------------------------
	340	def __repr__(self):
	341	whatprint = 'plot object. \"help(plot)\" for more information\n'
	342	return whatprint
	343
	344	# default settings
	345	# -- user can define settings by two methods.
	346	# -- 1. yeah = plot2d(title="foo")
	347	# -- 2. yeah = pp() ; yeah.title = "foo"
	348	# -------------------------------
	349	def __init__(self,\
	350	var=None,\
[1050]	351	f=None,\
	352	x=None,\
[910]	353	xlabel="",\
	354	ylabel="",\
	355	div=20,\
	356	logx=False,\
	357	logy=False,\
	358	swap=False,\
	359	swaplab=True,\
	360	invert=False,\
	361	xcoeff=1.,\
	362	ycoeff=1.,\
[933]	363	fmt=None,\
[1050]	364	colorbar="jet",\
[933]	365	units="",\
[991]	366	modx=None,\
[1004]	367	xmin=None,\
	368	ymin=None,\
	369	xmax=None,\
	370	ymax=None,\
[1029]	371	nxticks=10,\
	372	nyticks=10,\
[1084]	373	cbticks=None,\
[910]	374	title=""):
	375	## what could be defined by the user
	376	self.var = var
[1050]	377	self.f = f
	378	self.x = x
[910]	379	self.xlabel = xlabel
	380	self.ylabel = ylabel
	381	self.title = title
	382	self.div = div
	383	self.logx = logx
	384	self.logy = logy
	385	self.swap = swap
	386	self.swaplab = swaplab # NB: swaplab only used if swap=True
	387	self.invert = invert
	388	self.xcoeff = xcoeff
	389	self.ycoeff = ycoeff
[933]	390	self.fmt = fmt
	391	self.units = units
[1050]	392	self.colorbar = colorbar
[991]	393	self.modx = modx
[1004]	394	self.xmin = xmin
	395	self.ymin = ymin
	396	self.xmax = xmax
	397	self.ymax = ymax
[1029]	398	self.nxticks = nxticks
	399	self.nyticks = nyticks
[1084]	400	self.cbticks = cbticks
[910]	401	## other useful arguments
	402	## ... not used here in ppplot but need to be attached to plot object
	403	self.axisbg = "white"
	404	self.superpose = False
	405
	406	# check
	407	# -------------------------------
	408	def check(self):
[1050]	409	if self.f is None: print "!! ERROR !! Please define a field to be plotted" ; exit()
	410	else: self.f = np.array(self.f) # ensure this is a numpy array
[910]	411
	412	# define_from_var
	413	# ... this uses settings in set_var.txt
	414	# -------------------------------
	415	def define_from_var(self):
	416	if self.var is not None:
	417	if self.var.upper() in vl.keys():
	418	self.title = vl[self.var.upper()]
[933]	419	self.units = vu[self.var.upper()]
[910]	420
	421	# make
	422	# this is generic to all plots
	423	# -------------------------------
	424	def make(self):
	425	self.check()
[928]	426	# labels, title, etc...
[910]	427	mpl.xlabel(self.xlabel)
	428	mpl.ylabel(self.ylabel)
	429	if self.swap:
	430	if self.swaplab:
	431	mpl.xlabel(self.ylabel)
	432	mpl.ylabel(self.xlabel)
	433	mpl.title(self.title)
	434	# if masked array, set masked values to filled values (e.g. np.nan) for plotting purposes
[1050]	435	if type(self.f).__name__ in 'MaskedArray':
	436	self.f[self.f.mask] = self.f.fill_value
[910]	437
	438	################################
	439	# a subclass to make a 1D plot #
	440	################################
	441	class plot1d(plot):
	442
	443	# print out a help string when help is invoked on the object
	444	# -------------------------------
	445	def __repr__(self):
	446	whatprint = 'plot1d object. \"help(plot1d)\" for more information\n'
	447	return whatprint
	448
	449	# default settings
	450	# -- user can define settings by two methods.
	451	# -- 1. yeah = plot1d(title="foo")
	452	# -- 2. yeah = pp() ; yeah.title = "foo"
	453	# -------------------------------
	454	def __init__(self,\
[1050]	455	linestyle=None,\
[947]	456	color=None,\
[923]	457	marker='x',\
[1050]	458	legend=None):
[910]	459	## get initialization from parent class
	460	plot.__init__(self)
	461	## what could be defined by the user
[1050]	462	self.linestyle = linestyle
[910]	463	self.color = color
	464	self.marker = marker
[1050]	465	self.legend = legend
[910]	466
	467	# define_from_var
	468	# ... this uses settings in set_var.txt
	469	# -------------------------------
	470	def define_from_var(self):
	471	# get what is done in the parent class
	472	plot.define_from_var(self)
	473	# add specific stuff
	474	if self.var is not None:
	475	if self.var.upper() in vl.keys():
[943]	476	self.ylabel = vl[self.var.upper()] + " (" + vu[self.var.upper()] + ")"
[910]	477	self.title = ""
[933]	478	self.fmt = vf[self.var.upper()]
[910]	479
	480	# make
	481	# -------------------------------
	482	def make(self):
	483	# get what is done in the parent class
	484	plot.make(self)
[933]	485	if self.fmt is None: self.fmt = '%.0f'
[910]	486	# add specific stuff
	487	mpl.grid(color='grey')
[1050]	488	if self.linestyle == "": self.linestyle = " " # to allow for no line at all with ""
[977]	489	# set dummy x axis if not defined
[1050]	490	if self.x is None:
	491	self.x = np.array(range(self.f.shape[0]))
[977]	492	print "!! WARNING !! dummy coordinates on x axis"
[1002]	493	else:
[1050]	494	self.x = np.array(self.x) # ensure this is a numpy array
[910]	495	# swapping if requested
[1050]	496	if self.swap: x = self.f ; y = self.x
	497	else: x = self.x ; y = self.f
[910]	498	# coefficients on axis
	499	x=xself.xcoeff ; y=yself.ycoeff
	500	# check axis
	501	if x.size != y.size:
	502	print "!! ERROR !! x and y sizes don't match on 1D plot.", x.size, y.size
	503	exit()
	504	# make the 1D plot
	505	# either request linestyle or let matplotlib decide
[1050]	506	if self.linestyle is not None and self.color is not None:
	507	mpl.plot(x,y,self.color+self.linestyle,marker=self.marker,label=self.legend)
[942]	508	elif self.color is not None:
[1050]	509	mpl.plot(x,y,color=self.color,marker=self.marker,label=self.legend)
	510	elif self.linestyle is not None:
	511	mpl.plot(x,y,linestyle=self.linestyle,marker=self.marker,label=self.legend)
[910]	512	else:
[1050]	513	mpl.plot(x,y,marker=self.marker,label=self.legend)
[1027]	514	# AXES
	515	ax = mpl.gca()
[910]	516	# make log axes and/or invert ordinate
	517	# ... this must be after plot so that axis bounds are well-defined
	518	# ... also inverting must be after making the thing logarithmic
[977]	519	if self.logx: mpl.xscale("log") # not mpl.semilogx() because excludes log on y
	520	if self.logy: mpl.yscale("log") # not mpl.semilogy() because excludes log on x
[1027]	521	if self.invert: ax.set_ylim(ax.get_ylim()[::-1])
	522	if self.xmin is not None and self.xmax is not None:
	523	if self.xmin > self.xmax:
	524	ax.set_xlim(ax.get_xlim()[::-1])
	525	self.xmin,self.xmax = self.xmax,self.xmin
[923]	526	# add a label for line(s)
[1050]	527	if self.legend is not None:
	528	if self.legend != "":
[923]	529	mpl.legend(loc="best",fancybox=True)
[933]	530	# format labels
[936]	531	if self.swap: ax.xaxis.set_major_formatter(FormatStrFormatter(self.fmt))
	532	else: ax.yaxis.set_major_formatter(FormatStrFormatter(self.fmt))
[933]	533	# plot limits: ensure that no curve would be outside the window
	534	# x-axis
	535	x1, x2 = ax.get_xbound()
	536	xmin = ppcompute.min(x)
	537	xmax = ppcompute.max(x)
	538	if xmin < x1: x1 = xmin
	539	if xmax > x2: x2 = xmax
[977]	540	if self.xmin is not None: x1 = self.xmin
	541	if self.xmax is not None: x2 = self.xmax
[933]	542	ax.set_xbound(lower=x1,upper=x2)
	543	# y-axis
	544	y1, y2 = ax.get_ybound()
	545	ymin = ppcompute.min(y)
	546	ymax = ppcompute.max(y)
	547	if ymin < y1: y1 = ymin
	548	if ymax > y2: y2 = ymax
[977]	549	if self.ymin is not None: y1 = self.ymin
	550	if self.ymax is not None: y2 = self.ymax
[933]	551	ax.set_ybound(lower=y1,upper=y2)
[1084]	552	## set with .div the number of ticks.
[981]	553	if not self.logx:
[1084]	554	ax.xaxis.set_major_locator(MaxNLocator(self.nxticks))
[981]	555	else:
	556	print "!! WARNING. in logx mode, ticks are set automatically."
[1084]	557	if not self.logy:
	558	ax.yaxis.set_major_locator(MaxNLocator(self.nyticks))
	559	else:
	560	print "!! WARNING. in logy mode, ticks are set automatically."
[991]	561	## specific modulo labels
	562	if self.modx is not None:
	563	ax = labelmodulo(ax,self.modx)
[910]	564
[1050]	565	# makeshow = make + show
	566	# -------------------------------
	567	def makeshow(self):
	568	self.make()
	569	mpl.show()
	570
[910]	571	################################
	572	# a subclass to make a 2D plot #
	573	################################
	574	class plot2d(plot):
	575
	576	# print out a help string when help is invoked on the object
	577	# -------------------------------
	578	def __repr__(self):
	579	whatprint = 'plot2d object. \"help(plot2d)\" for more information\n'
	580	return whatprint
	581
	582	# default settings
	583	# -- user can define settings by two methods.
	584	# -- 1. yeah = plot2d(title="foo")
	585	# -- 2. yeah = pp() ; yeah.title = "foo"
	586	# -------------------------------
	587	def __init__(self,\
[1050]	588	y=None,\
[910]	589	mapmode=False,\
	590	proj="cyl",\
	591	back=None,\
	592	trans=1.0,\
[1050]	593	vx=None,\
	594	vy=None,\
	595	c=None,\
[910]	596	blon=None,\
	597	blat=None,\
	598	area=None,\
	599	vmin=None,\
	600	vmax=None,\
[964]	601	showcb=True,\
[972]	602	wscale=None,\
[1050]	603	svx=1,\
	604	svy=1,\
[1040]	605	leftcorrect=False,\
[910]	606	colorvec="black"):
	607	## get initialization from parent class
	608	plot.__init__(self)
	609	## what could be defined by the user
[1050]	610	self.y = y
[910]	611	self.mapmode = mapmode
	612	self.proj = proj
	613	self.back = back
	614	self.trans = trans
[1050]	615	self.vx = vx
	616	self.vy = vy
[910]	617	self.colorvec = colorvec
[1050]	618	self.c = c
[910]	619	self.blon = blon ; self.blat = blat
	620	self.area = area
	621	self.vmin = vmin ; self.vmax = vmax
[964]	622	self.showcb = showcb
[972]	623	self.wscale = wscale
[1050]	624	self.svx = svx
	625	self.svy = svy
[1040]	626	self.leftcorrect = leftcorrect
[910]	627
	628	# define_from_var
	629	# ... this uses settings in set_var.txt
	630	# -------------------------------
	631	def define_from_var(self):
	632	# get what is done in the parent class
	633	plot.define_from_var(self)
	634	# add specific stuff
	635	if self.var is not None:
	636	if self.var.upper() in vl.keys():
[1050]	637	self.colorbar = vc[self.var.upper()]
[910]	638	self.fmt = vf[self.var.upper()]
	639
	640	# make
	641	# -------------------------------
	642	def make(self):
	643	# get what is done in the parent class...
	644	plot.make(self)
[977]	645	if self.fmt is None: self.fmt = "%.1e"
[910]	646	# ... then add specific stuff
	647	############################################################################################
	648	### PRE-SETTINGS
	649	############################################################################################
[960]	650	# set dummy xy axis if not defined
[1050]	651	if self.x is None:
	652	self.x = np.array(range(self.f.shape[0]))
[960]	653	self.mapmode = False
	654	print "!! WARNING !! dummy coordinates on x axis"
[1050]	655	if self.y is None:
	656	self.y = np.array(range(self.f.shape[1]))
[960]	657	self.mapmode = False
	658	print "!! WARNING !! dummy coordinates on y axis"
[1064]	659	# check sizes
[1065]	660	if self.c is not None:
	661	if self.c.ndim != 2:
	662	print "!! WARNING !! Contour is not a 2D field. No contour.",self.c.ndim
	663	self.c = None
	664	if self.f.ndim != 2:
[1064]	665	print "!! ERROR !! Field is not two-dimensional" ; exit()
[977]	666	# transposing if necessary
[1050]	667	shape = self.f.shape
[977]	668	if shape[0] != shape[1]:
[1050]	669	if len(self.x) == shape[0] and len(self.y) == shape[1]:
[977]	670	print "!! WARNING !! Transposing axes"
[1050]	671	self.f = np.transpose(self.f)
[910]	672	# bound field
[1050]	673	zevmin, zevmax = calculate_bounds(self.f,vmin=self.vmin,vmax=self.vmax)
	674	what_I_plot = bounds(self.f,zevmin,zevmax)
[910]	675	# define contour field levels. define color palette
	676	ticks = self.div + 1
	677	zelevels = np.linspace(zevmin,zevmax,ticks)
[1050]	678	palette = get_cmap(name=self.colorbar)
[910]	679	# do the same thing for possible contourline entries
[1050]	680	if self.c is not None:
[910]	681	# if masked array, set masked values to filled values (e.g. np.nan) for plotting purposes
[1050]	682	if type(self.c).__name__ in 'MaskedArray':
	683	self.c[self.c.mask] = self.c.fill_value
	684	zevminc, zevmaxc = calculate_bounds(self.c)
	685	what_I_contour = bounds(self.c,zevminc,zevmaxc)
[910]	686	ticks = self.div + 1
	687	zelevelsc = np.linspace(zevminc,zevmaxc,ticks)
	688	############################################################################################
	689	### MAIN PLOT
	690	### NB: contour lines are done before contour shades otherwise colorar error
	691	############################################################################################
	692	if not self.mapmode:
	693	## A SIMPLE 2D PLOT
	694	###################
	695	# swapping if requested
[1050]	696	if self.swap: x = self.y ; y = self.x
	697	else: x = self.x ; y = self.y
[910]	698	# coefficients on axis
	699	x=xself.xcoeff ; y=yself.ycoeff
	700	# make shaded and line contours
[1050]	701	if self.c is not None:
[943]	702	objC = mpl.contour(x, y, what_I_contour, \
[910]	703	zelevelsc, colors = ccol, linewidths = cline)
[1007]	704	#mpl.clabel(objC, inline=1, fontsize="small",\
	705	# inline_spacing=1,fmt="%.0f")
[910]	706	mpl.contourf(x, y, \
[1050]	707	self.f, \
[910]	708	zelevels, cmap=palette)
[1007]	709	#mpl.pcolor(x,y,\
[1050]	710	# self.f, \
[1007]	711	# cmap=palette)
[910]	712	# make log axes and/or invert ordinate
[1004]	713	ax = mpl.gca()
[910]	714	if self.logx: mpl.semilogx()
	715	if self.logy: mpl.semilogy()
[1004]	716	if self.invert: ax.set_ylim(ax.get_ylim()[::-1])
[1027]	717	if self.xmin is not None and self.xmax is not None:
	718	if self.xmin > self.xmax:
	719	ax.set_xlim(ax.get_xlim()[::-1])
	720	self.xmin,self.xmax = self.xmax,self.xmin
[1004]	721	if self.xmin is not None: ax.set_xbound(lower=self.xmin)
	722	if self.xmax is not None: ax.set_xbound(upper=self.xmax)
	723	if self.ymin is not None: ax.set_ybound(lower=self.ymin)
	724	if self.ymax is not None: ax.set_ybound(upper=self.ymax)
[1029]	725	# set the number of ticks
	726	if not self.logx:
	727	ax.xaxis.set_major_locator(MaxNLocator(self.nxticks))
	728	else:
	729	print "!! WARNING. in logx mode, ticks are set automatically."
[1084]	730	if not self.logy:
	731	ax.yaxis.set_major_locator(MaxNLocator(self.nyticks))
	732	else:
	733	print "!! WARNING. in logy mode, ticks are set automatically."
[1027]	734	## specific modulo labels
	735	if self.modx is not None:
	736	ax = labelmodulo(ax,self.modx)
[910]	737	else:
	738	## A 2D MAP USING PROJECTIONS (basemap)
	739	#######################################
	740	mpl.xlabel("") ; mpl.ylabel("")
	741	# additional security in case self.proj is None here
	742	# ... we set cylindrical projection (the simplest one)
	743	if self.proj is None: self.proj = "cyl"
	744	# get lon and lat in 2D version.
	745	# (but first ensure we do have 2D coordinates)
[1050]	746	if self.x.ndim == 1: [self.x,self.y] = np.meshgrid(self.x,self.y)
	747	elif self.x.ndim > 2: print "!! ERROR !! lon and lat arrays must be 1D or 2D"
[910]	748	# get lat lon intervals and associated settings
[1050]	749	wlon = [np.min(self.x),np.max(self.x)]
	750	wlat = [np.min(self.y),np.max(self.y)]
[910]	751	# -- area presets are in set_area.txt
	752	if self.area is not None:
	753	if self.area in area.keys():
	754	wlon, wlat = area[self.area]
	755	# -- settings for meridians and parallels
[936]	756	steplon = int(abs(wlon[1]-wlon[0])/6.)
	757	steplat = int(abs(wlat[1]-wlat[0])/3.)
	758	#mertab = np.r_[wlon[0]:wlon[1]:steplon] ; merlab = [0,0,0,1]
	759	#partab = np.r_[wlat[0]:wlat[1]:steplat] ; parlab = [1,0,0,0]
[1159]	760	if steplon < 1: steplon = 1
	761	if steplat < 1: steplat = 1
[1091]	762	if np.abs(wlon[0]) < 180.1 and np.abs(wlon[1]) < 180.1:
	763	mertab = np.r_[-180.:180.:steplon]
	764	else:
	765	mertab = np.r_[0.:360.:steplon]
	766	merlab = [0,0,0,1]
[1131]	767	partab = np.r_[-90.:90.+steplat:steplat] ; parlab = [1,0,0,0]
[910]	768	format = '%.1f'
	769	# -- center of domain and bounding lats
	770	lon_0 = 0.5*(wlon[0]+wlon[1])
	771	lat_0 = 0.5*(wlat[0]+wlat[1])
	772	# some tests, bug fixes, and good-looking settings
	773	# ... cyl is good for global and regional
	774	if self.proj == "cyl":
[917]	775	format = '%.0f'
[910]	776	# ... global projections
	777	elif self.proj in ["ortho","moll","robin"]:
	778	wlat[0] = None ; wlat[1] = None ; wlon[0] = None ; wlon[1] = None
[1091]	779	lon_0 = np.ceil(lon_0) # reverse map if lon_0 is slightly below 180 with [0,360]
[910]	780	steplon = 30. ; steplat = 30.
[1062]	781	if self.proj in ["moll"]: steplon = 60.
	782	if self.proj in ["robin"]: steplon = 90.
[910]	783	mertab = np.r_[-360.:360.:steplon]
[1131]	784	partab = np.r_[-90.:90.+steplat:steplat]
[910]	785	if self.proj == "ortho":
	786	merlab = [0,0,0,0] ; parlab = [0,0,0,0]
	787	# in ortho projection, blon and blat can be used to set map center
	788	if self.blon is not None: lon_0 = self.blon
	789	if self.blat is not None: lat_0 = self.blat
	790	elif self.proj == "moll":
	791	merlab = [0,0,0,0]
	792	format = '%.0f'
	793	# ... regional projections
	794	elif self.proj in ["lcc","laea","merc"]:
	795	if self.proj in ["lcc","laea"] and wlat[0] == -wlat[1]:
	796	print "!! ERROR !! with Lambert lat1 must be different than lat2" ; exit()
	797	if wlat[0] < -80. and wlat[1] > 80.:
	798	print "!! ERROR !! set an area (not global)" ; exit()
	799	format = '%.0f'
	800	elif self.proj in ["npstere","spstere"]:
	801	# in polar projections, blat gives the bounding lat
	802	# if not set, set something reasonable
	803	if self.blat is None: self.blat = 60.
	804	# help the user who forgets self.blat would better be negative in spstere
	805	# (this actually serves for the default setting just above)
	806	if self.proj == "spstere" and self.blat > 0: self.blat = -self.blat
[932]	807	# labels
	808	mertab = np.r_[-360.:360.:15.]
	809	partab = np.r_[-90.:90.:5.]
[910]	810	# ... unsupported projections
	811	else:
	812	print "!! ERROR !! unsupported projection. supported: "+\
	813	"cyl, npstere, spstere, ortho, moll, robin, lcc, laea, merc"
	814	# finally define projection
	815	m = Basemap(projection=self.proj,\
	816	lat_0=lat_0,lon_0=lon_0,\
	817	boundinglat=self.blat,\
	818	llcrnrlat=wlat[0],urcrnrlat=wlat[1],\
	819	llcrnrlon=wlon[0],urcrnrlon=wlon[1])
[1040]	820	# in some case need to translated to the left for colorbar + labels
	821	# TBD: break stuff. a better solution should be found.
	822	if self.leftcorrect:
[1039]	823	ax = mpl.gca()
	824	pos = ax.get_position().bounds
	825	newpos = [0.,pos[1],pos[2],pos[3]]
	826	ax.set_position(newpos)
[910]	827	# draw meridians and parallels
	828	ft = int(mpl.rcParams['font.size']*3./4.)
[932]	829	zelatmax = 85.
	830	m.drawmeridians(mertab,labels=merlab,color='grey',linewidth=0.75,fontsize=ft,fmt=format,latmax=zelatmax)
	831	m.drawparallels(partab,labels=parlab,color='grey',linewidth=0.75,fontsize=ft,fmt=format,latmax=zelatmax)
[910]	832	# define background (see set_back.txt)
	833	if self.back is not None:
	834	if self.back in back.keys():
	835	print "**** info: loading a background, please wait.",self.back
[1029]	836	if self.back not in ["coast","sea"]:
	837	try: m.warpimage(back[self.back],scale=0.75)
	838	except: print "!! ERROR !! no background image could be loaded. probably not connected to the internet?"
	839	elif self.back == "coast":
	840	m.drawcoastlines()
	841	elif self.back == "sea":
	842	m.drawlsmask(land_color='white',ocean_color='aqua')
[910]	843	else:
	844	print "!! ERROR !! requested background not defined. change name or fill in set_back.txt" ; exit()
	845	# define x and y given the projection
[1050]	846	x, y = m(self.x, self.y)
[910]	847	# contour field. first line contour then shaded contour.
[1050]	848	if self.c is not None:
[943]	849	objC2 = m.contour(x, y, what_I_contour, \
[910]	850	zelevelsc, colors = ccol, linewidths = cline)
[1054]	851	#mpl.clabel(objC2, inline=1, fontsize=10,manual=True,fmt='-%2.0f$^{\circ}$C',colors='r')
[1062]	852	m.contourf(x, y, what_I_plot, zelevels, cmap = palette, alpha = self.trans, antialiased=True)
[910]	853	############################################################################################
	854	### COLORBAR
	855	############################################################################################
[964]	856	if self.trans > 0. and self.showcb:
[910]	857	## draw colorbar. settings are different with projections. or if not mapmode.
[933]	858	if not self.mapmode: orientation=zeorientation ; frac = 0.075 ; pad = 0.03 ; lu = 0.5
[942]	859	elif self.proj in ['moll']: orientation="horizontal" ; frac = 0.08 ; pad = 0.03 ; lu = 1.0
[933]	860	elif self.proj in ['cyl']: orientation="vertical" ; frac = 0.023 ; pad = 0.03 ; lu = 0.5
	861	else: orientation = zeorientation ; frac = zefrac ; pad = 0.03 ; lu = 0.5
[1084]	862	if self.cbticks is None:
	863	self.cbticks = min([ticks/2+1,21])
	864	zelevpal = np.linspace(zevmin,zevmax,num=self.cbticks)
[917]	865	zecb = mpl.colorbar(fraction=frac,pad=pad,\
[910]	866	format=self.fmt,orientation=orientation,\
	867	ticks=zelevpal,\
[933]	868	extend='neither',spacing='proportional')
[910]	869	if zeorientation == "horizontal": zecb.ax.set_xlabel(self.title) ; self.title = ""
[933]	870	# colorbar title --> units
	871	if self.units not in ["dimless",""]:
	872	zecb.ax.set_title("["+self.units+"]",fontsize=3.*mpl.rcParams['font.size']/4.,x=lu,y=1.025)
	873
[910]	874	############################################################################################
	875	### VECTORS. must be after the colorbar. we could also leave possibility for streamlines.
	876	############################################################################################
[1050]	877	### not expecting NaN in self.vx and self.vy. masked arrays is just enough.
	878	if self.vx is not None and self.vy is not None:
[982]	879	# vectors on map projection or simple 2D mapping
[1002]	880	if self.mapmode:
[1050]	881	[vecx,vecy] = m.rotate_vector(self.vx,self.vy,self.x,self.y) # for metwinds only ?
[1002]	882	else:
[1050]	883	vecx,vecy = self.vx,self.vy
[1002]	884	if x.ndim < 2 and y.ndim < 2: x,y = np.meshgrid(x,y)
[910]	885	# reference vector is scaled
[1159]	886	if self.wscale is None:
[1050]	887	self.wscale = ppcompute.mean(np.sqrt(self.vxself.vx+self.vyself.vy))
[910]	888	# make vector field
[981]	889	if self.mapmode:
[1050]	890	q = m.quiver( x[::self.svy,::self.svx],y[::self.svy,::self.svx],\
	891	vecx[::self.svy,::self.svx],vecy[::self.svy,::self.svx],\
[981]	892	angles='xy',color=self.colorvec,pivot='tail',\
	893	scale=self.wscale*reducevec,width=widthvec )
	894	else:
[1050]	895	q = mpl.quiver( x[::self.svy,::self.svx],y[::self.svy,::self.svx],\
	896	vecx[::self.svy,::self.svx],vecy[::self.svy,::self.svx],\
[981]	897	angles='xy',color=self.colorvec,pivot='tail',\
	898	scale=self.wscale*reducevec,width=widthvec )
[930]	899	# make vector key.
[933]	900	#keyh = 1.025 ; keyv = 1.05 # upper right corner over colorbar
[972]	901	keyh = 0.97 ; keyv = 1.05
[930]	902	#keyh = -0.03 ; keyv = 1.08 # upper left corner
[923]	903	p = mpl.quiverkey(q,keyh,keyv,\
[972]	904	self.wscale,str(int(self.wscale)),\
	905	fontproperties={'size': 'small'},\
[910]	906	color='black',labelpos='S',labelsep = 0.07)
[942]	907	############################################################################################
	908	### TEXT. ANYWHERE. add_text.txt should be present with lines x ; y ; text ; color
	909	############################################################################################
	910	try:
	911	f = open("add_text.txt", 'r')
	912	for line in f:
	913	if "#" in line: pass
	914	else:
	915	userx, usery, usert, userc = line.strip().split(';')
	916	userc = userc.strip()
	917	usert = usert.strip()
	918	userx = float(userx.strip())
	919	usery = float(usery.strip())
	920	if self.mapmode: userx,usery = m(userx,usery)
	921	mpl.text(userx,usery,usert,\
	922	color = userc,\
	923	horizontalalignment='center',\
	924	verticalalignment='center')
	925	f.close()
	926	except IOError:
	927	pass
[1050]	928
	929	# makeshow = make + show
	930	# -------------------------------
	931	def makeshow(self):
	932	self.make()
	933	mpl.show()

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

Download in other formats: