''' Plotting functions for emma glue results'''
from gemmaclass import gemma
from sampleclass import sample
from gemmafiles import read_gemma_stream
import numpy as np
from scipy import stats
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.ticker import MaxNLocator
import cPickle as pickle
import pandas
try:
import petq_hmm as plot_data
except:
pass
import datetime as dt
import gc
import string
mpl.rcParams['mathtext.default'] = 'regular'
colcol = {'INLET': 'b','GWAQD':'darkorange','GWPH':'g','PRECIP':'lightblue',
'GWAQS':'yellow','GWDB':'yellow','BOIL':'darkorange'}
colbl = {'INLET': 'k','GWAQD':'k','GWPH':'k','PRECIP':'k',
'GWAQS':'k','GWDB':'k','BOIL':'k'}
#labels = {'INLET': 'Inlet water','GWAQD':'Deeper aquifer groundwater','GWPH':'Soil water / phreatic groundwater',
# 'PRECIP':'Precipitation','GWAQS':'Shallow aquifer groundwater','GWDB':'Groundwater at bottom of ditch',
# 'BOIL':'Boil discharge'}
_labels = {'INLET': 'IL','GWAQD':'AD','GWPH':'SL','PRECIP':'PR','GWAQS':'AS','GWDB':'BD',
'BOIL':'BL', 'MIX':'Stream', 'SO4':'SO_4^{2-}'}
def _constr_box(plot_petq, separate, params, pad=.05):
# box dimensions:
bot = .1
top = .95
left = .05
right = .95
height=[]
if separate:
if plot_petq:
height = [1.3]
height+=[1. for x in params]
height = np.array(height)
height/= float(height.sum())
else:
if plot_petq:
height= np.array([.35,.65])
else:
height= np.array([1.])
height = height*(top-bot)
hcs = height.cumsum()
for irow in xrange(len(height)):
yield [left,top-hcs[irow],right-left,(1.-pad)*height[irow]]
def _plot_text(ax, s, **kwargs):
props = dict(facecolor='w',edgecolor='k',pad=10,zorder=100, fill=True)
defkwargs = {'x': 0.01,
'y': 0.95,
'fontsize': 16,
'bbox':props,
'horizontalalignment': 'left',
'verticalalignment': 'top'}
kwargs = dict(defkwargs.items() + kwargs.items())
ax.text(s=s, transform=ax.transAxes, zorder=110, **kwargs)
[docs]def create_colblack(params):
'''Create color dictionary for use in params_time for a black/white plot
Parameters
----------
params : list of parameters to color
Returns
-------
col : dict of params -> colors (in this case black)
'''
col = {}
for p in params:
col[p] = 'k'
return col
[docs]def create_colcolor(params, dcol=colcol):
'''Create color dictionary for use in params_time for a color plot
Solutes part of end-members receive color of end-member
Parameters
----------
params : list of parameters to color
dcol : dict of end-member or solute names to colors
Returns
-------
col : dict of params -> colors
'''
col = {}
for p in params:
if '_' in p:
e = p[:p.index('_')]
else: e = p
col[p] = dcol[e]
return col
[docs]def create_labels(params, dlabel=_labels):
'''Create label dictionary for use in params_time
Solutes part of end-members receive label of end-member - solute
Parameters
----------
params : list of parameters to label
dlabel : dict of parameters -> labels
Returns
-------
col : dict of params -> colors
'''
lab = {}
for p in params:
if '_' in p:
e = p[:p.index('_')]
if e in dlabel:
e = dlabel[e]
s = p[p.index('_')+1:]
if s in dlabel:
s2 = dlabel[s]
else:
s2 = ''
for i,x in enumerate(s):
try:
t=int(x)
s2+='_'+x
except ValueError:
s2+=x
lab[p] = e+' - '+r'$'+s2+'$'
else:
if p in _labels:
lab[p] = dlabel[p]
else:
lab[p] = p
return lab
[docs]def params_time(params=None, pct=(0,50,100),result='gemma_output.bin',
separate=True, stacked=False, calcload=False, hiatuses=[],
plot_peq=True, fname=None, xlim=None, ylim=None,
col=None, labels=None, labelstr='fraction', unitstr='-',
f_conv=1., returnfig=False, figsize=None):
'''Workhorse function for plotting a lineplot of parameter values over time,
including uncertainty
Parameters
----------
params : list of parameters to plot
pct : list of three or more percentiles: lower bound, median value, upper bound. If more than three are given, additional uncertainty bands are drawn
result : g-emma result, either filename or gemma object
separate : bool, plot separate plots for each param (default True)
stacked : bool, stack plots, eg when contributions are additional (default: False)
calcload : bool, calculate loads for solutes (default False)
hiatuses : list of lists of two datetime objects. Draws vertical grey bands to exclude period from figure
plot_peq : bool, plot precip, evap and discharge data (default True). Assumes a module+function plot_data.create_petq(afig=fig, box=box.next())
fname : filename to write plot to. If None (default), output to screen
xlim : list of limits in x (default None)
ylim : list of limits in y (default None)
col : dictionary of colors for each param
labels : dictionary of labels for each param,
labelstr : string, additional part of label, denoting figure content (default 'fraction')
unitstr : string, additional part of label, denoting unit of data (default '-')
f_conv : float, conversion factor for load calculation (default 1.)
returnfig : bool, return figure object and list of axes
figsize : size of figure
'''
if pct[-1] > 1:
pct = [x/100. for x in pct]
if separate:
stacked = False
if col == None: col = create_colblack(params)
if labels == None: labels = create_labels(params)
if isinstance(ylim,(tuple,list)):
ylim=np.atleast_2d(ylim)
if len(ylim) < len(params):
ylim = np.tile(ylim, (len(params),1))
if not isinstance(result, gemma):
result = gemma(result)
data = result.percentiles(pct, params).copy() # use local copy
x = data['Date']
# if stacked in one subplot: add values to lower in order
if not separate and stacked:
prev = None
for i,e in enumerate(params):
if prev != None:
for p in pct:
data[e,p] = data[e,p] + data[params[prev],p]
prev = i
# if load: multiply by discharge
if calcload:
p,et,q = plot_data.get_petq()
q = q * f_conv
dates = data['Date'].map(lambda x: x.replace(hour=0,minute=0,second=0))
q_date = q['Gemaal Heye'][dates].fillna(0)
q_date.index = dates.index
cols = [(e,p) for p in pct for e in params]
data[cols] = data[cols].mul(q_date, 'index')
# x limits
x1 = 0
x2 = -1
if xlim != None:
for i, xi in enumerate(x):
if x1 == 0 and xi > xlim[0]:
x1 = i
if x2 == -1 and xi > xlim[1]:
x2 = i
if figsize is None:
figheight = int(plot_peq)*5
if separate:
if len(params) < 6:
figheight += len(params)*5
else:
figheight += len(params)*(25.+figheight)/len(params)
else: figheight += 10
figsize = (20,figheight)
fig = plt.figure(facecolor='w',figsize=figsize)
ax = []
box = _constr_box(plot_peq, separate, params)
if plot_peq:
ax += plot_data.create_petq(afig=fig, box=box.next())
ax[0].set_xlim(x[x1],x[x2])
ax[1].set_xlim(x[x1],x[x2])
_plot_text(ax[0],'a) net precipitation [mm/d]', y=0.9)
_plot_text(ax[1],'b) discharge [mm/d]',y=.94)
ax += [fig.add_axes(box.next(), sharex=ax[0])]
peqplots=2
else:
ax += [fig.add_axes(box.next())]
peqplots=0
for i,p in enumerate(params):
if separate and i:
ax += [fig.add_axes(box.next(), sharex=ax[0])]
ax[-1].plot(x[x1:x2], data[p,pct[1]][x1:x2], color=col[p], label=labels[p],
linewidth=2, zorder=0)
if stacked:
if i:
ax[-1].fill_between(x[x1:x2], data[params[i-1],pct[1]][x1:x2], data[p,pct[1]][x1:x2],
facecolor=col[p], alpha=0.25, zorder=0)
else:
ax[-1].fill_between(x[x1:x2], 0, data[p,pct[1]][x1:x2],
facecolor=col[p], alpha=0.25, zorder=0)
else:
ax[-1].fill_between(x[x1:x2], data[p,pct[0]][x1:x2], data[p,pct[2]][x1:x2],
facecolor=col[p], alpha=0.25, zorder=0)
if len(pct) > 3:
for j in xrange(3, len(pct), 2):
ax[-1].fill_between(x[x1:x2], data[p,pct[j]][x1:x2], data[p,pct[j+1]][x1:x2],
facecolor=col[p], alpha=0.25, zorder=0)
ax[-1].set_xlim(x[x1],x[x2])
ax[-1].set_xbound(x[x1],x[x2])
ax[-1].grid()
if ylim != None:
ax[-1].set_ylim(ylim[i])
ax[-1].set_ybound(ylim[i])
labstr = string.lowercase[i+peqplots]+')'
if labelstr != None:
labstr += ' '+labelstr.strip()+' '+labels[p]
if unitstr != None:
labstr += ' ['+unitstr.strip()+']'
_plot_text(ax[-1], labstr)
xd = []
for yr in xrange(x[x1].year, x[x2].year+1):
for m in xrange(1,13):
d = dt.datetime(yr,m,1)
if d >= x[x1] and d <= x[x2]:
xd += [d]
ticklabels = [d.strftime('%b-%Y') for d in xd]
for i,cax in enumerate(ax):
cax.set_xticks(xd)
if i:
cax.yaxis.set_major_locator(MaxNLocator(nbins=4, prune='upper'))
else:
cax.yaxis.set_major_locator(MaxNLocator(nbins=4))
cax.fmt_xdata = mdates.DateFormatter('%d-%m-%Y %H:%M:%S')
# draw measurement hiatuses, not on precip/q plots
if hiatuses and i>=peqplots:
for h in hiatuses:
cax.axvspan(h[0], h[1], facecolor='w', linewidth=0, zorder=1)#'lightgrey'
if cax == ax[-1]:
# only bottom subplots:
cax.set_xticklabels(labels=ticklabels)#, rotation='vertical', visible=True)
else:
plt.setp(cax.get_xticklabels(),visible=False)
plt.setp(cax.get_xticklabels(), fontsize=16)
plt.setp(cax.get_yticklabels(), fontsize=16)
#fig.suptitle('Endmember fractions', fontsize=18)
if not returnfig:
if (fname==None):
plt.show()
plt.close()
else:
plt.savefig(fname, dpi=96)
plt.close()
else:
return fig,ax
[docs]def endmembers_time(pct=(0,50,100),result='gemma_output.bin',
separate=False, stacked=True, fracload='frac', ylim=None,
col=None, labels=None, labelstr=None, unitstr=None,
f_conv=86400. * 1000. / 10e6, **kwargs):
'''Convenience function for plotting endmembers over time,
including uncertainty. Params and param dicts are set to end-members,
labelstr and unitstr are also set.
Parameters
----------
fracload : plot either fractions plot or load plot (default 'frac')
others : see params_time
Examples
--------
.. plot:: pyplots/plot_em_time.py
'''
if pct[-1] > 1:
pct = [x/100. for x in pct]
if separate:
stacked = False
if not isinstance(result, gemma):
result = gemma(result)
if col is None:
if stacked:
col = create_colcolor(result.endmembers)
else:
col = create_colblack(result.endmembers)
if labels is None:
labels=create_labels(result.endmembers)
if labelstr is None:
if fracload == 'frac':
labelstr = 'fraction'
else:
labelstr = 'discharge'
if unitstr is None:
if fracload == 'frac':
unitstr = '-'
else:
unitstr = 'mm/d'
if fracload == 'frac':
ylim = (0,1)
calcload=False
else:
calcload=True
# present end-members in order
order = ['GWAQD','BOIL','GWAQS','GWDB','GWPH','INLET','PRECIP']
if set(result.endmembers).issubset(order):
order = [p for p in order if p in result.endmembers]
return params_time(order,pct=pct,result=result,separate=separate,
stacked=stacked,calcload=calcload,col=col,ylim=ylim,
labels=labels,labelstr=labelstr,unitstr=unitstr,f_conv=f_conv,
**kwargs)
[docs]def plot_em_solutes(em, pct=(0,50,100), result='gemma_output.bin',solutes=None,
plot_peq=False,calcload=False, hiatuses=[], xlim=None,
ylim=None,fname=None):
'''Convenience function to plot all solutes of a given endmember.
including uncertainty. Params and param dicts are set to end-members,
labelstr and unitstr are also set.
Parameters
----------
em : end-member
others : see params_time
'''
if not isinstance(result, gemma):
result = gemma(result)
if solutes == None:
solutes=result.solutes
params = [em+'_'+s for s in solutes]
if not calcload:
labelstr = 'concentration'
unitstr = 'mg/l'
f_conv=1.
else:
labelstr = 'load'
unitstr = 'g/d'
f_conv=86400. #mg/l*m3/s
labels = create_labels(params)
col = create_colblack(params)
params_time(params, pct=pct,result=result, separate=True,
hiatuses=hiatuses,plot_peq=plot_peq,xlim=xlim,ylim=ylim,
col=col,labels=labels,labelstr=labelstr,unitstr=unitstr,
f_conv=f_conv,fname=fname)
[docs]def plot_calc_meas(fname_stream='gemma_stream.txt',pct=(0,50,100),
result='gemma_output.bin',hiatuses=[],plot_peq=True,
fname=None,returnfig=None, col=None,labels=None, **kwargs):
'''Function to plot calculated concentrations against measured concentrations
Parameters
----------
fname_stream : filename of g-emma stream input file
pct : list-like of percentiles (default: 0,50,100)
result : g-emma result, either filename or gemma object
hiatuses : list of lists of two datetime objects. Draws vertical grey bands to exclude period from figure
fname : filename to write plot to. If None (default), output to screen
returnfig : bool, return figure object and list of axes
figsize : size of figure
kwargs : additional parameters to pass to params_time
Examples
--------
.. plot:: pyplots/plot_calc_meas.py
'''
if not isinstance(result, gemma):
result = gemma(result)
stream = read_gemma_stream(fname_stream)
params = ['MIX_'+s for s in result.solutes]
if col is None: col = create_colblack(params)
if labels is None: labels = create_labels(params)
fig,ax = params_time(params,pct,result=result,labelstr='concentration',
hiatuses=hiatuses,plot_peq=plot_peq,col=col,
labels=labels,unitstr='mg/l',returnfig=True,**kwargs)
if plot_peq: add = 2
else: add=0
for i,s in enumerate(result.solutes):
ax[i+add].errorbar(stream['Datetime'],stream[s],stream[s+'_sd'],fmt='.',c='k')
if not returnfig:
if (fname==None):
plt.show()
plt.close()
else:
plt.savefig(fname, dpi=96)
plt.close()
else:
return fig,ax
